Validate CUDA NMS mask size

[Akshay Sonawane (apsonawane)]

This pull request improves the robustness and error handling of the CUDA NonMaxSuppression (NMS) implementation in ONNX Runtime. The changes ensure that large input sizes which could exceed the range of supported integer indices are properly checked and handled, preventing potential overflows or crashes. Additionally, a new CUDA-specific test has been added to verify this behavior.

Error handling and robustness improvements:

• Updated the calculation of max_nms_mask_size in non_max_suppression_impl.cu to use SafeInt<size_t> for safer arithmetic and added a check to ensure the mask size does not exceed the maximum value of an int. If the limit is exceeded, an error is returned. (onnxruntime/core/providers/cuda/object_detection/non_max_suppression_impl.cu)
• Included necessary headers for safe integer operations and numeric limits (<limits>, core/common/safeint.h>) to support the above changes. (onnxruntime/core/providers/cuda/object_detection/non_max_suppression_impl.cu)

Testing enhancements:

• Added a CUDA-specific unit test that verifies the operator correctly rejects inputs that would require a mask size outside the supported integer range, ensuring the new error handling works as intended. (onnxruntime/test/providers/cpu/object_detection/non_max_suppression_test.cc)
• Included conditional compilation and necessary CUDA test utilities to enable the new test only when CUDA is available. (onnxruntime/test/providers/cpu/object_detection/non_max_suppression_test.cc)

[Copilot]

Pull request overview

This pull request hardens the CUDA NonMaxSuppression (NMS) implementation against integer overflow by validating the computed NMS mask size before allocating/zeroing the mask buffer, and adds a CUDA-scoped unit test to ensure oversized mask requests are rejected with an expected error.

Changes:

• Compute max_nms_mask_size using SafeInt<size_t> and fail fast when the mask size would exceed int indexability.
• Replace the ceil(float(...))-based grid size computation for SetZero with integer math and pass an int-validated element count into the kernel.
• Add a CUDA-only unit test that constructs an input size where num_boxes * bit_mask_len crosses INT_MAX, asserting the new failure behavior.

Reviewed changes

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

File	Description
onnxruntime/core/providers/cuda/object_detection/non_max_suppression_impl.cu	Adds safe arithmetic + explicit int-range validation for the NMS mask sizing/allocation/zeroing path to prevent overflow-driven misallocation.
onnxruntime/test/providers/cpu/object_detection/non_max_suppression_test.cc	Adds a USE_CUDA-guarded test that runs NMS on CUDA and asserts failure when the required mask size exceeds the supported int range.

---

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