diff --git a/cpp/bench/common/ml_benchmark.hpp b/cpp/bench/common/ml_benchmark.hpp index ebe2d5ef49..5723ec2429 100644 --- a/cpp/bench/common/ml_benchmark.hpp +++ b/cpp/bench/common/ml_benchmark.hpp @@ -7,6 +7,7 @@ #include #include +#include #include @@ -158,7 +159,7 @@ class Fixture : public ::benchmark::Fixture { template void alloc(T*& ptr, size_t len, bool init = false) { - auto nBytes = len * sizeof(T); + auto nBytes = ML::checked_mul(len, sizeof(T)); auto d_alloc = rmm::mr::get_current_device_resource_ref(); ptr = (T*)d_alloc.allocate(stream, nBytes); if (init) { RAFT_CUDA_TRY(cudaMemsetAsync(ptr, 0, nBytes, stream)); } @@ -168,7 +169,8 @@ class Fixture : public ::benchmark::Fixture { void dealloc(T* ptr, size_t len) { auto d_alloc = rmm::mr::get_current_device_resource_ref(); - d_alloc.deallocate(stream, ptr, len * sizeof(T)); + auto nBytes = ML::checked_mul(len, sizeof(T)); + d_alloc.deallocate(stream, ptr, nBytes); } cudaStream_t stream = 0; diff --git a/cpp/include/cuml/prims/opg/matrix/data.hpp b/cpp/include/cuml/prims/opg/matrix/data.hpp index 93edc4529e..d91734aae7 100644 --- a/cpp/include/cuml/prims/opg/matrix/data.hpp +++ b/cpp/include/cuml/prims/opg/matrix/data.hpp @@ -5,39 +5,58 @@ * SPDX-License-Identifier: Apache-2.0 */ +#include + +#include #include namespace CUML_EXPORT MLCommon { namespace Matrix { -/** - * @brief This is a *helper* wrapper around the multi-gpu data blocks owned - * by a worker. It's design is NOT final. Its so written this way to get - * something concrete in a short span of time. - * @todo add support for custom allocators - */ -template -struct Data { - Data() : ptr(nullptr), totalSize(0) {} - Data(Type* _ptr, size_t _n_elements) : ptr(_ptr), totalSize(_n_elements * sizeof(Type)) {} - - /** - * actual data block. This is just a linearly laid out buffer of all blocks - * owned by this worker - */ - Type* ptr = nullptr; - - /** - * total size (in bytes) of this buffer. In future, this will be passed - * to the dealloc function underneath - */ - size_t totalSize = (size_t)0; - - /** - * Return the number of elements of Type in ptr. - */ - size_t numElements() const { return totalSize / sizeof(Type); } -}; + /** + * @brief This is a *helper* wrapper around the multi-gpu data blocks owned + * by a worker. It's design is NOT final. Its so written this way to get + * something concrete in a short span of time. + * @todo add support for custom allocators + */ + template + struct Data { + Data() : ptr(nullptr), nElements(0), totalSize(0) {} + Data(Type* _ptr, size_t _n_elements) + : ptr(_ptr), + nElements(_n_elements), + totalSize(ML::checked_mul(_n_elements, sizeof(Type))) + { + } + + void setNumElements(size_t _n_elements) + { + nElements = _n_elements; + totalSize = ML::checked_mul(_n_elements, sizeof(Type)); + } + + /** + * actual data block. This is just a linearly laid out buffer of all blocks + * owned by this worker + */ + Type* ptr = nullptr; + + /** + * number of elements in this buffer. + */ + size_t nElements = 0; + + /** + * total size (in bytes) of this buffer. In future, this will be passed + * to the dealloc function underneath + */ + size_t totalSize = (size_t)0; + + /** + * Return the number of elements of Type in ptr. + */ + size_t numElements() const { return nElements; } + }; typedef Data floatData_t; typedef Data doubleData_t; diff --git a/cpp/src/glm/ridge_mg.cu b/cpp/src/glm/ridge_mg.cu index 9ed62b7b29..f680e1eba6 100644 --- a/cpp/src/glm/ridge_mg.cu +++ b/cpp/src/glm/ridge_mg.cu @@ -72,7 +72,7 @@ void ridgeSolve(const raft::handle_t& handle, raft::make_device_vector_view(S, UDesc.N)); MLCommon::Matrix::Data S_nnz_data; - S_nnz_data.totalSize = UDesc.N; + S_nnz_data.setNumElements(UDesc.N); S_nnz_data.ptr = S_nnz; MLCommon::LinAlg::opg::mv_aTb(handle, S_nnz_data, U, UDesc, b, streams, n_streams); @@ -125,7 +125,7 @@ void ridgeEig(raft::handle_t& handle, for (std::size_t i = 0; i < partsToRanks.size(); i++) { MLCommon::Matrix::Data d; - d.totalSize = partsToRanks[i]->size; + d.setNumElements(partsToRanks[i]->size); d.ptr = curr_ptr; curr_ptr = curr_ptr + (partsToRanks[i]->size * ADesc.N); U_temp.push_back(d); diff --git a/cpp/src/solver/cd_mg.cu b/cpp/src/solver/cd_mg.cu index 2be635bebf..718eeabf43 100644 --- a/cpp/src/solver/cd_mg.cu +++ b/cpp/src/solver/cd_mg.cu @@ -116,11 +116,11 @@ int fit_impl(raft::handle_t& handle, MLCommon::Matrix::Data* rs_data = new MLCommon::Matrix::Data(); rs_data->ptr = rs; - rs_data->totalSize = partsToRanks[i]->size; + rs_data->setNumElements(partsToRanks[i]->size); residual_temp.push_back(rs_data); MLCommon::Matrix::Data* temp_data = new MLCommon::Matrix::Data(); - temp_data->totalSize = partsToRanks[i]->size; + temp_data->setNumElements(partsToRanks[i]->size); input_data_temp.push_back(temp_data); rs += partsToRanks[i]->size; @@ -156,7 +156,7 @@ int fit_impl(raft::handle_t& handle, input_col_loc = input_data[k]->ptr + (ci * partsToRanks[k]->size); input_data_temp[k]->ptr = input_col_loc; - input_data_temp[k]->totalSize = partsToRanks[k]->size; + input_data_temp[k]->setNumElements(partsToRanks[k]->size); raft::linalg::multiplyScalar( pred_loc, input_col_loc, h_coef[ci], partsToRanks[k]->size, streams[k % n_streams]); @@ -173,7 +173,7 @@ int fit_impl(raft::handle_t& handle, } coef_loc_data.ptr = coef_loc; - coef_loc_data.totalSize = size_t(1); + coef_loc_data.setNumElements(size_t(1)); MLCommon::LinAlg::opg::mv_aTb( handle, coef_loc_data, input_data_temp, input_desc_temp, residual_temp, streams, n_streams); diff --git a/cpp/src_prims/opg/linalg/lstsq.cu b/cpp/src_prims/opg/linalg/lstsq.cu index 5f40f268e2..5283c15fc2 100644 --- a/cpp/src_prims/opg/linalg/lstsq.cu +++ b/cpp/src_prims/opg/linalg/lstsq.cu @@ -50,7 +50,7 @@ void lstsqEig_impl(const raft::handle_t& handle, for (size_t i = 0; i < partsToRanks.size(); i++) { Matrix::Data d; - d.totalSize = partsToRanks[i]->size; + d.setNumElements(partsToRanks[i]->size); d.ptr = curr_ptr; curr_ptr = curr_ptr + (partsToRanks[i]->size * ADesc.N); U_temp.push_back(d); @@ -66,7 +66,7 @@ void lstsqEig_impl(const raft::handle_t& handle, Matrix::Data w_out; w_out.ptr = tmp_vector.data(); - w_out.totalSize = ADesc.N; + w_out.setNumElements(ADesc.N); mv_aTb(handle, w_out, U, ADesc, b, streams, n_streams); diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt index 3bb0cc4325..534cefc54b 100644 --- a/cpp/tests/CMakeLists.txt +++ b/cpp/tests/CMakeLists.txt @@ -260,6 +260,8 @@ if(BUILD_PRIMS_TESTS) ConfigureTest(PREFIX PRIMS NAME LINEARREG_TEST prims/linearReg.cu) ConfigureTest(PREFIX PRIMS NAME LOG_TEST prims/log.cu) ConfigureTest(PREFIX PRIMS NAME LOGISTICREG_TEST prims/logisticReg.cu) + ConfigureTest(PREFIX PRIMS NAME ML_BENCHMARK_TEST prims/ml_benchmark.cpp ML_INCLUDE) + ConfigureTest(PREFIX PRIMS NAME MATRIX_DATA_TEST prims/matrix_data.cpp ML_INCLUDE) ConfigureTest(PREFIX PRIMS NAME MAKE_ARIMA_TEST prims/make_arima.cu) ConfigureTest(PREFIX PRIMS NAME PENALTY_TEST prims/penalty.cu) ConfigureTest(PREFIX PRIMS NAME SIGMOID_TEST prims/sigmoid.cu) diff --git a/cpp/tests/prims/matrix_data.cpp b/cpp/tests/prims/matrix_data.cpp new file mode 100644 index 0000000000..b5e6cc5b21 --- /dev/null +++ b/cpp/tests/prims/matrix_data.cpp @@ -0,0 +1,42 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include + +#include + +#include + +#include +#include + +namespace MLCommon { +namespace Matrix { + +TEST(MatrixData, ComputesBytesAndElements) +{ + float value = 1.0f; + Data data(&value, size_t(4)); + + EXPECT_EQ(data.numElements(), 4u); + EXPECT_EQ(data.totalSize, data.numElements() * sizeof(float)); + + data.setNumElements(2); + EXPECT_EQ(data.numElements(), 2u); + EXPECT_EQ(data.totalSize, data.numElements() * sizeof(float)); +} + +TEST(MatrixData, ThrowsOnElementCountOverflowForByteSize) +{ + float* ptr = nullptr; + auto max_elements = std::numeric_limits::max() / sizeof(float); + + EXPECT_THROW(Data data(ptr, max_elements + 1), raft::exception); + Data data(ptr, size_t(0)); + EXPECT_THROW(data.setNumElements(max_elements + 1), raft::exception); +} + +} // namespace Matrix +} // namespace MLCommon diff --git a/cpp/tests/prims/ml_benchmark.cpp b/cpp/tests/prims/ml_benchmark.cpp new file mode 100644 index 0000000000..532d6fdd38 --- /dev/null +++ b/cpp/tests/prims/ml_benchmark.cpp @@ -0,0 +1,51 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "../bench/common/ml_benchmark.hpp" + +#include + +#include + +#include + +#include +#include + +namespace MLCommon { +namespace Bench { + +class TestFixture : public Fixture { + public: + TestFixture() : Fixture("MLBenchmarkFixtureTest") {} + + void runBenchmark(::benchmark::State&) override {} + + template + void testAlloc(T*& ptr, size_t len, bool init = false) + { + alloc(ptr, len, init); + } + + template + void testDealloc(T* ptr, size_t len) + { + dealloc(ptr, len); + } +}; + +TEST(MlBenchmarkFixtureAllocator, ThrowOnHugeAllocationOrDeallocationLength) +{ + TestFixture fixture; + + int* ptr = nullptr; + auto const len = std::numeric_limits::max() / sizeof(int) + 1; + + EXPECT_THROW(fixture.testAlloc(ptr, len), raft::exception); + EXPECT_THROW(fixture.testDealloc(ptr, len), raft::exception); +} + +} // namespace Bench +} // namespace MLCommon