Experimental: Allow users to enable all features of default delegates in python

XNNPack is the only default delegate and this will allow users to benefit from all flag protected features. Other delegates can use this in the future.

PiperOrigin-RevId: 618111218

Author: alankelly

RELEASE.md

@@ -75,6 +75,9 @@
     * C API:
         * The experimental `TfLiteRegistrationExternal` type has been renamed as
           `TfLiteOperator`, and likewise for the corresponding API functions.
+    * The Python TF Lite Interpreter bindings now have an option
+      `experimental_default_delegate_latest_features` to enable all default
+      delegate features.
 
 ## Thanks to our Contributors
 

tensorflow/lite/python/interpreter.py

@@ -396,6 +396,7 @@ def __init__(
       experimental_op_resolver_type=OpResolverType.AUTO,
       experimental_preserve_all_tensors=False,
       experimental_disable_delegate_clustering=False,
+      experimental_default_delegate_latest_features=False,
   ):
     """Constructor.
 
@@ -437,6 +438,8 @@ def __init__(
         this flag is currently experimental, and it might be removed/updated if
         the TF Lite converter doesn't drop such control dependencies in the
         model. Default is False.
+      experimental_default_delegate_latest_features: If true, default delegates
+        may enable all flag protected features. Default is False;
 
     Raises:
       ValueError: If the interpreter was unable to create.
@@ -454,6 +457,12 @@ def __init__(
       raise ValueError('Unrecognized passed in op resolver type: {}'.format(
           experimental_op_resolver_type))
 
+    if num_threads is not None:
+      if not isinstance(num_threads, int):
+        raise ValueError('type of num_threads should be int')
+      if num_threads < 1:
+        raise ValueError('num_threads should >= 1')
+
     if model_path and not model_content:
       custom_op_registerers_by_name = [
           x for x in self._custom_op_registerers if isinstance(x, str)
@@ -468,6 +477,8 @@ def __init__(
           custom_op_registerers_by_func,
           experimental_preserve_all_tensors,
           experimental_disable_delegate_clustering,
+          int(num_threads or 1),
+          experimental_default_delegate_latest_features,
       )
       if not self._interpreter:
         raise ValueError('Failed to open {}'.format(model_path))
@@ -489,19 +500,14 @@ def __init__(
           custom_op_registerers_by_func,
           experimental_preserve_all_tensors,
           experimental_disable_delegate_clustering,
+          int(num_threads or 1),
+          experimental_default_delegate_latest_features,
       )
     elif not model_content and not model_path:
       raise ValueError('`model_path` or `model_content` must be specified.')
     else:
       raise ValueError('Can\'t both provide `model_path` and `model_content`')
 
-    if num_threads is not None:
-      if not isinstance(num_threads, int):
-        raise ValueError('type of num_threads should be int')
-      if num_threads < 1:
-        raise ValueError('num_threads should >= 1')
-      self._interpreter.SetNumThreads(num_threads)
-
     # Each delegate is a wrapper that owns the delegates that have been loaded
     # as plugins. The interpreter wrapper will be using them, but we need to
     # hold them in a list so that the lifetime is preserved at least as long as

tensorflow/lite/python/interpreter_wrapper/BUILD

@@ -39,6 +39,7 @@ cc_library(
         "//tensorflow/lite/core/api",
         "//tensorflow/lite/core/c:common",
         "//tensorflow/lite/core/kernels:builtin_ops",
+        "//tensorflow/lite/delegates/xnnpack:xnnpack_delegate",
         "//tensorflow/lite/kernels:reference_ops",
         "//tensorflow/lite/kernels/internal:compatibility",
         "//third_party/python_runtime:headers",  # buildcleaner: keep

tensorflow/lite/python/interpreter_wrapper/interpreter_wrapper.cc

@@ -32,6 +32,7 @@ limitations under the License.
 #include "tensorflow/lite/core/interpreter.h"
 #include "tensorflow/lite/core/kernels/register.h"
 #include "tensorflow/lite/core/model.h"
+#include "tensorflow/lite/delegates/xnnpack/xnnpack_delegate.h"
 #include "tensorflow/lite/kernels/internal/compatibility.h"
 #include "tensorflow/lite/kernels/register_ref.h"
 #include "tensorflow/lite/mutable_op_resolver.h"
@@ -85,18 +86,31 @@ using python_utils::PyDecrefDeleter;
 std::unique_ptr<Interpreter> CreateInterpreter(
     const InterpreterWrapper::Model* model,
     const tflite::MutableOpResolver& resolver, bool preserve_all_tensors,
-    bool disable_delegate_clustering) {
+    bool disable_delegate_clustering, int num_threads,
+    bool default_delegate_latest_features) {
   if (!model) {
     return nullptr;
   }
 
   ::tflite::python::ImportNumpy();
 
+  TfLiteDelegate* xnnpack_delegate = nullptr;
+  if (default_delegate_latest_features) {
+    auto opts = TfLiteXNNPackDelegateOptionsDefault();
+    opts.flags |= TFLITE_XNNPACK_DELEGATE_FLAG_ENABLE_LATEST_OPERATORS;
+    opts.flags |= TFLITE_XNNPACK_DELEGATE_FLAG_ENABLE_SUBGRAPH_RESHAPING;
+    opts.num_threads = num_threads;
+    xnnpack_delegate = TfLiteXNNPackDelegateCreate(&opts);
+  }
   std::unique_ptr<Interpreter> interpreter;
   InterpreterOptions options;
   options.SetPreserveAllTensors(preserve_all_tensors);
   options.SetDisableDelegateClustering(disable_delegate_clustering);
   InterpreterBuilder builder(*model, resolver, &options);
+  if (default_delegate_latest_features) {
+    builder.AddDelegate(xnnpack_delegate);
+  }
+  builder.SetNumThreads(num_threads);
   if (builder(&interpreter) != kTfLiteOk) {
     return nullptr;
   }
@@ -200,29 +214,36 @@ InterpreterWrapper* InterpreterWrapper::CreateInterpreterWrapper(
     const std::vector<std::string>& registerers_by_name,
     const std::vector<std::function<void(uintptr_t)>>& registerers_by_func,
     std::string* error_msg, bool preserve_all_tensors,
-    bool disable_delegate_clustering) {
+    bool disable_delegate_clustering, int num_threads,
+    bool default_delegate_latest_features) {
   if (!model) {
     *error_msg = error_reporter->message();
     return nullptr;
   }
 
   std::unique_ptr<tflite::MutableOpResolver> resolver;
-  switch (op_resolver_id) {
-    case kBuiltinOpResolver:
-      resolver = std::make_unique<tflite::ops::builtin::BuiltinOpResolver>();
-      break;
-    case kBuiltinRefOpResolver:
-      resolver = std::make_unique<tflite::ops::builtin::BuiltinRefOpResolver>();
-      break;
-    case kBuiltinOpResolverWithoutDefaultDelegates:
-      resolver = std::make_unique<
-          tflite::ops::builtin::BuiltinOpResolverWithoutDefaultDelegates>();
-      break;
-    default:
-      // This should not never happen because the eventual caller in
-      // interpreter.py should have passed a valid id here.
-      TFLITE_DCHECK(false);
-      return nullptr;
+  if (default_delegate_latest_features) {
+    resolver = std::make_unique<
+        tflite::ops::builtin::BuiltinOpResolverWithoutDefaultDelegates>();
+  } else {
+    switch (op_resolver_id) {
+      case kBuiltinOpResolver:
+        resolver = std::make_unique<tflite::ops::builtin::BuiltinOpResolver>();
+        break;
+      case kBuiltinRefOpResolver:
+        resolver =
+            std::make_unique<tflite::ops::builtin::BuiltinRefOpResolver>();
+        break;
+      case kBuiltinOpResolverWithoutDefaultDelegates:
+        resolver = std::make_unique<
+            tflite::ops::builtin::BuiltinOpResolverWithoutDefaultDelegates>();
+        break;
+      default:
+        // This should not never happen because the eventual caller in
+        // interpreter.py should have passed a valid id here.
+        TFLITE_DCHECK(false);
+        return nullptr;
+    }
   }
 
   for (const auto& registerer : registerers_by_name) {
@@ -232,9 +253,9 @@ InterpreterWrapper* InterpreterWrapper::CreateInterpreterWrapper(
   for (const auto& registerer : registerers_by_func) {
     registerer(reinterpret_cast<uintptr_t>(resolver.get()));
   }
-  auto interpreter =
-      CreateInterpreter(model.get(), *resolver, preserve_all_tensors,
-                        disable_delegate_clustering);
+  auto interpreter = CreateInterpreter(
+      model.get(), *resolver, preserve_all_tensors, disable_delegate_clustering,
+      num_threads, default_delegate_latest_features);
   if (!interpreter) {
     *error_msg = error_reporter->message();
     return nullptr;
@@ -806,14 +827,16 @@ InterpreterWrapper* InterpreterWrapper::CreateWrapperCPPFromFile(
     const std::vector<std::string>& registerers_by_name,
     const std::vector<std::function<void(uintptr_t)>>& registerers_by_func,
     std::string* error_msg, bool preserve_all_tensors,
-    bool disable_delegate_clustering) {
+    bool disable_delegate_clustering, int num_threads,
+    bool default_delegate_latest_features) {
   std::unique_ptr<PythonErrorReporter> error_reporter(new PythonErrorReporter);
   std::unique_ptr<InterpreterWrapper::Model> model =
       Model::BuildFromFile(model_path, error_reporter.get());
   return CreateInterpreterWrapper(
       std::move(model), op_resolver_id, std::move(error_reporter),
       registerers_by_name, registerers_by_func, error_msg, preserve_all_tensors,
-      disable_delegate_clustering);
+      disable_delegate_clustering, num_threads,
+      default_delegate_latest_features);
 }
 
 InterpreterWrapper* InterpreterWrapper::CreateWrapperCPPFromFile(
@@ -822,15 +845,17 @@ InterpreterWrapper* InterpreterWrapper::CreateWrapperCPPFromFile(
     bool preserve_all_tensors, bool disable_delegate_clustering) {
   return CreateWrapperCPPFromFile(
       model_path, op_resolver_id, registerers, {} /*registerers_by_func*/,
-      error_msg, preserve_all_tensors, disable_delegate_clustering);
+      error_msg, preserve_all_tensors, disable_delegate_clustering,
+      /*num_threads=*/1, /*default_delegate_latest_features=*/false);
 }
 
 InterpreterWrapper* InterpreterWrapper::CreateWrapperCPPFromBuffer(
     PyObject* data, int op_resolver_id,
     const std::vector<std::string>& registerers_by_name,
     const std::vector<std::function<void(uintptr_t)>>& registerers_by_func,
     std::string* error_msg, bool preserve_all_tensors,
-    bool disable_delegate_clustering) {
+    bool disable_delegate_clustering, int num_threads,
+    bool default_delegate_latest_features) {
   char* buf = nullptr;
   Py_ssize_t length;
   std::unique_ptr<PythonErrorReporter> error_reporter(new PythonErrorReporter);
@@ -843,16 +868,18 @@ InterpreterWrapper* InterpreterWrapper::CreateWrapperCPPFromBuffer(
   return CreateInterpreterWrapper(
       std::move(model), op_resolver_id, std::move(error_reporter),
       registerers_by_name, registerers_by_func, error_msg, preserve_all_tensors,
-      disable_delegate_clustering);
+      disable_delegate_clustering, num_threads,
+      default_delegate_latest_features);
 }
 
 InterpreterWrapper* InterpreterWrapper::CreateWrapperCPPFromBuffer(
     PyObject* data, int op_resolver_id,
     const std::vector<std::string>& registerers, std::string* error_msg,
     bool preserve_all_tensors, bool disable_delegate_clustering) {
-  return CreateWrapperCPPFromBuffer(data, op_resolver_id, registerers, {},
-                                    error_msg, preserve_all_tensors,
-                                    disable_delegate_clustering);
+  return CreateWrapperCPPFromBuffer(
+      data, op_resolver_id, registerers, {}, error_msg, preserve_all_tensors,
+      disable_delegate_clustering, /*num_threads=*/1,
+      /*default_delegate_latest_features=*/false);
 }
 
 PyObject* InterpreterWrapper::ResetVariableTensors() {

tensorflow/lite/python/interpreter_wrapper/interpreter_wrapper.h

@@ -57,7 +57,8 @@ class InterpreterWrapper {
       const std::vector<std::string>& registerers_by_name,
       const std::vector<std::function<void(uintptr_t)>>& registerers_by_func,
       std::string* error_msg, bool preserve_all_tensors,
-      bool disable_delegate_clustering);
+      bool disable_delegate_clustering, int num_threads,
+      bool default_delegate_latest_features);
 
   // SWIG caller takes ownership of pointer.
   static InterpreterWrapper* CreateWrapperCPPFromBuffer(
@@ -69,7 +70,8 @@ class InterpreterWrapper {
       const std::vector<std::string>& registerers_by_name,
       const std::vector<std::function<void(uintptr_t)>>& registerers_by_func,
       std::string* error_msg, bool preserve_all_tensors,
-      bool disable_delegate_clustering);
+      bool disable_delegate_clustering, int num_threads,
+      bool default_delegate_latest_features);
 
   ~InterpreterWrapper();
   PyObject* AllocateTensors(int subgraph_index);
@@ -126,7 +128,8 @@ class InterpreterWrapper {
       const std::vector<std::string>& registerers_by_name,
       const std::vector<std::function<void(uintptr_t)>>& registerers_by_func,
       std::string* error_msg, bool preserve_all_tensors,
-      bool disable_delegate_clustering);
+      bool disable_delegate_clustering, int num_threads,
+      bool default_delegate_latest_features);
 
   InterpreterWrapper(std::unique_ptr<Model> model,
                      std::unique_ptr<PythonErrorReporter> error_reporter,

tensorflow/lite/python/interpreter_wrapper/interpreter_wrapper_pybind11.cc

@@ -53,12 +53,14 @@ PYBIND11_MODULE(_pywrap_tensorflow_interpreter_wrapper, m) {
       [](const std::string& model_path, int op_resolver_id,
          const std::vector<std::string>& registerers_by_name,
          const std::vector<std::function<void(uintptr_t)>>& registerers_by_func,
-         bool preserve_all_tensors, bool disable_delegate_clustering) {
+         bool preserve_all_tensors, bool disable_delegate_clustering,
+         int num_threads, bool default_delegate_latest_features) {
         std::string error;
         auto* wrapper = ::InterpreterWrapper::CreateWrapperCPPFromFile(
             model_path.c_str(), op_resolver_id, registerers_by_name,
             registerers_by_func, &error, preserve_all_tensors,
-            disable_delegate_clustering);
+            disable_delegate_clustering, num_threads,
+            default_delegate_latest_features);
         if (!wrapper) {
           throw std::invalid_argument(error);
         }
@@ -82,12 +84,14 @@ PYBIND11_MODULE(_pywrap_tensorflow_interpreter_wrapper, m) {
       [](const py::bytes& data, int op_resolver_id,
          const std::vector<std::string>& registerers_by_name,
          const std::vector<std::function<void(uintptr_t)>>& registerers_by_func,
-         bool preserve_all_tensors, bool disable_delegate_clustering) {
+         bool preserve_all_tensors, bool disable_delegate_clustering,
+         int num_threads, bool default_delegate_latest_features) {
         std::string error;
         auto* wrapper = ::InterpreterWrapper::CreateWrapperCPPFromBuffer(
             data.ptr(), op_resolver_id, registerers_by_name,
             registerers_by_func, &error, preserve_all_tensors,
-            disable_delegate_clustering);
+            disable_delegate_clustering, num_threads,
+            default_delegate_latest_features);
         if (!wrapper) {
           throw std::invalid_argument(error);
         }

tensorflow/tools/api/golden/v1/tensorflow.lite.-interpreter.pbtxt

@@ -4,7 +4,7 @@ tf_class {
   is_instance: "<type \'object\'>"
   member_method {
     name: "__init__"
-    argspec: "args=[\'self\', \'model_path\', \'model_content\', \'experimental_delegates\', \'num_threads\', \'experimental_op_resolver_type\', \'experimental_preserve_all_tensors\', \'experimental_disable_delegate_clustering\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\', \'OpResolverType.AUTO\', \'False\', \'False\'], "
+    argspec: "args=[\'self\', \'model_path\', \'model_content\', \'experimental_delegates\', \'num_threads\', \'experimental_op_resolver_type\', \'experimental_preserve_all_tensors\', \'experimental_disable_delegate_clustering\', \'experimental_default_delegate_latest_features\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\', \'OpResolverType.AUTO\', \'False\', \'False\', \'False\'], "
   }
   member_method {
     name: "allocate_tensors"

tensorflow/tools/api/golden/v2/tensorflow.lite.-interpreter.pbtxt

@@ -4,7 +4,7 @@ tf_class {
   is_instance: "<type \'object\'>"
   member_method {
     name: "__init__"
-    argspec: "args=[\'self\', \'model_path\', \'model_content\', \'experimental_delegates\', \'num_threads\', \'experimental_op_resolver_type\', \'experimental_preserve_all_tensors\', \'experimental_disable_delegate_clustering\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\', \'OpResolverType.AUTO\', \'False\', \'False\'], "
+    argspec: "args=[\'self\', \'model_path\', \'model_content\', \'experimental_delegates\', \'num_threads\', \'experimental_op_resolver_type\', \'experimental_preserve_all_tensors\', \'experimental_disable_delegate_clustering\', \'experimental_default_delegate_latest_features\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\', \'OpResolverType.AUTO\', \'False\', \'False\', \'False\'], "
   }
   member_method {
     name: "allocate_tensors"