tflite: npu: Async Kernel API implementation

BUG=b:353367133
TEST=async_delegate_test
--stable_delegate_settings_file=/etc/tflite/settings.json

Change-Id: I7ee15eceb3d5f98ff16271340944bd13176be0ab
Signed-off-by: Ritul Jasuja <ritul.jasuja@intel.com>
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/platform/tflite/+/5872382
Reviewed-by: Tommy Chiang <ototot@google.com>
Tested-by: Tommy Chiang <ototot@google.com>
Reviewed-by: Shik Chen <shik@chromium.org>
Commit-Queue: Tommy Chiang <ototot@google.com>
diff --git a/delegate/intel_openvino/BUILD.bazel b/delegate/intel_openvino/BUILD.bazel
index 2cafd43..3da761a 100644
--- a/delegate/intel_openvino/BUILD.bazel
+++ b/delegate/intel_openvino/BUILD.bazel
@@ -20,6 +20,25 @@
 )
 
 cc_library_with_tflite(
+    name = "openvino_async_kernel",
+    srcs = ["openvino_async_kernel.cc"],
+    hdrs = ["openvino_async_kernel.h"],
+    copts = tflite_copts(),
+    deps = [
+        ":openvino_delegate_kernel",
+        "//android:hardware_buffer",
+        "@com_google_absl//absl/synchronization",
+        "@org_tensorflow//tensorflow/lite/async:backend_async_kernel_interface",
+        "@org_tensorflow//tensorflow/lite/async/c:async_kernel",
+        "@org_tensorflow//tensorflow/lite/async/c:types",
+        "@org_tensorflow//tensorflow/lite/async/interop/c:attribute_map",
+        "@org_tensorflow//tensorflow/lite/core/async/c:task",
+        "@org_tensorflow//tensorflow/lite/core/async/interop/c:constants",
+        "@org_tensorflow//tensorflow/lite/delegates/utils:async_type_helpers",
+    ],
+)
+
+cc_library_with_tflite(
     name = "openvino_graph_builder",
     srcs = ["openvino_graph_builder.cc"],
     hdrs = ["openvino_graph_builder.h"],
@@ -49,6 +68,7 @@
     deps = [
         ":log",
         ":openvino_graph_builder",
+        "//common:simple_async_delegate",
         "@intel_openvino//:openvino",
         "@org_tensorflow//tensorflow/lite:kernel_api",
         "@org_tensorflow//tensorflow/lite/c:c_api",
@@ -64,6 +84,7 @@
     name = "openvino_delegate_kernel",
     srcs = ["openvino_delegate_kernel.cc"],
     hdrs = [
+        "openvino_async_kernel.h",
         "openvino_delegate.h",
         "openvino_delegate_kernel.h",
     ],
@@ -71,14 +92,24 @@
     deps = [
         ":log",
         ":openvino_delegate_core",
+        "//android:hardware_buffer",
+        "@com_google_absl//absl/synchronization",
         "@intel_openvino//:openvino",
         "@org_tensorflow//tensorflow/lite:kernel_api",
+        "@org_tensorflow//tensorflow/lite/async:backend_async_kernel_interface",
+        "@org_tensorflow//tensorflow/lite/async/c:async_kernel",
+        "@org_tensorflow//tensorflow/lite/async/interop/c:attribute_map",
+        "@org_tensorflow//tensorflow/lite/async/interop/c:types",
         "@org_tensorflow//tensorflow/lite/c:c_api",
         "@org_tensorflow//tensorflow/lite/c:c_api_experimental",
         "@org_tensorflow//tensorflow/lite/c:c_api_types",
         "@org_tensorflow//tensorflow/lite/c:common",
+        "@org_tensorflow//tensorflow/lite/core/async/c:task",
+        "@org_tensorflow//tensorflow/lite/core/async/interop/c:constants",
+        "@org_tensorflow//tensorflow/lite/delegates/utils:async_type_helpers",
         "@org_tensorflow//tensorflow/lite/delegates/utils:simple_delegate",
         "@org_tensorflow//tensorflow/lite/delegates/utils:simple_opaque_delegate",
+        "@org_tensorflow//tensorflow/lite/delegates/utils:sync_fence",
         "@org_tensorflow//tensorflow/lite/kernels:kernel_util",
     ],
 )
@@ -90,13 +121,18 @@
     copts = tflite_copts(),
     deps = [
         ":log",
+        ":openvino_async_kernel",
         ":openvino_delegate_kernel",
+        "//android:hardware_buffer",
         "@intel_openvino//:openvino",
         "@org_tensorflow//tensorflow/lite:kernel_api",
         "@org_tensorflow//tensorflow/lite/c:c_api",
         "@org_tensorflow//tensorflow/lite/c:c_api_experimental",
         "@org_tensorflow//tensorflow/lite/c:c_api_types",
         "@org_tensorflow//tensorflow/lite/c:common",
+        "@org_tensorflow//tensorflow/lite/core/async/c:async_kernel",
+        "@org_tensorflow//tensorflow/lite/core/async/interop/c:constants",
+        "@org_tensorflow//tensorflow/lite/delegates/utils:async_type_helpers",
         "@org_tensorflow//tensorflow/lite/delegates/utils:simple_opaque_delegate",
         "@org_tensorflow//tensorflow/lite/kernels:builtin_ops",
         "@org_tensorflow//tensorflow/lite/kernels:padding",
@@ -114,6 +150,7 @@
     deps = [
         ":log",
         ":openvino_delegate",
+        "//common:simple_async_delegate",
         "@org_tensorflow//tensorflow/lite/acceleration/configuration/c:delegate_plugin",
         "@org_tensorflow//tensorflow/lite/acceleration/configuration/c:stable_delegate",
         "@org_tensorflow//tensorflow/lite/c:c_api",
@@ -129,12 +166,14 @@
 
 tflite_cc_shared_object(
     name = "tensorflowlite_intel_openvino_delegate",
+    srcs = ["//android:libnativewindow.so"],
     linkopts = [
         "-Wl,--no-undefined",
         "-Wl,--version-script,$(location @org_tensorflow//tensorflow/lite/delegates/utils/experimental/stable_delegate:version_script.lds)",
     ],
     visibility = ["//:__pkg__"],
     deps = [
+        ":openvino_async_kernel",
         ":openvino_delegate_external",
         "@org_tensorflow//tensorflow/lite/delegates/utils/experimental/stable_delegate:version_script.lds",
     ],
@@ -143,7 +182,10 @@
 cc_binary(
     name = "openvino_delegate_core_test",
     testonly = True,
-    srcs = ["openvino_delegate_core_test.cc"],
+    srcs = [
+        "openvino_delegate_core_test.cc",
+        "//android:libnativewindow.so",
+    ],
     copts = tflite_copts(),
     data = [
         ":tensorflowlite_intel_openvino_delegate",
@@ -185,7 +227,10 @@
 cc_binary(
     name = "openvino_delegate_test",
     testonly = True,
-    srcs = ["openvino_delegate_test.cc"],
+    srcs = [
+        "openvino_delegate_test.cc",
+        "//android:libnativewindow.so",
+    ],
     copts = tflite_copts(),
     data = [
         ":tensorflowlite_intel_openvino_delegate",
diff --git a/delegate/intel_openvino/openvino_async_kernel.cc b/delegate/intel_openvino/openvino_async_kernel.cc
new file mode 100644
index 0000000..57e078e
--- /dev/null
+++ b/delegate/intel_openvino/openvino_async_kernel.cc
@@ -0,0 +1,136 @@
+/*
+ * Copyright (C) 2024 Intel Corporation
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#include "delegate/intel_openvino/openvino_async_kernel.h"
+
+#include <sys/mman.h>
+
+#include "delegate/intel_openvino/openvino_delegate_kernel.h"
+#include "tensorflow/lite/builtin_ops.h"
+#include "tensorflow/lite/c/c_api.h"
+#include "tensorflow/lite/c/c_api_opaque.h"
+#include "tensorflow/lite/core/async/c/task.h"
+#include "tensorflow/lite/kernels/internal/compatibility.h"
+#include "tensorflow/lite/util.h"
+
+namespace tflite::openvinodelegate {
+
+using tflite::delegates::utils::ReadBufferAttrs;
+using tflite::delegates::utils::ReadSyncAttrs;
+
+OVDelegateAsyncKernel::OVDelegateAsyncKernel(OpenVINOAsyncDelegateKernel* core)
+    : core_(*core) {}
+
+TfLiteStatus OVDelegateAsyncKernel::RegisterBuffer(
+    TfLiteOpaqueContext* opaque_context, TfLiteIoType io_type,
+    const TfLiteBackendBuffer* buffer, const TfLiteAttributeMap* attrs,
+    TfLiteBufferHandle handle) {
+  absl::MutexLock lock(&mutex_);
+  if (TfLiteAttributeMapIsBufferAttributeMap(attrs)) {
+    auto buffer_attrs = ReadBufferAttrs(attrs);
+    size_t buffer_size = buffer_attrs.size.value();
+    if (buffer_attrs.buffer_type != BufferType::kAHardwareBufferBlob)
+      return kTfLiteError;
+
+    auto* ptr =
+        static_cast<AHardwareBuffer*>(TfLiteBackendBufferGetPtr(buffer));
+
+    if (core_.RegisterBuffer(handle, ptr, buffer_size) != kTfLiteOk)
+      return kTfLiteError;
+    return kTfLiteOk;
+  }
+  return kTfLiteError;
+}
+
+TfLiteStatus OVDelegateAsyncKernel::UnregisterBuffer(
+    TfLiteOpaqueContext* opaque_context, TfLiteBufferHandle handle) {
+  absl::MutexLock lock(&mutex_);
+  if (core_.UnregisterBuffer(handle) != kTfLiteOk) return kTfLiteError;
+  return kTfLiteOk;
+}
+
+bool OVDelegateAsyncKernel::ReconcileRestrictions(
+    const TfLiteOpaqueContext* opaque_context,
+    const TfLiteOpaqueNode* opaque_node, int tensor_index,
+    const TfLiteAttributeMap* user_provided_attributes,
+    TfLiteAttributeMap* merged, TfLiteAttributeMap* conflict) const {
+  absl::MutexLock lock(&mutex_);
+  TfLiteOpaqueTensor* tensor =
+      TfLiteOpaqueContextGetOpaqueTensor(opaque_context, tensor_index);
+  if (TfLiteAttributeMapIsBufferAttributeMap(user_provided_attributes)) {
+    BufferAttributes conflict_attrs{};
+    auto buffer_attrs = ReadBufferAttrs(user_provided_attributes);
+
+    auto buffer_type =
+        buffer_attrs.buffer_type.value_or(BufferType::kAHardwareBufferBlob);
+    if (buffer_type != BufferType::kAHardwareBufferBlob) {
+      conflict_attrs.buffer_type = BufferType::kAHardwareBufferBlob;
+      delegates::utils::WriteBufferAttrs(conflict_attrs, conflict);
+      return false;
+    }
+    size_t tensor_size = TfLiteOpaqueTensorByteSize(tensor);
+    buffer_attrs.size = std::max(buffer_attrs.size.value_or(0), tensor_size);
+    delegates::utils::WriteBufferAttrs(buffer_attrs, merged);
+    return true;
+  } else if (TfLiteAttributeMapIsSyncAttributeMap(user_provided_attributes)) {
+    SyncAttributes conflict_attrs;
+    SyncAttributes merged_attrs;
+    auto sync_attrs = delegates::utils::ReadSyncAttrs(user_provided_attributes);
+    auto sync_type =
+        sync_attrs.sync_type.value_or(delegates::utils::SyncType::kNoSyncObj);
+    if (sync_type == delegates::utils::SyncType::kUnknown) {
+      conflict_attrs.sync_type = delegates::utils::SyncType::kNoSyncObj;
+      delegates::utils::WriteSyncAttrs(conflict_attrs, conflict);
+      return false;
+    }
+    merged_attrs.sync_type = sync_type;
+    delegates::utils::WriteSyncAttrs(merged_attrs, merged);
+  } else {
+    return false;
+  }
+  return true;
+}
+
+TfLiteStatus OVDelegateAsyncKernel::SetAttributes(
+    TfLiteOpaqueContext* opaque_context, TfLiteOpaqueNode* opaque_node,
+    int tensor_index, const TfLiteAttributeMap* attrs) {
+  // TODO: This function is a no-op for now
+  return kTfLiteOk;
+}
+
+TfLiteStatus OVDelegateAsyncKernel::Prepare(TfLiteOpaqueContext* opaque_context,
+                                            TfLiteOpaqueNode* opaque_node) {
+  // TODO: This function is a no-op for now
+  return kTfLiteOk;
+}
+
+TfLiteStatus OVDelegateAsyncKernel::Eval(TfLiteOpaqueContext* opaque_context,
+                                         TfLiteOpaqueNode* opaque_node,
+                                         TfLiteExecutionTask* task) {
+  absl::MutexLock lock(&mutex_);
+  TfLiteStatus status = core_.EvalAsyncImpl(opaque_context, opaque_node, task);
+  task_status_map_.insert_or_assign(task, status);
+  return status;
+}
+
+TfLiteStatus OVDelegateAsyncKernel::Wait(TfLiteOpaqueContext* opaque_context,
+                                         TfLiteExecutionTask* task) {
+  absl::MutexLock lock(&mutex_);
+  auto it = task_status_map_.find(task);
+  if (it == task_status_map_.end()) {
+    return kTfLiteError;
+  }
+
+  return it->second;
+}
+
+TfLiteStatus OVDelegateAsyncKernel::Finish(TfLiteOpaqueContext* opaque_context,
+                                           TfLiteExecutionTask* task) {
+  absl::MutexLock lock(&mutex_);
+  size_t ret = task_status_map_.erase(task);
+  return ret == 1 ? kTfLiteOk : kTfLiteError;
+}
+
+}  // namespace tflite::openvinodelegate
diff --git a/delegate/intel_openvino/openvino_async_kernel.h b/delegate/intel_openvino/openvino_async_kernel.h
new file mode 100644
index 0000000..0c110f6
--- /dev/null
+++ b/delegate/intel_openvino/openvino_async_kernel.h
@@ -0,0 +1,116 @@
+/*
+ * Copyright (C) 2024 Intel Corporation
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#ifndef DELEGATE_OPENVINO_ASYNC_KERNEL_H_
+#define DELEGATE_OPENVINO_ASYNC_KERNEL_H_
+
+#include <algorithm>
+#include <map>
+#include <vector>
+
+#include "absl/synchronization/mutex.h"
+#include "android/hardware_buffer.h"
+#include "delegate/intel_openvino/log.h"
+#include "delegate/intel_openvino/openvino_delegate_kernel.h"
+#include "tensorflow/lite/async/backend_async_kernel_interface.h"
+#include "tensorflow/lite/core/async/interop/c/constants.h"
+#include "tensorflow/lite/delegates/utils/async_type_helpers.h"
+
+namespace tflite::openvinodelegate {
+
+using tflite::delegates::utils::BufferAttributes;
+using tflite::delegates::utils::BufferType;
+using tflite::delegates::utils::kBufferTypeAHardwareBufferBlob;
+using tflite::delegates::utils::SyncAttributes;
+using BackendAsyncKernelInterface =
+    ::tflite::delegates::BackendAsyncKernelInterface;
+
+class OpenVINOAsyncDelegateKernel;
+class OVDelegateAsyncKernel : public BackendAsyncKernelInterface {
+ public:
+  explicit OVDelegateAsyncKernel(OpenVINOAsyncDelegateKernel* core);
+  ~OVDelegateAsyncKernel() override = default;
+
+  // Buffer operations
+  TfLiteStatus RegisterBuffer(TfLiteOpaqueContext* opaque_context,
+                              TfLiteIoType io_type,
+                              const TfLiteBackendBuffer* buffer,
+                              const TfLiteAttributeMap* attrs,
+                              TfLiteBufferHandle handle) override;
+  TfLiteStatus RegisterBufferSlice(TfLiteOpaqueContext* context,
+                                   TfLiteBufferHandle buffer_pool,
+                                   const TfLiteAttributeMap* attrs,
+                                   TfLiteBufferHandle handle) override {
+    // TODO: implement the interface when required later
+    TFLITE_LOG(INFO)
+        << "OVDelegateAsyncKernel::RegisterBufferSlice is not implemented";
+    return kTfLiteError;
+  }
+  TfLiteStatus UnregisterBuffer(TfLiteOpaqueContext* opaque_context,
+                                TfLiteBufferHandle handle) override;
+
+  // Reconciliations
+  const std::vector<const char*>& SupportedBufferTypes(
+      TfLiteIoType io_type) const override {
+    return supported_buffer_types_;
+  }
+  const std::vector<const char*>& SupportedSynchronizations(
+      TfLiteIoType io_type) const override {
+    return supported_synchronizations_;
+  }
+  bool ReconcileRestrictions(const TfLiteOpaqueContext* opaque_context,
+                             const TfLiteOpaqueNode* opaque_node,
+                             int tensor_index,
+                             const TfLiteAttributeMap* user_provided_attributes,
+                             TfLiteAttributeMap* merged,
+                             TfLiteAttributeMap* conflict) const override;
+  TfLiteStatus SetAttributes(TfLiteOpaqueContext* context,
+                             TfLiteOpaqueNode* node, int tensor_index,
+                             const TfLiteAttributeMap* attrs) override;
+  TfLiteStatus SetBufferAttributes(const TfLiteBackendBuffer* buffer,
+                                   const TfLiteAttributeMap* attrs) override {
+    // TODO: implement the interface when required later
+    TFLITE_LOG(INFO)
+        << "OVDelegateAsyncKernel::SetBufferAttributes is not supported";
+    return kTfLiteError;
+  }
+
+  TfLiteStatus GetBufferAttributes(const TfLiteBackendBuffer* buffer,
+                                   TfLiteAttributeMap* attrs) override {
+    // TODO: implement the interface when required later
+    TFLITE_LOG(INFO)
+        << "OVDelegateAsyncKernel::GetBufferAttributes is not supported";
+    return kTfLiteError;
+  }
+  TfLiteStatus Prepare(TfLiteOpaqueContext* context,
+                       TfLiteOpaqueNode* node) override;
+
+  // Execution methods
+  TfLiteStatus Eval(TfLiteOpaqueContext* opaque_context,
+                    TfLiteOpaqueNode* opaque_node,
+                    TfLiteExecutionTask* task) override;
+  TfLiteStatus Wait(TfLiteOpaqueContext* opaque_context,
+                    TfLiteExecutionTask* task) override;
+  TfLiteStatus Finish(TfLiteOpaqueContext* opaque_context,
+                      TfLiteExecutionTask* task) override;
+
+ private:
+  mutable absl::Mutex mutex_;
+  OpenVINOAsyncDelegateKernel& core_ ABSL_GUARDED_BY(mutex_);
+  // For SupportedBufferTypes and SupportedSynchronizations
+  const std::vector<const char*> supported_buffer_types_ = {
+      kBufferTypeAHardwareBufferBlob};
+  const std::vector<const char*> supported_synchronizations_ = {
+      kTfLiteSyncTypeNoSyncObj,
+      ::tflite::delegates::utils::kSyncTypeSyncFenceFd};
+  std::map<TfLiteExecutionTask*, TfLiteStatus> task_status_map_
+      ABSL_GUARDED_BY(mutex_);
+  std::map<TfLiteBufferHandle, AHardwareBuffer*> ahwb_buffer_map_
+      ABSL_GUARDED_BY(mutex_);
+};
+
+}  // namespace tflite::openvinodelegate
+
+#endif  // DELEGATE_OPENVINO_ASYNC_KERNEL_H_
diff --git a/delegate/intel_openvino/openvino_delegate.cc b/delegate/intel_openvino/openvino_delegate.cc
index 7cf7f42..54cd3ba 100644
--- a/delegate/intel_openvino/openvino_delegate.cc
+++ b/delegate/intel_openvino/openvino_delegate.cc
@@ -197,6 +197,11 @@
                  {{kTfLiteFloat32}, {kTfLiteInt32, kTfLiteInt64}}) &&
              CheckDims(context, node, {{1, 2, 3, 4}, {2}});
     }
+    case kTfLiteBuiltinSub: {
+      return CheckDataTypeSupported(context, node,
+                                    {{kTfLiteFloat32}, {kTfLiteFloat32}}) &&
+             CheckDims(context, node, {{1, 2, 3, 4}, {1, 2, 3, 4}});
+    }
     default:
       return false;
   }
@@ -218,10 +223,10 @@
   return "OpenVINO SimpleOpaqueDelegate";
 }
 
-std::unique_ptr<tflite::SimpleOpaqueDelegateKernelInterface>
+std::unique_ptr<cros::SimpleAsyncDelegateKernelInterface>
 OpenVINODelegate::CreateDelegateKernelInterface() {
-  return std::unique_ptr<tflite::openvinodelegate::OpenVINODelegateKernel>(
-      new tflite::openvinodelegate::OpenVINODelegateKernel(options_));
+  return std::unique_ptr<tflite::openvinodelegate::OpenVINOAsyncDelegateKernel>(
+      new tflite::openvinodelegate::OpenVINOAsyncDelegateKernel(options_));
 }
 
 }  // namespace openvinodelegate
diff --git a/delegate/intel_openvino/openvino_delegate.h b/delegate/intel_openvino/openvino_delegate.h
index dbea435..486e69c 100644
--- a/delegate/intel_openvino/openvino_delegate.h
+++ b/delegate/intel_openvino/openvino_delegate.h
@@ -10,6 +10,7 @@
 #include <string>
 #include <vector>
 
+#include "common/simple_async_delegate.h"
 #include "delegate/intel_openvino/log.h"
 #include "tensorflow/lite/delegates/utils/simple_opaque_delegate.h"
 
@@ -32,7 +33,7 @@
 // forward declaration
 class OpenVINODelegateTestPeer;
 
-class OpenVINODelegate : public SimpleOpaqueDelegateInterface {
+class OpenVINODelegate : public cros::SimpleAsyncDelegateInterface {
  public:
   explicit OpenVINODelegate(const TfLiteOpenVINODelegateOptions *options) {
     if (options == nullptr) {
@@ -52,7 +53,7 @@
 
   const char *Name() const override;
 
-  std::unique_ptr<SimpleOpaqueDelegateKernelInterface>
+  std::unique_ptr<cros::SimpleAsyncDelegateKernelInterface>
   CreateDelegateKernelInterface() override;
 
  private:
diff --git a/delegate/intel_openvino/openvino_delegate_core.h b/delegate/intel_openvino/openvino_delegate_core.h
index f4d37e6..2e615e2 100644
--- a/delegate/intel_openvino/openvino_delegate_core.h
+++ b/delegate/intel_openvino/openvino_delegate_core.h
@@ -7,6 +7,7 @@
 #define DELEGATE_INTEL_OPENVINO_OPENVINO_DELEGATE_CORE_H_
 
 #include <openvino/openvino.hpp>
+#include <openvino/runtime/remote_context.hpp>
 
 #include <memory>
 #include <string>
@@ -34,6 +35,10 @@
                            const TfLiteOpenVINODelegateOptions *options);
   TfLiteStatus CompileAndInfer();
 
+  ov::RemoteContext get_context() {
+    return ov_core_.get_default_context("NPU");
+  }
+
  private:
   TfLiteStatus InitializeModel(TfLiteOpaqueContext *context,
                                const TfLiteOpaqueDelegateParams *params,
diff --git a/delegate/intel_openvino/openvino_delegate_external.cc b/delegate/intel_openvino/openvino_delegate_external.cc
index 69a43bd..e6d2ef5 100644
--- a/delegate/intel_openvino/openvino_delegate_external.cc
+++ b/delegate/intel_openvino/openvino_delegate_external.cc
@@ -37,13 +37,13 @@
 
   auto delegate =
       std::make_unique<tflite::openvinodelegate::OpenVINODelegate>(&options);
-  return tflite::TfLiteOpaqueDelegateFactory::CreateSimpleDelegate(
+  return tflite::cros::SimpleAsyncDelegateFactory::CreateAsyncDelegate(
       std::move(delegate));
 }
 
 void OpenVINOStableDelegateDestroyFunc(
     TfLiteOpaqueDelegate *openvino_stable_delegate) {
-  tflite::TfLiteOpaqueDelegateFactory::DeleteSimpleDelegate(
+  tflite::cros::SimpleAsyncDelegateFactory::DeleteAsyncDelegate(
       openvino_stable_delegate);
 }
 
diff --git a/delegate/intel_openvino/openvino_delegate_kernel.cc b/delegate/intel_openvino/openvino_delegate_kernel.cc
index 24ef9bd..a82d01a 100644
--- a/delegate/intel_openvino/openvino_delegate_kernel.cc
+++ b/delegate/intel_openvino/openvino_delegate_kernel.cc
@@ -6,15 +6,24 @@
 #include "delegate/intel_openvino/openvino_delegate_kernel.h"
 
 #include <openvino/runtime/core.hpp>
+#include <openvino/runtime/intel_npu/level_zero/level_zero.hpp>
+#include <openvino/runtime/remote_context.hpp>
+#include <sys/mman.h>
 
+#include <cerrno>
 #include <vector>
 
 #include "delegate/intel_openvino/log.h"
+#include "tensorflow/lite/core/async/c/task.h"
+#include "tensorflow/lite/core/async/interop/c/attribute_map.h"
+#include "tensorflow/lite/delegates/utils/sync_fence.h"
+
+using tflite::delegates::utils::WaitForAllFds;
 
 namespace tflite {
 namespace openvinodelegate {
 
-TfLiteStatus OpenVINODelegateKernel::Init(
+TfLiteStatus OpenVINOAsyncDelegateKernel::Init(
     TfLiteOpaqueContext *context, const TfLiteOpaqueDelegateParams *params) {
   TfLiteStatus init_status = ov_delegate_core_->Init();
   if (init_status != kTfLiteOk) return init_status;
@@ -26,16 +35,21 @@
   return kTfLiteOk;
 }
 
-TfLiteStatus OpenVINODelegateKernel::Prepare(TfLiteOpaqueContext *context,
-                                             TfLiteOpaqueNode *node) {
-  TFLITE_LOG(INFO) << "inside Prepare";
+TfLiteAsyncKernel *OpenVINOAsyncDelegateKernel::AsyncKernel(
+    TfLiteOpaqueContext *context, TfLiteOpaqueNode *node) {
+  return async_kernel_->kernel();
+}
+
+TfLiteStatus OpenVINOAsyncDelegateKernel::Prepare(TfLiteOpaqueContext *context,
+                                                  TfLiteOpaqueNode *node) {
+  absl::MutexLock lock(&prep_mutex_);
   TfLiteStatus set_status = ov_delegate_core_->CompileAndInfer();
   if (set_status != kTfLiteOk) return set_status;
   return kTfLiteOk;
 }
 
-TfLiteStatus OpenVINODelegateKernel::Eval(TfLiteOpaqueContext *context,
-                                          TfLiteOpaqueNode *node) {
+TfLiteStatus OpenVINOAsyncDelegateKernel::Eval(TfLiteOpaqueContext *context,
+                                               TfLiteOpaqueNode *node) {
   std::vector<int> compute_inputs = ov_delegate_core_->getComputeInputs();
   for (int i = 0; i < compute_inputs.size(); i++) {
     int t = compute_inputs[i];
@@ -53,7 +67,7 @@
 
   ov_delegate_core_->getInferRequest().start_async();
   if (!ov_delegate_core_->getInferRequest().wait_for(
-          std::chrono::milliseconds(10000))) {
+          std::chrono::milliseconds(kInferRequestTimeout))) {
     TFLITE_LOG(ERROR) << "Infer request failed";
     return kTfLiteError;
   }
@@ -74,5 +88,114 @@
   return kTfLiteOk;
 }
 
+TfLiteStatus OpenVINOAsyncDelegateKernel::RegisterBuffer(
+    TfLiteBufferHandle handle, AHardwareBuffer *ptr, size_t buffer_size) {
+  size_t numElements = (buffer_size / 4);
+  auto context = ov_delegate_core_->get_context()
+                     .as<ov::intel_npu::level_zero::ZeroContext>();
+  AHardwareBuffer_acquire(ptr);
+  if (ptr == NULL) return kTfLiteError;
+
+  const native_handle_t *buffer_handle = AHardwareBuffer_getNativeHandle(ptr);
+  if (buffer_handle == nullptr) return kTfLiteError;
+  if (buffer_handle->numFds != 1) return kTfLiteError;
+  int fd = buffer_handle->data[0];
+  if (fd == -1) return kTfLiteError;
+  auto mmap_ret =
+      mmap(NULL, buffer_size, PROT_WRITE | PROT_READ, MAP_SHARED, fd, 0);
+  if (mmap_ret == MAP_FAILED) {
+    return kTfLiteError;
+  }
+  mmap_ret_map_.emplace(handle, mmap_ret);
+  buffer_size_map_.emplace(handle, buffer_size);
+  auto remote_tensor =
+      context.create_tensor(ov::element::f32, ov::Shape{1, numElements}, fd);
+  buffer_map_.emplace(handle, remote_tensor);
+  ahwb_buffer_map_.emplace(handle, ptr);
+  return kTfLiteOk;
+}
+
+TfLiteStatus OpenVINOAsyncDelegateKernel::UnregisterBuffer(
+    TfLiteBufferHandle handle) {
+  auto tensor = buffer_map_.find(handle);
+  if (tensor == buffer_map_.end()) return kTfLiteError;
+  auto it = ahwb_buffer_map_.find(handle);
+  if (it == ahwb_buffer_map_.end()) return kTfLiteError;
+  auto mmap_it = mmap_ret_map_.find(handle);
+  if (mmap_it == mmap_ret_map_.end()) return kTfLiteError;
+  auto buffer_size_it = buffer_size_map_.find(handle);
+  if (buffer_size_it == buffer_size_map_.end()) return kTfLiteError;
+
+  munmap(mmap_it->second, buffer_size_it->second);
+  tensor->second = {};
+  AHardwareBuffer_release(it->second);
+  ahwb_buffer_map_.erase(handle);
+  buffer_map_.erase(handle);
+  mmap_ret_map_.erase(handle);
+  buffer_size_map_.erase(handle);
+  return kTfLiteOk;
+}
+
+TfLiteStatus OpenVINOAsyncDelegateKernel::EvalAsyncImpl(
+    TfLiteOpaqueContext *context, TfLiteOpaqueNode *node,
+    TfLiteExecutionTask *task) {
+  absl::MutexLock lock(&eval_mutex_);
+  std::vector<int> compute_inputs = ov_delegate_core_->getComputeInputs();
+  for (int i = 0; i < compute_inputs.size(); i++) {
+    TfLiteSynchronization *sync_ptr =
+        TfLiteExecutionTaskGetSyncByIndex(task, compute_inputs[i]);
+    if (sync_ptr == nullptr) continue;
+    auto sync_obj = TfLiteSynchronizationGetPtr(sync_ptr);
+    if (sync_obj == nullptr) continue;
+
+    input_sync_fence_fds_.push_back(*(reinterpret_cast<int *>(sync_obj)));
+  }
+
+  auto wait_resp = WaitForAllFds(input_sync_fence_fds_);
+  if (!wait_resp.has_value()) return kTfLiteError;
+
+  for (int i = 0; i < compute_inputs.size(); i++) {
+    TfLiteBufferHandle buffer_handle =
+        TfLiteExecutionTaskGetBufferByIndex(task, compute_inputs[i]);
+    const TfLiteOpaqueTensor *tfl_tensor =
+        TfLiteOpaqueNodeGetInput(context, node, i);
+    int32_t num_dims = TfLiteOpaqueTensorNumDims(tfl_tensor);
+    std::vector<int> dims(num_dims);
+    for (int j = 0; j < num_dims; j++) {
+      dims[j] = TfLiteOpaqueTensorDim(tfl_tensor, j);
+    }
+    ov::Tensor remote_tensor = buffer_map_.at(buffer_handle);
+    remote_tensor.set_shape(ov::Shape(dims.begin(), dims.end()));
+    ov_delegate_core_->getInferRequest().set_input_tensor(compute_inputs[i],
+                                                          remote_tensor);
+  }
+
+  std::vector<int> outputs = ov_delegate_core_->getOutputs();
+  for (int o = 0; o < outputs.size(); o++) {
+    TfLiteBufferHandle buffer_handle =
+        TfLiteExecutionTaskGetBufferByIndex(task, outputs[o]);
+    const TfLiteOpaqueTensor *tfl_tensor =
+        TfLiteOpaqueNodeGetOutput(context, node, o);
+    int32_t num_dims = TfLiteOpaqueTensorNumDims(tfl_tensor);
+    std::vector<int> dims(num_dims);
+    for (int j = 0; j < num_dims; j++) {
+      dims[j] = TfLiteOpaqueTensorDim(tfl_tensor, j);
+    }
+
+    ov::Tensor remote_tensor = buffer_map_.at(buffer_handle);
+    remote_tensor.set_shape(ov::Shape(dims.begin(), dims.end()));
+    ov_delegate_core_->getInferRequest().set_output_tensor(o, remote_tensor);
+  }
+
+  ov_delegate_core_->getInferRequest().start_async();
+  if (!ov_delegate_core_->getInferRequest().wait_for(
+          std::chrono::milliseconds(kInferRequestTimeout))) {
+    TFLITE_LOG(ERROR) << "Infer request failed";
+    return kTfLiteError;
+  }
+
+  return kTfLiteOk;
+}
+
 }  // namespace openvinodelegate
 }  // namespace tflite
diff --git a/delegate/intel_openvino/openvino_delegate_kernel.h b/delegate/intel_openvino/openvino_delegate_kernel.h
index f10fb9a..8786512 100644
--- a/delegate/intel_openvino/openvino_delegate_kernel.h
+++ b/delegate/intel_openvino/openvino_delegate_kernel.h
@@ -8,19 +8,29 @@
 
 #include <openvino/openvino.hpp>
 
+#include <map>
 #include <memory>
+#include <vector>
 
+#include "absl/synchronization/mutex.h"
+#include "common/simple_async_delegate.h"
+#include "delegate/intel_openvino/openvino_async_kernel.h"
 #include "delegate/intel_openvino/openvino_delegate.h"
 #include "delegate/intel_openvino/openvino_delegate_core.h"
 #include "tensorflow/lite/delegates/utils/simple_opaque_delegate.h"
 
 namespace tflite {
 namespace openvinodelegate {
-class OpenVINODelegateKernel : public SimpleOpaqueDelegateKernelInterface {
+
+class OVDelegateAsyncKernel;
+
+class OpenVINOAsyncDelegateKernel
+    : public cros::SimpleAsyncDelegateKernelInterface {
  public:
-  OpenVINODelegateKernel(TfLiteOpenVINODelegateOptions options)
+  explicit OpenVINOAsyncDelegateKernel(TfLiteOpenVINODelegateOptions options)
       : ov_delegate_core_(std::make_unique<OpenVINODelegateCore>(
-            "/etc/openvino/plugins.xml")) {
+            "/etc/openvino/plugins.xml")),
+        async_kernel_(std::make_unique<OVDelegateAsyncKernel>(this)) {
     options_ = options;
   }
 
@@ -33,9 +43,27 @@
   TfLiteStatus Eval(TfLiteOpaqueContext *context,
                     TfLiteOpaqueNode *node) override;
 
+  TfLiteStatus EvalAsyncImpl(TfLiteOpaqueContext *context,
+                             TfLiteOpaqueNode *node, TfLiteExecutionTask *task);
+
+  TfLiteStatus RegisterBuffer(TfLiteBufferHandle handle, AHardwareBuffer *ptr,
+                              size_t buffer_size);
+  TfLiteStatus UnregisterBuffer(TfLiteBufferHandle handle);
+  TfLiteAsyncKernel *AsyncKernel(TfLiteOpaqueContext *context,
+                                 TfLiteOpaqueNode *node) override;
+
  private:
+  static constexpr int kInferRequestTimeout = 10000;
   std::unique_ptr<OpenVINODelegateCore> ov_delegate_core_;
+  std::unique_ptr<OVDelegateAsyncKernel> async_kernel_;
   TfLiteOpenVINODelegateOptions options_;
+  std::vector<int> input_sync_fence_fds_ ABSL_GUARDED_BY(eval_mutex_);
+  mutable absl::Mutex eval_mutex_;
+  mutable absl::Mutex prep_mutex_;
+  std::map<TfLiteBufferHandle, ov::RemoteTensor> buffer_map_;
+  std::map<TfLiteBufferHandle, AHardwareBuffer *> ahwb_buffer_map_;
+  std::map<TfLiteBufferHandle, void *> mmap_ret_map_;
+  std::map<TfLiteBufferHandle, size_t> buffer_size_map_;
 };
 
 }  // namespace openvinodelegate
diff --git a/delegate/intel_openvino/openvino_graph_builder.cc b/delegate/intel_openvino/openvino_graph_builder.cc
index f64b235..6e06b82 100644
--- a/delegate/intel_openvino/openvino_graph_builder.cc
+++ b/delegate/intel_openvino/openvino_graph_builder.cc
@@ -24,6 +24,7 @@
 #include "delegate/intel_openvino/operations/include/reshape.h"
 #include "delegate/intel_openvino/operations/include/resize_bilinear.h"
 #include "delegate/intel_openvino/operations/include/softmax.h"
+#include "delegate/intel_openvino/operations/include/sub.h"
 #include "delegate/intel_openvino/operations/include/tanh.h"
 #include "delegate/intel_openvino/operations/include/transpose_conv.h"
 #include "delegate/intel_openvino/operations/utility.h"
@@ -172,6 +173,9 @@
     case kTfLiteBuiltinSoftmax: {
       return std::make_shared<Softmax>();
     }
+    case kTfLiteBuiltinSub: {
+      return std::make_shared<Sub>();
+    }
     case kTfLiteBuiltinTanh: {
       return std::make_shared<Tanh>();
     }
diff --git a/delegate/intel_openvino/operations/BUILD.bazel b/delegate/intel_openvino/operations/BUILD.bazel
index 6addfe8..94b70fa 100644
--- a/delegate/intel_openvino/operations/BUILD.bazel
+++ b/delegate/intel_openvino/operations/BUILD.bazel
@@ -29,6 +29,7 @@
         "src/reshape.cc",
         "src/resize_bilinear.cc",
         "src/softmax.cc",
+        "src/sub.cc",
         "src/tanh.cc",
         "src/transpose_conv.cc",
     ],
@@ -50,6 +51,7 @@
         "include/reshape.h",
         "include/resize_bilinear.h",
         "include/softmax.h",
+        "include/sub.h",
         "include/tanh.h",
         "include/transpose_conv.h",
         "openvino_node_manager.h",
diff --git a/delegate/intel_openvino/operations/include/sub.h b/delegate/intel_openvino/operations/include/sub.h
new file mode 100644
index 0000000..5860d79
--- /dev/null
+++ b/delegate/intel_openvino/operations/include/sub.h
@@ -0,0 +1,22 @@
+/*
+ * Copyright (C) 2024 Intel Corporation
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#ifndef DELEGATE_INTEL_OPENVINO_OPERATIONS_INCLUDE_SUB_H_
+#define DELEGATE_INTEL_OPENVINO_OPERATIONS_INCLUDE_SUB_H_
+
+#include "delegate/intel_openvino/operations/operations_base.h"
+
+namespace tflite {
+namespace openvinodelegate {
+
+class Sub final : public OperationsBase {
+ public:
+  Sub() {}
+  TfLiteStatus CreateNode() override;
+};
+
+}  // namespace openvinodelegate
+}  // namespace tflite
+#endif  // DELEGATE_INTEL_OPENVINO_OPERATIONS_INCLUDE_SUB_H_
diff --git a/delegate/intel_openvino/operations/src/sub.cc b/delegate/intel_openvino/operations/src/sub.cc
new file mode 100644
index 0000000..5ff5ef5
--- /dev/null
+++ b/delegate/intel_openvino/operations/src/sub.cc
@@ -0,0 +1,31 @@
+/*
+ * Copyright (C) 2024 Intel Corporation
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#include "delegate/intel_openvino/operations/include/sub.h"
+
+namespace tflite {
+namespace openvinodelegate {
+
+TfLiteStatus Sub::CreateNode() {
+  auto *sub_params = GetBuiltinData<TfLiteSubParams>();
+  auto input_node_1 = getInputNode(tensor_indices_[INPUT_NODE_1]);
+  if (input_node_1 == nullptr) {
+    TFLITE_LOG(INFO) << "input node 1 is null";
+    return kTfLiteError;
+  }
+  auto input_node_2 = getInputNode(tensor_indices_[INPUT_NODE_2]);
+  if (input_node_2 == nullptr) {
+    TFLITE_LOG(INFO) << "input Node 2 is null";
+    return kTfLiteError;
+  }
+
+  auto sub_node = std::make_shared<ov::opset8::Subtract>(
+      input_node_1, input_node_2, ov::op::AutoBroadcastType::NUMPY);
+  output_node_ = ApplyActivation(sub_node, sub_params->activation);
+  return kTfLiteOk;
+}
+
+}  // namespace openvinodelegate
+}  // namespace tflite