[NNAPI] Add int32_t as supported input data type and other minor gather op updates (#12171)

* update (including commented out code for gather)

* update tests etc.

* update

* minor updates

* fix typo

* fix build

* minor update

* address pr comment

* refine comments

* address pr comment

* update condition check and UTs

* refine code comments

* address lint warning
This commit is contained in:
Rachel Guo 2022-07-20 12:07:46 -07:00 committed by GitHub
parent 5651d91c32
commit 471dbfc250
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
4 changed files with 89 additions and 76 deletions

View file

@ -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

View file

@ -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<uint8_t> 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<uint8_t> 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<uint32_t>(indices_shape[i]));
}
std::vector<int32_t> 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<int32_t>(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<uint32_t>(indices_shape[i]);
indices_dimen.push_back(static_cast<uint32_t>(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<int32_t>(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<int32_t> 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<int32_t>(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<int32_t>(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]);

View file

@ -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;

View file

@ -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<int64_t>("axis", 0LL);
test.AddInput<float>("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<int32_t>("indices", {1}, {1}, indices_is_initializer);
test.AddOutput<float>("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<int64_t>("axis", 0LL);
test.AddInput<float>("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<int32_t>("indices", {1}, {1});
test.AddOutput<float>("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<int64_t>("axis", 0LL);
test.AddInput<float>("data", {3, 3},
{0.0f, 0.1f, 0.2f,
1.0f, 1.1f, 1.2f,
2.0f, 2.1f, 2.2f});
test.AddInput<int32_t>("indices", {2, 2},
{1, 0,
2, 1},
indices_is_initializer);
test.AddOutput<float>("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<int64_t>("axis", 0LL);
test.AddInput<float>("data", {3, 3},
{0.0f, 0.1f, 0.2f,
1.0f, 1.1f, 1.2f,
2.0f, 2.1f, 2.2f});
test.AddInput<int32_t>("indices", {2, 2},
{1, 0,
2, 1});
test.AddOutput<float>("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) {