Support Hexagon delegate to have an option to specify the max number of...

Support Hexagon delegate to have an option to specify the max number of partitions to be delegated, and changed the benchmark tool to support this feature.

PiperOrigin-RevId: 306560419
Change-Id: I75148012edb7109e84fee213720c66a2385a09cb

tensorflow/lite/experimental/delegates/hexagon/BUILD

@@ -82,6 +82,7 @@ cc_library(
         "//tensorflow/lite:kernel_api",
         "//tensorflow/lite:minimal_logging",
         "//tensorflow/lite/c:common",
+        "//tensorflow/lite/delegates:utils",
     ],
 )
 

tensorflow/lite/experimental/delegates/hexagon/hexagon_delegate.cc

@@ -20,6 +20,7 @@ limitations under the License.
 
 #include "tensorflow/lite/c/common.h"
 #include "tensorflow/lite/context_util.h"
+#include "tensorflow/lite/delegates/utils.h"
 #include "tensorflow/lite/experimental/delegates/hexagon/hexagon_delegate_kernel.h"
 #include "tensorflow/lite/experimental/delegates/hexagon/hexagon_implementation.h"
 #include "tensorflow/lite/experimental/delegates/hexagon/utils.h"
@@ -29,6 +30,8 @@ namespace tflite {
 namespace {
 // Should be > 0. > 16 causes problems.
 constexpr int kMaxHexagonGraphs = 4;
+constexpr int kMaxMaxHexagonGraphs = 16;
+constexpr int kMinNodesPerHexagonGraph = 2;
 
 TfLiteRegistration GetHexagonKernelRegistration() {
   // This is the registration for the Delegate Node that gets added to
@@ -80,50 +83,57 @@ TfLiteRegistration GetHexagonKernelRegistration() {
 }
 
 TfLiteStatus DelegatePrepare(TfLiteContext* context, TfLiteDelegate* delegate) {
-  // Reserve 1 element, since we need first element to be size, will be updated
-  // later.
+  delegates::IsNodeSupportedFn node_supported_fn =
+      [=](TfLiteContext* context, TfLiteNode* node,
+          TfLiteRegistration* registration,
+          std::string* unsupported_details) -> bool {
+    return IsNodeSupportedByHexagon(registration, node, context);
+  };
+  delegates::GraphPartitionHelper helper(context, node_supported_fn);
+  TF_LITE_ENSURE_STATUS(helper.Partition(nullptr));
+
+  TfLiteHexagonDelegateOptions* params =
+      static_cast<TfLiteHexagonDelegateOptions*>(delegate->data_);
+  const auto delegate_partitions = helper.GetFirstNLargestPartitions(
+      params->max_delegated_partitions, kMinNodesPerHexagonGraph);
+
+  // To avoid creating a new TfLiteIntArray and free it later, we reserve one
+  // element to represent TfLiteIntArray.size which is the 1st element of
+  // TfLiteIntArray C struct.
   std::vector<int> supported_nodes(1);
-  TfLiteIntArray* plan;
-  TF_LITE_ENSURE_STATUS(context->GetExecutionPlan(context, &plan));
-  TfLiteNode* node;
-  TfLiteRegistration* registration;
-
-  // Rudimentary mechanism to check how many Hexagon graphs we initialize.
-  int num_components = 1;
-  int last_index = -1;
-  for (int node_index : TfLiteIntArrayView(plan)) {
-    TF_LITE_ENSURE_STATUS(context->GetNodeAndRegistration(
-        context, node_index, &node, &registration));
-    if (IsNodeSupportedByHexagon(registration, node, context)) {
-      // If there is a 'break' in node indices, a new subgraph (and therefore, a
-      // new Hexagon graph) will be created.
-      if (last_index != -1 && node_index != last_index + 1) {
-        if (num_components == kMaxHexagonGraphs) {
-          break;
-        }
-        ++num_components;
-      }
-      supported_nodes.push_back(node_index);
-      last_index = node_index;
-    }
+  for (const auto partition : delegate_partitions) {
+    auto* nodes = partition->nodes_to_replace;
+    supported_nodes.insert(supported_nodes.end(), nodes->data,
+                           nodes->data + nodes->size);
   }
   // Set first element to the number of nodes to replace.
   supported_nodes[0] = supported_nodes.size() - 1;
   TFLITE_LOG_PROD(tflite::TFLITE_LOG_INFO,
-                  "Hexagon delegate: %d nodes delegated out of %d nodes.\n",
-                  supported_nodes[0], plan->size);
-  TfLiteRegistration hexagon_kernel_registration =
-      GetHexagonKernelRegistration();
+                  "Hexagon delegate: %d nodes delegated out of %d nodes with "
+                  "%d partitions.\n",
+                  supported_nodes[0], helper.num_total_nodes(),
+                  delegate_partitions.size());
 
   return context->ReplaceNodeSubsetsWithDelegateKernels(
-      context, hexagon_kernel_registration,
+      context, GetHexagonKernelRegistration(),
       reinterpret_cast<TfLiteIntArray*>(supported_nodes.data()), delegate);
 }
 
 class HexagonDelegate : public TfLiteDelegate {
  public:
   explicit HexagonDelegate(const TfLiteHexagonDelegateOptions* params)
-      : params_(params != nullptr ? *params : TfLiteHexagonDelegateOptions()) {}
+      : params_(params != nullptr ? *params
+                                  : TfLiteHexagonDelegateOptions({0})) {
+    if (params_.max_delegated_partitions <= 0) {
+      params_.max_delegated_partitions = kMaxHexagonGraphs;
+    } else if (params_.max_delegated_partitions > kMaxMaxHexagonGraphs) {
+      TFLITE_LOG_PROD(tflite::TFLITE_LOG_WARNING,
+                      "Hexagon delegate: cannot have this many %d partitions, "
+                      "and will cap to at most %d partitions.\n",
+                      params_.max_delegated_partitions, kMaxMaxHexagonGraphs);
+      params_.max_delegated_partitions = kMaxMaxHexagonGraphs;
+    }
+  }
 
   TfLiteHexagonDelegateOptions* params() { return &params_; }
 

tensorflow/lite/experimental/delegates/hexagon/hexagon_delegate.h

@@ -51,6 +51,10 @@ struct TFL_CAPI_EXPORT TfLiteHexagonDelegateOptions {
   // the DSP. Combining with 'debug_level' more information will be printed.
   // WARNING: Experimental and subject to change anytime.
   bool print_graph_debug;
+  // This sets the maximum number of Hexagon graphs created with
+  // hexagon_nn_init. Each graph corresponds to one delegated node subset in the
+  // TFLite model.
+  int max_delegated_partitions;
 };
 
 // Return a delegate that uses Hexagon SDK for ops execution.

tensorflow/lite/tools/benchmark/BUILD

@@ -298,7 +298,15 @@ cc_library(
     deps = [
         ":delegate_provider_hdr",
         "//tensorflow/lite/tools/evaluation:utils",
-    ],
+    ] + select({
+        "//tensorflow:android_arm": [
+            "//tensorflow/lite/experimental/delegates/hexagon:hexagon_delegate",
+        ],
+        "//tensorflow:android_arm64": [
+            "//tensorflow/lite/experimental/delegates/hexagon:hexagon_delegate",
+        ],
+        "//conditions:default": [],
+    }),
     alwayslink = 1,
 )
 

tensorflow/lite/tools/benchmark/hexagon_delegate_provider.cc

@@ -22,6 +22,10 @@ limitations under the License.
 #define TFLITE_ENABLE_HEXAGON
 #endif
 
+#if defined(TFLITE_ENABLE_HEXAGON)
+#include "tensorflow/lite/experimental/delegates/hexagon/hexagon_delegate.h"
+#endif
+
 namespace tflite {
 namespace benchmark {
 
@@ -82,9 +86,12 @@ TfLiteDelegatePtr HexagonDelegateProvider::CreateTfLiteDelegate(
   TfLiteDelegatePtr delegate(nullptr, [](TfLiteDelegate*) {});
 #if defined(TFLITE_ENABLE_HEXAGON)
   if (params.Get<bool>("use_hexagon")) {
+    TfLiteHexagonDelegateOptions options = {0};
+    options.print_graph_profile = params.Get<bool>("hexagon_profiling");
+    options.max_delegated_partitions =
+        params.Get<int>("max_delegated_partitions");
     delegate = evaluation::CreateHexagonDelegate(
-        params.Get<std::string>("hexagon_lib_path"),
-        params.Get<bool>("hexagon_profiling"));
+        &options, params.Get<std::string>("hexagon_lib_path"));
 
     if (!delegate.get()) {
       TFLITE_LOG(WARN)

tensorflow/lite/tools/evaluation/utils.cc

@@ -139,16 +139,25 @@ TfLiteDelegatePtr CreateGPUDelegate() {
 TfLiteDelegatePtr CreateHexagonDelegate(
     const std::string& library_directory_path, bool profiling) {
 #if defined(__ANDROID__) && (defined(__arm__) || defined(__aarch64__))
+  TfLiteHexagonDelegateOptions options = {0};
+  options.print_graph_profile = profiling;
+  return CreateHexagonDelegate(&options, library_directory_path);
+#else
+  return CreateNullDelegate();
+#endif  // defined(__ANDROID__)
+}
+
+#if defined(__ANDROID__) && (defined(__arm__) || defined(__aarch64__))
+TfLiteDelegatePtr CreateHexagonDelegate(
+    const TfLiteHexagonDelegateOptions* options,
+    const std::string& library_directory_path) {
   if (library_directory_path.empty()) {
     TfLiteHexagonInit();
   } else {
     TfLiteHexagonInitWithPath(library_directory_path.c_str());
   }
 
-  const TfLiteHexagonDelegateOptions options = {
-      /*debug_level=*/0, /*powersave_level=*/0, profiling,
-      /*print_graph_debug=*/false};
-  TfLiteDelegate* delegate = TfLiteHexagonDelegateCreate(&options);
+  TfLiteDelegate* delegate = TfLiteHexagonDelegateCreate(options);
   if (!delegate) {
     TfLiteHexagonTearDown();
     return CreateNullDelegate();
@@ -157,10 +166,8 @@ TfLiteDelegatePtr CreateHexagonDelegate(
     TfLiteHexagonDelegateDelete(delegate);
     TfLiteHexagonTearDown();
   });
-#else
-  return CreateNullDelegate();
-#endif  // defined(__ANDROID__)
 }
+#endif
 
 // TODO(b/149248802): include XNNPACK delegate when the issue is resolved.
 #if defined(__Fuchsia__) || defined(TFLITE_WITHOUT_XNNPACK)

tensorflow/lite/tools/evaluation/utils.h

@@ -71,6 +71,11 @@ TfLiteDelegatePtr CreateGPUDelegate(TfLiteGpuDelegateOptionsV2* options);
 
 TfLiteDelegatePtr CreateHexagonDelegate(
     const std::string& library_directory_path, bool profiling);
+#if defined(__ANDROID__) && (defined(__arm__) || defined(__aarch64__))
+TfLiteDelegatePtr CreateHexagonDelegate(
+    const TfLiteHexagonDelegateOptions* options,
+    const std::string& library_directory_path);
+#endif
 
 // TODO(b/149248802): include XNNPACK delegate when the issue is resolved.
 #if !defined(__Fuchsia__) || defined(TFLITE_WITHOUT_XNNPACK)