tflite: Add cros-mtk-pre-compile custom op support We will use this custom op to run the pre-compile node in the tflite. BUG=b:372342175 TEST=Run the mixed model Change-Id: I9abec6dacb28fa938602badab31d52440b983955 Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/platform/tflite/+/5802655 Reviewed-by: Pochun Lin <pochun.lin@mediatek.corp-partner.google.com> Reviewed-by: Shik Chen <shik@chromium.org> Commit-Queue: Yi Chou <yich@google.com> Tested-by: Yi Chou <yich@google.com>

diff --git a/WORKSPACE.bazel b/WORKSPACE.bazel
index 2067b42..7904a02 100644
--- a/WORKSPACE.bazel
+++ b/WORKSPACE.bazel

@@ -37,6 +37,7 @@
         "//patch:0014-xnnpack-disable-latest-ops.patch",
         "//patch:0015-cl-fast-relaxed-math.patch",
         "//patch:0016-dts-fp16-mode.patch",
+        "//patch:0017-Add-cros-mtk-pre-compile-custom-op.patch",
     ],
     sha256 = "9cc4d5773b8ee910079baaecb4086d0c28939f024dd74b33fc5e64779b6533dc",
     strip_prefix = "tensorflow-{}".format(TENSORFLOW_VERSION),

diff --git a/delegate/mtk_neuron/neuron_delegate_builder.h b/delegate/mtk_neuron/neuron_delegate_builder.h
index 16b9e2d..10e85b6 100644
--- a/delegate/mtk_neuron/neuron_delegate_builder.h
+++ b/delegate/mtk_neuron/neuron_delegate_builder.h

@@ -45,6 +45,8 @@
 
 static constexpr const char* kExtensionGeneralOpration =
     "com.mediatek.general_operation";
+static constexpr const char* kExtensionCompiledNetwork =
+    "com.mediatek.compiled_network";
 
 // Extension operand
 enum {
@@ -186,6 +188,43 @@
     return AddScalarOperand<float>(value, NEURON_FLOAT32);
   }
 
+  TfLiteStatus AddPreCompileExtensionOperand(const uint8_t* value,
+                                             uint32_t length) {
+    uint16_t dla_ext_operand_type = 0x0100;
+    NeuronRuntimeVersion version;
+    neuronapi_->Neuron_getVersion(&version);
+    if (version.major >= 8) {
+      dla_ext_operand_type = 0x0200;
+    }
+
+    int32_t operand_type;
+    RETURN_TFLITE_ERROR_IF_NEURON_ERROR(
+        context_,
+        neuronapi_->NeuronModel_getExtensionOperandType(
+            nn_model_, kExtensionCompiledNetwork, dla_ext_operand_type,
+            &operand_type),
+        "getting extension operand type");
+
+    NeuronOperandType operand_ext_type = {};
+    operand_ext_type.type = operand_type;
+    operand_ext_type.scale = 0.0f;
+    operand_ext_type.zeroPoint = 0;
+    operand_ext_type.dimensionCount = 0;
+
+    RETURN_TFLITE_ERROR_IF_NEURON_ERROR(
+        context_,
+        neuronapi_->NeuronModel_addOperand(nn_model_, &operand_ext_type),
+        "adding extension operand");
+    const int neuron_index = operand_mapping_->add_new_non_tensor_operand();
+    RETURN_TFLITE_ERROR_IF_NEURON_ERROR(
+        context_,
+        neuronapi_->NeuronModel_setOperandValue(nn_model_, neuron_index, value,
+                                                length),
+        "setting new extension operand value");
+    augmented_inputs_.push_back(neuron_index);
+    return kTfLiteOk;
+  }
+
   TfLiteStatus AddScalarExtensionOperand(const uint8_t* value,
                                          uint32_t length) {
     int32_t operand_type;
@@ -641,6 +680,18 @@
     return kTfLiteOk;
   }
 
+  TfLiteStatus AddPreCompileExtensionOperand(NeuronOperationType* type) {
+    int32_t operation_type = 0;
+    RETURN_TFLITE_ERROR_IF_NEURON_ERROR(
+        context_,
+        neuronapi_->NeuronModel_getExtensionOperationType(
+            nn_model_, kExtensionCompiledNetwork,
+            ADAPTER_EXTENSION_GENERAL_OPERATION_TYPE, &operation_type),
+        "getting extension operation type");
+    *type = static_cast<NeuronOperationType>(operation_type);
+    return kTfLiteOk;
+  }
+
   // Finish emitting the op (of type `type`) into the Neuron.
   TfLiteStatus FinalizeAddOperation(NeuronOperationType type) {
     // Actually add a Neuron operation

diff --git a/delegate/mtk_neuron/neuron_delegate_kernel.cc b/delegate/mtk_neuron/neuron_delegate_kernel.cc
index 1cb8817..1d89652 100644
--- a/delegate/mtk_neuron/neuron_delegate_kernel.cc
+++ b/delegate/mtk_neuron/neuron_delegate_kernel.cc

@@ -1145,6 +1145,18 @@
         mapping_args.builder->AddTensorOutput(outputs[3]);
         break;  // needed
       }
