- Update SimpleDelegate to allow providing options for the delegate which controls delegate max_number of graphs/ min nodes per subgraph.

- Update Hexagon Delegate to use SimpleDelegate Interface.

PiperOrigin-RevId: 319064977
Change-Id: Icdeeba44f24109fd09f26e372f9668a8c00cedaf

tensorflow/lite/delegates/flex/delegate.h

@@ -77,6 +77,11 @@ class FlexDelegate : public SimpleDelegateInterface {
 
   TfLiteStatus Initialize(TfLiteContext* context) override;
 
+  SimpleDelegateInterface::Options DelegateOptions() const override {
+    // Use default options.
+    return SimpleDelegateInterface::Options();
+  }
+
   std::unique_ptr<SimpleDelegateKernelInterface> CreateDelegateKernelInterface()
       override;
 

tensorflow/lite/delegates/hexagon/BUILD

@@ -57,9 +57,9 @@ cc_library(
         "//tensorflow/lite:kernel_api",
         "//tensorflow/lite/c:common",
         "//tensorflow/lite/core/api",
-        "//tensorflow/lite/delegates:utils",
         "//tensorflow/lite/delegates/hexagon/builders:op_builder",
         "//tensorflow/lite/delegates/hexagon/hexagon_nn:hexagon_nn_header",
+        "//tensorflow/lite/delegates/utils:simple_delegate",
         "//tensorflow/lite/kernels:kernel_util",
         "//tensorflow/lite/schema:schema_fbs",
         "@hexagon_nn//:hexagon_nn_ops",
@@ -81,7 +81,7 @@ cc_library(
         "//tensorflow/lite:kernel_api",
         "//tensorflow/lite:minimal_logging",
         "//tensorflow/lite/c:common",
-        "//tensorflow/lite/delegates:utils",
+        "//tensorflow/lite/delegates/utils:simple_delegate",
     ],
 )
 

tensorflow/lite/delegates/hexagon/hexagon_delegate.cc

@@ -20,10 +20,10 @@ 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/delegates/hexagon/hexagon_delegate_kernel.h"
 #include "tensorflow/lite/delegates/hexagon/hexagon_implementation.h"
 #include "tensorflow/lite/delegates/hexagon/utils.h"
+#include "tensorflow/lite/delegates/utils/simple_delegate.h"
 #include "tensorflow/lite/minimal_logging.h"
 
 namespace tflite {
@@ -33,56 +33,7 @@ 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
-  // the TFLite graph instead of the subGraph it replaces it.
-  // It is treated as a an OP node. But in our case
-  // Init will initialize the delegate
-  // Invoke will run the delegate graph.
-  // Prepare for prearing the delegate.
-  // Free for any cleaning needed by the delegate.
-  TfLiteRegistration kernel_registration;
-  kernel_registration.profiling_string = nullptr;
-  kernel_registration.builtin_code = kTfLiteBuiltinDelegate;
-  kernel_registration.custom_name = "TfLiteHexagonDelegate";
-  kernel_registration.free = [](TfLiteContext* context, void* buffer) -> void {
-    delete reinterpret_cast<HexagonDelegateKernel*>(buffer);
-  };
-  kernel_registration.init = [](TfLiteContext* context, const char* buffer,
-                                size_t length) -> void* {
-    const TfLiteDelegateParams* params =
-        reinterpret_cast<const TfLiteDelegateParams*>(buffer);
-    auto hexagon_kernel = std::make_unique<HexagonDelegateKernel>();
-    if (hexagon_kernel->Init(context, params) != kTfLiteOk) {
-      return nullptr;
-    }
-    return hexagon_kernel.release();
-  };
-  kernel_registration.invoke = [](TfLiteContext* context,
-                                  TfLiteNode* node) -> TfLiteStatus {
-    HexagonDelegateKernel* kernel =
-        reinterpret_cast<HexagonDelegateKernel*>(node->user_data);
-    if (!kernel) {
-      context->ReportError(context, "Hexagon Kernel was not initialized");
-      return kTfLiteError;
-    }
-    return kernel->Invoke(context, node);
-  };
-  kernel_registration.prepare = [](TfLiteContext* context,
-                                   TfLiteNode* node) -> TfLiteStatus {
-    if (node->user_data == nullptr) {
-      context->ReportError(context, "Hexagon Kernel was not initialized");
-      return kTfLiteError;
-    }
-    HexagonDelegateKernel* kernel =
-        reinterpret_cast<HexagonDelegateKernel*>(node->user_data);
-    return kernel->Prepare(context, node);
-  };
-
-  return kernel_registration;
-}
-
-class HexagonDelegate : public TfLiteDelegate {
+class HexagonDelegate : public SimpleDelegateInterface {
  public:
   explicit HexagonDelegate(const TfLiteHexagonDelegateOptions* params)
       : params_(params != nullptr ? *params
@@ -101,113 +52,40 @@ class HexagonDelegate : public TfLiteDelegate {
     }
   }
 
-  TfLiteHexagonDelegateOptions* params() { return &params_; }
-
-  bool VerifyDelegate() {
-    auto* hexagon_nn = HexagonNNImplementation();
-    if (hexagon_nn == nullptr) {
-      return false;
-    }
-    if (hexagon_nn->hexagon_nn_version != nullptr &&
-        hexagon_nn->hexagon_nn_hexagon_interface_version) {
-      int hexagon_nn_version = -1;
-      int hexagon_interface_version =
-          hexagon_nn->hexagon_nn_hexagon_interface_version();
-      if (hexagon_nn->hexagon_nn_version(&hexagon_nn_version) != 0) {
-        TFLITE_LOG_PROD(tflite::TFLITE_LOG_WARNING,
-                        "Failed to fetch Hexagon NN version. This might be "
-                        "because you're using incompatible versions of "
-                        "libhexagon_interface and libhexagon_nn_skel. "
-                        "You must use compatible versions. "
-                        "Refer to Tensorflow Lite Hexagon Delegate Guide.");
-        return false;
-      }
-      if (hexagon_nn_version != hexagon_interface_version) {
-        TFLITE_LOG_PROD(
-            tflite::TFLITE_LOG_WARNING,
-            "Incompatible versions between interface library and "
-            "libhexagon_skel %d vs %d. You must use compatible versions. "
-            "Refer to Tensorflow Lite Hexagon Delegate Guide.",
-            hexagon_interface_version, hexagon_nn_version);
-        return false;
-      }
-    }
-    return hexagon_nn->hexagon_nn_is_device_supported &&
-           hexagon_nn->hexagon_nn_is_device_supported();
+  bool IsNodeSupportedByDelegate(const TfLiteRegistration* registration,
+                                 const TfLiteNode* node,
+                                 TfLiteContext* context) const override {
+    return IsNodeSupportedByHexagon(registration, node, context);
   }
 
-  ~HexagonDelegate() {
-    TfLiteIntArrayFree(params_.input_batch_dimensions);
-    TfLiteIntArrayFree(params_.output_batch_dimensions);
+  TfLiteStatus Initialize(TfLiteContext* context) override { return kTfLiteOk; }
+
+  const char* Name() const override { return "TfLiteHexagonDelegate"; }
+
+  std::unique_ptr<SimpleDelegateKernelInterface> CreateDelegateKernelInterface()
+      override {
+    return std::make_unique<HexagonDelegateKernel>(params_);
+  }
+
+  SimpleDelegateInterface::Options DelegateOptions() const override {
+    auto options = SimpleDelegateInterface::Options();
+    options.max_delegated_partitions = params_.max_delegated_partitions;
+    options.min_nodes_per_partition = params_.min_nodes_per_partition;
+    return options;
   }
 
  private:
   TfLiteHexagonDelegateOptions params_;
 };
 
-TfLiteStatus DelegatePrepare(TfLiteContext* context, TfLiteDelegate* delegate) {
-  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_);
-  std::vector<int> supported_nodes = helper.GetNodesOfFirstNLargestPartitions(
-      params->max_delegated_partitions, params->min_nodes_per_partition);
-
-  auto* hexagon_delegate = static_cast<HexagonDelegate*>(delegate);
-  // Make sure dynamic batch is requested on fully delegated graph only.
-  if (supported_nodes.size() != helper.num_total_nodes() &&
-      hexagon_delegate != nullptr &&
-      hexagon_delegate->params()->enable_dynamic_batch_size) {
-    TF_LITE_KERNEL_LOG(
-        context, "Dynamic batch requested on non-fully delegated graph !!.");
-    return kTfLiteError;
-  }
-  TFLITE_LOG_PROD(tflite::TFLITE_LOG_INFO,
-                  "Hexagon delegate: %d nodes delegated out of %d nodes with "
-                  "%d partitions.\n",
-                  supported_nodes.size(), helper.num_total_nodes(),
-                  helper.num_partitions());
-
-  return context->ReplaceNodeSubsetsWithDelegateKernels(
-      context, GetHexagonKernelRegistration(),
-      BuildTfLiteIntArray(supported_nodes).get(), delegate);
-}
-
-TfLiteDelegate* CreateDelegate(const TfLiteHexagonDelegateOptions* params) {
-  TfLiteDelegate* delegate = new HexagonDelegate(params);
-  if (!static_cast<HexagonDelegate*>(delegate)->VerifyDelegate()) {
-    delete delegate;
-    TFLITE_LOG_PROD_ONCE(tflite::TFLITE_LOG_INFO,
-                         "Hexagon Delegate is not supported.\n");
-    return nullptr;
-  }
-
-  delegate->data_ = static_cast<HexagonDelegate*>(delegate)->params();
-  delegate->flags = kTfLiteDelegateFlagsAllowDynamicTensors;
-  delegate->Prepare = &DelegatePrepare;
-  delegate->CopyFromBufferHandle = nullptr;
-  delegate->CopyToBufferHandle = nullptr;
-  delegate->FreeBufferHandle = nullptr;
-
-  TFLITE_LOG_PROD_ONCE(tflite::TFLITE_LOG_INFO,
-                       "Created TensorFlow Lite delegate for Hexagon.");
-
-  return delegate;
-}
-
 }  // namespace
 }  // namespace tflite
 
 TfLiteDelegate* TfLiteHexagonDelegateCreate(
     const TfLiteHexagonDelegateOptions* options) {
-  return tflite::CreateDelegate(options);
+  // return tflite::CreateDelegate(options);
+  return tflite::TfLiteDelegateFactory::CreateSimpleDelegate(
+      std::make_unique<tflite::HexagonDelegate>(options));
 }
 
 TfLiteHexagonDelegateOptions TfLiteHexagonDelegateOptionsDefault() {
@@ -215,7 +93,9 @@ TfLiteHexagonDelegateOptions TfLiteHexagonDelegateOptionsDefault() {
   return result;
 }
 
-void TfLiteHexagonDelegateDelete(TfLiteDelegate* delegate) { delete delegate; }
+void TfLiteHexagonDelegateDelete(TfLiteDelegate* delegate) {
+  tflite::TfLiteDelegateFactory::DeleteSimpleDelegate(delegate);
+}
 
 void TfLiteHexagonInit() { tflite::HexagonDelegateKernel::InitState(); }
 

tensorflow/lite/delegates/hexagon/hexagon_delegate_kernel.cc

@@ -20,7 +20,6 @@ limitations under the License.
 #include "tensorflow/lite/c/builtin_op_data.h"
 #include "tensorflow/lite/c/common.h"
 #include "tensorflow/lite/context_util.h"
-#include "tensorflow/lite/delegates/utils.h"
 #include "tensorflow/lite/delegates/hexagon/hexagon_implementation.h"
 #include "tensorflow/lite/delegates/hexagon/utils.h"
 #include "tensorflow/lite/kernels/kernel_util.h"
@@ -51,13 +50,6 @@ TfLiteStatus HexagonDelegateKernel::Init(TfLiteContext* context,
     TF_LITE_KERNEL_LOG(context, "Hexagon interface not available.");
     return kTfLiteError;
   }
-  if (params != nullptr && params->delegate != nullptr) {
-    const ::TfLiteHexagonDelegateOptions* options_ptr =
-        reinterpret_cast<const ::TfLiteHexagonDelegateOptions*>(
-            params->delegate->data_);
-    params_ = (options_ptr == nullptr ? ::TfLiteHexagonDelegateOptions()
-                                      : *options_ptr);
-  }
 
   // Ensure Hexagon NNLib is ready to start working.
   int error = hexagon_nn_->hexagon_nn_config();
@@ -94,8 +86,8 @@ TfLiteStatus HexagonDelegateKernel::Init(TfLiteContext* context,
   return kTfLiteOk;
 }
 
-TfLiteStatus HexagonDelegateKernel::Invoke(TfLiteContext* context,
-                                           TfLiteNode* node) {
+TfLiteStatus HexagonDelegateKernel::Eval(TfLiteContext* context,
+                                         TfLiteNode* node) {
   if (hexagon_nn_ == nullptr) {
     TF_LITE_KERNEL_LOG(context, "Hexagon interface not available.");
     return kTfLiteError;

tensorflow/lite/delegates/hexagon/hexagon_delegate_kernel.h

@@ -31,6 +31,7 @@ limitations under the License.
 #include "tensorflow/lite/delegates/hexagon/hexagon_delegate.h"
 #include "tensorflow/lite/delegates/hexagon/hexagon_implementation.h"
 #include "tensorflow/lite/delegates/hexagon/hexagon_nn/hexagon_nn.h"
+#include "tensorflow/lite/delegates/utils/simple_delegate.h"
 #include "tensorflow/lite/schema/schema_generated.h"
 
 namespace tflite {
@@ -38,18 +39,22 @@ namespace tflite {
 // Represents an abstraction of a Hexagon NNLib graph with functionality to
 // initialize, prepare and invoke it based on the TFLite subgraph to be
 // delegated.
-class HexagonDelegateKernel {
+class HexagonDelegateKernel : public SimpleDelegateKernelInterface {
  public:
+  explicit HexagonDelegateKernel(const ::TfLiteHexagonDelegateOptions& params)
+      : params_(params) {}
+
   // Initialize the Hexagon graph and add required nodes.
-  TfLiteStatus Init(TfLiteContext* context, const TfLiteDelegateParams* params);
+  TfLiteStatus Init(TfLiteContext* context,
+                    const TfLiteDelegateParams* params) override;
 
   // Prepare the Hexagon graph with hexagon_nn_prepare.
-  TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node);
+  TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) override;
 
   // Allocate Hexagon tensordefs for graph I/O & execute it.
-  TfLiteStatus Invoke(TfLiteContext* context, TfLiteNode* node);
+  TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) override;
 
-  ~HexagonDelegateKernel();
+  ~HexagonDelegateKernel() override;
 
   // Sets the environment required for Hexagon execution: DSP attributes,
   // rpcmem, etc.

tensorflow/lite/delegates/utils/dummy_delegate/dummy_delegate.cc

@@ -68,6 +68,11 @@ class DummyDelegate : public SimpleDelegateInterface {
     return std::make_unique<DummyDelegateKernel>(options_);
   }
 
+  SimpleDelegateInterface::Options DelegateOptions() const override {
+    // Use default options.
+    return SimpleDelegateInterface::Options();
+  }
+
  private:
   const DummyDelegateOptions options_;
 };

tensorflow/lite/delegates/utils/simple_delegate.cc

@@ -14,6 +14,7 @@ limitations under the License.
 ==============================================================================*/
 #include "tensorflow/lite/delegates/utils/simple_delegate.h"
 
+#include <limits>
 #include <memory>
 #include <vector>
 
@@ -78,6 +79,10 @@ TfLiteStatus DelegatePrepare(TfLiteContext* context,
                              TfLiteDelegate* base_delegate) {
   auto* delegate =
       reinterpret_cast<SimpleDelegateInterface*>(base_delegate->data_);
+  auto delegate_options = delegate->DelegateOptions();
+  if (delegate_options.max_delegated_partitions <= 0)
+    delegate_options.max_delegated_partitions = std::numeric_limits<int>::max();
+
   TF_LITE_ENSURE_STATUS(delegate->Initialize(context));
   delegates::IsNodeSupportedFn node_supported_fn =
       [=](TfLiteContext* context, TfLiteNode* node,
@@ -89,7 +94,9 @@ TfLiteStatus DelegatePrepare(TfLiteContext* context,
   delegates::GraphPartitionHelper helper(context, node_supported_fn);
   TF_LITE_ENSURE_STATUS(helper.Partition(nullptr));
 
-  std::vector<int> supported_nodes = helper.GetNodesOfFirstNLargestPartitions();
+  std::vector<int> supported_nodes = helper.GetNodesOfFirstNLargestPartitions(
+      delegate_options.max_delegated_partitions,
+      delegate_options.min_nodes_per_partition);
 
   TFLITE_LOG_PROD(tflite::TFLITE_LOG_INFO,
                   "%s delegate: %d nodes delegated out of %d nodes with "

tensorflow/lite/delegates/utils/simple_delegate.h

@@ -70,7 +70,15 @@ class SimpleDelegateKernelInterface {
 // - CreateDelegateKernelInterface
 class SimpleDelegateInterface {
  public:
-  SimpleDelegateInterface() {}
+  // Options for configuring a delegate.
+  struct Options {
+    // Maximum number of delegated subgraph, values <=0 means unlimited.
+    int max_delegated_partitions = 0;
+
+    // The minimum number of nodes allowed in a delegated graph, values <=0
+    // means unlimited.
+    int min_nodes_per_partition = 0;
+  };
 
   virtual ~SimpleDelegateInterface() {}
 
@@ -95,6 +103,9 @@ class SimpleDelegateInterface {
   // Caller takes ownership of the returned object.
   virtual std::unique_ptr<SimpleDelegateKernelInterface>
   CreateDelegateKernelInterface() = 0;
+
+  // Returns SimpleDelegateInterface::Options which has the delegate options.
+  virtual SimpleDelegateInterface::Options DelegateOptions() const = 0;
 };
 
 // Factory class that provides static methods to deal with SimpleDelegate