Add API to interpreter experimental to retrieve async signature runner.

PiperOrigin-RevId: 507869775

Author: sirakiin

tensorflow/lite/CMakeLists.txt

@@ -245,6 +245,7 @@ if(NOT "${CMAKE_SYSTEM_NAME}" STREQUAL "iOS")
 endif()
 populate_tflite_source_vars("core" TFLITE_CORE_SRCS)
 populate_tflite_source_vars("core/api" TFLITE_CORE_API_SRCS)
+populate_tflite_source_vars("core/async" TFLITE_CORE_ASYNC_SRCS)
 populate_tflite_source_vars("core/c" TFLITE_CORE_C_SRCS)
 populate_tflite_source_vars("core/experimental/acceleration/configuration" TFLITE_CORE_EXPERIMENTAL_SRCS)
 populate_tflite_source_vars("core/kernels" TFLITE_CORE_KERNELS_SRCS)
@@ -512,6 +513,7 @@ set(_ALL_TFLITE_SRCS
   ${TFLITE_CORE_EXPERIMENTAL_SRCS}
   ${TFLITE_CORE_KERNELS_SRCS}
   ${TFLITE_CORE_SRCS}
+  ${TFLITE_CORE_ASYNC_SRCS}
   ${TFLITE_CORE_TOOLS_SRCS}
   ${TFLITE_C_SRCS}
   ${TFLITE_DELEGATES_FLEX_SRCS}

tensorflow/lite/core/BUILD

@@ -45,6 +45,7 @@ cc_library(
         "//tensorflow/lite:type_to_tflitetype",
         "//tensorflow/lite/core/api",
         "//tensorflow/lite/core/api:error_reporter",
+        "//tensorflow/lite/core/async:async_signature_runner",
         "//tensorflow/lite/core/c:common",
         "//tensorflow/lite/experimental/remat:metadata_util",
         "//tensorflow/lite/experimental/resource",
@@ -109,6 +110,7 @@ cc_library(
         "//tensorflow/lite/c:common_internal",
         "//tensorflow/lite/core/api",
         "//tensorflow/lite/core/api:verifier",
+        "//tensorflow/lite/core/async:async_signature_runner",
         "//tensorflow/lite/core/c:common",
         "//tensorflow/lite/experimental/remat:metadata_util",
         "//tensorflow/lite/experimental/resource",
@@ -180,6 +182,7 @@ cc_library(
         "//tensorflow/lite/c:common_internal",
         "//tensorflow/lite/core/api",
         "//tensorflow/lite/core/api:verifier",
+        "//tensorflow/lite/core/async:async_signature_runner",
         "//tensorflow/lite/core/c:common",
         "//tensorflow/lite/experimental/remat:metadata_util",
         "//tensorflow/lite/experimental/resource",
@@ -237,6 +240,7 @@ cc_library(
         "//tensorflow/lite/c:common_internal",
         "//tensorflow/lite/core/api",
         "//tensorflow/lite/core/api:verifier",
+        "//tensorflow/lite/core/async:async_signature_runner",
         "//tensorflow/lite/core/c:c_api_types",
         "//tensorflow/lite/core/c:common",
         "//tensorflow/lite/delegates:telemetry",
@@ -313,6 +317,7 @@ cc_library(
         "//tensorflow/lite/c:common_internal",
         "//tensorflow/lite/core/api",
         "//tensorflow/lite/core/api:verifier",
+        "//tensorflow/lite/core/async:async_signature_runner",
         "//tensorflow/lite/core/c:c_api_types",
         "//tensorflow/lite/core/c:common",
         "//tensorflow/lite/experimental/remat:metadata_util",

tensorflow/lite/core/async/BUILD

@@ -54,12 +54,13 @@ cc_library(
     name = "async_subgraph",
     srcs = ["async_subgraph.cc"],
     hdrs = ["async_subgraph.h"],
+    compatible_with = get_compatible_with_portable(),
     deps = [
         ":async_kernel_internal",
         ":common",
         ":task_internal",
-        "//tensorflow/lite:framework",
         "//tensorflow/lite:minimal_logging",
+        "//tensorflow/lite/core:subgraph",
         "//tensorflow/lite/core/async/interop/c:types",
         "//tensorflow/lite/core/c:c_api_types",
         "//tensorflow/lite/core/c:common",
@@ -103,14 +104,13 @@ cc_library(
     name = "async_signature_runner",
     srcs = ["async_signature_runner.cc"],
     hdrs = ["async_signature_runner.h"],
+    compatible_with = get_compatible_with_portable(),
     deps = [
         ":async_kernel_internal",
         ":async_subgraph",
         ":common",
         ":task_internal",
-        "//tensorflow/lite:framework",
-        "//tensorflow/lite/c:c_api_without_op_resolver",
-        "//tensorflow/lite/core/c:c_api",
+        "//tensorflow/lite/core:subgraph",
         "//tensorflow/lite/core/c:c_api_types",
         "//tensorflow/lite/core/c:common",
         "//tensorflow/lite/internal:signature_def",
@@ -126,7 +126,7 @@ cc_test(
         ":backend_async_kernel_interface",
         ":common",
         "//tensorflow/lite:framework",
-        "//tensorflow/lite/c:c_api_for_testing",
+        "//tensorflow/lite:interpreter_test_util",
         "//tensorflow/lite/core:headers",
         "//tensorflow/lite/core/async/c:task",
         "//tensorflow/lite/core/async/c:types",

tensorflow/lite/core/async/async_signature_runner.cc

@@ -26,23 +26,8 @@ limitations under the License.
 #include "tensorflow/lite/core/async/async_subgraph.h"
 #include "tensorflow/lite/core/async/common.h"
 #include "tensorflow/lite/core/async/task_internal.h"
-#include "tensorflow/lite/signature_runner.h"
 
 namespace tflite {
-
-// This is a temporary helper class that will be removed after this API is
-// moved out of experimental.
-class SignatureRunnerHelper {
- public:
-  static Subgraph* GetSubgraph(SignatureRunner* runner) {
-    return runner->subgraph_;
-  }
-  static const internal::SignatureDef* GetSignatureDef(
-      SignatureRunner* runner) {
-    return runner->signature_def_;
-  }
-};
-
 namespace async {
 
 namespace {
@@ -79,11 +64,6 @@ int AsyncSignatureRunner::GetTensorIndex(TfLiteIoType io_type,
   return tensor_index;
 }
 
-AsyncSignatureRunner::AsyncSignatureRunner(SignatureRunner* signature_runner)
-    : AsyncSignatureRunner(
-          SignatureRunnerHelper::GetSignatureDef(signature_runner),
-          SignatureRunnerHelper::GetSubgraph(signature_runner)) {}
-
 AsyncSignatureRunner::AsyncSignatureRunner(
     const internal::SignatureDef* signature_def, Subgraph* subgraph)
     : signature_def_(signature_def), subgraph_(subgraph) {

tensorflow/lite/core/async/async_signature_runner.h

@@ -18,14 +18,12 @@ limitations under the License.
 #include <memory>
 #include <vector>
 
-#include "tensorflow/lite/core/c/c_api.h"
 #include "tensorflow/lite/core/c/common.h"
 #include "tensorflow/lite/core/subgraph.h"
 #include "tensorflow/lite/core/async/async_kernel_internal.h"
 #include "tensorflow/lite/core/async/async_subgraph.h"
 #include "tensorflow/lite/core/async/common.h"
 #include "tensorflow/lite/internal/signature_def.h"
-#include "tensorflow/lite/signature_runner.h"
 
 namespace tflite {
 namespace async {
@@ -39,20 +37,8 @@ class AsyncSignatureRunnerTest;
 // SignatureDef.
 class AsyncSignatureRunner {
  public:
-  // TODO(b/191883048): Move ctor to private and use `Create` function as
-  // factory method.
-  // Currently we don't have way to expose signature def from interpreter
-  // without changes to interpreter.
-  //
-  // static AsyncSignatureRunner* Create(const TfLiteInterpreter* interpreter,
-  //                                     const char* signature_key);
-  // WARNING: This is a temporary constructor before we stablize the API.
-  // This if for avoiding making intrusive changes to non experimental code.
-  // For now, users can construct AsyncSignatureRunner as follows:
-  //  std::unique_ptr<tflite::Interpreter> interpreter;
-  //  InterpreterBuilder(model, resolver)(&interpreter);
-  //  AsyncSignatureRunner runner(interpreter->GetSignatureRunner("func"));
-  explicit AsyncSignatureRunner(SignatureRunner* signature_runner);
+  // Builds the AsyncSignatureRunner given the provided signature_def and
+  // subgraph.
   AsyncSignatureRunner(const internal::SignatureDef* signature_def,
                        Subgraph* subgraph);
 

tensorflow/lite/core/async/async_signature_runner_test.cc

@@ -18,22 +18,23 @@ limitations under the License.
 
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
-#include "tensorflow/lite/core/c/c_api_types.h"
-#include "tensorflow/lite/core/c/common.h"
-#include "tensorflow/lite/core/interpreter.h"
 #include "tensorflow/lite/core/async/async_kernel_internal.h"
 #include "tensorflow/lite/core/async/backend_async_kernel_interface.h"
 #include "tensorflow/lite/core/async/c/task.h"
 #include "tensorflow/lite/core/async/c/types.h"
 #include "tensorflow/lite/core/async/common.h"
 #include "tensorflow/lite/core/async/testing/mock_async_kernel.h"
 #include "tensorflow/lite/core/async/testing/test_backend.h"
+#include "tensorflow/lite/core/c/c_api_types.h"
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/core/interpreter.h"
+#include "tensorflow/lite/interpreter_test_util.h"
 #include "tensorflow/lite/kernels/builtin_op_kernels.h"
 
 namespace tflite {
 namespace async {
 
-class AsyncSignatureRunnerTest : public ::testing::Test {
+class AsyncSignatureRunnerTest : public InterpreterTest {
  protected:
   void SetUp() override {
     kernel_ =
@@ -53,41 +54,39 @@ class AsyncSignatureRunnerTest : public ::testing::Test {
     void* builtin_data_1 = malloc(sizeof(int));
     interpreter_->AddNodeWithParameters({0, 0}, {1}, nullptr, 0, builtin_data_1,
                                         reg);
-    signature_def_.signature_key = "serving_default";
-    signature_def_.inputs["input"] = 0;
-    signature_def_.outputs["output"] = 1;
-    signature_def_.subgraph_index = 0;
-  }
-
-  void BuildAsyncSignatureRunner() {
+    const char kSignatureKey[] = "serving_default";
+    BuildSignature(kSignatureKey, {{"input", 0}}, {{"output", 1}});
     interpreter_->ModifyGraphWithDelegate(backend_->get_delegate());
-    signature_runner_ = std::make_unique<AsyncSignatureRunner>(
-        &signature_def_, interpreter_->subgraph(0));
   }
 
   int GetTensorIndex(TfLiteIoType io_type, const char* name) {
     return signature_runner_->GetTensorIndex(io_type, name);
   }
 
-  void TearDown() override { signature_runner_.reset(); }
-
  protected:
   std::unique_ptr<::testing::StrictMock<testing::MockAsyncKernel>> kernel_;
   std::unique_ptr<testing::TestBackend> backend_;
-  std::unique_ptr<Interpreter> interpreter_;
   internal::SignatureDef signature_def_;
-  std::unique_ptr<AsyncSignatureRunner> signature_runner_;
+  AsyncSignatureRunner* signature_runner_ = nullptr;
 };
 
+TEST_F(AsyncSignatureRunnerTest, GetAsyncSignatureRunner) {
+  EXPECT_EQ(nullptr, signature_runner_);
+  signature_runner_ = interpreter_->GetAsyncSignatureRunner("serving_default");
+  EXPECT_NE(nullptr, signature_runner_);
+
+  EXPECT_EQ(nullptr, interpreter_->GetAsyncSignatureRunner("foo"));
+}
+
 TEST_F(AsyncSignatureRunnerTest, InputNameTest) {
-  BuildAsyncSignatureRunner();
+  signature_runner_ = interpreter_->GetAsyncSignatureRunner("serving_default");
   EXPECT_EQ(0, GetTensorIndex(TfLiteIoType::kTfLiteIoInput, "input"));
   EXPECT_EQ(-1, GetTensorIndex(TfLiteIoType::kTfLiteIoInput, "output"));
   EXPECT_EQ(-1, GetTensorIndex(TfLiteIoType::kTfLiteIoInput, "foo"));
 }
 
 TEST_F(AsyncSignatureRunnerTest, OutputNameTest) {
-  BuildAsyncSignatureRunner();
+  signature_runner_ = interpreter_->GetAsyncSignatureRunner("serving_default");
   EXPECT_EQ(1, GetTensorIndex(TfLiteIoType::kTfLiteIoOutput, "output"));
   EXPECT_EQ(-1, GetTensorIndex(TfLiteIoType::kTfLiteIoOutput, "input"));
   EXPECT_EQ(-1, GetTensorIndex(TfLiteIoType::kTfLiteIoOutput, "foo"));
@@ -96,7 +95,7 @@ TEST_F(AsyncSignatureRunnerTest, OutputNameTest) {
 TEST_F(AsyncSignatureRunnerTest, CreateTaskTest) {
   EXPECT_CALL(*kernel_, Finish(::testing::_, ::testing::_));
 
-  BuildAsyncSignatureRunner();
+  signature_runner_ = interpreter_->GetAsyncSignatureRunner("serving_default");
   auto* task = signature_runner_->CreateTask();
   EXPECT_NE(nullptr, task);
 

tensorflow/lite/core/async/async_subgraph.cc

@@ -47,7 +47,7 @@ TfLiteAsyncKernel* AsyncSubgraph::async_kernel() const {
 AsyncSubgraph::AsyncSubgraph(Subgraph* subgraph) : subgraph_(subgraph) {
   // Currently we only support one delegate and fully delegated subgph.
   if (!IsFullyDelegated()) {
-    subgraph->ReportError("Model is no fully delegated by 1 backend.");
+    subgraph->ReportError("Model is not fully delegated by 1 backend.");
     return;
   }
   // TODO(b/191883048): Add/Check delegate flag to indicate kernel support.

tensorflow/lite/core/interpreter.h

@@ -41,6 +41,7 @@ limitations under the License.
 
 #include "tensorflow/lite/allocation.h"
 #include "tensorflow/lite/core/api/error_reporter.h"
+#include "tensorflow/lite/core/async/async_signature_runner.h"
 #include "tensorflow/lite/core/api/profiler.h"
 #include "tensorflow/lite/core/c/common.h"  // IWYU pragma: export
 #include "tensorflow/lite/core/subgraph.h"
@@ -344,6 +345,14 @@ class Interpreter {
   /// and the SignatureRunner class is *not* thread-safe.
   SignatureRunner* GetSignatureRunner(const char* signature_key);
 
+  /// \warning Experimental interface, subject to change. \n
+  /// \brief Returns a pointer to the AsyncSignatureRunner instance to run the
+  /// part of the graph identified by a SignatureDef. The nullptr is returned if
+  /// the given signature key is not valid.
+  /// The async delegate should be applied before calling this function.
+  async::AsyncSignatureRunner* GetAsyncSignatureRunner(
+      const char* signature_key);
+
   /// \warning Experimental interface, subject to change. \n
   /// \brief Return the subgraph index that corresponds to a SignatureDef,
   /// defined by 'signature_key'.
@@ -946,6 +955,12 @@ class Interpreter {
   // its SignatureDef.
   std::map<std::string, SignatureRunner> signature_runner_map_;
 
+  // Map of signature key to its corresponding AsyncSignatureRunner object.
+  // An AsyncSignatureRunner is basically a wrapper of the AsyncSubgraph
+  // corresponding to its SignatureDef.
+  std::map<std::string, async::AsyncSignatureRunner>
+      async_signature_runner_map_;
+
   // Model metadata stored as mapping of name (key) to buffer (value).
   // Data is mapped from the Metadata in TFLite flatbuffer model.
   std::map<std::string, std::string> metadata_;

tensorflow/lite/core/interpreter_experimental.cc

@@ -27,6 +27,7 @@ limitations under the License.
 #include "tensorflow/lite/c/common_internal.h"
 #include "tensorflow/lite/core/api/error_reporter.h"
 #include "tensorflow/lite/core/api/profiler.h"
+#include "tensorflow/lite/core/async/async_signature_runner.h"
 #include "tensorflow/lite/core/c/c_api_types.h"
 #include "tensorflow/lite/core/c/common.h"
 #include "tensorflow/lite/core/interpreter.h"
@@ -165,4 +166,23 @@ SignatureRunner* Interpreter::GetSignatureRunner(const char* signature_key) {
   return nullptr;
 }
 
+async::AsyncSignatureRunner* Interpreter::GetAsyncSignatureRunner(
+    const char* signature_key) {
+  auto iter = async_signature_runner_map_.find(signature_key);
+  if (iter != async_signature_runner_map_.end()) {
+    return &(iter->second);
+  }
+
+  for (const auto& signature : signature_defs_) {
+    if (signature.signature_key == signature_key) {
+      auto status = async_signature_runner_map_.insert(
+          {signature_key, async::AsyncSignatureRunner(
+                              &signature, subgraph(signature.subgraph_index))});
+      return &(status.first->second);
+    }
+  }
+
+  return nullptr;
+}
+
 }  // namespace tflite