From 23e16501e0ea914496a0f4c57310c4252f5c607c Mon Sep 17 00:00:00 2001 From: Rory Mitchell Date: Tue, 8 Feb 2022 02:39:09 +0100 Subject: [PATCH] Fix integer overflow in distances (#490) Fix for https://github.com/rapidsai/cuml/issues/4552. Authors: - Rory Mitchell (https://github.com/RAMitchell) Approvers: - Corey J. Nolet (https://github.com/cjnolet) URL: https://github.com/rapidsai/raft/pull/490 --- .../detail/pairwise_distance_base.cuh | 5 ++- cpp/test/distance/dist_canberra.cu | 6 ++- cpp/test/distance/dist_chebyshev.cu | 6 ++- cpp/test/distance/dist_correlation.cu | 6 ++- cpp/test/distance/dist_cos.cu | 6 ++- cpp/test/distance/dist_euc_exp.cu | 5 ++- cpp/test/distance/dist_euc_unexp.cu | 5 ++- cpp/test/distance/dist_hamming.cu | 6 ++- cpp/test/distance/dist_hellinger.cu | 6 ++- cpp/test/distance/dist_jensen_shannon.cu | 6 ++- cpp/test/distance/dist_kl_divergence.cu | 6 ++- cpp/test/distance/dist_l1.cu | 6 ++- cpp/test/distance/dist_minkowski.cu | 5 ++- cpp/test/distance/dist_russell_rao.cu | 6 ++- cpp/test/distance/distance_base.cuh | 37 ++++++++++++++++++- 15 files changed, 101 insertions(+), 16 deletions(-) diff --git a/cpp/include/raft/distance/detail/pairwise_distance_base.cuh b/cpp/include/raft/distance/detail/pairwise_distance_base.cuh index bfca731443..996cc544a6 100644 --- a/cpp/include/raft/distance/detail/pairwise_distance_base.cuh +++ b/cpp/include/raft/distance/detail/pairwise_distance_base.cuh @@ -1,5 +1,5 @@ /* - * Copyright (c) 2021, NVIDIA CORPORATION. + * Copyright (c) 2022, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -276,7 +276,8 @@ struct PairwiseDistances : public BaseClass { for (int j = 0; j < P::AccColsPerTh; ++j) { auto colId = startx + j * P::AccThCols; if (rowId < this->m && colId < this->n) { - dOutput[rowId * this->n + colId] = fin_op(acc[i][j], 0); + // Promote to 64 bit index for final write, as output array can be > 2^31 + dOutput[std::size_t(rowId) * this->n + colId] = fin_op(acc[i][j], 0); } } } diff --git a/cpp/test/distance/dist_canberra.cu b/cpp/test/distance/dist_canberra.cu index ca90907779..1f368fbee8 100644 --- a/cpp/test/distance/dist_canberra.cu +++ b/cpp/test/distance/dist_canberra.cu @@ -1,5 +1,5 @@ /* - * Copyright (c) 2018-2021, NVIDIA CORPORATION. + * Copyright (c) 2018-2022, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -64,5 +64,9 @@ TEST_P(DistanceCanberraD, Result) } INSTANTIATE_TEST_CASE_P(DistanceTests, DistanceCanberraD, ::testing::ValuesIn(inputsd)); +class BigMatrixCanberra : public BigMatrixDistanceTest { +}; +TEST_F(BigMatrixCanberra, Result) {} + } // end namespace distance } // end namespace raft diff --git a/cpp/test/distance/dist_chebyshev.cu b/cpp/test/distance/dist_chebyshev.cu index 641b958d72..8f506601ca 100644 --- a/cpp/test/distance/dist_chebyshev.cu +++ b/cpp/test/distance/dist_chebyshev.cu @@ -1,5 +1,5 @@ /* - * Copyright (c) 2018-2021, NVIDIA CORPORATION. + * Copyright (c) 2018-2022, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -64,5 +64,9 @@ TEST_P(DistanceLinfD, Result) } INSTANTIATE_TEST_CASE_P(DistanceTests, DistanceLinfD, ::testing::ValuesIn(inputsd)); +class BigMatrixLinf : public BigMatrixDistanceTest { +}; +TEST_F(BigMatrixLinf, Result) {} + } // end namespace distance } // end namespace raft diff --git a/cpp/test/distance/dist_correlation.cu b/cpp/test/distance/dist_correlation.cu index 72df5b10f4..77d770b4d1 100644 --- a/cpp/test/distance/dist_correlation.cu +++ b/cpp/test/distance/dist_correlation.cu @@ -1,5 +1,5 @@ /* - * Copyright (c) 2021, NVIDIA CORPORATION. + * Copyright (c) 2021-2022, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -65,5 +65,9 @@ TEST_P(DistanceCorrelationD, Result) } INSTANTIATE_TEST_CASE_P(DistanceTests, DistanceCorrelationD, ::testing::ValuesIn(inputsd)); +class BigMatrixCorrelation + : public BigMatrixDistanceTest { +}; +TEST_F(BigMatrixCorrelation, Result) {} } // end namespace distance } // end namespace raft diff --git a/cpp/test/distance/dist_cos.cu b/cpp/test/distance/dist_cos.cu index a085e82705..900a71e514 100644 --- a/cpp/test/distance/dist_cos.cu +++ b/cpp/test/distance/dist_cos.cu @@ -1,5 +1,5 @@ /* - * Copyright (c) 2018-2021, NVIDIA CORPORATION. + * Copyright (c) 2018-2022, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -64,5 +64,9 @@ TEST_P(DistanceExpCosD, Result) } INSTANTIATE_TEST_CASE_P(DistanceTests, DistanceExpCosD, ::testing::ValuesIn(inputsd)); +class BigMatrixCos : public BigMatrixDistanceTest { +}; +TEST_F(BigMatrixCos, Result) {} + } // end namespace distance } // end namespace raft diff --git a/cpp/test/distance/dist_euc_exp.cu b/cpp/test/distance/dist_euc_exp.cu index f840a91bec..ff142da7fa 100644 --- a/cpp/test/distance/dist_euc_exp.cu +++ b/cpp/test/distance/dist_euc_exp.cu @@ -1,5 +1,5 @@ /* - * Copyright (c) 2018-2021, NVIDIA CORPORATION. + * Copyright (c) 2018-2022, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -64,5 +64,8 @@ TEST_P(DistanceEucExpTestD, Result) } INSTANTIATE_TEST_CASE_P(DistanceTests, DistanceEucExpTestD, ::testing::ValuesIn(inputsd)); +class BigMatrixEucExp : public BigMatrixDistanceTest { +}; +TEST_F(BigMatrixEucExp, Result) {} } // end namespace distance } // end namespace raft diff --git a/cpp/test/distance/dist_euc_unexp.cu b/cpp/test/distance/dist_euc_unexp.cu index 6d374f3332..81e6be7116 100644 --- a/cpp/test/distance/dist_euc_unexp.cu +++ b/cpp/test/distance/dist_euc_unexp.cu @@ -1,5 +1,5 @@ /* - * Copyright (c) 2018-2021, NVIDIA CORPORATION. + * Copyright (c) 2018-2022, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -65,5 +65,8 @@ TEST_P(DistanceEucUnexpTestD, Result) } INSTANTIATE_TEST_CASE_P(DistanceTests, DistanceEucUnexpTestD, ::testing::ValuesIn(inputsd)); +class BigMatrixEucUnexp : public BigMatrixDistanceTest { +}; +TEST_F(BigMatrixEucUnexp, Result) {} } // end namespace distance } // end namespace raft diff --git a/cpp/test/distance/dist_hamming.cu b/cpp/test/distance/dist_hamming.cu index e0f1efc3f7..616ce8f729 100644 --- a/cpp/test/distance/dist_hamming.cu +++ b/cpp/test/distance/dist_hamming.cu @@ -1,5 +1,5 @@ /* - * Copyright (c) 2018-2021, NVIDIA CORPORATION. + * Copyright (c) 2018-2022, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -65,5 +65,9 @@ TEST_P(DistanceHammingD, Result) } INSTANTIATE_TEST_CASE_P(DistanceTests, DistanceHammingD, ::testing::ValuesIn(inputsd)); +class BigMatrixHamming + : public BigMatrixDistanceTest { +}; +TEST_F(BigMatrixHamming, Result) {} } // end namespace distance } // end namespace raft diff --git a/cpp/test/distance/dist_hellinger.cu b/cpp/test/distance/dist_hellinger.cu index caa96f189d..d6f994aaf6 100644 --- a/cpp/test/distance/dist_hellinger.cu +++ b/cpp/test/distance/dist_hellinger.cu @@ -1,5 +1,5 @@ /* - * Copyright (c) 2021, NVIDIA CORPORATION. + * Copyright (c) 2022, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -65,5 +65,9 @@ TEST_P(DistanceHellingerExpD, Result) } INSTANTIATE_TEST_CASE_P(DistanceTests, DistanceHellingerExpD, ::testing::ValuesIn(inputsd)); +class BigMatrixHellingerExp + : public BigMatrixDistanceTest { +}; +TEST_F(BigMatrixHellingerExp, Result) {} } // end namespace distance } // end namespace raft diff --git a/cpp/test/distance/dist_jensen_shannon.cu b/cpp/test/distance/dist_jensen_shannon.cu index 74b02ef18d..43e4f3aa0f 100644 --- a/cpp/test/distance/dist_jensen_shannon.cu +++ b/cpp/test/distance/dist_jensen_shannon.cu @@ -1,5 +1,5 @@ /* - * Copyright (c) 2021, NVIDIA CORPORATION. + * Copyright (c) 2022, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -65,5 +65,9 @@ TEST_P(DistanceJensenShannonD, Result) } INSTANTIATE_TEST_CASE_P(DistanceTests, DistanceJensenShannonD, ::testing::ValuesIn(inputsd)); +class BigMatrixJensenShannon + : public BigMatrixDistanceTest { +}; +TEST_F(BigMatrixJensenShannon, Result) {} } // end namespace distance } // end namespace raft diff --git a/cpp/test/distance/dist_kl_divergence.cu b/cpp/test/distance/dist_kl_divergence.cu index e551eda0ab..6a5fe8d7ac 100644 --- a/cpp/test/distance/dist_kl_divergence.cu +++ b/cpp/test/distance/dist_kl_divergence.cu @@ -1,5 +1,5 @@ /* - * Copyright (c) 2021, NVIDIA CORPORATION. + * Copyright (c) 2022, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -65,5 +65,9 @@ TEST_P(DistanceKLDivergenceD, Result) } INSTANTIATE_TEST_CASE_P(DistanceTests, DistanceKLDivergenceD, ::testing::ValuesIn(inputsd)); +class BigMatrixKLDivergence + : public BigMatrixDistanceTest { +}; +TEST_F(BigMatrixKLDivergence, Result) {} } // end namespace distance } // end namespace raft diff --git a/cpp/test/distance/dist_l1.cu b/cpp/test/distance/dist_l1.cu index ac2ee024f6..322fb52d5c 100644 --- a/cpp/test/distance/dist_l1.cu +++ b/cpp/test/distance/dist_l1.cu @@ -1,5 +1,5 @@ /* - * Copyright (c) 2018-2021, NVIDIA CORPORATION. + * Copyright (c) 2018-2022, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -64,5 +64,9 @@ TEST_P(DistanceUnexpL1D, Result) } INSTANTIATE_TEST_CASE_P(DistanceTests, DistanceUnexpL1D, ::testing::ValuesIn(inputsd)); +class BigMatrixUnexpL1 : public BigMatrixDistanceTest { +}; +TEST_F(BigMatrixUnexpL1, Result) {} + } // end namespace distance } // end namespace raft diff --git a/cpp/test/distance/dist_minkowski.cu b/cpp/test/distance/dist_minkowski.cu index f0a6833f2b..3e0a2ead92 100644 --- a/cpp/test/distance/dist_minkowski.cu +++ b/cpp/test/distance/dist_minkowski.cu @@ -1,5 +1,5 @@ /* - * Copyright (c) 2018-2021, NVIDIA CORPORATION. + * Copyright (c) 2018-2022, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -64,5 +64,8 @@ TEST_P(DistanceLpUnexpD, Result) } INSTANTIATE_TEST_CASE_P(DistanceTests, DistanceLpUnexpD, ::testing::ValuesIn(inputsd)); +class BigMatrixLpUnexp : public BigMatrixDistanceTest { +}; +TEST_F(BigMatrixLpUnexp, Result) {} } // end namespace distance } // end namespace raft diff --git a/cpp/test/distance/dist_russell_rao.cu b/cpp/test/distance/dist_russell_rao.cu index 42234d4f0b..e92a01c70a 100644 --- a/cpp/test/distance/dist_russell_rao.cu +++ b/cpp/test/distance/dist_russell_rao.cu @@ -1,5 +1,5 @@ /* - * Copyright (c) 2021, NVIDIA CORPORATION. + * Copyright (c) 2021-2022, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -65,5 +65,9 @@ TEST_P(DistanceRussellRaoD, Result) } INSTANTIATE_TEST_CASE_P(DistanceTests, DistanceRussellRaoD, ::testing::ValuesIn(inputsd)); +class BigMatrixRussellRao + : public BigMatrixDistanceTest { +}; +TEST_F(BigMatrixRussellRao, Result) {} } // end namespace distance } // end namespace raft diff --git a/cpp/test/distance/distance_base.cuh b/cpp/test/distance/distance_base.cuh index 8f0de29eed..4c9d5b11cc 100644 --- a/cpp/test/distance/distance_base.cuh +++ b/cpp/test/distance/distance_base.cuh @@ -1,5 +1,5 @@ /* - * Copyright (c) 2018-2021, NVIDIA CORPORATION. + * Copyright (c) 2018-2022, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -476,5 +476,40 @@ class DistanceTest : public ::testing::TestWithParam> { rmm::device_uvector x, y, dist_ref, dist, dist2; }; +template +class BigMatrixDistanceTest : public ::testing::Test { + public: + BigMatrixDistanceTest() + : x(m * k, handle.get_stream()), dist(std::size_t(m) * m, handle.get_stream()){}; + void SetUp() override + { + auto testInfo = testing::UnitTest::GetInstance()->current_test_info(); + common::nvtx::range fun_scope("test::%s/%s", testInfo->test_suite_name(), testInfo->name()); + + size_t worksize = raft::distance::getWorkspaceSize( + x.data(), x.data(), m, n, k); + rmm::device_uvector workspace(worksize, handle.get_stream()); + raft::distance::distance(x.data(), + x.data(), + dist.data(), + m, + n, + k, + workspace.data(), + worksize, + handle.get_stream(), + true, + 0.0f); + + RAFT_CUDA_TRY(cudaStreamSynchronize(handle.get_stream())); + } + + protected: + int m = 48000; + int n = 48000; + int k = 1; + raft::handle_t handle; + rmm::device_uvector x, dist; +}; } // end namespace distance } // end namespace raft