diff --git a/onnxruntime/core/providers/cpu/tensor/gather_elements.cc b/onnxruntime/core/providers/cpu/tensor/gather_elements.cc index d4172ee6f9..0fd6926c81 100644 --- a/onnxruntime/core/providers/cpu/tensor/gather_elements.cc +++ b/onnxruntime/core/providers/cpu/tensor/gather_elements.cc @@ -9,9 +9,10 @@ namespace onnxruntime { ONNX_CPU_OPERATOR_KERNEL( GatherElements, 11, - KernelDefBuilder().TypeConstraint("T", DataTypeImpl::AllTensorTypes()) - .TypeConstraint("Tind", std::vector{DataTypeImpl::GetTensorType(), - DataTypeImpl::GetTensorType()}), + KernelDefBuilder() + .TypeConstraint("T", DataTypeImpl::AllTensorTypes()) + .TypeConstraint("Tind", std::vector{DataTypeImpl::GetTensorType(), + DataTypeImpl::GetTensorType()}), GatherElements); // Some helpers needed for GatherElements op - @@ -57,7 +58,7 @@ static inline void increment_over_inner_dim(std::vector& current_dims, int64_t rank = static_cast(current_dims.size()); // 'reset' innermost dimension value - current_dims[0] = 0; + current_dims[rank - 1] = 0; // nothing to increment over if (rank == 1) { @@ -77,7 +78,7 @@ static inline void increment_over_inner_dim(std::vector& current_dims, // parse indices_tensor and along the way validate its shape and contents static std::vector parse_and_validate_indices_tensor(const Tensor* indices_tensor, - int64_t axis, const TensorShape& input_shape) { + int64_t axis, const TensorShape& input_shape) { const auto& indices_shape = indices_tensor->Shape().GetDims(); int64_t indices_rank = static_cast(indices_shape.size()); for (int64_t i = 0; i < indices_rank; ++i) { @@ -86,8 +87,10 @@ static std::vector parse_and_validate_indices_tensor(const Tensor* indi // value if within bounds of the corresponding 'data' shape if (i != axis) { if (indices_shape[i] < 0 || indices_shape[i] > input_shape[i]) - ORT_THROW("GatherElements op: 'indices' shape should have values within bounds of 'data' shape. " - "Invalid value in indices shape is: ", indices_shape[i]); + ORT_THROW( + "GatherElements op: 'indices' shape should have values within bounds of 'data' shape. " + "Invalid value in indices shape is: ", + indices_shape[i]); } } @@ -131,21 +134,20 @@ static std::vector parse_and_validate_indices_tensor(const Tensor* indi #pragma GCC diagnostic ignored "-Wclass-memaccess" #endif #endif -template +template static void core_impl(const Tensor* input_tensor, const Tensor* indices_tensor, - Tensor* output_tensor, int64_t axis) { - - // get pointer to input data - // optimizer will remove the redundant if/else block based on 'is_string' template parameter + Tensor* output_tensor, int64_t axis) { + // get pointer to input data + // optimizer will remove the redundant if/else block based on 'is_string' template parameter const T* input_data = nullptr; if (is_string) { - input_data = input_tensor->Data(); + input_data = input_tensor->Data(); } else { - input_data = reinterpret_cast(input_tensor->DataRaw()); + input_data = reinterpret_cast(input_tensor->DataRaw()); } // get pointer to output data - // optimizer will remove the redundant if/else block based on 'is_string' template parameter + // optimizer will remove the redundant if/else block based on 'is_string' template parameter T* output_data = nullptr; if (is_string) { output_data = output_tensor->MutableData(); @@ -176,18 +178,19 @@ static void core_impl(const Tensor* input_tensor, const Tensor* indices_tensor, // process 1 chunk of 'inner dimension' length for (int64_t i = 0; i < inner_dim_size; ++i) { - // optimizer will remove the redundant if/else block based on 'is_string' template parameter + // optimizer will remove the redundant if/else block based on 'is_string' template parameter if (is_string) { - output_data[++output_counter] = input_data[base_offset + (indices_data[++indices_counter] * input_shape_pitches[axis]) + i]; + output_data[++output_counter] = input_data[base_offset + (indices_data[++indices_counter] * input_shape_pitches[axis]) + i]; } else { - memcpy(output_data, input_data + (base_offset + (indices_data[++indices_counter] * input_shape_pitches[axis]) + i) * element_size, element_size); - output_data += element_size; + memcpy(output_data, + input_data + (base_offset + (indices_data[++indices_counter] * input_shape_pitches[axis]) + i) * element_size, element_size); + output_data += element_size; } } increment_over_inner_dim(process_dims, indices_shape); } - } + } // we special-case inner dim as we can weed-out some unnecessary computations in element offset calculations else { while (num_inner_dim-- != 0) { @@ -196,12 +199,12 @@ static void core_impl(const Tensor* input_tensor, const Tensor* indices_tensor, // process 1 chunk of 'inner dimension' length for (int64_t i = 0; i < inner_dim_size; ++i) { // for innermost axis, input_shape_pitches[axis] = 1 (so no need to multiply) - // optimizer will remove the redundant if/else block based on 'is_string' template parameter - if (is_string) { - output_data[++output_counter] = input_data[base_offset + indices_data[++indices_counter]]; + // optimizer will remove the redundant if/else block based on 'is_string' template parameter + if (is_string) { + output_data[++output_counter] = input_data[base_offset + indices_data[++indices_counter]]; } else { memcpy(output_data, input_data + (base_offset + indices_data[++indices_counter]) * element_size, element_size); - output_data += element_size; + output_data += element_size; } } diff --git a/onnxruntime/test/providers/cpu/tensor/gather_elements_op_test.cc b/onnxruntime/test/providers/cpu/tensor/gather_elements_op_test.cc index 4c3b50d64a..0ba7734af8 100644 --- a/onnxruntime/test/providers/cpu/tensor/gather_elements_op_test.cc +++ b/onnxruntime/test/providers/cpu/tensor/gather_elements_op_test.cc @@ -8,46 +8,44 @@ namespace onnxruntime { namespace test { template -void RunTypedTest() -{ +void RunTypedTest() { // int32_t indices - axis 0 OpTester test1("GatherElements", 11); test1.AddAttribute("axis", 0LL); test1.AddInput("data", {2, 3}, - {0, 1, 2, 3, 4, 5}); + {0, 1, 2, 3, 4, 5}); test1.AddInput("indices", {1, 2}, {0, 1}); test1.AddOutput("output", {1, 2}, - {0, 4}); + {0, 4}); test1.Run(); // int32_t indices - axis 1 OpTester test2("GatherElements", 11); test2.AddAttribute("axis", 1LL); test2.AddInput("data", {2, 2}, - {1, 2, - 3, 4}); + {1, 2, + 3, 4}); test2.AddInput("indices", {2, 2}, - {0, 0, - 1, 0}); + {0, 0, + 1, 0}); test2.AddOutput("output", {2, 2}, - {1, 1, - 4, 3}); + {1, 1, + 4, 3}); test2.Run(); - // int64_t indices - axis 1 OpTester test3("GatherElements", 11); test3.AddAttribute("axis", 1LL); test3.AddInput("data", {2, 2}, - {1, 2, - 3, 4}); + {1, 2, + 3, 4}); test3.AddInput("indices", {2, 2}, - {0, 0, - 1, 0}); + {0, 0, + 1, 0}); test3.AddOutput("output", {2, 2}, - {1, 1, - 4, 3}); + {1, 1, + 4, 3}); test3.Run(); // negative indices - axis 1 @@ -106,36 +104,69 @@ void RunTypedTest() {0, 1}); test7.AddOutput("output", {1, 2, 1}, {1, 4}); test7.Run(); + + // 2D input - axis 1 + OpTester test8("GatherElements", 11); + test8.AddAttribute("axis", 1LL); + test8.AddInput("data", {3, 3}, + {1, 2, 3, + 4, 5, 6, + 7, 8, 9}); + test8.AddInput("indices", {3, 2}, + {1, 0, 0, 1, 0, 1}); + test8.AddOutput("output", {3, 2}, {2, 1, 4, 5, 7, 8}); + test8.Run(); + + // 2D input - axis 1 + OpTester test9("GatherElements", 11); + test9.AddAttribute("axis", 0LL); + test9.AddInput("data", {3, 3}, + {1, 2, 3, + 4, 5, 6, + 7, 8, 9}); + test9.AddInput("indices", {3, 2}, + {1, 0, 0, 1, 0, 1}); + test9.AddOutput("output", {3, 2}, {4, 2, 1, 5, 1, 5}); + test9.Run(); + + // 1D input - axis 0 + OpTester test10("GatherElements", 11); + test10.AddAttribute("axis", 0LL); + test10.AddInput("data", {3}, + {1, 2, 3}); + test10.AddInput("indices", {2}, + {1, 0}); + test10.AddOutput("output", {2}, {2, 1}); + test10.Run(); } template <> void RunTypedTest() { - // int32_t indices - axis 0 OpTester test1("GatherElements", 11); test1.AddAttribute("axis", 0LL); test1.AddInput("data", {2, 3}, - {"a", "b", "c", "d", "e", "f"}); + {"a", "b", "c", "d", "e", "f"}); test1.AddInput("indices", {1, 2}, {0, 1}); test1.AddOutput("output", {1, 2}, - {"a", "e"}); + {"a", "e"}); test1.Run(); // int32_t indices - axis 1 OpTester test2("GatherElements", 11); test2.AddAttribute("axis", 1LL); test2.AddInput("data", {2, 2}, - {"a", "b", - "c", "d"}); + {"a", "b", + "c", "d"}); test2.AddInput("indices", {2, 2}, - {0, 0, - 1, 0}); + {0, 0, + 1, 0}); test2.AddOutput("output", {2, 2}, - {"a", "a", - "d", "c"}); + {"a", "a", + "d", "c"}); test2.Run(); - // negative indices - axis 1 + // negative indices - axis 1 OpTester test3("GatherElements", 11); test3.AddAttribute("axis", 1LL); test3.AddInput("data", {2, 2}, @@ -193,6 +224,30 @@ void RunTypedTest() { test6.AddOutput("output", {1, 2, 1}, {"a", "d"}); test6.Run(); + + // 2D input - axis 1 + OpTester test7("GatherElements", 11); + test7.AddAttribute("axis", 1LL); + test7.AddInput("data", {3, 3}, + {"a", "b", "c", + "d", "e", "f", + "g", "h", "i"}); + test7.AddInput("indices", {3, 2}, + {1, 0, 0, 1, 0, 1}); + test7.AddOutput("output", {3, 2}, {"b", "a", "d", "e", "g", "h"}); + test7.Run(); + + // 2D input - axis 2 + OpTester test8("GatherElements", 11); + test8.AddAttribute("axis", 0LL); + test8.AddInput("data", {3, 3}, + {"a", "b", "c", + "d", "e", "f", + "g", "h", "i"}); + test8.AddInput("indices", {3, 2}, + {1, 0, 0, 1, 0, 1}); + test8.AddOutput("output", {3, 2}, {"d", "b", "a", "e", "a", "e"}); + test8.Run(); } TEST(GatherElementsOpTest, int8_t) {