Validate CUDA GatherND indices bounds in release mode - #31646
Open
apsonawane wants to merge 6 commits into
Open
Conversation
Add host-side indices range validation in GatherND CUDA PrepareCompute for both CPU- and GPU-resident indices, matching CPU semantics. Add CUDA regression tests for invalid int64 and contrib int32 indices.
Contributor
There was a problem hiding this comment.
Pull request overview
This PR strengthens robustness of the CUDA GatherND implementation by validating indices bounds before launching CUDA kernels, and adds new tests intended to verify out-of-bounds indices are reported as user-facing errors (instead of causing invalid device accesses/crashes).
Changes:
- Added host-side index bounds validation in CUDA
GatherNDBase::PrepareCompute, including copying indices from device to host for validation. - Added new (CUDA-gated) unit tests to assert invalid indices produce the expected failure message.
- Added
<algorithm>include forstd::copy_n.
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 4 comments.
| File | Description |
|---|---|
| onnxruntime/core/providers/cuda/tensor/gather_nd.cc | Adds pre-kernel bounds validation for indices (with device→host copy) in CUDA GatherND. |
| onnxruntime/test/providers/cpu/tensor/gather_nd_op_test.cc | Adds new tests for invalid-index error reporting, including CUDA-only coverage. |
Addresses 3 of 4 open review comments on PR #31646: 1. Division-by-Zero Risk (Comment 1): - Added validation for num_batches != 0 before computing num_slices_per_batch - Added divisibility check: num_slices % num_batches == 0 - Prevents crash on edge cases with zero batch dimensions - Consistent with CPU kernel error handling 2. Device Memory Access and Explicit cudaMemcpyDeviceToHost (Comment 2): - Changed cudaMemcpyDefault to explicit cudaMemcpyDeviceToHost for clarity - Added explicit check for OrtDevice::CUDA type - Added error handling for unsupported device types - Prevents potential device access violations 3. Contrib Op Test Correction (Comment 3): - Renamed test from 'GatherND_contrib_int32_invalid_index_cuda_error' to 'GatherND_contrib_int32_invalid_index_error' - Changed EP from CUDA-only to CPU-only (contrib op has no CUDA kernel) - Test now runs on CPU path consistently - Removed incorrect CUDA environment check that would skip the test Note: Comment 4 (Performance Optimization) requires on-device validation kernel implementation, which is a larger optimization deferred to a follow-up. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
…ent 4)
Performance optimization for large indices tensors:
1. New ValidateIndicesKernel (_ValidateIndicesKernel):
- CUDA kernel that validates indices bounds in parallel on GPU
- Uses atomic compare-and-swap (atomicCAS) to record first invalid index
- Stores result in 1-element device buffer instead of full tensor copy
- Avoids expensive device-to-host transfer and stream synchronization
2. New ValidateIndicesAndReturnFirstInvalidIndex function:
- Wrapper that allocates 1-element device buffer (not full indices tensor)
- Launches validation kernel on all slices in parallel
- Copies back only the error status (8 bytes instead of entire indices)
- Returns -1 if all indices valid, or the first invalid index found
3. Updated PrepareCompute in gather_nd.cc:
- CPU indices: Continue using host-side validation (unchanged)
- CUDA indices: Call on-device kernel instead of full D2H copy
- Eliminates full D2H transfer + cudaStreamSynchronize overhead
- Reduces performance impact for large indices tensors
Performance Impact:
- For small indices: Negligible change (kernel launch overhead minimal)
- For large indices (millions): Significant improvement
* Before: D2H copy entire tensor + stream sync + O(n) host scan
* After: Parallel validation on GPU + 1-element D2H copy + stream sync
* Avoids blocking upstream GPU work from overlapping with validation
Added template instantiations for int32_t and int64_t index types.
Addresses PR #31646 Comment 4 feedback from Copilot reviewer.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
apsonawane
enabled auto-merge (squash)
August 5, 2026 17:26
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
Use the repo's GPU device enum in the CUDA GatherND indices validation path so the plugin build compiles, and drop the unused num_indices local that was triggering -Werror. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This pull request adds improved validation for indices in the CUDA implementation of the
GatherNDoperator and introduces new CUDA-specific tests to ensure invalid indices are properly detected and reported. The main goal is to catch out-of-bounds indices on the host before launching CUDA kernels, improving error handling and robustness.Validation improvements:
GatherNDBase::PrepareComputeto check that all indices are within valid bounds, both for CPU and CUDA tensors, preventing invalid memory accesses during CUDA execution.cudaMemcpyAsyncandcudaStreamSynchronizeto copy indices from device to host for validation when running on CUDA.Testing enhancements:
gather_nd_op_test.ccto verify that invalid indices are properly detected and result in expected error messages for both standard and contrib ops.Code maintenance:
<algorithm>header for use ofstd::copy_ningather_nd.cc.