Harden OPG Data byte-size computation against overflow

[fallintoplace]

What changed

• Add direct include and use ML::checked_mul<size_t> when computing OPG matrix buffer byte size in cpp/include/cuml/prims/opg/matrix/data.hpp.
• Extend MLCommon::Matrix::Data to explicitly store element count in nElements, and keep numElements() returning that value.
• Add setNumElements() helper and migrate internal assignments to avoid writing totalSize directly with element counts.
• Add regression tests in cpp/tests/prims/matrix_data.cpp and register via cpp/tests/CMakeLists.txt.

Why

Data previously multiplied by sizeof(Type) without overflow checks, allowing silent wraparound on large block sizes. This also made byte-size vs element-count handling easy to get wrong; storing both avoids the confusion and makes unit size bookkeeping explicit.

Notes

I did not run tests in this pass.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: c88766e3-03d3-45f3-a93f-8675f3fe2c76

📥 Commits

Reviewing files that changed from the base of the PR and between 4a0991b and 6b94e65.

📒 Files selected for processing (1)

• cpp/tests/prims/matrix_data.cpp

🚧 Files skipped from review as they are similar to previous changes (1)

• cpp/tests/prims/matrix_data.cpp

---

📝 Walkthrough

Summary by CodeRabbit

• Bug Fixes

  • Improved safety when calculating matrix and allocation byte sizes with overflow-checked arithmetic.
  • Kept element counts and total memory sizes consistent across matrix operations, solvers, and benchmarks.
  • Oversized allocation and deallocation requests now fail safely instead of risking invalid memory calculations.

• Tests

  • Added coverage for matrix element-count and byte-size consistency.
  • Added tests confirming oversized benchmark allocation and deallocation requests are rejected.

Walkthrough

This PR adds overflow-checked byte-size calculations to Matrix::Data<Type> and the benchmark Fixture allocator. It adds explicit element-count tracking, updates solver callers, and adds tests for size consistency and overflow exceptions.

Changes

Checked-arithmetic sizing update

Layer / File(s)	Summary
Data struct rework: nElements storage and checked sizing cpp/include/cuml/prims/opg/matrix/data.hpp	Data<Type> stores nElements, computes totalSize with ML::checked_mul, adds setNumElements, and returns the stored count from numElements().
Caller migration to setNumElements cpp/src/glm/ridge_mg.cu, cpp/src/solver/cd_mg.cu, cpp/src_prims/opg/linalg/lstsq.cu	Solver code replaces direct totalSize assignments with setNumElements calls.
Benchmark Fixture checked allocation cpp/bench/common/ml_benchmark.hpp	Fixture::alloc and Fixture::dealloc use checked multiplication for byte counts.
Unit tests and build wiring cpp/tests/prims/matrix_data.cpp, cpp/tests/prims/ml_benchmark.cpp, cpp/tests/CMakeLists.txt	New tests cover element-count consistency and overflow exceptions. CMake registers both test executables.

Estimated code review effort: 2 (Simple) | ~15 minutes

Possibly related PRs

• rapidsai/cuml#8371: Also modifies Fixture allocation and deallocation behavior in ml_benchmark.hpp.

Suggested reviewers: dantegd, jcrist, jinsolp

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly and concisely describes the primary change: overflow protection for OPG Data byte-size computation.
Description check	✅ Passed	The description accurately explains the overflow fix, explicit element tracking, migrated assignments, and regression tests.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

Actionable comments posted: 1

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: e9ca7533-1d52-4d76-84a1-eac1cca07ef5

📥 Commits

Reviewing files that changed from the base of the PR and between 857cc5a and ad89c05.

📒 Files selected for processing (8)

• cpp/bench/common/ml_benchmark.hpp
• cpp/include/cuml/prims/opg/matrix/data.hpp
• cpp/src/glm/ridge_mg.cu
• cpp/src/solver/cd_mg.cu
• cpp/src_prims/opg/linalg/lstsq.cu
• cpp/tests/CMakeLists.txt
• cpp/tests/prims/matrix_data.cpp
• cpp/tests/prims/ml_benchmark.cpp

[chyunsu3]

/ok to test 32e43a3

[coderabbitai]

Actionable comments posted: 1

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@cpp/tests/prims/matrix_data.cpp`:
- Line 36: Declare the zero-length Data<float> instance before EXPECT_THROW,
then invoke setNumElements on that named object inside the macro so the braced
initializer comma does not split the macro arguments.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 7964be35-570b-46c7-b81b-5b261f659209

📥 Commits

Reviewing files that changed from the base of the PR and between e59a96e and 4a0991b.

📒 Files selected for processing (8)

• cpp/bench/common/ml_benchmark.hpp
• cpp/include/cuml/prims/opg/matrix/data.hpp
• cpp/src/glm/ridge_mg.cu
• cpp/src/solver/cd_mg.cu
• cpp/src_prims/opg/linalg/lstsq.cu
• cpp/tests/CMakeLists.txt
• cpp/tests/prims/matrix_data.cpp
• cpp/tests/prims/ml_benchmark.cpp

🚧 Files skipped from review as they are similar to previous changes (6)

• cpp/src/solver/cd_mg.cu
• cpp/src/glm/ridge_mg.cu
• cpp/bench/common/ml_benchmark.hpp
• cpp/tests/CMakeLists.txt
• cpp/src_prims/opg/linalg/lstsq.cu
• cpp/include/cuml/prims/opg/matrix/data.hpp