Harden OPG Data byte-size computation against overflow

cpp/bench/common/ml_benchmark.hpp

@@ -7,6 +7,7 @@
 
 #include <cuml/common/logger.hpp>
 #include <cuml/common/utils.hpp>
+#include <cuml/common/checked_arithmetic.hpp>
 
 #include <raft/util/cudart_utils.hpp>
 
@@ -158,7 +159,7 @@ class Fixture : public ::benchmark::Fixture {
   template <typename T>
   void alloc(T*& ptr, size_t len, bool init = false)
   {
-    auto nBytes  = len * sizeof(T);
+    auto nBytes = ML::checked_mul<size_t>(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<size_t>(len, sizeof(T));
+    d_alloc.deallocate(stream, ptr, nBytes);
   }
 
   cudaStream_t stream = 0;

cpp/include/cuml/prims/opg/matrix/data.hpp

@@ -5,39 +5,58 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
+#include <cstddef>
+
+#include <cuml/common/checked_arithmetic.hpp>
 #include <cuml/common/export.hpp>
 
 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 <typename Type>
-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 <typename Type>
+    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<size_t>(_n_elements, sizeof(Type)))
+      {
+      }
+
+      void setNumElements(size_t _n_elements)
+      {
+        nElements = _n_elements;
+        totalSize = ML::checked_mul<size_t>(_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<float> floatData_t;
 typedef Data<double> doubleData_t;

cpp/src/glm/ridge_mg.cu

@@ -72,7 +72,7 @@ void ridgeSolve(const raft::handle_t& handle,
     raft::make_device_vector_view<const T, std::size_t>(S, UDesc.N));
 
   MLCommon::Matrix::Data<T> 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<T> 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);

cpp/src/solver/cd_mg.cu

@@ -116,11 +116,11 @@ int fit_impl(raft::handle_t& handle,
 
     MLCommon::Matrix::Data<T>* rs_data = new MLCommon::Matrix::Data<T>();
     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<T>* temp_data = new MLCommon::Matrix::Data<T>();
-    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);
 

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<T> 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<T> 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);
 

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)

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 <cuml/prims/opg/matrix/data.hpp>
+
+#include <raft/core/error.hpp>
+
+#include <gtest/gtest.h>
+
+#include <cstddef>
+#include <limits>
+
+namespace MLCommon {
+namespace Matrix {
+
+TEST(MatrixData, ComputesBytesAndElements)
+{
+  float value = 1.0f;
+  Data<float> 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<size_t>::max() / sizeof(float);
+
+  EXPECT_THROW(Data<float> data(ptr, max_elements + 1), raft::exception);
+  Data<float> data(ptr, size_t(0));
+  EXPECT_THROW(data.setNumElements(max_elements + 1), raft::exception);
+}
+
+}  // namespace Matrix
+}  // namespace MLCommon

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 <raft/core/error.hpp>
+
+#include <gtest/gtest.h>
+
+#include <benchmark/benchmark.h>
+
+#include <cstddef>
+#include <limits>
+
+namespace MLCommon {
+namespace Bench {
+
+class TestFixture : public Fixture {
+ public:
+  TestFixture() : Fixture("MLBenchmarkFixtureTest") {}
+
+  void runBenchmark(::benchmark::State&) override {}
+
+  template <typename T>
+  void testAlloc(T*& ptr, size_t len, bool init = false)
+  {
+    alloc(ptr, len, init);
+  }
+
+  template <typename T>
+  void testDealloc(T* ptr, size_t len)
+  {
+    dealloc(ptr, len);
+  }
+};
+
+TEST(MlBenchmarkFixtureAllocator, ThrowOnHugeAllocationOrDeallocationLength)
+{
+  TestFixture fixture;
+
+  int* ptr      = nullptr;
+  auto const len = std::numeric_limits<size_t>::max() / sizeof(int) + 1;
+
+  EXPECT_THROW(fixture.testAlloc(ptr, len), raft::exception);
+  EXPECT_THROW(fixture.testDealloc(ptr, len), raft::exception);
+}
+
+}  // namespace Bench
+}  // namespace MLCommon