Avoid overflow in CUDA reduction scans

[Akshay Sonawane (apsonawane)]

This pull request refactors and improves the CUDA reduction functions to make index calculations safer and more robust, especially for large tensor sizes. The main changes introduce utility functions to avoid integer overflow and clarify the logic for stepping through elements during reductions. Additionally, new tests are added to ensure correct behavior at boundary conditions.

Reduction logic improvements:

• Added two new utility functions, reduction_scan_delta_is_valid and advance_reduction_scan, in reduction_utils.cuh to encapsulate index validation and advancement logic for reduction scans, reducing the risk of integer overflow and improving code clarity.
• Updated the main reduction loop in reduce_all and reduce_matrix_rows_kernel in reduction_functions.cu to use these new utility functions for safer and clearer iteration over elements and rows. [1] [2]

Bug fixes and correctness:

• Fixed a bug in reduce_matrix_rows_kernel where output writes could occur for invalid columns by adding a valid_col check before atomic writes.

Testing improvements:

• Added a new unit test, ScanIndexArithmeticAtIntMax, in reduction_functions_test.cc to verify correct behavior of the new utility functions at integer limit edge cases, ensuring robustness for very large tensors.
• Included the new utility header reduction_utils.cuh in the test file.

Code cleanup:

• Removed an unused local variable in reduce_matrix_rows_kernel for clarity.

[Copilot]

Pull request overview

Improves CUDA reduction scan arithmetic near INT_MAX to prevent index overflow and invalid output writes.

Changes:

• Adds overflow-safe scan helpers.
• Refactors scalar and matrix-row reduction loops.
• Adds boundary arithmetic tests.

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated no comments.

File	Description
reduction_functions.cu	Uses safe scan advancement and validates output columns.
reduction_utils.cuh	Adds scan validation and advancement helpers.
reduction_functions_test.cc	Tests scan arithmetic near INT_MAX.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Tianlei Wu (tianleiwu)]

The overflow-safe runtime loop changes look sound, but both Linux CUDA and TensorRT CI currently fail while compiling the new unit-test include. Details are inline.

[Tianlei Wu (tianleiwu)]

Reviewed the current head. The remaining-count checks avoid constructing overflowing absolute indices, and scan advancement only adds a step when the result stays below the limit. The block-uniform column loop keeps every thread participating in each synchronization while valid_col guards partial-block loads and writes. The host-compatible helper split resolves the earlier compile blocker, and all current CI checks pass. I do not see any remaining issues.