diff --git a/tensorflow/lite/delegates/nnapi/nnapi_delegate.cc b/tensorflow/lite/delegates/nnapi/nnapi_delegate.cc index 474885fb025..296088cbac7 100644 --- a/tensorflow/lite/delegates/nnapi/nnapi_delegate.cc +++ b/tensorflow/lite/delegates/nnapi/nnapi_delegate.cc @@ -1760,6 +1760,25 @@ class NNAPIDelegateKernel { }; } break; + case kTfLiteBuiltinExpandDims: { + const auto input_type = context->tensors[node->inputs->data[0]].type; + const auto axis = context->tensors[node->inputs->data[1]]; + if (version == 1 && android_sdk_version >= kMinSdkVersionForNNAPI12 && + (input_type == kTfLiteFloat16 || input_type == kTfLiteFloat32 || + input_type == kTfLiteInt32 || input_type == kTfLiteUInt8) && + // TFLite supports axis also as int64 but NNAPI only int32 + (axis.type == kTfLiteInt32 && + axis.allocation_type == kTfLiteMmapRo)) { + return [](const NNAPIOpMappingArgs& mapping_args) + -> ANeuralNetworksOperationType { + const TfLiteTensor& axis_param = + mapping_args.context + ->tensors[mapping_args.node->inputs->data[1]]; + mapping_args.builder->AddScalarInt32Operand(*axis_param.data.i32); + return ANEURALNETWORKS_EXPAND_DIMS; + }; + } + } break; default: // All other operators are not mapped. return nullptr; @@ -2224,6 +2243,10 @@ class NNAPIDelegateKernel { // Everything is added during Map since input tensors // have different order. continue; + } else if (reg->builtin_code == kTfLiteBuiltinExpandDims && + input_pos == 1) { + // The axis param is added during Map + continue; } else { TF_LITE_ENSURE_STATUS(builder.AddTensorInput(input_index, hybrid_op, input_tensor_flags)); diff --git a/tensorflow/lite/kernels/BUILD b/tensorflow/lite/kernels/BUILD index ff79521daa2..5fc0b4f15c6 100644 --- a/tensorflow/lite/kernels/BUILD +++ b/tensorflow/lite/kernels/BUILD @@ -1370,6 +1370,7 @@ cc_test( name = "expand_dims_test", size = "small", srcs = ["expand_dims_test.cc"], + tags = ["tflite_nnapi"], deps = [ ":builtin_ops", ":test_main", diff --git a/tensorflow/lite/kernels/expand_dims_test.cc b/tensorflow/lite/kernels/expand_dims_test.cc index 388a8adea2c..db1c9f462da 100644 --- a/tensorflow/lite/kernels/expand_dims_test.cc +++ b/tensorflow/lite/kernels/expand_dims_test.cc @@ -25,22 +25,39 @@ namespace { using ::testing::ElementsAreArray; +enum class TestType { + CONST = 0, + DYNAMIC = 1, +}; + +template class ExpandDimsOpModel : public SingleOpModel { public: - ExpandDimsOpModel(std::initializer_list input_shape, - TensorType input_type) { - input_ = AddInput(input_type); - axis_ = AddInput(TensorType_INT32); - output_ = AddOutput(input_type); + ExpandDimsOpModel(int axis, std::initializer_list input_shape, + std::initializer_list input_data, + TestType input_tensor_types) { + if (input_tensor_types == TestType::DYNAMIC) { + input_ = AddInput(GetTensorType()); + axis_ = AddInput(TensorType_INT32); + } else { + input_ = + AddConstInput(GetTensorType(), input_data, input_shape); + axis_ = AddConstInput(TensorType_INT32, {axis}, {1}); + } + output_ = AddOutput(GetTensorType()); SetBuiltinOp(BuiltinOperator_EXPAND_DIMS, BuiltinOptions_ExpandDimsOptions, 0); + BuildInterpreter({input_shape, {1}}); + + if (input_tensor_types == TestType::DYNAMIC) { + PopulateTensor(input_, input_data); + PopulateTensor(axis_, {axis}); + } } - void SetInputFloat(std::initializer_list data) { - PopulateTensor(input_, data); + std::vector GetValues() { + return ExtractVector(output_); } - void SetAxis(int axis) { PopulateTensor(axis_, {axis}); } - std::vector GetValuesFloat() { return ExtractVector(output_); } std::vector GetOutputShape() { return GetTensorShape(output_); } protected: @@ -49,29 +66,37 @@ class ExpandDimsOpModel : public SingleOpModel { int output_; }; -TEST(ExpandDimsOpTest, DifferentAxis) { - ExpandDimsOpModel m({2, 2}, TensorType_FLOAT32); +class ExpandDimsOpTest : public ::testing::TestWithParam {}; + +TEST_P(ExpandDimsOpTest, PositiveAxis) { std::initializer_list values = {-1.f, 1.f, -2.f, 2.f}; - m.SetInputFloat(values); - m.SetAxis(0); - m.Invoke(); - EXPECT_THAT(m.GetValuesFloat(), ElementsAreArray(values)); - EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1, 2, 2})); - m.SetAxis(1); - m.Invoke(); - EXPECT_THAT(m.GetValuesFloat(), ElementsAreArray(values)); - EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({2, 1, 2})); + ExpandDimsOpModel axis_0(0, {2, 2}, values, GetParam()); + axis_0.Invoke(); + EXPECT_THAT(axis_0.GetValues(), ElementsAreArray(values)); + EXPECT_THAT(axis_0.GetOutputShape(), ElementsAreArray({1, 2, 2})); - m.SetAxis(2); - m.Invoke(); - EXPECT_THAT(m.GetValuesFloat(), ElementsAreArray(values)); - EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({2, 2, 1})); + ExpandDimsOpModel axis_1(1, {2, 2}, values, GetParam()); + axis_1.Invoke(); + EXPECT_THAT(axis_1.GetValues(), ElementsAreArray(values)); + EXPECT_THAT(axis_1.GetOutputShape(), ElementsAreArray({2, 1, 2})); - m.SetAxis(-1); + ExpandDimsOpModel axis_2(2, {2, 2}, values, GetParam()); + axis_2.Invoke(); + EXPECT_THAT(axis_2.GetValues(), ElementsAreArray(values)); + EXPECT_THAT(axis_2.GetOutputShape(), ElementsAreArray({2, 2, 1})); +} + +TEST_P(ExpandDimsOpTest, NegativeAxis) { + std::initializer_list values = {-1.f, 1.f, -2.f, 2.f}; + + ExpandDimsOpModel m(-1, {2, 2}, values, GetParam()); m.Invoke(); - EXPECT_THAT(m.GetValuesFloat(), ElementsAreArray(values)); + EXPECT_THAT(m.GetValues(), ElementsAreArray(values)); EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({2, 2, 1})); } + +INSTANTIATE_TEST_SUITE_P(ExpandDimsOpTest, ExpandDimsOpTest, + ::testing::Values(TestType::DYNAMIC, TestType::CONST)); } // namespace } // namespace tflite diff --git a/tensorflow/lite/nnapi/NeuralNetworksTypes.h b/tensorflow/lite/nnapi/NeuralNetworksTypes.h index 0bc81019ff3..f618c677db9 100644 --- a/tensorflow/lite/nnapi/NeuralNetworksTypes.h +++ b/tensorflow/lite/nnapi/NeuralNetworksTypes.h @@ -93,6 +93,7 @@ enum { ANEURALNETWORKS_BIDIRECTIONAL_SEQUENCE_LSTM = 42, ANEURALNETWORKS_EQUAL = 48, ANEURALNETWORKS_EXP = 49, + ANEURALNETWORKS_EXPAND_DIMS = 50, ANEURALNETWORKS_GATHER = 51, ANEURALNETWORKS_GREATER = 53, ANEURALNETWORKS_GREATER_EQUAL = 54,