Project import generated by Copybara.

GitOrigin-RevId: 27c70b5fe62ab71189d358ca122ee4b19c817a8f

mediapipe/calculators/tensor/inference_calculator.cc

@@ -33,7 +33,7 @@ class InferenceCalculatorSelectorImpl
   absl::StatusOr<CalculatorGraphConfig> GetConfig(
       const CalculatorGraphConfig::Node& subgraph_node) {
     const auto& options =
-        Subgraph::GetOptions<::mediapipe::InferenceCalculatorOptions>(
+        Subgraph::GetOptions<mediapipe::InferenceCalculatorOptions>(
             subgraph_node);
     std::vector<absl::string_view> impls;
     const bool should_use_gpu =

mediapipe/calculators/tensor/inference_calculator.h

@@ -99,8 +99,13 @@ class InferenceCalculator : public NodeIntf {
       kSideInCustomOpResolver{"CUSTOM_OP_RESOLVER"};
   static constexpr SideInput<TfLiteModelPtr>::Optional kSideInModel{"MODEL"};
   static constexpr Output<std::vector<Tensor>> kOutTensors{"TENSORS"};
+  static constexpr SideInput<std::string>::Optional kNnApiDelegateCacheDir{
+      "NNAPI_CACHE_DIR"};
+  static constexpr SideInput<std::string>::Optional kNnApiDelegateModelToken{
+      "NNAPI_MODEL_TOKEN"};
   MEDIAPIPE_NODE_CONTRACT(kInTensors, kSideInCustomOpResolver, kSideInModel,
-                          kOutTensors);
+                          kOutTensors, kNnApiDelegateCacheDir,
+                          kNnApiDelegateModelToken);
 
  protected:
   using TfLiteDelegatePtr =

mediapipe/calculators/tensor/inference_calculator.proto

@@ -67,9 +67,32 @@ message InferenceCalculatorOptions {
       // Only available for OpenCL delegate on Android.
       // Kernel caching will only be enabled if this path is set.
       optional string cached_kernel_path = 2;
+
+      // Encapsulated compilation/runtime tradeoffs.
+      enum InferenceUsage {
+        UNSPECIFIED = 0;
+
+        // InferenceRunner will be used only once. Therefore, it is important to
+        // minimize bootstrap time as well.
+        FAST_SINGLE_ANSWER = 1;
+
+        // Prefer maximizing the throughput. Same inference runner will be used
+        // repeatedly on different inputs.
+        SUSTAINED_SPEED = 2;
+      }
+      optional InferenceUsage usage = 5 [default = SUSTAINED_SPEED];
     }
+
     // Android only.
-    message Nnapi {}
+    message Nnapi {
+      // Directory to store compilation cache. If unspecified, NNAPI will not
+      // try caching the compilation.
+      optional string cache_dir = 1;
+      // Unique token identifying the model. It is the caller's responsibility
+      // to ensure there is no clash of the tokens. If unspecified, NNAPI will
+      // not try caching the compilation.
+      optional string model_token = 2;
+    }
     message Xnnpack {
       // Number of threads for XNNPACK delegate. (By default, calculator tries
       // to choose optimal number of threads depending on the device.)

mediapipe/calculators/tensor/inference_calculator_cpu.cc

@@ -181,9 +181,21 @@ absl::Status InferenceCalculatorCpuImpl::LoadDelegate(CalculatorContext* cc) {
     // Attempt to use NNAPI.
     // If not supported, the default CPU delegate will be created and used.
     interpreter_->SetAllowFp16PrecisionForFp32(1);
-    delegate_ = TfLiteDelegatePtr(tflite::NnApiDelegate(), [](TfLiteDelegate*) {
-      // No need to free according to tflite::NnApiDelegate() documentation.
-    });
+    tflite::StatefulNnApiDelegate::Options options;
+    const auto& nnapi = calculator_opts.delegate().nnapi();
+    // Set up cache_dir and model_token for NNAPI compilation cache.
+    options.cache_dir =
+        nnapi.has_cache_dir() ? nnapi.cache_dir().c_str() : nullptr;
+    if (!kNnApiDelegateCacheDir(cc).IsEmpty()) {
+      options.cache_dir = kNnApiDelegateCacheDir(cc).Get().c_str();
+    }
+    options.model_token =
+        nnapi.has_model_token() ? nnapi.model_token().c_str() : nullptr;
+    if (!kNnApiDelegateModelToken(cc).IsEmpty()) {
+      options.model_token = kNnApiDelegateModelToken(cc).Get().c_str();
+    }
+    delegate_ = TfLiteDelegatePtr(new tflite::StatefulNnApiDelegate(options),
+                                  [](TfLiteDelegate*) {});
     RET_CHECK_EQ(interpreter_->ModifyGraphWithDelegate(delegate_.get()),
                  kTfLiteOk);
     return absl::OkStatus();

mediapipe/calculators/tensor/inference_calculator_gl.cc

@@ -18,6 +18,7 @@
 #include <vector>
 
 #include "absl/memory/memory.h"
+#include "absl/status/status.h"
 #include "mediapipe/calculators/tensor/inference_calculator.h"
 #include "mediapipe/util/tflite/config.h"
 
@@ -65,6 +66,8 @@ class InferenceCalculatorGlImpl
   bool allow_precision_loss_ = false;
   mediapipe::InferenceCalculatorOptions::Delegate::Gpu::Api
       tflite_gpu_runner_api_;
+  mediapipe::InferenceCalculatorOptions::Delegate::Gpu::InferenceUsage
+      tflite_gpu_runner_usage_;
 #endif  // MEDIAPIPE_TFLITE_GL_INFERENCE
 
 #if MEDIAPIPE_TFLITE_GPU_SUPPORTED
@@ -96,6 +99,7 @@ absl::Status InferenceCalculatorGlImpl::Open(CalculatorContext* cc) {
                           options.delegate().gpu().use_advanced_gpu_api();
   allow_precision_loss_ = options.delegate().gpu().allow_precision_loss();
   tflite_gpu_runner_api_ = options.delegate().gpu().api();
+  tflite_gpu_runner_usage_ = options.delegate().gpu().usage();
   use_kernel_caching_ = use_advanced_gpu_api_ &&
                         options.delegate().gpu().has_cached_kernel_path();
   use_gpu_delegate_ = !use_advanced_gpu_api_;
@@ -253,9 +257,27 @@ absl::Status InferenceCalculatorGlImpl::InitTFLiteGPURunner(
                           : tflite::gpu::InferencePriority::MAX_PRECISION;
   options.priority2 = tflite::gpu::InferencePriority::AUTO;
   options.priority3 = tflite::gpu::InferencePriority::AUTO;
-  options.usage = tflite::gpu::InferenceUsage::SUSTAINED_SPEED;
+  switch (tflite_gpu_runner_usage_) {
+    case mediapipe::InferenceCalculatorOptions::Delegate::Gpu::
+        FAST_SINGLE_ANSWER: {
+      options.usage = tflite::gpu::InferenceUsage::FAST_SINGLE_ANSWER;
+      break;
+    }
+    case mediapipe::InferenceCalculatorOptions::Delegate::Gpu::
+        SUSTAINED_SPEED: {
+      options.usage = tflite::gpu::InferenceUsage::SUSTAINED_SPEED;
+      break;
+    }
+    case mediapipe::InferenceCalculatorOptions::Delegate::Gpu::UNSPECIFIED: {
+      return absl::InternalError("inference usage need to be specified.");
+    }
+  }
   tflite_gpu_runner_ = std::make_unique<tflite::gpu::TFLiteGPURunner>(options);
   switch (tflite_gpu_runner_api_) {
+    case mediapipe::InferenceCalculatorOptions::Delegate::Gpu::ANY: {
+      // Do not need to force any specific API.
+      break;
+    }
     case mediapipe::InferenceCalculatorOptions::Delegate::Gpu::OPENGL: {
       tflite_gpu_runner_->ForceOpenGL();
       break;
@@ -264,10 +286,6 @@ absl::Status InferenceCalculatorGlImpl::InitTFLiteGPURunner(
       tflite_gpu_runner_->ForceOpenCL();
       break;
     }
-    case mediapipe::InferenceCalculatorOptions::Delegate::Gpu::ANY: {
-      // Do not need to force any specific API.
-      break;
-    }
   }
   MP_RETURN_IF_ERROR(tflite_gpu_runner_->InitializeWithModel(
       model, op_resolver, /*allow_quant_ops=*/true));