From bc1a5128879ac7dcc6fc3ef411bbe38afe182929 Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Wed, 12 Aug 2026 10:33:56 -0700 Subject: [PATCH 1/4] Validate ScatterND index depth --- onnxruntime/core/providers/cpu/tensor/scatter_nd.h | 5 +++++ .../test/providers/cpu/tensor/scatter_nd_op_test.cc | 9 +++++++++ 2 files changed, 14 insertions(+) diff --git a/onnxruntime/core/providers/cpu/tensor/scatter_nd.h b/onnxruntime/core/providers/cpu/tensor/scatter_nd.h index 6560ece07b28f..a4cde0caaeba2 100644 --- a/onnxruntime/core/providers/cpu/tensor/scatter_nd.h +++ b/onnxruntime/core/providers/cpu/tensor/scatter_nd.h @@ -31,6 +31,11 @@ inline Status ValidateShapes(const TensorShape& input_shape, } auto last_indice_dimension = indice_shape[indice_rank - 1]; + if (last_indice_dimension < 1) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "last dimension of indices must be at least 1, got ", + last_indice_dimension); + } if (last_indice_dimension > static_cast(input_rank)) { return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "last dimension of indices must not be larger than rank of input tensor"); diff --git a/onnxruntime/test/providers/cpu/tensor/scatter_nd_op_test.cc b/onnxruntime/test/providers/cpu/tensor/scatter_nd_op_test.cc index e6d113e1e4dca..4a96fa935a9c3 100644 --- a/onnxruntime/test/providers/cpu/tensor/scatter_nd_op_test.cc +++ b/onnxruntime/test/providers/cpu/tensor/scatter_nd_op_test.cc @@ -246,5 +246,14 @@ TEST(ScatterNDOpTest, ScatterND_empty_indices) { test1.Run(OpTester::ExpectResult::kExpectSuccess, "", {kDmlExecutionProvider}); } +TEST(ScatterNDOpTest, ScatterND_rejects_zero_index_depth) { + OpTester test("ScatterND", 18); + test.AddInput("data", {2, 3}, std::vector(6)); + test.AddInput("indices", {1, 0}, {}); + test.AddInput("updates", {1, 2, 3}, std::vector(6)); + test.AddOutput("output", {2, 3}, std::vector(6)); + test.Run(OpTester::ExpectResult::kExpectFailure, "last dimension of indices must be at least 1"); +} + } // namespace test } // namespace onnxruntime From 3a9cf23eb5e56a4d28a23c3e90f933d5f75d70b2 Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Thu, 13 Aug 2026 13:58:25 -0700 Subject: [PATCH 2/4] Support zero-depth ScatterND indices --- .../core/providers/cpu/tensor/scatter_nd.cc | 15 ++++++++-- .../core/providers/cpu/tensor/scatter_nd.h | 5 ---- .../core/providers/cuda/tensor/scatter_nd.cc | 30 +++++++++++-------- .../cpu/tensor/scatter_nd_op_test.cc | 29 ++++++++++++++---- 4 files changed, 54 insertions(+), 25 deletions(-) diff --git a/onnxruntime/core/providers/cpu/tensor/scatter_nd.cc b/onnxruntime/core/providers/cpu/tensor/scatter_nd.cc index 833ee1b774f30..af8ef7f62156c 100644 --- a/onnxruntime/core/providers/cpu/tensor/scatter_nd.cc +++ b/onnxruntime/core/providers/cpu/tensor/scatter_nd.cc @@ -61,11 +61,13 @@ struct Prepare { TData* output_base; uint64_t element_to_copy; std::vector element_offsets; + bool serialize_updates; Prepare() : input_base(nullptr), output_base(nullptr), element_to_copy(0), - element_offsets(0) {} + element_offsets(0), + serialize_updates(false) {} }; // struct Prepare template @@ -89,7 +91,7 @@ Status PrepareForCompute(OpKernelContext* context, Prepare& p) { auto last_indice_dimension = indice_shape[indice_shape.NumDimensions() - 1]; // Re-use input for output. If input/output Tensor* are the same, do not copy. - if (src_base != dst_base) { + if (src_base != dst_base && input_tensor->Shape().Size() > 0) { if (is_string_type) { const auto* str_begin = input_tensor->Data(); const std::string* str_end = str_begin + input_shape.Size(); @@ -108,8 +110,9 @@ Status PrepareForCompute(OpKernelContext* context, Prepare& p) { } p.element_to_copy = input_shape.SizeFromDimension(onnxruntime::narrow(last_indice_dimension)); + p.serialize_updates = last_indice_dimension == 0; const int64_t* indice_offset = indice_tensor->Data(); - auto offset_count = indice_shape.Size() / last_indice_dimension; // Times to copy + auto offset_count = indice_shape.SizeToDimension(indice_shape.NumDimensions() - 1); // Times to copy p.element_offsets.assign(onnxruntime::narrow(offset_count), 0LL); p.input_base = update_tensor->Data(); @@ -302,6 +305,12 @@ struct ScatterNDDispatchTarget { Status operator()(OpKernelContext* context, concurrency::ThreadPool* tp, ScatterND::Reduction reduction) const { Prepare prepare; ORT_RETURN_IF_ERROR(PrepareForCompute(context, prepare)); + if (prepare.element_to_copy == 0 || prepare.element_offsets.empty()) { + return Status::OK(); + } + if (prepare.serialize_updates) { + tp = nullptr; + } auto lambda = [&](ptrdiff_t i) { switch (reduction) { diff --git a/onnxruntime/core/providers/cpu/tensor/scatter_nd.h b/onnxruntime/core/providers/cpu/tensor/scatter_nd.h index a4cde0caaeba2..6560ece07b28f 100644 --- a/onnxruntime/core/providers/cpu/tensor/scatter_nd.h +++ b/onnxruntime/core/providers/cpu/tensor/scatter_nd.h @@ -31,11 +31,6 @@ inline Status ValidateShapes(const TensorShape& input_shape, } auto last_indice_dimension = indice_shape[indice_rank - 1]; - if (last_indice_dimension < 1) { - return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, - "last dimension of indices must be at least 1, got ", - last_indice_dimension); - } if (last_indice_dimension > static_cast(input_rank)) { return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "last dimension of indices must not be larger than rank of input tensor"); diff --git a/onnxruntime/core/providers/cuda/tensor/scatter_nd.cc b/onnxruntime/core/providers/cuda/tensor/scatter_nd.cc index e6359cc048048..530160b94f43c 100644 --- a/onnxruntime/core/providers/cuda/tensor/scatter_nd.cc +++ b/onnxruntime/core/providers/cuda/tensor/scatter_nd.cc @@ -90,18 +90,21 @@ Status ScatterNDDisjointAndNoReduction::ComputeInternal(OpKernelContext* context const void* input_data = input_tensor->DataRaw(); void* output_data = output_tensor->MutableDataRaw(); - if (input_data != output_data) { + if (input_data != output_data && input_tensor->SizeInBytes() > 0) { // TODO: Run benchmarks to determine if a dedicated kernel doing data copy will be faster than invoking cudaMemcpy ? CUDA_RETURN_IF_ERROR( cudaMemcpyAsync(output_data, input_data, input_tensor->SizeInBytes(), cudaMemcpyDeviceToDevice, Stream(context))); } + const auto num_indices = indices_shape.SizeToDimension(indices_shape.NumDimensions() - 1); + auto last_index_dimension = indices_shape[indices_shape.NumDimensions() - 1]; + const auto num_update_elements = input_shape.SizeFromDimension(last_index_dimension); + // Bail out early - if (indices_shape.Size() == 0) { + if (num_indices == 0 || num_update_elements == 0) { return Status::OK(); } - auto last_index_dimension = indices_shape[indices_shape.NumDimensions() - 1]; size_t element_size = input_tensor->DataType()->Size(); // We need element counts for each dimension and the input dim value for each dimension @@ -118,12 +121,12 @@ Status ScatterNDDisjointAndNoReduction::ComputeInternal(OpKernelContext* context Stream(context), output_data, element_size, - indices_shape.Size() / static_cast(last_index_dimension), + onnxruntime::narrow(num_indices), indices_tensor->Data(), // only int64_t is supported for indices as per the onnx spec last_index_dimension, element_counts_and_input_dims, updates_tensor->DataRaw(), - input_shape.SizeFromDimension(last_index_dimension))); + onnxruntime::narrow(num_update_elements))); return Status::OK(); } @@ -145,7 +148,7 @@ Status ScatterNDWithAtomicReduction::ComputeInternal(OpKernelContext* context) c const void* input_data = input_tensor->DataRaw(); void* output_data = output_tensor->MutableDataRaw(); - if (input_data != output_data) { + if (input_data != output_data && input_tensor->SizeInBytes() > 0) { // TODO: Run benchmarks to determine if a dedicated kernel doing data copy will // be faster than invoking cudaMemcpy ? CUDA_RETURN_IF_ERROR( @@ -153,12 +156,15 @@ Status ScatterNDWithAtomicReduction::ComputeInternal(OpKernelContext* context) c cudaMemcpyDeviceToDevice, Stream(context))); } + const auto num_indices = indices_shape.SizeToDimension(indices_shape.NumDimensions() - 1); + auto last_index_dimension = indices_shape[indices_shape.NumDimensions() - 1]; + const auto num_update_elements = input_shape.SizeFromDimension(last_index_dimension); + // Bail out early - if (indices_shape.Size() == 0) { + if (num_indices == 0 || num_update_elements == 0) { return Status::OK(); } - auto last_index_dimension = indices_shape[indices_shape.NumDimensions() - 1]; ElementCountsAndInputDimsSpanOrGpu element_counts_and_input_dims; CudaAsyncBuffer element_counts_and_input_dims_gpu(this); ORT_RETURN_IF_ERROR(InitializeElementCountsAndInputDimsSpanOrGpu(last_index_dimension, input_shape, @@ -173,12 +179,12 @@ Status ScatterNDWithAtomicReduction::ComputeInternal(OpKernelContext* context) c Stream(context), output_data, element_size, - indices_shape.Size() / static_cast(last_index_dimension), + onnxruntime::narrow(num_indices), indices_tensor->Data(), // only int64_t is supported for indices as per the onnx spec last_index_dimension, element_counts_and_input_dims, updates_tensor->DataRaw(), - input_shape.SizeFromDimension(last_index_dimension))); + onnxruntime::narrow(num_update_elements))); } break; case ScatterNDReduction::Add: case ScatterNDReduction::Min: @@ -189,12 +195,12 @@ Status ScatterNDWithAtomicReduction::ComputeInternal(OpKernelContext* context) c Stream(context), output_data, element_type, - indices_shape.Size() / static_cast(last_index_dimension), + onnxruntime::narrow(num_indices), indices_tensor->Data(), // only int64_t is supported for indices as per the onnx spec last_index_dimension, element_counts_and_input_dims, updates_tensor->DataRaw(), - input_shape.SizeFromDimension(last_index_dimension), + onnxruntime::narrow(num_update_elements), reduction_)); } break; default: diff --git a/onnxruntime/test/providers/cpu/tensor/scatter_nd_op_test.cc b/onnxruntime/test/providers/cpu/tensor/scatter_nd_op_test.cc index 4a96fa935a9c3..86848843fa5b1 100644 --- a/onnxruntime/test/providers/cpu/tensor/scatter_nd_op_test.cc +++ b/onnxruntime/test/providers/cpu/tensor/scatter_nd_op_test.cc @@ -246,13 +246,32 @@ TEST(ScatterNDOpTest, ScatterND_empty_indices) { test1.Run(OpTester::ExpectResult::kExpectSuccess, "", {kDmlExecutionProvider}); } -TEST(ScatterNDOpTest, ScatterND_rejects_zero_index_depth) { +TEST(ScatterNDOpTest, ScatterND_zero_index_depth_updates_entire_tensor) { OpTester test("ScatterND", 18); - test.AddInput("data", {2, 3}, std::vector(6)); + test.AddInput("data", {2, 3}, {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f}); test.AddInput("indices", {1, 0}, {}); - test.AddInput("updates", {1, 2, 3}, std::vector(6)); - test.AddOutput("output", {2, 3}, std::vector(6)); - test.Run(OpTester::ExpectResult::kExpectFailure, "last dimension of indices must be at least 1"); + test.AddInput("updates", {1, 2, 3}, {10.0f, 20.0f, 30.0f, 40.0f, 50.0f, 60.0f}); + test.AddOutput("output", {2, 3}, {10.0f, 20.0f, 30.0f, 40.0f, 50.0f, 60.0f}); + test.Run(); +} + +TEST(ScatterNDOpTest, ScatterND_zero_index_depth_adds_multiple_updates) { + OpTester test("ScatterND", 18); + test.AddAttribute("reduction", "add"); + test.AddInput("data", {2}, {1.0f, 2.0f}); + test.AddInput("indices", {2, 0}, {}); + test.AddInput("updates", {2, 2}, {10.0f, 20.0f, 100.0f, 200.0f}); + test.AddOutput("output", {2}, {111.0f, 222.0f}); + test.Run(); +} + +TEST(ScatterNDOpTest, ScatterND_zero_index_depth_empty_data) { + OpTester test("ScatterND", 18); + test.AddInput("data", {0, 3}, {}); + test.AddInput("indices", {1, 0}, {}); + test.AddInput("updates", {1, 0, 3}, {}); + test.AddOutput("output", {0, 3}, {}); + test.Run(); } } // namespace test From 6c253a0cac9ef8b7cf2ff87c9a14fe24f00ef543 Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Thu, 13 Aug 2026 16:15:41 -0700 Subject: [PATCH 3/4] Exclude WebGPU from zero-depth ScatterND tests --- onnxruntime/test/providers/cpu/tensor/scatter_nd_op_test.cc | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/onnxruntime/test/providers/cpu/tensor/scatter_nd_op_test.cc b/onnxruntime/test/providers/cpu/tensor/scatter_nd_op_test.cc index 86848843fa5b1..aeb9da14fd653 100644 --- a/onnxruntime/test/providers/cpu/tensor/scatter_nd_op_test.cc +++ b/onnxruntime/test/providers/cpu/tensor/scatter_nd_op_test.cc @@ -252,7 +252,7 @@ TEST(ScatterNDOpTest, ScatterND_zero_index_depth_updates_entire_tensor) { test.AddInput("indices", {1, 0}, {}); test.AddInput("updates", {1, 2, 3}, {10.0f, 20.0f, 30.0f, 40.0f, 50.0f, 60.0f}); test.AddOutput("output", {2, 3}, {10.0f, 20.0f, 30.0f, 40.0f, 50.0f, 60.0f}); - test.Run(); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kWebGpuExecutionProvider}); } TEST(ScatterNDOpTest, ScatterND_zero_index_depth_adds_multiple_updates) { @@ -262,7 +262,7 @@ TEST(ScatterNDOpTest, ScatterND_zero_index_depth_adds_multiple_updates) { test.AddInput("indices", {2, 0}, {}); test.AddInput("updates", {2, 2}, {10.0f, 20.0f, 100.0f, 200.0f}); test.AddOutput("output", {2}, {111.0f, 222.0f}); - test.Run(); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kWebGpuExecutionProvider}); } TEST(ScatterNDOpTest, ScatterND_zero_index_depth_empty_data) { @@ -271,7 +271,7 @@ TEST(ScatterNDOpTest, ScatterND_zero_index_depth_empty_data) { test.AddInput("indices", {1, 0}, {}); test.AddInput("updates", {1, 0, 3}, {}); test.AddOutput("output", {0, 3}, {}); - test.Run(); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kWebGpuExecutionProvider}); } } // namespace test From c86bbe5bde8bb539b93a33c326a964e5ac6b396b Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Mon, 17 Aug 2026 13:03:02 -0700 Subject: [PATCH 4/4] Exclude TensorRT from zero-depth ScatterND tests --- .../test/providers/cpu/tensor/scatter_nd_op_test.cc | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/onnxruntime/test/providers/cpu/tensor/scatter_nd_op_test.cc b/onnxruntime/test/providers/cpu/tensor/scatter_nd_op_test.cc index aeb9da14fd653..717ffa6b56d9e 100644 --- a/onnxruntime/test/providers/cpu/tensor/scatter_nd_op_test.cc +++ b/onnxruntime/test/providers/cpu/tensor/scatter_nd_op_test.cc @@ -252,7 +252,8 @@ TEST(ScatterNDOpTest, ScatterND_zero_index_depth_updates_entire_tensor) { test.AddInput("indices", {1, 0}, {}); test.AddInput("updates", {1, 2, 3}, {10.0f, 20.0f, 30.0f, 40.0f, 50.0f, 60.0f}); test.AddOutput("output", {2, 3}, {10.0f, 20.0f, 30.0f, 40.0f, 50.0f, 60.0f}); - test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kWebGpuExecutionProvider}); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", + {kTensorrtExecutionProvider, kWebGpuExecutionProvider}); } TEST(ScatterNDOpTest, ScatterND_zero_index_depth_adds_multiple_updates) { @@ -262,7 +263,8 @@ TEST(ScatterNDOpTest, ScatterND_zero_index_depth_adds_multiple_updates) { test.AddInput("indices", {2, 0}, {}); test.AddInput("updates", {2, 2}, {10.0f, 20.0f, 100.0f, 200.0f}); test.AddOutput("output", {2}, {111.0f, 222.0f}); - test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kWebGpuExecutionProvider}); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", + {kTensorrtExecutionProvider, kWebGpuExecutionProvider}); } TEST(ScatterNDOpTest, ScatterND_zero_index_depth_empty_data) { @@ -271,7 +273,8 @@ TEST(ScatterNDOpTest, ScatterND_zero_index_depth_empty_data) { test.AddInput("indices", {1, 0}, {}); test.AddInput("updates", {1, 0, 3}, {}); test.AddOutput("output", {0, 3}, {}); - test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kWebGpuExecutionProvider}); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", + {kTensorrtExecutionProvider, kWebGpuExecutionProvider}); } } // namespace test