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/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 e6d113e1e4dca..717ffa6b56d9e 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,36 @@ TEST(ScatterNDOpTest, ScatterND_empty_indices) { test1.Run(OpTester::ExpectResult::kExpectSuccess, "", {kDmlExecutionProvider}); } +TEST(ScatterNDOpTest, ScatterND_zero_index_depth_updates_entire_tensor) { + OpTester test("ScatterND", 18); + 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}, {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, "", + {kTensorrtExecutionProvider, kWebGpuExecutionProvider}); +} + +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(OpTester::ExpectResult::kExpectSuccess, "", + {kTensorrtExecutionProvider, kWebGpuExecutionProvider}); +} + +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(OpTester::ExpectResult::kExpectSuccess, "", + {kTensorrtExecutionProvider, kWebGpuExecutionProvider}); +} + } // namespace test } // namespace onnxruntime