Validate CUDA GatherND indices bounds in release mode - #31646
Open
apsonawane wants to merge 1 commit 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. |
Comment on lines
58
to
+62
| const auto input_batch_stride = input_shape.SizeFromDimension(batch_dims); | ||
| const auto num_slices_per_batch = num_slices / num_batches; | ||
|
|
||
| const TIndex* const indices_data = indices_tensor->Data<TIndex>(); | ||
| const size_t num_indices = static_cast<size_t>(indices_shape.Size()); |
Comment on lines
+65
to
+71
| if (indices_tensor->Location().device.Type() == OrtDevice::CPU) { | ||
| std::copy_n(indices_data, num_indices, indices_data_host.data()); | ||
| } else { | ||
| CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(indices_data_host.data(), indices_data, num_indices * sizeof(TIndex), | ||
| cudaMemcpyDefault, cuda_stream)); | ||
| CUDA_RETURN_IF_ERROR(cudaStreamSynchronize(cuda_stream)); | ||
| } |
Comment on lines
+421
to
+439
| TEST(GatherNDOpTest, GatherND_contrib_int32_invalid_index_cuda_error) { | ||
| if (!HasCudaEnvironment(0)) { | ||
| GTEST_SKIP() << "CUDA not available"; | ||
| } | ||
|
|
||
| OpTester test("GatherND", 1, kMSDomain); | ||
| test.AddInput<float>("data", {4, 4, 4}, std::vector<float>(64, 0.0f)); | ||
| test.AddInput<int32_t>("indices", {1, 2}, {1048576, 0}); | ||
| test.AddOutput<float>("output", {1, 4}, std::vector<float>(4, 0.0f)); | ||
|
|
||
| std::vector<std::unique_ptr<IExecutionProvider>> cuda_only_ep; | ||
| cuda_only_ep.push_back(DefaultCudaExecutionProvider()); | ||
|
|
||
| test.Run(OpTester::ExpectResult::kExpectFailure, | ||
| "invalid index found, index = 1048576", | ||
| {}, | ||
| nullptr, | ||
| &cuda_only_ep); | ||
| } |
Comment on lines
+73
to
+77
| const size_t num_slices_size_t = static_cast<size_t>(num_slices); | ||
| const size_t num_slice_dims_size_t = static_cast<size_t>(num_slice_dims); | ||
| for (size_t slice_idx = 0; slice_idx < num_slices_size_t; ++slice_idx) { | ||
| const size_t slice_base = slice_idx * num_slice_dims_size_t; | ||
| for (size_t dim_idx = 0; dim_idx < num_slice_dims_size_t; ++dim_idx) { |
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.