+      if (strcmp(TfLiteRegistrationExternalGetCustomName(
+                     mapping_args.registration),
+                 "cros-mtk-pre-compile") == 0) {
+        const void* init_data = nullptr;
+        int size = 0;
+        TF_LITE_ENSURE_STATUS(TfLiteOpaqueNodeGetCustomInitialData(
+            mapping_args.node, &init_data, &size));
+        mapping_args.builder->AddPreCompileExtensionOperand(
+            static_cast<const uint8_t*>(init_data), size);
+        mapping_args.builder->AddPreCompileExtensionOperand(nn_op_type);
+        break;  // needed
+      }
     }  // no break, pass to default
     default:
       // All other operators are not mapped.

diff --git a/delegate/mtk_neuron/neuron_delegate_validation.cc b/delegate/mtk_neuron/neuron_delegate_validation.cc
index 3092b7a..93c812f 100644
--- a/delegate/mtk_neuron/neuron_delegate_validation.cc
+++ b/delegate/mtk_neuron/neuron_delegate_validation.cc

@@ -922,6 +922,10 @@
     } break;
 #ifndef MTK_INTERNAL_USE
     case kTfLiteBuiltinCustom: {
+      if (strcmp(TfLiteRegistrationExternalGetCustomName(registration),
+                 "cros-mtk-pre-compile") == 0) {
+        break;
+      }
       int32_t magic_number = -1;
       // get NP magic number
       if (neuronapi->Neuron_getNeuroPilotMagicNumber != nullptr) {

diff --git a/patch/0017-Add-cros-mtk-pre-compile-custom-op.patch b/patch/0017-Add-cros-mtk-pre-compile-custom-op.patch
new file mode 100644
index 0000000..4394a56
--- /dev/null
+++ b/patch/0017-Add-cros-mtk-pre-compile-custom-op.patch

@@ -0,0 +1,37 @@
+From 460b2a9a6446913b1dc8a75987940dd21ef74fdc Mon Sep 17 00:00:00 2001
+From: Yi Chou <yich@google.com>
+Date: Fri, 23 Aug 2024 13:13:34 +0800
+Subject: [PATCH] Add cros-mtk-pre-compile custom op
+
+---
+ tensorflow/lite/tools/benchmark/benchmark_tflite_model.cc        | 1 +
+ .../lite/tools/evaluation/stages/tflite_inference_stage.cc       | 1 +
+ 2 files changed, 2 insertions(+)
+
+diff --git a/tensorflow/lite/tools/benchmark/benchmark_tflite_model.cc b/tensorflow/lite/tools/benchmark/benchmark_tflite_model.cc
+index b229a1b47c6..1ac054f6334 100644
+--- a/tensorflow/lite/tools/benchmark/benchmark_tflite_model.cc
++++ b/tensorflow/lite/tools/benchmark/benchmark_tflite_model.cc
+@@ -67,6 +67,7 @@ RegisterSelectedOps(::tflite::MutableOpResolver* resolver) {
+   static TfLiteRegistration reg = { nullptr, nullptr, nullptr, nullptr };
+   resolver->AddCustom("Convolution2DTransposeBias",
+                       &reg);
++  resolver->AddCustom("cros-mtk-pre-compile", &reg);
+ }
+ 
+ namespace tflite {
+diff --git a/tensorflow/lite/tools/evaluation/stages/tflite_inference_stage.cc b/tensorflow/lite/tools/evaluation/stages/tflite_inference_stage.cc
+index bd22f5d1f09..f68dccd75dd 100644
+--- a/tensorflow/lite/tools/evaluation/stages/tflite_inference_stage.cc
++++ b/tensorflow/lite/tools/evaluation/stages/tflite_inference_stage.cc
+@@ -142,6 +142,7 @@ TfLiteStatus TfliteInferenceStage::Init(
+   }
+   TfLiteRegistration reg = { nullptr, nullptr, nullptr, nullptr };
+   resolver_->AddCustom("Convolution2DTransposeBias", &reg);
++  resolver_->AddCustom("cros-mtk-pre-compile", &reg);
+   RegisterSelectedOps(resolver_.get());
+   InterpreterBuilder(*model_, *resolver_)(&interpreter_);
+   if (!interpreter_) {
+-- 
+2.46.0.295.g3b9ea8a38a-goog
+