patch/0008-nnapi-Convolution2DTransposeBias.patch

From 4aa6669b85746436929c19dee88e69e6f716746f Mon Sep 17 00:00:00 2001
From: TFLite Patcher <tflite@localhost>
Date: Fri, 18 Oct 2024 18:05:59 +0800
Subject: [PATCH] Implement Convolution2DTransposeBias in NNAPI delegate

PATCH_NAME=nnapi-Convolution2DTransposeBias
---
 .../lite/delegates/nnapi/nnapi_delegate.cc    | 47 +++++++++++++++++++
 1 file changed, 47 insertions(+)

diff --git a/tensorflow/lite/delegates/nnapi/nnapi_delegate.cc b/tensorflow/lite/delegates/nnapi/nnapi_delegate.cc
index a9a5cd1a..b6c60ff7 100644
--- a/tensorflow/lite/delegates/nnapi/nnapi_delegate.cc
+++ b/tensorflow/lite/delegates/nnapi/nnapi_delegate.cc
@@ -3630,6 +3630,8 @@ bool NNAPIDelegateKernel::Validate(
              NNAPIValidationFailureType::kUnsupportedOperatorVariant,
              "Expecting 2 inputs", &val_ctx);
     } break;
+    case kTfLiteBuiltinCustom: {
+    } break;
     default:
       // All other operators are not mapped.
       AddValidationFailure(NNAPIValidationFailureType::kUnsupportedOperator,
@@ -4544,6 +4546,49 @@ TfLiteStatus NNAPIDelegateKernel::Map(
     case kTfLiteBuiltinReverseV2: {
       *nn_op_type = ANEURALNETWORKS_REVERSE;
     } break;
+    case kTfLiteBuiltinCustom: {
+      TfLiteNode* node;
+      TfLiteRegistration* reg;
+      context->GetNodeAndRegistration(context, mapping_args.node_index, &node, &reg);
+      if (strcmp(reg->custom_name, "Convolution2DTransposeBias") == 0) {
+        // Convolution2DTransposeBias can be processed as TRANSPOSE_CONV_2D
+        int input_tensor_flags = 0;
+        const int input_tensor_id =
+            mapping_args.node->inputs->data[0];
+        const int weight_tensor_id =
+            mapping_args.node->inputs->data[1];
+        const int bias_tensor_id =
+            mapping_args.node->inputs->data[2];
+
+        // Transpose convolution doesn't have hybrid variation.
+        const bool hybrid_op = false;
+
+        mapping_args.builder->AddTensorInput(input_tensor_id, hybrid_op, input_tensor_flags);   // 1
+        mapping_args.builder->AddTensorInput(weight_tensor_id, hybrid_op, input_tensor_flags);  // 2
+
+        // Get params from custom_initial_data
+        TfLiteTransposeConvParams deconv_params = {kTfLitePaddingUnknown};
+        std::memcpy(&deconv_params, mapping_args.node->custom_initial_data,
+                    mapping_args.node->custom_initial_data_size);
+
+        mapping_args.builder->AddTensorInput(bias_tensor_id, hybrid_op, input_tensor_flags);    // 3
+        mapping_args.builder->AddScalarInt32Operand(0);                           // 4
+        mapping_args.builder->AddScalarInt32Operand(0);                           // 5
+        mapping_args.builder->AddScalarInt32Operand(0);                           // 6
+        mapping_args.builder->AddScalarInt32Operand(0);                           // 7
+        mapping_args.builder->AddScalarInt32Operand(deconv_params.stride_width);  // 8
+        mapping_args.builder->AddScalarInt32Operand(deconv_params.stride_height); // 9
+
+        // ANEURALNETWORKS_FUSED_NONE
+        mapping_args.builder->AddScalarInt32Operand(0);
+        // Use NHWC layout for input and output.
+        mapping_args.builder->AddScalarBoolOperand(false);
+
+        *nn_op_type = ANEURALNETWORKS_TRANSPOSE_CONV;
+      } else {
+        return kTfLiteError;
+      }
+    } break;
     default:
       // All other operators are not mapped.
       return kTfLiteError;
@@ -6157,6 +6202,8 @@ TfLiteStatus NNAPIDelegateKernel::AddOpsAndTensors(
               return kTfLiteError;
           }
         }
+      } else if (reg->builtin_code == kTfLiteBuiltinCustom) {
+        continue;
       } else {
         TF_LITE_ENSURE_STATUS(
             builder.AddTensorInput(input_index, hybrid_op, input_tensor_flags));