tflite: Extract core computation from sample delegate kernel This CL extracts the core computation part from the sample delegate kernel. This common core would be used by both sync and async kernels. A new method SetExternalTensorMemory() is added to allow providing cpu-accessible memory for external tensors, either allocated by TFLite runtime or registered by application via async kernel API. BUG=b:305999697 TEST=Pass DTS locally with sample delegate. Change-Id: I4ae8d42af92ebe1ce41c6a8e39fee587a28a2704 Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/platform/tflite/+/5367488 Commit-Queue: Shik Chen <shik@chromium.org> Reviewed-by: Tommy Chiang <ototot@google.com> Tested-by: Shik Chen <shik@chromium.org>
diff --git a/delegate/sample/BUILD.bazel b/delegate/sample/BUILD.bazel index 09534ef..e452fe8 100644 --- a/delegate/sample/BUILD.bazel +++ b/delegate/sample/BUILD.bazel
@@ -6,6 +6,23 @@ load("@org_tensorflow//tensorflow/lite/core/shims:cc_library_with_tflite.bzl", "cc_library_with_tflite") cc_library_with_tflite( + name = "core", + srcs = [ + "core.cc", + ], + hdrs = [ + "core.h", + ], + copts = tflite_copts(), + generate_opaque_delegate_target = True, + tflite_deps = [ + "@org_tensorflow//tensorflow/lite/c:c_api", + "@org_tensorflow//tensorflow/lite/c:c_api_experimental", + "@org_tensorflow//tensorflow/lite/c:common", + ], +) + +cc_library_with_tflite( name = "kernel", srcs = [ "kernel.cc", @@ -16,6 +33,7 @@ copts = tflite_copts(), generate_opaque_delegate_target = True, tflite_deps = [ + ":core", "//common:simple_async_delegate", "@org_tensorflow//tensorflow/lite/c:c_api", "@org_tensorflow//tensorflow/lite/c:c_api_experimental",
diff --git a/delegate/sample/core.cc b/delegate/sample/core.cc new file mode 100644 index 0000000..2a3c336 --- /dev/null +++ b/delegate/sample/core.cc
@@ -0,0 +1,125 @@ +/* + * Copyright 2024 The ChromiumOS Authors + * Use of this source code is governed by a BSD-style license that can be + * found in the LICENSE file. + */ + +#include "delegate/sample/core.h" + +#include <algorithm> +#include <map> +#include <set> +#include <vector> + +#include "tensorflow/lite/core/c/c_api_opaque.h" + +namespace tflite::cros { + +namespace { + +float AddImpl(float a, float b) { + return a + b; +} + +float SubImpl(float a, float b) { + return a - b; +} + +int CalculateNumElements(const TfLiteOpaqueTensor* tensor) { + int num_elements = 1; + int num_dims = TfLiteOpaqueTensorNumDims(tensor); + for (int i = 0; i < num_dims; ++i) { + num_elements *= TfLiteOpaqueTensorDim(tensor, i); + } + return num_elements; +} + +} // namespace + +TfLiteStatus CrosSampleDelegateCore::Init( + TfLiteOpaqueContext* context, + const TfLiteOpaqueDelegateParams* params) { + int num_nodes = params->nodes_to_replace->size; + node_infos_.reserve(num_nodes); + for (int i = 0; i < num_nodes; ++i) { + TfLiteOpaqueNode* node = nullptr; + TfLiteRegistrationExternal* registration = nullptr; + int node_index = params->nodes_to_replace->data[i]; + TfLiteOpaqueContextGetNodeAndRegistration(context, node_index, &node, + ®istration); + + NodeInfo info = { + .op = TfLiteRegistrationExternalGetBuiltInCode(registration), + .input1 = TfLiteOpaqueNodeGetInput(context, node, 0), + .input2 = TfLiteOpaqueNodeGetInput(context, node, 1), + .output = TfLiteOpaqueNodeGetOutput(context, node, 0), + }; + node_infos_.push_back(info); + } + + using TensorSet = std::set<const TfLiteOpaqueTensor*>; + TensorSet all_inputs; + TensorSet all_outputs; + for (const auto& info : node_infos_) { + all_inputs.insert(info.input1); + all_inputs.insert(info.input2); + all_outputs.insert(info.output); + } + // If the input of some node is an output of some node in the same delegated + // subgraph, it's an internal tensor for us. + std::set_intersection(all_inputs.begin(), all_inputs.end(), + all_outputs.begin(), all_outputs.end(), + std::back_inserter(internal_tensors_)); + + return kTfLiteOk; +} + +TfLiteStatus CrosSampleDelegateCore::Prepare() { + // Allocate memory for internal tensors. For external tensors, the memory will + // be provided with SetExternalTensorMemory(). + for (const auto& tensor : internal_tensors_) { + int size = CalculateNumElements(tensor); + internal_tensors_memory_[tensor].resize(size); + } + + return kTfLiteOk; +} + +float* CrosSampleDelegateCore::GetRawDataSource( + const TfLiteOpaqueTensor* tensor) { + if (auto it = internal_tensors_memory_.find(tensor); + it != internal_tensors_memory_.end()) { + return it->second.data(); + } + + if (auto it = external_tensors_memory_.find(tensor); + it != external_tensors_memory_.end()) { + return static_cast<float*>(it->second); + } + + // Fall back to use data pointer inside the tensor. Normally this should not + // happen. + return reinterpret_cast<float*>(TfLiteOpaqueTensorData(tensor)); +} + +void CrosSampleDelegateCore::SetExternalTensorMemory( + const TfLiteOpaqueTensor* tensor, + void* memory) { + external_tensors_memory_.insert_or_assign(tensor, memory); +} + +TfLiteStatus CrosSampleDelegateCore::Eval() { + for (const auto& info : node_infos_) { + float* input1 = GetRawDataSource(info.input1); + float* input2 = GetRawDataSource(info.input2); + float* output = GetRawDataSource(info.output); + + // The input/output tensors have the same size. + int size = CalculateNumElements(info.output); + std::transform(input1, input1 + size, input2, output, + info.op == kTfLiteBuiltinAdd ? AddImpl : SubImpl); + } + return kTfLiteOk; +} + +} // namespace tflite::cros
diff --git a/delegate/sample/core.h b/delegate/sample/core.h new file mode 100644 index 0000000..f6c4574 --- /dev/null +++ b/delegate/sample/core.h
@@ -0,0 +1,83 @@ +/* + * Copyright 2024 The ChromiumOS Authors + * Use of this source code is governed by a BSD-style license that can be + * found in the LICENSE file. + */ + +#ifndef DELEGATE_SAMPLE_CORE_H_ +#define DELEGATE_SAMPLE_CORE_H_ + +// TODO(shik): Note that the upstream sample delegate use flat_hash_{map,set} +// from absl/container. We haven't decided whether we would like to allow +// absl/libchrome usage here. +#include <map> +#include <vector> + +#include "tensorflow/lite/core/c/c_api_opaque.h" + +namespace tflite::cros { + +// The core of delegate kernel that implement the computation of addition and +// subtraction operations. +// The typical flow would be: +// 1. Initialize the core with Init(). +// 2. Prepare the memory for internal tensors with Prepare(). +// 3. Provide the memory for external tensors with SetExternalTensorMemory(). +// 4. Run the model inference with Eval(). +// +// The steps 3 and 4 can be performed multiple times. It's ok to skip step 3 if +// the backing memories are the same as the previous run. +// +// This class is thread-compatible. +class CrosSampleDelegateCore { + public: + // Move-only. + CrosSampleDelegateCore() = default; + CrosSampleDelegateCore(CrosSampleDelegateCore&& other) = default; + CrosSampleDelegateCore& operator=(CrosSampleDelegateCore&& other) = default; + CrosSampleDelegateCore(const CrosSampleDelegateCore&) = delete; + CrosSampleDelegateCore& operator=(const CrosSampleDelegateCore&) = delete; + + // Initializes the kernel by extracting necessary information into + // `node_infos_` and deriving internal tensors. + TfLiteStatus Init(TfLiteOpaqueContext* context, + const TfLiteOpaqueDelegateParams* params); + + // Prepares the memory for internal tensors. Note that this need be called + // everytime tensors are resized so the memory could grow/shrink accordingly. + TfLiteStatus Prepare(); + + // Sets the memory pointer for an external tensor. The memory must be + // cpu-accessible and synchronized. + // TODO(shik): Expose an API for fine-grained synchronizations. + void SetExternalTensorMemory(const TfLiteOpaqueTensor* tensor, void* memory); + + // Computes the inference result. Note that the data pointer of external + // tensors may change for every inference. + TfLiteStatus Eval(); + + private: + struct NodeInfo { + TfLiteBuiltinOperator op; + const TfLiteOpaqueTensor* input1; + const TfLiteOpaqueTensor* input2; + const TfLiteOpaqueTensor* output; + }; + + float* GetRawDataSource(const TfLiteOpaqueTensor* tensor); + + std::vector<NodeInfo> node_infos_; + std::vector<const TfLiteOpaqueTensor*> internal_tensors_; + + // Owned memory for internal tensors. + std::map<const TfLiteOpaqueTensor*, std::vector<float>> + internal_tensors_memory_; + + // Non-owned memory for external tensors provided by + // SetExternalTensorMemory(). + std::map<const TfLiteOpaqueTensor*, void*> external_tensors_memory_; +}; + +} // namespace tflite::cros + +#endif // DELEGATE_SAMPLE_CORE_H_
diff --git a/delegate/sample/kernel.cc b/delegate/sample/kernel.cc index a177ce6..a80947e 100644 --- a/delegate/sample/kernel.cc +++ b/delegate/sample/kernel.cc
@@ -6,109 +6,38 @@ #include "delegate/sample/kernel.h" -#include <algorithm> -#include <map> -#include <set> -#include <vector> - #include "tensorflow/lite/core/c/c_api_opaque.h" namespace tflite::cros { -namespace { - -float AddImpl(float a, float b) { - return a + b; -} - -float SubImpl(float a, float b) { - return a - b; -} - -int CalculateNumElements(const TfLiteOpaqueTensor* tensor) { - int num_elements = 1; - int num_dims = TfLiteOpaqueTensorNumDims(tensor); - for (int i = 0; i < num_dims; ++i) { - num_elements *= TfLiteOpaqueTensorDim(tensor, i); - } - return num_elements; -} - -} // namespace - TfLiteStatus CrosSampleDelegateKernel::Init( TfLiteOpaqueContext* context, const TfLiteOpaqueDelegateParams* params) { - int num_nodes = params->nodes_to_replace->size; - node_infos_.reserve(num_nodes); - for (int i = 0; i < num_nodes; ++i) { - TfLiteOpaqueNode* node = nullptr; - TfLiteRegistrationExternal* registration = nullptr; - int node_index = params->nodes_to_replace->data[i]; - TfLiteOpaqueContextGetNodeAndRegistration(context, node_index, &node, - ®istration); - - NodeInfo info = { - .op = TfLiteRegistrationExternalGetBuiltInCode(registration), - .input1 = TfLiteOpaqueNodeGetInput(context, node, 0), - .input2 = TfLiteOpaqueNodeGetInput(context, node, 1), - .output = TfLiteOpaqueNodeGetOutput(context, node, 0), - }; - node_infos_.push_back(info); - } - - using TensorSet = std::set<const TfLiteOpaqueTensor*>; - TensorSet all_inputs; - TensorSet all_outputs; - for (const auto& info : node_infos_) { - all_inputs.insert(info.input1); - all_inputs.insert(info.input2); - all_outputs.insert(info.output); - } - // If the input of some node is an output of some node in the same delegated - // subgraph, it's an internal tensor for us. - std::set_intersection(all_inputs.begin(), all_inputs.end(), - all_outputs.begin(), all_outputs.end(), - std::back_inserter(internal_tensors_)); - - return kTfLiteOk; + return core_.Init(context, params); } TfLiteStatus CrosSampleDelegateKernel::Prepare(TfLiteOpaqueContext* context, TfLiteOpaqueNode* node) { - // The TFLite runtime takes care of external tensors. - // For internal tensors, we need to allocate memory for them. - for (const auto& tensor : internal_tensors_) { - int size = CalculateNumElements(tensor); - internal_tensors_memory_[tensor].resize(size); - } - - return kTfLiteOk; -} - -float* CrosSampleDelegateKernel::GetRawDataSource( - const TfLiteOpaqueTensor* tensor) { - auto it = internal_tensors_memory_.find(tensor); - if (it != internal_tensors_memory_.end()) { - return it->second.data(); - } else { - return reinterpret_cast<float*>(TfLiteOpaqueTensorData(tensor)); - } + return core_.Prepare(); } TfLiteStatus CrosSampleDelegateKernel::Eval(TfLiteOpaqueContext* context, TfLiteOpaqueNode* node) { - for (const auto& info : node_infos_) { - float* input1 = GetRawDataSource(info.input1); - float* input2 = GetRawDataSource(info.input2); - float* output = GetRawDataSource(info.output); - - // The input/output tensors have the same size. - int size = CalculateNumElements(info.output); - std::transform(input1, input1 + size, input2, output, - info.op == kTfLiteBuiltinAdd ? AddImpl : SubImpl); + int num_inputs = TfLiteOpaqueNodeNumberOfInputs(node); + for (int i = 0; i < num_inputs; ++i) { + auto tensor = TfLiteOpaqueNodeGetInput(context, node, i); + auto data = TfLiteOpaqueTensorData(tensor); + core_.SetExternalTensorMemory(tensor, data); } - return kTfLiteOk; + + int num_outputs = TfLiteOpaqueNodeNumberOfOutputs(node); + for (int i = 0; i < num_outputs; ++i) { + auto tensor = TfLiteOpaqueNodeGetOutput(context, node, i); + auto data = TfLiteOpaqueTensorData(tensor); + core_.SetExternalTensorMemory(tensor, data); + } + + return core_.Eval(); } TfLiteAsyncKernel* CrosSampleDelegateKernel::AsyncKernel(
diff --git a/delegate/sample/kernel.h b/delegate/sample/kernel.h index 5c6b9b5..de1effb 100644 --- a/delegate/sample/kernel.h +++ b/delegate/sample/kernel.h
@@ -14,49 +14,37 @@ #include <vector> #include "common/simple_async_delegate.h" +#include "delegate/sample/core.h" #include "tensorflow/lite/core/c/c_api_opaque.h" namespace tflite::cros { // A simple delegate kernel that implement the computation of addition and -// subtraction operations. +// subtraction operations. This class is thread-compatible. class CrosSampleDelegateKernel : public SimpleAsyncDelegateKernelInterface { public: - // Initializes the kernel by extracting necessary information into - // `node_infos_` and deriving internal tensors. + // Move-only. + CrosSampleDelegateKernel() = default; + CrosSampleDelegateKernel(CrosSampleDelegateKernel&& other) = default; + CrosSampleDelegateKernel& operator=(CrosSampleDelegateKernel&& other) = + default; + CrosSampleDelegateKernel(const CrosSampleDelegateKernel&) = delete; + CrosSampleDelegateKernel& operator=(const CrosSampleDelegateKernel&) = delete; + + // Implementation of SimpleAsyncDelegateKernelInterface. TfLiteStatus Init(TfLiteOpaqueContext* context, const TfLiteOpaqueDelegateParams* params) override; - - // Prepares the memory for internal tensors. Note that this will be called - // everytime tensors are resized. TfLiteStatus Prepare(TfLiteOpaqueContext* context, TfLiteOpaqueNode* node) override; - - // Computes the inference result. Note that the data pointer of external - // tensors may change for every inference. TfLiteStatus Eval(TfLiteOpaqueContext* context, TfLiteOpaqueNode* node) override; - - // Retrieves the async kernel. // Returns nullptr for now, since the delegate does not support asynchronous // execution yet. TfLiteAsyncKernel* AsyncKernel(TfLiteOpaqueContext* context, TfLiteOpaqueNode* node) override; private: - struct NodeInfo { - TfLiteBuiltinOperator op; - const TfLiteOpaqueTensor* input1; - const TfLiteOpaqueTensor* input2; - const TfLiteOpaqueTensor* output; - }; - - float* GetRawDataSource(const TfLiteOpaqueTensor* tensor); - - std::vector<NodeInfo> node_infos_; - std::vector<const TfLiteOpaqueTensor*> internal_tensors_; - std::map<const TfLiteOpaqueTensor*, std::vector<float>> - internal_tensors_memory_; + CrosSampleDelegateCore core_; }; } // namespace tflite::cros