[CUDA] Add Validation of batch_indices in RoiAlign
Merged
TI
8mo ago+90+90+01234567890123456789-29-29-01234567890123456789
Description
This PR implements a device-side bounds check for batch_indices in the RoiAlign CUDA operator. This is a follow-up to https://github.com/microsoft/onnxruntime/pull/27543, which fixed the same vulnerability in the CPU implementation.
Previously, CheckROIAlignValidInput() only validated batch_indices when they were accessible on the host (CPU). For the CUDA EP, batch_indices reside in GPU memory, so host-side validation would require an expensive GPU-to-CPU copy, which could also break CUDA graph capture.
This change:
- Passes
batch_sizefrom the host to the CUDA kernel. - Adds a check within the
RoIAlignForwardkernel to ensure0 <= batch_index < batch_size. - If an invalid
batch_indexis encountered, the kernel sets the output value for that specific RoI element to 0 and returns early for that thread.
Impact
- Vulnerability fixed: Heap out-of-bounds read on GPU.
- Performance: Negligible impact as it's a simple range check within the existing kernel.
- Compatibility: No changes to ONNX models or public APIs.
Validation
- Existing
RoiAlignTestsuite. - Added two new test cases:
BatchIndicesOutOfRange_CUDAandBatchIndicesNegative_CUDAto verify that the CUDA provider correctly handles out-of-range batch indices. - Verified that the CUDA provider handles opset 10 without falling back to the CPU EP for these tests.