Add transpose conv support to the NNAPI delegate

PiperOrigin-RevId: 260909964

tensorflow/lite/delegates/nnapi/nnapi_delegate.cc

@@ -557,7 +557,14 @@ class NNAPIOpBuilder {
   TfLiteStatus AddVectorInt32Operand(const int32_t* values,
                                      uint32_t num_values) {
     return AddVectorOperand<int32_t>(values, num_values,
-                                     ANEURALNETWORKS_TENSOR_INT32);
+                                     ANEURALNETWORKS_TENSOR_INT32,
+                                     /*scale=*/0.f, /*zero_point=*/0);
+  }
+
+  TfLiteStatus AddVectorInt32Operand(const int32_t* values, uint32_t num_values,
+                                     float scale, int32_t zero_point) {
+    return AddVectorOperand<int32_t>(
+        values, num_values, ANEURALNETWORKS_TENSOR_INT32, scale, zero_point);
   }
 
   TfLiteStatus AddVectorFloat32Operand(const float* values,
@@ -781,9 +788,13 @@ class NNAPIOpBuilder {
 
   template <typename T>
   TfLiteStatus AddVectorOperand(const T* values, uint32_t num_values,
-                                int32_t nn_type) {
+                                int32_t nn_type, float scale,
-    ANeuralNetworksOperandType operand_type{
+                                int32_t zero_point) {
-        .type = nn_type, .dimensionCount = 1, .dimensions = &num_values};
+    ANeuralNetworksOperandType operand_type{.type = nn_type,
+                                            .dimensionCount = 1,
+                                            .dimensions = &num_values,
+                                            .scale = scale,
+                                            .zeroPoint = zero_point};
 
     RETURN_TFLITE_ERROR_IF_NN_ERROR(
         context_,
@@ -797,6 +808,13 @@ class NNAPIOpBuilder {
     return kTfLiteOk;
   }
 
+  template <typename T>
+  TfLiteStatus AddVectorOperand(const T* values, uint32_t num_values,
+                                int32_t nn_type) {
+    return AddVectorOperand(values, num_values, nn_type, /*scale=*/0.f,
+                            /*zero_point=*/0);
+  }
+
   TfLiteStatus AddFloat32OutputTensor(uint32_t dimension_count,
                                       const uint32_t* dimension_data,
                                       int* ann_index_out) {
@@ -1812,6 +1830,91 @@ class NNAPIDelegateKernel {
           return BasicMappingFn<ANEURALNETWORKS_SIN>;
         }
         break;
+      case kTfLiteBuiltinTransposeConv:
+        if (version == 1 && android_sdk_version >= kMinSdkVersionForNNAPI12) {
+          return [](const NNAPIOpMappingArgs& mapping_args)
+                     -> ANeuralNetworksOperationType {
+            const bool hybrid_op = IsHybridOperator(mapping_args.context,
+                                                    kTfLiteBuiltinTransposeConv,
+                                                    mapping_args.node);
+            mapping_args.builder->AddTensorInput(/*kDataInputTensor*/ 2,
+                                                 hybrid_op);
+            mapping_args.builder->AddTensorInput(/*kWeightsTensor*/ 1,
+                                                 hybrid_op);
+
+            // NNAPI requires a bias tensor, so we allocate a new tensor to fill
+            // it with zeroes. It is deleted with other tensors in the context
+            // during subgraph destructor call.
+            int bias_index = -1;
+            mapping_args.context->AddTensors(mapping_args.context, 1,
+                                             &bias_index);
+            TfLiteTensor* bias_tensor =
+                &mapping_args.context->tensors[bias_index];
+            const auto input_type =
+                mapping_args.context
+                    ->tensors[mapping_args.node->inputs
+                                  ->data[/*kDataInputTensor*/ 2]]
+                    .type;
+            if (input_type == kTfLiteFloat32) {
+              bias_tensor->type = kTfLiteFloat32;
+            } else {
+              bias_tensor->type = kTfLiteInt32;
+            }
+
+            // Create an array with a required bias shape and resize the bias
+            // tensor.
+            TfLiteIntArray* bias_shape = TfLiteIntArrayCreate(1);
+            const TfLiteTensor& output_shape =
+                mapping_args.context->tensors
+                    [mapping_args.node->inputs->data[/*kOutputShapeTensor*/ 0]];
+            const int output_depth = output_shape.data.i32[3];
+            bias_shape->data[0] = output_depth;
+            bias_tensor->allocation_type = kTfLiteDynamic;
+            mapping_args.context->ResizeTensor(mapping_args.context,
+                                               bias_tensor, bias_shape);
+
+            // Set tensor's values to zeroes and add it using AddVector*, so
+            // that the values are copied to NNAPI. We don't use the AddTensor
+            // function because it doesn't copy values and the tensor we just
+            // created is not in the node->inputs.
+            if (input_type == kTfLiteFloat32) {
+              memset(bias_tensor->data.f, 0, output_depth * sizeof(float));
+              mapping_args.builder->AddVectorFloat32Operand(bias_tensor->data.f,
+                                                            output_depth);
+            } else {
+              memset(bias_tensor->data.i32, 0, output_depth * sizeof(int));
+              const TfLiteTensor& input_tensor =
+                  mapping_args.context->tensors
+                      [mapping_args.node->inputs->data[/*kDataInputTensor*/ 2]];
+              const TfLiteTensor& filter_tensor =
+                  mapping_args.context->tensors
+                      [mapping_args.node->inputs->data[/*kWeightsTensor*/ 1]];
+              // NNAPI requires bias scale to be a product of an input scale and
+              // a filter scale.
+              bias_tensor->params.scale =
+                  input_tensor.params.scale * filter_tensor.params.scale;
+              mapping_args.builder->AddVectorInt32Operand(
+                  bias_tensor->data.i32, output_depth,
+                  input_tensor.params.scale * filter_tensor.params.scale,
+                  /*zero_point=*/0);
+            }
+
+            mapping_args.builder->AddTensorInput(/*kOutputShapeTensor*/ 0,
+                                                 hybrid_op);
+
+            auto builtin = reinterpret_cast<TfLiteTransposeConvParams*>(
+                mapping_args.node->builtin_data);
+            mapping_args.builder->AddScalarInt32Operand(builtin->padding);
+            mapping_args.builder->AddScalarInt32Operand(builtin->stride_width);
+            mapping_args.builder->AddScalarInt32Operand(builtin->stride_height);
+            mapping_args.builder->AddScalarInt32Operand(
+                /*ANEURALNETWORKS_FUSED_NONE*/ 0);
+            // Use NHWC layout for input and output
+            mapping_args.builder->AddScalarBoolOperand(false);
+            return ANEURALNETWORKS_TRANSPOSE_CONV;
+          };
+        }
+        break;
       case kTfLiteBuiltinSqrt:
         if (version == 1 && android_sdk_version >= kMinSdkVersionForNNAPI12 &&
             IsFloat(context->tensors[node->inputs->data[0]].type)) {
@@ -3087,13 +3190,17 @@ class NNAPIDelegateKernel {
             continue;
           }
         }
-
         if ((reg->builtin_code == kTfLiteBuiltinSplit) &&
             (input_index == node->inputs->data[0])) {
           // Skip the axis input tensor; it will be added as a scalar operand
           // by the Map() mapping.
           continue;
         }
+        if (reg->builtin_code == kTfLiteBuiltinTransposeConv) {
+          // Everything is added during Map since input tensors
+          // have different order.
+          continue;
+        }
 
         // Pad and Padv2 have an optional parameter for a pad value which has
         // to be converted to a scalar type in NN API.

tensorflow/lite/kernels/BUILD

@@ -1460,6 +1460,7 @@ cc_test(
     name = "transpose_conv_test",
     size = "small",
     srcs = ["transpose_conv_test.cc"],
+    tags = ["tflite_nnapi"],
     deps = [
         ":builtin_ops",
         ":test_main",

tensorflow/lite/nnapi/NeuralNetworksTypes.h

@@ -130,6 +130,7 @@ enum {
   ANEURALNETWORKS_SQRT = 88,
   ANEURALNETWORKS_TILE = 89,
   ANEURALNETWORKS_TOPK_V2 = 90,
+  ANEURALNETWORKS_TRANSPOSE_CONV = 91,
   ANEURALNETWORKS_UNIDIRECTIONAL_SEQUENCE_LSTM = 92,
   ANEURALNETWORKS_UNIDIRECTIONAL_SEQUENCE_RNN = 93,
   ANEURALNETWORKS_RESIZE_NEAREST_NEIGHBOR = 94,