diff --git a/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/op_builder.cc b/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/op_builder.cc index de95987a03..ce9ea9c746 100644 --- a/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/op_builder.cc +++ b/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/op_builder.cc @@ -2436,6 +2436,82 @@ Status FlattenOpBuilder::AddToModelBuilderImpl(ModelBuilder& model_builder, cons #pragma endregion +#pragma region op_gather + +class GatherOpBuilder : public BaseOpBuilder { + public: + void AddInitializersToSkip(ModelBuilder& model_builder, const NodeUnit& node_unit) const override; + + private: + Status AddToModelBuilderImpl(ModelBuilder& model_builder, const NodeUnit& node_unit) const override; +}; + +void GatherOpBuilder::AddInitializersToSkip(ModelBuilder& model_builder, const NodeUnit& node_unit) const { + // Skip the second input `indices` for Gather + const auto& inputs = node_unit.Inputs(); + model_builder.AddInitializerToSkip(inputs[1].node_arg.Name()); // indices +} + +Status GatherOpBuilder::AddToModelBuilderImpl(ModelBuilder& model_builder, const NodeUnit& node_unit) const { + auto& shaper(model_builder.GetShaper()); + const auto& operand_indices(model_builder.GetOperandIndices()); + const auto& operand_types(model_builder.GetOperandTypes()); + const auto& initializers(model_builder.GetInitializerTensors()); + const auto& input1 = node_unit.Inputs()[0].node_arg.Name(); + const auto& input2 = node_unit.Inputs()[1].node_arg.Name(); // "indices" + const auto& output = node_unit.Outputs()[0].node_arg.Name(); + + NodeAttrHelper helper(node_unit); + int32_t rank = static_cast(shaper[input1].size()); + int32_t axis = static_cast(HandleNegativeAxis(helper.Get("axis", 0), rank)); + + std::vector input_indices; + input_indices.push_back(operand_indices.at(input1)); + ADD_SCALAR_OPERAND(model_builder, input_indices, axis); + + // Add indices operand into nnapi + const auto& indices_tensor = *initializers.at(input2); + std::vector unpacked_tensor; + ORT_RETURN_IF_ERROR(onnxruntime::utils::UnpackInitializerData(indices_tensor, unpacked_tensor)); + + const auto data_type = indices_tensor.data_type(); + const auto indices_shape = indices_tensor.dims(); + uint32_t size = 1; + Shape indices_dimen; + indices_dimen.reserve(indices_tensor.dims_size()); + for (auto i = 0; i < indices_tensor.dims_size(); i++) { + size *= indices_shape[i]; + indices_dimen.push_back(static_cast(indices_shape[i])); + } + + std::vector indices(size); + // see https://gist.github.com/shafik/848ae25ee209f698763cffee272a58f8#type-punning-arrays for the usage of memcpy here + if (data_type == ONNX_NAMESPACE::TensorProto_DataType_INT64) { + for (uint32_t i = 0; i < size; i++) { + int64_t index_i64; + memcpy(&index_i64, unpacked_tensor.data() + i * sizeof(int64_t), sizeof(int64_t)); + indices[i] = SafeInt(index_i64); + } + } else if (data_type == ONNX_NAMESPACE::TensorProto_DataType_INT32) { + for (uint32_t i = 0; i < size; i++) { + int32_t index; + memcpy(&index, unpacked_tensor.data() + i * sizeof(int32_t), sizeof(int32_t)); + indices[i] = SafeInt(index); + } + } + + OperandType indices_operand_type(Type::TENSOR_INT32, indices_dimen); + ORT_RETURN_IF_ERROR(model_builder.AddOperandFromPersistMemoryBuffer(input2, indices.data(), indices_operand_type)); + input_indices.push_back(operand_indices.at(input2)); + ORT_RETURN_IF_ERROR(shaper.Gather(input1, input2, axis, output)); + const OperandType output_operand_type(operand_types.at(input1).type, shaper[output]); + + return model_builder.AddOperation(ANEURALNETWORKS_GATHER, input_indices, + {output}, {output_operand_type}); +} + +#pragma endregion + #pragma region op_minmax class MinMaxOpBuilder : public BaseOpBuilder { @@ -2721,6 +2797,7 @@ static OpBuilderRegistrations CreateOpBuilderRegistrations() { NNAPI_EP_ADD_SINGLE_OP_BUILDER("DequantizeLinear", DequantizeLinearOpBuilder); NNAPI_EP_ADD_SINGLE_OP_BUILDER("Elu", EluOpBuilder); NNAPI_EP_ADD_SINGLE_OP_BUILDER("Flatten", FlattenOpBuilder); + NNAPI_EP_ADD_SINGLE_OP_BUILDER("Gather", GatherOpBuilder); NNAPI_EP_ADD_SINGLE_OP_BUILDER("Identity", IdentityOpBuilder); NNAPI_EP_ADD_SINGLE_OP_BUILDER("LRN", LRNOpBuilder); NNAPI_EP_ADD_SINGLE_OP_BUILDER("QuantizeLinear", QuantizeLinearOpBuilder); diff --git a/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/op_support_checker.cc b/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/op_support_checker.cc index 3aa9d1cbd2..eea6f2bd27 100644 --- a/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/op_support_checker.cc +++ b/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/op_support_checker.cc @@ -2067,6 +2067,50 @@ bool FlattenOpSupportChecker::IsOpSupportedImpl(const InitializedTensorSet& /* i #pragma endregion +#pragma region op_gather + +class GatherOpSupportChecker : public BaseOpSupportChecker { + private: + int32_t GetMinSupportedNNAPIFeatureLevel(const NodeUnit& /* node_unit */, + const OpSupportCheckParams& /* params */) const override { + return ANEURALNETWORKS_FEATURE_LEVEL_3; + } + + bool IsOpSupportedImpl(const InitializedTensorSet& initializers, const NodeUnit& node_unit, + const OpSupportCheckParams& params) const override; +}; + +bool GatherOpSupportChecker::IsOpSupportedImpl(const InitializedTensorSet& initializers, const NodeUnit& node_unit, + const OpSupportCheckParams& /* params */) const { + const auto& inputs = node_unit.Inputs(); + Shape input_shape; + + if (!GetShape(inputs[0].node_arg, input_shape)) { + return false; + } + + if (input_shape.size() > 4 || input_shape.empty()) { + LOGS_DEFAULT(VERBOSE) << "Gather only supports up to 1-4d shape, input is " + << input_shape.size() << "d shape"; + return false; + } + + if (std::any_of(input_shape.cbegin(), input_shape.cend(), [](int32_t i) { return i == 0; })) { + LOGS_DEFAULT(VERBOSE) << "Gather doesn't support dynamic input shape"; + return false; + } + + const auto& indices_name = inputs[1].node_arg.Name(); + if (!Contains(initializers, indices_name)) { + LOGS_DEFAULT(VERBOSE) << "Indices of Gather must be known"; + return false; + } + + return true; +} + +#pragma endregion + #pragma region op_minmax class MinMaxOpSupportChecker : public BaseOpSupportChecker { @@ -2210,6 +2254,7 @@ static OpSupportCheckerRegistrations CreateOpSupportCheckerRegistrations() { NNAPI_EP_ADD_SINGLE_OP_SUPPORT_CHECKER("DequantizeLinear", DequantizeLinearOpSupportChecker); NNAPI_EP_ADD_SINGLE_OP_SUPPORT_CHECKER("Elu", EluOpSupportChecker); NNAPI_EP_ADD_SINGLE_OP_SUPPORT_CHECKER("Flatten", FlattenOpSupportChecker); + NNAPI_EP_ADD_SINGLE_OP_SUPPORT_CHECKER("Gather", GatherOpSupportChecker); NNAPI_EP_ADD_SINGLE_OP_SUPPORT_CHECKER("LRN", LRNOpSupportChecker); NNAPI_EP_ADD_SINGLE_OP_SUPPORT_CHECKER("QuantizeLinear", QuantizeLinearOpSupportChecker); NNAPI_EP_ADD_SINGLE_OP_SUPPORT_CHECKER("Reshape", ReshapeOpSupportChecker); diff --git a/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/shaper.cc b/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/shaper.cc index e962147b53..68db233472 100644 --- a/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/shaper.cc +++ b/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/shaper.cc @@ -136,6 +136,13 @@ Status Shaper::DepthToSpace(const std::string& input_name, SHAPER_FUNC(DepthToSpace, input_name, blocksize, nchw, output_name); } +Status Shaper::Gather(const std::string& input_name1, + const std::string& input_name2, + const int32_t axis, + const std::string& output_name) { + SHAPER_FUNC(Gather, input_name1, input_name2, axis, output_name); +} + Status Shaper::ResizeUsingScales(const std::string& input_name, const float scale_h, const float scale_w, bool nchw, @@ -428,6 +435,30 @@ Status Shaper::DepthToSpaceImpl(const std::string& input_name, return Status::OK(); } +Status Shaper::GatherImpl(const std::string& input_name1, + const std::string& input_name2, + const int32_t axis, + const std::string& output_name) { + const Shape& input_dimen = shape_map_.at(input_name1); + const Shape& indices_dimen = shape_map_.at(input_name2); + + std::vector output_dimen; + output_dimen.reserve(indices_dimen.size() + input_dimen.size() - 1); + + // Calculate the output dim + for (int32_t i = 0; i < axis; ++i) + output_dimen.push_back(input_dimen[i]); + + for (const auto dim : indices_dimen) + output_dimen.push_back(dim); + + for (size_t i = axis + 1; i < input_dimen.size(); ++i) + output_dimen.push_back(input_dimen[i]); + + shape_map_[output_name] = output_dimen; + return Status::OK(); +} + Status Shaper::ResizeUsingScalesImpl(const std::string& input_name, const float scale_h, const float scale_w, bool nchw, diff --git a/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/shaper.h b/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/shaper.h index 44bcdd35c2..c3f2d45e60 100644 --- a/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/shaper.h +++ b/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/shaper.h @@ -64,6 +64,11 @@ class Shaper { bool nchw, const std::string& output_name); + common::Status Gather(const std::string& input_name1, + const std::string& input_name2, + const int32_t axis, + const std::string& output_name); + common::Status ResizeUsingScales(const std::string& input_name, const float scale_h, const float scale_w, bool nchw, @@ -117,6 +122,10 @@ class Shaper { const int32_t blocksize, bool nchw, const std::string& output_name); + common::Status GatherImpl(const std::string& input_name1, + const std::string& input_name2, + const int32_t axis, + const std::string& output_name); common::Status ResizeUsingScalesImpl(const std::string& input_name, const float scale_h, const float scale_w, bool nchw, diff --git a/onnxruntime/test/providers/cpu/tensor/gather_op_test.cc b/onnxruntime/test/providers/cpu/tensor/gather_op_test.cc index d5229c0113..703d185b7c 100644 --- a/onnxruntime/test/providers/cpu/tensor/gather_op_test.cc +++ b/onnxruntime/test/providers/cpu/tensor/gather_op_test.cc @@ -13,58 +13,76 @@ namespace test { // Those tests will fallback to other EPs TEST(GatherOpTest, Gather_axis0) { - OpTester test("Gather"); - test.AddAttribute("axis", 0LL); - test.AddInput("data", {2, 3, 4}, - {0.0f, 0.1f, 0.2f, 0.3f, - 1.0f, 1.1f, 1.2f, 1.3f, - 2.0f, 2.1f, 2.2f, 2.3f, - 10.0f, 10.1f, 10.2f, 10.3f, - 11.0f, 11.1f, 11.2f, 11.3f, - 12.0f, 12.1f, 12.2f, 12.3f}); - test.AddInput("indices", {1}, {1LL}); - test.AddOutput("output", {1, 3, 4}, - {10.0f, 10.1f, 10.2f, 10.3f, - 11.0f, 11.1f, 11.2f, 11.3f, - 12.0f, 12.1f, 12.2f, 12.3f}); - test.Run(); + // To test for NNAPI EP, we need the indices to be initializers + auto run_test = [](bool indices_is_initializer) { + OpTester test("Gather"); + test.AddAttribute("axis", 0LL); + test.AddInput("data", {2, 3, 4}, + {0.0f, 0.1f, 0.2f, 0.3f, + 1.0f, 1.1f, 1.2f, 1.3f, + 2.0f, 2.1f, 2.2f, 2.3f, + 10.0f, 10.1f, 10.2f, 10.3f, + 11.0f, 11.1f, 11.2f, 11.3f, + 12.0f, 12.1f, 12.2f, 12.3f}); + test.AddInput("indices", {1}, {1LL}, indices_is_initializer); + test.AddOutput("output", {1, 3, 4}, + {10.0f, 10.1f, 10.2f, 10.3f, + 11.0f, 11.1f, 11.2f, 11.3f, + 12.0f, 12.1f, 12.2f, 12.3f}); + test.Run(); + }; + + run_test(false); + run_test(true); } TEST(GatherOpTest, Gather_negative_axis) { - OpTester test("Gather"); - test.AddAttribute("axis", -3LL); - test.AddInput("data", {2, 3, 4}, - {0.0f, 0.1f, 0.2f, 0.3f, - 1.0f, 1.1f, 1.2f, 1.3f, - 2.0f, 2.1f, 2.2f, 2.3f, - 10.0f, 10.1f, 10.2f, 10.3f, - 11.0f, 11.1f, 11.2f, 11.3f, - 12.0f, 12.1f, 12.2f, 12.3f}); - test.AddInput("indices", {1}, {1LL}); - test.AddOutput("output", {1, 3, 4}, - {10.0f, 10.1f, 10.2f, 10.3f, - 11.0f, 11.1f, 11.2f, 11.3f, - 12.0f, 12.1f, 12.2f, 12.3f}); - test.Run(); + // To test for NNAPI EP, we need the indices to be initializers + auto run_test = [](bool indices_is_initializer) { + OpTester test("Gather"); + test.AddAttribute("axis", -3LL); + test.AddInput("data", {2, 3, 4}, + {0.0f, 0.1f, 0.2f, 0.3f, + 1.0f, 1.1f, 1.2f, 1.3f, + 2.0f, 2.1f, 2.2f, 2.3f, + 10.0f, 10.1f, 10.2f, 10.3f, + 11.0f, 11.1f, 11.2f, 11.3f, + 12.0f, 12.1f, 12.2f, 12.3f}); + test.AddInput("indices", {1}, {1LL}, indices_is_initializer); + test.AddOutput("output", {1, 3, 4}, + {10.0f, 10.1f, 10.2f, 10.3f, + 11.0f, 11.1f, 11.2f, 11.3f, + 12.0f, 12.1f, 12.2f, 12.3f}); + test.Run(); + }; + + run_test(false); + run_test(true); } TEST(GatherOpTest, Gather_invalid_axis) { - OpTester test("Gather"); - // Invalid axis not in range [-r, r-1] - test.AddAttribute("axis", -10LL); - test.AddInput("data", {2, 3, 4}, - {0.0f, 0.1f, 0.2f, 0.3f, - 1.0f, 1.1f, 1.2f, 1.3f, - 2.0f, 2.1f, 2.2f, 2.3f, - 10.0f, 10.1f, 10.2f, 10.3f, - 11.0f, 11.1f, 11.2f, 11.3f, - 12.0f, 12.1f, 12.2f, 12.3f}); - test.AddInput("indices", {1}, {1LL}); - test.AddOutput("output", {1, 3, 4}, - {10.0f, 10.1f, 10.2f, 10.3f, - 11.0f, 11.1f, 11.2f, 11.3f, - 12.0f, 12.1f, 12.2f, 12.3f}); - test.Run(OpTester::ExpectResult::kExpectFailure, "axis must be in [-r, r-1]"); + // To test for NNAPI EP, we need the indices to be initializers + auto run_test = [](bool indices_is_initializer) { + OpTester test("Gather"); + // Invalid axis not in range [-r, r-1] + test.AddAttribute("axis", -10LL); + test.AddInput("data", {2, 3, 4}, + {0.0f, 0.1f, 0.2f, 0.3f, + 1.0f, 1.1f, 1.2f, 1.3f, + 2.0f, 2.1f, 2.2f, 2.3f, + 10.0f, 10.1f, 10.2f, 10.3f, + 11.0f, 11.1f, 11.2f, 11.3f, + 12.0f, 12.1f, 12.2f, 12.3f}); + test.AddInput("indices", {1}, {1LL}, indices_is_initializer); + test.AddOutput("output", {1, 3, 4}, + {10.0f, 10.1f, 10.2f, 10.3f, + 11.0f, 11.1f, 11.2f, 11.3f, + 12.0f, 12.1f, 12.2f, 12.3f}); + test.Run(OpTester::ExpectResult::kExpectFailure, "axis must be in [-r, r-1]"); + }; + + run_test(false); + run_test(true); } TEST(GatherOpTest, Gather_invalid_index_cpu) { @@ -107,76 +125,153 @@ TEST(GatherOpTest, Gather_invalid_index_gpu) { #endif TEST(GatherOpTest, Gather_axis1) { - OpTester test("Gather"); - test.AddAttribute("axis", 1LL); - test.AddInput("data", {2, 3, 4}, - {0.0f, 0.1f, 0.2f, 0.3f, - 1.0f, 1.1f, 1.2f, 1.3f, - 2.0f, 2.1f, 2.2f, 2.3f, - 10.0f, 10.1f, 10.2f, 10.3f, - 11.0f, 11.1f, 11.2f, 11.3f, - 12.0f, 12.1f, 12.2f, 12.3f}); - test.AddInput("indices", {2}, {2LL, 0LL}); - test.AddOutput("output", {2, 2, 4}, - {2.0f, 2.1f, 2.2f, 2.3f, - 0.0f, 0.1f, 0.2f, 0.3f, - 12.0f, 12.1f, 12.2f, 12.3f, - 10.0f, 10.1f, 10.2f, 10.3f}); - test.Run(); + // To test for NNAPI EP, we need the indices to be initializers + auto run_test = [](bool indices_is_initializer) { + OpTester test("Gather"); + test.AddAttribute("axis", 1LL); + test.AddInput("data", {2, 3, 4}, + {0.0f, 0.1f, 0.2f, 0.3f, + 1.0f, 1.1f, 1.2f, 1.3f, + 2.0f, 2.1f, 2.2f, 2.3f, + 10.0f, 10.1f, 10.2f, 10.3f, + 11.0f, 11.1f, 11.2f, 11.3f, + 12.0f, 12.1f, 12.2f, 12.3f}); + test.AddInput("indices", {2}, {2LL, 0LL}, indices_is_initializer); + test.AddOutput("output", {2, 2, 4}, + {2.0f, 2.1f, 2.2f, 2.3f, + 0.0f, 0.1f, 0.2f, 0.3f, + 12.0f, 12.1f, 12.2f, 12.3f, + 10.0f, 10.1f, 10.2f, 10.3f}); + test.Run(); + }; + + run_test(false); + run_test(true); } TEST(GatherOpTest, Gather_axis2) { - OpTester test("Gather"); - test.AddAttribute("axis", 2LL); - test.AddInput("data", {2, 3, 4}, - {0.0f, 0.1f, 0.2f, 0.3f, - 1.0f, 1.1f, 1.2f, 1.3f, - 2.0f, 2.1f, 2.2f, 2.3f, - 10.0f, 10.1f, 10.2f, 10.3f, - 11.0f, 11.1f, 11.2f, 11.3f, - 12.0f, 12.1f, 12.2f, 12.3f}); - test.AddInput("indices", {3}, {1LL, 0LL, 2LL}); - test.AddOutput("output", {2, 3, 3}, - {0.1f, 0.0f, 0.2f, - 1.1f, 1.0f, 1.2f, - 2.1f, 2.0f, 2.2f, - 10.1f, 10.0f, 10.2f, - 11.1f, 11.0f, 11.2f, - 12.1f, 12.0f, 12.2f}); - test.Run(); + // To test for NNAPI EP, we need the indices to be initializers + auto run_test = [](bool indices_is_initializer) { + OpTester test("Gather"); + test.AddAttribute("axis", 2LL); + test.AddInput("data", {2, 3, 4}, + {0.0f, 0.1f, 0.2f, 0.3f, + 1.0f, 1.1f, 1.2f, 1.3f, + 2.0f, 2.1f, 2.2f, 2.3f, + 10.0f, 10.1f, 10.2f, 10.3f, + 11.0f, 11.1f, 11.2f, 11.3f, + 12.0f, 12.1f, 12.2f, 12.3f}); + test.AddInput("indices", {3}, {1LL, 0LL, 2LL}, indices_is_initializer); + test.AddOutput("output", {2, 3, 3}, + {0.1f, 0.0f, 0.2f, + 1.1f, 1.0f, 1.2f, + 2.1f, 2.0f, 2.2f, + 10.1f, 10.0f, 10.2f, + 11.1f, 11.0f, 11.2f, + 12.1f, 12.0f, 12.2f}); + test.Run(); + }; + + run_test(false); + run_test(true); } TEST(GatherOpTest, Gather_axis0_indices2d) { - OpTester test("Gather"); - test.AddAttribute("axis", 0LL); - test.AddInput("data", {3, 3}, - {0.0f, 0.1f, 0.2f, - 1.0f, 1.1f, 1.2f, - 2.0f, 2.1f, 2.2f}); - test.AddInput("indices", {2LL, 2LL}, - {1LL, 0LL, - 2LL, 1LL}); - test.AddOutput("output", {2, 2, 3}, - {1.0f, 1.1f, 1.2f, 0.0f, 0.1f, 0.2f, - 2.0f, 2.1f, 2.2f, 1.0f, 1.1f, 1.2f}); - test.Run(); + // To test for NNAPI EP, we need the indices to be initializers + auto run_test = [](bool indices_is_initializer) { + OpTester test("Gather"); + test.AddAttribute("axis", 0LL); + test.AddInput("data", {3, 3}, + {0.0f, 0.1f, 0.2f, + 1.0f, 1.1f, 1.2f, + 2.0f, 2.1f, 2.2f}); + test.AddInput("indices", {2LL, 2LL}, + {1LL, 0LL, + 2LL, 1LL}, + indices_is_initializer); + test.AddOutput("output", {2, 2, 3}, + {1.0f, 1.1f, 1.2f, 0.0f, 0.1f, 0.2f, + 2.0f, 2.1f, 2.2f, 1.0f, 1.1f, 1.2f}); + test.Run(); + }; + + run_test(false); + run_test(true); } TEST(GatherOpTest, Gather_axis1_indices2d) { - OpTester test("Gather"); - test.AddAttribute("axis", 1LL); - test.AddInput("data", {3, 3}, - {0.0f, 0.1f, 0.2f, - 1.0f, 1.1f, 1.2f, - 2.0f, 2.1f, 2.2f}); - test.AddInput("indices", {2LL, 2LL}, - {1LL, 0LL, - 2LL, 1LL}); - test.AddOutput("output", {3, 2, 2}, - {0.1f, 0.0f, 0.2f, 0.1f, - 1.1f, 1.0f, 1.2f, 1.1f, - 2.1f, 2.0f, 2.2f, 2.1f}); - test.Run(); + // To test for NNAPI EP, we need the indices to be initializers + auto run_test = [](bool indices_is_initializer) { + OpTester test("Gather"); + test.AddAttribute("axis", 1LL); + test.AddInput("data", {3, 3}, + {0.0f, 0.1f, 0.2f, + 1.0f, 1.1f, 1.2f, + 2.0f, 2.1f, 2.2f}); + test.AddInput("indices", {2LL, 2LL}, + {1LL, 0LL, + 2LL, 1LL}, + indices_is_initializer); + test.AddOutput("output", {3, 2, 2}, + {0.1f, 0.0f, 0.2f, 0.1f, + 1.1f, 1.0f, 1.2f, 1.1f, + 2.1f, 2.0f, 2.2f, 2.1f}); + test.Run(); + }; + + run_test(false); + run_test(true); +} + +TEST(GatherOpTest, Gather_axis0_indicesInt32) { + // To test for NNAPI EP, we need the indices to be initializers + // NNAPI EP only supports float input data for now, + // the following two test cases cover int32_t indices with float input other than int64_t type for Nnapi + auto run_test = [](bool indices_is_initializer) { + OpTester test("Gather"); + test.AddAttribute("axis", 0LL); + test.AddInput("data", {2, 3, 4}, + {0.0f, 0.1f, 0.2f, 0.3f, + 1.0f, 1.1f, 1.2f, 1.3f, + 2.0f, 2.1f, 2.2f, 2.3f, + 10.0f, 10.1f, 10.2f, 10.3f, + 11.0f, 11.1f, 11.2f, 11.3f, + 12.0f, 12.1f, 12.2f, 12.3f}); + test.AddInput("indices", {1}, {1}, indices_is_initializer); + test.AddOutput("output", {1, 3, 4}, + {10.0f, 10.1f, 10.2f, 10.3f, + 11.0f, 11.1f, 11.2f, 11.3f, + 12.0f, 12.1f, 12.2f, 12.3f}); + + test.Run(); + }; + + run_test(false); + run_test(true); +} + +TEST(GatherOpTest, Gather_axis0_indices2dInt32) { + // To test for NNAPI EP, we need the indices to be initializers + auto run_test = [](bool indices_is_initializer) { + OpTester test("Gather"); + test.AddAttribute("axis", 0LL); + test.AddInput("data", {3, 3}, + {0.0f, 0.1f, 0.2f, + 1.0f, 1.1f, 1.2f, + 2.0f, 2.1f, 2.2f}); + test.AddInput("indices", {2, 2}, + {1, 0, + 2, 1}, + indices_is_initializer); + test.AddOutput("output", {2, 2, 3}, + {1.0f, 1.1f, 1.2f, 0.0f, 0.1f, 0.2f, + 2.0f, 2.1f, 2.2f, 1.0f, 1.1f, 1.2f}); + + test.Run(); + }; + + run_test(false); + run_test(true); } TEST(GatherOpTest, Gather_axis1_indices2d_int32) {