diff --git a/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/model_builder.cc b/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/model_builder.cc index 98f844ce0f..e15c5641e2 100644 --- a/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/model_builder.cc +++ b/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/model_builder.cc @@ -214,6 +214,9 @@ static Status GetInputDataType( initializers, *node_unit, name, scale, zero_point, ArgType::kInput)); break; } + case ONNX_NAMESPACE::TensorProto_DataType_INT32: + type = Type::TENSOR_INT32; + break; // case ONNX_NAMESPACE::TensorProto_DataType_INT8: // We also do not consider ONNX_NAMESPACE::TensorProto_DataType_INT8 case here, since that can only // be input 2 of Qlinear[Conv/MatMul], which has to be an initializer tensor and will be added 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 aa81b16569..7ef9607352 100644 --- a/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/op_builder.cc +++ b/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/op_builder.cc @@ -2494,9 +2494,15 @@ class GatherOpBuilder : public BaseOpBuilder { }; 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 + const auto& indices_name = inputs[1].node_arg.Name(); + int32_t indices_data_type; + GetType(node_unit.Inputs()[1].node_arg, indices_data_type); + if (Contains(model_builder.GetInitializerTensors(), indices_name) && + indices_data_type != ONNX_NAMESPACE::TensorProto_DataType_INT32) { + // Skip the second input `indices` for Gather if it is an initializer + model_builder.AddInitializerToSkip(indices_name); + } } Status GatherOpBuilder::AddToModelBuilderImpl(ModelBuilder& model_builder, const NodeUnit& node_unit) const { @@ -2516,39 +2522,44 @@ Status GatherOpBuilder::AddToModelBuilderImpl(ModelBuilder& model_builder, const 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)); + int32_t indices_data_type; + GetType(node_unit.Inputs()[1].node_arg, indices_data_type); + if (Contains(model_builder.GetInitializerTensors(), input2) && + indices_data_type != ONNX_NAMESPACE::TensorProto_DataType_INT32) { + // 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); + 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 *= SafeInt(indices_shape[i]); + indices_dimen.push_back(static_cast(indices_shape[i])); } - } 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)); + 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]); 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 aeccfa021d..092e5b451c 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 @@ -2138,10 +2138,22 @@ bool GatherOpSupportChecker::IsOpSupportedImpl(const InitializedTensorSet& initi return false; } + // Here in GatherOpSupportChecker::IsOpSupportedImpl, we removed the restriction that 2nd input "indices" must be an initializer + // to accommodate the support for some models such as mobileBERT. It doesn't need to be an initializer for int32 as NNAPI Gather + // uses int32 for indices so the type matches. + // However, we still require indices of other types to be an initializer as we convert the data to int32 during model building. + // TODO: We could potentially support non-initializer inputs for the other types if we inserted a cast. const auto& indices_name = inputs[1].node_arg.Name(); - if (!Contains(initializers, indices_name)) { - LOGS_DEFAULT(VERBOSE) << "Indices of Gather must be known"; + + int32_t indices_type; + if (!GetType(node_unit.Inputs()[1].node_arg, indices_type)) return false; + + if (indices_type != ONNX_NAMESPACE::TensorProto_DataType_INT32) { + if (!Contains(initializers, indices_name)) { + LOGS_DEFAULT(VERBOSE) << "Indices of Gather must be known."; + return false; + } } return true; diff --git a/onnxruntime/test/providers/cpu/tensor/gather_op_test.cc b/onnxruntime/test/providers/cpu/tensor/gather_op_test.cc index 703d185b7c..1b9fbd32e0 100644 --- a/onnxruntime/test/providers/cpu/tensor/gather_op_test.cc +++ b/onnxruntime/test/providers/cpu/tensor/gather_op_test.cc @@ -101,7 +101,7 @@ TEST(GatherOpTest, Gather_invalid_index_cpu) { ASSERT_STATUS_OK(so.config_options.AddConfigEntry(kOrtSessionOptionsConfigStrictShapeTypeInference, "0")); test.Run(so, OpTester::ExpectResult::kExpectFailure, "indices element out of data bounds, idx=1000 must be within the inclusive range [-3,2]", // On Cuda it is impossible to dereference indices memory on CPU so the check can not run - {kCudaExecutionProvider, kOpenVINOExecutionProvider, kDnnlExecutionProvider, kNupharExecutionProvider, kTensorrtExecutionProvider}); + {kCudaExecutionProvider, kOpenVINOExecutionProvider, kDnnlExecutionProvider, kNupharExecutionProvider, kTensorrtExecutionProvider, kNnapiExecutionProvider}); } #if defined(USE_CUDA) || defined(USE_ROCM) @@ -224,54 +224,41 @@ TEST(GatherOpTest, Gather_axis1_indices2d) { } 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}); + 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}); + 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.Run(); } 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}); + 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}); + 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.Run(); } TEST(GatherOpTest, Gather_axis1_indices2d_int32) {