chore: import upstream snapshot with attribution

tensorflow/lite/delegates/BUILD

@@ -0,0 +1,245 @@
+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+load("@rules_cc//cc:cc_library.bzl", "cc_library")
+load("@rules_cc//cc:cc_test.bzl", "cc_test")
+load("//tensorflow:tensorflow.default.bzl", "get_compatible_with_portable")
+load("//tensorflow/lite:build_def.bzl", "tflite_copts", "tflite_linkopts")
+load("//tensorflow/lite:special_rules.bzl", "tflite_portable_test_suite")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:LICENSE"],
+    default_visibility = ["//visibility:public"],
+    licenses = ["notice"],
+)
+
+config_setting(
+    name = "tflite_debug_delegate",
+    define_values = {"tflite_debug_delegate": "true"},
+)
+
+cc_library(
+    name = "telemetry",
+    srcs = ["telemetry.cc"],
+    hdrs = ["telemetry.h"],
+    compatible_with = get_compatible_with_portable(),
+    copts = tflite_copts(),
+    deps = [
+        "//tensorflow/lite/acceleration/configuration:configuration_fbs",
+        "//tensorflow/lite/core/api",
+        "//tensorflow/lite/core/c:common",
+    ],
+)
+
+cc_test(
+    name = "telemetry_test",
+    srcs = ["telemetry_test.cc"],
+    linkopts = tflite_linkopts(),
+    linkstatic = 1,
+    deps = [
+        ":telemetry",
+        "//tensorflow/lite/acceleration/configuration:configuration_fbs",
+        "//tensorflow/lite/core/api",
+        "//tensorflow/lite/core/c:common",
+        "//tensorflow/lite/profiling:profile_buffer",
+        "@com_google_googletest//:gtest_main",
+        "@flatbuffers",
+    ],
+)
+
+cc_library(
+    name = "utils",
+    srcs = ["utils.cc"],
+    hdrs = ["utils.h"],
+    compatible_with = get_compatible_with_portable(),
+    copts = tflite_copts(),
+    deps = [
+        "//tensorflow/lite:kernel_api",
+        "//tensorflow/lite:util",
+        "//tensorflow/lite/c:c_api_types",
+        "//tensorflow/lite/c:common",
+        "//tensorflow/lite/core:subgraph",
+        "//tensorflow/lite/core/c:common",
+        "//tensorflow/lite/kernels:kernel_util",
+    ],
+)
+
+cc_library(
+    name = "interpreter_utils",
+    srcs = ["interpreter_utils.cc"],
+    hdrs = ["interpreter_utils.h"],
+    compatible_with = get_compatible_with_portable(),
+    copts = tflite_copts(),
+    deps = [
+        "//tensorflow/lite:framework",
+        "//tensorflow/lite/c:c_api_types",
+        "//tensorflow/lite/c:common",
+        "//tensorflow/lite/core/api:error_reporter",
+    ],
+)
+
+cc_test(
+    name = "utils_test",
+    srcs = ["utils_test.cc"],
+    linkopts = tflite_linkopts(),
+    linkstatic = 1,
+    deps = [
+        ":utils",
+        "//tensorflow/lite/core:subgraph",
+        "//tensorflow/lite/core/c:common",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
+
+cc_test(
+    name = "interpreter_utils_test",
+    size = "small",
+    srcs = ["interpreter_utils_test.cc"],
+    features = ["-dynamic_link_test_srcs"],  # see go/dynamic_link_test_srcs
+    deps = [
+        ":delegate_test_util",
+        ":interpreter_utils",
+        "//tensorflow/lite:framework",
+        "//tensorflow/lite/core/c:common",
+        "//tensorflow/lite/kernels:builtin_ops",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
+
+cc_test(
+    name = "delegate_test",
+    size = "small",
+    srcs = ["delegate_test.cc"],
+    data = ["//tensorflow/lite:testdata/add.bin"],
+    features = ["-dynamic_link_test_srcs"],  # see go/dynamic_link_test_srcs.
+    deps = [
+        ":delegate_test_util",
+        "//tensorflow/compiler/mlir/lite/experimental/remat:metadata_util",
+        "//tensorflow/compiler/mlir/lite/schema:schema_conversion_utils",
+        "//tensorflow/lite:builtin_ops",
+        "//tensorflow/lite:framework",
+        "//tensorflow/lite:version",
+        "//tensorflow/lite/core:framework",
+        "//tensorflow/lite/core/c:c_api_experimental",
+        "//tensorflow/lite/core/c:c_api_types",
+        "//tensorflow/lite/core/c:common",
+        "//tensorflow/lite/core/kernels:builtin_ops",
+        "//tensorflow/lite/delegates/external:external_delegate",
+        "//tensorflow/lite/kernels:builtin_ops",
+        "//tensorflow/lite/kernels:kernel_util",
+        "//tensorflow/lite/schema:schema_fbs",
+        "@com_google_googletest//:gtest_main",
+        "@flatbuffers",
+    ],
+)
+
+cc_test(
+    name = "opaque_delegate_test",
+    size = "small",
+    srcs = ["opaque_delegate_test.cc"],
+    data = [
+        "//tensorflow/lite:testdata/add.bin",
+        "//tensorflow/lite:testdata/test_custom_node_with_init_data.bin",
+    ],
+    defines = [
+        "TFLITE_USE_OPAQUE_DELEGATE",
+    ],
+    deps = [
+        "//tensorflow/lite:builtin_ops",
+        "//tensorflow/lite:framework",
+        "//tensorflow/lite/c:c_api_experimental",
+        "//tensorflow/lite/c:common",
+        "//tensorflow/lite/core/c:c_api_types",
+        "//tensorflow/lite/core/c:common",
+        "//tensorflow/lite/kernels:builtin_ops",
+        "//tensorflow/lite/schema:schema_fbs",
+        "//tensorflow/lite/testing:util",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
+
+cc_test(
+    name = "opaque_delegate_strip_error_strings_test",
+    size = "small",
+    srcs = ["opaque_delegate_test.cc"],
+    data = [
+        "//tensorflow/lite:testdata/add.bin",
+        "//tensorflow/lite:testdata/test_custom_node_with_init_data.bin",
+    ],
+    defines = [
+        "TFLITE_USE_OPAQUE_DELEGATE",
+        "TF_LITE_STRIP_ERROR_STRINGS",
+    ],
+    deps = [
+        "//tensorflow/lite:builtin_ops",
+        "//tensorflow/lite:framework",
+        "//tensorflow/lite/c:c_api_experimental",
+        "//tensorflow/lite/c:common",
+        "//tensorflow/lite/core/c:c_api_types",
+        "//tensorflow/lite/core/c:common",
+        "//tensorflow/lite/kernels:builtin_ops",
+        "//tensorflow/lite/schema:schema_fbs",
+        "//tensorflow/lite/testing:util",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
+
+cc_library(
+    name = "delegate_test_util",
+    testonly = True,
+    srcs = ["delegate_test_util.cc"],
+    hdrs = ["delegate_test_util.h"],
+    deps = [
+        ":utils",
+        "//tensorflow/lite:builtin_ops",
+        "//tensorflow/lite:util",
+        "//tensorflow/lite/core:framework",
+        "//tensorflow/lite/core/c:common",
+        "//tensorflow/lite/core/kernels:builtin_ops",
+        "//tensorflow/lite/kernels:kernel_util",
+        "//tensorflow/lite/kernels/internal:compatibility",
+        "//tensorflow/lite/schema:schema_fbs",
+        "@com_google_googletest//:gtest",
+        "@eigen_archive//:eigen3",
+    ],
+)
+
+cc_library(
+    name = "serialization",
+    srcs = ["serialization.cc"],
+    hdrs = ["serialization.h"],
+    compatible_with = get_compatible_with_portable(),
+    copts = tflite_copts(),
+    deps = [
+        "//tensorflow/lite:minimal_logging",
+        "//tensorflow/lite/core/c:common",
+        "@farmhash_archive//:farmhash",
+    ],
+)
+
+cc_test(
+    name = "serialization_test",
+    srcs = ["serialization_test.cc"],
+    linkopts = tflite_linkopts(),
+    linkstatic = 1,
+    deps = [
+        ":serialization",
+        "//tensorflow/lite:util",
+        "//tensorflow/lite/core/c:common",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
+
+tflite_portable_test_suite()

tensorflow/lite/delegates/coreml/README.md

@@ -0,0 +1,197 @@
+# Tensorflow Lite Core ML Delegate
+
+TensorFlow Lite Core ML Delegate enables running TensorFlow Lite models on
+[Core ML framework](https://developer.apple.com/documentation/coreml),
+which results in faster model inference on iOS devices.
+
+[TOC]
+
+## Supported iOS versions and processors
+
+* iOS 12 and later. In the older iOS versions, Core ML delegate will
+  automatically fallback to CPU.
+* When running on iPhone Xs and later, it will use Neural Engine for faster
+  inference.
+
+## Update code to use Core ML delegate
+
+### Swift
+
+Initialize TensorFlow Lite interpreter with Core ML delegate.
+
+```swift
+let coreMlDelegate = CoreMLDelegate()
+let interpreter = try Interpreter(modelPath: modelPath,
+                                  delegates: [coreMLDelegate])
+```
+
+### Objective-C++
+
+#### Interpreter initialization
+
+Include `coreml_delegate.h`.
+
+```objectivec++
+#include "tensorflow/lite/experimental/delegates/coreml/coreml_delegate.h"
+```
+
+Modify code following interpreter initialization to apply delegate.
+
+```objectivec++
+// initializer interpreter with model.
+tflite::InterpreterBuilder(*model, resolver)(&interpreter);
+
+// Add following section to use Core ML delegate.
+TfLiteCoreMlDelegateOptions options = {};
+delegate = TfLiteCoreMlDelegateCreate(&options);
+interpreter->ModifyGraphWithDelegate(delegate);
+
+// Any calls to AllocateTensors must happen strictly AFTER all
+// ModifyGraphWithDelegate calls.
+
+// ...
+```
+
+#### Disposal
+
+Add this code to the section where you dispose of the delegate (e.g. `dealloc`
+of class).
+
+```objectivec++
+TfLiteCoreMlDelegateDelete(delegate);
+```
+
+## Supported ops
+
+Following ops are supported by the Core ML delegate.
+
+*   Add
+    *   Only certain shapes are broadcastable. In Core ML tensor layout,
+        following tensor shapes are broadcastable. `[B, C, H, W]`, `[B, C, 1,
+        1]`, `[B, 1, H, W]`, `[B, 1, 1, 1]`.
+*   AveragePool2D
+*   Concat
+*   Conv2D
+    *   Weights and bias should be constant.
+*   DepthwiseConv2D
+    *   Weights and bias should be constant.
+*   FullyConnected (aka Dense or InnerProduct)
+    *   Weights and bias (if present) should be constant.
+    *   Only supports single-batch case. Input dimensions should be 1, except
+        the last dimension.
+*   Hardswish
+*   Logistic (aka Sigmoid)
+*   MaxPool2D
+*   MirrorPad
+    *   Only 4D input with `REFLECT` mode is supported. Padding should be
+        constant, and is only allowed for H and W dimensions.
+*   Mul
+    *   Only certain shapes are broadcastable. In Core ML tensor layout,
+        following tensor shapes are broadcastable. `[B, C, H, W]`, `[B, C, 1,
+        1]`, `[B, 1, H, W]`, `[B, 1, 1, 1]`.
+*   Pad and PadV2
+    *   Only 4D input is supported. Padding should be constant, and is only
+        allowed for H and W dimensions.
+*   Relu
+*   ReluN1To1
+*   Relu6
+*   Reshape
+    *   Only supported when target Core ML version is 2, not supported when
+        targeting Core ML 3.
+*   ResizeBilinear
+*   SoftMax
+*   Tanh
+*   TransposeConv
+    *   Weights should be constant.
+
+## FAQ
+
+* Does Core ML delegate support fallback to CPU if a graph contains unsupported
+  ops?
+  * Yes.
+* Does Core ML delegate work on iOS Simulator?
+  * Yes. The library includes x86 and x86_64 targets so it can run on
+    a simulator, but you will not see performance boost over CPU.
+* Does TensorFlow Lite and Core ML delegate support macOS?
+  * TensorFlow Lite is only tested on iOS but not macOS.
+* Are custom TF Lite ops supported?
+  * No, CoreML delegate does not support custom ops and they will fallback to
+    CPU.
+
+## Appendix
+
+### Core ML delegate Swift API
+
+```swift
+/// A delegate that uses the `Core ML` framework for performing
+/// TensorFlow Lite graph operations.
+///
+/// - Important: This is an experimental interface that is subject to change.
+public final class CoreMLDelegate: Delegate {
+ /// The configuration options for the `CoreMLDelegate`.
+ public let options: Options
+
+ // Conformance to the `Delegate` protocol.
+ public private(set) var cDelegate: CDelegate
+
+ * /// Creates a new instance configured with the given `options`.
+ ///
+ /// - Parameters:
+ ///   - options: Configurations for the delegate. The default is a new instance of
+ ///       `CoreMLDelegate.Options` with the default configuration values.
+ public init(options: Options = Options()) {
+   self.options = options
+   var delegateOptions = TfLiteCoreMlDelegateOptions()
+   cDelegate = TfLiteCoreMlDelegateCreate(&delegateOptions)
+ }
+
+ deinit {
+   TfLiteCoreMlDelegateDelete(cDelegate)
+ }
+}
+
+extension CoreMLDelegate {
+ /// Options for configuring the `CoreMLDelegate`.
+ public struct Options: Equatable, Hashable {
+   /// Creates a new instance with the default values.
+   public init() {}
+ }
+}
+```
+
+### Core ML delegate C++ API
+
+```c++
+typedef struct {
+  // Only create delegate when Neural Engine is available on the device.
+  TfLiteCoreMlDelegateEnabledDevices enabled_devices;
+  // Specifies target Core ML version for model conversion.
+  // Core ML 3 come with a lot more ops, but some ops (e.g. reshape) is not
+  // delegated due to input rank constraint.
+  // if not set to one of the valid versions, the delegate will use highest
+  // version possible in the platform.
+  // Valid versions: (2, 3)
+  int coreml_version;
+  // This sets the maximum number of Core ML delegates created.
+  // Each graph corresponds to one delegated node subset in the
+  // TFLite model. Set this to 0 to delegate all possible partitions.
+  int max_delegated_partitions;
+  // This sets the minimum number of nodes per partition delegated with
+  // Core ML delegate. Defaults to 2.
+  int min_nodes_per_partition;
+#ifdef TFLITE_DEBUG_DELEGATE
+  // This sets the index of the first node that could be delegated.
+  int first_delegate_node_index;
+  // This sets the index of the last node that could be delegated.
+  int last_delegate_node_index;
+#endif
+} TfLiteCoreMlDelegateOptions;
+
+// Return a delegate that uses CoreML for ops execution.
+// Must outlive the interpreter.
+TfLiteDelegate* TfLiteCoreMlDelegateCreate(
+   const TfLiteCoreMlDelegateOptions* options);
+
+// Do any needed cleanup and delete 'delegate'.
+void TfLiteCoreMlDelegateDelete(TfLiteDelegate* delegate);
+```

tensorflow/lite/delegates/coreml/builders/activation_layer_builder.cc

@@ -0,0 +1,143 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include "tensorflow/lite/delegates/coreml/builders/activation_layer_builder.h"
+
+#include <cstdio>
+#include <string>
+
+#include "mlmodel/format/NeuralNetwork.pb.h"
+#include "tensorflow/lite/core/c/builtin_op_data.h"
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/delegates/coreml/builders/op_factory.h"
+#include "tensorflow/lite/delegates/coreml/builders/threshold_layer_builder.h"
+
+namespace tflite {
+namespace delegates {
+namespace coreml {
+
+const std::string& ActivationLayerBuilder::DebugName() {
+  if (debug_name_.empty()) SetDebugName("ActivationLayerBuilder", node_id_);
+  return debug_name_;
+}
+
+CoreML::Specification::NeuralNetworkLayer* ActivationLayerBuilder::Build() {
+  layer_->set_name(DebugName());
+  switch (activation_) {
+    // ActNone is used for sclalar multiplication (linear activation)
+    case kTfLiteActNone:
+      layer_->mutable_activation()->mutable_linear()->set_alpha(alpha_);
+      break;
+    case kTfLiteActRelu:
+      layer_->mutable_activation()->mutable_relu();
+      break;
+    // Relu1 and Relu6 layers are fully composed in PopulateSubgraph().
+    case kTfLiteActReluN1To1:  // clip(-1, 1)
+      layer_->mutable_unary()->set_alpha(-1);
+      layer_->mutable_unary()->set_type(
+          CoreML::Specification::UnaryFunctionLayerParams::THRESHOLD);
+      break;
+    case kTfLiteActRelu6:  // clip(0, 6)
+      layer_->mutable_activation()->mutable_relu();
+      break;
+    case kTfLiteActTanh:
+      layer_->mutable_activation()->mutable_tanh();
+      break;
+    case kTfLiteActSigmoid:
+      layer_->mutable_activation()->mutable_sigmoid();
+      break;
+    // TODO(taeheej): signbit is not implemented.
+    default:
+      fprintf(stderr, "Activation %d is not supported.\n", activation_);
+      break;
+  }
+  return layer_.release();
+}
+
+TfLiteStatus ActivationLayerBuilder::PopulateSubgraph(TfLiteContext* context) {
+  if (!(activation_ == kTfLiteActRelu6 || activation_ == kTfLiteActReluN1To1)) {
+    builder_output_ = AddOutput();
+    return kTfLiteOk;
+  }
+
+  // Relu1: Threshold(-1) -> Threshold(-1) with scale: -1 -> Negation
+  // Relu6: ReLU -> Threshold(-6) with scale: -1 -> Negation
+  const int relu_threshold = activation_ == kTfLiteActRelu6 ? 6 : 1;
+  ThresholdLayerBuilder* threshold_builder =
+      reinterpret_cast<ThresholdLayerBuilder*>(
+          graph_builder_->AddBuilder(CreateThresholdLayerBuilder, nullptr));
+
+  threshold_builder->SetAlpha(-relu_threshold);
+  threshold_builder->SetScale(-1);
+
+  threshold_builder->AddInput(AddOutput());
+
+  ActivationLayerBuilder* negation_builder =
+      reinterpret_cast<ActivationLayerBuilder*>(
+          graph_builder_->AddBuilder(CreateActivationLayerBuilder, nullptr));
+  negation_builder->SetActivation(kTfLiteActNone);
+  negation_builder->SetAlpha(-1);
+
+  negation_builder->AddInput(threshold_builder->AddOutput());
+  builder_output_ = negation_builder->AddOutput();
+  return kTfLiteOk;
+}
+
+TfLiteStatus ActivationLayerBuilder::RegisterInputs(
+    const TfLiteIntArray* inputs, TfLiteContext* context) {
+  if (inputs->size != 1) {
+    TF_LITE_KERNEL_LOG(context, "Activation: Wrong # of inputs!.");
+    return kTfLiteError;
+  }
+  AddInput(inputs->data[0]);
+  return kTfLiteOk;
+}
+
+TfLiteStatus ActivationLayerBuilder::RegisterOutputs(
+    const TfLiteIntArray* outputs, TfLiteContext* context) {
+  if (outputs->size != 1) {
+    TF_LITE_KERNEL_LOG(context, "Activation: Wrong # of outputs!.");
+    return kTfLiteError;
+  }
+  graph_builder_->AddTensorWithID(outputs->data[0], GetOutput(context));
+  return kTfLiteOk;
+}
+
+OpBuilder* CreateActivationLayerBuilder(GraphBuilder* graph_builder) {
+  return new ActivationLayerBuilder(graph_builder);
+}
+
+OpBuilder* CreateLogisticOpBuilder(GraphBuilder* graph_builder) {
+  return new ActivationLayerBuilder(graph_builder, kTfLiteActSigmoid);
+}
+
+OpBuilder* CreateReluOpBuilder(GraphBuilder* graph_builder) {
+  return new ActivationLayerBuilder(graph_builder, kTfLiteActRelu);
+}
+
+OpBuilder* CreateReluN1To1OpBuilder(GraphBuilder* graph_builder) {
+  return new ActivationLayerBuilder(graph_builder, kTfLiteActReluN1To1);
+}
+
+OpBuilder* CreateRelu6OpBuilder(GraphBuilder* graph_builder) {
+  return new ActivationLayerBuilder(graph_builder, kTfLiteActRelu6);
+}
+
+OpBuilder* CreateTanhOpBuilder(GraphBuilder* graph_builder) {
+  return new ActivationLayerBuilder(graph_builder, kTfLiteActTanh);
+}
+
+}  // namespace coreml
+}  // namespace delegates
+}  // namespace tflite

tensorflow/lite/delegates/coreml/builders/activation_layer_builder.h

@@ -0,0 +1,65 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_DELEGATES_COREML_BUILDERS_ACTIVATION_LAYER_BUILDER_H_
+#define TENSORFLOW_LITE_DELEGATES_COREML_BUILDERS_ACTIVATION_LAYER_BUILDER_H_
+
+#include <string>
+
+#include "mlmodel/format/NeuralNetwork.pb.h"
+#include "tensorflow/lite/c/c_api_types.h"
+#include "tensorflow/lite/c/common.h"
+#include "tensorflow/lite/core/c/builtin_op_data.h"
+#include "tensorflow/lite/delegates/coreml/builders/op_builder.h"
+
+namespace tflite {
+namespace delegates {
+namespace coreml {
+
+class ActivationLayerBuilder : public OpBuilder {
+ public:
+  explicit ActivationLayerBuilder(GraphBuilder* graph_builder)
+      : OpBuilder(graph_builder) {}
+
+  explicit ActivationLayerBuilder(GraphBuilder* graph_builder,
+                                  TfLiteFusedActivation activation)
+      : OpBuilder(graph_builder), activation_(activation) {}
+
+  const std::string& DebugName() override;
+
+  CoreML::Specification::NeuralNetworkLayer* Build() override;
+
+  void SetActivation(TfLiteFusedActivation activation) {
+    activation_ = activation;
+  }
+
+  void SetAlpha(float alpha) { alpha_ = alpha; }
+
+  TfLiteStatus PopulateSubgraph(TfLiteContext* context) override;
+
+  TfLiteStatus RegisterInputs(const TfLiteIntArray* inputs,
+                              TfLiteContext* context) override;
+
+  TfLiteStatus RegisterOutputs(const TfLiteIntArray* outputs,
+                               TfLiteContext* context) override;
+
+ private:
+  TfLiteFusedActivation activation_;
+  float alpha_ = 1.0f;
+};
+
+}  // namespace coreml
+}  // namespace delegates
+}  // namespace tflite
+#endif  // TENSORFLOW_LITE_DELEGATES_COREML_BUILDERS_ACTIVATION_LAYER_BUILDER_H_

tensorflow/lite/delegates/coreml/builders/add_op_builder.cc

@@ -0,0 +1,113 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include "tensorflow/lite/delegates/coreml/builders/add_op_builder.h"
+
+#include <memory>
+#include <string>
+
+#include "mlmodel/format/NeuralNetwork.pb.h"
+#include "tensorflow/lite/core/c/builtin_op_data.h"
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/delegates/coreml/builders/activation_layer_builder.h"
+#include "tensorflow/lite/delegates/coreml/builders/op_builder.h"
+#include "tensorflow/lite/delegates/coreml/builders/op_factory.h"
+#include "tensorflow/lite/delegates/coreml/builders/util.h"
+#include "tensorflow/lite/kernels/kernel_util.h"
+
+namespace tflite {
+namespace delegates {
+namespace coreml {
+const std::string& AddOpBuilder::DebugName() {
+  if (debug_name_.empty()) SetDebugName("AddOpBuilder", node_id_);
+  return debug_name_;
+}
+
+CoreML::Specification::NeuralNetworkLayer* AddOpBuilder::Build() {
+  if (layer_ == nullptr) {
+    layer_ = std::make_unique<CoreML::Specification::NeuralNetworkLayer>();
+  }
+  layer_->set_name(DebugName());
+  layer_->mutable_add();
+  if (alpha_ != 0.0f) {
+    layer_->mutable_add()->set_alpha(alpha_);
+  }
+
+  return layer_.release();
+}
+
+TfLiteStatus AddOpBuilder::PopulateSubgraph(TfLiteContext* context) {
+  TfLiteAddParams* params = reinterpret_cast<TfLiteAddParams*>(builtin_data_);
+
+  TfLiteFusedActivation activation = params->activation;
+  if (activation == kTfLiteActNone) {
+    builder_output_ = AddOutput();
+  } else {
+    ActivationLayerBuilder* activation_builder =
+        reinterpret_cast<ActivationLayerBuilder*>(
+            graph_builder_->AddBuilder(CreateActivationLayerBuilder, nullptr));
+    activation_builder->SetActivation(activation);
+    activation_builder->AddInput(AddOutput());
+    activation_builder->PopulateSubgraph(context);
+    builder_output_ = activation_builder->GetOutput(context);
+  }
+  return kTfLiteOk;
+}
+
+TfLiteStatus AddOpBuilder::RegisterInputs(const TfLiteIntArray* inputs,
+                                          TfLiteContext* context) {
+  // TODO(taeheej): support 1 input case if necessary. TFL add needs 2 inputs.
+  if (inputs->size != 2) {
+    TF_LITE_KERNEL_LOG(context, "Wrong # of inputs to add!.");
+    return kTfLiteError;
+  }
+  const auto* input_0 = &context->tensors[inputs->data[0]];
+  const auto* input_1 = &context->tensors[inputs->data[1]];
+  // store constant, scalar value into MultiplyLayerParams directly.
+  if (IsConstantTensor(input_0) && NumElements(input_0) == 1) {
+    AddInput(inputs->data[1]);
+    SetAlpha(GetScalarFloatFromTensor(input_0));
+  } else if (IsConstantTensor(input_1) && NumElements(input_1) == 1) {
+    AddInput(inputs->data[0]);
+    SetAlpha(GetScalarFloatFromTensor(input_1));
+  } else {
+    AddInput(inputs->data[0]);
+    AddInput(inputs->data[1]);
+  }
+  return kTfLiteOk;
+}
+
+TfLiteStatus AddOpBuilder::RegisterOutputs(const TfLiteIntArray* outputs,
+                                           TfLiteContext* context) {
+  if (outputs->size != 1) {
+    TF_LITE_KERNEL_LOG(context, "Wrong # of outputs to add!.");
+    return kTfLiteError;
+  }
+  TensorID output_tensor = GetOutput(context);
+  if (output_tensor.NodeID() == -1) {
+    TF_LITE_KERNEL_LOG(context, "Failed to build output tensor.");
+    return kTfLiteError;
+  }
+  graph_builder_->AddTensorWithID(outputs->data[0], output_tensor);
+  return kTfLiteOk;
+}
+
+void AddOpBuilder::SetAlpha(float alpha) { alpha_ = alpha; }
+
+OpBuilder* CreateAddOpBuilder(GraphBuilder* graph_builder) {
+  return new AddOpBuilder(graph_builder);
+}
+}  // namespace coreml
+}  // namespace delegates
+}  // namespace tflite

tensorflow/lite/delegates/coreml/builders/add_op_builder.h

@@ -0,0 +1,56 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_DELEGATES_COREML_BUILDERS_ADD_OP_BUILDER_H_
+#define TENSORFLOW_LITE_DELEGATES_COREML_BUILDERS_ADD_OP_BUILDER_H_
+
+#include <string>
+
+#include "mlmodel/format/NeuralNetwork.pb.h"
+#include "tensorflow/lite/c/c_api_types.h"
+#include "tensorflow/lite/c/common.h"
+#include "tensorflow/lite/delegates/coreml/builders/op_builder.h"
+
+namespace tflite {
+namespace delegates {
+namespace coreml {
+// Builder for Add op in CoreML.
+class AddOpBuilder : public OpBuilder {
+ public:
+  explicit AddOpBuilder(GraphBuilder* graph_builder)
+      : OpBuilder(graph_builder) {}
+  const std::string& DebugName() override;
+
+  CoreML::Specification::NeuralNetworkLayer* Build() override;
+
+  TfLiteStatus PopulateSubgraph(TfLiteContext* context) override;
+
+  TfLiteStatus RegisterInputs(const TfLiteIntArray* inputs,
+                              TfLiteContext* context) override;
+
+  TfLiteStatus RegisterOutputs(const TfLiteIntArray* outputs,
+                               TfLiteContext* context) override;
+
+  void SetAlpha(float alpha);
+
+ private:
+  // Used for unary add
+  float alpha_ = 0.0f;
+};
+
+}  // namespace coreml
+}  // namespace delegates
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_DELEGATES_COREML_BUILDERS_ADD_OP_BUILDER_H_

tensorflow/lite/delegates/coreml/builders/concatenation_op_builder.cc

@@ -0,0 +1,87 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include "tensorflow/lite/delegates/coreml/builders/concatenation_op_builder.h"
+
+#include <memory>
+
+#include "mlmodel/format/NeuralNetwork.pb.h"
+#include "tensorflow/lite/core/c/builtin_op_data.h"
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/delegates/coreml/builders/op_builder.h"
+#include "tensorflow/lite/delegates/coreml/builders/op_validator.h"
+
+namespace tflite {
+namespace delegates {
+namespace coreml {
+
+CoreML::Specification::NeuralNetworkLayer* ConcatenationOpBuilder::Build() {
+  if (layer_ == nullptr) {
+    layer_ = std::make_unique<CoreML::Specification::NeuralNetworkLayer>();
+  }
+  layer_->set_name(DebugName());
+  layer_->mutable_concat()->set_sequenceconcat(false);
+  return layer_.release();
+}
+
+TfLiteStatus ConcatenationOpBuilder::RegisterInputs(
+    const TfLiteIntArray* inputs, TfLiteContext* context) {
+  if (inputs->size < 2) {
+    TF_LITE_KERNEL_LOG(
+        context, "ConcatenationOpBuilder: at least 2 inputs are required.");
+    return kTfLiteError;
+  }
+  for (int i = 0; i < inputs->size; ++i) {
+    AddInput(inputs->data[i]);
+  }
+  return kTfLiteOk;
+}
+
+TfLiteStatus ConcatenationOpBuilder::RegisterOutputs(
+    const TfLiteIntArray* outputs, TfLiteContext* context) {
+  if (outputs->size != 1) {
+    TF_LITE_KERNEL_LOG(context, "Wrong # of outputs to Concat!.");
+    return kTfLiteError;
+  }
+  graph_builder_->AddTensorWithID(outputs->data[0], GetOutput(context));
+  return kTfLiteOk;
+}
+
+OpBuilder* CreateConcatenationOpBuilder(GraphBuilder* graph_builder) {
+  return new ConcatenationOpBuilder(graph_builder);
+}
+
+bool IsConcatenationOpSupported(const TfLiteRegistration* registration,
+                                const TfLiteNode* node,
+                                TfLiteContext* context) {
+  if (node->builtin_data == nullptr) return false;
+  auto params =
+      reinterpret_cast<TfLiteConcatenationParams*>(node->builtin_data);
+  int input_dims = context->tensors[node->inputs->data[0]].dims->size;
+
+  // Not supported in TfLite kernel.
+  if (params->activation != kTfLiteActNone) return false;
+  if (node->inputs->size < 2) return false;
+
+  // Only supports concatenation by channel. Core ML concatenation supports
+  // concatenation by channel and by sequence (axis -5) only.
+  // TODO(b/145642128): support stack layer here with Core ML 3 support.
+  if (input_dims == 3) return params->axis == 2 || params->axis == -1;
+  if (input_dims == 4) return params->axis == 3 || params->axis == -1;
+  return false;
+}
+
+}  // namespace coreml
+}  // namespace delegates
+}  // namespace tflite

tensorflow/lite/delegates/coreml/builders/concatenation_op_builder.h

@@ -0,0 +1,52 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_DELEGATES_COREML_BUILDERS_CONCATENATION_OP_BUILDER_H_
+#define TENSORFLOW_LITE_DELEGATES_COREML_BUILDERS_CONCATENATION_OP_BUILDER_H_
+
+#include <string>
+
+#include "mlmodel/format/NeuralNetwork.pb.h"
+#include "tensorflow/lite/c/c_api_types.h"
+#include "tensorflow/lite/c/common.h"
+#include "tensorflow/lite/delegates/coreml/builders/op_builder.h"
+
+namespace tflite {
+namespace delegates {
+namespace coreml {
+
+class ConcatenationOpBuilder : public OpBuilder {
+ public:
+  explicit ConcatenationOpBuilder(GraphBuilder* graph_builder)
+      : OpBuilder(graph_builder) {}
+
+  const std::string& DebugName() override {
+    if (debug_name_.empty()) SetDebugName("ConcatOpBuilder", node_id_);
+    return debug_name_;
+  }
+
+  CoreML::Specification::NeuralNetworkLayer* Build() override;
+
+  TfLiteStatus RegisterInputs(const TfLiteIntArray* inputs,
+                              TfLiteContext* context) override;
+
+  TfLiteStatus RegisterOutputs(const TfLiteIntArray* outputs,
+                               TfLiteContext* context) override;
+};
+
+}  // namespace coreml
+}  // namespace delegates
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_DELEGATES_COREML_BUILDERS_CONCATENATION_OP_BUILDER_H_

tensorflow/lite/delegates/coreml/builders/convolution_op_builder.cc

@@ -0,0 +1,381 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include "tensorflow/lite/delegates/coreml/builders/convolution_op_builder.h"
+
+#include <algorithm>
+#include <cstdint>
+#include <cstdio>
+#include <memory>
+#include <string>
+
+#include "mlmodel/format/NeuralNetwork.pb.h"
+#include "google/protobuf/repeated_field.h"
+#include "tensorflow/lite/builtin_ops.h"
+#include "tensorflow/lite/core/c/builtin_op_data.h"
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/delegates/coreml/builders/activation_layer_builder.h"
+#include "tensorflow/lite/delegates/coreml/builders/op_builder.h"
+#include "tensorflow/lite/delegates/coreml/builders/op_factory.h"
+#include "tensorflow/lite/delegates/coreml/builders/op_validator.h"
+#include "tensorflow/lite/kernels/internal/optimized/optimized_ops.h"
+#include "tensorflow/lite/kernels/internal/runtime_shape.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
+#include "tensorflow/lite/kernels/internal/types.h"
+#include "tensorflow/lite/kernels/kernel_util.h"
+
+namespace tflite {
+namespace delegates {
+namespace coreml {
+const std::string& ConvolutionOpBuilder::DebugName() {
+  if (debug_name_.empty()) SetDebugName("ConvolutionOpBuilder", node_id_);
+  return debug_name_;
+}
+
+void ConvolutionOpBuilder::SetWeights(TfLiteTensor* weights) {
+  weights_ = weights;
+}
+
+void ConvolutionOpBuilder::SetBias(TfLiteTensor* bias) { bias_ = bias; }
+
+void ConvolutionOpBuilder::SetOutputShape(TfLiteTensor* output_shape) {
+  output_shape_ = output_shape;
+}
+
+CoreML::Specification::NeuralNetworkLayer* ConvolutionOpBuilder::Build() {
+  if (layer_ == nullptr) {
+    layer_ = std::make_unique<CoreML::Specification::NeuralNetworkLayer>();
+  }
+  layer_->set_name(DebugName());
+
+  int stride_height = 0;
+  int stride_width = 0;
+  int dilation_height = 0;
+  int dilation_width = 0;
+  TfLitePadding padding;
+
+  switch (conv_type_) {
+    case ConvolutionType::kConv: {
+      const auto* conv_params =
+          reinterpret_cast<const TfLiteConvParams*>(builtin_data_);
+      stride_height = conv_params->stride_height;
+      stride_width = conv_params->stride_width;
+      dilation_height = conv_params->dilation_height_factor;
+      dilation_width = conv_params->dilation_width_factor;
+      padding = conv_params->padding;
+
+      layer_->mutable_convolution()->set_ngroups(1);
+      break;
+    }
+    case ConvolutionType::kDepthwiseConv: {
+      const auto* depthwise_conv_params =
+          reinterpret_cast<const TfLiteDepthwiseConvParams*>(builtin_data_);
+      stride_height = depthwise_conv_params->stride_height;
+      stride_width = depthwise_conv_params->stride_width;
+      dilation_height = depthwise_conv_params->dilation_height_factor;
+      dilation_width = depthwise_conv_params->dilation_width_factor;
+      padding = depthwise_conv_params->padding;
+
+      // n_groups = kernel_channel / depth_multiplier
+      layer_->mutable_convolution()->set_ngroups(
+          weights_->dims->data[3] / depthwise_conv_params->depth_multiplier);
+      break;
+    }
+    case ConvolutionType::kTransposeConv: {
+      const auto* transpose_conv_params =
+          reinterpret_cast<const TfLiteTransposeConvParams*>(builtin_data_);
+      const int height_index = 1;
+      const int width_index = 2;
+
+      stride_height = transpose_conv_params->stride_height;
+      stride_width = transpose_conv_params->stride_width;
+      padding = transpose_conv_params->padding;
+
+      layer_->mutable_convolution()->mutable_outputshape()->Add(
+          GetTensorData<int32_t>(output_shape_)[height_index]);
+      layer_->mutable_convolution()->mutable_outputshape()->Add(
+          GetTensorData<int32_t>(output_shape_)[width_index]);
+      break;
+    }
+  }
+
+  // If not set, it will default to (1,1)
+  if (stride_height) {
+    layer_->mutable_convolution()->add_stride(stride_height);
+    layer_->mutable_convolution()->add_stride(stride_width);
+  }
+
+  if (dilation_height) {
+    layer_->mutable_convolution()->add_dilationfactor(dilation_height);
+    layer_->mutable_convolution()->add_dilationfactor(dilation_width);
+  }
+
+  switch (padding) {
+    case kTfLitePaddingSame:
+      layer_->mutable_convolution()->mutable_same();
+      break;
+    case kTfLitePaddingValid:
+      layer_->mutable_convolution()->mutable_valid();
+      break;
+    case kTfLitePaddingUnknown:
+      fprintf(stderr, "Padding is unknown.\n");
+      break;
+  }
+
+  FillCoreMLWeights();
+  FillCoreMLBias();
+
+  return layer_.release();
+}
+
+void ConvolutionOpBuilder::FillCoreMLWeights() {
+  if (conv_type_ == ConvolutionType::kDepthwiseConv) {
+    layer_->mutable_convolution()->set_kernelchannels(1);
+    layer_->mutable_convolution()->set_outputchannels(weights_->dims->data[3]);
+  } else {
+    layer_->mutable_convolution()->set_kernelchannels(weights_->dims->data[3]);
+    layer_->mutable_convolution()->set_outputchannels(weights_->dims->data[0]);
+  }
+  layer_->mutable_convolution()->add_kernelsize(weights_->dims->data[1]);
+  layer_->mutable_convolution()->add_kernelsize(weights_->dims->data[2]);
+
+  TransposeKernelWeights();  // Should be called after CoreML shape is set.
+}
+
+void ConvolutionOpBuilder::TransposeKernelWeights() {
+  RuntimeShape tfl_shape(4, weights_->dims->data);
+  // CoreML kernel has shape of (C_out, C_in, H, W)
+  RuntimeShape coreml_shape(
+      {static_cast<int>(layer_->convolution().outputchannels()),
+       static_cast<int>(layer_->convolution().kernelchannels()),
+       static_cast<int>(layer_->convolution().kernelsize()[0]),
+       static_cast<int>(layer_->convolution().kernelsize()[1])});
+
+  TransposeParams params;
+
+  if (conv_type_ == ConvolutionType::kDepthwiseConv) {
+    // DepthwiseConv2D: TFL kernel has shape of (1, H, W, C_out),
+    // and CoreML kernel has shape of (C_out, 1, H, W)
+    params = {/*perm_count=*/4, /*perm=*/{3, 0, 1, 2}};
+  } else {
+    // Conv2D and TransposeConv: TFL kernel has shape of (C_out, H, W, C_in),
+    // and CoreML kernel has shape of (C_out, C_in, H, W)
+    params = {/*perm_count=*/4, /*perm=*/{0, 3, 1, 2}};
+  }
+
+  if (conv_type_ == ConvolutionType::kTransposeConv) {
+    layer_->mutable_convolution()->set_isdeconvolution(true);
+  }
+
+  if (weights_->type == kTfLiteFloat32) {
+    auto* coreml_weights =
+        layer_->mutable_convolution()->mutable_weights()->mutable_floatvalue();
+    coreml_weights->resize(NumElements(weights_), 0);
+
+    optimized_ops::Transpose<float>(params, tfl_shape, weights_->data.f,
+                                    coreml_shape,
+                                    coreml_weights->mutable_data());
+  } else if (weights_->type == kTfLiteFloat16) {
+    auto* coreml_weights = layer_->mutable_convolution()
+                               ->mutable_weights()
+                               ->mutable_float16value();
+    // float16value has type of bytes (std::string)
+    coreml_weights->resize(weights_->bytes, 0);
+
+    optimized_ops::Transpose<uint16_t>(
+        params, tfl_shape, reinterpret_cast<uint16_t*>(weights_->data.raw),
+        coreml_shape, reinterpret_cast<uint16_t*>(&coreml_weights->front()));
+  }
+}
+
+void ConvolutionOpBuilder::FillCoreMLBias() {
+  if (bias_ != nullptr) {
+    layer_->mutable_convolution()->set_hasbias(true);
+    if (bias_->type == kTfLiteFloat32) {
+      std::copy(bias_->data.f, bias_->data.f + NumElements(bias_->dims),
+                google::protobuf::RepeatedFieldBackInserter(layer_->mutable_convolution()
+                                                      ->mutable_bias()
+                                                      ->mutable_floatvalue()));
+    } else if (bias_->type == kTfLiteFloat16) {
+      // float16value has type of bytes (std::string)
+      layer_->mutable_convolution()
+          ->mutable_bias()
+          ->mutable_float16value()
+          ->assign(bias_->data.raw, bias_->bytes);
+    }
+  }
+}
+
+TfLiteStatus ConvolutionOpBuilder::PopulateSubgraph(TfLiteContext* context) {
+  TfLiteFusedActivation activation;
+  switch (conv_type_) {
+    case ConvolutionType::kConv: {
+      const auto* conv_params =
+          reinterpret_cast<const TfLiteConvParams*>(builtin_data_);
+      activation = conv_params->activation;
+      break;
+    }
+    case ConvolutionType::kDepthwiseConv: {
+      const auto* depthwise_conv_params =
+          reinterpret_cast<const TfLiteDepthwiseConvParams*>(builtin_data_);
+      activation = depthwise_conv_params->activation;
+      break;
+    }
+    case ConvolutionType::kTransposeConv: {
+      activation = kTfLiteActNone;
+      break;
+    }
+  }
+
+  if (activation == kTfLiteActNone) {
+    builder_output_ = AddOutput();
+  } else {
+    ActivationLayerBuilder* activation_builder =
+        reinterpret_cast<ActivationLayerBuilder*>(
+            graph_builder_->AddBuilder(CreateActivationLayerBuilder, nullptr));
+    activation_builder->SetActivation(activation);
+    activation_builder->AddInput(AddOutput());
+    activation_builder->PopulateSubgraph(context);
+    builder_output_ = activation_builder->GetOutput(context);
+  }
+  return kTfLiteOk;
+}
+
+TfLiteStatus ConvolutionOpBuilder::RegisterInputs(const TfLiteIntArray* inputs,
+                                                  TfLiteContext* context) {
+  if (conv_type_ == ConvolutionType::kTransposeConv) {
+    if (inputs->size != 3) {
+      TF_LITE_KERNEL_LOG(context,
+                         "Transpose Conv should have 3 inputs, %d given.",
+                         inputs->size);
+      return kTfLiteError;
+    }
+    AddInput(inputs->data[2]);
+    SetOutputShape(&context->tensors[inputs->data[0]]);
+  } else {
+    if (inputs->size != 2 && inputs->size != 3) {
+      TF_LITE_KERNEL_LOG(context,
+                         "Convolution and depthwise convolution should have 2 "
+                         "or 3 inputs, %d given.",
+                         inputs->size);
+      return kTfLiteError;
+    }
+    AddInput(inputs->data[0]);
+    if (inputs->size > 2) {
+      SetBias(&context->tensors[inputs->data[2]]);
+    }
+  }
+  SetWeights(&context->tensors[inputs->data[1]]);
+  return kTfLiteOk;
+}
+
+TfLiteStatus ConvolutionOpBuilder::RegisterOutputs(
+    const TfLiteIntArray* outputs, TfLiteContext* context) {
+  if (outputs->size != 1) {
+    TF_LITE_KERNEL_LOG(context, "Wrong # of outputs!.");
+    return kTfLiteError;
+  }
+  TensorID output_tensor = GetOutput(context);
+  if (output_tensor.NodeID() == -1) {
+    TF_LITE_KERNEL_LOG(context, "Failed to build output tensor.");
+    return kTfLiteError;
+  }
+  graph_builder_->AddTensorWithID(outputs->data[0], output_tensor);
+  return kTfLiteOk;
+}
+
+OpBuilder* CreateConvolutionOpBuilder(GraphBuilder* graph_builder) {
+  return new ConvolutionOpBuilder(graph_builder, ConvolutionType::kConv);
+}
+
+OpBuilder* CreateDepthwiseConvolutionOpBuilder(GraphBuilder* graph_builder) {
+  return new ConvolutionOpBuilder(graph_builder,
+                                  ConvolutionType::kDepthwiseConv);
+}
+
+OpBuilder* CreateTransposeConvolutionOpBuilder(GraphBuilder* graph_builder) {
+  return new ConvolutionOpBuilder(graph_builder,
+                                  ConvolutionType::kTransposeConv);
+}
+
+bool IsConvolutionOpSupported(const TfLiteRegistration* registration,
+                              const TfLiteNode* node, TfLiteContext* context) {
+  if (node->builtin_data == nullptr) return false;
+
+  TfLiteFusedActivation activation;
+
+  if (registration->builtin_code == kTfLiteBuiltinConv2d) {
+    const auto* conv_params =
+        reinterpret_cast<const TfLiteConvParams*>(node->builtin_data);
+    activation = conv_params->activation;
+  } else if (registration->builtin_code == kTfLiteBuiltinDepthwiseConv2d) {
+    const auto* depthwise_conv_params =
+        reinterpret_cast<const TfLiteDepthwiseConvParams*>(node->builtin_data);
+    activation = depthwise_conv_params->activation;
+  } else if (registration->builtin_code == kTfLiteBuiltinTransposeConv) {
+    activation = kTfLiteActNone;
+  } else {
+    TF_LITE_KERNEL_LOG(
+        context,
+        "Invalid op: op must be Conv2D, DepthwiseConv2D or TransposeConv.");
+    return false;
+  }
+
+  if (activation == kTfLiteActSignBit) {
+    return false;
+  }
+
+  const int kOutputShapeTensor = 0;  // Only used for TransposeConv
+  const int kWeightTensor = 1;
+  const int kBiasTensor = 2;  // Only used for non-TransposeConv
+  const TfLiteTensor* weights;
+  TF_LITE_ENSURE_OK(context,
+                    GetInputSafe(context, node, kWeightTensor, &weights));
+  const int max_kernel_size = 16384;
+  if (!IsConstantTensor(weights)) {
+    return false;
+  }
+  if (weights->dims->data[1] > max_kernel_size ||
+      weights->dims->data[2] > max_kernel_size) {
+    return false;
+  }
+  if (registration->builtin_code == kTfLiteBuiltinTransposeConv) {
+    if (!IsConstantTensor(GetInput(context, node, kOutputShapeTensor))) {
+      return false;
+    }
+  } else {
+    if (node->inputs->size >= kBiasTensor &&
+        !IsConstantTensor(GetInput(context, node, kBiasTensor))) {
+      return false;
+    }
+  }
+
+  return true;
+}
+
+bool IsDepthwiseConvolutionOpSupported(const TfLiteRegistration* registration,
+                                       const TfLiteNode* node,
+                                       TfLiteContext* context) {
+  return IsConvolutionOpSupported(registration, node, context);
+}
+
+bool IsTransposeConvolutionOpSupported(const TfLiteRegistration* registration,
+                                       const TfLiteNode* node,
+                                       TfLiteContext* context) {
+  return IsConvolutionOpSupported(registration, node, context);
+}
+
+}  // namespace coreml
+}  // namespace delegates
+}  // namespace tflite

tensorflow/lite/delegates/coreml/builders/convolution_op_builder.h

@@ -0,0 +1,88 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_DELEGATES_COREML_BUILDERS_CONVOLUTION_OP_BUILDER_H_
+#define TENSORFLOW_LITE_DELEGATES_COREML_BUILDERS_CONVOLUTION_OP_BUILDER_H_
+
+#include <string>
+
+#include "mlmodel/format/NeuralNetwork.pb.h"
+#include "tensorflow/lite/builtin_ops.h"
+#include "tensorflow/lite/core/c/builtin_op_data.h"
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/delegates/coreml/builders/op_builder.h"
+
+namespace tflite {
+namespace delegates {
+namespace coreml {
+
+enum class ConvolutionType { kConv, kDepthwiseConv, kTransposeConv };
+
+// Layer that provides convolution and depthwise convolution.
+class ConvolutionOpBuilder : public OpBuilder {
+ public:
+  explicit ConvolutionOpBuilder(GraphBuilder* graph_builder,
+                                ConvolutionType conv_type)
+      : OpBuilder(graph_builder), conv_type_(conv_type) {}
+
+  const std::string& DebugName() override;
+
+  CoreML::Specification::NeuralNetworkLayer* Build() override;
+
+  TfLiteStatus PopulateSubgraph(TfLiteContext* context) override;
+
+  void SetOutputChannels(uint64_t output_channels);
+
+  void SetNGroups(uint64_t n_groups);
+
+  void SetWeights(TfLiteTensor* weights);
+
+  void SetBias(TfLiteTensor* bias);
+
+  void SetOutputShape(TfLiteTensor* output_shape);
+
+  void SetParams(void* builtin_data);
+
+  TfLiteStatus RegisterInputs(const TfLiteIntArray* inputs,
+                              TfLiteContext* context) override;
+
+  TfLiteStatus RegisterOutputs(const TfLiteIntArray* outputs,
+                               TfLiteContext* context) override;
+
+ private:
+  void FillCoreMLWeights();
+  void FillCoreMLBias();
+
+  // Transpose TFLite kernel weights to CoreML kernel weights.
+  // Should be called after setting CoreML's kernel shapes.
+  void TransposeKernelWeights();
+
+  uint64_t output_channels_;
+  uint64_t n_groups_ = 1;
+
+  ConvolutionType conv_type_;
+
+  // using default dilation_factor (1, 1)
+  // CoreML ConvolutionLayerParams.isDeconvolution == false
+  TfLiteTensor* weights_ = nullptr;
+  TfLiteTensor* bias_ = nullptr;
+  // Only used for TransposeConv.
+  TfLiteTensor* output_shape_ = nullptr;
+};
+
+}  // namespace coreml
+}  // namespace delegates
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_DELEGATES_COREML_BUILDERS_CONVOLUTION_OP_BUILDER_H_

tensorflow/lite/delegates/coreml/builders/dummy_op_builder.cc

@@ -0,0 +1,57 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include "tensorflow/lite/delegates/coreml/builders/dummy_op_builder.h"
+
+#include <string>
+
+#include "mlmodel/format/NeuralNetwork.pb.h"
+#include "tensorflow/lite/c/c_api_types.h"
+#include "tensorflow/lite/c/common.h"
+#include "tensorflow/lite/delegates/coreml/builders/op_factory.h"
+
+namespace tflite {
+namespace delegates {
+namespace coreml {
+
+CoreML::Specification::NeuralNetworkLayer* DummyOpBuilder::Build() {
+  return nullptr;
+}
+
+const std::string& DummyOpBuilder::DebugName() {
+  SetDebugName("DummyOpBuilder", node_id_);
+  return debug_name_;
+}
+
+TfLiteStatus DummyOpBuilder::PopulateSubgraph(TfLiteContext* context) {
+  return kTfLiteOk;
+}
+
+OpBuilder* CreateDummyOpBuilder(GraphBuilder* graph_builder) {
+  return new DummyOpBuilder(graph_builder);
+}
+
+TfLiteStatus DummyOpBuilder::RegisterInputs(const TfLiteIntArray* inputs,
+                                            TfLiteContext* context) {
+  return kTfLiteOk;
+}
+
+TfLiteStatus DummyOpBuilder::RegisterOutputs(const TfLiteIntArray* outputs,
+                                             TfLiteContext* context) {
+  return kTfLiteOk;
+}
+
+}  // namespace coreml
+}  // namespace delegates
+}  // namespace tflite

tensorflow/lite/delegates/coreml/builders/dummy_op_builder.h

@@ -0,0 +1,50 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_DELEGATES_COREML_BUILDERS_DUMMY_OP_BUILDER_H_
+#define TENSORFLOW_LITE_DELEGATES_COREML_BUILDERS_DUMMY_OP_BUILDER_H_
+
+#include <string>
+
+#include "mlmodel/format/NeuralNetwork.pb.h"
+#include "tensorflow/lite/builtin_ops.h"
+#include "tensorflow/lite/core/c/builtin_op_data.h"
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/delegates/coreml/builders/op_builder.h"
+
+namespace tflite {
+namespace delegates {
+namespace coreml {
+
+// Dummy Opbuilder for nodes that are claimed but not used. ex) FP16 dequantize
+class DummyOpBuilder : public OpBuilder {
+ public:
+  explicit DummyOpBuilder(GraphBuilder* graph_builder)
+      : OpBuilder(graph_builder) {}
+  CoreML::Specification::NeuralNetworkLayer* Build() override;
+  TfLiteStatus PopulateSubgraph(TfLiteContext* context) override;
+  const std::string& DebugName() override;
+
+  TfLiteStatus RegisterInputs(const TfLiteIntArray* inputs,
+                              TfLiteContext* context) override;
+
+  TfLiteStatus RegisterOutputs(const TfLiteIntArray* outputs,
+                               TfLiteContext* context) override;
+};
+
+}  // namespace coreml
+}  // namespace delegates
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_DELEGATES_COREML_BUILDERS_DUMMY_OP_BUILDER_H_

tensorflow/lite/delegates/coreml/builders/fully_connected_op_builder.cc

@@ -0,0 +1,196 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include "tensorflow/lite/delegates/coreml/builders/fully_connected_op_builder.h"
+
+#include <algorithm>
+#include <memory>
+#include <string>
+
+#include "mlmodel/format/NeuralNetwork.pb.h"
+#include "google/protobuf/repeated_field.h"
+#include "tensorflow/lite/core/c/builtin_op_data.h"
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/delegates/coreml/builders/activation_layer_builder.h"
+#include "tensorflow/lite/delegates/coreml/builders/op_builder.h"
+#include "tensorflow/lite/delegates/coreml/builders/op_factory.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
+#include "tensorflow/lite/kernels/kernel_util.h"
+
+namespace tflite {
+namespace delegates {
+namespace coreml {
+const std::string& FullyConnectedOpBuilder::DebugName() {
+  if (debug_name_.empty()) SetDebugName("FullyConnectedOpBuilder", node_id_);
+  return debug_name_;
+}
+
+void FullyConnectedOpBuilder::SetWeights(TfLiteTensor* weights) {
+  weights_ = weights;
+}
+
+void FullyConnectedOpBuilder::SetBias(TfLiteTensor* bias) { bias_ = bias; }
+
+CoreML::Specification::NeuralNetworkLayer* FullyConnectedOpBuilder::Build() {
+  if (layer_ == nullptr) {
+    layer_ = std::make_unique<CoreML::Specification::NeuralNetworkLayer>();
+  }
+  layer_->set_name(DebugName());
+
+  FillCoreMLWeights();
+  FillCoreMLBias();
+
+  return layer_.release();
+}
+
+void FullyConnectedOpBuilder::FillCoreMLWeights() {
+  layer_->mutable_innerproduct()->set_inputchannels(weights_->dims->data[1]);
+  layer_->mutable_innerproduct()->set_outputchannels(weights_->dims->data[0]);
+  if (weights_->type == kTfLiteFloat32) {
+    const float* weights_data = GetTensorData<float>(weights_);
+    std::copy(weights_data, weights_data + NumElements(weights_),
+              google::protobuf::RepeatedFieldBackInserter(layer_->mutable_innerproduct()
+                                                    ->mutable_weights()
+                                                    ->mutable_floatvalue()));
+  } else if (weights_->type == kTfLiteFloat16) {
+    // float16value has type of bytes (std::string)
+    layer_->mutable_innerproduct()
+        ->mutable_weights()
+        ->mutable_float16value()
+        ->assign(weights_->data.raw, weights_->bytes);
+  }
+}
+
+void FullyConnectedOpBuilder::FillCoreMLBias() {
+  if (bias_ != nullptr) {
+    layer_->mutable_innerproduct()->set_hasbias(true);
+    if (bias_->type == kTfLiteFloat32) {
+      const float* bias_data = GetTensorData<float>(bias_);
+      std::copy(bias_data, bias_data + NumElements(bias_),
+                google::protobuf::RepeatedFieldBackInserter(layer_->mutable_innerproduct()
+                                                      ->mutable_bias()
+                                                      ->mutable_floatvalue()));
+    } else if (bias_->type == kTfLiteFloat16) {
+      // float16value has type of bytes (std::string)
+      layer_->mutable_innerproduct()
+          ->mutable_bias()
+          ->mutable_float16value()
+          ->assign(bias_->data.raw, bias_->bytes);
+    }
+  }
+}
+
+TfLiteStatus FullyConnectedOpBuilder::PopulateSubgraph(TfLiteContext* context) {
+  const auto* fc_params =
+      reinterpret_cast<const TfLiteFullyConnectedParams*>(builtin_data_);
+  TfLiteFusedActivation activation = fc_params->activation;
+
+  if (activation == kTfLiteActNone) {
+    builder_output_ = AddOutput();
+  } else {
+    ActivationLayerBuilder* activation_builder =
+        reinterpret_cast<ActivationLayerBuilder*>(
+            graph_builder_->AddBuilder(CreateActivationLayerBuilder, nullptr));
+    activation_builder->SetActivation(activation);
+    activation_builder->AddInput(AddOutput());
+    activation_builder->PopulateSubgraph(context);
+    builder_output_ = activation_builder->GetOutput(context);
+  }
+  return kTfLiteOk;
+}
+
+TfLiteStatus FullyConnectedOpBuilder::RegisterInputs(
+    const TfLiteIntArray* inputs, TfLiteContext* context) {
+  const int kInput = 0;
+  const int kWeights = 1;
+  const int kBias = 2;
+  AddInput(inputs->data[kInput]);
+  SetWeights(&context->tensors[inputs->data[kWeights]]);
+  if (inputs->size > 2) {
+    SetBias(&context->tensors[inputs->data[kBias]]);
+  }
+  return kTfLiteOk;
+}
+
+TfLiteStatus FullyConnectedOpBuilder::RegisterOutputs(
+    const TfLiteIntArray* outputs, TfLiteContext* context) {
+  if (outputs->size != 1) {
+    TF_LITE_KERNEL_LOG(context, "Wrong # of outputs!.");
+    return kTfLiteError;
+  }
+  TensorID output_tensor = GetOutput(context);
+  if (output_tensor.NodeID() == -1) {
+    TF_LITE_KERNEL_LOG(context, "Failed to build output tensor.");
+    return kTfLiteError;
+  }
+  graph_builder_->AddTensorWithID(outputs->data[0], output_tensor);
+  return kTfLiteOk;
+}
+
+OpBuilder* CreateFullyConnectedOpBuilder(GraphBuilder* graph_builder) {
+  return new FullyConnectedOpBuilder(graph_builder);
+}
+
+bool IsFloatType(TfLiteType type) {
+  return type == kTfLiteFloat32 || type == kTfLiteFloat16;
+}
+
+bool IsFullyConnectedOpSupported(const TfLiteRegistration* registration,
+                                 const TfLiteNode* node,
+                                 TfLiteContext* context) {
+  if (node->builtin_data == nullptr) return false;
+  const auto* fc_params =
+      reinterpret_cast<const TfLiteFullyConnectedParams*>(node->builtin_data);
+  const int kInput = 0;
+  const int kWeights = 1;
+  const int kBias = 2;
+
+  if (fc_params->weights_format != kTfLiteFullyConnectedWeightsFormatDefault) {
+    return false;
+  }
+  const TfLiteTensor* input;
+  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kInput, &input));
+  const TfLiteTensor* weights;
+  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kWeights, &weights));
+
+  if (!IsFloatType(input->type)) {
+    return false;
+  }
+  if (!IsFloatType(weights->type) || !IsConstantTensor(weights)) {
+    return false;
+  }
+  // Core ML 2 only supports single-batch fully connected layer, thus dimensions
+  // except the last one should be 1.
+  if (input->dims->data[input->dims->size - 1] != NumElements(input)) {
+    return false;
+  }
+
+  if (node->inputs->size > 2) {
+    const TfLiteTensor* bias;
+    TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kBias, &bias));
+    if (!IsFloatType(bias->type) || !IsConstantTensor(bias)) {
+      return false;
+    }
+  }
+
+  TfLiteFusedActivation activation = fc_params->activation;
+  if (activation == kTfLiteActSignBit) {
+    return false;
+  }
+  return true;
+}
+
+}  // namespace coreml
+}  // namespace delegates
+}  // namespace tflite

tensorflow/lite/delegates/coreml/builders/fully_connected_op_builder.h

@@ -0,0 +1,61 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_DELEGATES_COREML_BUILDERS_FULLY_CONNECTED_OP_BUILDER_H_
+#define TENSORFLOW_LITE_DELEGATES_COREML_BUILDERS_FULLY_CONNECTED_OP_BUILDER_H_
+
+#include <string>
+
+#include "mlmodel/format/NeuralNetwork.pb.h"
+#include "tensorflow/lite/c/c_api_types.h"
+#include "tensorflow/lite/c/common.h"
+#include "tensorflow/lite/delegates/coreml/builders/op_builder.h"
+
+namespace tflite {
+namespace delegates {
+namespace coreml {
+// Builder for InnerProductLayer in Core ML.
+class FullyConnectedOpBuilder : public OpBuilder {
+ public:
+  explicit FullyConnectedOpBuilder(GraphBuilder* graph_builder)
+      : OpBuilder(graph_builder) {}
+  const std::string& DebugName() override;
+
+  CoreML::Specification::NeuralNetworkLayer* Build() override;
+
+  TfLiteStatus PopulateSubgraph(TfLiteContext* context) override;
+
+  TfLiteStatus RegisterInputs(const TfLiteIntArray* inputs,
+                              TfLiteContext* context) override;
+
+  TfLiteStatus RegisterOutputs(const TfLiteIntArray* outputs,
+                               TfLiteContext* context) override;
+
+  void SetWeights(TfLiteTensor* weights);
+
+  void SetBias(TfLiteTensor* bias);
+
+ private:
+  void FillCoreMLWeights();
+  void FillCoreMLBias();
+
+  TfLiteTensor* weights_ = nullptr;
+  TfLiteTensor* bias_ = nullptr;
+};
+
+}  // namespace coreml
+}  // namespace delegates
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_DELEGATES_COREML_BUILDERS_FULLY_CONNECTED_OP_BUILDER_H_

tensorflow/lite/delegates/coreml/builders/hardswish_op_builder.cc

@@ -0,0 +1,92 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include "tensorflow/lite/delegates/coreml/builders/hardswish_op_builder.h"
+
+#include <string>
+
+#include "mlmodel/format/NeuralNetwork.pb.h"
+#include "tensorflow/lite/core/c/builtin_op_data.h"
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/delegates/coreml/builders/add_op_builder.h"
+#include "tensorflow/lite/delegates/coreml/builders/mul_op_builder.h"
+#include "tensorflow/lite/delegates/coreml/builders/op_builder.h"
+#include "tensorflow/lite/delegates/coreml/builders/op_factory.h"
+
+namespace tflite {
+namespace delegates {
+namespace coreml {
+const std::string& HardSwishOpBuilder::DebugName() {
+  if (debug_name_.empty()) SetDebugName("HardSwishOpBuilder", node_id_);
+  return debug_name_;
+}
+
+CoreML::Specification::NeuralNetworkLayer* HardSwishOpBuilder::Build() {
+  layer_->set_name(DebugName());
+  layer_->mutable_multiply()->set_alpha(1.0f / 6.0f);
+
+  return layer_.release();
+}
+
+TfLiteStatus HardSwishOpBuilder::PopulateSubgraph(TfLiteContext* context) {
+  // hswish(x) = (x/6) * ReLU6(x+3). main layer_ contains the first part, x/6.
+  // ReLU6(x +3) constructed as add op with fused ReLU6 activation.
+  AddOpBuilder* add_builder = reinterpret_cast<AddOpBuilder*>(
+      graph_builder_->AddBuilder(CreateAddOpBuilder, nullptr));
+  TfLiteAddParams add_param{kTfLiteActRelu6};
+  add_builder->SetBuiltinData(&add_param);
+  add_builder->SetAlpha(3.0f);
+  add_builder->AddInput(layer_->input(0));
+  add_builder->PopulateSubgraph(context);
+
+  // multiplies (x/6) from main layer_ and ReLU6(x+3) from the above code.
+  MulOpBuilder* mul_builder = reinterpret_cast<MulOpBuilder*>(
+      graph_builder_->AddBuilder(CreateMulOpBuilder, nullptr));
+  mul_builder->AddInput(AddOutput());
+  mul_builder->AddInput(add_builder->GetOutput(context));
+  builder_output_ = mul_builder->AddOutput();
+  return kTfLiteOk;
+}
+
+TfLiteStatus HardSwishOpBuilder::RegisterInputs(const TfLiteIntArray* inputs,
+                                                TfLiteContext* context) {
+  if (inputs->size != 1) {
+    TF_LITE_KERNEL_LOG(context, "Wrong # of inputs to hardswish!.");
+    return kTfLiteError;
+  }
+  AddInput(inputs->data[0]);
+  return kTfLiteOk;
+}
+
+TfLiteStatus HardSwishOpBuilder::RegisterOutputs(const TfLiteIntArray* outputs,
+                                                 TfLiteContext* context) {
+  if (outputs->size != 1) {
+    TF_LITE_KERNEL_LOG(context, "Wrong # of outputs to hardswish!.");
+    return kTfLiteError;
+  }
+  TensorID output_tensor = GetOutput(context);
+  if (output_tensor.NodeID() == -1) {
+    TF_LITE_KERNEL_LOG(context, "Failed to build output tensor.");
+    return kTfLiteError;
+  }
+  graph_builder_->AddTensorWithID(outputs->data[0], output_tensor);
+  return kTfLiteOk;
+}
+
+OpBuilder* CreateHardSwishOpBuilder(GraphBuilder* graph_builder) {
+  return new HardSwishOpBuilder(graph_builder);
+}
+}  // namespace coreml
+}  // namespace delegates
+}  // namespace tflite

tensorflow/lite/delegates/coreml/builders/hardswish_op_builder.h

@@ -0,0 +1,50 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_DELEGATES_COREML_BUILDERS_HARDSWISH_OP_BUILDER_H_
+#define TENSORFLOW_LITE_DELEGATES_COREML_BUILDERS_HARDSWISH_OP_BUILDER_H_
+
+#include <string>
+
+#include "mlmodel/format/NeuralNetwork.pb.h"
+#include "tensorflow/lite/c/c_api_types.h"
+#include "tensorflow/lite/c/common.h"
+#include "tensorflow/lite/delegates/coreml/builders/op_builder.h"
+
+namespace tflite {
+namespace delegates {
+namespace coreml {
+// hswish(x) = x * ReLU6(x + 3) / 6
+class HardSwishOpBuilder : public OpBuilder {
+ public:
+  explicit HardSwishOpBuilder(GraphBuilder* graph_builder)
+      : OpBuilder(graph_builder) {}
+  const std::string& DebugName() override;
+
+  CoreML::Specification::NeuralNetworkLayer* Build() override;
+
+  TfLiteStatus PopulateSubgraph(TfLiteContext* context) override;
+
+  TfLiteStatus RegisterInputs(const TfLiteIntArray* inputs,
+                              TfLiteContext* context) override;
+
+  TfLiteStatus RegisterOutputs(const TfLiteIntArray* outputs,
+                               TfLiteContext* context) override;
+};
+
+}  // namespace coreml
+}  // namespace delegates
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_DELEGATES_COREML_BUILDERS_HARDSWISH_OP_BUILDER_H_

tensorflow/lite/delegates/coreml/builders/mul_op_builder.cc

@@ -0,0 +1,117 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include "tensorflow/lite/delegates/coreml/builders/mul_op_builder.h"
+
+#include <memory>
+#include <string>
+
+#include "mlmodel/format/NeuralNetwork.pb.h"
+#include "tensorflow/lite/core/c/builtin_op_data.h"
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/delegates/coreml/builders/activation_layer_builder.h"
+#include "tensorflow/lite/delegates/coreml/builders/op_builder.h"
+#include "tensorflow/lite/delegates/coreml/builders/op_factory.h"
+#include "tensorflow/lite/delegates/coreml/builders/util.h"
+#include "tensorflow/lite/kernels/kernel_util.h"
+
+namespace tflite {
+namespace delegates {
+namespace coreml {
+const std::string& MulOpBuilder::DebugName() {
+  if (debug_name_.empty()) SetDebugName("MulOpBuilder", node_id_);
+  return debug_name_;
+}
+
+CoreML::Specification::NeuralNetworkLayer* MulOpBuilder::Build() {
+  if (layer_ == nullptr) {
+    layer_ = std::make_unique<CoreML::Specification::NeuralNetworkLayer>();
+  }
+  // MultiplyLayerParams only has limited broadcasting support. For example:
+  // [B, 1, 1, 1], [B, C, 1, 1], [B, 1, H, W], [B, C, H, W]. other shapes
+  // will make broadcasting fail.
+  layer_->set_name(DebugName());
+  layer_->mutable_multiply();
+  if (alpha_ != 1.0f) {
+    layer_->mutable_multiply()->set_alpha(alpha_);
+  }
+
+  return layer_.release();
+}
+
+TfLiteStatus MulOpBuilder::PopulateSubgraph(TfLiteContext* context) {
+  TfLiteMulParams* params = reinterpret_cast<TfLiteMulParams*>(builtin_data_);
+
+  TfLiteFusedActivation activation = params->activation;
+  if (activation == kTfLiteActNone) {
+    builder_output_ = AddOutput();
+  } else {
+    ActivationLayerBuilder* activation_builder =
+        reinterpret_cast<ActivationLayerBuilder*>(
+            graph_builder_->AddBuilder(CreateActivationLayerBuilder, nullptr));
+    activation_builder->SetActivation(activation);
+    activation_builder->AddInput(AddOutput());
+    activation_builder->PopulateSubgraph(context);
+    builder_output_ = activation_builder->GetOutput(context);
+  }
+  return kTfLiteOk;
+}
+
+TfLiteStatus MulOpBuilder::RegisterInputs(const TfLiteIntArray* inputs,
+                                          TfLiteContext* context) {
+  // TFL MUL op always has 2 inputs.
+  if (inputs->size != 2) {
+    TF_LITE_KERNEL_LOG(context, "Wrong # of inputs to mul!.");
+    return kTfLiteError;
+  }
+  const auto* input_0 = &context->tensors[inputs->data[0]];
+  const auto* input_1 = &context->tensors[inputs->data[1]];
+  // store constant, scalar value into MultiplyLayerParams directly.
+  if (IsConstantTensor(input_0) && NumElements(input_0) == 1) {
+    AddInput(inputs->data[1]);
+    SetAlpha(GetScalarFloatFromTensor(input_0));
+  } else if (IsConstantTensor(input_1) && NumElements(input_1) == 1) {
+    AddInput(inputs->data[0]);
+    SetAlpha(GetScalarFloatFromTensor(input_1));
+  } else {
+    AddInput(inputs->data[0]);
+    AddInput(inputs->data[1]);
+  }
+  return kTfLiteOk;
+}
+
+TfLiteStatus MulOpBuilder::RegisterOutputs(const TfLiteIntArray* outputs,
+                                           TfLiteContext* context) {
+  if (outputs->size != 1) {
+    TF_LITE_KERNEL_LOG(context, "Wrong # of outputs to mul!.");
+    return kTfLiteError;
+  }
+  TensorID output_tensor = GetOutput(context);
+  if (output_tensor.NodeID() == -1) {
+    TF_LITE_KERNEL_LOG(context, "Failed to build output tensor.");
+    return kTfLiteError;
+  }
+  graph_builder_->AddTensorWithID(outputs->data[0], output_tensor);
+  return kTfLiteOk;
+}
+
+void MulOpBuilder::SetAlpha(float alpha) { alpha_ = alpha; }
+
+OpBuilder* CreateMulOpBuilder(GraphBuilder* graph_builder) {
+  return new MulOpBuilder(graph_builder);
+}
+
+}  // namespace coreml
+}  // namespace delegates
+}  // namespace tflite

tensorflow/lite/delegates/coreml/builders/mul_op_builder.h

@@ -0,0 +1,56 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_DELEGATES_COREML_BUILDERS_MUL_OP_BUILDER_H_
+#define TENSORFLOW_LITE_DELEGATES_COREML_BUILDERS_MUL_OP_BUILDER_H_
+
+#include <string>
+
+#include "mlmodel/format/NeuralNetwork.pb.h"
+#include "tensorflow/lite/c/c_api_types.h"
+#include "tensorflow/lite/c/common.h"
+#include "tensorflow/lite/delegates/coreml/builders/op_builder.h"
+
+namespace tflite {
+namespace delegates {
+namespace coreml {
+// Builder for Mul op in CoreML.
+class MulOpBuilder : public OpBuilder {
+ public:
+  explicit MulOpBuilder(GraphBuilder* graph_builder)
+      : OpBuilder(graph_builder) {}
+  const std::string& DebugName() override;
+
+  CoreML::Specification::NeuralNetworkLayer* Build() override;
+
+  TfLiteStatus PopulateSubgraph(TfLiteContext* context) override;
+
+  TfLiteStatus RegisterInputs(const TfLiteIntArray* inputs,
+                              TfLiteContext* context) override;
+
+  TfLiteStatus RegisterOutputs(const TfLiteIntArray* outputs,
+                               TfLiteContext* context) override;
+
+  void SetAlpha(float alpha);
+
+ private:
+  // Used for unary mul
+  float alpha_ = 1.0f;
+};
+
+}  // namespace coreml
+}  // namespace delegates
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_DELEGATES_COREML_BUILDERS_MUL_OP_BUILDER_H_

tensorflow/lite/delegates/coreml/builders/op_builder.cc

@@ -0,0 +1,234 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include "tensorflow/lite/delegates/coreml/builders/op_builder.h"
+
+#include <cstdio>
+#include <functional>
+#include <memory>
+#include <string>
+
+#include "mlmodel/format/Model.pb.h"
+#include "mlmodel/format/NeuralNetwork.pb.h"
+#include "tensorflow/lite/builtin_ops.h"
+#include "tensorflow/lite/c/c_api_types.h"
+#include "tensorflow/lite/c/common.h"
+#include "tensorflow/lite/delegates/coreml/builders/op_factory.h"
+
+namespace tflite {
+namespace delegates {
+namespace coreml {
+
+std::string TensorID::ToString() const {
+  return std::to_string(node_) + "_" + std::to_string(output_id_);
+}
+
+int TensorID::NodeID() const { return node_; }
+
+int TensorID::OutputID() const { return output_id_; }
+
+OpBuilder* GraphBuilder::AddBuilder(int builtin_code, const TfLiteNode* node) {
+  switch (builtin_code) {
+    case kTfLiteBuiltinAdd:
+      return AddBuilder(CreateAddOpBuilder, node);
+    case kTfLiteBuiltinAveragePool2d:
+      return AddBuilder(CreateAveragePool2dOpBuilder, node);
+    case kTfLiteBuiltinConcatenation:
+      return AddBuilder(CreateConcatenationOpBuilder, node);
+    case kTfLiteBuiltinConv2d:
+      return AddBuilder(CreateConvolutionOpBuilder, node);
+    case kTfLiteBuiltinDepthwiseConv2d:
+      return AddBuilder(CreateDepthwiseConvolutionOpBuilder, node);
+    // TODO(b/141490853): Add proper dequantize OpBuilder for int8/uint8 inputs.
+    case kTfLiteBuiltinDequantize:
+      // FP16 dequantize is claimed by the delegate to prevent them from running
+      // on CPU, but don't need to be excuted on the Core ML delegate either.
+      return AddBuilder(CreateDummyOpBuilder, node);
+    case kTfLiteBuiltinFullyConnected:
+      return AddBuilder(CreateFullyConnectedOpBuilder, node);
+    case kTfLiteBuiltinLogistic:
+      return AddBuilder(CreateLogisticOpBuilder, node);
+    case kTfLiteBuiltinMaxPool2d:
+      return AddBuilder(CreateMaxPool2dOpBuilder, node);
+    case kTfLiteBuiltinMean:
+      return AddBuilder(CreateMeanOpBuilder, node);
+    case kTfLiteBuiltinMirrorPad:
+      return AddBuilder(CreateMirrorPadOpBuilder, node);
+    case kTfLiteBuiltinMul:
+      return AddBuilder(CreateMulOpBuilder, node);
+    case kTfLiteBuiltinPad:
+    case kTfLiteBuiltinPadv2:
+      return AddBuilder(CreatePadOpBuilder, node);
+    case kTfLiteBuiltinRelu:
+      return AddBuilder(CreateReluOpBuilder, node);
+    case kTfLiteBuiltinReluN1To1:
+      return AddBuilder(CreateReluN1To1OpBuilder, node);
+    case kTfLiteBuiltinRelu6:
+      return AddBuilder(CreateRelu6OpBuilder, node);
+    case kTfLiteBuiltinReshape:
+      return AddBuilder(CreateReshapeOpBuilder, node);
+    case kTfLiteBuiltinResizeBilinear:
+      return AddBuilder(CreateResizeBilinearOpBuilder, node);
+    case kTfLiteBuiltinSoftmax:
+      return AddBuilder(CreateSoftmaxOpBuilder, node);
+    case kTfLiteBuiltinTanh:
+      return AddBuilder(CreateTanhOpBuilder, node);
+    case kTfLiteBuiltinTransposeConv:
+      return AddBuilder(CreateTransposeConvolutionOpBuilder, node);
+    case kTfLiteBuiltinHardSwish:
+      return AddBuilder(CreateHardSwishOpBuilder, node);
+    default:
+      return nullptr;
+  }
+}
+
+OpBuilder* GraphBuilder::AddBuilder(
+    const std::function<OpBuilder*(GraphBuilder*)>& builder,
+    const TfLiteNode* node) {
+  if (builder == nullptr) {
+    fprintf(stderr, "builder should be set.\n");
+    return nullptr;
+  }
+  OpBuilder* op = builder(this);
+
+  builders_.emplace_back(op);
+  op->SetNodeID(builders_.size());
+  if (node != nullptr) {
+    op->SetBuiltinData(node->builtin_data);
+    op->SetTfLiteNode(node);
+  }
+  return builders_.back().get();
+}
+
+CoreML::Specification::Model* GraphBuilder::BuildModel() {
+  CoreML::Specification::Model* model = new CoreML::Specification::Model();
+  if (coreml_version_ == 2) {  // Core ML 2, iOS >= 12.0
+    model->set_specificationversion(3);
+  } else if (coreml_version_ == 3) {  // Core ML 3, iOS >= 13.0
+    model->set_specificationversion(4);
+    model->mutable_neuralnetwork()->set_arrayinputshapemapping(
+        CoreML::Specification::EXACT_ARRAY_MAPPING);
+  } else {
+    fprintf(stderr, "Unsupported Core ML version: %d\n", coreml_version_);
+    delete model;
+    return nullptr;
+  }
+  auto* neural_network = model->mutable_neuralnetwork();
+  for (auto& builder : builders_) {
+    CoreML::Specification::NeuralNetworkLayer* layer = builder->Build();
+    if (layer == nullptr) {
+      fprintf(stderr, "Null layer returned from builder: %s\n",
+              builder->DebugName().c_str());
+      continue;
+    }
+    neural_network->mutable_layers()->AddAllocated(layer);
+  }
+  return model;
+}
+
+void GraphBuilder::AddTensorWithID(int tf_tensor_id,
+                                   const TensorID& tensor_id) {
+  if (tensors_.size() <= tf_tensor_id) {
+    tensors_.resize(tf_tensor_id + 1);
+    used_tensor_.resize(tf_tensor_id + 1);
+  }
+  tensors_[tf_tensor_id] = tensor_id;
+}
+
+std::string GraphBuilder::GetTensorName(int tensor_id) {
+  return GetTensorID(tensor_id).ToString();
+}
+
+const TensorID GraphBuilder::GetTensorID(int tensor_id) {
+  if (!HasTensor(tensor_id)) {
+    // TODO(karimnosseir): Double check if this happened, if we are
+    // adding in execution order it shouldn't happen.
+    fprintf(stderr, "index out of range...!!! Requested index %d , size %d\n",
+            tensor_id, static_cast<int>(tensors_.size()));
+    // Return invalid ID.
+    return TensorID(-1, -1);
+  }
+  used_tensor_[tensor_id] = true;
+  return tensors_[tensor_id];
+}
+
+bool GraphBuilder::HasTensor(int tflite_tensor_index) {
+  if (tensors_.size() <= tflite_tensor_index) {
+    return false;
+  }
+  return tensors_[tflite_tensor_index].NodeID() != -1;
+}
+
+bool GraphBuilder::IsTensorUsed(int tflite_tensor_index) {
+  if (!HasTensor(tflite_tensor_index)) return false;
+  return used_tensor_[tflite_tensor_index];
+}
+
+CoreML::Specification::NeuralNetworkLayer* OpBuilder::Build() {
+  layer_->set_name(DebugName());
+  return layer_.release();
+}
+
+TfLiteStatus OpBuilder::PopulateSubgraph(TfLiteContext* context) {
+  builder_output_ = AddOutput();
+  return kTfLiteOk;
+}
+
+void OpBuilder::SetBuiltinData(void* builtin_data) {
+  builtin_data_ = builtin_data;
+}
+
+void OpBuilder::SetNodeID(int id) { node_id_ = id; }
+
+void OpBuilder::SetTfLiteNode(const TfLiteNode* node) { tflite_node_ = node; }
+
+int OpBuilder::GetID() const { return node_id_; }
+
+TensorID OpBuilder::GetOutput(TfLiteContext* context) {
+  if (builder_output_.NodeID() != -1) {
+    return builder_output_;
+  }
+  // builder_output_ is not set when PopulateSubgraph is not called.
+  builder_output_ = AddOutput();
+  return builder_output_;
+}
+
+void OpBuilder::AddInput(const std::string& input_name) {
+  if (layer_ == nullptr) {
+    layer_ = std::make_unique<CoreML::Specification::NeuralNetworkLayer>();
+  }
+  *layer_->mutable_input()->Add() = input_name;
+}
+
+void OpBuilder::AddInput(const TensorID& input_id) {
+  AddInput(input_id.ToString());
+}
+
+void OpBuilder::AddInput(int tf_input_id) {
+  AddInput(graph_builder_->GetTensorName(tf_input_id));
+}
+
+TensorID OpBuilder::AddOutput() {
+  auto tensor_id = TensorID(GetID(), num_outputs_++);
+  *layer_->mutable_output()->Add() = tensor_id.ToString();
+  return tensor_id;
+}
+
+void OpBuilder::SetDebugName(const char* name, int id) {
+  debug_name_ = std::string(name) + "_" + std::to_string(id);
+}
+
+}  // namespace coreml
+}  // namespace delegates
+}  // namespace tflite

tensorflow/lite/delegates/coreml/builders/op_builder.h

@@ -0,0 +1,172 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_DELEGATES_COREML_BUILDERS_OP_BUILDER_H_
+#define TENSORFLOW_LITE_DELEGATES_COREML_BUILDERS_OP_BUILDER_H_
+
+#include <functional>
+#include <string>
+
+#include "mlmodel/format/Model.pb.h"
+#include "mlmodel/format/NeuralNetwork.pb.h"
+#include "tensorflow/lite/core/c/common.h"
+
+namespace tflite {
+namespace delegates {
+namespace coreml {
+class OpBuilder;
+
+// A class represents an ID in the coreML graph.
+// A node is represented by a pair (node_id, and output_index)
+// API is experimental and subject to change.
+class TensorID {
+ public:
+  TensorID() {}
+  TensorID(int node, int output_id) : node_(node), output_id_(output_id) {}
+
+  std::string ToString() const;
+
+  int NodeID() const;
+
+  int OutputID() const;
+
+ private:
+  int node_ = -1;
+  int output_id_ = -1;
+};
+
+// Builder for the whole graph.
+// All op builders should be added using AddBuilder
+// and then BuildModel should be called to return the CoreML generated.
+//
+// API is experimental and subject to change.
+class GraphBuilder {
+ public:
+  explicit GraphBuilder(int coreml_version) : coreml_version_(coreml_version) {}
+
+  // Returns pointer to the created builder. Ownership still belongs
+  // to the GraphBuilder.
+  OpBuilder* AddBuilder(int builtin_code, const TfLiteNode* node);
+
+  // Returns pointer to the created builder with op builder function provided.
+  OpBuilder* AddBuilder(const std::function<OpBuilder*(GraphBuilder*)>& builder,
+                        const TfLiteNode* node);
+
+  // Builds Model instance and returns it.
+  CoreML::Specification::Model* BuildModel();
+
+  // Returns string representing tensor 'tensor_id' in coreML.
+  // tensor_id should have been added before calling this method.
+  std::string GetTensorName(int tensor_id);
+
+  // Returns Core ML Tensor ID for TFL 'tensor_id'.
+  // tensor_id should have been added before calling this method.
+  const TensorID GetTensorID(int tensor_id);
+
+  void AddTensorWithID(int tf_tensor_id, const TensorID& tensor_id);
+
+  // Return true if this tensor was added before to the graph.
+  bool HasTensor(int tflite_tensor_index);
+  // Return if this tensor is used in the graph (not as data).
+  // This information is used to mark constant tensors that are used as input.
+  bool IsTensorUsed(int tflite_tensor_index);
+
+  const int coreml_version_;
+
+ private:
+  std::vector<std::unique_ptr<OpBuilder>> builders_;
+  // Index in the vector is the tflite_tensor_index, the value
+  // is the ID in the coreml graph.
+  std::vector<TensorID> tensors_;
+  std::vector<bool> used_tensor_;
+};
+
+// Interface for all op layers
+// API is experimental and subject to change.
+class OpBuilder {
+ public:
+  explicit OpBuilder(GraphBuilder* graph_builder)
+      : graph_builder_(graph_builder) {}
+  virtual ~OpBuilder() {}
+
+  // Returns the Layer this builder responsible for.
+  // Ownership is transferred to caller.
+  virtual CoreML::Specification::NeuralNetworkLayer* Build();
+
+  // Associates TfLite input tensors to Core ML layer's inputs and properties.
+  // Verification for input constraints should happen here.
+  virtual TfLiteStatus RegisterInputs(const TfLiteIntArray* inputs,
+                                      TfLiteContext* context) = 0;
+
+  // Associates TFLite output tensor with the node's output. If the OpBuilder
+  // has subgraphs, The final output of that subgraph should be associated with
+  // the output tensor.
+  virtual TfLiteStatus RegisterOutputs(const TfLiteIntArray* outputs,
+                                       TfLiteContext* context) = 0;
+
+  // Adds additional required OpBuilders, and populate builder_output_ with
+  // Actual output that corresponds to output tensor of TFL Node.
+  // Clients need to override this in cases where the nodes can be used for
+  // composing other ops. For example, Relu6 in TfLite can be converted to
+  // Relu -> Threshold -> Neg.
+  // TODO(b/147211734): have this called automatically when necessary.
+  virtual TfLiteStatus PopulateSubgraph(TfLiteContext* context);
+
+  virtual const std::string& DebugName() = 0;
+
+  void SetBuiltinData(void* builtin_data);
+
+  void SetNodeID(int id);
+
+  void SetTfLiteNode(const TfLiteNode* node);
+
+  int GetID() const;
+
+  // Adds input with tensor name.
+  void AddInput(const std::string& input_name);
+
+  // Adds input with CoreML tensor ID.
+  void AddInput(const TensorID& input_id);
+
+  // Adds input with TF Lite tensor ID.
+  // TODO(taeheej): cleanup AddInput use cases and used tensor tracking.
+  void AddInput(int tf_input_id);
+
+  // Simply adds new output to the underlying layer.
+  TensorID AddOutput();
+
+  // Should set builder_output_ (if unset) and return it as the output of
+  // this node. To be used by clients that needs the output of the node.
+  virtual TensorID GetOutput(TfLiteContext* context);
+
+ protected:
+  // Sets layer's name.
+  void SetDebugName(const char* layer_name, int id);
+
+  GraphBuilder* graph_builder_ = nullptr;
+  // Data needed by this node.
+  void* builtin_data_ = nullptr;
+  int node_id_ = -1;
+  int num_outputs_ = 0;
+  const TfLiteNode* tflite_node_ = nullptr;
+  TensorID builder_output_;
+  std::string debug_name_;
+  std::unique_ptr<CoreML::Specification::NeuralNetworkLayer> layer_;
+};
+
+}  // namespace coreml
+}  // namespace delegates
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_DELEGATES_COREML_BUILDERS_OP_BUILDER_H_

tensorflow/lite/delegates/coreml/builders/op_factory.h

@@ -0,0 +1,58 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_DELEGATES_COREML_BUILDERS_OP_FACTORY_H_
+#define TENSORFLOW_LITE_DELEGATES_COREML_BUILDERS_OP_FACTORY_H_
+
+#include "tensorflow/lite/core/c/builtin_op_data.h"
+
+namespace tflite {
+namespace delegates {
+namespace coreml {
+class GraphBuilder;
+class OpBuilder;
+
+OpBuilder* CreateAddOpBuilder(GraphBuilder* graph_builder);
+OpBuilder* CreateAveragePool2dOpBuilder(GraphBuilder* graph_builder);
+OpBuilder* CreateConcatenationOpBuilder(GraphBuilder* graph_builder);
+OpBuilder* CreateConvolutionOpBuilder(GraphBuilder* graph_builder);
+OpBuilder* CreateDepthwiseConvolutionOpBuilder(GraphBuilder* graph_builder);
+OpBuilder* CreateFullyConnectedOpBuilder(GraphBuilder* graph_builder);
+OpBuilder* CreateHardSwishOpBuilder(GraphBuilder* graph_builder);
+OpBuilder* CreateLogisticOpBuilder(GraphBuilder* graph_builder);
+OpBuilder* CreateMaxPool2dOpBuilder(GraphBuilder* graph_builder);
+OpBuilder* CreateMeanOpBuilder(GraphBuilder* graph_builder);
+OpBuilder* CreateMirrorPadOpBuilder(GraphBuilder* graph_builder);
+OpBuilder* CreateMulOpBuilder(GraphBuilder* graph_builder);
+// PAD handles PAD and PADV2 together.
+OpBuilder* CreatePadOpBuilder(GraphBuilder* graph_builder);
+OpBuilder* CreateReluOpBuilder(GraphBuilder* graph_builder);
+OpBuilder* CreateReluN1To1OpBuilder(GraphBuilder* graph_builder);
+OpBuilder* CreateRelu6OpBuilder(GraphBuilder* graph_builder);
+OpBuilder* CreateReshapeOpBuilder(GraphBuilder* graph_builder);
+OpBuilder* CreateResizeBilinearOpBuilder(GraphBuilder* graph_builder);
+OpBuilder* CreateSoftmaxOpBuilder(GraphBuilder* graph_builder);
+OpBuilder* CreateTanhOpBuilder(GraphBuilder* graph_builder);
+OpBuilder* CreateTransposeConvolutionOpBuilder(GraphBuilder* graph_builder);
+
+OpBuilder* CreateActivationLayerBuilder(GraphBuilder* graph_builder);
+OpBuilder* CreateThresholdLayerBuilder(GraphBuilder* graph_builder);
+// Dummy Opbuilder for nodes that are claimed but not used. ex) FP16 dequantize
+OpBuilder* CreateDummyOpBuilder(GraphBuilder* graph_builder);
+
+}  // namespace coreml
+}  // namespace delegates
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_DELEGATES_COREML_BUILDERS_OP_FACTORY_H_

tensorflow/lite/delegates/coreml/builders/op_validator.h

@@ -0,0 +1,53 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_DELEGATES_COREML_BUILDERS_OP_VALIDATOR_H_
+#define TENSORFLOW_LITE_DELEGATES_COREML_BUILDERS_OP_VALIDATOR_H_
+
+#include "tensorflow/lite/core/c/builtin_op_data.h"
+
+namespace tflite {
+namespace delegates {
+namespace coreml {
+// Follow the ordering of TfLiteBuiltinOperator enum.
+bool IsConcatenationOpSupported(const TfLiteRegistration* registration,
+                                const TfLiteNode* node, TfLiteContext* context);
+bool IsConvolutionOpSupported(const TfLiteRegistration* registration,
+                              const TfLiteNode* node, TfLiteContext* context);
+bool IsDepthwiseConvolutionOpSupported(const TfLiteRegistration* registration,
+                                       const TfLiteNode* node,
+                                       TfLiteContext* context);
+bool IsFullyConnectedOpSupported(const TfLiteRegistration* registration,
+                                 const TfLiteNode* node,
+                                 TfLiteContext* context);
+bool IsMeanOpSupported(const TfLiteRegistration* registration,
+                       const TfLiteNode* node, TfLiteContext* context);
+bool IsMirrorPadOpSupported(const TfLiteRegistration* registration,
+                            const TfLiteNode* node, TfLiteContext* context);
+bool IsPadOpSupported(const TfLiteRegistration* registration,
+                      const TfLiteNode* node, TfLiteContext* context);
+bool IsReshapeOpSupported(const TfLiteRegistration* registration,
+                          const TfLiteNode* node, TfLiteContext* context,
+                          int coreml_version);
+bool IsResizeBilinearOpSupported(const TfLiteRegistration* registration,
+                                 const TfLiteNode* node,
+                                 TfLiteContext* context);
+bool IsTransposeConvolutionOpSupported(const TfLiteRegistration* registration,
+                                       const TfLiteNode* node,
+                                       TfLiteContext* context);
+}  // namespace coreml
+}  // namespace delegates
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_DELEGATES_COREML_BUILDERS_OP_VALIDATOR_H_

tensorflow/lite/delegates/coreml/builders/pad_op_builder.cc

@@ -0,0 +1,147 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include "tensorflow/lite/delegates/coreml/builders/pad_op_builder.h"
+
+#include <cstdint>
+#include <string>
+
+#include "mlmodel/format/NeuralNetwork.pb.h"
+#include "tensorflow/lite/core/c/builtin_op_data.h"
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/delegates/coreml/builders/op_factory.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
+#include "tensorflow/lite/kernels/kernel_util.h"
+
+namespace tflite {
+namespace delegates {
+namespace coreml {
+
+const std::string& PadOpBuilder::DebugName() {
+  if (!debug_name_.empty()) return debug_name_;
+  SetDebugName(padding_type_ == PadType::kPad ? "PadOpBuilder (PAD)"
+                                              : "PadOpBuilder (MIRROR_PAD)",
+               node_id_);
+  return debug_name_;
+}
+
+CoreML::Specification::NeuralNetworkLayer* PadOpBuilder::Build() {
+  layer_->set_name(DebugName());
+  if (padding_type_ == PadType::kPad) {
+    layer_->mutable_padding()->mutable_constant();
+  } else if (padding_type_ == PadType::kMirrorPad) {
+    layer_->mutable_padding()->mutable_reflection();
+  }
+  return layer_.release();
+}
+
+// padding is d x 2 tensor, where d is the dimension of input.
+// only paddings for width and height are considered.
+void PadOpBuilder::SetPadding(const TfLiteTensor* padding) {
+  const int32_t* padding_data = GetTensorData<int32_t>(padding);
+  for (int i = 1; i <= 2; ++i) {
+    auto* borderamount = layer_->mutable_padding()
+                             ->mutable_paddingamounts()
+                             ->add_borderamounts();
+    borderamount->set_startedgesize(padding_data[i * 2]);
+    borderamount->set_endedgesize(padding_data[i * 2 + 1]);
+  }
+}
+
+void PadOpBuilder::SetConstantValue(const TfLiteTensor* constant_value) {
+  layer_->mutable_padding()->mutable_constant()->set_value(
+      GetTensorData<float>(constant_value)[0]);
+}
+
+TfLiteStatus PadOpBuilder::RegisterInputs(const TfLiteIntArray* inputs,
+                                          TfLiteContext* context) {
+  if (!(inputs->size == 2 || inputs->size == 3)) {
+    TF_LITE_KERNEL_LOG(context, "Wrong # of inputs to Padding!.");
+    return kTfLiteError;
+  }
+  AddInput(inputs->data[0]);
+  SetPadding(GetInput(context, tflite_node_, 1));
+  if (inputs->size == 3) {
+    SetConstantValue(GetInput(context, tflite_node_, 2));
+  }
+
+  return kTfLiteOk;
+}
+
+TfLiteStatus PadOpBuilder::RegisterOutputs(const TfLiteIntArray* outputs,
+                                           TfLiteContext* context) {
+  if (outputs->size != 1) {
+    TF_LITE_KERNEL_LOG(context, "Wrong # of outputs to Padding!.");
+    return kTfLiteError;
+  }
+  graph_builder_->AddTensorWithID(outputs->data[0], GetOutput(context));
+  return kTfLiteOk;
+}
+
+OpBuilder* CreatePadOpBuilder(GraphBuilder* graph_builder) {
+  return new PadOpBuilder(graph_builder, PadType::kPad);
+}
+
+OpBuilder* CreateMirrorPadOpBuilder(GraphBuilder* graph_builder) {
+  return new PadOpBuilder(graph_builder, PadType::kMirrorPad);
+}
+
+bool IsPadOpSupported(const TfLiteRegistration* registration,
+                      const TfLiteNode* node, TfLiteContext* context) {
+  // padding is d x 2 tensor, where d is the dimension of input.
+  const TfLiteTensor* padding;
+  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, 1, &padding));
+  if (!IsConstantTensor(padding)) {
+    TF_LITE_KERNEL_LOG(context,
+                       "%s: Only constant padding is supported for PAD.",
+                       padding->name);
+    return false;
+  }
+  if (padding->dims->data[0] != 4 || padding->dims->data[1] != 2) {
+    TF_LITE_KERNEL_LOG(context, "%s: Only 4D inputs are supported for PAD.",
+                       padding->name);
+    return false;
+  }
+  const int32_t* padding_data = GetTensorData<int32_t>(padding);
+  if (!(padding_data[0] == 0 && padding_data[1] == 0)) {
+    TF_LITE_KERNEL_LOG(
+        context, "%s: Padding for batch dimension is not supported in PAD.",
+        padding->name);
+    return false;
+  }
+
+  if (!(padding_data[6] == 0 && padding_data[7] == 0)) {
+    TF_LITE_KERNEL_LOG(
+        context, "%s: Padding for channel dimension is not supported in PAD.",
+        padding->name);
+    return false;
+  }
+  return true;
+}
+
+bool IsMirrorPadOpSupported(const TfLiteRegistration* registration,
+                            const TfLiteNode* node, TfLiteContext* context) {
+  auto* params =
+      reinterpret_cast<TfLiteMirrorPaddingParams*>(node->builtin_data);
+  if (params->mode != kTfLiteMirrorPaddingReflect) {
+    TF_LITE_KERNEL_LOG(context,
+                       "Only REFLECT mode is supported for MIRROR_PAD.");
+    return false;
+  }
+  return IsPadOpSupported(registration, node, context);
+}
+
+}  // namespace coreml
+}  // namespace delegates
+}  // namespace tflite

tensorflow/lite/delegates/coreml/builders/pad_op_builder.h

@@ -0,0 +1,60 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_DELEGATES_COREML_BUILDERS_PAD_OP_BUILDER_H_
+#define TENSORFLOW_LITE_DELEGATES_COREML_BUILDERS_PAD_OP_BUILDER_H_
+
+#include <string>
+
+#include "mlmodel/format/NeuralNetwork.pb.h"
+#include "tensorflow/lite/builtin_ops.h"
+#include "tensorflow/lite/c/c_api_types.h"
+#include "tensorflow/lite/c/common.h"
+#include "tensorflow/lite/delegates/coreml/builders/op_builder.h"
+
+namespace tflite {
+namespace delegates {
+namespace coreml {
+
+enum class PadType { kPad, kMirrorPad };
+
+// Supports PAD, PADV2, MIRROR_PAD
+class PadOpBuilder : public OpBuilder {
+ public:
+  explicit PadOpBuilder(GraphBuilder* graph_builder, PadType padding_type)
+      : OpBuilder(graph_builder), padding_type_(padding_type) {}
+
+  const std::string& DebugName() override;
+
+  CoreML::Specification::NeuralNetworkLayer* Build() override;
+
+  TfLiteStatus RegisterInputs(const TfLiteIntArray* inputs,
+                              TfLiteContext* context) override;
+
+  TfLiteStatus RegisterOutputs(const TfLiteIntArray* outputs,
+                               TfLiteContext* context) override;
+
+  void SetPadding(const TfLiteTensor* padding);
+
+  void SetConstantValue(const TfLiteTensor* constant_value);
+
+ private:
+  PadType padding_type_;
+};
+
+}  // namespace coreml
+}  // namespace delegates
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_DELEGATES_COREML_BUILDERS_PAD_OP_BUILDER_H_

tensorflow/lite/delegates/coreml/builders/pooling_layer_builder.cc

@@ -0,0 +1,170 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include "tensorflow/lite/delegates/coreml/builders/pooling_layer_builder.h"
+
+#include <cstdio>
+#include <string>
+#include <vector>
+
+#include "mlmodel/format/NeuralNetwork.pb.h"
+#include "tensorflow/lite/builtin_ops.h"
+#include "tensorflow/lite/core/c/builtin_op_data.h"
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/delegates/coreml/builders/op_factory.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
+#include "tensorflow/lite/kernels/kernel_util.h"
+
+namespace tflite {
+namespace delegates {
+namespace coreml {
+
+const std::string& PoolingLayerBuilder::DebugName() {
+  if (!debug_name_.empty()) return debug_name_;
+  switch (pooling_type_) {
+    case kTfLiteBuiltinAveragePool2d:
+      SetDebugName("PoolingLayerBuilder (AVERAGE)", node_id_);
+      break;
+    case kTfLiteBuiltinMaxPool2d:
+      SetDebugName("PoolingLayerBuilder (MAX)", node_id_);
+      break;
+    case kTfLiteBuiltinL2Pool2d:
+      SetDebugName("PoolingLayerBuilder (L2, unsupported)", node_id_);
+      break;
+    case kTfLiteBuiltinMean:
+      SetDebugName("PoolingLayerBuilder (MEAN)", node_id_);
+      break;
+    default:
+      SetDebugName("PoolingLayerBuilder (ERROR)", node_id_);
+  }
+  return debug_name_;
+}
+
+CoreML::Specification::NeuralNetworkLayer* PoolingLayerBuilder::Build() {
+  layer_->set_name(DebugName());
+  auto* pooling_params = layer_->mutable_pooling();
+
+  if (pooling_type_ == kTfLiteBuiltinMean) {
+    pooling_params->set_type(
+        CoreML::Specification::PoolingLayerParams::AVERAGE);
+    pooling_params->set_globalpooling(true);
+    return layer_.release();
+  }
+
+  const TfLitePoolParams* params =
+      reinterpret_cast<const TfLitePoolParams*>(builtin_data_);
+  pooling_params->mutable_stride()->Add(params->stride_height);
+  pooling_params->mutable_stride()->Add(params->stride_width);
+  pooling_params->mutable_kernelsize()->Add(params->filter_height);
+  pooling_params->mutable_kernelsize()->Add(params->filter_width);
+
+  if (params->padding == kTfLitePaddingSame) {
+    pooling_params->mutable_same();
+  } else {
+    pooling_params->mutable_valid();
+  }
+
+  switch (pooling_type_) {
+    case kTfLiteBuiltinAveragePool2d:
+      pooling_params->set_type(
+          CoreML::Specification::PoolingLayerParams::AVERAGE);
+      pooling_params->set_avgpoolexcludepadding(true);
+      break;
+    case kTfLiteBuiltinMaxPool2d:
+      pooling_params->set_type(CoreML::Specification::PoolingLayerParams::MAX);
+      break;
+    case kTfLiteBuiltinL2Pool2d:
+      // TODO(b/145873272) implement L2 pooling
+      // NOLINTNEXTLINE: minimize absl usage
+      fprintf(stderr, "L2 pooling is not supported yet.\n");
+      return nullptr;
+    default:
+      // NOLINTNEXTLINE: minimize absl usage
+      fprintf(stderr, "Unexpected pooling type.\n");  // Should not reach here.
+      return nullptr;
+  }
+
+  // TODO(b/145582958): Add padding values.
+  // TODO(b/145582958): Handle fused activation function.
+  return layer_.release();
+}
+
+TfLiteStatus PoolingLayerBuilder::RegisterInputs(const TfLiteIntArray* inputs,
+                                                 TfLiteContext* context) {
+  if (pooling_type_ == kTfLiteBuiltinMean) {
+    if (inputs->size != 2) {
+      TF_LITE_KERNEL_LOG(context, "Wrong # of inputs to Mean!.");
+      return kTfLiteError;
+    }
+  } else if (inputs->size != 1) {
+    TF_LITE_KERNEL_LOG(context, "Wrong # of inputs to Pooling!.");
+    return kTfLiteError;
+  }
+  AddInput(inputs->data[0]);
+  return kTfLiteOk;
+}
+
+TfLiteStatus PoolingLayerBuilder::RegisterOutputs(const TfLiteIntArray* outputs,
+                                                  TfLiteContext* context) {
+  if (outputs->size != 1) {
+    TF_LITE_KERNEL_LOG(context, "Wrong # of outputs to Pooling!.");
+    return kTfLiteError;
+  }
+  graph_builder_->AddTensorWithID(outputs->data[0], GetOutput(context));
+  return kTfLiteOk;
+}
+
+OpBuilder* CreateAveragePool2dOpBuilder(GraphBuilder* graph_builder) {
+  return new PoolingLayerBuilder(graph_builder, kTfLiteBuiltinAveragePool2d);
+}
+
+OpBuilder* CreateMaxPool2dOpBuilder(GraphBuilder* graph_builder) {
+  return new PoolingLayerBuilder(graph_builder, kTfLiteBuiltinMaxPool2d);
+}
+
+OpBuilder* CreateMeanOpBuilder(GraphBuilder* graph_builder) {
+  return new PoolingLayerBuilder(graph_builder, kTfLiteBuiltinMean);
+}
+
+// Only supports averaging over H and W dimensions, as
+bool IsMeanOpSupported(const TfLiteRegistration* registration,
+                       const TfLiteNode* node, TfLiteContext* context) {
+  const TfLiteTensor* input = GetInput(context, node, 0);
+  const TfLiteTensor* axis = GetInput(context, node, 1);
+  const auto* params =
+      reinterpret_cast<TfLiteReducerParams*>(node->builtin_data);
+
+  if (!params->keep_dims) {
+    TF_LITE_KERNEL_LOG(context, "keep_dims should be true for Mean op.");
+    return false;
+  }
+  if (input->dims->size != 4) {
+    TF_LITE_KERNEL_LOG(context, "Mean op is only supported for 4D input.");
+    return false;
+  }
+  const int* axis_data = GetTensorData<int>(axis);
+  std::vector<bool> axis_mask = {false, true, true, false};
+  for (int i = 0; i < axis->dims->data[0]; ++i) {
+    if (!axis_mask[(axis_data[i] + 4) % 4]) {
+      TF_LITE_KERNEL_LOG(context,
+                         "Mean op should reduce for H and W dimensions.");
+      return false;
+    }
+  }
+  return true;
+}
+
+}  // namespace coreml
+}  // namespace delegates
+}  // namespace tflite

tensorflow/lite/delegates/coreml/builders/pooling_layer_builder.h

@@ -0,0 +1,55 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_DELEGATES_COREML_BUILDERS_POOLING_LAYER_BUILDER_H_
+#define TENSORFLOW_LITE_DELEGATES_COREML_BUILDERS_POOLING_LAYER_BUILDER_H_
+
+#include <string>
+
+#include "mlmodel/format/NeuralNetwork.pb.h"
+#include "tensorflow/lite/builtin_ops.h"
+#include "tensorflow/lite/c/c_api_types.h"
+#include "tensorflow/lite/c/common.h"
+#include "tensorflow/lite/delegates/coreml/builders/op_builder.h"
+
+namespace tflite {
+namespace delegates {
+namespace coreml {
+
+class PoolingLayerBuilder : public OpBuilder {
+ public:
+  explicit PoolingLayerBuilder(GraphBuilder* graph_builder,
+                               TfLiteBuiltinOperator pooling_type)
+      : OpBuilder(graph_builder), pooling_type_(pooling_type) {}
+
+  const std::string& DebugName() override;
+
+  CoreML::Specification::NeuralNetworkLayer* Build() override;
+
+  TfLiteStatus RegisterInputs(const TfLiteIntArray* inputs,
+                              TfLiteContext* context) override;
+
+  TfLiteStatus RegisterOutputs(const TfLiteIntArray* outputs,
+                               TfLiteContext* context) override;
+
+ private:
+  // Should be one of pooling types.
+  TfLiteBuiltinOperator pooling_type_;
+};
+
+}  // namespace coreml
+}  // namespace delegates
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_DELEGATES_COREML_BUILDERS_POOLING_LAYER_BUILDER_H_

tensorflow/lite/delegates/coreml/builders/reshape_op_builder.cc

@@ -0,0 +1,155 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include "tensorflow/lite/delegates/coreml/builders/reshape_op_builder.h"
+
+#include <algorithm>
+#include <cstdint>
+#include <cstring>
+#include <iterator>
+#include <memory>
+#include <string>
+
+#include "mlmodel/format/NeuralNetwork.pb.h"
+#include "tensorflow/lite/core/c/builtin_op_data.h"
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/delegates/coreml/builders/op_builder.h"
+#include "tensorflow/lite/delegates/coreml/builders/op_factory.h"
+#include "tensorflow/lite/delegates/coreml/builders/op_validator.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
+#include "tensorflow/lite/kernels/kernel_util.h"
+
+namespace tflite {
+namespace delegates {
+namespace coreml {
+
+const std::string& ReshapeOpBuilder::DebugName() {
+  if (debug_name_.empty()) {
+    SetDebugName("ReshapeOpBuilder", node_id_);
+  }
+  return debug_name_;
+}
+
+CoreML::Specification::NeuralNetworkLayer* ReshapeOpBuilder::Build() {
+  if (layer_ == nullptr) {
+    layer_ = std::make_unique<CoreML::Specification::NeuralNetworkLayer>();
+  }
+  layer_->set_name(DebugName());
+  for (int dim : shape_) {
+    layer_->mutable_reshape()->add_targetshape(dim);
+  }
+  if (need_transpose_)
+    layer_->mutable_reshape()->set_mode(
+        CoreML::Specification::ReshapeLayerParams::CHANNEL_LAST);
+  return layer_.release();
+}
+
+void ReshapeOpBuilder::SetShapeFromTensor(const TfLiteTensor* output_shape,
+                                          const TfLiteIntArray* input_shape) {
+  TfLiteIntArray* shape = TfLiteIntArrayCreate(output_shape->dims->data[0]);
+  std::memcpy(shape->data, GetTensorData<int>(output_shape),
+              shape->size * sizeof(int));
+
+  SetShapeFromIntArray(shape, input_shape);
+  TfLiteIntArrayFree(shape);
+}
+
+void ReshapeOpBuilder::SetShapeFromIntArray(const TfLiteIntArray* output_shape,
+                                            const TfLiteIntArray* input_shape) {
+  // ignore first dimension (batch)
+  std::copy(output_shape->data + 1, output_shape->data + output_shape->size,
+            std::back_inserter(shape_));
+
+  int64_t reshape_size = 1;
+  int negative_index = -1;
+  for (int i = 0; i < shape_.size(); ++i) {
+    if (shape_[i] == -1) {
+      negative_index = i;
+    } else {
+      reshape_size *= shape_[i];
+    }
+  }
+  if (negative_index >= 0) {
+    int64_t input_size = NumElements(input_shape);
+    shape_[negative_index] = input_size / reshape_size;
+  }
+
+  if (shape_.size() == 2) {
+    shape_ = {shape_[1], 1, shape_[0]};
+  } else if (shape_.size() == 3) {
+    shape_ = {shape_[2], shape_[0], shape_[1]};
+  }
+  // When channel dimension is changed, reshape should be done with HWC layout.
+  if (shape_[0] != input_shape->data[input_shape->size - 1]) {
+    need_transpose_ = true;
+  }
+}
+
+TfLiteStatus ReshapeOpBuilder::RegisterInputs(const TfLiteIntArray* inputs,
+                                              TfLiteContext* context) {
+  AddInput(inputs->data[0]);
+
+  if (inputs->size == 2) {
+    SetShapeFromTensor(&context->tensors[inputs->data[1]],
+                       context->tensors[inputs->data[0]].dims);
+  } else {
+    const auto* params = reinterpret_cast<TfLiteReshapeParams*>(builtin_data_);
+    TfLiteIntArray* output_shape = TfLiteIntArrayCreate(params->num_dimensions);
+    std::memcpy(output_shape->data, params->shape,
+                params->num_dimensions * sizeof(int));
+
+    SetShapeFromIntArray(output_shape, context->tensors[inputs->data[0]].dims);
+    TfLiteIntArrayFree(output_shape);
+  }
+  return kTfLiteOk;
+}
+
+TfLiteStatus ReshapeOpBuilder::RegisterOutputs(const TfLiteIntArray* outputs,
+                                               TfLiteContext* context) {
+  graph_builder_->AddTensorWithID(outputs->data[0], GetOutput(context));
+  return kTfLiteOk;
+}
+
+bool IsReshapeOpSupported(const TfLiteRegistration* registration,
+                          const TfLiteNode* node, TfLiteContext* context,
+                          int coreml_version) {
+  if (coreml_version >= 3) {
+    return false;
+  }
+  if (node->inputs->size == 1) {
+    const auto* params =
+        reinterpret_cast<TfLiteReshapeParams*>(node->builtin_data);
+    return params->num_dimensions == 3 || params->num_dimensions == 4;
+  }
+
+  const int kShapeTensor = 1;
+  const TfLiteTensor* shape;
+  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kShapeTensor, &shape));
+  if (shape->allocation_type != kTfLiteMmapRo) {
+    TF_LITE_KERNEL_LOG(context, "Reshape has non-const shape.");
+    return false;
+  }
+  const bool is_shape_tensor =
+      shape->dims->size == 1 && shape->type == kTfLiteInt32;
+  return is_shape_tensor &&
+         (shape->dims->data[0] == 3 || shape->dims->data[0] == 4);
+}
+
+OpBuilder* CreateReshapeOpBuilder(GraphBuilder* graph_builder) {
+  return new ReshapeOpBuilder(graph_builder);
+}
+
+}  // namespace coreml
+}  // namespace delegates
+}  // namespace tflite

tensorflow/lite/delegates/coreml/builders/reshape_op_builder.h

@@ -0,0 +1,59 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_DELEGATES_COREML_BUILDERS_RESHAPE_OP_BUILDER_H_
+#define TENSORFLOW_LITE_DELEGATES_COREML_BUILDERS_RESHAPE_OP_BUILDER_H_
+
+#include <string>
+
+#include "mlmodel/format/NeuralNetwork.pb.h"
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/delegates/coreml/builders/op_builder.h"
+
+namespace tflite {
+namespace delegates {
+namespace coreml {
+// Builder for Reshape op in CoreML.
+class ReshapeOpBuilder : public OpBuilder {
+ public:
+  explicit ReshapeOpBuilder(GraphBuilder* graph_builder)
+      : OpBuilder(graph_builder) {}
+  const std::string& DebugName() override;
+
+  CoreML::Specification::NeuralNetworkLayer* Build() override;
+
+  TfLiteStatus RegisterInputs(const TfLiteIntArray* inputs,
+                              TfLiteContext* context) override;
+  TfLiteStatus RegisterOutputs(const TfLiteIntArray* outputs,
+                               TfLiteContext* context) override;
+
+  // Sets output shape of the Core ML reshape layer, given output shape and
+  // the input tensor's shape.
+  void SetShapeFromTensor(const TfLiteTensor* output_shape,
+                          const TfLiteIntArray* input_shape);
+  void SetShapeFromIntArray(const TfLiteIntArray* output_shape,
+                            const TfLiteIntArray* input_shape);
+
+ private:
+  std::vector<int> shape_;
+  // When channel dimension is changed, reshape should be done with HWC layout,
+  // thus transpose is required. (set with ReshapeLayerParams.mode)
+  bool need_transpose_ = false;
+};
+
+}  // namespace coreml
+}  // namespace delegates
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_DELEGATES_COREML_BUILDERS_RESHAPE_OP_BUILDER_H_

tensorflow/lite/delegates/coreml/builders/resize_bilinear_op_builder.cc

@@ -0,0 +1,108 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include "tensorflow/lite/delegates/coreml/builders/resize_bilinear_op_builder.h"
+
+#include <cstdint>
+#include <memory>
+#include <string>
+
+#include "mlmodel/format/NeuralNetwork.pb.h"
+#include "tensorflow/lite/core/c/builtin_op_data.h"
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/delegates/coreml/builders/op_factory.h"
+#include "tensorflow/lite/delegates/coreml/builders/op_validator.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
+#include "tensorflow/lite/kernels/kernel_util.h"
+
+namespace tflite {
+namespace delegates {
+namespace coreml {
+
+const std::string& ResizeBilinearOpBuilder::DebugName() {
+  if (!debug_name_.empty()) return debug_name_;
+  SetDebugName("ResizeBilinearOpBuilder", node_id_);
+  return debug_name_;
+}
+
+CoreML::Specification::NeuralNetworkLayer* ResizeBilinearOpBuilder::Build() {
+  if (layer_ == nullptr) {
+    layer_ = std::make_unique<CoreML::Specification::NeuralNetworkLayer>();
+  }
+  layer_->set_name(DebugName());
+  const TfLiteResizeBilinearParams* params =
+      reinterpret_cast<const TfLiteResizeBilinearParams*>(builtin_data_);
+
+  layer_->mutable_resizebilinear()->mutable_targetsize()->Add(height_);
+  layer_->mutable_resizebilinear()->mutable_targetsize()->Add(width_);
+
+  // align_corners makes last sampling position to be aligned with last index of
+  // input. This is the same behavior as STRICT_ALIGN_ENDPOINTS_MODE in Core ML
+  // sampling mode. When not set, the sampling positions are the same as
+  // UPSAMPLE_MODE. (indices are in [0, (input_size-1)/output_size])
+  if (params->align_corners) {
+    layer_->mutable_resizebilinear()->mutable_mode()->set_samplingmethod(
+        CoreML::Specification::SamplingMode::STRICT_ALIGN_ENDPOINTS_MODE);
+  } else {
+    layer_->mutable_resizebilinear()->mutable_mode()->set_samplingmethod(
+        CoreML::Specification::SamplingMode::UPSAMPLE_MODE);
+  }
+  return layer_.release();
+}
+
+TfLiteStatus ResizeBilinearOpBuilder::RegisterInputs(
+    const TfLiteIntArray* inputs, TfLiteContext* context) {
+  if (inputs->size != 2) {
+    TF_LITE_KERNEL_LOG(context, "Wrong # of inputs to ResizeBilinear!.");
+    return kTfLiteError;
+  }
+  AddInput(inputs->data[0]);
+  TfLiteTensor* size = &context->tensors[inputs->data[1]];
+  height_ = GetTensorData<int32_t>(size)[0];
+  width_ = GetTensorData<int32_t>(size)[1];
+  return kTfLiteOk;
+}
+
+TfLiteStatus ResizeBilinearOpBuilder::RegisterOutputs(
+    const TfLiteIntArray* outputs, TfLiteContext* context) {
+  if (outputs->size != 1) {
+    TF_LITE_KERNEL_LOG(context, "Wrong # of outputs to ResizeBilinear!.");
+    return kTfLiteError;
+  }
+  graph_builder_->AddTensorWithID(outputs->data[0], GetOutput(context));
+  return kTfLiteOk;
+}
+
+OpBuilder* CreateResizeBilinearOpBuilder(GraphBuilder* graph_builder) {
+  return new ResizeBilinearOpBuilder(graph_builder);
+}
+
+bool IsResizeBilinearOpSupported(const TfLiteRegistration* registration,
+                                 const TfLiteNode* node,
+                                 TfLiteContext* context) {
+  if (node->builtin_data == nullptr) {
+    return false;
+  }
+  const int kOutputSize = 1;
+  if (!IsConstantTensor(GetInput(context, node, kOutputSize))) {
+    TF_LITE_KERNEL_LOG(context,
+                       "Output size of ResizeBilinear should be constant.");
+    return false;
+  }
+  return true;
+}
+
+}  // namespace coreml
+}  // namespace delegates
+}  // namespace tflite

tensorflow/lite/delegates/coreml/builders/resize_bilinear_op_builder.h

@@ -0,0 +1,54 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_DELEGATES_COREML_BUILDERS_RESIZE_BILINEAR_OP_BUILDER_H_
+#define TENSORFLOW_LITE_DELEGATES_COREML_BUILDERS_RESIZE_BILINEAR_OP_BUILDER_H_
+
+#include <string>
+
+#include "mlmodel/format/NeuralNetwork.pb.h"
+#include "tensorflow/lite/builtin_ops.h"
+#include "tensorflow/lite/c/c_api_types.h"
+#include "tensorflow/lite/c/common.h"
+#include "tensorflow/lite/delegates/coreml/builders/op_builder.h"
+
+namespace tflite {
+namespace delegates {
+namespace coreml {
+
+class ResizeBilinearOpBuilder : public OpBuilder {
+ public:
+  explicit ResizeBilinearOpBuilder(GraphBuilder* graph_builder)
+      : OpBuilder(graph_builder) {}
+
+  const std::string& DebugName() override;
+
+  CoreML::Specification::NeuralNetworkLayer* Build() override;
+
+  TfLiteStatus RegisterInputs(const TfLiteIntArray* inputs,
+                              TfLiteContext* context) override;
+
+  TfLiteStatus RegisterOutputs(const TfLiteIntArray* outputs,
+                               TfLiteContext* context) override;
+
+ private:
+  int height_;
+  int width_;
+};
+
+}  // namespace coreml
+}  // namespace delegates
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_DELEGATES_COREML_BUILDERS_RESIZE_BILINEAR_OP_BUILDER_H_

tensorflow/lite/delegates/coreml/builders/softmax_op_builder.cc

@@ -0,0 +1,72 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include "tensorflow/lite/delegates/coreml/builders/softmax_op_builder.h"
+
+#include <memory>
+#include <string>
+
+#include "mlmodel/format/NeuralNetwork.pb.h"
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/delegates/coreml/builders/op_builder.h"
+
+namespace tflite {
+namespace delegates {
+namespace coreml {
+const std::string& SoftmaxOpBuilder::DebugName() {
+  if (debug_name_.empty()) SetDebugName("SoftmaxOpBuilder", node_id_);
+  return debug_name_;
+}
+
+CoreML::Specification::NeuralNetworkLayer* SoftmaxOpBuilder::Build() {
+  if (layer_ == nullptr) {
+    layer_ = std::make_unique<CoreML::Specification::NeuralNetworkLayer>();
+  }
+  layer_->set_name(DebugName());
+  layer_->mutable_softmax();
+
+  return layer_.release();
+}
+
+TfLiteStatus SoftmaxOpBuilder::RegisterInputs(const TfLiteIntArray* inputs,
+                                              TfLiteContext* context) {
+  if (inputs->size != 1) {
+    TF_LITE_KERNEL_LOG(context, "Wrong # of inputs to softmax!.");
+    return kTfLiteError;
+  }
+  AddInput(inputs->data[0]);
+  return kTfLiteOk;
+}
+
+TfLiteStatus SoftmaxOpBuilder::RegisterOutputs(const TfLiteIntArray* outputs,
+                                               TfLiteContext* context) {
+  if (outputs->size != 1) {
+    TF_LITE_KERNEL_LOG(context, "Wrong # of outputs to softmax!.");
+    return kTfLiteError;
+  }
+  TensorID output_tensor = GetOutput(context);
+  if (output_tensor.NodeID() == -1) {
+    TF_LITE_KERNEL_LOG(context, "Failed to build output tensor.");
+    return kTfLiteError;
+  }
+  graph_builder_->AddTensorWithID(outputs->data[0], output_tensor);
+  return kTfLiteOk;
+}
+
+OpBuilder* CreateSoftmaxOpBuilder(GraphBuilder* graph_builder) {
+  return new SoftmaxOpBuilder(graph_builder);
+}
+}  // namespace coreml
+}  // namespace delegates
+}  // namespace tflite

tensorflow/lite/delegates/coreml/builders/softmax_op_builder.h

@@ -0,0 +1,48 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_DELEGATES_COREML_BUILDERS_SOFTMAX_OP_BUILDER_H_
+#define TENSORFLOW_LITE_DELEGATES_COREML_BUILDERS_SOFTMAX_OP_BUILDER_H_
+
+#include <string>
+
+#include "mlmodel/format/NeuralNetwork.pb.h"
+#include "tensorflow/lite/c/c_api_types.h"
+#include "tensorflow/lite/c/common.h"
+#include "tensorflow/lite/delegates/coreml/builders/op_builder.h"
+
+namespace tflite {
+namespace delegates {
+namespace coreml {
+// Builder for Softmax op in CoreML.
+class SoftmaxOpBuilder : public OpBuilder {
+ public:
+  explicit SoftmaxOpBuilder(GraphBuilder* graph_builder)
+      : OpBuilder(graph_builder) {}
+  const std::string& DebugName() override;
+
+  CoreML::Specification::NeuralNetworkLayer* Build() override;
+
+  TfLiteStatus RegisterInputs(const TfLiteIntArray* inputs,
+                              TfLiteContext* context) override;
+
+  TfLiteStatus RegisterOutputs(const TfLiteIntArray* outputs,
+                               TfLiteContext* context) override;
+};
+
+}  // namespace coreml
+}  // namespace delegates
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_DELEGATES_COREML_BUILDERS_SOFTMAX_OP_BUILDER_H_

tensorflow/lite/delegates/coreml/builders/test_util.h

@@ -0,0 +1,55 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_DELEGATES_COREML_BUILDERS_TEST_UTIL_H_
+#define TENSORFLOW_LITE_DELEGATES_COREML_BUILDERS_TEST_UTIL_H_
+
+#include "tensorflow/lite/delegates/coreml/coreml_delegate.h"
+#include "tensorflow/lite/kernels/test_util.h"
+
+#import <XCTest/XCTest.h>
+
+namespace tflite {
+namespace delegates {
+namespace coreml {
+class SingleOpModelWithCoreMlDelegate : public tflite::SingleOpModel {
+ public:
+  static const char kDelegateName[];
+
+  SingleOpModelWithCoreMlDelegate();
+  tflite::Interpreter* interpreter() { return interpreter_.get(); }
+
+ protected:
+  using SingleOpModel::builder_;
+
+ private:
+  TfLiteCoreMlDelegateOptions params_ = {
+      .enabled_devices = TfLiteCoreMlDelegateAllDevices,
+      .min_nodes_per_partition = 1,
+  };
+};
+
+}  // namespace coreml
+}  // namespace delegates
+}  // namespace tflite
+
+@interface BaseOpTest : XCTestCase
+@property tflite::delegates::coreml::SingleOpModelWithCoreMlDelegate* model;
+- (void)validateInterpreter:(tflite::Interpreter*)interpreter;
+- (void)checkInterpreterNotDelegated:(tflite::Interpreter*)interpreter;
+- (void)invokeAndValidate;
+- (void)invokeAndCheckNotDelegated;
+@end
+
+#endif  // TENSORFLOW_LITE_DELEGATES_COREML_BUILDERS_TEST_UTIL_H_

tensorflow/lite/delegates/coreml/builders/test_util.mm

@@ -0,0 +1,78 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include "tensorflow/lite/delegates/coreml/builders/test_util.h"
+
+namespace tflite {
+namespace delegates {
+namespace coreml {
+
+const char SingleOpModelWithCoreMlDelegate::kDelegateName[] = "TfLiteCoreMlDelegate";
+
+SingleOpModelWithCoreMlDelegate::SingleOpModelWithCoreMlDelegate() {
+  auto* delegate_ptr = TfLiteCoreMlDelegateCreate(&params_);
+  EXPECT_TRUE(delegate_ptr != nullptr);
+
+  // Note that tflite::SingleOpModel::BuildInterpreter(...) will apply the delegate that's set here
+  // to the model graph.
+  SetDelegate(tflite::Interpreter::TfLiteDelegatePtr(
+      delegate_ptr, [](TfLiteDelegate* delegate) { TfLiteCoreMlDelegateDelete(delegate); }));
+}
+
+}  // namespace coreml
+}  // namespace delegates
+}  // namespace tflite
+
+@implementation BaseOpTest
+- (void)validateInterpreter:(tflite::Interpreter*)interpreter {
+  // Make sure we have valid interpreter.
+  XCTAssertTrue(interpreter != nullptr);
+  // Make sure graph has one Op which is the delegate node.
+  XCTAssertEqual(interpreter->execution_plan().size(), 1);
+  const int node_index = interpreter->execution_plan()[0];
+  const auto* node_and_reg = interpreter->node_and_registration(node_index);
+  XCTAssertTrue(node_and_reg != nullptr);
+  XCTAssertTrue(node_and_reg->second.custom_name != nullptr);
+  XCTAssertTrue(
+      node_and_reg->second.custom_name ==
+      std::string(tflite::delegates::coreml::SingleOpModelWithCoreMlDelegate::kDelegateName));
+}
+
+- (void)checkInterpreterNotDelegated:(tflite::Interpreter*)interpreter {
+  // Make sure we have valid interpreter.
+  XCTAssertTrue(interpreter != nullptr);
+  for (int node_idx : interpreter->execution_plan()) {
+    // Make sure no node is delegated.
+    XCTAssertEqual(interpreter->execution_plan().size(), 1);
+    const auto* node_and_reg = interpreter->node_and_registration(node_idx);
+    XCTAssertTrue(node_and_reg != nullptr);
+    if (node_and_reg->second.custom_name != nullptr) {
+      XCTAssertTrue(
+          node_and_reg->second.custom_name !=
+          std::string(tflite::delegates::coreml::SingleOpModelWithCoreMlDelegate::kDelegateName));
+    }
+  }
+}
+
+- (void)invokeAndValidate {
+  _model->Invoke();
+  [self validateInterpreter:_model->interpreter()];
+}
+
+- (void)invokeAndCheckNotDelegated {
+  _model->Invoke();
+  [self checkInterpreterNotDelegated:_model->interpreter()];
+}
+
+@end

tensorflow/lite/delegates/coreml/builders/threshold_layer_builder.cc

@@ -0,0 +1,71 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include "tensorflow/lite/delegates/coreml/builders/threshold_layer_builder.h"
+
+#include <memory>
+#include <string>
+
+#include "mlmodel/format/NeuralNetwork.pb.h"
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/delegates/coreml/builders/op_builder.h"
+
+namespace tflite {
+namespace delegates {
+namespace coreml {
+
+const std::string& ThresholdLayerBuilder::DebugName() {
+  if (debug_name_.empty()) SetDebugName("ThresholdLayerBuilder", node_id_);
+  return debug_name_;
+}
+
+CoreML::Specification::NeuralNetworkLayer* ThresholdLayerBuilder::Build() {
+  if (layer_ == nullptr) {
+    layer_ = std::make_unique<CoreML::Specification::NeuralNetworkLayer>();
+  }
+  layer_->set_name(DebugName());
+  layer_->mutable_unary()->set_alpha(alpha_);
+  layer_->mutable_unary()->set_scale(scale_);
+  layer_->mutable_unary()->set_type(
+      CoreML::Specification::UnaryFunctionLayerParams::THRESHOLD);
+  return layer_.release();
+}
+
+TfLiteStatus ThresholdLayerBuilder::RegisterInputs(const TfLiteIntArray* inputs,
+                                                   TfLiteContext* context) {
+  if (inputs->size != 1) {
+    TF_LITE_KERNEL_LOG(context, "Threshold: Wrong # of inputs!.");
+    return kTfLiteError;
+  }
+  AddInput(inputs->data[0]);
+  return kTfLiteOk;
+}
+
+TfLiteStatus ThresholdLayerBuilder::RegisterOutputs(
+    const TfLiteIntArray* outputs, TfLiteContext* context) {
+  if (outputs->size != 1) {
+    TF_LITE_KERNEL_LOG(context, "Threshold: Wrong # of outputs!.");
+    return kTfLiteError;
+  }
+  graph_builder_->AddTensorWithID(outputs->data[0], GetOutput(context));
+  return kTfLiteOk;
+}
+
+OpBuilder* CreateThresholdLayerBuilder(GraphBuilder* graph_builder) {
+  return new ThresholdLayerBuilder(graph_builder);
+}
+
+}  // namespace coreml
+}  // namespace delegates
+}  // namespace tflite

tensorflow/lite/delegates/coreml/builders/threshold_layer_builder.h

@@ -0,0 +1,61 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_DELEGATES_COREML_BUILDERS_THRESHOLD_LAYER_BUILDER_H_
+#define TENSORFLOW_LITE_DELEGATES_COREML_BUILDERS_THRESHOLD_LAYER_BUILDER_H_
+
+#include <string>
+
+#include "mlmodel/format/NeuralNetwork.pb.h"
+#include "tensorflow/lite/c/c_api_types.h"
+#include "tensorflow/lite/c/common.h"
+#include "tensorflow/lite/delegates/coreml/builders/op_builder.h"
+
+namespace tflite {
+namespace delegates {
+namespace coreml {
+
+// Layer that provides threshold operation. Depending on scale, this can be used
+// as max (scale > 0) or min (scale < 0), in combination with another negative
+// linear activation layer) operation.
+// TODO(karimnosseir): Generalize to other unary operators.
+class ThresholdLayerBuilder : public OpBuilder {
+ public:
+  explicit ThresholdLayerBuilder(GraphBuilder* graph_builder)
+      : OpBuilder(graph_builder) {}
+
+  const std::string& DebugName() override;
+
+  CoreML::Specification::NeuralNetworkLayer* Build() override;
+
+  void SetAlpha(float alpha) { alpha_ = alpha; }
+
+  void SetScale(float scale) { scale_ = scale; }
+
+  TfLiteStatus RegisterInputs(const TfLiteIntArray* inputs,
+                              TfLiteContext* context) override;
+
+  TfLiteStatus RegisterOutputs(const TfLiteIntArray* outputs,
+                               TfLiteContext* context) override;
+
+ private:
+  float alpha_ = 0.0f;
+  float scale_ = 1.0f;
+};
+
+}  // namespace coreml
+}  // namespace delegates
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_DELEGATES_COREML_BUILDERS_THRESHOLD_LAYER_BUILDER_H_

tensorflow/lite/delegates/coreml/builders/util.cc

@@ -0,0 +1,110 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include "tensorflow/lite/delegates/coreml/builders/util.h"
+
+#include <cstdint>
+#include <vector>
+
+#include "fp16.h"  // from @FP16
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
+#include "tensorflow/lite/kernels/kernel_util.h"
+
+namespace tflite {
+namespace delegates {
+namespace coreml {
+namespace {
+void Get4DShape(const TfLiteTensor* tensor, std::vector<int>* shape) {
+  const int rank = tensor->dims->size;
+  shape->resize(4);
+  for (int i = 0; i < 4; ++i) {
+    if (i < 4 - rank) {
+      (*shape)[i] = 1;
+    } else {
+      (*shape)[i] = tensor->dims->data[i - (4 - rank)];
+    }
+  }
+}
+}  // namespace
+
+// Determines if two tensor shapes are broadcastable. See comment of
+// IsBinaryOpSupported for more info.
+bool IsBroadcastable(const TfLiteTensor* input_0, const TfLiteTensor* input_1) {
+  if (input_0->dims->size > 4 || input_1->dims->size > 4) {
+    return false;
+  }
+  std::vector<int> shape_0;
+  std::vector<int> shape_1;
+  Get4DShape(input_0, &shape_0);
+  Get4DShape(input_1, &shape_1);
+  const int B_0 = shape_0[0];
+  const int B_1 = shape_1[0];
+  const int H_0 = shape_0[1];
+  const int H_1 = shape_1[1];
+  const int W_0 = shape_0[2];
+  const int W_1 = shape_1[2];
+  const int C_0 = shape_0[3];
+  const int C_1 = shape_1[3];
+
+  // TFL tensor has [B, H, W, C] format.
+  // comparing B: shape[0], (H, W): (shape[1], shape[2]), C: shape[3].
+
+  // When B is different, it's not supported unless
+  // one of the tensor is size 1 constant tensor.
+  if (B_0 != B_1) {
+    if (!((IsConstantTensor(input_0) && NumElements(input_0) == 1) ||
+          (IsConstantTensor(input_1) && NumElements(input_1) == 1)))
+      return false;
+  }
+
+  // When (H, W) are different, one of the (H, W) should be (1, 1).
+  if (H_0 != H_1 || W_0 != W_1) {
+    if (!((H_0 == 1 && W_0 == 1) || (H_1 == 1 && W_1 == 1))) {
+      return false;
+    }
+  }
+
+  // When C is different, one of the C should be 1.
+  if (C_0 != C_1) {
+    if (C_0 != 1 && C_1 != 1) return false;
+  }
+  return true;
+}
+
+bool IsBinaryOpSupportedType(const TfLiteTensor* tensor) {
+  return tensor->type == kTfLiteFloat32 ||
+         (tensor->type == kTfLiteFloat16 && IsConstantTensor(tensor));
+}
+
+bool IsBinaryOpSupported(const TfLiteRegistration* registration,
+                         const TfLiteNode* node, TfLiteContext* context) {
+  const auto* input_0 = GetInput(context, node, 0);
+  const auto* input_1 = GetInput(context, node, 1);
+  if (IsBinaryOpSupportedType(input_0) && IsBinaryOpSupportedType(input_1)) {
+    return IsBroadcastable(input_0, input_1);
+  }
+  return false;
+}
+
+float GetScalarFloatFromTensor(const TfLiteTensor* tensor) {
+  if (tensor->type == kTfLiteFloat16) {
+    return fp16_ieee_to_fp32_value(GetTensorData<uint16_t>(tensor)[0]);
+  }
+  return GetTensorData<float>(tensor)[0];
+}
+
+}  // namespace coreml
+}  // namespace delegates
+}  // namespace tflite

tensorflow/lite/delegates/coreml/builders/util.h

@@ -0,0 +1,40 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_DELEGATES_COREML_BUILDERS_UTIL_H_
+#define TENSORFLOW_LITE_DELEGATES_COREML_BUILDERS_UTIL_H_
+
+#include "tensorflow/lite/core/c/common.h"
+
+namespace tflite {
+namespace delegates {
+namespace coreml {
+
+// Checks if Binary ops have supported broadcastable shapes.
+// Core ml arithmetic ops - Add and Mul support broadcasts among
+// [B, 1, 1, 1], [B, C, 1, 1], [B, 1, H, W], [B, C, H, W].
+// other shapes should be rejected. Unless it is a constant tensor of size 1,
+// which will be added as data.
+
+bool IsBinaryOpSupported(const TfLiteRegistration* registration,
+                         const TfLiteNode* node, TfLiteContext* context);
+
+// Gets the float scalar value from the given tensor. The tensor should be a
+// constant float32/float16 tensor of size 1.
+float GetScalarFloatFromTensor(const TfLiteTensor* tensor);
+
+}  // namespace coreml
+}  // namespace delegates
+}  // namespace tflite
+#endif  // TENSORFLOW_LITE_DELEGATES_COREML_BUILDERS_UTIL_H_

tensorflow/lite/delegates/coreml/builders/util_test.cc

@@ -0,0 +1,156 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include "tensorflow/lite/delegates/coreml/builders/util.h"
+
+#include <algorithm>
+#include <cstdint>
+#include <vector>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "Eigen/Core"  // from @eigen_archive
+#include "tensorflow/lite/core/c/common.h"
+
+using tflite::delegates::coreml::IsBinaryOpSupported;
+
+namespace {
+
+class IsBinaryOpSupportedTest : public testing::Test {
+ protected:
+  void SetUp() override {
+    const int input_size = 2;
+    tensors_.resize(input_size);
+    context_.tensors = tensors_.data();
+    node_.inputs = TfLiteIntArrayCreate(input_size);
+    for (int i = 0; i < input_size; ++i) {
+      node_.inputs->data[i] = i;
+    }
+
+    for (auto& tensor : tensors_) {
+      tensor.allocation_type = kTfLiteArenaRw;
+      tensor.type = kTfLiteFloat32;
+      tensor.dims = nullptr;
+    }
+  }
+
+  void TearDown() override {
+    FreeInputShapes();
+    TfLiteIntArrayFree(node_.inputs);
+  }
+
+  void SetInputShapes(const std::vector<std::vector<int>>& shapes) {
+    for (int i = 0; i < tensors_.size(); ++i) {
+      tensors_[i].dims = TfLiteIntArrayCreate(shapes[i].size());
+      std::copy(shapes[i].begin(), shapes[i].end(), tensors_[i].dims->data);
+    }
+  }
+
+  void FreeInputShapes() {
+    for (auto& tensor : tensors_) {
+      if (tensor.dims != nullptr) {
+        TfLiteIntArrayFree(tensor.dims);
+        tensor.dims = nullptr;
+      }
+    }
+  }
+
+  TfLiteContext context_ = {};
+  TfLiteNode node_;
+  std::vector<TfLiteTensor> tensors_;
+};
+
+TEST_F(IsBinaryOpSupportedTest, BroadcastTest) {
+  std::vector<int> base_shape = {2, 2, 3, 4};
+  std::vector<std::vector<int>> shapes = {
+      {2, 2, 3, 4}, {2, 1, 1, 4}, {2, 2, 3, 1}, {2, 1, 1, 1}};
+  std::vector<TfLiteTensor> inputs(2);
+  for (const auto& shape : shapes) {
+    SetInputShapes({base_shape, shape});
+    ASSERT_TRUE(IsBinaryOpSupported(nullptr, &node_, &context_));
+    FreeInputShapes();
+  }
+}
+
+TEST_F(IsBinaryOpSupportedTest, LessThan4DTest) {
+  std::vector<int> base_shape = {1, 2, 3, 4};
+  std::vector<std::vector<int>> shapes = {{4}, {2, 3, 1}, {1, 1, 1, 1}};
+  for (const auto& shape : shapes) {
+    SetInputShapes({base_shape, shape});
+    ASSERT_TRUE(IsBinaryOpSupported(nullptr, &node_, &context_));
+    FreeInputShapes();
+  }
+}
+
+TEST_F(IsBinaryOpSupportedTest, ConstScalarTest) {
+  std::vector<int> base_shape = {2, 2, 3, 4};
+  tensors_[1].allocation_type = kTfLiteMmapRo;
+  SetInputShapes({base_shape, {1}});
+  ASSERT_TRUE(IsBinaryOpSupported(nullptr, &node_, &context_));
+  FreeInputShapes();
+}
+
+TEST_F(IsBinaryOpSupportedTest, NotSupportedBroadcastTest) {
+  std::vector<int> base_shape = {2, 2, 3, 4};
+  std::vector<std::vector<int>> shapes = {
+      {2, 2, 1, 4}, {2, 1, 2, 4}, {1, 2, 3, 1}, {1, 1, 1, 1}};
+  for (const auto& shape : shapes) {
+    SetInputShapes({base_shape, shape});
+    ASSERT_FALSE(IsBinaryOpSupported(nullptr, &node_, &context_));
+    FreeInputShapes();
+  }
+}
+
+TEST_F(IsBinaryOpSupportedTest, RankGreaterThan4Test) {
+  std::vector<int> base_shape = {1, 2, 3, 4};
+  std::vector<std::vector<int>> shapes = {{1, 2, 3, 4, 5}, {2, 1, 1, 4}};
+  SetInputShapes({base_shape, shapes[0]});
+  ASSERT_FALSE(IsBinaryOpSupported(nullptr, &node_, &context_));
+  FreeInputShapes();
+
+  SetInputShapes({shapes[0], base_shape});
+  ASSERT_FALSE(IsBinaryOpSupported(nullptr, &node_, &context_));
+  FreeInputShapes();
+}
+
+TEST_F(IsBinaryOpSupportedTest, SupportConstFloat16InputTest) {
+  tensors_[0].allocation_type = kTfLiteMmapRo;
+  tensors_[0].type = kTfLiteFloat16;
+  SetInputShapes({{2, 2, 3}, {2, 2, 3}});
+  ASSERT_TRUE(IsBinaryOpSupported(nullptr, &node_, &context_));
+}
+
+TEST_F(IsBinaryOpSupportedTest, NotSupportNonConstFloat16InputTest) {
+  tensors_[0].allocation_type = kTfLiteArenaRw;
+  tensors_[0].type = kTfLiteFloat16;
+  SetInputShapes({{2, 2, 3}, {2, 2, 3}});
+  ASSERT_FALSE(IsBinaryOpSupported(nullptr, &node_, &context_));
+}
+
+TEST(UtilTest, GetScalarFloatFromTensorTest) {
+  std::vector<Eigen::half> half{Eigen::half{-535.54f}};
+  std::vector<float> float32{-535.54f};
+  TfLiteTensor tensor;
+  tensor.type = kTfLiteFloat16;
+  tensor.data.data = reinterpret_cast<uint16_t*>(half.data());
+  EXPECT_THAT(tflite::delegates::coreml::GetScalarFloatFromTensor(&tensor),
+              testing::FloatNear(-535.54f, 0.1f));
+
+  tensor.type = kTfLiteFloat32;
+  tensor.data.data = float32.data();
+  EXPECT_THAT(tflite::delegates::coreml::GetScalarFloatFromTensor(&tensor),
+              testing::FloatNear(-535.54f, 0.f));
+}
+
+}  // namespace

tensorflow/lite/delegates/coreml/coreml_delegate.h

@@ -0,0 +1,72 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_DELEGATES_COREML_COREML_DELEGATE_H_
+#define TENSORFLOW_LITE_DELEGATES_COREML_COREML_DELEGATE_H_
+
+#include "tensorflow/lite/core/c/common.h"
+
+// LINT.IfChange
+
+#ifdef __cplusplus
+extern "C" {
+#endif  // __cplusplus
+typedef enum {
+  // Create Core ML delegate only on devices with Apple Neural Engine.
+  // Returns nullptr otherwise.
+  TfLiteCoreMlDelegateDevicesWithNeuralEngine,
+  // Always create Core ML delegate
+  TfLiteCoreMlDelegateAllDevices
+} TfLiteCoreMlDelegateEnabledDevices;
+
+typedef struct {
+  // Only create delegate when Neural Engine is available on the device.
+  TfLiteCoreMlDelegateEnabledDevices enabled_devices;
+  // Specifies target Core ML version for model conversion.
+  // Core ML 3 come with a lot more ops, but some ops (e.g. reshape) is not
+  // delegated due to input rank constraint.
+  // if not set to one of the valid versions, the delegate will use highest
+  // version possible in the platform.
+  // Valid versions: (2, 3)
+  int coreml_version;
+  // This sets the maximum number of Core ML delegates created.
+  // Each graph corresponds to one delegated node subset in the
+  // TFLite model. Set this to 0 to delegate all possible partitions.
+  int max_delegated_partitions;
+  // This sets the minimum number of nodes per partition delegated with
+  // Core ML delegate. Defaults to 2.
+  int min_nodes_per_partition;
+#ifdef TFLITE_DEBUG_DELEGATE
+  // This sets the index of the first node that could be delegated.
+  int first_delegate_node_index;
+  // This sets the index of the last node that could be delegated.
+  int last_delegate_node_index;
+#endif
+} TfLiteCoreMlDelegateOptions;
+
+// Return a delegate that uses CoreML for ops execution.
+// Must outlive the interpreter.
+TfLiteDelegate* TfLiteCoreMlDelegateCreate(
+    const TfLiteCoreMlDelegateOptions* options);
+
+// Do any needed cleanup and delete 'delegate'.
+void TfLiteCoreMlDelegateDelete(TfLiteDelegate* delegate);
+
+#ifdef __cplusplus
+}
+#endif  // __cplusplus
+
+// LINT.ThenChange(README.md)
+
+#endif  // TENSORFLOW_LITE_DELEGATES_COREML_COREML_DELEGATE_H_

tensorflow/lite/delegates/coreml/coreml_delegate.mm

@@ -0,0 +1,354 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include "tensorflow/lite/delegates/coreml/coreml_delegate.h"
+
+#include <string.h>
+#include <sys/utsname.h>
+#include <limits>
+#include <vector>
+
+#include "tensorflow/lite/builtin_ops.h"
+#include "tensorflow/lite/context_util.h"
+#include "tensorflow/lite/core/c/builtin_op_data.h"
+#include "tensorflow/lite/delegates/coreml/builders/op_validator.h"
+#include "tensorflow/lite/delegates/coreml/builders/util.h"
+#include "tensorflow/lite/delegates/coreml/coreml_delegate_kernel.h"
+#include "tensorflow/lite/delegates/utils.h"
+#include "tensorflow/lite/kernels/kernel_util.h"
+#include "tensorflow/lite/minimal_logging.h"
+
+namespace tflite {
+namespace {
+constexpr int kMinNodesPerCoreMlDelegate = 2;
+
+using delegates::coreml::CoreMlDelegateKernel;
+
+bool IsNodeSupportedByDelegate(const TfLiteRegistration* registration, const TfLiteNode* node,
+                               TfLiteContext* context, const TfLiteCoreMlDelegateOptions* options) {
+  if (@available(iOS 11.0, *)) {
+  } else {
+    return false;
+  }
+
+  // For most ops, only version 1 is supported.
+  if (registration->version > 1) {
+    switch (registration->builtin_code) {
+      case kTfLiteBuiltinDepthwiseConv2d:
+        if (registration->version > 2) return false;
+        break;
+      // FullyConnected without bias is supported starting from version 6.
+      case kTfLiteBuiltinFullyConnected:
+        if (registration->version > 6) return false;
+        break;
+      default:
+        return false;
+    }
+  }
+
+  // The model should not be full-integer quantized. For ops supported by Core ML delegate,
+  // Testing if the first input is float is sufficient to filter full-integer quantized ops.
+  int input_tensor_index = 0;
+  // TransposeConv input: (output_shape, filters, input)
+  if (registration->builtin_code == kTfLiteBuiltinTransposeConv) {
+    input_tensor_index = 2;
+  }
+  if (GetInput(context, node, input_tensor_index)->type != kTfLiteFloat32) {
+    return false;
+  }
+
+  // TODO(b/149179044): Add extra validation if this is not sufficient.
+
+  // TODO(karimnossier): Refactor this function.
+  // TODO(karimnosseir): Add
+  // 1) Checks for versioning.
+  // 2) Checks for input constraints.
+  // Follow the ordering of TfLiteBuiltinOperator enum.
+  switch (registration->builtin_code) {
+    case kTfLiteBuiltinAdd: {
+      return node->builtin_data != nullptr &&
+             delegates::coreml::IsBinaryOpSupported(registration, node, context);
+    }
+    case kTfLiteBuiltinAveragePool2d: {
+      const auto* params = reinterpret_cast<const TfLitePoolParams*>(node->builtin_data);
+      return params != nullptr && params->activation == kTfLiteActNone;
+    }
+    case kTfLiteBuiltinConcatenation: {
+      return delegates::coreml::IsConcatenationOpSupported(registration, node, context);
+    }
+    case kTfLiteBuiltinConv2d: {
+      return delegates::coreml::IsConvolutionOpSupported(registration, node, context);
+    }
+    case kTfLiteBuiltinDepthwiseConv2d: {
+      return delegates::coreml::IsDepthwiseConvolutionOpSupported(registration, node, context);
+    }
+    case kTfLiteBuiltinFullyConnected: {
+      return delegates::coreml::IsFullyConnectedOpSupported(registration, node, context);
+    }
+    case kTfLiteBuiltinHardSwish: {
+      return true;
+    }
+    case kTfLiteBuiltinLogistic: {
+      return true;
+    }
+    case kTfLiteBuiltinMaxPool2d: {
+      const auto* params = reinterpret_cast<const TfLitePoolParams*>(node->builtin_data);
+      return params != nullptr && params->activation == kTfLiteActNone;
+    }
+    case kTfLiteBuiltinMirrorPad: {
+      return delegates::coreml::IsMirrorPadOpSupported(registration, node, context);
+    }
+    case kTfLiteBuiltinMean: {
+      return delegates::coreml::IsMeanOpSupported(registration, node, context);
+    }
+    case kTfLiteBuiltinMul: {
+      return node->builtin_data != nullptr &&
+             delegates::coreml::IsBinaryOpSupported(registration, node, context);
+    }
+    case kTfLiteBuiltinPad:
+    case kTfLiteBuiltinPadv2: {
+      return delegates::coreml::IsPadOpSupported(registration, node, context);
+    }
+    case kTfLiteBuiltinRelu: {
+      return true;
+    }
+    case kTfLiteBuiltinReluN1To1: {
+      return true;
+    }
+    case kTfLiteBuiltinRelu6: {
+      return true;
+    }
+    case kTfLiteBuiltinReshape: {
+      return delegates::coreml::IsReshapeOpSupported(registration, node, context,
+                                                     options->coreml_version);
+    }
+    case kTfLiteBuiltinResizeBilinear: {
+      return delegates::coreml::IsResizeBilinearOpSupported(registration, node, context);
+    }
+    case kTfLiteBuiltinSoftmax: {
+      // Only supports when beta is 1.0 for now.
+      const auto* softmax_params = reinterpret_cast<const TfLiteSoftmaxParams*>(node->builtin_data);
+      return softmax_params != nullptr && softmax_params->beta == 1.0;
+    }
+    case kTfLiteBuiltinTanh: {
+      return true;
+    }
+    case kTfLiteBuiltinTransposeConv: {
+      return delegates::coreml::IsTransposeConvolutionOpSupported(registration, node, context);
+    }
+    default:
+      return false;
+  }
+  return false;
+}
+
+class CoreMlDelegate : public TfLiteDelegate {
+ public:
+  explicit CoreMlDelegate(const TfLiteCoreMlDelegateOptions* params)
+      : params_(params != nullptr ? *params : TfLiteCoreMlDelegateOptions()) {
+    {
+      if (@available(iOS 13.0, *)) {
+        if (params_.coreml_version != 2 && params_.coreml_version != 3) {
+          NSLog(@"coreml_version must be 2 or 3. Setting to 3.");
+          params_.coreml_version = 3;
+        }
+      } else if (@available(iOS 12.0, *)) {
+        if (params_.coreml_version != 2) {
+          NSLog(@"coreml_version must be 2 - using Core ML version 2.");
+          params_.coreml_version = 2;
+        }
+      }
+      if (params_.max_delegated_partitions <= 0) {
+        params_.max_delegated_partitions = std::numeric_limits<int>::max();
+      }
+      if (params_.min_nodes_per_partition <= 0) {
+        params_.min_nodes_per_partition = kMinNodesPerCoreMlDelegate;
+      }
+#ifdef TFLITE_DEBUG_DELEGATE
+      if (params_.first_delegate_node_index < 0) {
+        params_.first_delegate_node_index = 0;
+      }
+      if (params->last_delegate_node_index <= 0) {
+        params_.last_delegate_node_index = std::numeric_limits<int>::max();
+      }
+#endif
+    }
+  }
+
+  TfLiteCoreMlDelegateOptions* params() { return &params_; }
+
+  bool VerifyDelegate() { return true; }
+
+ private:
+  TfLiteCoreMlDelegateOptions params_;
+};
+
+TfLiteRegistration GetCoreMlKernelRegistration() {
+  // 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 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 = "TfLiteCoreMlDelegate";
+  kernel_registration.free = [](TfLiteContext* context, void* buffer) -> void {
+    delete reinterpret_cast<CoreMlDelegateKernel*>(buffer);
+  };
+  kernel_registration.init = [](TfLiteContext* context, const char* buffer,
+                                size_t length) -> void* {
+    const auto* params = reinterpret_cast<const TfLiteDelegateParams*>(buffer);
+    const auto* coreml_options = (reinterpret_cast<CoreMlDelegate*>(params->delegate))->params();
+    CoreMlDelegateKernel* coreml_kernel = new CoreMlDelegateKernel(coreml_options->coreml_version);
+    if (coreml_kernel->Init(context, params) != kTfLiteOk) {
+      delete coreml_kernel;
+      return nullptr;
+    }
+    return coreml_kernel;
+  };
+  kernel_registration.invoke = [](TfLiteContext* context, TfLiteNode* node) -> TfLiteStatus {
+    CoreMlDelegateKernel* kernel = reinterpret_cast<CoreMlDelegateKernel*>(node->user_data);
+    if (!kernel) {
+      TF_LITE_KERNEL_LOG(context, "CoreMl Kernel was not initialized");
+      return kTfLiteError;
+    }
+    return kernel->Invoke(context, node);
+  };
+  kernel_registration.prepare = [](TfLiteContext* context, TfLiteNode* node) -> TfLiteStatus {
+    CoreMlDelegateKernel* kernel = reinterpret_cast<CoreMlDelegateKernel*>(node->user_data);
+    if (kernel == nullptr) {
+      TF_LITE_KERNEL_LOG(context, "CoreMl Kernel was not initialized");
+      return kTfLiteError;
+    }
+    return kernel->Prepare(context, node);
+  };
+
+  return kernel_registration;
+}
+
+TfLiteStatus DelegatePrepare(TfLiteContext* context, TfLiteDelegate* delegate) {
+  const auto* params = reinterpret_cast<TfLiteCoreMlDelegateOptions*>(delegate->data_);
+
+  delegates::IsNodeSupportedFn node_supported_fn = [=](TfLiteContext* context, TfLiteNode* node,
+                                                       TfLiteRegistration* registration,
+                                                       std::string* unsupported_details) -> bool {
+    return IsNodeSupportedByDelegate(registration, node, context, params);
+  };
+
+  delegates::FP16GraphPartitionHelper partition_helper(context, node_supported_fn);
+#ifndef TFLITE_DEBUG_DELEGATE
+  TF_LITE_ENSURE_STATUS(partition_helper.Partition(nullptr));
+#else
+  TF_LITE_ENSURE_STATUS(partition_helper.Partition(nullptr, params->first_delegate_node_index,
+                                                   params->last_delegate_node_index));
+#endif
+
+  std::vector<int> delegated_nodes = partition_helper.GetNodesOfFirstNLargestPartitions(
+      params->max_delegated_partitions, params->min_nodes_per_partition);
+  TFLITE_LOG_PROD(tflite::TFLITE_LOG_INFO,
+                  "CoreML delegate: %d nodes delegated out of %d nodes, "
+                  "with %d partitions.\n",
+                  delegated_nodes.size(), partition_helper.num_total_nodes(),
+                  partition_helper.num_partitions());
+  return context->ReplaceNodeSubsetsWithDelegateKernels(
+      context, GetCoreMlKernelRegistration(), BuildTfLiteArray(delegated_nodes).get(), delegate);
+}
+
+TfLiteDelegate* CreateCoreMlDelegate(const TfLiteCoreMlDelegateOptions* options) {
+  TfLiteDelegate* delegate = new CoreMlDelegate(options);
+  if (!static_cast<CoreMlDelegate*>(delegate)->VerifyDelegate()) {
+    delete delegate;
+    return nullptr;
+  }
+
+  delegate->data_ = static_cast<tflite::CoreMlDelegate*>(delegate)->params();
+  delegate->flags = kTfLiteDelegateFlagsNone;
+  delegate->Prepare = &DelegatePrepare;
+  delegate->CopyFromBufferHandle = nullptr;
+  delegate->CopyToBufferHandle = nullptr;
+  delegate->FreeBufferHandle = nullptr;
+
+  return delegate;
+}
+}  // namespace
+}  // namespace tflite
+
+namespace {
+
+bool IsNeuralEngineAvailable() {
+  if (@available(iOS 17.0, macOS 14.0, watchOS 10.0, tvOS 17.0, *)) {
+    // This is a safe way to check if Neural Engine is available.
+    NSArray<id<MLComputeDeviceProtocol>>* devices = MLAllComputeDevices();
+    Class cls = [MLNeuralEngineComputeDevice class];
+    for (id<MLComputeDeviceProtocol> device in devices) {
+      if ([device isKindOfClass:cls]) {
+        return true;
+      }
+    }
+    return false;
+  } else {
+#if TARGET_OS_IPHONE
+    // utsname.machine has device identifier. ex: the identifier for iPhone Xs is "iPhone11,2".
+    // Since Neural Engine is only available for use on A12 and later, major device version in the
+    // identifier is checked for these models:
+    // A12: iPhone XS (11,2), iPad Mini - 5th Gen (11,1)
+    // A12X: iPad Pro - 3rd Gen (8,1)
+    // For more information, see https://www.theiphonewiki.com/wiki/Models
+    struct utsname system_info;
+    uname(&system_info);
+    if (strncmp("iPad", system_info.machine, 4) == 0) {
+      const int major_version = atoi(system_info.machine + 4);
+      return major_version >= 8;  // There are no device between iPad 8 and 11.
+    } else if (strncmp("iPhone", system_info.machine, 6) == 0) {
+      const int major_version = atoi(system_info.machine + 6);
+      return major_version >= 11;
+    }
+#elif TARGET_OS_OSX
+#if defined(__aarch64__)
+    // All Apple Silicon based Macs have Neural Engine.`
+    return true;
+#else
+    // Intel based Macs do not have Neural Engine.
+    return false;
+#endif  // defined(__arm__)
+#else
+#error "Unsupported platform. Not recognized as iPhone or Mac."
+#endif  // TARGET_OS_IPHONE
+  }
+  return false;
+}
+
+}  // namespace
+
+TfLiteDelegate* TfLiteCoreMlDelegateCreate(const TfLiteCoreMlDelegateOptions* options) {
+  if (@available(iOS 12.0, macOS 10.13, *)) {
+    if (options->enabled_devices == TfLiteCoreMlDelegateDevicesWithNeuralEngine &&
+        !IsNeuralEngineAvailable()) {
+      NSLog(@"This device does not have Neural Engine, so Core ML delegate will not be enabled. "
+             "If you want to run Core ML delegate anyway, set enabled_devices option to "
+             "TfLiteCoreMlDelegateAllDevices (or enabledDevices to .allDevices in Swift).");
+      return nullptr;
+    }
+    return tflite::CreateCoreMlDelegate(options);
+  } else {
+    NSLog(@"Core ML delegate is not supported in this iOS version. "
+           "Minimum required iOS/macOS version is 12.0/10.13.");
+    return nullptr;
+  }
+}
+
+void TfLiteCoreMlDelegateDelete(TfLiteDelegate* delegate) { delete delegate; }

tensorflow/lite/delegates/coreml/coreml_delegate_kernel.h

@@ -0,0 +1,79 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_DELEGATES_COREML_COREML_DELEGATE_KERNEL_H_
+#define TENSORFLOW_LITE_DELEGATES_COREML_COREML_DELEGATE_KERNEL_H_
+
+#include <memory>
+#include <vector>
+
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/delegates/coreml/builders/op_builder.h"
+#import "tensorflow/lite/delegates/coreml/coreml_executor.h"
+
+namespace tflite {
+namespace delegates {
+namespace coreml {
+
+// Represents a subgraph in TFLite that will be delegated to CoreML.
+// It is abstracted as a single kernel node in the main TFLite graph and
+// implements Init/Prepare/Invoke as TFLite kernel nodes.
+class CoreMlDelegateKernel {
+ public:
+  explicit CoreMlDelegateKernel(int coreml_version)
+      : coreml_version_(coreml_version) {}
+
+  CoreMlDelegateKernel(const CoreMlDelegateKernel&) = delete;
+  CoreMlDelegateKernel& operator=(const CoreMlDelegateKernel&) = delete;
+
+  // Initialize the delegated graph and add required nodes.
+  TfLiteStatus Init(TfLiteContext* context,
+                    const TfLiteDelegateParams* delegate_params);
+
+  // Any preparation work needed for the delegated graph.
+  TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node);
+
+  // Allocates delegated tensordefs for graph I/O & execute it.
+  TfLiteStatus Invoke(TfLiteContext* context, TfLiteNode* node);
+
+  ~CoreMlDelegateKernel();
+
+ private:
+  // Builds the ML Model protocol buffer
+  TfLiteStatus BuildModel(TfLiteContext* context,
+                          const TfLiteDelegateParams* delegate_params);
+
+  // Adds the output tensors to the model generated.
+  TfLiteStatus AddOutputTensors(const TfLiteIntArray* output_tensors,
+                                TfLiteContext* context);
+
+  // Adds the input tensors to the model generated.
+  void AddInputTensors(const TfLiteIntArray* input_tensors,
+                       TfLiteContext* context);
+
+  std::unique_ptr<delegates::coreml::GraphBuilder> builder_;
+  std::unique_ptr<CoreML::Specification::Model> model_;
+  ::CoreMlExecutor* executor_ = nullptr;
+  int coreml_version_;
+
+  std::vector<int> input_tensor_ids_;
+  std::vector<TensorData> inputs_;
+  std::vector<TensorData> outputs_;
+};
+
+}  // namespace coreml
+}  // namespace delegates
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_DELEGATES_COREML_COREML_DELEGATE_KERNEL_H_

tensorflow/lite/delegates/coreml/coreml_delegate_kernel.mm

@@ -0,0 +1,269 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include "tensorflow/lite/delegates/coreml/coreml_delegate_kernel.h"
+
+#include "tensorflow/lite/context_util.h"
+#include "tensorflow/lite/types/half.h"
+#include "tensorflow/lite/kernels/internal/optimized/optimized_ops.h"
+#include "tensorflow/lite/kernels/internal/types.h"
+#include "tensorflow/lite/kernels/kernel_util.h"
+
+#import "tensorflow/lite/delegates/coreml/coreml_executor.h"
+
+namespace tflite {
+namespace delegates {
+namespace coreml {
+namespace {
+// TODO(karimnosseir): Move to util library
+TfLiteStatus GetDims(int* batch_size, int* height_size, int* width_size, int* depth_size,
+                     const TfLiteIntArray* dims) {
+  if (dims == nullptr || dims->size > 4) {
+    return kTfLiteError;
+  }
+  int* dim[] = {batch_size, height_size, width_size, depth_size};
+  for (int i = 0; i < 4; ++i) *(dim[i]) = 1;
+  const int offset = 4 - dims->size;
+  for (int i = 0; i < dims->size; ++i) {
+    *dim[offset + i] = dims->data[i];
+  }
+  return kTfLiteOk;
+}
+
+TfLiteStatus TransposeToCHW(const float* hwc, float* chw, const TfLiteIntArray* hwc_dims) {
+  int batch_size, height_size, width_size, depth_size;
+  TF_LITE_ENSURE_STATUS(GetDims(&batch_size, &height_size, &width_size, &depth_size, hwc_dims));
+  RuntimeShape hwc_shape({batch_size, height_size, width_size, depth_size});
+  RuntimeShape chw_shape({batch_size, depth_size, height_size, width_size});
+  TransposeParams params = {/*perm_count=*/4, /*perm=*/{0, 3, 1, 2}};
+  optimized_ops::Transpose<float>(params, hwc_shape, hwc, chw_shape, chw);
+  return kTfLiteOk;
+}
+
+TfLiteStatus TransposeToHWC(const float* chw, float* hwc, const TfLiteIntArray* hwc_dims) {
+  int batch_size, height_size, width_size, depth_size;
+  TF_LITE_ENSURE_STATUS(GetDims(&batch_size, &height_size, &width_size, &depth_size, hwc_dims));
+  RuntimeShape hwc_shape({batch_size, height_size, width_size, depth_size});
+  RuntimeShape chw_shape({batch_size, depth_size, height_size, width_size});
+  TransposeParams params = {/*perm_count=*/4, /*perm=*/{0, 2, 3, 1}};
+  optimized_ops::Transpose<float>(params, chw_shape, chw, hwc_shape, hwc);
+  return kTfLiteOk;
+}
+}  // namespace
+
+TfLiteStatus CoreMlDelegateKernel::Init(TfLiteContext* context,
+                                        const TfLiteDelegateParams* delegate_params) {
+  if (@available(iOS 12.0, *)) {
+    executor_ = [[::CoreMlExecutor alloc] init];
+    TF_LITE_ENSURE_STATUS(BuildModel(context, delegate_params));
+    // Serialize the model protocol buffer and compile it.
+    NSURL* model_url = [executor_ saveModel:model_.get()];
+    model_.reset();
+    if (model_url == nil) {
+      TF_LITE_KERNEL_LOG(context, "Failed to save model.");
+      return kTfLiteError;
+    }
+    if (![executor_ build:model_url]) {
+      TF_LITE_KERNEL_LOG(context, "Failed to compile and save model.");
+      return kTfLiteError;
+    }
+    return kTfLiteOk;
+  } else {
+    TF_LITE_KERNEL_LOG(context, "Minimum required iOS version is 12.0.");
+    return kTfLiteError;
+  }
+}
+
+void CoreMlDelegateKernel::AddInputTensors(const TfLiteIntArray* input_tensors,
+                                           TfLiteContext* context) {
+  int num_inputs = 0;
+  for (int i = 0; i < input_tensors->size; ++i) {
+    const int tensor_id = input_tensors->data[i];
+    builder_->AddTensorWithID(tensor_id, delegates::coreml::TensorID(0, num_inputs++));
+  }
+}
+
+TfLiteStatus CoreMlDelegateKernel::AddOutputTensors(const TfLiteIntArray* output_tensors,
+                                                    TfLiteContext* context) {
+  auto* model_description = model_->mutable_description();
+  for (int i = 0; i < output_tensors->size; ++i) {
+    const int tensor_id = output_tensors->data[i];
+    const auto& tensor = context->tensors[tensor_id];
+
+    auto* output = model_description->mutable_output()->Add();
+    output->set_name(builder_->GetTensorName(tensor_id));
+    auto* multi_array = output->mutable_type()->mutable_multiarraytype();
+    int batch_size, height_size, width_size, depth_size;
+    TF_LITE_ENSURE_STATUS(
+        GetDims(&batch_size, &height_size, &width_size, &depth_size, tensor.dims));
+    multi_array->set_datatype(CoreML::Specification::ArrayFeatureType::FLOAT32);
+    if (coreml_version_ >= 3) {
+      multi_array->mutable_shape()->Add(batch_size);
+    }
+    multi_array->mutable_shape()->Add(depth_size);
+    multi_array->mutable_shape()->Add(height_size);
+    multi_array->mutable_shape()->Add(width_size);
+  }
+  return kTfLiteOk;
+}
+
+TfLiteStatus CoreMlDelegateKernel::BuildModel(TfLiteContext* context,
+                                              const TfLiteDelegateParams* delegate_params) {
+  builder_ = std::make_unique<delegates::coreml::GraphBuilder>(coreml_version_);
+  // Add Inputs
+  AddInputTensors(delegate_params->input_tensors, context);
+  // Build all ops.
+  for (int node_index : TfLiteIntArrayView(delegate_params->nodes_to_replace)) {
+    TfLiteNode* node;
+    TfLiteRegistration* reg;
+    TF_LITE_ENSURE_STATUS(context->GetNodeAndRegistration(context, node_index, &node, &reg));
+    auto* op_builder = builder_->AddBuilder(reg->builtin_code, node);
+    if (op_builder == nullptr) {
+      TF_LITE_KERNEL_LOG(context, "Failed to build node %d.", node_index);
+      return kTfLiteError;
+    }
+    if (op_builder->RegisterInputs(node->inputs, context) != kTfLiteOk) {
+      TF_LITE_KERNEL_LOG(context, "Failed to add inputs for node %d.", node_index);
+      return kTfLiteError;
+    }
+    if (op_builder->PopulateSubgraph(context) != kTfLiteOk) {
+      TF_LITE_KERNEL_LOG(context, "Failed to add sub-builders for node %d.", node_index);
+      return kTfLiteError;
+    }
+    if (op_builder->RegisterOutputs(node->outputs, context) != kTfLiteOk) {
+      TF_LITE_KERNEL_LOG(context, "Failed to add outputs for node %d.", node_index);
+      return kTfLiteError;
+    }
+  }
+  model_.reset(builder_->BuildModel());
+  if (model_ == nullptr) {
+    TF_LITE_KERNEL_LOG(context, "Failed to build Model");
+    return kTfLiteError;
+  }
+  TF_LITE_ENSURE_STATUS(AddOutputTensors(delegate_params->output_tensors, context));
+  auto* model_description = model_->mutable_description();
+  for (int i = 0; i < delegate_params->input_tensors->size; ++i) {
+    const int tensor_id = delegate_params->input_tensors->data[i];
+    if (builder_->IsTensorUsed(tensor_id)) {
+      const auto& tensor = context->tensors[tensor_id];
+      auto* input = model_description->mutable_input()->Add();
+      input->set_name(builder_->GetTensorName(tensor_id));
+      // TODO(karimnosseir): Handle different types ?
+      auto* multi_array = input->mutable_type()->mutable_multiarraytype();
+      int batch_size, height_size, width_size, depth_size;
+      TF_LITE_ENSURE_STATUS(
+          GetDims(&batch_size, &height_size, &width_size, &depth_size, tensor.dims));
+      multi_array->set_datatype(CoreML::Specification::ArrayFeatureType::FLOAT32);
+      if (coreml_version_ >= 3) {
+        multi_array->mutable_shape()->Add(batch_size);
+      }
+      multi_array->mutable_shape()->Add(depth_size);
+      multi_array->mutable_shape()->Add(height_size);
+      multi_array->mutable_shape()->Add(width_size);
+    }
+  }
+
+  return kTfLiteOk;
+}
+
+TfLiteStatus CoreMlDelegateKernel::Prepare(TfLiteContext* context, TfLiteNode* node) {
+  input_tensor_ids_.clear();
+  inputs_.clear();
+  outputs_.clear();
+
+  for (int tensor_index : TfLiteIntArrayView(node->inputs)) {
+    if (builder_->IsTensorUsed(tensor_index)) {
+      input_tensor_ids_.push_back(tensor_index);
+    }
+  }
+
+  inputs_.reserve(input_tensor_ids_.size());
+  for (int tensor_index : input_tensor_ids_) {
+    TfLiteTensor* tensor = &context->tensors[tensor_index];
+    const int input_size = NumElements(tensor);
+    int batch_size, height_size, width_size, depth_size;
+    TF_LITE_ENSURE_STATUS(
+        GetDims(&batch_size, &height_size, &width_size, &depth_size, tensor->dims));
+
+    std::vector<int> input_shape = {depth_size, height_size, width_size};
+    if (coreml_version_ >= 3) {
+      input_shape.insert(input_shape.begin(), batch_size);
+    }
+    inputs_.push_back(
+        {std::vector<float>(input_size), builder_->GetTensorName(tensor_index), input_shape});
+  }
+
+  outputs_.reserve(node->outputs->size);
+  for (int tensor_index : TfLiteIntArrayView(node->outputs)) {
+    TfLiteTensor* tensor = &context->tensors[tensor_index];
+    const int output_size = NumElements(tensor);
+
+    outputs_.push_back({std::vector<float>(output_size), builder_->GetTensorName(tensor_index)});
+  }
+
+  return kTfLiteOk;
+}
+
+TfLiteStatus CoreMlDelegateKernel::Invoke(TfLiteContext* context, TfLiteNode* node) {
+  if (@available(iOS 12.0, *)) {
+    @autoreleasepool {
+      for (size_t i = 0; i < input_tensor_ids_.size(); ++i) {
+        const int tensor_id = input_tensor_ids_[i];
+        TfLiteTensor* tensor = &context->tensors[tensor_id];
+        const float* float_ptr = tensor->data.f;
+        std::vector<float> floats;
+        if (tensor->type == kTfLiteFloat16) {
+          floats = std::vector<float>(NumElements(tensor));
+          const tflite::half* float16_ptr = reinterpret_cast<const tflite::half*>(tensor->data.f16);
+          for (size_t j = 0; j < floats.size(); j++) {
+            floats[j] = float16_ptr[j];
+          }
+          float_ptr = floats.data();
+        }
+        // Transpose input to CHW.
+        // TODO(b/143992544): try adding transpose op for inputs.
+        TF_LITE_ENSURE_STATUS(TransposeToCHW(float_ptr, inputs_[i].data.data(), tensor->dims));
+      }
+
+      if (![executor_ invokeWithInputs:inputs_ outputs:outputs_]) {
+        return kTfLiteError;
+      }
+      for (int i = 0; i < node->outputs->size; ++i) {
+        TfLiteTensor* output_tensor = GetOutput(context, node, i);
+        if (output_tensor->type == kTfLiteFloat16) {
+          std::vector<float> temp_fp32(NumElements(output_tensor));
+          TF_LITE_ENSURE_STATUS(
+              TransposeToHWC(outputs_[i].data.data(), temp_fp32.data(), output_tensor->dims));
+          tflite::half* float16_ptr = reinterpret_cast<tflite::half*>(output_tensor->data.f16);
+          for (size_t j = 0; j < temp_fp32.size(); j++) {
+            float16_ptr[j] = tflite::half(temp_fp32[j]);
+          }
+        } else {
+          TF_LITE_ENSURE_STATUS(
+              TransposeToHWC(outputs_[i].data.data(), output_tensor->data.f, output_tensor->dims));
+        }
+      }
+    }
+    return kTfLiteOk;
+  } else {
+    TF_LITE_KERNEL_LOG(context, "Minimum required iOS version is 12.0.");
+    return kTfLiteError;
+  }
+}
+
+CoreMlDelegateKernel::~CoreMlDelegateKernel() { [executor_ cleanup]; }
+
+}  // namespace coreml
+}  // namespace delegates
+}  // namespace tflite

tensorflow/lite/delegates/coreml/coreml_executor.h

@@ -0,0 +1,48 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#import <CoreML/CoreML.h>
+
+#include <string>
+#include <vector>
+
+#include "mlmodel/format/Model.pb.h"
+
+// Data for input/output tensors.
+struct TensorData {
+  std::vector<float> data;
+  const std::string name;
+  std::vector<int> shape;  // only required for input tensor.
+};
+
+// Responsible for:
+// - Compiling and constructing MLModel from a serialized MlModel
+//   protocol buffer.
+// - Invoking predictions on the built model.
+// Usage: Construct object, call Build() and Invoke() for inference.
+@interface CoreMlExecutor : NSObject
+
+- (bool)invokeWithInputs:(const std::vector<TensorData>&)inputs
+                 outputs:(const std::vector<TensorData>&)outputs API_AVAILABLE(ios(11));
+
+- (NSURL*)saveModel:(CoreML::Specification::Model*)model API_AVAILABLE(ios(11));
+- (bool)build:(NSURL*)modelUrl API_AVAILABLE(ios(11));
+
+- (bool)cleanup;
+
+@property MLModel* model API_AVAILABLE(ios(11));
+@property NSString* mlModelFilePath;
+@property NSString* compiledModelFilePath;
+@property(nonatomic, readonly) int coreMlVersion;
+@end

tensorflow/lite/delegates/coreml/coreml_executor.mm

@@ -0,0 +1,222 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#import "tensorflow/lite/delegates/coreml/coreml_executor.h"
+
+#import <CoreML/CoreML.h>
+#import <Foundation/Foundation.h>
+
+#include <fstream>
+#include <iostream>
+
+namespace {
+// Returns NSURL for a temporary file.
+NSURL* createTemporaryFile() {
+  // Get temporary directory.
+  NSURL* temporaryDirectoryURL = [NSURL fileURLWithPath:NSTemporaryDirectory() isDirectory:YES];
+  // Generate a Unique file name to use.
+  NSString* temporaryFilename = [[NSProcessInfo processInfo] globallyUniqueString];
+  // Create URL to that file.
+  NSURL* temporaryFileURL = [temporaryDirectoryURL URLByAppendingPathComponent:temporaryFilename];
+
+  return temporaryFileURL;
+}
+}  // namespace
+
+@interface MultiArrayFeatureProvider : NSObject <MLFeatureProvider> {
+  const std::vector<TensorData>* _inputs;
+  NSSet* _featureNames;
+}
+
+- (instancetype)initWithInputs:(const std::vector<TensorData>*)inputs
+                 coreMlVersion:(int)coreMlVersion;
+- (MLFeatureValue*)featureValueForName:(NSString*)featureName API_AVAILABLE(ios(11));
+- (NSSet<NSString*>*)featureNames;
+
+@property(nonatomic, readonly) int coreMlVersion;
+
+@end
+
+@implementation MultiArrayFeatureProvider
+
+- (instancetype)initWithInputs:(const std::vector<TensorData>*)inputs
+                 coreMlVersion:(int)coreMlVersion {
+  self = [super init];
+  _inputs = inputs;
+  _coreMlVersion = coreMlVersion;
+  for (auto& input : *_inputs) {
+    if (input.name.empty()) {
+      return nil;
+    }
+  }
+  return self;
+}
+
+- (NSSet<NSString*>*)featureNames {
+  if (_featureNames == nil) {
+    NSMutableArray* names = [[NSMutableArray alloc] init];
+    for (auto& input : *_inputs) {
+      [names addObject:[NSString stringWithCString:input.name.c_str()
+                                          encoding:[NSString defaultCStringEncoding]]];
+    }
+    _featureNames = [NSSet setWithArray:names];
+  }
+  return _featureNames;
+}
+
+- (MLFeatureValue*)featureValueForName:(NSString*)featureName {
+  for (auto& input : *_inputs) {
+    if ([featureName cStringUsingEncoding:NSUTF8StringEncoding] == input.name) {
+      // TODO(b/141492326): Update shape handling for higher ranks
+      NSArray* shape = @[
+        @(input.shape[0]),
+        @(input.shape[1]),
+        @(input.shape[2]),
+      ];
+      NSArray* strides = @[
+        @(input.shape[1] * input.shape[2]),
+        @(input.shape[2]),
+        @1,
+      ];
+
+      if ([self coreMlVersion] >= 3) {
+        shape = @[
+          @(input.shape[0]),
+          @(input.shape[1]),
+          @(input.shape[2]),
+          @(input.shape[3]),
+        ];
+        strides = @[
+          @(input.shape[1] * input.shape[2] * input.shape[3]),
+          @(input.shape[2] * input.shape[3]),
+          @(input.shape[3]),
+          @1,
+        ];
+      };
+      NSError* error = nil;
+      MLMultiArray* mlArray = [[MLMultiArray alloc] initWithDataPointer:(float*)input.data.data()
+                                                                  shape:shape
+                                                               dataType:MLMultiArrayDataTypeFloat32
+                                                                strides:strides
+                                                            deallocator:(^(void* bytes){
+                                                                        })error:&error];
+      if (error != nil) {
+        NSLog(@"Failed to create MLMultiArray for feature %@ error: %@", featureName,
+              [error localizedDescription]);
+        return nil;
+      }
+      auto* mlFeatureValue = [MLFeatureValue featureValueWithMultiArray:mlArray];
+      return mlFeatureValue;
+    }
+  }
+
+  NSLog(@"Feature %@ not found", featureName);
+  return nil;
+}
+@end
+
+@implementation CoreMlExecutor
+- (bool)invokeWithInputs:(const std::vector<TensorData>&)inputs
+                 outputs:(const std::vector<TensorData>&)outputs {
+  if (_model == nil) {
+    return NO;
+  }
+  NSError* error = nil;
+  MultiArrayFeatureProvider* inputFeature =
+      [[MultiArrayFeatureProvider alloc] initWithInputs:&inputs coreMlVersion:[self coreMlVersion]];
+  if (inputFeature == nil) {
+    NSLog(@"inputFeature is not initialized.");
+    return NO;
+  }
+  MLPredictionOptions* options = [[MLPredictionOptions alloc] init];
+  id<MLFeatureProvider> outputFeature = [_model predictionFromFeatures:inputFeature
+                                                               options:options
+                                                                 error:&error];
+  if (error != nil) {
+    NSLog(@"Error executing model: %@", [error localizedDescription]);
+    return NO;
+  }
+  NSSet<NSString*>* outputFeatureNames = [outputFeature featureNames];
+  for (auto& output : outputs) {
+    NSString* outputName = [NSString stringWithCString:output.name.c_str()
+                                              encoding:[NSString defaultCStringEncoding]];
+    MLFeatureValue* outputValue =
+        [outputFeature featureValueForName:[outputFeatureNames member:outputName]];
+    auto* data = [outputValue multiArrayValue];
+    float* outputData = (float*)data.dataPointer;
+    if (outputData == nullptr) {
+      return NO;
+    }
+    memcpy((float*)output.data.data(), outputData, output.data.size() * sizeof(output.data[0]));
+  }
+  return YES;
+}
+
+- (bool)cleanup {
+  NSError* error = nil;
+  [[NSFileManager defaultManager] removeItemAtPath:_mlModelFilePath error:&error];
+  if (error != nil) {
+    NSLog(@"Failed cleaning up model: %@", [error localizedDescription]);
+    return NO;
+  }
+  [[NSFileManager defaultManager] removeItemAtPath:_compiledModelFilePath error:&error];
+  if (error != nil) {
+    NSLog(@"Failed cleaning up compiled model: %@", [error localizedDescription]);
+    return NO;
+  }
+  return YES;
+}
+
+- (NSURL*)saveModel:(CoreML::Specification::Model*)model {
+  NSURL* modelUrl = createTemporaryFile();
+  NSString* modelPath = [modelUrl path];
+  if (model->specificationversion() == 3) {
+    _coreMlVersion = 2;
+  } else if (model->specificationversion() == 4) {
+    _coreMlVersion = 3;
+  } else {
+    NSLog(@"Only Core ML models with specification version 3 or 4 are supported");
+    return nil;
+  }
+  // Flush data to file.
+  // TODO(karimnosseir): Can we mmap this instead of actual writing it to phone ?
+  std::ofstream file_stream([modelPath UTF8String], std::ios::out | std::ios::binary);
+  model->SerializeToOstream(&file_stream);
+  return modelUrl;
+}
+
+- (bool)build:(NSURL*)modelUrl {
+  NSError* error = nil;
+  NSURL* compileUrl = [MLModel compileModelAtURL:modelUrl error:&error];
+  if (error != nil) {
+    NSLog(@"Error compiling model %@", [error localizedDescription]);
+    return NO;
+  }
+  _mlModelFilePath = [modelUrl path];
+  _compiledModelFilePath = [compileUrl path];
+
+  if (@available(iOS 12.0, *)) {
+    MLModelConfiguration* config = [MLModelConfiguration alloc];
+    config.computeUnits = MLComputeUnitsAll;
+    _model = [MLModel modelWithContentsOfURL:compileUrl configuration:config error:&error];
+  } else {
+    _model = [MLModel modelWithContentsOfURL:compileUrl error:&error];
+  }
+  if (error != NULL) {
+    NSLog(@"Error Creating MLModel %@", [error localizedDescription]);
+    return NO;
+  }
+  return YES;
+}
+@end

tensorflow/lite/delegates/delegate_test_util.cc

@@ -0,0 +1,598 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/lite/delegates/delegate_test_util.h"
+
+#include <stdint.h>
+#include <stdlib.h>
+#include <string.h>
+
+#include <memory>
+#include <string>
+#include <vector>
+
+#include <gtest/gtest.h>
+#include "Eigen/Core"  // from @eigen_archive
+#include "tensorflow/lite/builtin_ops.h"
+#include "tensorflow/lite/core/c/builtin_op_data.h"
+#include "tensorflow/lite/core/interpreter.h"
+#include "tensorflow/lite/core/kernels/builtin_op_kernels.h"
+#include "tensorflow/lite/core/subgraph.h"
+#include "tensorflow/lite/delegates/utils.h"
+#include "tensorflow/lite/kernels/internal/compatibility.h"
+#include "tensorflow/lite/kernels/kernel_util.h"
+#include "tensorflow/lite/schema/schema_generated.h"
+#include "tensorflow/lite/util.h"
+
+namespace tflite {
+namespace delegates {
+namespace test_utils {
+
+TfLiteRegistration AddOpRegistration() {
+  TfLiteRegistration reg = {nullptr, nullptr, nullptr, nullptr};
+
+  reg.custom_name = "my_add";
+  reg.builtin_code = tflite::BuiltinOperator_CUSTOM;
+
+  reg.prepare = [](TfLiteContext* context, TfLiteNode* node) {
+    const TfLiteTensor* input1;
+    TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, 0, &input1));
+    const TfLiteTensor* input2;
+    TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, 1, &input2));
+    TfLiteTensor* output;
+    TF_LITE_ENSURE_OK(context, GetOutputSafe(context, node, 0, &output));
+
+    // Verify that the two inputs have the same shape.
+    TF_LITE_ENSURE_EQ(context, input1->dims->size, input2->dims->size);
+    for (int i = 0; i < input1->dims->size; ++i) {
+      TF_LITE_ENSURE_EQ(context, input1->dims->data[i], input2->dims->data[i]);
+    }
+
+    // Set output shape to match input shape.
+    TF_LITE_ENSURE_STATUS(context->ResizeTensor(
+        context, output, TfLiteIntArrayCopy(input1->dims)));
+    return kTfLiteOk;
+  };
+
+  reg.invoke = [](TfLiteContext* context, TfLiteNode* node) {
+    const TfLiteTensor* a0;
+    TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, 0, &a0));
+    TF_LITE_ENSURE(context, a0);
+    TF_LITE_ENSURE(context, a0->data.f);
+    const TfLiteTensor* a1;
+    TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, 1, &a1));
+    TF_LITE_ENSURE(context, a1);
+    TF_LITE_ENSURE(context, a1->data.f);
+    TfLiteTensor* out;
+    TF_LITE_ENSURE_OK(context, GetOutputSafe(context, node, 0, &out));
+    TF_LITE_ENSURE(context, out);
+    TF_LITE_ENSURE(context, out->data.f);
+    // Set output data to element-wise sum of input data.
+    int num = a0->dims->data[0];
+    for (int i = 0; i < num; i++) {
+      out->data.f[i] = a0->data.f[i] + a1->data.f[i];
+    }
+    return kTfLiteOk;
+  };
+  return reg;
+}
+
+void TestDelegation::SetUpSubgraph(Subgraph* subgraph) {
+  subgraph->AddTensors(5);
+  subgraph->SetInputs({0, 1});
+  subgraph->SetOutputs({3, 4});
+  std::vector<int> dims({3});
+  TfLiteQuantization quant{kTfLiteNoQuantization, nullptr};
+  subgraph->SetTensorParametersReadWrite(0, kTfLiteFloat32, "", dims.size(),
+                                         dims.data(), quant, false);
+  subgraph->SetTensorParametersReadWrite(1, kTfLiteFloat32, "", dims.size(),
+                                         dims.data(), quant, false);
+  subgraph->SetTensorParametersReadWrite(2, kTfLiteFloat32, "", dims.size(),
+                                         dims.data(), quant, false);
+  subgraph->SetTensorParametersReadWrite(3, kTfLiteFloat32, "", dims.size(),
+                                         dims.data(), quant, false);
+  subgraph->SetTensorParametersReadWrite(4, kTfLiteFloat32, "", dims.size(),
+                                         dims.data(), quant, false);
+  TfLiteRegistration reg = AddOpRegistration();
+  int node_index_ignored;
+  subgraph->AddNodeWithParameters({0, 0}, {2}, {}, nullptr, 0, nullptr, &reg,
+                                  &node_index_ignored);
+  subgraph->AddNodeWithParameters({1, 1}, {3}, {}, nullptr, 0, nullptr, &reg,
+                                  &node_index_ignored);
+  subgraph->AddNodeWithParameters({2, 1}, {4}, {}, nullptr, 0, nullptr, &reg,
+                                  &node_index_ignored);
+}
+
+void TestDelegation::AddSubgraphs(int subgraphs_to_add,
+                                  int* first_new_subgraph_index) {
+  interpreter_->AddSubgraphs(subgraphs_to_add, first_new_subgraph_index);
+}
+
+void TestDelegate::SetUp() {
+  interpreter_ = TestDelegation::NewInterpreterWithDefaultDelegates();
+  SetUpSubgraph(&interpreter_->primary_subgraph());
+}
+
+void TestDelegate::TearDown() {
+  // Interpreter relies on delegate to free the resources properly. Thus
+  // the life cycle of delegate must be longer than interpreter.
+  interpreter_.reset();
+  delegate_.reset();
+  delegate2_.reset();
+}
+
+void TestTwoDelegates::SetUp() {
+  interpreter_ = TestDelegation::NewInterpreterWithDefaultDelegates();
+  SetUpSubgraph(&interpreter_->primary_subgraph());
+}
+
+void TestTwoDelegates::TearDown() {
+  // Interpreter relies on delegate to free the resources properly. Thus
+  // the life cycle of delegate must be longer than interpreter.
+  interpreter_.reset();
+  delegate_.reset();
+  delegate2_.reset();
+}
+
+SimpleDelegate::SimpleDelegate(const std::vector<int>& nodes,
+                               int64_t delegate_flags, Options options,
+                               int min_ops_per_subset)
+    : nodes_(nodes),
+      fail_delegate_node_init_(options & Options::kFailOnInit),
+      fail_delegate_node_prepare_(options & Options::kFailOnPrepare),
+      fail_delegate_node_invoke_(options & Options::kFailOnInvoke),
+      min_ops_per_subset_(min_ops_per_subset),
+      automatic_shape_propagation_(options ==
+                                   Options::kAutomaticShapePropagation),
+      custom_op_(!(options & Options::kNoCustomOp)),
+      set_output_tensor_dynamic_(options & Options::kSetOutputTensorDynamic) {
+  delegate_ = TfLiteDelegateCreate();
+  delegate_.Prepare = [](TfLiteContext* context,
+                         TfLiteDelegate* delegate) -> TfLiteStatus {
+    auto* simple = static_cast<SimpleDelegate*>(delegate->data_);
+    TfLiteIntArray* nodes_to_separate =
+        TfLiteIntArrayCreate(simple->nodes_.size());
+    // Mark nodes that we want in TfLiteIntArray* structure.
+    int index = 0;
+    for (auto node_index : simple->nodes_) {
+      nodes_to_separate->data[index++] = node_index;
+      // make sure node is added
+      TfLiteNode* node;
+      TfLiteRegistration* reg;
+      context->GetNodeAndRegistration(context, node_index, &node, &reg);
+      if (simple->custom_op_) {
+        TFLITE_CHECK_EQ(reg->builtin_code, tflite::BuiltinOperator_CUSTOM);
+        TFLITE_CHECK_EQ(strcmp(reg->custom_name, "my_add"), 0);
+      } else {
+        TFLITE_CHECK_EQ(reg->builtin_code, tflite::BuiltinOperator_ADD);
+      }
+    }
+    // Check that all nodes are available
+    TfLiteIntArray* execution_plan;
+    TF_LITE_ENSURE_STATUS(context->GetExecutionPlan(context, &execution_plan));
+    for (int exec_index = 0; exec_index < execution_plan->size; exec_index++) {
+      int node_index = execution_plan->data[exec_index];
+      TfLiteNode* node;
+      TfLiteRegistration* reg;
+      context->GetNodeAndRegistration(context, node_index, &node, &reg);
+      if (exec_index == node_index) {
+        // Check op details only if it wasn't delegated already.
+        if (simple->custom_op_) {
+          TFLITE_CHECK_EQ(reg->builtin_code, tflite::BuiltinOperator_CUSTOM);
+          TFLITE_CHECK_EQ(strcmp(reg->custom_name, "my_add"), 0);
+        } else {
+          TFLITE_CHECK_EQ(reg->builtin_code, tflite::BuiltinOperator_ADD);
+        }
+      }
+    }
+
+    // Get preview of delegate partitioning from the context.
+    TfLiteDelegateParams* params_array;
+    int num_partitions;
+    TFLITE_CHECK_EQ(
+        context->PreviewDelegatePartitioning(context, nodes_to_separate,
+                                             &params_array, &num_partitions),
+        kTfLiteOk);
+
+    if (simple->min_ops_per_subset() > 0) {
+      // Build a new vector of ops from subsets with at least the minimum
+      // size.
+      std::vector<int> allowed_ops;
+      for (int idx = 0; idx < num_partitions; ++idx) {
+        const auto* nodes_in_subset = params_array[idx].nodes_to_replace;
+        if (nodes_in_subset->size < simple->min_ops_per_subset()) continue;
+        allowed_ops.insert(allowed_ops.end(), nodes_in_subset->data,
+                           nodes_in_subset->data + nodes_in_subset->size);
+      }
+
+      // Free existing nodes_to_separate & initialize a new array with
+      // allowed_ops.
+      TfLiteIntArrayFree(nodes_to_separate);
+      nodes_to_separate = TfLiteIntArrayCreate(allowed_ops.size());
+      memcpy(nodes_to_separate->data, allowed_ops.data(),
+             sizeof(int) * nodes_to_separate->size);
+    }
+
+    // Another call to PreviewDelegatePartitioning should be okay, since
+    // partitioning memory is managed by context.
+    TFLITE_CHECK_EQ(
+        context->PreviewDelegatePartitioning(context, nodes_to_separate,
+                                             &params_array, &num_partitions),
+        kTfLiteOk);
+
+    TfLiteStatus res = context->ReplaceNodeSubsetsWithDelegateKernels(
+        context, simple->FakeFusedRegistration(), nodes_to_separate, delegate);
+    TfLiteIntArrayFree(nodes_to_separate);
+    return res;
+  };
+  delegate_.CopyToBufferHandle = [](TfLiteContext* context,
+                                    TfLiteDelegate* delegate,
+                                    TfLiteBufferHandle buffer_handle,
+                                    TfLiteTensor* tensor) -> TfLiteStatus {
+    // TODO(b/156586986): Implement tests to test buffer copying logic.
+    return kTfLiteOk;
+  };
+  delegate_.CopyFromBufferHandle = [](TfLiteContext* context,
+                                      TfLiteDelegate* delegate,
+                                      TfLiteBufferHandle buffer_handle,
+                                      TfLiteTensor* output) -> TfLiteStatus {
+    TFLITE_CHECK_GE(buffer_handle, -1);
+    TFLITE_CHECK_EQ(output->buffer_handle, buffer_handle);
+    const float floats[] = {6., 6., 6.};
+    int num = output->dims->data[0];
+    for (int i = 0; i < num; i++) {
+      output->data.f[i] = floats[i];
+    }
+    return kTfLiteOk;
+  };
+
+  delegate_.FreeBufferHandle =
+      [](TfLiteContext* context, TfLiteDelegate* delegate,
+         TfLiteBufferHandle* handle) { *handle = kTfLiteNullBufferHandle; };
+  // Store type-punned data SimpleDelegate structure.
+  delegate_.data_ = static_cast<void*>(this);
+  delegate_.flags = delegate_flags;
+}
+
+TfLiteRegistration SimpleDelegate::FakeFusedRegistration() {
+  TfLiteRegistration reg = {nullptr};
+  reg.custom_name = "fake_fused_op";
+
+  if (fail_delegate_node_init_) {
+    reg.init = [](TfLiteContext* context, const char* buffer,
+                  size_t length) -> void* { return TfLiteKernelInitFailed(); };
+  }
+
+  // Different flavors of the delegate kernel's Invoke(), dependent on
+  // testing parameters.
+  if (fail_delegate_node_invoke_) {
+    reg.invoke = [](TfLiteContext* context, TfLiteNode* node) -> TfLiteStatus {
+      return kTfLiteError;
+    };
+  } else {
+    reg.invoke = [](TfLiteContext* context, TfLiteNode* node) -> TfLiteStatus {
+      // Compute output data as elementwise sum of the two input arguments:
+      //   func(x, y) = x + y
+      // or for a single argument compute 2 * x:
+      //   func(x) = x + x
+      const TfLiteTensor* a0;
+      const TfLiteTensor* a1;
+      if (node->inputs->size == 2) {
+        a0 = GetInput(context, node, 0);
+        a1 = GetInput(context, node, 1);
+      } else {
+        a0 = GetInput(context, node, 0);
+        a1 = a0;
+      }
+      TfLiteTensor* out;
+      TF_LITE_ENSURE_OK(context, GetOutputSafe(context, node, 0, &out));
+      int num = 1;
+      for (int i = 0; i < a0->dims->size; ++i) {
+        num *= a0->dims->data[i];
+      }
+      for (int i = 0; i < num; i++) {
+        out->data.f[i] = a0->data.f[i] + a1->data.f[i];
+      }
+      if (out->buffer_handle != kTfLiteNullBufferHandle) {
+        // Make the data stale so that CopyFromBufferHandle can be invoked
+        out->data_is_stale = true;
+      }
+      return kTfLiteOk;
+    };
+  }
+
+  // Different flavors of the delegate kernel's Prepare(), dependent on
+  // testing parameters.
+  if (automatic_shape_propagation_) {
+    reg.prepare = [](TfLiteContext* context, TfLiteNode* node) {
+      // Shapes should already by propagated by the runtime, just need to
+      // check.
+      const TfLiteTensor* input1;
+      TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, 0, &input1));
+      TfLiteTensor* output;
+      TF_LITE_ENSURE_OK(context, GetOutputSafe(context, node, 0, &output));
+      const int input_dims_size = input1->dims->size;
+      TF_LITE_ENSURE(context, output->dims->size == input_dims_size);
+      for (int i = 0; i < input_dims_size; ++i) {
+        TF_LITE_ENSURE(context, output->dims->data[i] == input1->dims->data[i]);
+      }
+      return kTfLiteOk;
+    };
+  } else if (fail_delegate_node_prepare_) {
+    reg.prepare = [](TfLiteContext* context, TfLiteNode* node) {
+      return kTfLiteError;
+    };
+  } else if (set_output_tensor_dynamic_) {
+    reg.prepare = [](TfLiteContext* context, TfLiteNode* node) {
+      TfLiteTensor* output;
+      TF_LITE_ENSURE_OK(context, GetOutputSafe(context, node, 0, &output));
+      SetTensorToDynamic(output);
+      return kTfLiteOk;
+    };
+  } else {
+    reg.prepare = [](TfLiteContext* context, TfLiteNode* node) {
+      // Set output size to input size
+      const TfLiteTensor* input1;
+      const TfLiteTensor* input2;
+      if (node->inputs->size == 2) {
+        input1 = GetInput(context, node, 0);
+        input2 = GetInput(context, node, 1);
+      } else {
+        input1 = GetInput(context, node, 0);
+        input2 = input1;
+      }
+      TfLiteTensor* output;
+      TF_LITE_ENSURE_OK(context, GetOutputSafe(context, node, 0, &output));
+
+      TF_LITE_ENSURE_STATUS(context->ResizeTensor(
+          context, output, TfLiteIntArrayCopy(input1->dims)));
+      return kTfLiteOk;
+    };
+  }
+
+  return reg;
+}
+
+std::unique_ptr<SimpleDelegate>
+SimpleDelegate::DelegateWithRuntimeShapePropagation(
+    const std::vector<int>& nodes, int64_t delegate_flags,
+    int min_ops_per_subset) {
+  return std::make_unique<SimpleDelegate>(
+      nodes, delegate_flags,
+      SimpleDelegate::Options::kAutomaticShapePropagation,
+      /*min_ops_per_subset=*/min_ops_per_subset);
+}
+
+std::unique_ptr<SimpleDelegate> SimpleDelegate::DelegateWithDynamicOutput(
+    const std::vector<int>& nodes) {
+  // All params default except nodes & set_output_tensor_dynamic.
+  return std::make_unique<SimpleDelegate>(
+      nodes, kTfLiteDelegateFlagsAllowDynamicTensors,
+      SimpleDelegate::Options::kSetOutputTensorDynamic);
+}
+
+void TestFP16Delegation::SetUp() {
+  interpreter_ = TestDelegation::NewInterpreterWithDefaultDelegates();
+  interpreter_->AddTensors(13);
+  interpreter_->SetInputs({0});
+  interpreter_->SetOutputs({12});
+
+  float16_const_ = Eigen::half(2.f);
+
+  // TENSORS.
+  TfLiteQuantizationParams quant;
+  // Input.
+  interpreter_->SetTensorParametersReadWrite(0, kTfLiteFloat32, "", {1}, quant);
+  // fp16 constant, dequantize output, Add0 output.
+  interpreter_->SetTensorParametersReadOnly(
+      1, kTfLiteFloat16, "", {1}, quant,
+      reinterpret_cast<const char*>(&float16_const_), sizeof(TfLiteFloat16));
+  interpreter_->SetTensorParametersReadWrite(2, kTfLiteFloat32, "", {1}, quant);
+  interpreter_->SetTensorParametersReadWrite(3, kTfLiteFloat32, "", {1}, quant);
+  // fp16 constant, dequantize output, Add1 output.
+  interpreter_->SetTensorParametersReadOnly(
+      4, kTfLiteFloat16, "", {1}, quant,
+      reinterpret_cast<const char*>(&float16_const_), sizeof(TfLiteFloat16));
+  interpreter_->SetTensorParametersReadWrite(5, kTfLiteFloat32, "", {1}, quant);
+  interpreter_->SetTensorParametersReadWrite(6, kTfLiteFloat32, "", {1}, quant);
+  // fp16 constant, dequantize output, Mul0 output.
+  interpreter_->SetTensorParametersReadOnly(
+      7, kTfLiteFloat16, "", {1}, quant,
+      reinterpret_cast<const char*>(&float16_const_), sizeof(TfLiteFloat16));
+  interpreter_->SetTensorParametersReadWrite(8, kTfLiteFloat32, "", {1}, quant);
+  interpreter_->SetTensorParametersReadWrite(9, kTfLiteFloat32, "", {1}, quant);
+  // fp16 constant, dequantize output, Add2 output.
+  interpreter_->SetTensorParametersReadOnly(
+      10, kTfLiteFloat16, "", {1}, quant,
+      reinterpret_cast<const char*>(&float16_const_), sizeof(TfLiteFloat16));
+  interpreter_->SetTensorParametersReadWrite(11, kTfLiteFloat32, "", {1},
+                                             quant);
+  interpreter_->SetTensorParametersReadWrite(12, kTfLiteFloat32, "", {1},
+                                             quant);
+
+  // NODES.
+  auto* add_reg = ops::builtin::Register_ADD();
+  auto* mul_reg = ops::builtin::Register_MUL();
+  auto* deq_reg = ops::builtin::Register_DEQUANTIZE();
+  add_reg->builtin_code = kTfLiteBuiltinAdd;
+  deq_reg->builtin_code = kTfLiteBuiltinDequantize;
+  mul_reg->builtin_code = kTfLiteBuiltinMul;
+  TfLiteAddParams* builtin_data0 =
+      reinterpret_cast<TfLiteAddParams*>(malloc(sizeof(TfLiteAddParams)));
+  TfLiteAddParams* builtin_data1 =
+      reinterpret_cast<TfLiteAddParams*>(malloc(sizeof(TfLiteAddParams)));
+  TfLiteMulParams* builtin_data2 =
+      reinterpret_cast<TfLiteMulParams*>(malloc(sizeof(TfLiteMulParams)));
+  TfLiteAddParams* builtin_data3 =
+      reinterpret_cast<TfLiteAddParams*>(malloc(sizeof(TfLiteAddParams)));
+  builtin_data0->activation = kTfLiteActNone;
+  builtin_data1->activation = kTfLiteActNone;
+  builtin_data2->activation = kTfLiteActNone;
+  builtin_data3->activation = kTfLiteActNone;
+  interpreter_->AddNodeWithParameters({1}, {2}, nullptr, 0, nullptr, deq_reg);
+  interpreter_->AddNodeWithParameters({0, 2}, {3}, nullptr, 0, builtin_data0,
+                                      add_reg);
+  interpreter_->AddNodeWithParameters({4}, {5}, nullptr, 0, nullptr, deq_reg);
+  interpreter_->AddNodeWithParameters({3, 5}, {6}, nullptr, 0, builtin_data1,
+                                      add_reg);
+  interpreter_->AddNodeWithParameters({7}, {8}, nullptr, 0, nullptr, deq_reg);
+  interpreter_->AddNodeWithParameters({6, 8}, {9}, nullptr, 0, builtin_data2,
+                                      mul_reg);
+  interpreter_->AddNodeWithParameters({10}, {11}, nullptr, 0, nullptr, deq_reg);
+  interpreter_->AddNodeWithParameters({9, 11}, {12}, nullptr, 0, builtin_data3,
+                                      add_reg);
+}
+
+void TestFP16Delegation::VerifyInvoke() {
+  std::vector<float> input = {3.0f};
+  std::vector<float> expected_output = {16.0f};
+
+  const int input_tensor_idx = interpreter_->inputs()[0];
+  const int output_tensor_idx = interpreter_->outputs()[0];
+
+  memcpy(interpreter_->typed_tensor<float>(input_tensor_idx), input.data(),
+         sizeof(float));
+  ASSERT_EQ(interpreter_->Invoke(), kTfLiteOk);
+  TfLiteTensor* output_tensor = interpreter_->tensor(output_tensor_idx);
+  for (int i = 0; i < 1; ++i) {
+    EXPECT_EQ(output_tensor->data.f[i], expected_output[i]) << i;
+  }
+}
+
+TestFP16Delegation::FP16Delegate::FP16Delegate(int num_delegated_subsets,
+                                               bool fail_node_prepare,
+                                               bool fail_node_invoke)
+    : num_delegated_subsets_(num_delegated_subsets),
+      fail_delegate_node_prepare_(fail_node_prepare),
+      fail_delegate_node_invoke_(fail_node_invoke) {
+  delegate_.Prepare = [](TfLiteContext* context,
+                         TfLiteDelegate* delegate) -> TfLiteStatus {
+    auto* fp16_delegate = static_cast<FP16Delegate*>(delegate->data_);
+    // FP16 graph partitioning.
+    delegates::IsNodeSupportedFn node_supported_fn =
+        [=](TfLiteContext* context, TfLiteNode* node,
+            TfLiteRegistration* registration,
+            std::string* unsupported_details) -> bool {
+      return registration->builtin_code == kTfLiteBuiltinAdd;
+    };
+    delegates::FP16GraphPartitionHelper partition_helper(context,
+                                                         node_supported_fn);
+    TfLiteIntArray* nodes_to_separate = nullptr;
+    if (partition_helper.Partition(nullptr) != kTfLiteOk) {
+      nodes_to_separate = TfLiteIntArrayCreate(0);
+    } else {
+      std::vector<int> ops_to_replace =
+          partition_helper.GetNodesOfFirstNLargestPartitions(
+              fp16_delegate->num_delegated_subsets());
+      nodes_to_separate = ConvertVectorToTfLiteIntArray(ops_to_replace);
+    }
+
+    context->ReplaceNodeSubsetsWithDelegateKernels(
+        context, fp16_delegate->FakeFusedRegistration(), nodes_to_separate,
+        delegate);
+    TfLiteIntArrayFree(nodes_to_separate);
+    return kTfLiteOk;
+  };
+  delegate_.CopyFromBufferHandle =
+      [](TfLiteContext* context, TfLiteDelegate* delegate,
+         TfLiteBufferHandle buffer_handle,
+         TfLiteTensor* output) -> TfLiteStatus { return kTfLiteOk; };
+  delegate_.FreeBufferHandle = nullptr;
+  delegate_.CopyToBufferHandle = nullptr;
+  // Store type-punned data SimpleDelegate structure.
+  delegate_.data_ = static_cast<void*>(this);
+  delegate_.flags = kTfLiteDelegateFlagsNone;
+}
+
+TfLiteRegistration TestFP16Delegation::FP16Delegate::FakeFusedRegistration() {
+  TfLiteRegistration reg = {nullptr};
+  reg.custom_name = "fake_fp16_add_op";
+
+  // Different flavors of the delegate kernel's Invoke(), dependent on
+  // testing parameters.
+  if (fail_delegate_node_invoke_) {
+    reg.invoke = [](TfLiteContext* context, TfLiteNode* node) -> TfLiteStatus {
+      return kTfLiteError;
+    };
+  } else {
+    reg.invoke = [](TfLiteContext* context, TfLiteNode* node) -> TfLiteStatus {
+      float output = 0;
+      for (int i = 0; i < node->inputs->size; ++i) {
+        const TfLiteTensor* input_tensor = GetInput(context, node, i);
+        if (input_tensor->type == kTfLiteFloat32) {
+          output += input_tensor->data.f[0];
+        } else {
+          // All constants are 2.
+          output += 2;
+        }
+      }
+      TfLiteTensor* out = GetOutput(context, node, 0);
+      out->data.f[0] = output;
+      return kTfLiteOk;
+    };
+  }
+
+  // Different flavors of the delegate kernel's Prepare(), dependent on
+  // testing parameters.
+  if (fail_delegate_node_prepare_) {
+    reg.prepare = [](TfLiteContext* context, TfLiteNode* node) {
+      return kTfLiteError;
+    };
+  } else {
+    reg.prepare = [](TfLiteContext* context, TfLiteNode* node) {
+      // Set output size to input size
+      const TfLiteTensor* input = GetInput(context, node, 0);
+      TfLiteTensor* output = GetOutput(context, node, 0);
+      TF_LITE_ENSURE_STATUS(context->ResizeTensor(
+          context, output, TfLiteIntArrayCopy(input->dims)));
+      return kTfLiteOk;
+    };
+  }
+
+  return reg;
+}
+
+void TestDelegateWithControlEdges::SetUpSubgraph(Subgraph* subgraph) {
+  subgraph->AddTensors(5);
+  subgraph->SetInputs({0});
+  subgraph->SetOutputs({4});
+  std::vector<int> dims({3});
+  const TfLiteQuantization quant{kTfLiteNoQuantization, nullptr};
+  subgraph->SetTensorParametersReadWrite(0, kTfLiteFloat32, "", dims.size(),
+                                         dims.data(), quant, false);
+  subgraph->SetTensorParametersReadWrite(1, kTfLiteFloat32, "", dims.size(),
+                                         dims.data(), quant, false);
+  subgraph->SetTensorParametersReadWrite(2, kTfLiteFloat32, "", dims.size(),
+                                         dims.data(), quant, false);
+  subgraph->SetTensorParametersReadWrite(3, kTfLiteFloat32, "", dims.size(),
+                                         dims.data(), quant, false);
+  subgraph->SetTensorParametersReadWrite(4, kTfLiteFloat32, "", dims.size(),
+                                         dims.data(), quant, false);
+
+  TfLiteRegistration reg = AddOpRegistration();
+  int node_index_ignored;
+  subgraph->AddNodeWithParameters({0, 0}, {1}, {}, nullptr, 0, nullptr, &reg,
+                                  &node_index_ignored);
+  subgraph->AddNodeWithParameters({1, 1}, {2}, {}, nullptr, 0, nullptr, &reg,
+                                  &node_index_ignored);
+  subgraph->AddNodeWithParameters({1, 1}, {3}, {}, nullptr, 0, nullptr, &reg,
+                                  &node_index_ignored);
+  subgraph->AddNodeWithParameters({2, 3}, {4}, {}, nullptr, 0, nullptr, &reg,
+                                  &node_index_ignored);
+}
+
+}  // namespace test_utils
+}  // namespace delegates
+}  // namespace tflite

tensorflow/lite/delegates/delegate_test_util.h

@@ -0,0 +1,246 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_DELEGATES_DELEGATE_TEST_UTIL_H_
+#define TENSORFLOW_LITE_DELEGATES_DELEGATE_TEST_UTIL_H_
+
+#include <stdint.h>
+
+#include <map>
+#include <memory>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include <gtest/gtest.h>
+#include "Eigen/Core"  // from @eigen_archive
+#include "tensorflow/lite/core/interpreter.h"
+#include "tensorflow/lite/core/kernels/register.h"
+#include "tensorflow/lite/core/subgraph.h"
+#include "tensorflow/lite/kernels/internal/compatibility.h"
+
+namespace tflite {
+namespace delegates {
+namespace test_utils {
+
+// Build a kernel registration for a custom addition op that adds its two
+// tensor inputs to produce a tensor output.
+TfLiteRegistration AddOpRegistration();
+
+class SimpleDelegate {
+ public:
+  enum Options {
+    kNone = 0,
+    kFailOnInit = 1 << 0,  // To simulate failure of Delegate node's Init().
+    kFailOnPrepare =
+        1 << 1,  // To simulate failure of Delegate node's Prepare().
+    kFailOnInvoke = 1 << 2,  // To simulate failure of Delegate node's Invoke().
+    kAutomaticShapePropagation =
+        1 << 3,  // This assumes that the runtime will propagate
+                 // shapes using the original execution plan.
+    kNoCustomOp =
+        1 << 4,  // If not set, the graph nodes specified in the 'nodes'
+                 // parameter should be custom ops with name "my_add";
+                 // if set, they should be the builtin ADD operator.
+    kSetOutputTensorDynamic = 1
+                              << 5,  // If set, this delegate sets output tensor
+                                     // to as dynamic during kernel Prepare.
+  };
+
+  // Create a simple implementation of a TfLiteDelegate. We use the C++ class
+  // SimpleDelegate and it can produce a handle TfLiteDelegate that is
+  // value-copyable and compatible with TfLite.
+  //
+  // Parameters:
+  //   nodes: Indices of the graph nodes that the delegate will handle.
+  //   min_ops_per_subset: If >0, partitioning preview is used to choose only
+  //     those subsets with min_ops_per_subset number of nodes.
+  //   options: Options to control the behavior of the delegate.
+  explicit SimpleDelegate(const std::vector<int>& nodes,
+                          int64_t delegate_flags = kTfLiteDelegateFlagsNone,
+                          Options options = Options::kNone,
+                          int min_ops_per_subset = 0);
+
+  static std::unique_ptr<SimpleDelegate> DelegateWithRuntimeShapePropagation(
+      const std::vector<int>& nodes, int64_t delegate_flags,
+      int min_ops_per_subset);
+
+  static std::unique_ptr<SimpleDelegate> DelegateWithDynamicOutput(
+      const std::vector<int>& nodes);
+
+  TfLiteRegistration FakeFusedRegistration();
+
+  TfLiteDelegate* get_tf_lite_delegate() { return &delegate_; }
+
+  int min_ops_per_subset() { return min_ops_per_subset_; }
+
+ private:
+  std::vector<int> nodes_;
+  TfLiteDelegate delegate_;
+  bool fail_delegate_node_init_ = false;
+  bool fail_delegate_node_prepare_ = false;
+  bool fail_delegate_node_invoke_ = false;
+  int min_ops_per_subset_ = 0;
+  bool automatic_shape_propagation_ = false;
+  bool custom_op_ = true;
+  bool set_output_tensor_dynamic_ = false;
+};
+
+// Overload bitwise AND operator for SimpleDelegate::Options
+inline SimpleDelegate::Options operator&(SimpleDelegate::Options a,
+                                         SimpleDelegate::Options b) {
+  return static_cast<SimpleDelegate::Options>(static_cast<unsigned int>(a) &
+                                              static_cast<unsigned int>(b));
+}
+
+// Overload bitwise OR operator for SimpleDelegate::Options
+inline SimpleDelegate::Options operator|(SimpleDelegate::Options a,
+                                         SimpleDelegate::Options b) {
+  return static_cast<SimpleDelegate::Options>(static_cast<unsigned int>(a) |
+                                              static_cast<unsigned int>(b));
+}
+
+// Base class for single/multiple delegate tests.
+// Friend of Interpreter to access private methods.
+class TestDelegation {
+ public:
+  virtual ~TestDelegation() = default;
+
+  // Returns an empty interpreter that uses the same default delegates that are
+  // normally enabled by default.
+  static std::unique_ptr<impl::Interpreter>
+  NewInterpreterWithDefaultDelegates() {
+    auto interpreter = std::make_unique<impl::Interpreter>();
+    interpreter->lazy_delegate_providers_ =
+        tflite::ops::builtin::BuiltinOpResolver().GetDelegateCreators();
+    return interpreter;
+  }
+
+ protected:
+  TfLiteStatus RemoveAllDelegates() {
+    return interpreter_->RemoveAllDelegates();
+  }
+  void SetMetadata(const std::map<std::string, std::string>& metadata) {
+    interpreter_->SetMetadata(metadata);
+  }
+
+  virtual void SetUpSubgraph(Subgraph* subgraph);
+  void AddSubgraphs(int subgraphs_to_add,
+                    int* first_new_subgraph_index = nullptr);
+
+  std::unique_ptr<impl::Interpreter> interpreter_;
+};
+
+// Tests scenarios involving a single delegate.
+class TestDelegate : public TestDelegation, public ::testing::Test {
+ protected:
+  void SetUp() override;
+
+  void TearDown() override;
+
+  TfLiteBufferHandle last_allocated_handle_ = kTfLiteNullBufferHandle;
+
+  TfLiteBufferHandle AllocateBufferHandle() { return ++last_allocated_handle_; }
+
+  std::unique_ptr<SimpleDelegate> delegate_, delegate2_;
+};
+
+// Tests scenarios involving a single delegate and control edges.
+// Subgraph 0 has the form
+//
+//         /---OP2---\
+//        /           \
+// >---OP0             OP3--->
+//        \           /
+//         \---OP1---/
+//
+// Delegating OP0, OP2 will generate an execution graph with a "super-node"
+// {OP0->OP2}, which can be disabled by adding (in metadata) a control edge
+// between OP1 and OP2:
+//
+//         /->-OP2---\
+//        /     ^     \
+// >---OP0      ^      OP3--->
+//        \     ^     /
+//         \---OP1---/
+//
+class TestDelegateWithControlEdges : public TestDelegate {
+ protected:
+  void SetUpSubgraph(Subgraph* subgraph) override;
+};
+
+// Tests scenarios involving two delegates, parametrized by the first & second
+// delegate's flags.
+class TestTwoDelegates
+    : public TestDelegation,
+      public ::testing::TestWithParam<
+          std::pair<TfLiteDelegateFlags, TfLiteDelegateFlags>> {
+ protected:
+  void SetUp() override;
+
+  void TearDown() override;
+
+  std::unique_ptr<SimpleDelegate> delegate_, delegate2_;
+};
+
+// Tests delegate functionality related to FP16 graphs.
+// Model architecture:
+// 1->DEQ->2   4->DEQ->5   7->DEQ->8   10->DEQ->11
+//         |           |           |            |
+// 0----->ADD->3----->ADD->6----->MUL->9------>ADD-->12
+// Input: 0, Output:12.
+// All constants are 2, so the function is: (x + 2 + 2) * 2 + 2 = 2x + 10
+//
+// Delegate only supports ADD, so can have up to two delegated partitions.
+// TODO(b/156707497): Add more cases here once we have landed CPU kernels
+// supporting FP16.
+class TestFP16Delegation : public ::testing::TestWithParam<int> {
+ protected:
+  void SetUp() override;
+
+  void VerifyInvoke();
+
+  void TearDown() override { interpreter_.reset(); }
+
+ protected:
+  class FP16Delegate {
+   public:
+    // Uses FP16GraphPartitionHelper to accept ADD nodes with fp16 input.
+    explicit FP16Delegate(int num_delegated_subsets,
+                          bool fail_node_prepare = false,
+                          bool fail_node_invoke = false);
+
+    TfLiteRegistration FakeFusedRegistration();
+
+    TfLiteDelegate* get_tf_lite_delegate() { return &delegate_; }
+
+    int num_delegated_subsets() { return num_delegated_subsets_; }
+
+   private:
+    TfLiteDelegate delegate_;
+    int num_delegated_subsets_;
+    bool fail_delegate_node_prepare_ = false;
+    bool fail_delegate_node_invoke_ = false;
+  };
+
+  std::unique_ptr<impl::Interpreter> interpreter_;
+  std::unique_ptr<FP16Delegate> delegate_;
+  Eigen::half float16_const_;
+};
+
+}  // namespace test_utils
+}  // namespace delegates
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_DELEGATES_DELEGATE_TEST_UTIL_H_

tensorflow/lite/delegates/external/BUILD

@@ -0,0 +1,50 @@
+# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+load("@rules_cc//cc:cc_library.bzl", "cc_library")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:LICENSE"],
+    default_visibility = [
+        "//visibility:public",
+    ],
+    licenses = ["notice"],
+)
+
+cc_library(
+    name = "external_delegate_interface",
+    hdrs = ["external_delegate_interface.h"],
+    deps = [
+        "//tensorflow/lite/c:common",
+    ],
+)
+
+cc_library(
+    name = "external_delegate",
+    srcs = ["external_delegate.cc"],
+    hdrs = ["external_delegate.h"],
+    deps = [
+        ":external_delegate_interface",
+        "//tensorflow/lite:minimal_logging",
+        "//tensorflow/lite:shared_library",
+        "//tensorflow/lite/c:c_api_types",
+        "//tensorflow/lite/c:common",
+        "//tensorflow/lite/core/c:common",
+    ],
+)
+
+exports_files([
+    "external_delegate.h",
+])

tensorflow/lite/delegates/external/README.md

@@ -0,0 +1,33 @@
+# What is an External Delegate?
+An external delegate is a special Tensorflow Lite delegate that is simply
+initialized from loading a dynamic library which encapsulates an actual
+Tensorflow Lite delegate implementation. The actual delegate exposes the
+following two creation and deletion C APIs:
+
+* __tflite_plugin_create_delegate__ (declaration seen below) creates a delegate
+object based on provided key-value options. It may return NULL to indicate an
+error with the detailed information reported by calling `report_error` if
+provided. Each option key and value should be null-terminated.
+
+```
+TfLiteDelegate* tflite_plugin_create_delegate(
+  char** options_keys, char** options_values, size_t num_options,
+  void (*report_error)(const char *))
+```
+
+* __tflite_plugin_destroy_delegate__ (declaration seen below) destroys the
+delegate object that is created by the previous API. NULL as an argument value
+is allowed.
+
+```
+void tflite_plugin_destroy_delegate(TfLiteDelegate* delegate)
+```
+
+The external delegate provides an opaque and transparent way to utilize a
+Tensorflow Lite delegate when performing inference. In other words, one may
+replace the actual Tensorflow Lite delegate by simply updating the dynamic
+library without changing the application code. We developed this mainly for
+delegate evaluation.
+
+Note, this delegate is the corresponding C++ implementation to the one for
+Tensorflow Lite Python binding as shown [here](https://github.com/tensorflow/tensorflow/blob/7145fc0e49be01ef6943f4df386ce38567e37797/tensorflow/lite/python/interpreter.py#L42).

tensorflow/lite/delegates/external/external_delegate.cc

@@ -0,0 +1,227 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include "tensorflow/lite/delegates/external/external_delegate.h"
+
+#include <locale>
+#include <memory>
+#include <string>
+#include <vector>
+
+#include "tensorflow/lite/c/c_api_types.h"
+#include "tensorflow/lite/c/common.h"
+#include "tensorflow/lite/delegates/external/external_delegate_interface.h"
+#include "tensorflow/lite/logger.h"
+#include "tensorflow/lite/minimal_logging.h"
+#include "tensorflow/lite/shared_library.h"
+
+namespace tflite {
+namespace {
+
+// TODO(b/245168068): Add support for `TfLiteOpaqueDelegateBuilder`.
+
+// External delegate library construct
+struct ExternalLib {
+  struct wchar_codecvt : public std::codecvt<wchar_t, char, std::mbstate_t> {};
+
+  // Open a given delegate library and load the create/destroy symbols
+  bool load(const std::string library) {
+#if defined(_WIN32)
+    std::wstring_convert<wchar_codecvt> converter;
+    void* handle = SharedLibrary::LoadLibrary(
+        converter.from_bytes(library.c_str()).c_str());
+#else
+    void* handle = SharedLibrary::LoadLibrary(library.c_str());
+#endif  // defined(_WIN32)
+    if (handle == nullptr) {
+      TFLITE_LOG(TFLITE_LOG_INFO,
+                 "Unable to load external delegate from : %s (%s)",
+                 library.c_str(), SharedLibrary::GetError());
+    } else {
+      create = reinterpret_cast<decltype(&tflite_plugin_create_delegate)>(
+          SharedLibrary::GetLibrarySymbol(handle,
+                                          "tflite_plugin_create_delegate"));
+      destroy = reinterpret_cast<decltype(&tflite_plugin_destroy_delegate)>(
+          SharedLibrary::GetLibrarySymbol(handle,
+                                          "tflite_plugin_destroy_delegate"));
+      return create && destroy;
+    }
+    return false;
+  }
+
+  decltype(&tflite_plugin_create_delegate) create{nullptr};
+  decltype(&tflite_plugin_destroy_delegate) destroy{nullptr};
+};
+
+// An ExternalDelegateWrapper is responsibile to manage a TFLite delegate
+// initialized from a shared library. It creates a delegate from the given
+// option and storages it to external_delegate_ member variable. On the
+// destruction, it conducts necessary clean up process.
+class ExternalDelegateWrapper {
+ public:
+  explicit ExternalDelegateWrapper(
+      const TfLiteExternalDelegateOptions* options);
+  ~ExternalDelegateWrapper();
+
+  // Return a TfLiteDelegate which is created from
+  // tflite_plugin_create_delegate() of an external delegate logic.
+  TfLiteDelegate* tflite_external_delegate() { return external_delegate_; }
+
+  // Return a TfLiteDelegate which is convertibile to this class.
+  TfLiteDelegate* tflite_wrapper_delegate() { return &wrapper_delegate_; }
+
+ private:
+  ExternalLib external_lib_;
+
+  // external delegate instance owned by external delegate logic.
+  // It's created by "tflite_plugin_destroy_delegate()" function in the external
+  // delegate logic And it should be released by
+  // "tflite_plugin_destroy_delegate()" function.
+  TfLiteDelegate* external_delegate_;
+
+  // TfLiteDelegate representation of this ExternalDelegateWrapper object.
+  TfLiteDelegate wrapper_delegate_ = {};
+};
+
+// Converts the given TfLiteDelegate to an ExternalDelegateWrapper instance.
+inline ExternalDelegateWrapper* GetExternalDelegateWrapper(
+    TfLiteDelegate* delegate) {
+  return reinterpret_cast<ExternalDelegateWrapper*>(delegate->data_);
+}
+
+// Relay Prepare() call to the associated external TfLiteDelegate object.
+TfLiteStatus DelegatePrepare(TfLiteContext* context, TfLiteDelegate* delegate) {
+  auto external_delegate_wrapper = GetExternalDelegateWrapper(delegate);
+  TfLiteDelegate* external_delegate =
+      external_delegate_wrapper->tflite_external_delegate();
+  return external_delegate->Prepare(context, external_delegate);
+}
+
+// Relay CopyFromBufferHandle() call to the associated external TfLiteDelegate
+// object.
+TfLiteStatus DelegateCopyFromBufferHandle(TfLiteContext* context,
+                                          struct TfLiteDelegate* delegate,
+                                          TfLiteBufferHandle buffer_handle,
+                                          TfLiteTensor* tensor) {
+  auto external_delegate_wrapper = GetExternalDelegateWrapper(delegate);
+  TfLiteDelegate* external_delegate =
+      external_delegate_wrapper->tflite_external_delegate();
+  return external_delegate->CopyFromBufferHandle(context, delegate,
+                                                 buffer_handle, tensor);
+}
+
+// Relay CopyToBufferHandle() call to the associated external TfLiteDelegate
+// object.
+TfLiteStatus DelegateCopyToBufferHandle(TfLiteContext* context,
+                                        struct TfLiteDelegate* delegate,
+                                        TfLiteBufferHandle buffer_handle,
+                                        TfLiteTensor* tensor) {
+  auto external_delegate_wrapper = GetExternalDelegateWrapper(delegate);
+  TfLiteDelegate* external_delegate =
+      external_delegate_wrapper->tflite_external_delegate();
+  return external_delegate->CopyToBufferHandle(context, delegate, buffer_handle,
+                                               tensor);
+}
+
+// Relay FreeBufferHandle() call to the associated external TfLiteDelegate
+// object.
+void DelegateFreeBufferHandle(TfLiteContext* context,
+                              struct TfLiteDelegate* delegate,
+                              TfLiteBufferHandle* handle) {
+  auto external_delegate_wrapper = GetExternalDelegateWrapper(delegate);
+  TfLiteDelegate* external_delegate =
+      external_delegate_wrapper->tflite_external_delegate();
+  return external_delegate->FreeBufferHandle(context, delegate, handle);
+}
+
+ExternalDelegateWrapper::ExternalDelegateWrapper(
+    const TfLiteExternalDelegateOptions* options) {
+  external_delegate_ = nullptr;
+  if (external_lib_.load(options->lib_path)) {
+    std::vector<const char*> ckeys, cvalues;
+    for (int i = 0; i < options->count; i++) {
+      ckeys.push_back(options->keys[i]);
+      cvalues.push_back(options->values[i]);
+    }
+
+    external_delegate_ = external_lib_.create(ckeys.data(), cvalues.data(),
+                                              ckeys.size(), nullptr);
+    if (external_delegate_) {
+      wrapper_delegate_.data_ = reinterpret_cast<void*>(this);
+      wrapper_delegate_.Prepare = DelegatePrepare;
+      wrapper_delegate_.CopyFromBufferHandle = nullptr;
+      wrapper_delegate_.CopyToBufferHandle = nullptr;
+      wrapper_delegate_.FreeBufferHandle = nullptr;
+      wrapper_delegate_.flags = external_delegate_->flags;
+      if (external_delegate_->CopyFromBufferHandle) {
+        wrapper_delegate_.CopyFromBufferHandle = DelegateCopyFromBufferHandle;
+      }
+      if (external_delegate_->CopyToBufferHandle) {
+        wrapper_delegate_.CopyToBufferHandle = DelegateCopyToBufferHandle;
+      }
+      if (external_delegate_->FreeBufferHandle) {
+        wrapper_delegate_.FreeBufferHandle = DelegateFreeBufferHandle;
+      }
+    }
+  }
+}
+
+ExternalDelegateWrapper::~ExternalDelegateWrapper() {
+  if (external_delegate_ != nullptr) {
+    external_lib_.destroy(external_delegate_);
+  }
+}
+
+}  // namespace
+}  // namespace tflite
+
+// TfLiteExternalDelegateOptionsInsert adds key/value to the given
+// TfLiteExternalDelegateOptions instance.
+TfLiteStatus TfLiteExternalDelegateOptionsInsert(
+    TfLiteExternalDelegateOptions* options, const char* key,
+    const char* value) {
+  if (options->count >= kExternalDelegateMaxOptions) {
+    return kTfLiteError;
+  }
+  options->keys[options->count] = key;
+  options->values[options->count] = value;
+  options->count++;
+  return kTfLiteOk;
+}
+
+TfLiteExternalDelegateOptions TfLiteExternalDelegateOptionsDefault(
+    const char* lib_path) {
+  // As 'keys' and 'values' don't need to be set here, using designated
+  // initializers may cause a compiling error as "non-trivial designated
+  // initializers not supported" by some compiler.
+  TfLiteExternalDelegateOptions options;
+  options.lib_path = lib_path;
+  options.count = 0;
+  options.insert = TfLiteExternalDelegateOptionsInsert;
+  return options;
+}
+
+TfLiteDelegate* TfLiteExternalDelegateCreate(
+    const TfLiteExternalDelegateOptions* options) {
+  auto external_delegate_wrapper =
+      std::make_unique<tflite::ExternalDelegateWrapper>(options);
+  if (external_delegate_wrapper->tflite_external_delegate() != nullptr) {
+    return external_delegate_wrapper.release()->tflite_wrapper_delegate();
+  }
+  return nullptr;
+}
+
+void TfLiteExternalDelegateDelete(TfLiteDelegate* delegate) {
+  delete tflite::GetExternalDelegateWrapper(delegate);
+}

tensorflow/lite/delegates/external/external_delegate.h

@@ -0,0 +1,57 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_LITE_DELEGATES_EXTERNAL_EXTERNAL_DELEGATE_H_
+#define TENSORFLOW_LITE_DELEGATES_EXTERNAL_EXTERNAL_DELEGATE_H_
+
+#include "tensorflow/lite/core/c/common.h"
+
+#ifdef __cplusplus
+extern "C" {
+#endif  // __cplusplus
+
+// TfLiteExternalDelegateOptions is a structure of key/value options to create
+// an external delegate.
+#define kExternalDelegateMaxOptions 256
+typedef struct TfLiteExternalDelegateOptions {
+  const char* lib_path;
+  int count;
+  const char* keys[kExternalDelegateMaxOptions];
+  const char* values[kExternalDelegateMaxOptions];
+  TfLiteStatus (*insert)(struct TfLiteExternalDelegateOptions* options,
+                         const char* key, const char* value);
+} TfLiteExternalDelegateOptions;
+
+// Insert key/value to the options.
+TfLiteStatus TfLiteExternalDelegateOptionsInsert(
+    TfLiteExternalDelegateOptions* options, const char* key, const char* value);
+
+// Populates TfLiteExternalDelegateOptions with the given shared library path.
+TfLiteExternalDelegateOptions TfLiteExternalDelegateOptionsDefault(
+    const char* lib_path);
+
+// Creates a new delegate instance that need to be destroyed with
+// `TfLiteExternalDelegateDelete` when delegate is no longer used by TFLite.
+TfLiteDelegate* TfLiteExternalDelegateCreate(
+    const TfLiteExternalDelegateOptions* options);
+
+// Destroys a delegate created with `TfLiteExternalDelegateCreate` call.
+void TfLiteExternalDelegateDelete(TfLiteDelegate* delegate);
+
+#ifdef __cplusplus
+}
+#endif  // __cplusplus
+
+#endif  // TENSORFLOW_LITE_DELEGATES_EXTERNAL_EXTERNAL_DELEGATE_H_

tensorflow/lite/delegates/external/external_delegate_interface.h

@@ -0,0 +1,79 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_DELEGATES_EXTERNAL_EXTERNAL_DELEGATE_INTERFACE_H_
+#define TENSORFLOW_LITE_DELEGATES_EXTERNAL_EXTERNAL_DELEGATE_INTERFACE_H_
+
+#include "tensorflow/lite/c/common.h"
+
+// This header file declares the interface that external delegate shared
+// libraries need to implement.  The functions declared here are not defined
+// in TF Lite itself -- this just declares the interface to functions that
+// are defined elsewhere, in a shared library that TfLiteExternalDelegate will
+// dynamically load.
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+// Define TFL_EXTERNAL_DELEGATE_EXPORT macro to export an external delegate API
+// function properly with a shared library.
+#ifdef SWIG
+#define TFL_EXTERNAL_DELEGATE_EXPORT
+#else  // !defined SWIG
+#ifdef _WIN32
+// On Windows, the TFL_EXTERNAL_DELEGATE_COMPILE_LIBRARY macro should be
+// defined when _building_ an external delegate shared library, but should not
+// be defined when _using_ an external delegate shared library.
+#ifdef TFL_EXTERNAL_DELEGATE_COMPILE_LIBRARY
+#define TFL_EXTERNAL_DELEGATE_EXPORT __declspec(dllexport)
+#else  // !defined TFL_EXTERNAL_DELEGATE_COMPILE_LIBRARY
+// We may not actually need dllimport,
+// since the symbols will looked up dynamically?
+#define TFL_EXTERNAL_DELEGATE_EXPORT __declspec(dllimport)
+#endif  // !defined TFL_EXTERNAL_DELEGATE_COMPILE_LIBRARY
+#else   // !defined _WIN32
+#define TFL_EXTERNAL_DELEGATE_EXPORT __attribute__((visibility("default")))
+#endif  // !defined _WIN32
+#endif  // !defined SWIG
+
+// Creates a delegate object based on provided key-value options.
+//
+// The delegate is initialized using the option settings specified by the
+// names in `options_keys` and the corresponding values in `options_values`,
+// which are both arrays of length `num_options` of NUL-terminated C strings.
+// This function *should not* modify those arrays, but the caller must not rely
+// on that. `options_keys` and `options_values` may be null if `num_options` is
+// zero.
+//
+// On success, returns a non-null value that should be deallocated with
+// tflite_plugin_destroy_delegate when no longer needed.
+// On failure, returns NULL to indicate an error, with the detailed information
+// reported by calling `report_error` if provided.
+extern TFL_EXTERNAL_DELEGATE_EXPORT TfLiteDelegate*
+tflite_plugin_create_delegate(const char* const* options_keys,
+                              const char* const* options_values,
+                              size_t num_options,
+                              void (*report_error)(const char*));
+
+// Destroys a delegate object that was created by tflite_plugin_create_delegate.
+// Calling this with nullptr as the argument value is allowed and has no effect.
+extern TFL_EXTERNAL_DELEGATE_EXPORT void tflite_plugin_destroy_delegate(
+    TfLiteDelegate* delegate);
+
+#ifdef __cplusplus
+}  // extern "C"
+#endif
+
+#endif  // TENSORFLOW_LITE_DELEGATES_EXTERNAL_EXTERNAL_DELEGATE_INTERFACE_H_

tensorflow/lite/delegates/flex/BUILD

@@ -0,0 +1,438 @@
+load("@rules_cc//cc:cc_library.bzl", "cc_library")
+load(
+    "//tensorflow:tensorflow.bzl",
+    "if_android",
+    "if_mobile",
+    "if_not_mobile",
+    "tf_cc_test",
+    "tf_opts_nortti_if_lite_protos",
+    "tf_opts_nortti_if_mobile",
+)
+load("//tensorflow:tensorflow.default.bzl", "get_compatible_with_portable")
+load("//tensorflow/lite:build_def.bzl", "tflite_copts")
+load("//tensorflow/lite:special_rules.bzl", "internal_visibility_allowlist")
+load("//tensorflow/lite/delegates/flex:build_def.bzl", "tflite_flex_cc_library", "tflite_flex_shared_library")
+
+default_visibility = [
+    "//tensorflow/compiler/mlir/lite:__subpackages__",
+    "//tensorflow/lite/android:__subpackages__",
+    "//tensorflow/lite/toco/tflite:__subpackages__",
+]
+
+#
+# This is a TF Lite delegate that is powered by TensorFlow's Eager.
+#
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:LICENSE"],
+    default_visibility = default_visibility,
+    licenses = ["notice"],
+)
+
+exports_files([
+    "delegate.h",
+    "exported_symbols.lds",
+    "version_script.lds",
+])
+
+cc_library(
+    name = "buffer_map",
+    srcs = ["buffer_map.cc"],
+    hdrs = ["buffer_map.h"],
+    compatible_with = get_compatible_with_portable(),
+    copts = tf_opts_nortti_if_lite_protos(),
+    deps = [
+        ":buffer_map_util",
+        ":util",
+        "//tensorflow/lite:string",
+        "//tensorflow/lite/core/c:common",
+        "//tensorflow/lite/kernels/internal:compatibility",
+    ] + if_mobile([
+        "//tensorflow/core:portable_tensorflow_lib_lite",
+    ]) + if_not_mobile([
+        "//tensorflow/c:c_api_internal",
+        "//tensorflow/core:framework",
+    ]),
+)
+
+cc_library(
+    name = "buffer_map_util",
+    srcs = ["buffer_map_util.cc"],
+    hdrs = ["buffer_map_util.h"],
+    compatible_with = get_compatible_with_portable(),
+    copts = tf_opts_nortti_if_lite_protos(),
+    deps = [
+        ":util",
+        "//tensorflow/core/platform:tstring",
+        "//tensorflow/lite:string_util",
+        "//tensorflow/lite:util",
+        "//tensorflow/lite/c:c_api_types",
+        "//tensorflow/lite/core/c:common",
+        "//tensorflow/lite/experimental/resource",
+        "@com_google_absl//absl/status",
+        "@eigen_archive//:eigen3",
+    ] + if_mobile([
+        "//tensorflow/core:portable_tensorflow_lib_lite",
+    ]) + if_not_mobile([
+        "//tensorflow/c:c_api_internal",
+        "//tensorflow/core:framework",
+        "//tensorflow/c:tf_tensor_internal",
+        "//tensorflow/core:protos_all_cc",
+        "//tensorflow/core/platform:status",
+    ]),
+)
+
+tf_cc_test(
+    name = "buffer_map_test",
+    size = "small",
+    srcs = ["buffer_map_test.cc"],
+    deps = [
+        ":buffer_map",
+        ":buffer_map_util",
+        ":util",
+        "//tensorflow/core:framework",
+        "//tensorflow/lite:framework",
+        "//tensorflow/lite:string_util",
+        "//tensorflow/lite:util",
+        "//tensorflow/lite/core/c:c_api_types",
+        "//tensorflow/lite/testing:util",
+        "@com_google_absl//absl/status",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
+
+# Define the standard flex delegate library, that pulls in the standard set
+# of TensorFlow ops and kernels, using tflite_flex_cc_library with no
+# models parameter. Custom flex delegate can be defined with
+# tflite_flex_cc_library if the parameter models is provided. Tensorflow
+# user-provided ops could also be supported by passing to additional_deps.
+# Ex:
+# tflite_flex_cc_library(
+#   name = "sample_delegate",
+#   models = ["model1.tflite", "model2.tflite"],
+#   additional_deps = ["your_custom_ops_lib"],
+# )
+tflite_flex_cc_library(
+    name = "delegate",
+    compatible_with = get_compatible_with_portable(),
+    visibility = ["//visibility:public"],
+)
+
+# Compared to the library above, this one doesn't define a strong symbol for
+# AcquireFlexDelegate(). This is useful if one doesn't want the default flex
+# delegate to be automatically applied when building the interpreter.
+tflite_flex_cc_library(
+    name = "delegate_without_symbol",
+    link_symbol = False,
+    visibility = ["//visibility:public"],
+)
+
+# Shared lib target for convenience, pulls in the standard set of TensorFlow
+# ops and kernels. The output library name is platform dependent:
+#   - Linux/Android: `libtensorflowlite_flex.so`
+#   - Mac: `libtensorflowlite_flex.dylib`
+#   - Windows: `tensorflowlite_flex.dll`
+tflite_flex_shared_library(
+    name = "tensorflowlite_flex",
+)
+
+cc_library(
+    name = "delegate_symbol",
+    srcs = [
+        "delegate_symbol.cc",
+    ],
+    compatible_with = get_compatible_with_portable(),
+    copts = tflite_copts(),
+    visibility = ["//visibility:public"],
+    deps = [
+        ":delegate_only_runtime",
+        "//tensorflow/lite/core/c:c_api_types",
+    ],
+    alwayslink = 1,
+)
+
+# Delegate implementation that does *not* pull in the standard set of TensorFlow
+# ops and kernels.
+cc_library(
+    name = "delegate_only_runtime",
+    srcs = [
+        "delegate.cc",
+        "kernel.cc",
+        "kernel.h",
+    ],
+    hdrs = [
+        "delegate.h",
+    ],
+    compatible_with = get_compatible_with_portable(),
+    copts = tflite_copts() + tf_opts_nortti_if_mobile(),
+    visibility = ["//visibility:public"],
+    deps = [
+        ":buffer_map",
+        ":delegate_data",
+        ":tflite_subgraph_execute",
+        ":util",
+        "//tensorflow/core:session_options",
+        "//tensorflow/core/tfrt/fallback:op_kernel_runner",
+        "//tensorflow/lite:kernel_api",
+        "//tensorflow/lite:macros",
+        "//tensorflow/lite:minimal_logging",
+        "//tensorflow/lite:string",
+        "//tensorflow/lite:string_util",
+        "//tensorflow/lite:util",
+        "//tensorflow/lite/core:subgraph",
+        "//tensorflow/lite/core/api",
+        "//tensorflow/lite/core/c:common",
+        "//tensorflow/lite/delegates/utils:simple_delegate",
+        "//tensorflow/lite/kernels:kernel_util",
+        "@com_google_absl//absl/container:flat_hash_set",
+        "@com_google_absl//absl/container:inlined_vector",
+        "@com_google_absl//absl/log:check",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/strings",
+        "@flatbuffers",
+    ] + if_mobile([
+        "//tensorflow/core:portable_tensorflow_lib_lite",
+    ]) + if_not_mobile([
+        "//tensorflow/core/common_runtime/eager:context",
+        "//tensorflow/core:lib",
+        "//tensorflow/core:protos_all_cc",
+        "//tensorflow/core:framework",
+    ]),
+    alwayslink = 1,
+)
+
+tf_cc_test(
+    name = "delegate_test",
+    size = "small",
+    srcs = ["delegate_test.cc"],
+    tags = [
+        "no_gpu",  # GPU + flex is not officially supported.
+    ],
+    deps = [
+        ":delegate",
+        ":test_util",
+        "//tensorflow/lite:shared_library",
+        "//tensorflow/lite/kernels:test_util",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
+
+cc_library(
+    name = "delegate_data",
+    srcs = ["delegate_data.cc"],
+    hdrs = ["delegate_data.h"],
+    compatible_with = get_compatible_with_portable(),
+    copts = tf_opts_nortti_if_mobile(),
+    visibility = ["//visibility:public"],
+    deps = [
+        ":buffer_map",
+        ":subgraph_resource",
+        ":util",
+        "//tensorflow/lite:cc_api_experimental",
+        "//tensorflow/lite:util",
+        "//tensorflow/lite/core/c:common",
+        "//tensorflow/lite/schema:schema_fbs",
+        "@com_google_absl//absl/memory",
+        "@com_google_absl//absl/strings",
+        "@flatbuffers",
+    ] + if_mobile([
+        "//tensorflow/core:portable_tensorflow_lib_lite",
+    ]) + if_not_mobile([
+        "//tensorflow/core/common_runtime/eager:context",
+        "//tensorflow/core/common_runtime/eager:core_no_xla",
+        "//tensorflow/core:core_cpu",
+        "//tensorflow/core:framework",
+        "//tensorflow/core:lib",
+        "//tensorflow/core:protos_all_cc",
+    ]),
+)
+
+tf_cc_test(
+    name = "delegate_data_test",
+    size = "small",
+    srcs = ["delegate_data_test.cc"],
+    deps = [
+        ":delegate_data",
+        "//tensorflow/core:test",
+        "//tensorflow/core/common_runtime/eager:context",
+        "//tensorflow/core/platform:mutex",
+        "//tensorflow/core/platform:protobuf",
+        "//tensorflow/core/platform:status",
+        "//tensorflow/lite:framework",
+        "//tensorflow/lite/core/api:error_reporter",
+        "//tensorflow/lite/core/c:common",
+        "//tensorflow/lite/kernels:subgraph_test_util",
+        "//tensorflow/lite/testing:util",
+        "@com_google_absl//absl/memory",
+        "@com_google_absl//absl/strings",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
+
+cc_library(
+    name = "subgraph_resource",
+    hdrs = ["subgraph_resource.h"],
+    compatible_with = get_compatible_with_portable(),
+    copts = if_android(["-Wno-private-header"]),
+    deps = [
+        "//tensorflow/core/platform:mutex",
+        "//tensorflow/core/platform:thread_annotations",
+        "//tensorflow/lite:cc_api_experimental",
+        "//tensorflow/lite/core/c:common",
+    ] + if_mobile([
+        "//tensorflow/core:portable_tensorflow_lib_lite",
+    ]) + if_not_mobile([
+        "//tensorflow/core:framework",
+        "//tensorflow/core:lib",
+    ]),
+)
+
+tf_cc_test(
+    name = "kernel_test",
+    size = "small",
+    srcs = [
+        "kernel.h",
+        "kernel_test.cc",
+    ],
+    tags = ["no_gpu"],  # GPU + flex is not officially supported.
+    deps = [
+        ":delegate",
+        ":delegate_data",
+        ":test_util",
+        "//tensorflow/core/platform:status",
+        "//tensorflow/core/tfrt/fallback:op_kernel_runner",
+        "//tensorflow/lite/core/c:c_api_types",
+        "//tensorflow/lite/core/c:common",
+        "//tensorflow/lite/delegates/utils:simple_delegate",
+        "//tensorflow/lite/kernels:kernel_util",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
+
+cc_library(
+    name = "test_util",
+    testonly = True,
+    srcs = ["test_util.cc"],
+    hdrs = ["test_util.h"],
+    visibility = internal_visibility_allowlist(),
+    deps = [
+        "//tensorflow/c:c_api_internal",
+        "//tensorflow/lite:string",
+        "//tensorflow/lite/kernels:test_util",
+        "@com_google_absl//absl/memory",
+        "@flatbuffers",
+    ],
+)
+
+cc_library(
+    name = "util",
+    srcs = ["util.cc"],
+    hdrs = ["util.h"],
+    compatible_with = get_compatible_with_portable(),
+    #TODO(b/206038955): Consider restrict the visibility to '//third_party/fcp/client:__subpackages__'.
+    visibility = ["//visibility:public"],
+    deps = [
+        "//tensorflow/core:protos_all_cc",
+        "//tensorflow/lite:kernel_api",
+        "//tensorflow/lite:string_util",
+        "//tensorflow/lite:util",
+        "//tensorflow/lite/c:common",
+        "//tensorflow/lite/core/c:c_api_types",
+        "//tensorflow/lite/core/c:common",
+        "//tensorflow/lite/kernels/internal:tensor",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/strings:str_format",
+    ] + if_mobile([
+        "//tensorflow/core:portable_tensorflow_lib_lite",
+    ]) + if_not_mobile([
+        "//tensorflow/c:c_api_internal",
+        "//tensorflow/core:lib",
+        "//tensorflow/core:framework",
+        "//tensorflow/core/protobuf:error_codes_proto_impl_cc",
+    ]),
+)
+
+tf_cc_test(
+    name = "util_test",
+    size = "small",
+    srcs = ["util_test.cc"],
+    deps = [
+        ":util",
+        "//tensorflow/c:tf_datatype",
+        "//tensorflow/core:framework",
+        "//tensorflow/core:portable_gif_internal",
+        "//tensorflow/core:protos_all_cc",
+        "//tensorflow/core/platform:errors",
+        "//tensorflow/core/platform:status",
+        "//tensorflow/core/platform:tstring",
+        "//tensorflow/core/protobuf:error_codes_proto_impl_cc",
+        "//tensorflow/lite:string",
+        "//tensorflow/lite:string_util",
+        "//tensorflow/lite:util",
+        "//tensorflow/lite/core/c:c_api_types",
+        "//tensorflow/lite/core/c:common",
+        "@com_google_absl//absl/status",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
+
+tf_cc_test(
+    name = "allowlisted_flex_ops_test",
+    size = "small",
+    srcs = [
+        "allowlisted_flex_ops_test.cc",
+    ],
+    extra_copts = if_android(["-Wno-private-header"]),
+    deps = [
+        ":delegate",
+        "//tensorflow/compiler/mlir/lite/delegates/flex:allowlisted_flex_ops_lib",
+        "@com_google_googletest//:gtest_main",
+    ] + if_mobile([
+        "//tensorflow/core:portable_tensorflow_lib_lite",
+    ]) + if_not_mobile([
+        "//tensorflow/core:framework",
+    ]),
+)
+
+# Alias to support selective build of image ops.
+# TODO(b/163285312): Remove after tensorflow/core refactoring completed.
+cc_library(
+    name = "portable_images_lib",
+    visibility = ["//visibility:public"],
+    deps = [
+        "//tensorflow/core:portable_gif_internal",
+        "//tensorflow/core:portable_jpeg_internal",
+        "//tensorflow/core/lib/jxl:jxl_io",
+        "//tensorflow/core/lib/png:png_io",
+        "//tensorflow/core/lib/webp:webp_io",
+    ],
+)
+
+cc_library(
+    name = "tflite_subgraph_execute",
+    srcs = ["tflite_subgraph_execute.cc"],
+    compatible_with = get_compatible_with_portable(),
+    copts = tf_opts_nortti_if_mobile(),
+    deps = [
+        ":buffer_map_util",
+        ":subgraph_resource",
+        ":util",
+        "//tensorflow/lite:cc_api_experimental",
+        "//tensorflow/lite:string_util",
+        "//tensorflow/lite/core/c:c_api_types",
+        "//tensorflow/lite/core/c:common",
+        "//tensorflow/lite/kernels:builtin_ops",
+        "//tensorflow/lite/kernels:kernel_util",
+        "//tensorflow/lite/kernels/internal:tensor",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/strings:str_format",
+    ] + if_mobile([
+        "//tensorflow/core:portable_tensorflow_lib_lite",
+    ]) + if_not_mobile([
+        "//tensorflow/core:framework",
+        "//tensorflow/core:lib",
+        "//tensorflow/c:tf_tensor_internal",
+    ]),
+    alwayslink = 1,
+)

tensorflow/lite/delegates/flex/allowlisted_flex_ops_test.cc

@@ -0,0 +1,67 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include "tensorflow/compiler/mlir/lite/delegates/flex/allowlisted_flex_ops.h"
+
+#include <set>
+#include <string>
+
+#include <gtest/gtest.h>
+#include "tensorflow/compiler/mlir/lite/delegates/flex/allowlisted_flex_ops_internal.h"
+#include "tensorflow/core/framework/op_kernel.h"
+
+namespace tflite {
+namespace flex {
+
+// Get all cpu kernels registered in Tensorflow.
+std::set<std::string> GetAllCpuKernels() {
+  auto is_cpu_kernel = [](const tensorflow::KernelDef& def) {
+    return (def.device_type() == "CPU" || def.device_type() == "DEFAULT");
+  };
+
+  tensorflow::KernelList kernel_list =
+      tensorflow::GetFilteredRegisteredKernels(is_cpu_kernel);
+  std::set<std::string> result;
+
+  for (int i = 0; i < kernel_list.kernel_size(); ++i) {
+    tensorflow::KernelDef kernel_def = kernel_list.kernel(i);
+    result.insert(kernel_def.op());
+  }
+  return result;
+}
+
+// Test if every flex op has their kernel included in the flex delegate library.
+// This test must be run on both Linux and Android.
+TEST(AllowlistedFlexOpsTest, EveryOpHasKernel) {
+  const std::set<std::string>& allowlist = GetFlexAllowlist();
+  std::set<std::string> all_kernels = GetAllCpuKernels();
+
+  for (const std::string& op_name : allowlist) {
+    EXPECT_EQ(all_kernels.count(op_name), 1)
+        << op_name << " op is added to flex allowlist "
+        << "but its kernel is not found.";
+  }
+}
+
+TEST(TfTextUtilsTest, TestFlexOpAllowed) {
+  // Expect false since ConstrainedSequence kernel is not registered.
+  EXPECT_FALSE(IsAllowedTFTextOpForFlex("ConstrainedSequence"));
+}
+
+TEST(TfTextUtilsTest, TestFlexOpNotAllowed) {
+  EXPECT_FALSE(IsAllowedTFTextOpForFlex("ngrams"));
+}
+
+}  // namespace flex
+}  // namespace tflite

tensorflow/lite/delegates/flex/buffer_map.cc

@@ -0,0 +1,57 @@
+/* Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include "tensorflow/lite/delegates/flex/buffer_map.h"
+
+#include <utility>
+
+#include "tensorflow/c/c_api_internal.h"
+#include "tensorflow/lite/delegates/flex/buffer_map_util.h"
+#include "tensorflow/lite/delegates/flex/util.h"
+#include "tensorflow/lite/kernels/internal/compatibility.h"
+#include "tensorflow/lite/string_type.h"
+
+namespace tflite {
+namespace flex {
+
+BufferMap::BufferMap() {}
+
+BufferMap::~BufferMap() {}
+
+bool BufferMap::HasTensor(int tensor_index) const {
+  return id_to_tensor_.count(tensor_index) != 0;
+}
+
+tensorflow::Tensor BufferMap::GetTensor(int tensor_index) const {
+  return id_to_tensor_.at(tensor_index);
+}
+
+const tensorflow::Tensor* BufferMap::GetTensorPtr(int tensor_index) const {
+  auto& tensor = id_to_tensor_.at(tensor_index);
+  return &tensor;
+}
+
+void BufferMap::SetFromTfLite(int tensor_index, const TfLiteTensor* tensor,
+                              bool allow_reusing) {
+  TFLITE_CHECK(
+      SetTfTensorFromTfLite(tensor, &id_to_tensor_[tensor_index], allow_reusing)
+          .ok());
+}
+
+void BufferMap::SetFromTensorFlow(int tensor_index, tensorflow::Tensor tensor) {
+  id_to_tensor_[tensor_index] = std::move(tensor);
+}
+
+}  // namespace flex
+}  // namespace tflite

tensorflow/lite/delegates/flex/buffer_map.h

@@ -0,0 +1,72 @@
+/* Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_DELEGATES_FLEX_BUFFER_MAP_H_
+#define TENSORFLOW_LITE_DELEGATES_FLEX_BUFFER_MAP_H_
+
+#include <map>
+
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/lite/core/c/common.h"
+
+namespace tflite {
+namespace flex {
+
+// Maps a TF Lite tensor index into a TensorFlow tensor.
+//
+// The TF Lite interpreter assigns integer indices to each of its tensors, but
+// the Flex delegate deals in terms of TensorFlow tensors. This class maps
+// from indices to tensors and allows the creation of new tensors to be
+// associated with a given index.
+class BufferMap {
+ public:
+  BufferMap();
+  ~BufferMap();
+
+  // Returns true if the given 'tensor_index' has a corresponding
+  // tensorflow::Tensor.
+  bool HasTensor(int tensor_index) const;
+
+  // Returns the tensorflow::Tensor associated with the given 'tensor_index'.
+  // Precondition: HasTensor() is true.
+  tensorflow::Tensor GetTensor(int tensor_index) const;
+
+  // Returns the const pointer to tensorflow::Tensor associated with the given
+  // 'tensor_index'.
+  // Precondition: HasTensor() is true.
+  const tensorflow::Tensor* GetTensorPtr(int tensor_index) const;
+
+  // Associates the given tensorflow::Tensor with the given 'tensor_index'.
+  // Note that TensorFlow Tensors share data buffers, so this method is only a
+  // shallow copy.
+  void SetFromTensorFlow(int tensor_index, tensorflow::Tensor tensor);
+
+  // Same as above but creates a new tensorflow::Tensor with a copy of the
+  // given TfLiteTensor's data. If `allow_reusing=false`, then we explicitly
+  // disallow reusing the TF Lite tensor buffer when constructing the new
+  // tensorflow Tensor.
+  void SetFromTfLite(int tensor_index, const TfLiteTensor* tensor,
+                     bool allow_reusing = true);
+
+ private:
+  // Mapping from TL Lite tensor ID to TensorFlow's Tensor. All tensors that
+  // are inputs or outputs of a subgraph will be added here, irrespective of
+  // whether their data are managed by TF Lite or TensorFlow.
+  std::map<int, tensorflow::Tensor> id_to_tensor_;
+};
+
+}  // namespace flex
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_DELEGATES_FLEX_BUFFER_MAP_H_

tensorflow/lite/delegates/flex/buffer_map_test.cc

@@ -0,0 +1,318 @@
+/* Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include "tensorflow/lite/delegates/flex/buffer_map.h"
+
+#include <sys/types.h>
+
+#include <functional>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "absl/status/status.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/lite/core/c/c_api_types.h"
+#include "tensorflow/lite/delegates/flex/buffer_map_util.h"
+#include "tensorflow/lite/delegates/flex/util.h"
+#include "tensorflow/lite/interpreter.h"
+#include "tensorflow/lite/string_util.h"
+#include "tensorflow/lite/testing/util.h"
+#include "tensorflow/lite/util.h"
+
+namespace tflite {
+namespace flex {
+namespace {
+
+using ::testing::ElementsAre;
+using ::testing::HasSubstr;
+using ::testing::status::StatusIs;
+
+// A bit of RAII to simplify handling of TfLiteTensors in the tests.
+using UniqueTfLiteTensor =
+    std::unique_ptr<TfLiteTensor, std::function<void(TfLiteTensor*)>>;
+
+template <typename T>
+UniqueTfLiteTensor MakeLiteTensor(const std::vector<int>& shape,
+                                  const std::vector<T>& data) {
+  auto tensor = UniqueTfLiteTensor(new TfLiteTensor(), [](TfLiteTensor* t) {
+    TfLiteTensorDataFree(t);
+    TfLiteIntArrayFree(t->dims);
+    delete t;
+  });
+  tensor->allocation_type = kTfLiteDynamic;
+  tensor->type = typeToTfLiteType<T>();
+  tensor->dims = ConvertVectorToTfLiteIntArray(shape);
+  TfLiteTensorRealloc(data.size() * sizeof(T), tensor.get());
+  memcpy(tensor->data.raw, data.data(), data.size() * sizeof(T));
+  return tensor;
+}
+
+template <>
+UniqueTfLiteTensor MakeLiteTensor<string>(const std::vector<int>& shape,
+                                          const std::vector<string>& data) {
+  auto tensor = UniqueTfLiteTensor(new TfLiteTensor(), [](TfLiteTensor* t) {
+    TfLiteTensorDataFree(t);
+    TfLiteIntArrayFree(t->dims);
+    delete t;
+  });
+  tensor->allocation_type = kTfLiteDynamic;
+  tensor->type = typeToTfLiteType<string>();
+  tensor->dims = ConvertVectorToTfLiteIntArray(shape);
+  TfLiteTensorRealloc(data.size() * sizeof(string), tensor.get());
+
+  DynamicBuffer b;
+  for (const string& s : data) {
+    b.AddString(s.data(), s.size());
+  }
+  b.WriteToTensor(tensor.get(), ConvertVectorToTfLiteIntArray(shape));
+  return tensor;
+}
+
+template <typename T>
+tensorflow::Tensor MakeTensor(const std::vector<int64_t>& shape,
+                              const std::vector<T>& data,
+                              tensorflow::DataType dtype) {
+  tensorflow::Tensor tensor(dtype, tensorflow::TensorShape(shape));
+  memcpy(tensor.data(), data.data(), data.size() * sizeof(T));
+  return tensor;
+}
+
+std::vector<int64_t> GetTensorShape(const tensorflow::Tensor& t) {
+  std::vector<int64_t> shape(t.dims());
+  for (int i = 0; i < t.dims(); ++i) {
+    shape[i] = t.dim_size(i);
+  }
+  return shape;
+}
+
+template <typename T>
+std::vector<T> GetTensorData(const tensorflow::Tensor& t) {
+  const T* data = t.flat<T>().data();
+  return std::vector<T>(data, data + t.NumElements());
+}
+
+TEST(BufferMapTest, EmptyBuffer) {
+  BufferMap buffer_map;
+  EXPECT_FALSE(buffer_map.HasTensor(0));
+}
+
+TEST(BufferMapTest, SetFromTfLite) {
+  BufferMap buffer_map;
+
+  UniqueTfLiteTensor t =
+      MakeLiteTensor<float>({1, 2, 1, 3}, {0, 0, 0, 0.123f, 0, 0});
+  buffer_map.SetFromTfLite(0, t.get());
+  ASSERT_TRUE(buffer_map.HasTensor(0));
+
+  EXPECT_THAT(GetTensorData<float>(buffer_map.GetTensor(0)),
+              ElementsAre(0, 0, 0, 0.123f, 0, 0));
+
+  // Also check details of the tensor.
+  tensorflow::Tensor out_tensor = buffer_map.GetTensor(0);
+  ASSERT_EQ(out_tensor.dtype(), tensorflow::DT_FLOAT);
+  ASSERT_EQ(out_tensor.NumElements(), 6);
+  ASSERT_THAT(GetTensorShape(out_tensor), ElementsAre(1, 2, 1, 3));
+}
+
+TEST(BufferMapTest, SetFromTfLiteString) {
+  BufferMap buffer_map;
+
+  UniqueTfLiteTensor t =
+      MakeLiteTensor<string>({1, 2, 1, 3}, {"", "", "", "str1", "", ""});
+  buffer_map.SetFromTfLite(0, t.get());
+  ASSERT_TRUE(buffer_map.HasTensor(0));
+
+  EXPECT_THAT(GetTensorData<tensorflow::tstring>(buffer_map.GetTensor(0)),
+              ElementsAre("", "", "", "str1", "", ""));
+
+  // Also check details of the tensor.
+  tensorflow::Tensor out_tensor = buffer_map.GetTensor(0);
+  ASSERT_EQ(out_tensor.dtype(), tensorflow::DT_STRING);
+  ASSERT_EQ(out_tensor.NumElements(), 6);
+  ASSERT_THAT(GetTensorShape(out_tensor), ElementsAre(1, 2, 1, 3));
+}
+
+TEST(BufferMapTest, SetFromTfLiteTwice) {
+  UniqueTfLiteTensor t1 =
+      MakeLiteTensor<float>({1, 2, 1, 3}, {0, 0, 0, 0.123f, 0, 0});
+  UniqueTfLiteTensor t2 =
+      MakeLiteTensor<int>({1, 2, 4}, {0, 0, 0, 3, 0, 0, 1, 2});
+
+  BufferMap buffer_map;
+  buffer_map.SetFromTfLite(0, t1.get());
+  buffer_map.SetFromTfLite(0, t2.get());
+
+  EXPECT_THAT(GetTensorData<int>(buffer_map.GetTensor(0)),
+              ElementsAre(0, 0, 0, 3, 0, 0, 1, 2));
+}
+
+TEST(BufferMapTest, SetFromTfLiteStringTwice) {
+  UniqueTfLiteTensor t1 =
+      MakeLiteTensor<float>({1, 2, 1, 3}, {0, 0, 0, 0.123f, 0, 0});
+  UniqueTfLiteTensor t2 =
+      MakeLiteTensor<string>({1, 2, 4}, {"", "", "", "s3", "", "", "s1", "s2"});
+
+  BufferMap buffer_map;
+  buffer_map.SetFromTfLite(0, t1.get());
+  buffer_map.SetFromTfLite(0, t2.get());
+
+  EXPECT_THAT(GetTensorData<tensorflow::tstring>(buffer_map.GetTensor(0)),
+              ElementsAre("", "", "", "s3", "", "", "s1", "s2"));
+}
+
+TEST(BufferMapTest, SetFromTfLiteBuiltinResource) {
+  BufferMap buffer_map;
+
+  // Constructs a fake resource tensor.
+  auto tensor = UniqueTfLiteTensor(new TfLiteTensor(), [](TfLiteTensor* t) {
+    TfLiteTensorDataFree(t);
+    TfLiteIntArrayFree(t->dims);
+    delete t;
+  });
+  tensor->allocation_type = kTfLiteDynamic;
+  tensor->type = kTfLiteResource;
+  tensor->dims = ConvertVectorToTfLiteIntArray({1});
+  TfLiteTensorRealloc(sizeof(int32_t), tensor.get());
+  tensor->delegate = nullptr;
+  tensor->data.i32[0] = 1;
+
+  buffer_map.SetFromTfLite(0, tensor.get());
+  // Also check details of the tensor.
+  tensorflow::Tensor out_tensor = buffer_map.GetTensor(0);
+  ASSERT_EQ(out_tensor.dtype(), tensorflow::DT_RESOURCE);
+  ASSERT_EQ(out_tensor.NumElements(), 1);
+  tensorflow::ResourceHandle handle =
+      out_tensor.flat<tensorflow::ResourceHandle>()(0);
+  EXPECT_EQ(handle.name(), "tflite_resource_variable:1");
+}
+
+TEST(BufferMapTest, SetTfTensorFromTfLiteResourceOrVariantInvalidData) {
+  auto tensor = UniqueTfLiteTensor(new TfLiteTensor(), [](TfLiteTensor* t) {
+    TfLiteTensorDataFree(t);
+    TfLiteIntArrayFree(t->dims);
+    delete t;
+  });
+  tensor->allocation_type = kTfLiteDynamic;
+  tensor->type = kTfLiteVariant;
+  tensor->dims = ConvertVectorToTfLiteIntArray({1});
+  TfLiteTensorRealloc(tflite::flex::kTensorflowResourceTensorBytes,
+                      tensor.get());
+  tensor->delegate = nullptr;
+  tensor->data_is_stale = false;
+
+  tensorflow::Tensor out_tensor;
+  EXPECT_THAT(
+      SetTfTensorFromTfLite(tensor.get(), &out_tensor),
+      StatusIs(absl::StatusCode::kInvalidArgument,  // NOLINT
+               HasSubstr("Input tensor has resource or variant type but "
+                         "is not managed by the Flex delegate.")));
+}
+
+TEST(BufferMapTest, SetFromTensorFlow) {
+  tensorflow::Tensor t1 = MakeTensor<float>(
+      {1, 2, 1, 3}, {0, 0, 0, 0.123f, 0, 0}, tensorflow::DT_FLOAT);
+
+  BufferMap buffer_map;
+  buffer_map.SetFromTensorFlow(0, t1);
+
+  EXPECT_THAT(GetTensorData<float>(buffer_map.GetTensor(0)),
+              ElementsAre(0, 0, 0, 0.123f, 0, 0));
+
+  // Also check details of the tensor.
+  tensorflow::Tensor out_tensor = buffer_map.GetTensor(0);
+  ASSERT_EQ(out_tensor.dtype(), tensorflow::DT_FLOAT);
+  ASSERT_EQ(out_tensor.NumElements(), 6);
+  ASSERT_THAT(GetTensorShape(out_tensor), ElementsAre(1, 2, 1, 3));
+}
+
+TEST(BufferMapTest, SetFromTensorFlowTwice) {
+  tensorflow::Tensor t1 = MakeTensor<float>(
+      {1, 2, 1, 3}, {0, 0, 0, 0.123f, 0, 0}, tensorflow::DT_FLOAT);
+  tensorflow::Tensor t2 = MakeTensor<int>({1, 2, 4}, {0, 0, 0, 3, 0, 0, 1, 2},
+                                          tensorflow::DT_INT32);
+  BufferMap buffer_map;
+  buffer_map.SetFromTensorFlow(0, t1);
+  buffer_map.SetFromTensorFlow(0, t2);
+
+  EXPECT_THAT(GetTensorData<int>(buffer_map.GetTensor(0)),
+              ElementsAre(0, 0, 0, 3, 0, 0, 1, 2));
+}
+
+TEST(BufferMapTest, TfLiteOverwritesTensorFlow) {
+  tensorflow::Tensor t1 = MakeTensor<float>(
+      {1, 2, 1, 3}, {0, 0, 0, 0.123f, 0, 0}, tensorflow::DT_FLOAT);
+  UniqueTfLiteTensor t2 =
+      MakeLiteTensor<int>({1, 2, 4}, {0, 0, 0, 3, 0, 0, 1, 2});
+
+  BufferMap buffer_map;
+  buffer_map.SetFromTensorFlow(0, t1);
+  buffer_map.SetFromTfLite(0, t2.get());
+
+  EXPECT_THAT(GetTensorData<int>(buffer_map.GetTensor(0)),
+              ElementsAre(0, 0, 0, 3, 0, 0, 1, 2));
+}
+
+TEST(BufferMapTest, TensorFlowOverwritesTfLite) {
+  tensorflow::Tensor t1 = MakeTensor<float>(
+      {1, 2, 1, 3}, {0, 0, 0, 0.123f, 0, 0}, tensorflow::DT_FLOAT);
+  UniqueTfLiteTensor t2 =
+      MakeLiteTensor<int>({1, 2, 4}, {0, 0, 0, 3, 0, 0, 1, 2});
+  BufferMap buffer_map;
+  buffer_map.SetFromTfLite(0, t2.get());
+  buffer_map.SetFromTensorFlow(0, t1);
+
+  EXPECT_THAT(GetTensorData<float>(buffer_map.GetTensor(0)),
+              ElementsAre(0, 0, 0, 0.123f, 0, 0));
+}
+
+TEST(BufferMapTest, TensorflowBufferReuse) {
+  const int kAllocationSize = 1000;
+  TfLiteTensor tensor;
+  tensor.allocation_type = kTfLiteDynamic;
+  tensor.data.raw = nullptr;
+  TfLiteTensorRealloc(kAllocationSize, &tensor);
+  CHECK(tensor.data.raw);
+  EXPECT_EQ(tensor.bytes, kAllocationSize);
+
+  TfLiteTensorBuffer* tensor_buffer_reused = new TfLiteTensorBuffer(&tensor);
+  // Checks that the underlying buffer is reused.
+  EXPECT_TRUE(tensor_buffer_reused->BufferReusedFromTfLiteTensor());
+  EXPECT_EQ(tensor_buffer_reused->data(), tensor.data.raw);
+  tensor_buffer_reused->Unref();
+
+  TfLiteTensorDataFree(&tensor);
+}
+
+TEST(BufferMapTest, ExplicitlyDisableBufferReuse) {
+  TfLiteTensor tensor;
+  tensor.allocation_type = kTfLiteDynamic;
+  tensor.data.raw = nullptr;
+  TfLiteTensorRealloc(10, &tensor);
+  CHECK(tensor.data.raw);
+  EXPECT_EQ(tensor.bytes, 10);
+
+  TfLiteTensorBuffer* tensor_buffer =
+      new TfLiteTensorBuffer(&tensor, /*=allow_reusing*/ false);
+  // Checks that the underlying buffer is not reused.
+  EXPECT_FALSE(tensor_buffer->BufferReusedFromTfLiteTensor());
+  EXPECT_NE(tensor_buffer->data(), tensor.data.raw);
+  tensor_buffer->Unref();
+
+  TfLiteTensorDataFree(&tensor);
+}
+
+}  // namespace
+}  // namespace flex
+}  // namespace tflite

tensorflow/lite/delegates/flex/buffer_map_util.cc

@@ -0,0 +1,275 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include "tensorflow/lite/delegates/flex/buffer_map_util.h"
+
+#include <cstddef>
+#include <cstdint>
+#include <cstring>
+#include <utility>
+
+#include "absl/status/status.h"
+#include "tensorflow/c/tf_tensor_internal.h"
+#include "tensorflow/core/framework/allocator.h"
+#include "tensorflow/core/framework/log_memory.h"
+#include "tensorflow/core/framework/resource_handle.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/framework/typed_allocator.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/platform/tstring.h"
+#include "tensorflow/lite/c/c_api_types.h"
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/delegates/flex/util.h"
+#include "tensorflow/lite/string_util.h"
+#include "tensorflow/lite/util.h"
+
+namespace tflite {
+namespace flex {
+
+namespace {
+// Returns a boolean to indicate whether we should reuse memory from the
+// TfLiteTensor.
+inline bool ShouldReuseTensorMemory(const TfLiteTensor* tensor) {
+  // TODO(b/205153246): Currently arena-allocated memory could not be reused
+  // since it might be invalid after the original arena grow in size and copied
+  // over to a new memory block.
+  // First check alignment is consistent with Tensorflow.
+  if (EIGEN_MAX_ALIGN_BYTES != 0 &&  // NOLINT(misc-include-cleaner)
+      reinterpret_cast<intptr_t>(tensor->data.raw) %
+          EIGEN_MAX_ALIGN_BYTES) {  // NOLINT(misc-include-cleaner)
+    return false;
+  }
+  return tensor->allocation_type != kTfLiteArenaRw;
+}
+}  // namespace
+
+void BaseTfLiteTensorBuffer::FillAllocationDescription(
+    tensorflow::AllocationDescription* proto) const {
+  int64_t rb = size();
+  proto->set_requested_bytes(rb);
+  proto->set_allocator_name(tensorflow::cpu_allocator()->Name());
+}
+
+void BaseTfLiteTensorBuffer::LogAllocation() {
+  if (tensorflow::LogMemory::IsEnabled() && data() != nullptr) {
+    tensorflow::LogMemory::RecordRawAllocation(
+        "TfLiteTensorBuffer_New",
+        tensorflow::LogMemory::EXTERNAL_TENSOR_ALLOCATION_STEP_ID, size(),
+        data(), tensorflow::cpu_allocator());
+  }
+}
+void BaseTfLiteTensorBuffer::LogDeallocation() {
+  if (tensorflow::LogMemory::IsEnabled() && data() != nullptr) {
+    tensorflow::LogMemory::RecordRawDeallocation(
+        "TfLiteTensorBuffer_Delete",
+        tensorflow::LogMemory::EXTERNAL_TENSOR_ALLOCATION_STEP_ID, data(),
+        tensorflow::cpu_allocator(), false);
+  }
+}
+
+void* TfLiteTensorBuffer::MaybeAllocateTensorflowBuffer(
+    const TfLiteTensor* tensor, bool allow_reusing) const {
+  if (allow_reusing && ShouldReuseTensorMemory(tensor)) {
+    return tensor->data.raw;
+  }
+  return tensorflow::cpu_allocator()->AllocateRaw(
+      EIGEN_MAX_ALIGN_BYTES,  // NOLINT(misc-include-cleaner)
+      tensor->bytes);
+}
+
+TfLiteTensorBuffer::TfLiteTensorBuffer(const TfLiteTensor* tensor,
+                                       bool allow_reusing)
+    : BaseTfLiteTensorBuffer(
+          MaybeAllocateTensorflowBuffer(tensor, allow_reusing)) {
+  len_ = tensor->bytes;
+
+  reused_buffer_from_tflite_ = allow_reusing && ShouldReuseTensorMemory(tensor);
+
+  if (data() && !reused_buffer_from_tflite_) {
+    LogAllocation();
+    std::memcpy(data(), tensor->data.raw, tensor->bytes);
+  }
+}
+
+TfLiteTensorBuffer::~TfLiteTensorBuffer() {
+  if (!reused_buffer_from_tflite_) {
+    LogDeallocation();
+    // Only deallocate tensor memory if it's allocated via Tensorflow's CPU
+    // allocator.
+    tensorflow::cpu_allocator()->DeallocateRaw(data());
+  }
+}
+
+StringTfLiteTensorBuffer::StringTfLiteTensorBuffer(const TfLiteTensor* tensor)
+    : StringTfLiteTensorBuffer(
+          tensor, tensor->data.raw != nullptr ? GetStringCount(tensor) : 0) {}
+
+StringTfLiteTensorBuffer::~StringTfLiteTensorBuffer() {
+  LogDeallocation();
+  tensorflow::TypedAllocator::Deallocate<tensorflow::tstring>(
+      tensorflow::cpu_allocator(), static_cast<tensorflow::tstring*>(data()),
+      num_strings_);
+}
+
+StringTfLiteTensorBuffer::StringTfLiteTensorBuffer(const TfLiteTensor* tensor,
+                                                   int num_strings)
+    : BaseTfLiteTensorBuffer(
+          num_strings != 0
+              ? tensorflow::TypedAllocator::Allocate<tensorflow::tstring>(
+                    tensorflow::cpu_allocator(), num_strings,
+                    tensorflow::AllocationAttributes())
+              : nullptr),
+      num_strings_(num_strings) {
+  LogAllocation();
+
+  if (data()) {
+    tensorflow::tstring* p = static_cast<tensorflow::tstring*>(data());
+    for (size_t i = 0; i < num_strings_; ++p, ++i) {
+      auto ref = GetString(tensor, i);
+      p->assign(ref.str, ref.len);
+    }
+  }
+}
+
+absl::Status SetTfTensorFromTfLite(const TfLiteTensor* tensor,
+                                   tensorflow::Tensor* tf_tensor,
+                                   bool allow_reusing) {
+  if (tensor->type == kTfLiteResource &&
+      tensor->bytes != kTensorflowResourceTensorBytes) {
+    // If this is native TF Lite resource variable, then we create a TF resource
+    // tensor where the tensor handle encodes the identifier of the TF Lite
+    // resource.
+    // This approach assumes that there is only a single model being invoked
+    // via the Interpreter instance, so that the resource IDs won't have any
+    // collisions. If we plan to support concurrent execution in the future, we
+    // should make sure the resource ID being encoded is unique between
+    // different executions.
+    tensorflow::Tensor t(tensorflow::DT_RESOURCE, tensorflow::TensorShape({}));
+    tensorflow::ResourceHandle handle;
+    handle.set_name(TfLiteResourceIdentifier(tensor));
+    t.flat<tensorflow::ResourceHandle>()(0) = handle;
+    *tf_tensor = t;
+    return absl::OkStatus();
+  } else if (IsResourceOrVariant(tensor)) {
+    // Resource and Variant tensors are expected to be managed by the Flex
+    // delegate, which sets up the tensor->data.raw to point to a
+    // tensorflow::Tensor**. If tensor->delegate is nullptr, it means this
+    // tensor is not managed by the Flex delegate, and we cannot interpret its
+    // data as a TensorFlow tensor pointer.
+    if (tensor->delegate == nullptr) {
+      return absl::InvalidArgumentError(  // NOLINT
+          "Input tensor has resource or variant type but is not managed by "
+          "the Flex delegate.");
+    }
+    // TODO(b/179094265): This is an experimental implementation, subject to
+    // change. This can be re-implemented with life cycle management mechanism
+    // like reference counting.
+    // In a different subgraph, it can load the TensorFlow tensor pointer of the
+    // given TensorFlow Lite tensor, which is stored in the `data` field. The
+    // memory management cycle of the shared TensorFlow's tensor will be managed
+    // by the buffer maps since the loaded tensors always will be kept in the
+    // buffer map.
+    //
+    // The life cycle of the pointer will be managed by the reference counting
+    // in the TensorFlow world and the pointer will be freed when all the buffer
+    // maps, who own it, are gone.
+    const tensorflow::Tensor** tf_tensor_ptr =
+        reinterpret_cast<const tensorflow::Tensor**>(tensor->data.raw);
+    *tf_tensor = **tf_tensor_ptr;
+    return absl::OkStatus();
+  }
+
+  tensorflow::TensorShape shape;
+  int num_dims = tensor->dims->size;
+  for (int i = 0; i < num_dims; ++i) {
+    shape.AddDim(tensor->dims->data[i]);
+  }
+  // TODO(b/152916533): We assume this is a new tensor and allocate a new buffer
+  // for it. This is not always the best approach. For example, this might
+  // be a reallocation after resizing tensors. In that case it would be
+  // preferable to somehow reuse the buffer.
+  BaseTfLiteTensorBuffer* buf;
+  if (tensor->type == kTfLiteString) {
+    if (tensor->data.raw != nullptr) {
+      const char* raw_data = static_cast<const char*>(tensor->data.raw);
+      if (tensor->bytes < sizeof(int32_t)) {
+        return absl::InvalidArgumentError(
+            "String tensor buffer too small for string count");
+      }
+      int num_strings = GetStringCount(tensor);
+      if (num_strings < 0 || num_strings > INT32_MAX - 2) {
+        return absl::InvalidArgumentError(
+            "Invalid string count in string tensor");
+      }
+      uint64_t required_bytes =
+          sizeof(int32_t) * (static_cast<uint64_t>(num_strings) + 2);
+      if (tensor->bytes < required_bytes) {
+        return absl::InvalidArgumentError(
+            "String tensor buffer too small for implied structure");
+      }
+
+      auto read_int32 = [](const char* p) {
+        int32_t val;
+        std::memcpy(&val, p, sizeof(int32_t));
+        return val;
+      };
+
+      for (int i = 1; i <= num_strings; ++i) {
+        int32_t offset_i = read_int32(raw_data + i * sizeof(int32_t));
+        if (offset_i < 0) {
+          return absl::InvalidArgumentError("Invalid string start offset");
+        }
+        uint64_t u_offset_i = static_cast<uint64_t>(offset_i);
+        if (u_offset_i < required_bytes ||
+            u_offset_i > static_cast<uint64_t>(tensor->bytes)) {
+          return absl::InvalidArgumentError("Invalid string start offset");
+        }
+        if (i < num_strings) {
+          int32_t offset_i_plus_1 =
+              read_int32(raw_data + (i + 1) * sizeof(int32_t));
+          if (offset_i > offset_i_plus_1) {
+            return absl::InvalidArgumentError("Non-monotonic string offsets");
+          }
+        }
+      }
+      int32_t offset_last =
+          read_int32(raw_data + (num_strings + 1) * sizeof(int32_t));
+      if (offset_last < 0) {
+        return absl::InvalidArgumentError(
+            "Invalid total buffer length in string tensor");
+      }
+      uint64_t u_offset_last = static_cast<uint64_t>(offset_last);
+      int32_t offset_num_strings =
+          read_int32(raw_data + num_strings * sizeof(int32_t));
+      if (offset_last < offset_num_strings || u_offset_last < required_bytes ||
+          u_offset_last > static_cast<uint64_t>(tensor->bytes)) {
+        return absl::InvalidArgumentError(
+            "Invalid total buffer length in string tensor");
+      }
+    }
+    buf = new StringTfLiteTensorBuffer(tensor);
+  } else {
+    buf = new TfLiteTensorBuffer(tensor, allow_reusing);
+  }
+  tensorflow::Tensor t = tensorflow::TensorCApi::MakeTensor(
+      GetTensorFlowDataType(tensor->type), shape, buf);
+  buf->Unref();
+
+  *tf_tensor = std::move(t);
+  return absl::OkStatus();
+}
+
+}  // namespace flex
+}  // namespace tflite

tensorflow/lite/delegates/flex/buffer_map_util.h

@@ -0,0 +1,110 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_DELEGATES_FLEX_BUFFER_MAP_UTIL_H_
+#define TENSORFLOW_LITE_DELEGATES_FLEX_BUFFER_MAP_UTIL_H_
+
+#include "tensorflow/core/framework/allocation_description.pb.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/platform/status.h"
+#include "tensorflow/lite/core/c/common.h"
+
+namespace tflite {
+namespace flex {
+
+// A tensor buffer that is allocated, deallocated and populated by TF Lite.
+class BaseTfLiteTensorBuffer : public tensorflow::TensorBuffer {
+  using tensorflow::TensorBuffer::TensorBuffer;
+
+  inline TensorBuffer* root_buffer() override { return this; }
+
+  void FillAllocationDescription(
+      tensorflow::AllocationDescription* proto) const override;
+
+  // Prevents input forwarding from mutating this buffer.
+  inline bool OwnsMemory() const override { return false; }
+
+ protected:
+  void LogAllocation();
+  void LogDeallocation();
+};
+
+// A tensor buffer for most data types. Numeric types have exactly the same
+// representation in TFLITE and TF, so we just need use memcpy().
+// For memory efficiency, this TensorBuffer can possibly reuse memory from the
+// TfLiteTensor, hence caller should ensure that the TfLiteTensor always outlive
+// this TensorBuffer.
+class TfLiteTensorBuffer : public BaseTfLiteTensorBuffer {
+ public:
+  // If `allow_reusing=false`, then the tensor buffer won't be reused from the
+  // TfLiteTensor.
+  explicit TfLiteTensorBuffer(const TfLiteTensor* tensor,
+                              bool allow_reusing = true);
+
+  ~TfLiteTensorBuffer() override;
+
+  inline size_t size() const override { return len_; }
+
+  // Indicates that `TfLiteTensorBuffer` is responsible for deallocating its
+  // underlying buffer. This buffer must have been allocated by
+  // `tensorflow::cpu_allocator`
+  inline void TakeOwnershipOfBuffer() { reused_buffer_from_tflite_ = false; }
+
+  inline bool BufferReusedFromTfLiteTensor() const {
+    return reused_buffer_from_tflite_;
+  }
+
+  // This function will check if the underlying buffer in `tensor` can be
+  // reused by the tensorflow::Tensor. If it can reuse, it will return
+  // `tensor->data.raw`, otherwise it will create new tensor buffer using
+  // tensorflow's CPU allocator.
+  // TODO(b/205153246): Also consider reusing memory to avoid copying from
+  // tensorflow::Tensor to TfLiteTensor.
+  void* MaybeAllocateTensorflowBuffer(const TfLiteTensor* tensor,
+                                      bool allow_reusing) const;
+
+ private:
+  size_t len_;
+  bool reused_buffer_from_tflite_;
+};
+
+// A string buffer. TFLITE string tensor format is different than
+// TF's so we need perform the conversion here.
+class StringTfLiteTensorBuffer : public BaseTfLiteTensorBuffer {
+ public:
+  explicit StringTfLiteTensorBuffer(const TfLiteTensor* tensor);
+
+  ~StringTfLiteTensorBuffer() override;
+
+  inline size_t size() const override {
+    return num_strings_ * sizeof(tensorflow::tstring);
+  }
+
+ private:
+  StringTfLiteTensorBuffer(const TfLiteTensor* tensor, int num_strings);
+
+  int num_strings_;
+};
+
+// Sets the `tensorflow::Tensor` content from `TfLiteTensor` object. If
+// `allow_reusing=false`, then we explicitly disallow reusing the TF Lite
+// tensor buffer when constructing the new tensorflow Tensor.
+absl::Status SetTfTensorFromTfLite(const TfLiteTensor* tensor,
+                                   tensorflow::Tensor* tf_tensor,
+                                   bool allow_reusing = true);
+
+}  // namespace flex
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_DELEGATES_FLEX_BUFFER_MAP_UTIL_H_

tensorflow/lite/delegates/flex/build_def.bzl

@@ -0,0 +1,367 @@
+"""Generate custom flex delegate library."""
+
+load("@build_bazel_rules_android//android:rules.bzl", "android_library")
+load("@rules_cc//cc:cc_library.bzl", "cc_library")
+load(
+    "//tensorflow:tensorflow.bzl",
+    "if_android",
+    "if_ios",
+    "if_mobile",
+    "tf_cc_binary",
+    "tf_copts",
+    "tf_defines_nortti_if_lite_protos",
+    "tf_features_nomodules_if_mobile",
+    "tf_opts_nortti_if_lite_protos",
+    "tf_portable_full_lite_protos",
+)
+load(
+    "//tensorflow/lite:build_def.bzl",
+    "clean_dep",
+    "tflite_cc_shared_object",
+    "tflite_copts",
+    "tflite_jni_binary",
+    "tflite_jni_linkopts",
+)
+load(
+    "//tensorflow/lite:special_rules.bzl",
+    "flex_portable_tensorflow_deps",
+    "flex_portable_tensorflow_hdrs",
+)
+
+def generate_flex_kernel_header(
+        name,
+        models,
+        testonly = 0,
+        additional_deps = []):
+    """A rule to generate a header file listing only used operators.
+
+    Args:
+      name: Name of the generated library.
+      models: TFLite models to interpret.
+      testonly: Should be marked as true if additional_deps is testonly.
+      additional_deps: Dependencies for additional TF ops.
+
+    Returns:
+      A struct with 'header' and 'include_path' fields that
+      contain the generated header and the required include entry.
+    """
+    include_path = "%s_tf_generated_kernel_header" % name
+    header = include_path + "/ops_to_register.h"
+
+    if type(models) != type([]):
+        models = [models]
+
+    # List all flex ops from models.
+    model_file_args = " --graphs=%s" % ",".join(
+        ["$(location %s)" % f for f in models],
+    )
+    list_ops_output = include_path + "/list_flex_ops"
+    list_ops_tool = clean_dep("//tensorflow/lite/tools:list_flex_ops_main")
+    if additional_deps:
+        tf_cc_binary(
+            name = "%s_list_flex_ops_main" % name,
+            deps = [
+                clean_dep("//tensorflow/lite/tools:list_flex_ops_main_lib"),
+            ] + additional_deps,
+            testonly = testonly,
+        )
+        list_ops_tool = ":%s_list_flex_ops_main" % name
+    native.genrule(
+        name = "%s_list_flex_ops" % name,
+        srcs = models,
+        outs = [list_ops_output],
+        tools = [list_ops_tool],
+        message = "Listing flex ops from %s..." % ",".join(models),
+        cmd = ("$(location " + list_ops_tool + ")" +
+               model_file_args + " > \"$@\""),
+        testonly = testonly,
+    )
+
+    # Generate the kernel registration header file from list of flex ops.
+    tool = clean_dep("//tensorflow/python/tools:print_selective_registration_header")
+    native.genrule(
+        name = "%s_kernel_registration" % name,
+        srcs = [list_ops_output],
+        outs = [header],
+        tools = [tool],
+        message = "Processing %s..." % list_ops_output,
+        cmd = ("$(location " + tool + ")" +
+               " --default_ops=\"\"" +
+               " --proto_fileformat=ops_list" +
+               " --graphs=" + "$(location " + list_ops_output + ") > \"$@\""),
+    )
+    return struct(include_path = include_path, header = header)
+
+def tflite_flex_cc_library(
+        name,
+        models = [],
+        additional_deps = [],
+        testonly = 0,
+        visibility = ["//visibility:public"],
+        link_symbol = True,
+        compatible_with = None):
+    """A rule to generate a flex delegate with only ops to run listed models.
+
+    Args:
+      name: Name of the generated flex delegate.
+      models: TFLite models to interpret. The library will only include ops and kernels
+          to support these models. If empty, the library will include all Tensorflow
+          ops and kernels.
+      additional_deps: Dependencies for additional TF ops.
+      testonly: Mark this library as testonly if true.
+      visibility: visibility of the generated rules.
+      link_symbol: If true, add delegate_symbol to deps.
+      compatible_with: The standard compatible_with attribute.
+    """
+    portable_tensorflow_lib = clean_dep("//tensorflow/core:portable_tensorflow_lib")
+    if models:
+        CUSTOM_KERNEL_HEADER = generate_flex_kernel_header(
+            name = "%s_tf_op_headers" % name,
+            models = models,
+            additional_deps = additional_deps,
+            testonly = testonly,
+        )
+
+        # Define a custom tensorflow_lib with selective registration.
+        # The library will only contain ops exist in provided models.
+        cc_library(
+            name = "%s_tensorflow_lib" % name,
+            srcs = if_mobile([
+                clean_dep("//tensorflow/core/ops:mobile_srcs"),
+                clean_dep("//tensorflow/core/kernels:portable_core_ops_srcs"),
+                clean_dep("//tensorflow/core/kernels:portable_extended_ops"),
+            ]) + [CUSTOM_KERNEL_HEADER.header],
+            hdrs = flex_portable_tensorflow_hdrs() + [
+                clean_dep("//tensorflow/core/ops:mobile_hdrs"),
+                clean_dep("//tensorflow/core/kernels:portable_core_ops_hdrs"),
+                clean_dep("//tensorflow/core/kernels:portable_extended_ops_headers"),
+                clean_dep("//tensorflow/core/kernels:portable_quantized_ops_hdrs"),
+            ],
+            copts = tf_copts(android_optimization_level_override = None) + tf_opts_nortti_if_lite_protos() + if_ios(["-Os"]),
+            compatible_with = compatible_with,
+            defines = [
+                "SELECTIVE_REGISTRATION",
+                "SUPPORT_SELECTIVE_REGISTRATION",
+                "EIGEN_NEON_GEBP_NR=4",
+            ] + tf_portable_full_lite_protos(
+                full = [],
+                lite = ["TENSORFLOW_LITE_PROTOS"],
+            ) + tf_defines_nortti_if_lite_protos(),
+            features = tf_features_nomodules_if_mobile() + if_mobile(["-parse_headers"]),
+            linkopts = if_android(["-lz"]) + if_ios(["-lz"]),
+            includes = [
+                CUSTOM_KERNEL_HEADER.include_path,
+            ],
+            textual_hdrs = [
+                clean_dep("//tensorflow/core/kernels:portable_all_ops_textual_hdrs"),
+            ],
+            visibility = visibility,
+            deps = flex_portable_tensorflow_deps() + [
+                clean_dep("@ducc//:fft_wrapper"),
+                clean_dep("//tensorflow/core:protos_all_cc"),
+                clean_dep("//tensorflow/core:portable_tensorflow_lib_lite"),
+                clean_dep("//tensorflow/core/platform:strong_hash"),
+                clean_dep("//tensorflow/lite/delegates/flex:portable_images_lib"),
+            ],
+            alwayslink = 1,
+            testonly = testonly,
+        )
+        portable_tensorflow_lib = ":%s_tensorflow_lib" % name
+
+    delegate_symbol = []
+    if link_symbol:
+        delegate_symbol.append(clean_dep("//tensorflow/lite/delegates/flex:delegate_symbol"))
+
+    # Define a custom flex delegate with above tensorflow_lib.
+    cc_library(
+        name = name,
+        hdrs = [
+            clean_dep("//tensorflow/lite/delegates/flex:delegate.h"),
+        ],
+        compatible_with = compatible_with,
+        visibility = visibility,
+        deps = [
+            clean_dep("//tensorflow/lite/delegates/flex:delegate_data"),
+            clean_dep("//tensorflow/lite/delegates/flex:delegate_only_runtime"),
+            clean_dep("//tensorflow/lite/delegates/utils:simple_delegate"),
+        ] + select({
+            clean_dep("//tensorflow:mobile"): [
+                portable_tensorflow_lib,
+            ],
+            "//conditions:default": [
+                clean_dep("//tensorflow/core:tensorflow"),
+            ],
+        }) + [
+            clean_dep("//tensorflow/lite/core/c:private_common"),
+        ] + additional_deps + delegate_symbol,
+        testonly = testonly,
+        alwayslink = 1,
+    )
+
+def tflite_flex_shared_library(
+        name,
+        models = [],
+        additional_deps = [],
+        testonly = 0,
+        visibility = ["//visibility:private"]):
+    """A rule to generate a flex delegate shared library with only ops to run listed models.
+
+    The output library name is platform dependent:
+    - Linux/Android: `lib{name}.so`
+    - Mac: `lib{name}.dylib`
+    - Windows: `lib{name}.dll`
+
+    Args:
+      name: Name of the library.
+      models: TFLite models to interpret. The library will only include ops and kernels
+          to support these models. If empty, the library will include all Tensorflow
+          ops and kernels.
+      additional_deps: Dependencies for additional TF ops.
+      testonly: Mark this library as testonly if true.
+      visibility: visibility of the generated rules.
+    """
+    tflite_flex_cc_library(
+        name = "%s_flex_delegate" % name,
+        models = models,
+        additional_deps = additional_deps,
+        testonly = testonly,
+        visibility = visibility,
+    )
+
+    tflite_cc_shared_object(
+        name = name,
+        linkopts = select({
+            "//tensorflow:macos": [
+                "-Wl,-exported_symbols_list,$(location //tensorflow/lite/delegates/flex:exported_symbols.lds)",
+            ],
+            "//tensorflow:windows": [],
+            "//conditions:default": [
+                "-Wl,-z,defs",
+                "-Wl,--version-script,$(location //tensorflow/lite/delegates/flex:version_script.lds)",
+            ],
+        }),
+        per_os_targets = True,
+        deps = [
+            "%s_flex_delegate" % name,
+            "//tensorflow/lite/delegates/flex:exported_symbols.lds",
+            "//tensorflow/lite/delegates/flex:version_script.lds",
+        ],
+    )
+
+def tflite_flex_jni_library(
+        name,
+        models = [],
+        additional_deps = [],
+        testonly = 0,
+        visibility = ["//visibility:private"]):
+    """A rule to generate a jni library listing only used operators.
+
+    The libtensorflowlite_flex_jni.so name is fixed due to a limitation in JNI
+    Java wrapper, so please make sure there is no naming conflicts.
+
+    Args:
+      name: Prefix of the generated libraries.
+      models: TFLite models to interpret. The library will only include ops and kernels
+          to support these models. If empty, the library will include all Tensorflow
+          ops and kernels.
+      additional_deps: Dependencies for additional TF ops.
+      testonly: Mark this library as testonly if true.
+      visibility: visibility of the generated rules.
+    """
+
+    # Define a custom flex_delegate that depends on above tensorflow_lib.
+    # This will reduce the binary size comparing to the original flex delegate.
+    tflite_flex_cc_library(
+        name = "%s_flex_delegate" % name,
+        models = models,
+        additional_deps = additional_deps,
+        testonly = testonly,
+        visibility = visibility,
+    )
+
+    # Define a custom flex_native that depends on above flex_delegate.
+    cc_library(
+        name = "%s_flex_native" % name,
+        srcs = [
+            clean_dep("//tensorflow/lite/testing:init_tensorflow.cc"),
+            clean_dep("//tensorflow/lite/delegates/flex/java/src/main/native:flex_delegate_jni.cc"),
+        ],
+        hdrs = [
+            clean_dep("//tensorflow/lite/testing:init_tensorflow.h"),
+        ],
+        copts = tflite_copts(),
+        testonly = testonly,
+        visibility = visibility,
+        deps = [
+            ":%s_flex_delegate" % name,
+            clean_dep("//tensorflow/lite/java/jni"),
+            clean_dep("//tensorflow/lite/delegates/utils:simple_delegate"),
+        ] + select({
+            clean_dep("//tensorflow:mobile"): [
+                clean_dep("//tensorflow/core:portable_tensorflow_lib_lite"),
+            ],
+            "//conditions:default": [
+                clean_dep("//tensorflow/core:lib"),
+            ],
+        }),
+        alwayslink = 1,
+    )
+
+    # Build the jni binary based on the above flex_native.
+    # The library name is fixed as libtensorflowlite_flex_jni.so in FlexDelegate.java.
+    tflite_jni_binary(
+        name = "libtensorflowlite_flex_jni.so",
+        linkopts = tflite_jni_linkopts(),
+        testonly = testonly,
+        deps = [
+            ":%s_flex_native" % name,
+        ],
+    )
+
+def tflite_flex_android_library(
+        name,
+        models = [],
+        additional_deps = [],
+        custom_package = "org.tensorflow.lite.flex",
+        testonly = 0,
+        visibility = ["//visibility:private"]):
+    """A rule to generate an android library based on the selective-built jni library.
+
+    Args:
+      name: name of android library.
+      models: TFLite models used for selective build. The library will only include ops
+          and kernels to support these models. If empty, the library will include all
+          Tensorflow ops and kernels.
+      additional_deps: Dependencies for additional TF ops.
+      custom_package: Java package for which java sources will be generated.
+      testonly: Mark this library as testonly if true.
+      visibility: visibility of the generated rules.
+    """
+    tflite_flex_jni_library(
+        name = name,
+        models = models,
+        additional_deps = additional_deps,
+        testonly = testonly,
+        visibility = visibility,
+    )
+    cc_library(
+        name = "%s_native" % name,
+        srcs = ["libtensorflowlite_flex_jni.so"],
+        testonly = testonly,
+        visibility = visibility,
+    )
+
+    android_library(
+        name = name,
+        srcs = [clean_dep("//tensorflow/lite/delegates/flex/java/src/main/java/org/tensorflow/lite/flex:flex_delegate")],
+        manifest = clean_dep("//tensorflow/lite/java:AndroidManifest.xml"),
+        proguard_specs = [clean_dep("//tensorflow/lite/java:proguard.flags")],
+        custom_package = custom_package,
+        testonly = testonly,
+        deps = [
+            ":%s_native" % name,
+            clean_dep("//tensorflow/lite/java:tensorflowlite_java"),
+            clean_dep("@org_checkerframework_qual"),
+        ],
+        visibility = visibility,
+    )

tensorflow/lite/delegates/flex/delegate.cc

@@ -0,0 +1,188 @@
+/* Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include "tensorflow/lite/delegates/flex/delegate.h"
+
+#include <cstddef>
+#include <cstdlib>
+#include <cstring>
+#include <memory>
+#include <utility>
+
+#include "absl/status/status.h"
+#include "absl/strings/string_view.h"
+#include "tensorflow/core/framework/cancellation.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/platform/tstring.h"
+#include "tensorflow/core/public/session_options.h"
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/core/subgraph.h"
+#include "tensorflow/lite/delegates/flex/buffer_map.h"
+#include "tensorflow/lite/delegates/flex/kernel.h"
+#include "tensorflow/lite/delegates/flex/util.h"
+#include "tensorflow/lite/delegates/utils/simple_delegate.h"
+#include "tensorflow/lite/logger.h"
+#include "tensorflow/lite/minimal_logging.h"
+#include "tensorflow/lite/string_util.h"
+#include "tensorflow/lite/util.h"
+
+namespace tflite {
+
+TfLiteDelegateUniquePtr FlexDelegate::Create(
+    std::unique_ptr<FlexDelegate> base_delegate) {
+  TFLITE_LOG_PROD_ONCE(TFLITE_LOG_INFO,
+                       "Created TensorFlow Lite delegate for select TF ops.");
+  if (base_delegate == nullptr) {
+    base_delegate.reset(new FlexDelegate());
+  }
+  auto flex_delegate = TfLiteDelegateFactory::Create(std::move(base_delegate));
+  flex_delegate->flags |= kTfLiteDelegateFlagsAllowDynamicTensors;
+  // NOMUTANTS -- this flag has effects in profiler that disable the profiling
+  // of the macro operator "TfLiteFlexDelegate", which only shows in profiler
+  // output string. Adding flag check in Flex tests is currently not necessary.
+  flex_delegate->flags |= kTfLiteDelegateFlagsPerOperatorProfiling;
+  reinterpret_cast<FlexDelegate*>(flex_delegate->data_)->base_delegate_ =
+      flex_delegate.get();
+  return flex_delegate;
+}
+
+TfLiteStatus FlexDelegate::Initialize(TfLiteContext* context) {
+  // If the TensorFlow Lite thread count is explicitly configured, use it,
+  // otherwise rely on the default TensorFlow threading behavior.
+  tensorflow::SessionOptions session_options;
+  // We don't run multiple ops at the same time, so prefer using
+  // 1 thread for inter-op parallelism.
+  // Negative value means all are done on the caller thread.
+  session_options.config.set_inter_op_parallelism_threads(-1);
+  if (context->recommended_num_threads > 0) {
+    session_options.config.set_intra_op_parallelism_threads(
+        context->recommended_num_threads);
+  }
+
+  auto status = delegate_data_.Prepare(
+      session_options, reinterpret_cast<Subgraph*>(context->impl_),
+      base_delegate_);
+  if (!status.ok()) {
+    TF_LITE_KERNEL_LOG(context, "Failed to initialize TensorFlow context: %s",
+                       absl::StatusMessageAsCStr(status));
+    return kTfLiteError;
+  }
+
+  // Initializes the cancellation manager.
+  if (!cancellation_manager_) {
+    cancellation_manager_ = std::make_unique<tensorflow::CancellationManager>();
+    delegate_data_.SetCancellationManager(cancellation_manager_.get());
+  }
+
+  return kTfLiteOk;
+}
+
+const char* FlexDelegate::Name() const {
+  static constexpr char kName[] = "TfLiteFlexDelegate";
+  return kName;
+}
+
+bool FlexDelegate::IsNodeSupportedByDelegate(
+    const TfLiteRegistration* registration, const TfLiteNode* node,
+    TfLiteContext* context) const {
+  return IsFlexOp(registration->custom_name);
+}
+
+std::unique_ptr<SimpleDelegateKernelInterface>
+FlexDelegate::CreateDelegateKernelInterface() {
+  return std::unique_ptr<SimpleDelegateKernelInterface>(
+      new tflite::flex::DelegateKernel());
+}
+
+TfLiteStatus FlexDelegate::CopyFromBufferHandle(
+    TfLiteContext* context, TfLiteBufferHandle buffer_handle,
+    TfLiteTensor* output) {
+  flex::BufferMap* buffer_map = delegate_data_.GetBufferMap(context);
+
+  if (!buffer_map->HasTensor(buffer_handle)) {
+    TF_LITE_KERNEL_LOG(context, "Invalid tensor index %d.", buffer_handle);
+    return kTfLiteError;
+  }
+
+  tensorflow::Tensor t = buffer_map->GetTensor(buffer_handle);
+
+  if (output->type == kTfLiteString) {
+    if (t.dtype() != tensorflow::DT_STRING) {
+      TF_LITE_KERNEL_LOG(context,
+                         "Inconsistent type for TF string tensor index %d.",
+                         buffer_handle);
+      return kTfLiteError;
+    }
+    DynamicBuffer dynamic_buffer;
+
+    auto tf_data = t.flat<tensorflow::tstring>();
+    for (int i = 0; i < t.NumElements(); ++i) {
+      dynamic_buffer.AddString(tf_data(i).data(), tf_data(i).size());
+    }
+
+    dynamic_buffer.WriteToTensor(output, /*new_shape=*/nullptr);
+    return kTfLiteOk;
+  }
+
+  // TODO(b/179094265): This is an experimental implementation, subject to
+  // change. This can be re-implemented with life cycle management mechanism
+  // like reference counting.
+  // When copying resource and variant tensors from Flex delegate to TensorFlow
+  // Lite tensors, the CopyFromBufferHandle method of the Flex delegate is
+  // invoked and it will store the `data` field of the given TensorFlow Lite
+  // tensor and pass the TensorFlow Lite tensor pointer. Copying the `data`
+  // field will act as passing pointers between TensorFlow Lite tensors.
+  //
+  // The life cycle of the pointer will be managed by the reference counting in
+  // the TensorFlow world and the pointer will be freed when all the buffer
+  // maps, who own it, are gone.
+  if (IsResourceOrVariant(output)) {
+    const size_t required_bytes = tflite::flex::kTensorflowResourceTensorBytes;
+    const tensorflow::Tensor** tf_tensor_ptr =
+        reinterpret_cast<const tensorflow::Tensor**>(malloc(required_bytes));
+    *tf_tensor_ptr = buffer_map->GetTensorPtr(buffer_handle);
+
+    TfLiteTensorDataFree(output);
+    output->data.raw = reinterpret_cast<char*>(tf_tensor_ptr);
+    output->bytes = required_bytes;
+    output->data_is_stale = true;
+    return kTfLiteOk;
+  }
+
+  absl::string_view t_data = t.tensor_data();
+
+  if (output->bytes != t_data.size()) {
+    TF_LITE_KERNEL_LOG(context,
+                       "The given %zu bytes are not enough to store "
+                       "TensorFlow's aligned buffer of size %zu bytes.",
+                       output->bytes, t_data.size());
+    return kTfLiteError;
+  }
+
+  memcpy(output->data.raw, t_data.data(), t_data.size());
+  return kTfLiteOk;
+}
+
+void FlexDelegate::Cancel() { cancellation_manager_->StartCancel(); }
+
+bool FlexDelegate::HasCancelled(void* data) {
+  if (data == nullptr) {
+    return false;
+  }
+
+  auto* flex_delegate = static_cast<FlexDelegate*>(data);
+  return flex_delegate->cancellation_manager_->IsCancelled();
+}
+
+}  // namespace tflite

tensorflow/lite/delegates/flex/delegate.h

@@ -0,0 +1,125 @@
+/* Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_DELEGATES_FLEX_DELEGATE_H_
+#define TENSORFLOW_LITE_DELEGATES_FLEX_DELEGATE_H_
+
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/delegates/flex/delegate_data.h"
+#include "tensorflow/lite/delegates/utils/simple_delegate.h"
+
+namespace tflite {
+
+namespace flex {
+namespace testing {
+class KernelTest;
+}  // namespace testing
+}  // namespace flex
+
+// WARNING: This is an experimental interface that is subject to change.
+// Delegate that can be used to extract parts of a graph that are designed to be
+// executed by TensorFlow's runtime via Eager.
+//
+// The interpreter must be constructed after the FlexDelegate and destructed
+// before the FlexDelegate. This delegate may be used with multiple
+// interpreters, but it is *not* thread-safe.
+//
+// Usage:
+//   auto delegate = FlexDelegate::Create();
+//   ... build interpreter ...
+//
+//   if (delegate) {
+//     interpreter->ModifyGraphWithDelegate(delegate.get());
+//   }
+//
+//   void* delegate_data = delegate->data_;
+//   interpreter->SetCancellationFunction(
+//     delegate_data,
+//     FlexDelegate::HasCancelled);
+//
+//   ... run inference ...
+//
+//    static_cast<FlexDelegate*>(delegate_data)->Cancel();
+//
+//   ... destroy interpreter ...
+//   ... destroy delegate ...
+class FlexDelegate : public SimpleDelegateInterface {
+ public:
+  friend class flex::testing::KernelTest;
+
+  // Creates a delegate that supports TF ops.
+  static TfLiteDelegateUniquePtr Create() {
+    return Create(/*base_delegate*/ nullptr);
+  }
+
+  ~FlexDelegate() override {}
+
+  flex::DelegateData* mutable_data() { return &delegate_data_; }
+
+  // This method is thread safe. It does two things:
+  //   1. Calls the CancellationManager of the TF eager runtime to support
+  //      intra-op cancellation in TF.
+  //   2. Uses the CancellationManager to signal TFLite interpreter for inter-op
+  //      cancellation.
+  // Training is non-recoverable after calling this API.
+  void Cancel();
+
+  // The param `data` must be a pointer to a FlexDelegate instance.
+  static bool HasCancelled(void* data);
+
+ protected:
+  // We sometimes have to create certain stub data to test FlexDelegate. To
+  // achieve this, we will make a testing flex delegate class that inherits from
+  // FlexDelegate to override certain things for stub data creation. Therefore,
+  // this function accepts a FlexDelegate instance to initialize it properly for
+  // create a testing flex delegate in some cases, and it is only used in
+  // testing.
+  static TfLiteDelegateUniquePtr Create(
+      std::unique_ptr<FlexDelegate> base_delegate);
+
+  FlexDelegate() {}
+
+  const char* Name() const override;
+
+  bool IsNodeSupportedByDelegate(const TfLiteRegistration* registration,
+                                 const TfLiteNode* node,
+                                 TfLiteContext* context) const override;
+
+  TfLiteStatus Initialize(TfLiteContext* context) override;
+
+  SimpleDelegateInterface::Options DelegateOptions() const override {
+    // Use default options.
+    return SimpleDelegateInterface::Options();
+  }
+
+  std::unique_ptr<SimpleDelegateKernelInterface> CreateDelegateKernelInterface()
+      override;
+
+  TfLiteStatus CopyFromBufferHandle(TfLiteContext* context,
+                                    TfLiteBufferHandle buffer_handle,
+                                    TfLiteTensor* output) override;
+
+  flex::DelegateData delegate_data_;
+
+  // Pointer to the base TfLiteDelegate which is created from the Create call.
+  TfLiteDelegate* base_delegate_ = nullptr;
+
+ private:
+  // A cancellation manager.
+  std::unique_ptr<tensorflow::CancellationManager> cancellation_manager_;
+};
+
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_DELEGATES_FLEX_DELEGATE_H_

tensorflow/lite/delegates/flex/delegate_data.cc

@@ -0,0 +1,241 @@
+/* Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include "tensorflow/lite/delegates/flex/delegate_data.h"
+
+#include <functional>
+#include <memory>
+#include <set>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "absl/memory/memory.h"
+#include "absl/strings/str_cat.h"
+#include "flatbuffers/flexbuffers.h"  // from @flatbuffers
+#include "tensorflow/core/common_runtime/device_factory.h"
+#include "tensorflow/core/common_runtime/eager/context.h"
+#include "tensorflow/core/framework/function.h"
+#include "tensorflow/core/framework/node_def.pb.h"
+#include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/framework/resource_mgr.h"
+#include "tensorflow/core/graph/graph.h"
+#include "tensorflow/core/lib/core/status.h"
+#include "tensorflow/core/platform/errors.h"
+#include "tensorflow/core/platform/status.h"
+#include "tensorflow/core/platform/tstring.h"
+#include "tensorflow/core/protobuf/error_codes.pb.h"
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/core/subgraph.h"
+#include "tensorflow/lite/delegates/flex/util.h"
+#include "tensorflow/lite/schema/schema_generated.h"
+#include "tensorflow/lite/util.h"
+
+namespace tflite {
+namespace flex {
+
+namespace {
+
+// Builds a `FunctionDef` proto that contains two nodes:
+// The first node is a constant node which has the value of the resource key,
+// the second node is a `TfLiteSubgraphExecute` node which will take the
+// resource key, and the subgraph's inputs as arguments. The function's return
+// value is the return value of `TfLiteSubgraphExecute`.
+void BuildFunctionDefProto(const std::string& function_name,
+                           const Subgraph& subgraph,
+                           tensorflow::FunctionDef& fdef) {
+  // Map inputs/outputs to types.
+  std::vector<std::string> inputs, outputs;
+  inputs.reserve(subgraph.inputs().size());
+  outputs.reserve(subgraph.outputs().size());
+  for (int i = 0; i < subgraph.inputs().size(); ++i) {
+    inputs.push_back(absl::StrCat(
+        "args_", i, ": ",
+        TfLiteTypeToTfTypeName(subgraph.tensor(subgraph.inputs()[i])->type)));
+  }
+  for (int i = 0; i < subgraph.outputs().size(); ++i) {
+    outputs.push_back(absl::StrCat(
+        "res_", i, ": ",
+        TfLiteTypeToTfTypeName(subgraph.tensor(subgraph.outputs()[i])->type)));
+  }
+  std::vector<tensorflow::FunctionDefHelper::Node> nodes;
+  // The first node is a constant node containing the string value for the
+  // resource name.
+  nodes.push_back(tensorflow::FunctionDefHelper::Const<tensorflow::tstring>(
+      "SubgraphResourceKey", function_name));
+  // Builds the `TfLiteSubgraphExecute` node.
+  tensorflow::FunctionDefHelper::Node execute_node;
+  execute_node.ret.push_back("InvokeTfLite");
+  execute_node.op = "TfLiteSubgraphExecute";
+  execute_node.arg.push_back("SubgraphResourceKey:output:0");
+  for (int i = 0; i < subgraph.inputs().size(); ++i) {
+    execute_node.arg.push_back(absl::StrCat("args_", i));
+  }
+  nodes.push_back(execute_node);
+
+  std::vector<std::pair<std::string, std::string>> ret_def;
+  ret_def.reserve(subgraph.outputs().size());
+  for (int i = 0; i < subgraph.outputs().size(); ++i) {
+    ret_def.emplace_back(absl::StrCat("res_", i),
+                         absl::StrCat("InvokeTfLite:output:", i));
+  }
+  fdef = tensorflow::FunctionDefHelper::Create(function_name, inputs, outputs,
+                                               /*attr_def=*/{}, nodes, ret_def);
+  // Insert input/output type attrs.
+  tensorflow::AttrValue tin_attrs, tout_attrs;
+  for (int i = 0; i < subgraph.inputs().size(); ++i) {
+    TF_DataType dtype = tflite::flex::GetTensorFlowDataType(
+        subgraph.tensor(subgraph.inputs()[i])->type);
+    tin_attrs.mutable_list()->add_type(tensorflow::DataType(dtype));
+  }
+  for (int i = 0; i < subgraph.outputs().size(); ++i) {
+    TF_DataType dtype = tflite::flex::GetTensorFlowDataType(
+        subgraph.tensor(subgraph.outputs()[i])->type);
+    tout_attrs.mutable_list()->add_type(tensorflow::DataType(dtype));
+  }
+  fdef.mutable_node_def(1)->mutable_attr()->insert({"Tin", tin_attrs});
+  fdef.mutable_node_def(1)->mutable_attr()->insert({"Tout", tout_attrs});
+}
+
+// Returns a list of subgraph names which have associated function attributes.
+absl::Status GetSubgraphNamesForFunctionExecution(
+    const std::vector<std::unique_ptr<Subgraph>>& subgraphs,
+    std::set<std::string>* result) {
+  tensorflow::NodeDef node_def;
+  for (const auto& subgraph : subgraphs) {
+    for (const auto& node_and_reg : subgraph->nodes_and_registration()) {
+      if (node_and_reg.second.builtin_code != tflite::BuiltinOperator_CUSTOM) {
+        // If this isn't a custom op, skip.
+        continue;
+      }
+      const std::string custom_name = node_and_reg.second.custom_name;
+      if (custom_name.substr(0, strlen(tflite::kFlexCustomCodePrefix)) !=
+          tflite::kFlexCustomCodePrefix) {
+        // Skip if this is not a flex op.
+        continue;
+      }
+      // The flexbuffer contains a vector where the first elements is the
+      // op name and the second is a serialized NodeDef.
+      const flexbuffers::Vector& v =
+          flexbuffers::GetRoot(reinterpret_cast<const uint8_t*>(
+                                   node_and_reg.first.custom_initial_data),
+                               node_and_reg.first.custom_initial_data_size)
+              .AsVector();
+      // TODO(b/181352924): Use proto arena if we see performance regression.
+      if (!node_def.ParseFromString(v[1].AsString().str())) {
+        return absl::Status(absl::StatusCode::kInternal,
+                            "could not parse NodeDef");
+      }
+      // Loop through all the attributes in this node to check if it has
+      // function attribute.
+      for (const auto& attr : node_def.attr()) {
+        if (attr.second.has_func()) {
+          result->insert(attr.second.func().name());
+        }
+      }
+    }
+  }
+  return absl::OkStatus();
+}
+
+}  // namespace
+
+absl::Status RegisterFunctionDefForSubgraphs(
+    Subgraph& main_subgraph,
+    const std::function<absl::Status(
+        const std::vector<std::unique_ptr<Subgraph>>&, std::set<std::string>*)>&
+        select_subgraphs_to_register,
+    tensorflow::ResourceMgr* resource_mgr,
+    tensorflow::EagerContext* eager_context, TfLiteDelegate* flex_delegate) {
+  std::vector<std::unique_ptr<Subgraph>>* subgraphs =
+      main_subgraph.GetSubgraphs();
+  if (!subgraphs) {
+    // If there are no subgraphs associated with the main subgraph, we will
+    // return ok status because no FunctionDef needs to be registered.
+    return absl::OkStatus();
+  }
+  std::set<std::string> function_subgraphs;
+  TF_RETURN_IF_ERROR(
+      select_subgraphs_to_register(*subgraphs, &function_subgraphs));
+  for (int i = 0; i < subgraphs->size(); ++i) {
+    if (subgraphs->at(i)->GetName() == "main") {
+      continue;
+    }
+    const std::string subgraph_name = subgraphs->at(i)->GetName();
+    if (!function_subgraphs.count(subgraph_name)) {
+      continue;
+    }
+    // This is to ensure that we only register FunctionDefs for subgraphs that
+    // are used by TF ops to invoke functions.
+    auto* subgraph_resource =
+        new TFLiteSubgraphResource(*(subgraphs->at(i)), flex_delegate);
+    TF_RETURN_IF_ERROR(resource_mgr->Create<TFLiteSubgraphResource>(
+        "flex", subgraph_name, subgraph_resource));
+    tensorflow::FunctionDef fdef;
+    BuildFunctionDefProto(subgraph_name, *(subgraphs->at(i)), fdef);
+    TF_RETURN_IF_ERROR(eager_context->AddFunctionDef(fdef));
+  }
+  return absl::OkStatus();
+}
+
+DelegateData::DelegateData() {}
+
+DelegateData::~DelegateData() {
+  if (eager_context_) {
+    // Notify the eager context to clean up the resource being held before
+    // destructing the `DelegateData`.
+    eager_context_->HostCPU()->ClearResourceMgr();
+    eager_context_->Unref();
+  }
+}
+
+absl::Status DelegateData::Prepare(
+    const tensorflow::SessionOptions& session_options, Subgraph* main_subgraph,
+    TfLiteDelegate* flex_delegate) {
+  if (eager_context_) {
+    return absl::Status();
+  }
+  if (flex_delegate == nullptr && main_subgraph != nullptr) {
+    return absl::Status(
+        absl::StatusCode::kFailedPrecondition,
+        "flex_delegate must be non-null when main_subgraph is provided.");
+  }
+
+  std::vector<std::unique_ptr<tensorflow::Device>> devices;
+
+  TF_RETURN_IF_ERROR(tensorflow::DeviceFactory::AddDevices(
+      session_options, "/job:localhost/replica:0/task:0", &devices));
+
+  auto device_mgr =
+      std::make_unique<tensorflow::StaticDeviceMgr>(std::move(devices));
+  // Note that Rendezvous is ref-counted so it will be automatically deleted.
+  auto rendezvous = tsl::core::RefCountPtr<tensorflow::IntraProcessRendezvous>(
+      new tensorflow::IntraProcessRendezvous(device_mgr.get()));
+  eager_context_ = new tensorflow::EagerContext(
+      session_options,
+      tensorflow::ContextDevicePlacementPolicy::DEVICE_PLACEMENT_SILENT,
+      /*async=*/false, device_mgr.release(), /*device_mgr_owned*/ true,
+      std::move(rendezvous), nullptr);
+
+  if (main_subgraph) {
+    TF_RETURN_IF_ERROR(RegisterFunctionDefForSubgraphs(
+        *main_subgraph, GetSubgraphNamesForFunctionExecution,
+        eager_context_->HostCPU()->resource_manager(), eager_context_,
+        flex_delegate));
+  }
+  return absl::Status();
+}
+
+}  // namespace flex
+}  // namespace tflite

tensorflow/lite/delegates/flex/delegate_data.h

@@ -0,0 +1,107 @@
+/* Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_DELEGATES_FLEX_DELEGATE_DATA_H_
+#define TENSORFLOW_LITE_DELEGATES_FLEX_DELEGATE_DATA_H_
+
+#include <functional>
+#include <string>
+
+#include "tensorflow/core/common_runtime/eager/context.h"
+#include "tensorflow/core/framework/cancellation.h"
+#include "tensorflow/core/public/session_options.h"
+#include "tensorflow/lite/core/subgraph.h"
+#include "tensorflow/lite/delegates/flex/buffer_map.h"
+#include "tensorflow/lite/delegates/flex/subgraph_resource.h"
+
+namespace tflite {
+namespace flex {
+
+// Data kept by the Flex delegate for the lifetime of an Interpreter.
+//
+// Note: This class is *not* thread-safe; any dependent delegates should not be
+// used concurrently.
+class DelegateData {
+ public:
+  DelegateData();
+  ~DelegateData();
+
+  // Prepare the necessary EagerContext and data for execution.
+  // This must be called at least once before execution. After preparation
+  // succeeds, redundant calls will be ignored (even if the session_options
+  // differ).
+  // When `main_subgraph` parameter is provided, this function will register
+  // FunctionDefs associated with each of the subgraphs attached to the
+  // `main_subgraph` which is delegated by 'flex_delegate'.
+  // 'flex_delegate' should always be non-null when 'main_subgraph' is
+  // non-null.
+  absl::Status Prepare(const tensorflow::SessionOptions& session_options,
+                       Subgraph* main_subgraph = nullptr,
+                       TfLiteDelegate* flex_delegate = nullptr);
+
+  // The EagerContext that is required for execution of Flex Ops.
+  // Note: The context is lazily created after the first call to |Prepare()|.
+  tensorflow::EagerContext* GetEagerContext() { return eager_context_; }
+
+  tensorflow::CancellationManager* GetCancellationManager() {
+    return cancellation_manager_;
+  }
+
+  void SetCancellationManager(
+      tensorflow::CancellationManager* cancellation_manager) {
+    cancellation_manager_ = cancellation_manager;
+  }
+
+  // Map from TF Lite tensor index to TensorFlow tensor for a given context.
+  BufferMap* GetBufferMap(const TfLiteContext* context) {
+    return &buffer_map_[context];
+  }
+
+  // Returns the mapping between tensor index and last node index for a given
+  // context.
+  std::map<int, int>* GetTensorReleaseMap(const TfLiteContext* context) {
+    return &tensor_release_map_[context];
+  }
+
+ private:
+  // Will be null until Prepare() is called and completes successfully.
+  tensorflow::EagerContext* eager_context_ = nullptr;
+  // Not owned by DelegateData.
+  tensorflow::CancellationManager* cancellation_manager_ = nullptr;
+  // TODO(b/112439500): Clean up stale BufferMap instances after adding the
+  // necessary cleanup hook from a TfLiteContext to a TfLiteDelegate.
+  std::unordered_map<const TfLiteContext*, BufferMap> buffer_map_;
+  // Maps between context and the tensor release map. The map will be filled
+  // during delegate initialization, and queried during eval to look up tensor
+  // lifetime information.
+  std::unordered_map<const TfLiteContext*, std::map<int, int>>
+      tensor_release_map_;
+};
+
+// Creates a `TFLiteSubgraphResource` for each subgraph (execpt
+// for main subgraph) in the model and adds it in the eager context's resource
+// manager. It also registers FunctionDefs in the function library runtime for
+// subgraphs which are used by a list of flex ops.
+absl::Status RegisterFunctionDefForSubgraphs(
+    Subgraph& main_subgraph,
+    const std::function<absl::Status(
+        const std::vector<std::unique_ptr<Subgraph>>&,
+        std::set<std::string>* result)>& select_subgraphs_to_register,
+    tensorflow::ResourceMgr* resource_mgr,
+    tensorflow::EagerContext* eager_context, TfLiteDelegate* flex_delegate);
+
+}  // namespace flex
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_DELEGATES_FLEX_DELEGATE_DATA_H_

tensorflow/lite/delegates/flex/delegate_data_test.cc

@@ -0,0 +1,243 @@
+/* Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include "tensorflow/lite/delegates/flex/delegate_data.h"
+
+#include <memory>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "absl/memory/memory.h"
+#include "absl/strings/string_view.h"
+#include "tensorflow/core/common_runtime/eager/context.h"
+#include "tensorflow/core/lib/core/status_test_util.h"
+#include "tensorflow/core/platform/protobuf.h"
+#include "tensorflow/core/platform/status.h"
+#include "tensorflow/lite/core/api/error_reporter.h"
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/core/subgraph.h"
+#include "tensorflow/lite/kernels/subgraph_test_util.h"
+#include "tensorflow/lite/testing/util.h"
+
+namespace tflite {
+namespace flex {
+namespace {
+
+using ::tensorflow::protobuf::TextFormat;
+using ::tensorflow::protobuf::util::MessageDifferencer;
+
+TEST(DelegateDataTest, Basic) {
+  DelegateData data;
+  // We only check for success because it is hard to make initialization fail.
+  // It only happens if we manage to not link the CPU device factory into the
+  // binary.
+  tensorflow::SessionOptions session_options;
+  session_options.config.set_intra_op_parallelism_threads(2);
+  EXPECT_TRUE(data.Prepare(session_options).ok());
+
+  TfLiteContext dummy_context1 = {};
+  TfLiteContext dummy_context2 = {};
+  ASSERT_NE(data.GetEagerContext(), nullptr);
+  EXPECT_NE(data.GetBufferMap(&dummy_context1), nullptr);
+  EXPECT_NE(data.GetBufferMap(&dummy_context1),
+            data.GetBufferMap(&dummy_context2));
+}
+
+TEST(DelegateDataTest, CheckFunctionDef) {
+  tensorflow::StaticDeviceMgr device_mgr(tensorflow::DeviceFactory::NewDevice(
+      "CPU", {}, "/job:localhost/replica:0/task:0/device:CPU:0"));
+  tensorflow::EagerContext* eager_context = new tensorflow::EagerContext(
+      tensorflow::SessionOptions(),
+      tensorflow::ContextDevicePlacementPolicy::DEVICE_PLACEMENT_SILENT,
+      /*async=*/false, &device_mgr, /*device_mgr_owned*/ false, nullptr,
+      nullptr);
+
+  auto select_subgraphs_to_register =
+      [](const std::vector<std::unique_ptr<Subgraph>>& subgraphs,
+         std::set<std::string>* result) {
+        result->insert("add_subgraph");
+        result->insert("mul_subgraph");
+        return absl::OkStatus();
+      };
+
+  // Builds a TF Lite primary graph with two subgraphs.
+  subgraph_test_util::SubgraphBuilder builder;
+  std::unique_ptr<ErrorReporter> error_reporter =
+      std::make_unique<TestErrorReporter>();
+  auto add_subgraph = std::make_unique<Subgraph>(
+      error_reporter.get(), /*external_contexts=*/nullptr,
+      /*subgraphs=*/nullptr, /*resources=*/nullptr, /*resource_ids=*/nullptr,
+      /*initialization_status_map=*/nullptr);
+  add_subgraph->SetName("add_subgraph");
+  auto mul_subgraph = std::make_unique<Subgraph>(
+      error_reporter.get(), /*external_contexts=*/nullptr,
+      /*subgraphs=*/nullptr, /*resources=*/nullptr, /*resource_ids=*/nullptr,
+      /*initialization_status_map=*/nullptr);
+  mul_subgraph->SetName("mul_subgraph");
+  builder.BuildAddSubgraph(add_subgraph.get());
+  builder.BuildMulSubgraph(mul_subgraph.get());
+  std::vector<std::unique_ptr<Subgraph>> subgraphs;
+  subgraphs.push_back(std::move(add_subgraph));
+  subgraphs.push_back(std::move(mul_subgraph));
+  Subgraph main_subgraph(error_reporter.get(), nullptr, &subgraphs,
+                         /*resources=*/nullptr, /*resource_ids=*/nullptr,
+                         /*initialization_status_map=*/nullptr);
+  main_subgraph.SetName("main");
+  TF_ASSERT_OK(RegisterFunctionDefForSubgraphs(
+      main_subgraph, select_subgraphs_to_register,
+      eager_context->HostCPU()->resource_manager(), eager_context,
+      /*flex_delegate=*/nullptr));
+
+  const string add_fdef_txt = R"pb(
+    signature {
+      name: "add_subgraph"
+      input_arg { name: "args_0" type: DT_INT32 }
+      input_arg { name: "args_1" type: DT_INT32 }
+      output_arg { name: "res_0" type: DT_INT32 }
+      is_stateful: true
+    }
+    node_def {
+      name: "SubgraphResourceKey"
+      op: "Const"
+      attr {
+        key: "dtype"
+        value { type: DT_STRING }
+      }
+      attr {
+        key: "value"
+        value {
+          tensor {
+            dtype: DT_STRING
+            tensor_shape {}
+            string_val: "add_subgraph"
+          }
+        }
+      }
+    }
+    node_def {
+      name: "InvokeTfLite"
+      op: "TfLiteSubgraphExecute"
+      input: "SubgraphResourceKey:output:0"
+      input: "args_0"
+      input: "args_1"
+      attr {
+        key: "Tin"
+        value { list { type: DT_INT32 type: DT_INT32 } }
+      }
+      attr {
+        key: "Tout"
+        value { list { type: DT_INT32 } }
+      }
+    }
+    ret { key: "res_0" value: "InvokeTfLite:output:0" })pb";
+
+  const string mul_fdef_txt = R"pb(
+    signature {
+      name: "mul_subgraph"
+      input_arg { name: "args_0" type: DT_INT32 }
+      input_arg { name: "args_1" type: DT_INT32 }
+      output_arg { name: "res_0" type: DT_INT32 }
+      is_stateful: true
+    }
+    node_def {
+      name: "SubgraphResourceKey"
+      op: "Const"
+      attr {
+        key: "dtype"
+        value { type: DT_STRING }
+      }
+      attr {
+        key: "value"
+        value {
+          tensor {
+            dtype: DT_STRING
+            tensor_shape {}
+            string_val: "mul_subgraph"
+          }
+        }
+      }
+    }
+    node_def {
+      name: "InvokeTfLite"
+      op: "TfLiteSubgraphExecute"
+      input: "SubgraphResourceKey:output:0"
+      input: "args_0"
+      input: "args_1"
+      attr {
+        key: "Tin"
+        value { list { type: DT_INT32 type: DT_INT32 } }
+      }
+      attr {
+        key: "Tout"
+        value { list { type: DT_INT32 } }
+      }
+    }
+    ret { key: "res_0" value: "InvokeTfLite:output:0" })pb";
+
+  tensorflow::FunctionDef add_fdef, mul_fdef;
+  ASSERT_TRUE(TextFormat::ParseFromString(add_fdef_txt, &add_fdef));
+  ASSERT_TRUE(TextFormat::ParseFromString(mul_fdef_txt, &mul_fdef));
+  EXPECT_EQ(eager_context->GetFunctionDef("main"), nullptr);
+  ASSERT_NE(eager_context->GetFunctionDef("add_subgraph"), nullptr);
+  ASSERT_NE(eager_context->GetFunctionDef("mul_subgraph"), nullptr);
+  EXPECT_TRUE(MessageDifferencer::Equals(
+      *(eager_context->GetFunctionDef("add_subgraph")), add_fdef));
+  EXPECT_TRUE(MessageDifferencer::Equals(
+      *(eager_context->GetFunctionDef("mul_subgraph")), mul_fdef));
+
+  eager_context->Unref();
+}
+
+TEST(DelegateDataTest, CheckFunctionDefWithOnlyMainGraph) {
+  tensorflow::StaticDeviceMgr device_mgr(tensorflow::DeviceFactory::NewDevice(
+      "CPU", {}, "/job:localhost/replica:0/task:0/device:CPU:0"));
+  tensorflow::EagerContext* eager_context = new tensorflow::EagerContext(
+      tensorflow::SessionOptions(),
+      tensorflow::ContextDevicePlacementPolicy::DEVICE_PLACEMENT_SILENT,
+      /*async=*/false, &device_mgr, /*device_mgr_owned*/ false, nullptr,
+      nullptr);
+
+  auto select_subgraphs_to_register =
+      [](const std::vector<std::unique_ptr<Subgraph>>& subgraphs,
+         std::set<std::string>* result) {
+        result->insert("add_subgraph");
+        result->insert("mul_subgraph");
+        return absl::OkStatus();
+      };
+
+  // Builds a TF Lite primary graph with two subgraphs.
+  subgraph_test_util::SubgraphBuilder builder;
+  std::unique_ptr<ErrorReporter> error_reporter =
+      std::make_unique<TestErrorReporter>();
+  Subgraph main_subgraph(error_reporter.get(), /*external_contexts=*/nullptr,
+                         /*subgraphs=*/nullptr, /*resources=*/nullptr,
+                         /*resource_ids=*/nullptr,
+                         /*initialization_status_map=*/nullptr);
+  main_subgraph.SetName("main");
+  TF_ASSERT_OK(RegisterFunctionDefForSubgraphs(
+      main_subgraph, select_subgraphs_to_register,
+      eager_context->HostCPU()->resource_manager(), eager_context,
+      /*flex_delegate=*/nullptr));
+
+  EXPECT_EQ(eager_context->GetFunctionDef("main"), nullptr);
+
+  eager_context->Unref();
+}
+
+}  // namespace
+}  // namespace flex
+}  // namespace tflite

tensorflow/lite/delegates/flex/delegate_symbol.cc

@@ -0,0 +1,42 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/lite/core/c/c_api_types.h"
+#include "tensorflow/lite/delegates/flex/delegate.h"
+namespace tflite {
+// Corresponding weak declaration found in lite/core/interpreter_builder.cc.
+#if TFLITE_HAS_ATTRIBUTE_WEAK
+// If weak symbol is not supported (Windows), it can use
+// TF_AcquireFlexDelegate() path instead.
+TfLiteDelegateUniquePtr AcquireFlexDelegate() {
+  return tflite::FlexDelegate::Create();
+}
+#endif
+
+}  // namespace tflite
+
+// LINT.IfChange
+// Exported C interface function which is used by AcquireFlexDelegate() at
+// interpreter_builder.cc. To export the function name globally, the function
+// name must be matched with patterns in tf_version_script.lds. In Android, we
+// don't use this feature so skip building.
+#if !defined(__ANDROID__)
+extern "C" {
+TFL_CAPI_EXPORT tflite::TfLiteDelegateUniquePtr TF_AcquireFlexDelegate() {
+  return tflite::FlexDelegate::Create();
+}
+}  // extern "C"
+#endif  // !defined(__ANDROID__)
+// LINT.ThenChange(//tensorflow/lite/core/interpreter_builder.cc)

tensorflow/lite/delegates/flex/delegate_test.cc

@@ -0,0 +1,491 @@
+/* Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include "tensorflow/lite/delegates/flex/delegate.h"
+
+#include <cstdint>
+#include <memory>
+#include <vector>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "tensorflow/lite/delegates/flex/test_util.h"
+#include "tensorflow/lite/shared_library.h"
+
+namespace tflite {
+namespace flex {
+namespace {
+
+using ::testing::ElementsAre;
+
+class DelegateTest : public testing::FlexModelTest {
+ public:
+  DelegateTest() : delegate_(FlexDelegate::Create()) {
+    flex_delegate_ = static_cast<FlexDelegate*>(delegate_->data_);
+    interpreter_ = std::make_unique<Interpreter>(&error_reporter_);
+  }
+
+  ~DelegateTest() override {
+    // The delegate needs to be destructed after the interpreter because the
+    // interpreter references data contained in the delegate.
+    interpreter_.reset();
+    delegate_.reset();
+  }
+
+  void ConfigureDelegate() {
+    interpreter_->SetCancellationFunction(flex_delegate_,
+                                          FlexDelegate::HasCancelled);
+    ASSERT_EQ(interpreter_->ModifyGraphWithDelegate(delegate_.get()),
+              kTfLiteOk);
+  }
+
+  void Cancel() { flex_delegate_->Cancel(); }
+
+ private:
+  std::unique_ptr<TfLiteDelegate, void (*)(TfLiteDelegate*)> delegate_;
+  FlexDelegate* flex_delegate_;
+};
+
+TEST_F(DelegateTest, FullGraph) {
+  // Define the graph.
+  AddTensors(9, {0, 3}, {8}, kTfLiteFloat32, {3});
+
+  AddTfOp(testing::kUnpack, {0}, {1, 2});
+  AddTfOp(testing::kUnpack, {3}, {4, 5});
+  AddTfOp(testing::kAdd, {1, 4}, {6});
+  AddTfOp(testing::kAdd, {2, 5}, {7});
+  AddTfOp(testing::kMul, {6, 7}, {8});
+
+  // Apply the delegate.
+  ConfigureDelegate();
+
+  // Define inputs.
+  SetShape(0, {2, 2, 1});
+  SetValues(0, {1.1f, 2.2f, 3.3f, 4.4f});
+  SetShape(3, {2, 2, 1});
+  SetValues(3, {1.1f, 2.2f, 3.3f, 4.4f});
+
+  ASSERT_TRUE(Invoke());
+
+  ASSERT_THAT(GetShape(8), ElementsAre(2, 1));
+  ASSERT_THAT(GetValues(8), ElementsAre(14.52f, 38.72f));
+  ASSERT_EQ(GetType(8), kTfLiteFloat32);
+}
+
+TEST_F(DelegateTest, NonFloatTypeInference) {
+  AddTensors(3, {0, 1}, {2}, kTfLiteInt32, {2});
+
+  AddTfOp(testing::kAdd, {0, 1}, {2});
+
+  ConfigureDelegate();
+
+  SetShape(0, {2, 2});
+  SetTypedValues<int>(0, {1, 2, 3, 4});
+  SetShape(1, {2, 2});
+  SetTypedValues<int>(1, {4, 3, 2, 1});
+
+  ASSERT_TRUE(Invoke());
+
+  ASSERT_THAT(GetShape(2), ElementsAre(2, 2));
+  ASSERT_THAT(GetTypedValues<int>(2), ElementsAre(5, 5, 5, 5));
+  ASSERT_EQ(GetType(2), kTfLiteInt32);
+}
+
+TEST_F(DelegateTest, StringInference) {
+  AddTensors(3, {0, 1}, {2}, kTfLiteString, {2});
+
+  AddTfOp(testing::kAdd, {0, 1}, {2});
+
+  ConfigureDelegate();
+
+  SetShape(0, {2, 2});
+  SetStringValues(0, {"1", "2", "3", "4"});
+  SetShape(1, {2, 2});
+  SetStringValues(1, {"4", "3", "2", "1"});
+
+  ASSERT_TRUE(Invoke());
+
+  ASSERT_THAT(GetShape(2), ElementsAre(2, 2));
+  ASSERT_THAT(GetStringValues(2), ElementsAre("14", "23", "32", "41"));
+  ASSERT_EQ(GetType(2), kTfLiteString);
+}
+
+TEST_F(DelegateTest, MixedGraph) {
+  AddTensors(9, {0, 3}, {8}, kTfLiteFloat32, {3});
+
+  AddTfOp(testing::kUnpack, {0}, {1, 2});
+  AddTfOp(testing::kUnpack, {3}, {4, 5});
+  AddTfOp(testing::kAdd, {1, 4}, {6});
+  AddTfOp(testing::kAdd, {2, 5}, {7});
+  AddTfLiteMulOp({6, 7}, {8});
+
+  ConfigureDelegate();
+
+  SetShape(0, {2, 2, 1});
+  SetValues(0, {1.1f, 2.2f, 3.3f, 4.4f});
+  SetShape(3, {2, 2, 1});
+  SetValues(3, {1.1f, 2.2f, 3.3f, 4.4f});
+
+  ASSERT_TRUE(Invoke());
+
+  ASSERT_THAT(GetShape(8), ElementsAre(2, 1));
+  ASSERT_THAT(GetValues(8), ElementsAre(14.52f, 38.72f));
+}
+
+TEST_F(DelegateTest, SplitGraph) {
+  AddTensors(10, {0}, {9}, kTfLiteFloat32, {3});
+
+  AddTfOp(testing::kUnpack, {0}, {1, 2});
+  AddTfOp(testing::kAdd, {1, 2}, {3});
+  AddTfOp(testing::kUnpack, {3}, {4, 5});
+
+  AddTfLiteMulOp({4, 5}, {6});
+
+  AddTfOp(testing::kUnpack, {6}, {7, 8});
+  AddTfOp(testing::kAdd, {7, 8}, {9});
+
+  ConfigureDelegate();
+
+  SetShape(0, {2, 2, 2, 1});
+  SetValues(0, {3.0f, 1.0f, 0.5f, -1.0f, 0.0f, 1.0f, 1.5f, 3.0f});
+
+  ASSERT_TRUE(Invoke());
+
+  ASSERT_THAT(GetShape(9), ElementsAre(1));
+  ASSERT_THAT(GetValues(9), ElementsAre(10.0f));
+}
+
+TEST_F(DelegateTest, OnlyTFLite) {
+  // Only TFLite single op model.
+  AddTensors(10, {0, 1}, {2}, kTfLiteFloat32, {3});
+  AddTfLiteMulOp({0, 1}, {2});
+
+  ConfigureDelegate();
+
+  SetShape(0, {2, 2, 1});
+  SetValues(0, {1.1f, 2.2f, 3.3f, 4.4f});
+  SetShape(1, {2, 2, 1});
+  SetValues(1, {1.0f, 2.0f, 3.0f, 4.0f});
+
+  ASSERT_TRUE(Invoke());
+
+  ASSERT_THAT(GetShape(2), ElementsAre(2, 2, 1));
+  ASSERT_THAT(GetValues(2), ElementsAre(1.1f, 4.4f, 9.9f, 17.6f));
+}
+
+TEST_F(DelegateTest, MultipleInvokeCalls) {
+  // Call Invoke() multiple times on the same model.
+  AddTensors(10, {0, 1}, {2}, kTfLiteFloat32, {3});
+  AddTfLiteMulOp({0, 1}, {2});
+
+  ConfigureDelegate();
+
+  SetShape(0, {2, 2, 1});
+  SetValues(0, {1.1f, 2.2f, 3.3f, 4.4f});
+  SetShape(1, {2, 2, 1});
+  SetValues(1, {1.0f, 2.0f, 3.0f, 4.0f});
+
+  ASSERT_TRUE(Invoke());
+
+  ASSERT_THAT(GetShape(2), ElementsAre(2, 2, 1));
+  ASSERT_THAT(GetValues(2), ElementsAre(1.1f, 4.4f, 9.9f, 17.6f));
+
+  SetShape(0, {2, 2, 1});
+  SetValues(1, {4.0f, 3.0f, 2.0f, 1.0f});
+  SetShape(1, {2, 2, 1});
+  SetValues(0, {4.4f, 3.3f, 2.2f, 1.1f});
+
+  ASSERT_TRUE(Invoke());
+
+  ASSERT_THAT(GetShape(2), ElementsAre(2, 2, 1));
+  ASSERT_THAT(GetValues(2), ElementsAre(17.6f, 9.9f, 4.4f, 1.1f));
+}
+
+TEST_F(DelegateTest, MultipleInterpretersSameDelegate) {
+  // Build a graph, configure the delegate and set inputs.
+  {
+    AddTensors(9, {0, 3}, {8}, kTfLiteFloat32, {3});
+    AddTfOp(testing::kUnpack, {0}, {1, 2});
+    AddTfOp(testing::kUnpack, {3}, {4, 5});
+    AddTfOp(testing::kAdd, {1, 4}, {6});
+    AddTfOp(testing::kAdd, {2, 5}, {7});
+    AddTfOp(testing::kMul, {6, 7}, {8});
+    ConfigureDelegate();
+    SetShape(0, {2, 2, 1});
+    SetValues(0, {1.1f, 2.2f, 3.3f, 4.4f});
+    SetShape(3, {2, 2, 1});
+    SetValues(3, {1.1f, 2.2f, 3.3f, 4.4f});
+  }
+
+  // Create a new interpreter, inject into the test framework and build
+  // a different graph using the *same* delegate.
+  std::unique_ptr<Interpreter> interpreter(new Interpreter(&error_reporter_));
+  interpreter_.swap(interpreter);
+  {
+    AddTensors(10, {0}, {9}, kTfLiteFloat32, {3});
+    AddTfOp(testing::kUnpack, {0}, {1, 2});
+    AddTfOp(testing::kAdd, {1, 2}, {3});
+    AddTfOp(testing::kUnpack, {3}, {4, 5});
+    AddTfLiteMulOp({4, 5}, {6});
+    AddTfOp(testing::kUnpack, {6}, {7, 8});
+    AddTfOp(testing::kAdd, {7, 8}, {9});
+    ConfigureDelegate();
+    SetShape(0, {2, 2, 2, 1});
+    SetValues(0, {3.0f, 1.0f, 0.5f, -1.0f, 0.0f, 1.0f, 1.5f, 3.0f});
+  }
+
+  // Swap back in the first interpreter and validate inference.
+  interpreter_.swap(interpreter);
+  {
+    ASSERT_TRUE(Invoke());
+    EXPECT_THAT(GetShape(8), ElementsAre(2, 1));
+    EXPECT_THAT(GetValues(8), ElementsAre(14.52f, 38.72f));
+  }
+
+  // Swap in the second interpreter and validate inference.
+  interpreter_.swap(interpreter);
+  {
+    ASSERT_TRUE(Invoke());
+    EXPECT_THAT(GetShape(9), ElementsAre(1));
+    EXPECT_THAT(GetValues(9), ElementsAre(10.0f));
+  }
+}
+
+TEST_F(DelegateTest, SingleThreaded) {
+  AddTensors(9, {0, 3}, {8}, kTfLiteFloat32, {3});
+  AddTfOp(testing::kUnpack, {0}, {1, 2});
+  AddTfOp(testing::kUnpack, {3}, {4, 5});
+  AddTfOp(testing::kAdd, {1, 4}, {6});
+  AddTfOp(testing::kAdd, {2, 5}, {7});
+  AddTfOp(testing::kMul, {6, 7}, {8});
+
+  // Explicitly disable multi-threading before installing the delegate.
+  interpreter_->SetNumThreads(1);
+  ConfigureDelegate();
+
+  SetShape(0, {2, 2, 1});
+  SetValues(0, {1.1f, 2.2f, 3.3f, 4.4f});
+  SetShape(3, {2, 2, 1});
+  SetValues(3, {1.1f, 2.2f, 3.3f, 4.4f});
+
+  // Invocation should behave as expected.
+  ASSERT_TRUE(Invoke());
+
+  ASSERT_THAT(GetShape(8), ElementsAre(2, 1));
+  ASSERT_THAT(GetValues(8), ElementsAre(14.52f, 38.72f));
+  ASSERT_EQ(GetType(8), kTfLiteFloat32);
+}
+
+TEST_F(DelegateTest, MultiThreaded) {
+  AddTensors(9, {0, 3}, {8}, kTfLiteFloat32, {3});
+  AddTfOp(testing::kUnpack, {0}, {1, 2});
+  AddTfOp(testing::kUnpack, {3}, {4, 5});
+  AddTfOp(testing::kAdd, {1, 4}, {6});
+  AddTfOp(testing::kAdd, {2, 5}, {7});
+  AddTfOp(testing::kMul, {6, 7}, {8});
+
+  // Explicitly enable multi-threading before installing the delegate.
+  interpreter_->SetNumThreads(4);
+  ConfigureDelegate();
+
+  SetShape(0, {2, 2, 1});
+  SetValues(0, {1.1f, 2.2f, 3.3f, 4.4f});
+  SetShape(3, {2, 2, 1});
+  SetValues(3, {1.1f, 2.2f, 3.3f, 4.4f});
+
+  // Invocation should behave as expected.
+  ASSERT_TRUE(Invoke());
+
+  ASSERT_THAT(GetShape(8), ElementsAre(2, 1));
+  ASSERT_THAT(GetValues(8), ElementsAre(14.52f, 38.72f));
+  ASSERT_EQ(GetType(8), kTfLiteFloat32);
+}
+
+#if !defined(__ANDROID__)
+TEST_F(DelegateTest, TF_AcquireFlexDelegate) {
+  auto TF_AcquireFlexDelegate =
+      reinterpret_cast<Interpreter::TfLiteDelegatePtr (*)()>(
+          SharedLibrary::GetSymbol("TF_AcquireFlexDelegate"));
+  ASSERT_TRUE(TF_AcquireFlexDelegate);
+  auto delegate_ptr = TF_AcquireFlexDelegate();
+  ASSERT_TRUE(delegate_ptr != nullptr);
+}
+#endif  // !defined(__ANDROID__)
+
+TEST_F(DelegateTest, StaticOutput) {
+  // Define the graph with input, output shapes of [2].
+  AddTensors(7, {0, 1, 2, 3}, {6}, kTfLiteFloat32, {2});
+
+  AddTfOp(testing::kAdd, {0, 2}, {4});
+  AddTfOp(testing::kAdd, {1, 3}, {5});
+  AddTfOp(testing::kMul, {4, 5}, {6});
+
+  // Apply the delegate.
+  ConfigureDelegate();
+
+  // Define inputs which matech with the original shapes.
+  SetShape(0, {2});
+  SetShape(1, {2});
+  SetShape(2, {2});
+  SetShape(3, {2});
+  SetValues(0, {1.1f, 2.2f});
+  SetValues(1, {3.3f, 4.4f});
+  SetValues(2, {1.1f, 2.2f});
+  SetValues(3, {3.3f, 4.4f});
+
+  ASSERT_TRUE(Invoke());
+
+  ASSERT_THAT(GetShape(6), ElementsAre(2));
+  ASSERT_THAT(GetValues(6), ElementsAre(14.52f, 38.72f));
+  ASSERT_EQ(GetType(6), kTfLiteFloat32);
+  // Since shapes are consistent, static output tensor is used.
+  ASSERT_FALSE(IsDynamicTensor(6));
+}
+
+TEST_F(DelegateTest, StaticOutputRFFT) {
+  // Define the graph with input, output shapes of [3, 257].
+  AddTensors(4, {0, 1}, {3}, kTfLiteFloat32, {3, 257});
+  int32_t rfft_length[] = {512};
+  SetConstTensor(1, {1}, kTfLiteInt32,
+                 reinterpret_cast<const char*>(&rfft_length),
+                 sizeof(rfft_length));
+
+  AddTfOp(testing::kRfft, {0, 1}, {2});
+  AddTfOp(testing::kImag, {2}, {3});
+
+  // Apply the delegate.
+  ConfigureDelegate();
+
+  // Define inputs.
+  SetShape(0, {3, 512});
+  SetValues(0, std::vector<float>(3 * 512, 1.0f));
+
+  ASSERT_TRUE(Invoke());
+
+  ASSERT_EQ(GetType(3), kTfLiteFloat32);
+  // Since shapes are consistent, static output tensor is used.
+  ASSERT_FALSE(IsDynamicTensor(3));
+}
+
+TEST_F(DelegateTest, DynamicOutputAfterReshape) {
+  // Define the graph.
+  AddTensors(9, {0, 3}, {8}, kTfLiteFloat32, {3});
+
+  AddTfOp(testing::kUnpack, {0}, {1, 2});
+  AddTfOp(testing::kUnpack, {3}, {4, 5});
+  AddTfOp(testing::kAdd, {1, 4}, {6});
+  AddTfOp(testing::kAdd, {2, 5}, {7});
+  AddTfOp(testing::kMul, {6, 7}, {8});
+
+  // Apply the delegate.
+  ConfigureDelegate();
+
+  // Define inputs with reshape.
+  SetShape(0, {2, 2, 1});
+  SetValues(0, {1.1f, 2.2f, 3.3f, 4.4f});
+  SetShape(3, {2, 2, 1});
+  SetValues(3, {1.1f, 2.2f, 3.3f, 4.4f});
+
+  ASSERT_TRUE(Invoke());
+
+  ASSERT_THAT(GetShape(8), ElementsAre(2, 1));
+  ASSERT_THAT(GetValues(8), ElementsAre(14.52f, 38.72f));
+  ASSERT_EQ(GetType(8), kTfLiteFloat32);
+  // Since shapes are inconsistent, dynamic output tensor is used.
+  ASSERT_TRUE(IsDynamicTensor(8));
+}
+
+TEST_F(DelegateTest, TestCancellation1) {
+  AddTensors(3, {0, 1}, {2}, kTfLiteInt32, {2});
+
+  AddTfOp(testing::kAdd, {0, 1}, {2});
+
+  ConfigureDelegate();
+
+  SetShape(0, {2, 2});
+  SetTypedValues<int>(0, {1, 2, 3, 4});
+  SetShape(1, {2, 2});
+  SetTypedValues<int>(1, {4, 3, 2, 1});
+
+  ASSERT_TRUE(Invoke());
+
+  ASSERT_THAT(GetShape(2), ElementsAre(2, 2));
+  ASSERT_THAT(GetTypedValues<int>(2), ElementsAre(5, 5, 5, 5));
+  ASSERT_EQ(GetType(2), kTfLiteInt32);
+
+  Cancel();
+  // Op should be cancelled.
+  ASSERT_FALSE(Invoke());
+  // TODO(b/205345340): We shouldn't do raw string matching here. Instead we
+  // need to introduce fine-grained error codes to represent cancellation
+  // status.
+  EXPECT_EQ(error_reporter_.error_messages(),
+            "Client requested cancel during Invoke()");
+}
+
+TEST_F(DelegateTest, TestCancellation2) {
+  // Define the graph.
+  AddTensors(2, {0}, {1}, kTfLiteBool, {1});
+
+  // We need an op that checks the CancellationManager status.
+  AddTfOp(testing::kLoopCond, {0}, {1});
+
+  // Apply the delegate.
+  ConfigureDelegate();
+
+  // Define inputs.
+  SetShape(0, {1});
+
+  ASSERT_TRUE(Invoke());
+
+  Cancel();
+  // Op should be cancelled.
+  ASSERT_FALSE(Invoke());
+  // TODO(b/205345340): We shouldn't do raw string matching here. Instead we
+  // need to introduce fine-grained error codes to represent cancellation
+  // status.
+  EXPECT_EQ(error_reporter_.error_messages(),
+            "Client requested cancel during Invoke()");
+}
+
+TEST_F(DelegateTest, TestCancellationTwoThreads) {
+  AddTensors(3, {0, 1}, {2}, kTfLiteInt32, {2});
+
+  AddTfOp(testing::kAdd, {0, 1}, {2});
+
+  ConfigureDelegate();
+
+  SetShape(0, {2, 2});
+  SetTypedValues<int>(0, {1, 2, 3, 4});
+  SetShape(1, {2, 2});
+  SetTypedValues<int>(1, {4, 3, 2, 1});
+
+  std::thread invoke_thread([this]() {
+    bool result = true;
+    result = this->Invoke();
+    std::this_thread::sleep_for(std::chrono::milliseconds(1000));
+    result = this->Invoke();
+    ASSERT_FALSE(result);
+    // TODO(b/205345340): Check returned error code.
+  });
+
+  std::thread cancel_thread([this]() { this->Cancel(); });
+
+  invoke_thread.join();
+  cancel_thread.join();
+}
+
+}  // namespace
+}  // namespace flex
+}  // namespace tflite

tensorflow/lite/delegates/flex/exported_symbols.lds

@@ -0,0 +1 @@
+*AcquireFlexDelegate*

tensorflow/lite/delegates/flex/java/src/main/java/org/tensorflow/lite/flex/BUILD

@@ -0,0 +1,9 @@
+# copybara:uncomment package(default_applicable_licenses = ["//tensorflow:LICENSE"])
+
+licenses(["notice"])
+
+filegroup(
+    name = "flex_delegate",
+    srcs = ["FlexDelegate.java"],
+    visibility = ["//visibility:public"],
+)

tensorflow/lite/delegates/flex/java/src/main/java/org/tensorflow/lite/flex/FlexDelegate.java

@@ -0,0 +1,69 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+package org.tensorflow.lite.flex;
+
+import java.io.Closeable;
+import org.tensorflow.lite.Delegate;
+import org.tensorflow.lite.annotations.UsedByReflection;
+
+/** {@link Delegate} for using select TensorFlow ops. */
+@UsedByReflection("Interpreter")
+public class FlexDelegate implements Delegate, Closeable {
+
+  private static final long INVALID_DELEGATE_HANDLE = 0;
+  private static final String TFLITE_FLEX_LIB = "tensorflowlite_flex_jni";
+
+  private long delegateHandle;
+
+  @UsedByReflection("Interpreter")
+  public FlexDelegate() {
+    delegateHandle = nativeCreateDelegate();
+  }
+
+  @Override
+  @UsedByReflection("Interpreter")
+  public long getNativeHandle() {
+    return delegateHandle;
+  }
+
+  /**
+   * Releases native resources held by the delegate.
+   *
+   * <p>User is expected to call this method explicitly.
+   */
+  @Override
+  @UsedByReflection("Interpreter")
+  public void close() {
+    if (delegateHandle != INVALID_DELEGATE_HANDLE) {
+      nativeDeleteDelegate(delegateHandle);
+      delegateHandle = INVALID_DELEGATE_HANDLE;
+    }
+  }
+
+  public static void initTensorFlowForTesting() {
+    nativeInitTensorFlow();
+  }
+
+  static {
+    System.loadLibrary(TFLITE_FLEX_LIB);
+  }
+
+  private static native long nativeInitTensorFlow();
+
+  private static native long nativeCreateDelegate();
+
+  private static native void nativeDeleteDelegate(long delegateHandle);
+}

tensorflow/lite/delegates/flex/java/src/main/native/BUILD

@@ -0,0 +1,35 @@
+# Description:
+# Java Native Interface (JNI) library intended for implementing the
+# TensorFlow Lite Flex delegate for using TensorFlow ops with TensorFlow Lite.
+
+load("@rules_cc//cc:cc_library.bzl", "cc_library")
+load("//tensorflow/lite:build_def.bzl", "tflite_copts")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:LICENSE"],
+    default_visibility = ["//visibility:public"],
+)
+
+licenses(["notice"])
+
+exports_files(
+    srcs = ["flex_delegate_jni.cc"],
+    visibility = ["//visibility:public"],
+)
+
+cc_library(
+    name = "native",
+    srcs = ["flex_delegate_jni.cc"],
+    copts = tflite_copts(),
+    tags = [
+        "manual",
+        "notap",
+    ],
+    deps = [
+        "//tensorflow/lite/delegates/flex:delegate",
+        "//tensorflow/lite/delegates/utils:simple_delegate",
+        "//tensorflow/lite/java/jni",
+        "//tensorflow/lite/testing:init_tensorflow",
+    ],
+    alwayslink = 1,
+)

tensorflow/lite/delegates/flex/java/src/main/native/flex_delegate_jni.cc

@@ -0,0 +1,43 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <jni.h>
+
+#include "tensorflow/lite/delegates/flex/delegate.h"
+#include "tensorflow/lite/delegates/utils/simple_delegate.h"
+#include "tensorflow/lite/testing/init_tensorflow.h"
+
+extern "C" {
+
+JNIEXPORT void JNICALL
+Java_org_tensorflow_lite_flex_FlexDelegate_nativeInitTensorFlow(JNIEnv* env,
+                                                                jclass clazz) {
+  ::tflite::InitTensorFlow();
+}
+
+JNIEXPORT jlong JNICALL
+Java_org_tensorflow_lite_flex_FlexDelegate_nativeCreateDelegate(JNIEnv* env,
+                                                                jclass clazz) {
+  return reinterpret_cast<jlong>(tflite::FlexDelegate::Create().release());
+}
+
+JNIEXPORT void JNICALL
+Java_org_tensorflow_lite_flex_FlexDelegate_nativeDeleteDelegate(
+    JNIEnv* env, jclass clazz, jlong delegate) {
+  tflite::TfLiteDelegateFactory::DeleteSimpleDelegate(
+      reinterpret_cast<struct TfLiteDelegate*>(delegate));
+}
+
+}  // extern "C"

tensorflow/lite/delegates/flex/kernel.cc

@@ -0,0 +1,857 @@
+/* Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include "tensorflow/lite/delegates/flex/kernel.h"
+
+#include <inttypes.h>
+
+#include <algorithm>
+#include <cstdint>
+#include <cstring>
+#include <map>
+#include <memory>
+#include <set>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "absl/container/flat_hash_set.h"
+#include "absl/container/inlined_vector.h"
+#include "absl/log/check.h"
+#include "absl/status/status.h"
+#include "absl/strings/string_view.h"
+#include "flatbuffers/flexbuffers.h"  // from @flatbuffers
+#include "xla/tsl/platform/errors.h"
+#include "xla/tsl/platform/statusor.h"
+#include "tensorflow/core/common_runtime/eager/context.h"
+#include "tensorflow/core/framework/cancellation.h"
+#include "tensorflow/core/framework/node_def.pb.h"
+#include "tensorflow/core/framework/node_def_util.h"
+#include "tensorflow/core/framework/op.h"
+#include "tensorflow/core/framework/op_def_builder.h"
+#include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/framework/shape_inference.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/lib/core/errors.h"
+#include "tensorflow/core/protobuf/error_codes.pb.h"
+#include "tensorflow/core/public/version.h"
+#include "tensorflow/core/tfrt/fallback/op_kernel_runner.h"
+#include "tensorflow/lite/context_util.h"
+#include "tensorflow/lite/core/api/profiler.h"
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/delegates/flex/buffer_map.h"
+#include "tensorflow/lite/delegates/flex/delegate.h"
+#include "tensorflow/lite/delegates/flex/delegate_data.h"
+#include "tensorflow/lite/delegates/flex/util.h"
+#include "tensorflow/lite/kernels/kernel_util.h"
+#include "tensorflow/lite/logger.h"
+#include "tensorflow/lite/minimal_logging.h"
+#include "tensorflow/lite/string_type.h"
+#include "tensorflow/lite/util.h"
+
+// Note: this is part of TF Lite's Flex delegation code which is to be
+// completed soon.
+
+// This is the TF Lite op that is created by the flex delegate to handle
+// execution of a supported subgraph. The usual flow is that the delegate
+// informs the interpreter of supported nodes in a graph, and each supported
+// subgraph is replaced with one instance of this kernel.
+//
+// The kernel is initialized with TfLiteDelegateParams from which we retrieve
+// the global EagerContext and BufferMap, as well as a list of inputs and
+// outputs to the subgraph. Those are used to build the OpData, with a list of
+// TensorFlow Ops that should be executed in order (which we call an OpNode).
+//
+// For each node included in the subgraph, we query the interpreter and
+// retrieve the associated NodeDef, which is then used to configure the
+// corresponding TensorFlow OpKernel.
+
+using tensorflow::shape_inference::DimensionHandle;
+using tensorflow::shape_inference::InferenceContext;
+using tensorflow::shape_inference::ShapeAndType;
+using tensorflow::shape_inference::ShapeHandle;
+
+namespace tflite {
+namespace flex {
+
+constexpr char kReadVariableOp[] = "ReadVariableOp";
+constexpr char kInterOpParallelismAttrName[] = "use_inter_op_parallelism";
+
+struct OpNode;
+
+// Represents the origin of a given tensor as a reference to the output
+// of an upstream node.
+struct TensorSource {
+  OpNode* node;
+  int node_output_index;
+};
+
+// A list of inputs of a given node of the TensorFlow graph.
+class OpInputs {
+ public:
+  explicit OpInputs(const TfLiteIntArray* indexes) {
+    for (int index : TfLiteIntArrayView(indexes)) {
+      inputs_.push_back(index);
+    }
+    forwardable_.resize(inputs_.size());
+  }
+  ~OpInputs() = default;
+
+  int Size() const { return inputs_.size(); }
+
+  int TfLiteIndex(int i) const { return inputs_[i]; }
+
+  // Given a map relating tensors to the node that originates them, populate a
+  // list of sources for the tensors in this class.
+  void InitializeTensorSources(
+      const std::map<int, TensorSource>& tflite_tensor_sources) {
+    sources_.clear();
+    for (int i : inputs_) {
+      auto it = tflite_tensor_sources.find(i);
+      if (it == tflite_tensor_sources.end()) {
+        sources_.push_back({nullptr, 0});
+      } else {
+        sources_.push_back(it->second);
+      }
+    }
+  }
+
+  void SetForwardable(int i, bool v) { forwardable_[i] = v; }
+
+  bool IsForwardable(int i) const { return forwardable_[i]; }
+
+  TensorSource GetTensorSource(int i) const { return sources_[i]; }
+
+ private:
+  std::vector<int> inputs_;
+  std::vector<TensorSource> sources_;
+
+  // List of tensors that can be used by TF in its forwarding optimization.
+  // Doing so allows an input tensor to be modified and used as the output
+  // tensor. The delegate takes care of not holding any references to tensors
+  // in this list while the corresponding tensorflow::OpKernel is executed.
+  std::vector<int> forwardable_;
+};
+
+// A list of outputs of a given node of the TensorFlow graph, along with
+// the actual outputs of the tensorflow::OpKernel.
+class OpOutputs {
+ public:
+  explicit OpOutputs(const TfLiteIntArray* indexes) {
+    for (int index : TfLiteIntArrayView(indexes)) {
+      outputs_.push_back(index);
+    }
+    vector_.resize(outputs_.size());
+  }
+  ~OpOutputs() = default;
+
+  // Stores information about which of the tensors in this class are also
+  // outputs of the sugbraph.
+  void InitializeGraphOutputs(const std::set<int>& subgraph_outputs) {
+    subgraph_outputs_.clear();
+    for (int i : outputs_) {
+      subgraph_outputs_.push_back(subgraph_outputs.count(i) > 0);
+    }
+  }
+
+  // Returns true if the tensor given by index 'i' is an output of the entire
+  // subgraph.
+  bool IsSubgraphOutput(int i) const { return subgraph_outputs_[i]; }
+
+  const tensorflow::Tensor& GetTensor(int i) const { return vector_[i]; }
+  tensorflow::Tensor ReleaseTensor(int i) { return std::move(vector_[i]); }
+
+  int Size() const { return outputs_.size(); }
+
+  int TfLiteIndex(int i) const { return outputs_[i]; }
+
+  absl::InlinedVector<tensorflow::Tensor, 2UL>* GetTensors() {
+    return &vector_;
+  }
+
+ private:
+  std::vector<int> outputs_;
+  std::vector<bool> subgraph_outputs_;
+  absl::InlinedVector<tensorflow::Tensor, 2UL> vector_;
+};
+
+// This struct holds information such as tensor lifecycle and BufferMap which
+// needs to be shared between `OpNode` and DelegateKernel.
+struct OpDataInfo {
+  // Buffer map which stores the mapping between TfLiteTensor index to TF
+  // tensor.
+  BufferMap* buffer_map;
+  // Mapping information between TfLiteTensor index to last node which uses the
+  // tensor.
+  std::map<int, int>* tensor_release_map;
+  // For output tensors that don't need to be preserved in the BufferMap, we
+  // copy them to TF Lite tensors and keep the tensor indexes in this set.
+  std::set<int> already_transferred_outputs;
+};
+
+// A single node within the larger 'op'. Note that this kernel executes many
+// TensorFlow ops within a single TF Lite op.
+class OpNode {
+ public:
+  OpNode(const TfLiteIntArray* inputs, const TfLiteIntArray* outputs)
+      : inputs_(inputs), outputs_(outputs) {}
+  ~OpNode() = default;
+
+  const string& name() const { return name_; }
+  void set_name(const string& name) { name_ = name; }
+
+  int index() const { return index_; }
+  void set_index(int index) { index_ = index; }
+
+  const tensorflow::NodeDef& nodedef() const { return nodedef_; }
+  const tensorflow::OpRegistrationData* op_reg_data() const {
+    return op_reg_data_;
+  }
+
+  const OpInputs& inputs() const { return inputs_; }
+  OpInputs* mutable_inputs() { return &inputs_; }
+
+  const OpOutputs& outputs() const { return outputs_; }
+  OpOutputs* mutable_outputs() { return &outputs_; }
+
+  int NumInputs() const { return inputs_.Size(); }
+  int NumOutputs() const { return outputs_.Size(); }
+
+  const tensorflow::tfrt_stub::OpKernelRunner& op_kernel_runner() const {
+    return op_kernel_runner_;
+  }
+
+  absl::Status InitializeNodeDef(const void* custom_initial_data,
+                                 int custom_initial_data_size) {
+    if (!custom_initial_data) {
+      return tensorflow::errors::Internal(
+          "Cannot convert empty data into a valid NodeDef");
+    }
+    // The flexbuffer contains a vector where the first elements is the
+    // op name and the second is a serialized NodeDef.
+    const flexbuffers::Vector& v =
+        flexbuffers::GetRoot(
+            reinterpret_cast<const uint8_t*>(custom_initial_data),
+            custom_initial_data_size)
+            .AsVector();
+
+    name_ = v[0].AsString().str();
+    if (!nodedef_.ParseFromString(v[1].AsString().str())) {
+      nodedef_.Clear();
+      return tensorflow::errors::Internal(
+          "Failed to parse data into a valid NodeDef");
+    }
+
+    // Fill NodeDef with defaults if it's a valid op.
+    TF_RETURN_IF_ERROR(
+        tensorflow::OpRegistry::Global()->LookUp(nodedef_.op(), &op_reg_data_));
+    AddDefaultsToNodeDef(op_reg_data_->op_def, &nodedef_);
+
+    // Force disable the use of inter op parallelism to prevent deadlocks in
+    // Tensorflow Function Library Runtime when only one thread is allowed.
+    // This changes the threadpool that is used by TF's data ops by passing it
+    // to the CapturedFunction instantiate function.
+    //
+    // It should be ok to remove this when/if the tensorflow::Executor::Run
+    // function is changed not to call the RunAsync function and wait on its
+    // completion. See b/304799442 for more context.
+    const auto& op_def = op_reg_data_->op_def;
+    for (const auto& attr : op_def.attr()) {
+      if (attr.name() == kInterOpParallelismAttrName) {
+        (*nodedef_.mutable_attr())[kInterOpParallelismAttrName].set_b(false);
+        break;
+      }
+    }
+
+    return absl::OkStatus();
+  }
+
+  absl::Status BuildOpKernelRunner(tensorflow::EagerContext* eager_context) {
+    // Create tensorflow::OpKernel on host CPU.
+    TF_ASSIGN_OR_RETURN(op_kernel_runner_,
+                        tensorflow::tfrt_stub::OpKernelRunner::Create(
+                            name_, inputs_.Size(), /*attr_builder=*/
+                            [this](tensorflow::AttrValueMap* attr_value_map) {
+                              *attr_value_map = nodedef_.attr();
+                              return absl::OkStatus();
+                            },
+                            *eager_context->pflr(),
+                            eager_context->local_device_mgr()->HostCPU()));
+
+    return absl::OkStatus();
+  }
+
+  absl::Status BuildOpKernelInputs(
+      const BufferMap* buffer_map,
+      tensorflow::tfrt_stub::OpKernelRunState* run_state) {
+    run_state->input_tf_tensors.resize(inputs_.Size());
+    run_state->input_tf_tensor_values.resize(inputs_.Size());
+
+    for (int i = 0; i < inputs_.Size(); ++i) {
+      int input_index = inputs_.TfLiteIndex(i);
+      TensorSource s = inputs_.GetTensorSource(i);
+      if (!s.node) {
+        // This input is not produced by this TF subgraph (it could be a TF
+        // Lite native buffer, or could be produced by a separater subgraph). We
+        // need to fetch it from the delegate's buffer_map.
+        if (!buffer_map->HasTensor(input_index)) {
+          return tensorflow::errors::Internal(
+              "Cannot read from invalid tensor index ", input_index);
+        }
+        run_state->input_tf_tensors[i] = buffer_map->GetTensor(input_index);
+      } else {
+        // If this is a forwardable tensor, we will remove it from the previous
+        // op's list, giving TF the opportunity to reuse its buffer.
+        if (inputs_.IsForwardable(i)) {
+          run_state->input_tf_tensors[i] =
+              s.node->outputs_.ReleaseTensor(s.node_output_index);
+        } else {
+          run_state->input_tf_tensors[i] =
+              s.node->outputs_.GetTensor(s.node_output_index);
+        }
+      }
+      run_state->input_tf_tensor_values[i].tensor =
+          &run_state->input_tf_tensors[i];
+    }
+    return absl::OkStatus();
+  }
+
+  // Returns whether an output tensor should be preserved in the buffer map by
+  // checking its lifetime information.
+  // The eager tensor doesn't need to be persisted in the buffer map if it has
+  // no future uses in the graph.
+  bool ShouldPersistTensorflowTensor(TfLiteContext* context,
+                                     const OpDataInfo* shared_info,
+                                     int tensor_index, int node_index) {
+    TfLiteTensor* tensor = &context->tensors[tensor_index];
+    // Always persist variant|resource|string tensors since they have special
+    // storage requirement.
+    if (IsResourceOrVariant(tensor) || tensor->type == kTfLiteString) {
+      return true;
+    }
+
+    auto it = shared_info->tensor_release_map->find(tensor_index);
+    return it != shared_info->tensor_release_map->end() &&
+           it->second > node_index;
+  }
+
+  // Copies the data of Tensorflow tensor into the corresponding TfLite tensor,
+  // after copy is done release the original tensor so that memory could be
+  // released by TF runtime.
+  TfLiteStatus CopyToTfLiteTensor(TfLiteContext* context,
+                                  OpDataInfo* shared_info, TfLiteTensor* tensor,
+                                  tensorflow::Tensor* tf_tensor,
+                                  int tensor_index) const {
+    if (tensor->allocation_type == kTfLiteDynamic) {
+      // For dynamic tensors, update the TfLite tensor's shape information from
+      // the Tensorflow tensor.
+      CopyShapeAndType(context, *tf_tensor, tensor);
+    }
+    absl::string_view t_data = tf_tensor->tensor_data();
+    if (tf_tensor->NumElements() != NumElements(tensor) ||
+        tf_tensor->TotalBytes() != tensor->bytes) {
+      TF_LITE_KERNEL_LOG(context,
+                         "FlexDelegate: Tensor %s(%d) buffer size mismatch "
+                         "%zu(%" PRId64 ") != %zu(%" PRId64 ")",
+                         tensor->name, tensor_index, tf_tensor->TotalBytes(),
+                         tf_tensor->NumElements(), tensor->bytes,
+                         NumElements(tensor));
+      return kTfLiteError;
+    }
+    // Copy TF tensor's data content into TfLiteTensor, and release the tensor.
+    memcpy(tensor->data.raw, t_data.data(), t_data.size());
+    *tf_tensor = {};
+    shared_info->already_transferred_outputs.insert(tensor_index);
+    return kTfLiteOk;
+  }
+
+  // TODO(b/204479285): Release tensors from BufferMap if it has no future
+  // uses.
+  absl::Status MaybePersistTensorflowOutputs(TfLiteContext* context,
+                                             OpDataInfo* shared_info,
+                                             int node_index) {
+    auto* tensors = outputs_.GetTensors();
+
+    for (int i = 0; i < outputs_.Size(); ++i) {
+      if (outputs_.IsSubgraphOutput(i)) {
+        tensorflow::Tensor& tf_tensor = tensors->at(i);
+        const int tflite_index = outputs_.TfLiteIndex(i);
+        TfLiteTensor* tensor = &context->tensors[tflite_index];
+        if (!ShouldPersistTensorflowTensor(context, shared_info, tflite_index,
+                                           node_index)) {
+          if (CopyToTfLiteTensor(context, shared_info, tensor, &tf_tensor,
+                                 tflite_index) != kTfLiteOk) {
+            return absl::Status(absl::StatusCode::kInternal,
+                                "failed to copy data from TF tensor");
+          }
+        } else {
+          shared_info->buffer_map->SetFromTensorFlow(outputs_.TfLiteIndex(i),
+                                                     tf_tensor);
+        }
+      }
+    }
+    return absl::OkStatus();
+  }
+
+ private:
+  OpNode(const OpNode&) = delete;
+  OpNode& operator=(const OpNode&) = delete;
+
+  // The name of the TensorFlow op to execute.
+  string name_;
+  // Index of this node into TF Lite's operator list.
+  int index_;
+  // The corresponding NodeDef, containing the attributes for the op.
+  tensorflow::NodeDef nodedef_;
+  // The corresponding OpRegistrationData pointer.
+  const tensorflow::OpRegistrationData* op_reg_data_;
+  // List of inputs, as TF Lite tensor indices.
+  OpInputs inputs_;
+  // List of outputs, as TF Lite tensor indices.
+  OpOutputs outputs_;
+
+  tensorflow::tfrt_stub::OpKernelRunner op_kernel_runner_;
+};
+
+// The larger 'op', which contains all the nodes in a supported subgraph.
+struct OpData {
+  tensorflow::EagerContext* eager_context;
+  tensorflow::CancellationManager* cancellation_manager;
+  std::vector<std::unique_ptr<OpNode>> nodes;
+  std::vector<int> subgraph_inputs;
+  std::vector<int> subgraph_outputs;
+  std::set<int>
+      disable_reusing_buffer_tensors;  // A list of input tensor indexes which
+                                       // input buffer should not be reused by
+                                       // tensorflow::Tensor.
+  OpDataInfo shared_info;
+};
+
+absl::Status DelegateKernel::ExecuteOpKernelRunner(
+    tensorflow::tfrt_stub::OpKernelRunState* run_state, TfLiteContext* context,
+    OpNode* node_data) {
+  const auto& op_kernel_runner = node_data->op_kernel_runner();
+
+  if (op_kernel_runner.op_kernel()->num_outputs() != node_data->NumOutputs()) {
+    return absl::InternalError(
+        "Unexpected number of outputs from tensorflow::OpKernel");
+  }
+
+  TF_RETURN_IF_ERROR(node_data->BuildOpKernelInputs(
+      op_data_->shared_info.buffer_map, run_state));
+
+  run_state->params.inputs = run_state->input_tf_tensor_values;
+  run_state->params.op_kernel = op_kernel_runner.op_kernel();
+  run_state->params.input_alloc_attrs = op_kernel_runner.input_alloc_attrs();
+  run_state->params.output_attr_array =
+      op_kernel_runner.output_alloc_attrs().data();
+  run_state->params.function_library =
+      op_kernel_runner.function_library_runtime();
+
+  tensorflow::OpKernelContext tf_context(&run_state->params,
+                                         node_data->NumOutputs());
+  op_kernel_runner.Run(&tf_context);
+  TF_RETURN_IF_ERROR(tf_context.status());
+
+  auto& outputs = *node_data->mutable_outputs()->GetTensors();
+  for (int i = 0; i < tf_context.num_outputs(); ++i) {
+    outputs[i] = std::move(*tf_context.mutable_output(i));
+  }
+
+  return node_data->MaybePersistTensorflowOutputs(
+      context, &(op_data_->shared_info), node_data->index());
+}
+
+DelegateKernel::DelegateKernel() : op_data_(new OpData) {}
+DelegateKernel::~DelegateKernel() = default;
+
+TfLiteStatus DelegateKernel::Init(TfLiteContext* context,
+                                  const TfLiteDelegateParams* params) {
+  auto* flex_delegate_data =
+      reinterpret_cast<FlexDelegate*>(params->delegate->data_)->mutable_data();
+  op_data_->eager_context = flex_delegate_data->GetEagerContext();
+  op_data_->cancellation_manager = flex_delegate_data->GetCancellationManager();
+  op_data_->shared_info.buffer_map = flex_delegate_data->GetBufferMap(context);
+  op_data_->shared_info.tensor_release_map =
+      flex_delegate_data->GetTensorReleaseMap(context);
+
+  TF_LITE_ENSURE(context, params->output_tensors != nullptr);
+  std::set<int> output_set;
+  for (auto tensor_index : TfLiteIntArrayView(params->output_tensors)) {
+    op_data_->subgraph_outputs.push_back(tensor_index);
+    output_set.insert(tensor_index);
+  }
+
+  TF_LITE_ENSURE(context, params->input_tensors != nullptr);
+  for (auto tensor_index : TfLiteIntArrayView(params->input_tensors)) {
+    op_data_->subgraph_inputs.push_back(tensor_index);
+  }
+  std::set<int> subgraph_inputs(op_data_->subgraph_inputs.begin(),
+                                op_data_->subgraph_inputs.end());
+
+  op_data_->nodes.reserve(params->nodes_to_replace->size);
+
+  TF_LITE_ENSURE(context, params->nodes_to_replace != nullptr);
+  absl::Status status;
+
+  // Now we explicitly disable reusing TFLite tensor buffers for certain TF ops,
+  // since those ops might produce results which keep reference of the input
+  // tensors (buffer forwarding).
+  auto check_if_op_reuses_input = [](absl::string_view op_name) {
+    static const auto* const kReusingOps =
+        new absl::flat_hash_set<absl::string_view>(
+            {"TensorListPushBack", "TensorListSetItem", "SparseReshape",
+             "StridedSlice", "RaggedTensorToVariant", "TensorMapInsert",
+             "AssignVariableOp", "TensorArrayWriteV3", "QueueEnqueueV2"});
+    return kReusingOps->contains(op_name);
+    // TensorMapInsert hashes a tensor using a string_view of the key tensor.
+    // If the key tensor is shared with TfLite, the memory be reused. The string
+    // view will also change - it stores a ptr and a size, not the data so the
+    // data must be conserved or a false collision will be detected.
+  };
+
+  for (auto node_index : TfLiteIntArrayView(params->nodes_to_replace)) {
+    TfLiteNode* node;
+    TfLiteRegistration* reg;
+    context->GetNodeAndRegistration(context, node_index, &node, &reg);
+
+    op_data_->nodes.emplace_back(new OpNode(node->inputs, node->outputs));
+    OpNode& node_data = *op_data_->nodes.back();
+
+    node_data.set_index(node_index);
+    node_data.set_name("");
+
+    status = node_data.InitializeNodeDef(node->custom_initial_data,
+                                         node->custom_initial_data_size);
+    if (!status.ok()) break;
+    status = node_data.BuildOpKernelRunner(op_data_->eager_context);
+    if (!status.ok()) break;
+
+    // For each node handled by this delegate partition, record the mapping
+    // information between each input tensor and the node index. The node index
+    // is the index of the last node in execution order that uses this tensor.
+    // So the tensor is no longer needed after this last node is executed.
+    // Since we execute in order, then the maximum index is the index of the
+    // last node that needs this tensor.
+    for (auto tensor_index : TfLiteIntArrayView(node->inputs)) {
+      int node_id = node_index;
+      if (const std::map<int, int>::iterator it =
+              op_data_->shared_info.tensor_release_map->find(tensor_index);
+          it != op_data_->shared_info.tensor_release_map->end()) {
+        node_id = std::max(it->second, node_index);
+      }
+      (*op_data_->shared_info.tensor_release_map)[tensor_index] = node_id;
+
+      if (subgraph_inputs.count(tensor_index) &&
+          check_if_op_reuses_input(node_data.nodedef().op())) {
+        op_data_->disable_reusing_buffer_tensors.insert(tensor_index);
+      }
+    }
+  }
+
+  TF_LITE_ENSURE_STATUS(ConvertStatus(context, status));
+
+  // Given a TfLite tensor index, return the OpNode that produces it,
+  // along with it index into that OpNodes list of outputs.
+  std::map<int, TensorSource> tflite_tensor_sources;
+
+  // Find out how each tensor is produced. This does not account for
+  // tensors that are not produced by tensorflow::Opkernels.
+  for (auto& node_data : op_data_->nodes) {
+    node_data->mutable_outputs()->InitializeGraphOutputs(output_set);
+    for (int i = 0; i < node_data->outputs().Size(); ++i) {
+      int output_index = node_data->outputs().TfLiteIndex(i);
+      tflite_tensor_sources[output_index] = TensorSource{node_data.get(), i};
+    }
+  }
+
+  // For each node, resolve the inputs, so we can keep pointers to the nodes
+  // that produces them.
+  for (auto& node_data : op_data_->nodes) {
+    node_data->mutable_inputs()->InitializeTensorSources(tflite_tensor_sources);
+  }
+  return kTfLiteOk;
+}
+
+TfLiteStatus DelegateKernel::Prepare(TfLiteContext* context, TfLiteNode* node) {
+  TF_LITE_ENSURE_MSG(
+      context, op_data_->eager_context != nullptr,
+      "Failed to initialize eager context. This often happens when a CPU "
+      "device has not been registered, presumably because some symbols from "
+      "tensorflow/core:core_cpu_impl were not linked into the binary.");
+
+  // We will keep track of the number of references to each tensor in the
+  // graph, so we can make them "forwardable" if there is only one reference.
+  std::map<int, int> tensor_ref_count;
+
+  // Whenever we find a constant tensor, insert it in the buffer map.
+  BufferMap* buffer_map = op_data_->shared_info.buffer_map;
+  for (auto tensor_index : op_data_->subgraph_inputs) {
+    TfLiteTensor* tensor = &context->tensors[tensor_index];
+    if (IsConstantTensor(tensor)) {
+      if (!tensor->data_is_stale || !buffer_map->HasTensor(tensor_index)) {
+        buffer_map->SetFromTfLite(tensor_index, tensor);
+      }
+    }
+
+    // Input tensors should never be forwarded so we increment their ref counts
+    // twice: once for this graph and another for the possibility of them being
+    // used by another subgraph, or being an output of the full graph.
+    tensor_ref_count[tensor_index] += 2;
+  }
+
+  // Output shapes which may have initially been inferable may no longer be
+  // after ResizeInputTensor has been called, so it must be checked again.
+  if (shapes_are_valid_) {
+    shapes_are_valid_ =
+        (ValidateOutputTensorShapeConsistency(context) == kTfLiteOk);
+    if (shapes_are_valid_) {
+      TFLITE_LOG(tflite::TFLITE_LOG_INFO,
+                 "FlexDelegate: All tensor shapes are consistent.");
+    } else {
+      TFLITE_LOG(tflite::TFLITE_LOG_WARNING,
+                 "FlexDelegate: Some tensor shapes are inconsistent.");
+    }
+  }
+
+  // All output tensors are allocated by TensorFlow, so we mark them as
+  // kTfLiteDynamic.
+  for (auto tensor_index : op_data_->subgraph_outputs) {
+    if (!shapes_are_valid_) {
+      SetTensorToDynamic(&context->tensors[tensor_index]);
+    }
+    ++tensor_ref_count[tensor_index];
+  }
+
+  for (const auto& node_data : op_data_->nodes) {
+    if (node_data->nodedef().op().empty()) {
+      TF_LITE_KERNEL_LOG(context, "Invalid NodeDef in Flex op '%s'",
+                         node_data->name().c_str());
+      return kTfLiteError;
+    }
+    TF_LITE_ENSURE(context, node_data->op_kernel_runner());
+
+    for (int i = 0; i < node_data->inputs().Size(); ++i) {
+      ++tensor_ref_count[node_data->inputs().TfLiteIndex(i)];
+    }
+  }
+
+  // All tensors that are referenced exactly once are marked as "forwardable",
+  // meaning that we will allow TensorFlow to reuse its buffer as the output of
+  // an op.
+  for (auto& node_data : op_data_->nodes) {
+    for (int i = 0; i < node_data->inputs().Size(); ++i) {
+      bool f = (tensor_ref_count[node_data->inputs().TfLiteIndex(i)] == 1);
+      node_data->mutable_inputs()->SetForwardable(i, f);
+    }
+  }
+
+  return kTfLiteOk;
+}
+
+TfLiteStatus DelegateKernel::ValidateOutputTensorShapeConsistency(
+    TfLiteContext* context) const {
+  for (const auto& node_data : op_data_->nodes) {
+    auto op_name = node_data->name().c_str();
+    // Create an InferenceContext object.
+    auto num_inputs = node_data->inputs().Size();
+    std::vector<const tensorflow::Tensor*> input_tensors_vector(num_inputs,
+                                                                nullptr);
+    InferenceContext c(
+        TF_GRAPH_DEF_VERSION, node_data->nodedef(),
+        node_data->op_reg_data()->op_def, std::vector<ShapeHandle>(num_inputs),
+        input_tensors_vector, {},
+        std::vector<std::unique_ptr<std::vector<ShapeAndType>>>());
+
+    // Set input_shapes for ShapeInferenceFn.
+    for (int i = 0; i < num_inputs; ++i) {
+      const auto input_tensor_index = node_data->inputs().TfLiteIndex(i);
+      TfLiteTensor* tfl_tensor = &context->tensors[input_tensor_index];
+      // Provide constant input tensors since some op ("RFFT") needs it to
+      // calculate the output shape.
+      if (IsConstantTensor(tfl_tensor)) {
+        input_tensors_vector[i] =
+            op_data_->shared_info.buffer_map->GetTensorPtr(input_tensor_index);
+      }
+      const auto dims_array = tfl_tensor->dims;
+      std::vector<DimensionHandle> dims(dims_array->size);
+      for (int j = 0; j < dims_array->size; ++j) {
+        dims[j] = c.MakeDim(dims_array->data[j]);
+      }
+      c.SetInput(i, c.MakeShape(dims));
+    }
+    c.set_input_tensors(input_tensors_vector);
+
+    absl::Status status = c.construction_status();
+    if (!status.ok()) {
+      TFLITE_LOG(tflite::TFLITE_LOG_WARNING,
+                 "Shape construction failed for op '%s'", op_name);
+      return kTfLiteError;
+    }
+
+    // Run ShapeInferenceFn to calculate output shapes.
+    if (node_data->op_reg_data()->shape_inference_fn == nullptr) {
+      TFLITE_LOG(tflite::TFLITE_LOG_WARNING,
+                 "No shape inference function exists for op '%s'", op_name);
+      return kTfLiteError;
+    }
+    status = c.Run(node_data->op_reg_data()->shape_inference_fn);
+
+    // Compare calculated output shapes with node_data->outputs
+    auto num_outputs = node_data->outputs().Size();
+    if (num_outputs != c.num_outputs()) {
+      TFLITE_LOG(tflite::TFLITE_LOG_WARNING,
+                 "Number of output tensors are mismatched for op '%s' %d != %d",
+                 op_name, num_outputs, c.num_outputs());
+      return kTfLiteError;
+    }
+    for (int i = 0; i < num_outputs; ++i) {
+      const auto output_tensor_index = node_data->outputs().TfLiteIndex(i);
+      TfLiteTensor* tfl_tensor = &context->tensors[output_tensor_index];
+      // tfl_tensor->dims only has valid information if the given model is
+      // converted by the MLIR converter. Also when ResizeInputTensor() is
+      // called the dims information becomes invalid.
+      const std::string tfl_shape_string =
+          GetShapeDebugString(tfl_tensor->dims);
+      const std::string calculated_shape_string = c.DebugString(c.output(i));
+      // Getting a shape string via c.DebugString() is the easiest way to get
+      // the shape information of the given ShapeHandle for now.
+      // TODO(b/169017408): Find a better approach without using debug string.
+      if (tfl_shape_string != calculated_shape_string) {
+        if ((strcmp(op_name, kReadVariableOp) == 0) &&
+            (tfl_tensor->dims->size > 0)) {
+          // If ReadVariableOp has an output with valid shape, use it since
+          // ShapeInferenceFn of ReadVariableOp doesn't work well without having
+          // a valid resource handle.
+          continue;
+        }
+
+        TFLITE_LOG(tflite::TFLITE_LOG_WARNING,
+                   "op '%s' output%d tensor#%d shape mismatch for  %s != %s",
+                   op_name, i, output_tensor_index, tfl_shape_string.c_str(),
+                   calculated_shape_string.c_str());
+        return kTfLiteError;
+      }
+    }
+  }
+  return kTfLiteOk;
+}
+
+static tensorflow::CancellationManager* GetDefaultCancellationManager() {
+  static auto* const cancellation_manager = new tensorflow::CancellationManager;
+  return cancellation_manager;
+}
+
+TfLiteStatus DelegateKernel::Eval(TfLiteContext* context, TfLiteNode* node) {
+  BufferMap* buffer_map = op_data_->shared_info.buffer_map;
+
+  // Insert a tensor in the buffer map for all inputs that are not constant.
+  // Constants were handled in Prepare() already.
+  for (auto tensor_index : op_data_->subgraph_inputs) {
+    TfLiteTensor* tensor = &context->tensors[tensor_index];
+    if (!IsConstantTensor(tensor)) {
+      // If this tensor is part of an earlier TF subgraph we should not add it
+      // to the BufferMap again, because TF already knows about it and its
+      // contents are kept automatically up-to-date.
+      if (!tensor->data_is_stale || !buffer_map->HasTensor(tensor_index)) {
+        buffer_map->SetFromTfLite(
+            tensor_index, tensor,
+            !op_data_->disable_reusing_buffer_tensors.count(tensor_index));
+      }
+    }
+  }
+
+  auto& eager_context = *op_data_->eager_context;
+
+  {
+    tensorflow::tfrt_stub::OpKernelRunState run_state;
+
+    run_state.params.step_container = eager_context.StepContainer();
+    auto* device = eager_context.local_device_mgr()->HostCPU();
+    run_state.params.device = device;
+    run_state.params.resource_manager = device->resource_manager();
+    run_state.params.runner = eager_context.runner();
+    run_state.params.cancellation_manager =
+        op_data_->cancellation_manager ? op_data_->cancellation_manager
+                                       : GetDefaultCancellationManager();
+    // TODO(b/179048776): Set up remaining params such as collective and
+    // rendezvous.
+
+    // Execute the TensorFlow Ops sequentially.
+    for (auto& node_data : op_data_->nodes) {
+      TFLITE_SCOPED_DELEGATE_PROFILED_OPERATOR_PROFILE(
+          reinterpret_cast<Profiler*>(context->profiler),
+          node_data->name().c_str(), node_data->index());
+
+      if (op_data_->cancellation_manager != nullptr &&
+          op_data_->cancellation_manager->IsCancelled()) {
+        TF_LITE_KERNEL_LOG(
+            context, "Client requested cancel during DelegateKernel::Eval");
+        return kTfLiteError;
+      }
+
+      auto status = ExecuteOpKernelRunner(&run_state, context, node_data.get());
+      TF_LITE_ENSURE_OK(context, ConvertStatus(context, status));
+    }
+  }
+
+  for (auto tensor_index : op_data_->subgraph_outputs) {
+    if (op_data_->shared_info.already_transferred_outputs.count(tensor_index) !=
+        0) {
+      // Skip if a tensor output has already been copied to a TfLiteTensor.
+      continue;
+    }
+    if (!buffer_map->HasTensor(tensor_index)) {
+      TF_LITE_KERNEL_LOG(context, "Cannot write to invalid tensor index %d",
+                         tensor_index);
+      return kTfLiteError;
+    }
+
+    // Copy TF tensor data to TFL allocated buffer for non dynamic tensors.
+    // For dynamic tensors, copy shape and put buffer_handle for the later
+    // CopyFromBufferHandle() call.
+    TfLiteTensor* tensor = &context->tensors[tensor_index];
+    const tensorflow::Tensor& tf_tensor = buffer_map->GetTensor(tensor_index);
+    if (tensor->allocation_type == kTfLiteDynamic) {
+      TF_LITE_ENSURE_OK(context, CopyShapeAndType(context, tf_tensor, tensor));
+      tensor->buffer_handle = tensor_index;
+      tensor->data_is_stale = true;
+      continue;
+    }
+    // If the tensor isn't dynamic, we can copy data directly to the buffer of
+    // the tensor. Before copying the data, check if the target buffer has
+    // expected size.
+    if (tf_tensor.NumElements() != NumElements(tensor) ||
+        tf_tensor.TotalBytes() != tensor->bytes) {
+      TF_LITE_KERNEL_LOG(context,
+                         "FlexDelegate: Tensor %s(%d) buffer size mismatch "
+                         "%zu(%" PRId64 ") != %zu(%" PRId64 ")",
+                         tensor->name, tensor_index, tf_tensor.TotalBytes(),
+                         tf_tensor.NumElements(), tensor->bytes,
+                         NumElements(tensor));
+      return kTfLiteError;
+    }
+    absl::string_view t_data = tf_tensor.tensor_data();
+    memcpy(tensor->data.raw, t_data.data(), t_data.size());
+  }
+
+  return kTfLiteOk;
+}
+
+const std::map<int, int>& DelegateKernel::GetTensorReleaseMap() const {
+  return *(op_data_->shared_info.tensor_release_map);
+}
+
+}  // namespace flex
+}  // namespace tflite

tensorflow/lite/delegates/flex/kernel.h

@@ -0,0 +1,72 @@
+/* Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_DELEGATES_FLEX_KERNEL_H_
+#define TENSORFLOW_LITE_DELEGATES_FLEX_KERNEL_H_
+
+#include <memory>
+
+#include "tensorflow/core/platform/status.h"
+#include "tensorflow/core/tfrt/fallback/op_kernel_runner.h"
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/delegates/utils/simple_delegate.h"
+
+namespace tflite {
+namespace flex {
+
+namespace testing {
+class KernelTest;  // friend class declaration.
+}  // namespace testing
+
+struct OpData;
+struct OpNode;
+
+class DelegateKernel : public SimpleDelegateKernelInterface {
+ public:
+  DelegateKernel();
+  ~DelegateKernel() override;
+
+  TfLiteStatus Init(TfLiteContext* context,
+                    const TfLiteDelegateParams* params) override;
+  TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) override;
+  TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) override;
+
+ private:
+  friend class tflite::flex::testing::KernelTest;
+
+  // Validates that the computed output tensor shape for the Flex node matches
+  // the existing output shape assigned to the output tensor.
+  TfLiteStatus ValidateOutputTensorShapeConsistency(
+      TfLiteContext* context) const;
+
+  // Executes the Tensorflow op based on the inputs/outputs/attributes
+  // information represented in the `node_data`.
+  absl::Status ExecuteOpKernelRunner(
+      tensorflow::tfrt_stub::OpKernelRunState* run_state,
+      TfLiteContext* context, OpNode* node_data);
+
+  // Returns the tensor release map held in `op_data_`;
+  const std::map<int, int>& GetTensorReleaseMap() const;
+
+  std::unique_ptr<OpData> op_data_;
+
+  // Indicates that the output shapes may be inferred using the input shapes and
+  // May be allocated during Prepare.
+  bool shapes_are_valid_ = true;
+};
+
+}  // namespace flex
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_DELEGATES_FLEX_KERNEL_H_

tensorflow/lite/delegates/flex/kernel_test.cc

@@ -0,0 +1,516 @@
+/* Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include "tensorflow/lite/delegates/flex/kernel.h"
+
+#include <functional>
+#include <initializer_list>
+#include <memory>
+#include <utility>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "tensorflow/lite/core/c/c_api_types.h"
+#include "tensorflow/lite/delegates/flex/delegate.h"
+#include "tensorflow/lite/delegates/flex/delegate_data.h"
+#include "tensorflow/lite/delegates/flex/test_util.h"
+#include "tensorflow/lite/kernels/kernel_util.h"
+
+namespace tflite {
+namespace flex {
+namespace testing {
+
+using ::testing::ContainsRegex;
+using ::testing::ElementsAre;
+using ::testing::ElementsAreArray;
+using ::testing::Pair;
+using ::testing::UnorderedElementsAre;
+
+// A testing flex delegate that supports every node regardless whether it's
+// actually supported or not. It's only for testing certain scenarios.
+class TestFlexDelegate : public FlexDelegate {
+ protected:
+  bool IsNodeSupportedByDelegate(const TfLiteRegistration* registration,
+                                 const TfLiteNode* node,
+                                 TfLiteContext* context) const override {
+    return true;
+  }
+};
+
+class KernelTest : public testing::FlexModelTest {
+ public:
+  static constexpr int kOnes = 1;  // This is the index of a tensor of 1's.
+  static constexpr int kTwos = 2;  // This is the index of a tensor of 2's.
+  static constexpr int kMaxTensors = 30;
+
+  KernelTest() {
+    interpreter_ = std::make_unique<Interpreter>(&error_reporter_);
+  }
+
+  TfLiteStatus ApplyFlexDelegate(
+      std::unique_ptr<FlexDelegate> delegate = nullptr) {
+    auto flex_delegate = FlexDelegate::Create(std::move(delegate));
+    delegate_data_ =
+        reinterpret_cast<FlexDelegate*>(flex_delegate->data_)->mutable_data();
+    CHECK_OK(delegate_data_->Prepare(tensorflow::SessionOptions{}));
+    return interpreter_->ModifyGraphWithDelegate(std::move(flex_delegate));
+  }
+
+  const std::map<int, int>& GetTensorReleaseMap(DelegateKernel* kernel) {
+    return kernel->GetTensorReleaseMap();
+  }
+
+ protected:
+  tflite::flex::DelegateData* delegate_data_;
+};
+
+TEST_F(KernelTest, FullGraph) {
+  // Define the graph.
+  AddTensors(9, {0, 3}, {8}, kTfLiteFloat32, {3});
+
+  AddTfOp(testing::kUnpack, {0}, {1, 2});
+  AddTfOp(testing::kUnpack, {3}, {4, 5});
+  AddTfOp(testing::kAdd, {1, 4}, {6});
+  AddTfOp(testing::kAdd, {2, 5}, {7});
+  AddTfOp(testing::kMul, {6, 7}, {8});
+
+  ApplyFlexDelegate();
+
+  // Define inputs.
+  SetShape(0, {2, 2, 1});
+  SetValues(0, {1.1f, 2.2f, 3.3f, 4.4f});
+  SetShape(3, {2, 2, 1});
+  SetValues(3, {1.1f, 2.2f, 3.3f, 4.4f});
+
+  ASSERT_TRUE(Invoke());
+
+  ASSERT_THAT(GetShape(8), ElementsAre(2, 1));
+  ASSERT_THAT(GetValues(8), ElementsAre(14.52f, 38.72f));
+
+  // Try again with different inputs
+  SetShape(0, {2, 3, 1});
+  SetValues(0, {2.0f, 2.0f, 3.0f, 3.0f, 4.0f, 4.0f});
+  SetShape(3, {2, 3, 1});
+  SetValues(3, {2.0f, 2.0f, 3.0f, 3.0f, 4.0f, 4.0f});
+
+  ASSERT_TRUE(Invoke());
+
+  ASSERT_THAT(GetShape(8), ElementsAre(3, 1));
+  ASSERT_THAT(GetValues(8), ElementsAre(24.0f, 32.0f, 48.0f));
+}
+
+TEST_F(KernelTest, ValidateTensorReleaseMap) {
+  // Define the graph.
+  //        0           3
+  //        |           |
+  //      Unpack_0    Unpack_1
+  //       /  \        / \
+  //      1    2      4   5
+  //      |____|_______|__|
+  //         | |__________|
+  //         |      |
+  //        Add_2  Add_3
+  //         |      |
+  //         6      7
+  //         \______/
+  //             |
+  //            Mul_4
+  //             |
+  //             8
+  AddTensors(9, {0, 3}, {8}, kTfLiteFloat32, {3});
+  AddTfOp(testing::kUnpack, {0}, {1, 2});
+  AddTfOp(testing::kUnpack, {3}, {4, 5});
+  AddTfOp(testing::kAdd, {1, 4}, {6});
+  AddTfOp(testing::kAdd, {2, 5}, {7});
+  AddTfOp(testing::kMul, {6, 7}, {8});
+
+  ApplyFlexDelegate();
+
+  const int node_size = interpreter_->primary_subgraph().nodes_size();
+  const std::pair<TfLiteNode, TfLiteRegistration>* node_and_reg =
+      interpreter_->primary_subgraph().node_and_registration(node_size - 1);
+
+  DelegateKernel* delegate_kernel =
+      reinterpret_cast<DelegateKernel*>(node_and_reg->first.user_data);
+  const auto& tensor_release_map = GetTensorReleaseMap(delegate_kernel);
+  // Validate the tensor release mapping.
+  EXPECT_THAT(
+      tensor_release_map,
+      UnorderedElementsAre(Pair(0, 0), Pair(1, 2), Pair(2, 3), Pair(3, 1),
+                           Pair(4, 2), Pair(5, 3), Pair(6, 4), Pair(7, 4)));
+}
+
+TEST_F(KernelTest, PersistEagerTensor) {
+  // Define the graph.
+  //        0           3
+  //        |           |
+  //      Unpack_0    Unpack_1
+  //       /  \        / \
+  //      1    2      4   5
+  //      |____|_______|__|
+  //         | |__________|
+  //         |      |
+  //        Add_2  Add_3
+  //         |      |
+  //         6      7
+  //         | \   /
+  //         | TFL_MUL
+  //         |    |
+  //         |    8
+  //         |____|
+  //             AddN
+  //              |
+  //              9
+  AddTensors(10, {0, 3}, {9}, kTfLiteFloat32, {3});
+
+  AddTfOp(testing::kUnpack, {0}, {1, 2});
+  AddTfOp(testing::kUnpack, {3}, {4, 5});
+  AddTfOp(testing::kAdd, {1, 4}, {6});
+  AddTfOp(testing::kAdd, {2, 5}, {7});
+  AddTfLiteMulOp({6, 7}, {8});
+  AddTfOp(testing::kAdd, {6, 8}, {9});
+
+  ApplyFlexDelegate();
+
+  // Define inputs.
+  SetShape(0, {2, 2, 1});
+  SetValues(0, {1.1f, 2.2f, 3.3f, 4.4f});
+  SetShape(3, {2, 2, 1});
+  SetValues(3, {1.1f, 2.2f, 3.3f, 4.4f});
+
+  ASSERT_TRUE(Invoke());
+  // Validates that tensor 6 should be preserved in the buffer map.
+  auto* buffer_map =
+      delegate_data_->GetBufferMap(interpreter_->primary_subgraph().context());
+  EXPECT_TRUE(buffer_map->HasTensor(6));
+  EXPECT_FALSE(buffer_map->HasTensor(7));
+}
+
+TEST_F(KernelTest, BadTensorFlowOp) {
+  AddTensors(2, {0}, {1}, kTfLiteFloat32, {3});
+  AddTfOp(testing::kNonExistent, {0}, {1});
+
+  ASSERT_NE(
+      ApplyFlexDelegate(std::unique_ptr<FlexDelegate>(new TestFlexDelegate())),
+      kTfLiteOk);
+
+  ASSERT_THAT(error_reporter().error_messages(),
+              ContainsRegex("Op type not registered 'NonExistentOp'"));
+}
+
+TEST_F(KernelTest, BadNumberOfOutputs) {
+  AddTensors(3, {0}, {1, 2}, kTfLiteFloat32, {3});
+  AddTfOp(testing::kIdentity, {0}, {1, 2});
+
+  ApplyFlexDelegate();
+
+  SetShape(0, {2, 2, 1});
+  SetValues(0, {1.1f, 2.2f, 3.3f, 4.4f});
+
+  ASSERT_FALSE(Invoke());
+  ASSERT_THAT(error_reporter().error_messages(),
+              ContainsRegex("Unexpected number of outputs"));
+}
+
+TEST_F(KernelTest, IncompatibleNodeDef) {
+  AddTensors(2, {0}, {1}, kTfLiteFloat32, {3});
+
+  // Cast is a TF op, but we don't add the proper nodedef to it in AddTfOp.
+  AddTfOp(testing::kIncompatibleNodeDef, {0}, {1});
+
+  ASSERT_NE(ApplyFlexDelegate(), kTfLiteOk);
+
+  ASSERT_THAT(error_reporter().error_messages(),
+              ContainsRegex("No attr named 'SrcT' in NodeDef"));
+}
+
+TEST_F(KernelTest, WrongSetOfNodes) {
+  AddTensors(4, {0}, {3}, kTfLiteFloat32, {3});
+  AddTfOp(testing::kUnpack, {0}, {1, 2});
+  AddTfLiteMulOp({1, 2}, {3});
+
+  // Specify that testing::kMul (#1) is supported when it actually isn't so that
+  // we choose to use the TestFlexDelegate that supports every node regardless
+  // whether it's actually supported or not.
+  ASSERT_NE(
+      ApplyFlexDelegate(std::unique_ptr<FlexDelegate>(new TestFlexDelegate())),
+      kTfLiteOk);
+
+  ASSERT_THAT(error_reporter().error_messages(),
+              ContainsRegex("Cannot convert empty data into a valid NodeDef"));
+}
+
+TEST_F(KernelTest, MixedGraph) {
+  AddTensors(9, {0, 3}, {8}, kTfLiteFloat32, {3});
+
+  AddTfOp(testing::kUnpack, {0}, {1, 2});
+  AddTfOp(testing::kUnpack, {3}, {4, 5});
+  AddTfOp(testing::kAdd, {1, 4}, {6});
+  AddTfOp(testing::kAdd, {2, 5}, {7});
+  AddTfLiteMulOp({6, 7}, {8});
+
+  ApplyFlexDelegate();
+
+  SetShape(0, {2, 2, 1});
+  SetValues(0, {1.1f, 2.2f, 3.3f, 4.4f});
+  SetShape(3, {2, 2, 1});
+  SetValues(3, {1.1f, 2.2f, 3.3f, 4.4f});
+
+  ASSERT_TRUE(Invoke());
+
+  ASSERT_THAT(GetShape(8), ElementsAre(2, 1));
+  ASSERT_THAT(GetValues(8), ElementsAre(14.52f, 38.72f));
+}
+
+// We will build a complex graph where most of the ops are TF ops, but one
+// of them, right in the middle is handle natively by TF Lite. This results
+// in two flex subgraphs to handle the TF ops, and some of the tensors
+// connect those two subgraphs directly.
+TEST_F(KernelTest, SplitGraph) {
+  std::vector<float> a = {3.0f, 1.0f, 0.5f, -1.0f, 4.0f, -1.0f, -2.0f, 5.0f};
+  std::vector<float> b = {0.0f, 1.0f, 1.5f, 3.0f};
+
+  AddTensors(18, {0, 1}, {17}, kTfLiteFloat32, {3});
+
+  // Split the first input. Each branch below uses one half of it.
+  AddTfOp(testing::kUnpack, {0}, {2, 10});
+
+  // The left branch: l = (a0 + b0) * (a2 + b2) + (a1 + b1) * (a3 + b3) = 10
+  AddTfOp(testing::kAdd, {1, 2}, {3});     // => 3, 2, 2, 2
+  AddTfOp(testing::kUnpack, {3}, {4, 5});  // => 3, 2 --- 2, 2
+  AddTfLiteMulOp({4, 5}, {6});             // => 6, 4
+  AddTfOp(testing::kUnpack, {6}, {7, 8});  // => 6 -- 4
+  AddTfOp(testing::kAdd, {7, 8}, {9});     // => 10
+
+  // The right branch: r = (a4 + a6) + (a5 + a7) = 6
+  AddTfOp(testing::kUnpack, {10}, {11, 12});  // => 4, -1 --- -2, 5
+  AddTfOp(testing::kAdd, {11, 12}, {13});     // => 2, 4
+  AddTfOp(testing::kUnpack, {13}, {14, 15});  // => 2 --- 4
+  AddTfOp(testing::kAdd, {14, 15}, {16});     // => 6
+
+  // The two branches added together:
+  AddTfOp(testing::kAdd, {9, 16}, {17});  // => 16
+
+  ApplyFlexDelegate();
+
+  SetShape(0, {2, 2, 2, 1});
+  SetValues(0, a);
+  SetShape(1, {2, 2, 1});
+  SetValues(1, b);
+
+  ASSERT_TRUE(Invoke());
+
+  ASSERT_THAT(GetShape(17), ElementsAre(1));
+  ASSERT_THAT(GetValues(17), ElementsAre(16.0f));
+
+  // Same as above but with slightly different output.
+  // We still expect the result to be l + r where
+  //     l = (a0 + b0) * (a2 + b2) + (a1 + b1) * (a3 + b3)
+  //     r = (a4 + a6) + (a5 + a7)
+  SetShape(0, {2, 2, 2, 1});
+  SetValues(0, {4.0f, 1.0f, 1.5f, -2.0f, 2.0f, 0.0f, -2.0f, 3.0f});
+  SetShape(1, {2, 2, 1});
+  SetValues(1, {0.0f, 2.0f, 1.5f, 3.0f});
+  // So l = (4 + 0) * (1.5 + 1.5) + (1 + 2) * (-2 + 3) =  12 + 3 = 15
+  //    r = (2 - 2) + (0 + 3) = 3
+
+  ASSERT_TRUE(Invoke());
+
+  ASSERT_THAT(GetShape(17), ElementsAre(1));
+  ASSERT_THAT(GetValues(17), ElementsAre(18.0f));
+}
+
+class MultipleSubgraphsTest : public KernelTest {
+ public:
+  static constexpr int kInput = 0;
+
+  void PrepareInterpreter(const std::vector<float>& input) {
+    ApplyFlexDelegate();
+
+    SetShape(kOnes, {3});
+    SetValues(kOnes, {1.0f, 1.0f, 1.0f});
+    SetShape(kTwos, {3});
+    SetValues(kTwos, {2.0f, 2.0f, 2.0f});
+
+    SetValues(kInput, input);
+  }
+
+  std::vector<float> Apply(const std::vector<float>& input,
+                           std::function<float(float)> function) {
+    std::vector<float> result;
+    for (float f : input) {
+      result.push_back(function(f));
+    }
+    return result;
+  }
+};
+
+TEST_F(MultipleSubgraphsTest, ForwardabilityIsLocal) {
+  AddTensors(kMaxTensors, {kInput, kOnes, kTwos}, {12}, kTfLiteFloat32, {3});
+
+  // Only TF tensors can be forwarded, so we build a small first graph
+  // to produce tensor #10. Here #10 is forwardable, because it is only
+  // used once, as an output.
+  AddTfOp(testing::kAdd, {0, kOnes}, {3});
+  AddTfOp(testing::kAdd, {0, kOnes}, {10});
+
+  // The second TF graph, separated from the former by a TF Lite
+  // multiplication, will consume tensor #10, which is not forwardable here
+  // since it is used by more than one op. The existing code will forward the
+  // tensor anyway, because it was deemed to be forwardable by the previous
+  // subgraph.
+  AddTfLiteMulOp({3, kTwos}, {4});
+  AddTfOp(testing::kAdd, {10, 4}, {11});
+  AddTfOp(testing::kAdd, {11, 10}, {7});
+
+  // And a simple TF Lite op trying to access tensor #10, which was removed
+  // from the buffer map. It will cause Invoke() to fail.
+  AddTfLiteMulOp({10, 7}, {12});
+
+  auto input = {3.0f, 4.0f, 5.0f};
+  PrepareInterpreter(input);
+
+  ASSERT_TRUE(Invoke());
+  ASSERT_THAT(GetValues(12), ElementsAreArray(Apply(input, [](float in) {
+                return (4 * in + 4) * (in + 1);
+              })));
+}
+
+// Subgraphs should not remove input tensors from the buffer_map, since
+// they could be necessary for downstream graphs.
+TEST_F(MultipleSubgraphsTest, DoNotRemoveInputTensors) {
+  AddTensors(kMaxTensors, {kInput, kOnes, kTwos}, {12}, kTfLiteFloat32, {3});
+
+  // Only TF tensors can be removed, so we build a small first graph
+  // to produce tensor #10. We make sure it is used by more than one
+  // op, so it is not forwardable here.
+  AddTfOp(testing::kAdd, {0, kOnes}, {3});
+  AddTfOp(testing::kAdd, {0, kOnes}, {10});
+  AddTfOp(testing::kAdd, {10, kOnes}, {15});
+  AddTfOp(testing::kAdd, {10, kOnes}, {16});
+
+  // The second TF graph, separated from the former by a TF Lite
+  // multiplication, will consume tensor #10. The existing code will remove
+  // from the buffer_map all tensors that are not outputs, so #10 will
+  // disappear. Note that we are using #10 in two ops, so it is not forwardable
+  // either.
+  AddTfLiteMulOp({3, kTwos}, {4});
+  AddTfOp(testing::kAdd, {10, 4}, {11});
+  AddTfOp(testing::kAdd, {10, 11}, {7});
+
+  // And a simple TF Lite op trying to access tensor #10, which was removed
+  // from the buffer map. It will cause Invoke() to fail.
+  AddTfLiteMulOp({10, 7}, {12});
+
+  auto input = {3.0f, 4.0f, 5.0f};
+  PrepareInterpreter(input);
+
+  ASSERT_TRUE(Invoke());
+  ASSERT_THAT(GetValues(12), ElementsAreArray(Apply(input, [](float in) {
+                return (4 * in + 4) * (in + 1);
+              })));
+}
+
+// A tensor is deemed forwardable but it happens to be the input to
+// more than one subgraph. It should not be forwarded, otherwise its
+// contents will be overwritten.
+TEST_F(MultipleSubgraphsTest, DoNotForwardInputTensors) {
+  AddTensors(kMaxTensors, {kInput, kOnes, kTwos}, {12}, kTfLiteFloat32, {3});
+
+  // Only TF tensors can be forwarded, so we build a small first graph
+  // to produce tensor #10.
+  AddTfOp(testing::kAdd, {0, kOnes}, {3});
+  AddTfOp(testing::kAdd, {0, kOnes}, {10});
+
+  // The second TF graph, separated from the former by a TF Lite
+  // multiplication, will consume tensor #10 and will think it is forwardable
+  // because it is used by a single op. However, the subgraph doesn't have
+  // enough information to make that judgment, as the input tensor could be
+  // used by another graph further downstream. The existing code will forward
+  // the tensor and remove it from the buffer_map, causing a failure later.
+  AddTfLiteMulOp({3, kTwos}, {4});
+  AddTfOp(testing::kAdd, {10, 4}, {11});
+  AddTfOp(testing::kAdd, {11, 4}, {7});
+
+  // And a simple TF Lite op trying to access tensor #10, which was removed
+  // from the buffer map. It will cause Invoke() to fail.
+  AddTfLiteMulOp({10, 7}, {12});
+
+  auto input = {3.0f, 4.0f, 5.0f};
+  PrepareInterpreter(input);
+
+  ASSERT_TRUE(Invoke());
+  ASSERT_THAT(GetValues(12), ElementsAreArray(Apply(input, [](float in) {
+                return (5 * in + 5) * (in + 1);
+              })));
+}
+
+tensorflow::OpDef MakeOpDef(int num_inputs, int num_outputs) {
+  tensorflow::OpRegistrationData op_reg_data;
+  tensorflow::OpDefBuilder b("dummy");
+  for (int i = 0; i < num_inputs; ++i) {
+    b.Input(absl::StrCat("i", i, ": float"));
+  }
+  for (int i = 0; i < num_outputs; ++i) {
+    b.Output(absl::StrCat("o", i, ": float"));
+  }
+  CHECK_OK(b.Attr("foo:string").Finalize(&op_reg_data));
+  return op_reg_data.op_def;
+}
+
+tensorflow::PartialTensorShape S(std::initializer_list<int64_t> dims) {
+  return tensorflow::PartialTensorShape(dims);
+}
+
+TEST(ValidateOutputTensorShapeConsistencyTest, ShapeHandleDebugString) {
+  // Setup test to contain an input tensor list of size 3.
+  tensorflow::OpDef op_def = MakeOpDef(4, 1);
+  tensorflow::NodeDef def;
+  tensorflow::shape_inference::InferenceContext c(
+      0, def, op_def, {S({1}), S({2, 3}), S({4, 5, 6}), {}}, {}, {}, {});
+  c.SetInput(3, c.UnknownShape());
+
+  std::vector<tensorflow::shape_inference::ShapeHandle> shapes;
+  EXPECT_EQ("[1]", c.DebugString(c.input(0)));
+  EXPECT_EQ("[2,3]", c.DebugString(c.input(1)));
+  EXPECT_EQ("[4,5,6]", c.DebugString(c.input(2)));
+  // c.DebugString() returns "?" for the unknown shape which is different with
+  // "-1" of TFLite. But this is intended behavior since we should use dynamic
+  // tensor for unknown shape so the shape comparison must fail.
+  EXPECT_EQ("?", c.DebugString(c.input(3)));
+}
+
+TEST(ValidateOutputTensorShapeConsistencyTest, GetShapeDebugString) {
+  TfLiteIntArray* dims1 = TfLiteIntArrayCreate(1);
+  dims1->data[0] = 1;
+  EXPECT_EQ("[1]", GetShapeDebugString(dims1));
+  TfLiteIntArrayFree(dims1);
+
+  TfLiteIntArray* dims2 = TfLiteIntArrayCreate(2);
+  dims2->data[0] = 2;
+  dims2->data[1] = 3;
+  EXPECT_EQ("[2,3]", GetShapeDebugString(dims2));
+  TfLiteIntArrayFree(dims2);
+
+  TfLiteIntArray* dims3 = TfLiteIntArrayCreate(3);
+  dims3->data[0] = 4;
+  dims3->data[1] = 5;
+  dims3->data[2] = 6;
+  EXPECT_EQ("[4,5,6]", GetShapeDebugString(dims3));
+  TfLiteIntArrayFree(dims3);
+}
+
+}  // namespace testing
+}  // namespace flex
+}  // namespace tflite

tensorflow/lite/delegates/flex/subgraph_resource.h

@@ -0,0 +1,65 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_DELEGATES_FLEX_SUBGRAPH_RESOURCE_H_
+#define TENSORFLOW_LITE_DELEGATES_FLEX_SUBGRAPH_RESOURCE_H_
+
+#include <memory>
+#include <string>
+
+#include "tensorflow/core/framework/resource_mgr.h"
+#include "tensorflow/core/platform/mutex.h"
+#include "tensorflow/core/platform/thread_annotations.h"
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/core/subgraph.h"
+
+namespace tflite {
+namespace flex {
+
+// This object stores a pointer for a TfLite subgraph and the associated mutex
+// to access the subgraph. Before accessing the TF Lite subgraph, the caller
+// needs to first acquire a lock on the mutex object.
+class TFLiteSubgraphResource : public tensorflow::ResourceBase {
+ public:
+  explicit TFLiteSubgraphResource(Subgraph& subgraph, TfLiteDelegate* delegate)
+      : subgraph_(subgraph), delegate_(delegate) {}
+
+  std::string DebugString() const override { return "TFLiteSubgraphResource"; }
+
+  // Returns the TFLite subgraph. Before calling
+  // this method, the caller needs to acquire the underlying mutex lock.
+  Subgraph& GetSubgraphResource() const TF_EXCLUSIVE_LOCKS_REQUIRED(mutex_) {
+    return subgraph_;
+  }
+
+  tensorflow::mutex& GetExclusiveLock() TF_LOCK_RETURNED(mutex_) {
+    return mutex_;
+  }
+
+  // Returns a pointer to the TfLiteDelegate which this instance of subgraph
+  // is running as part of it.
+  TfLiteDelegate* GetFlexDelegate() TF_EXCLUSIVE_LOCKS_REQUIRED(mutex_) {
+    return delegate_;
+  }
+
+ private:
+  tensorflow::mutex mutex_;
+  Subgraph& subgraph_ TF_GUARDED_BY(mutex_);
+  TfLiteDelegate* delegate_ TF_GUARDED_BY(mutex_) = nullptr;
+};
+
+}  // namespace flex
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_DELEGATES_FLEX_SUBGRAPH_RESOURCE_H_

tensorflow/lite/delegates/flex/test/BUILD

@@ -0,0 +1,115 @@
+load("@bazel_skylib//rules:build_test.bzl", "build_test")
+load("@build_bazel_rules_apple//apple:ios.bzl", "ios_static_framework")
+load("@rules_java//java:defs.bzl", "java_library", "java_test")
+load("//tensorflow/java:build_defs.bzl", "JAVACOPTS")
+load("//tensorflow/lite/delegates/flex:build_def.bzl", "tflite_flex_jni_library")
+
+# Following targets are using for testing selective-built flex delegate
+# in Java. Please don't use them for other purposes.
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:LICENSE"],
+    licenses = ["notice"],
+)
+
+tflite_flex_jni_library(
+    name = "test",
+    testonly = 1,
+    additional_deps = ["//tensorflow/lite/python/testdata:double_op_and_kernels"],
+    models = [
+        "//tensorflow/lite:testdata/multi_add_flex.bin",
+        "//tensorflow/lite:testdata/double_flex.bin",
+    ],
+)
+
+java_library(
+    name = "test_tensorflowlitelib_flex",
+    testonly = 1,
+    srcs = ["//tensorflow/lite/delegates/flex/java/src/main/java/org/tensorflow/lite/flex:flex_delegate"],
+    javacopts = JAVACOPTS,
+    visibility = ["//visibility:private"],
+    deps = [
+        ":libtensorflowlite_flex_jni.so",
+        "//tensorflow/lite/java:tensorflowlitelib",
+        "@org_checkerframework_qual",
+    ],
+)
+
+java_test(
+    name = "SelectiveBuiltInterpreterFlexTest",
+    size = "small",
+    srcs = [
+        "//tensorflow/lite/java:portable_flex_tests",
+        "//tensorflow/lite/java:portable_test_utils",
+    ],
+    data = [
+        "//tensorflow/lite:testdata/multi_add_flex.bin",
+    ],
+    javacopts = JAVACOPTS,
+    tags = [
+        "no_cuda_on_cpu_tap",  # CUDA + flex is not officially supported.
+        "no_gpu",  # GPU + flex is not officially supported.
+        "no_oss",  # Currently requires --config=monolithic, b/118895218.
+    ],
+    test_class = "org.tensorflow.lite.InterpreterFlexTest",
+    visibility = ["//visibility:private"],
+    deps = [
+        ":test_tensorflowlitelib_flex",
+        "//tensorflow/lite/java:tensorflowlitelib",
+        "@com_google_truth",
+        "@junit",
+    ],
+)
+
+java_test(
+    name = "SelectiveBuiltInterpreterFlexWithCustomOpsTest",
+    size = "small",
+    srcs = [
+        "//tensorflow/lite/java:portable_flex_with_custom_ops_tests",
+        "//tensorflow/lite/java:portable_test_utils",
+    ],
+    data = [
+        "//tensorflow/lite:testdata/double_flex.bin",
+    ],
+    javacopts = JAVACOPTS,
+    tags = [
+        "no_cuda_on_cpu_tap",  # CUDA + flex is not officially supported.
+        "no_gpu",  # GPU + flex is not officially supported.
+        "no_oss",  # Currently requires --config=monolithic, b/118895218.
+    ],
+    test_class = "org.tensorflow.lite.InterpreterFlexWithCustomOpsTest",
+    visibility = ["//visibility:private"],
+    deps = [
+        ":test_tensorflowlitelib_flex",
+        "//tensorflow/lite/java:tensorflowlitelib",
+        "@com_google_truth",
+        "@junit",
+    ],
+)
+
+# For build test only.
+ios_static_framework(
+    name = "TestTensorFlowLiteSelectTfOps_framework",
+    testonly = 1,
+    avoid_deps = ["//tensorflow/lite/core/c:common"],
+    bundle_name = "TestTensorFlowLiteSelectTfOps",
+    minimum_os_version = "15.0",
+    deps = [
+        ":test_flex_delegate",
+    ],
+)
+
+build_test(
+    name = "framework_build_test",
+    # build_test targets are not meant to be run with sanitizers, and can also
+    # cause problems with coverage testing.
+    tags = [
+        "noasan",
+        "nomsan",
+        "notsan",
+        "nozapfhahn",
+    ],
+    targets = [
+        ":TestTensorFlowLiteSelectTfOps_framework",
+    ],
+)

tensorflow/lite/delegates/flex/test_util.cc

@@ -0,0 +1,216 @@
+/* Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/lite/delegates/flex/test_util.h"
+
+#include "absl/memory/memory.h"
+#include "flatbuffers/flexbuffers.h"  // from @flatbuffers
+#include "tensorflow/lite/string_type.h"
+
+namespace tflite {
+namespace flex {
+namespace testing {
+
+bool FlexModelTest::Invoke() { return interpreter_->Invoke() == kTfLiteOk; }
+
+void FlexModelTest::SetStringValues(int tensor_index,
+                                    const std::vector<string>& values) {
+  DynamicBuffer dynamic_buffer;
+  for (const string& s : values) {
+    dynamic_buffer.AddString(s.data(), s.size());
+  }
+  dynamic_buffer.WriteToTensor(interpreter_->tensor(tensor_index),
+                               /*new_shape=*/nullptr);
+}
+
+std::vector<string> FlexModelTest::GetStringValues(int tensor_index) const {
+  std::vector<string> result;
+
+  TfLiteTensor* tensor = interpreter_->tensor(tensor_index);
+  auto num_strings = GetStringCount(tensor);
+  for (size_t i = 0; i < num_strings; ++i) {
+    auto ref = GetString(tensor, i);
+    result.push_back(string(ref.str, ref.len));
+  }
+
+  return result;
+}
+
+void FlexModelTest::SetShape(int tensor_index, const std::vector<int>& values) {
+  ASSERT_EQ(interpreter_->ResizeInputTensor(tensor_index, values), kTfLiteOk);
+  ASSERT_EQ(interpreter_->AllocateTensors(), kTfLiteOk);
+}
+
+std::vector<int> FlexModelTest::GetShape(int tensor_index) {
+  std::vector<int> result;
+  auto* dims = interpreter_->tensor(tensor_index)->dims;
+  result.reserve(dims->size);
+  for (int i = 0; i < dims->size; ++i) {
+    result.push_back(dims->data[i]);
+  }
+  return result;
+}
+
+TfLiteType FlexModelTest::GetType(int tensor_index) {
+  return interpreter_->tensor(tensor_index)->type;
+}
+
+bool FlexModelTest::IsDynamicTensor(int tensor_index) {
+  return interpreter_->tensor(tensor_index)->allocation_type == kTfLiteDynamic;
+}
+
+void FlexModelTest::AddTensors(int num_tensors, const std::vector<int>& inputs,
+                               const std::vector<int>& outputs, TfLiteType type,
+                               const std::vector<int>& dims) {
+  interpreter_->AddTensors(num_tensors);
+  for (int i = 0; i < num_tensors; ++i) {
+    TfLiteQuantizationParams quant;
+    CHECK_EQ(interpreter_->SetTensorParametersReadWrite(i, type,
+                                                        /*name=*/"",
+                                                        /*dims=*/dims, quant),
+             kTfLiteOk);
+  }
+
+  CHECK_EQ(interpreter_->SetInputs(inputs), kTfLiteOk);
+  CHECK_EQ(interpreter_->SetOutputs(outputs), kTfLiteOk);
+}
+
+void FlexModelTest::SetConstTensor(int tensor_index,
+                                   const std::vector<int>& values,
+                                   TfLiteType type, const char* buffer,
+                                   size_t bytes) {
+  TfLiteQuantizationParams quant;
+  CHECK_EQ(interpreter_->SetTensorParametersReadOnly(tensor_index, type,
+                                                     /*name=*/"",
+                                                     /*dims=*/values, quant,
+                                                     buffer, bytes),
+           kTfLiteOk);
+}
+
+void FlexModelTest::AddTfLiteMulOp(const std::vector<int>& inputs,
+                                   const std::vector<int>& outputs) {
+  ++next_op_index_;
+
+  static TfLiteRegistration reg = {nullptr, nullptr, nullptr, nullptr};
+  reg.builtin_code = BuiltinOperator_MUL;
+  reg.prepare = [](TfLiteContext* context, TfLiteNode* node) {
+    auto* i0 = &context->tensors[node->inputs->data[0]];
+    auto* o = &context->tensors[node->outputs->data[0]];
+    return context->ResizeTensor(context, o, TfLiteIntArrayCopy(i0->dims));
+  };
+  reg.invoke = [](TfLiteContext* context, TfLiteNode* node) {
+    auto* i0 = &context->tensors[node->inputs->data[0]];
+    auto* i1 = &context->tensors[node->inputs->data[1]];
+    auto* o = &context->tensors[node->outputs->data[0]];
+    for (int i = 0; i < o->bytes / sizeof(float); ++i) {
+      o->data.f[i] = i0->data.f[i] * i1->data.f[i];
+    }
+    return kTfLiteOk;
+  };
+
+  CHECK_EQ(interpreter_->AddNodeWithParameters(inputs, outputs, nullptr, 0,
+                                               nullptr, &reg),
+           kTfLiteOk);
+}
+
+void FlexModelTest::AddTfOp(TfOpType op, const std::vector<int>& inputs,
+                            const std::vector<int>& outputs) {
+  tf_ops_.push_back(next_op_index_);
+  ++next_op_index_;
+
+  auto attr = [](const string& key, const string& value) {
+    return " attr{ key: '" + key + "' value {" + value + "}}";
+  };
+
+  string type_attribute;
+  switch (interpreter_->tensor(inputs[0])->type) {
+    case kTfLiteInt32:
+      type_attribute = attr("T", "type: DT_INT32");
+      break;
+    case kTfLiteFloat32:
+      type_attribute = attr("T", "type: DT_FLOAT");
+      break;
+    case kTfLiteString:
+      type_attribute = attr("T", "type: DT_STRING");
+      break;
+    case kTfLiteBool:
+      type_attribute = attr("T", "type: DT_BOOL");
+      break;
+    default:
+      // TODO(b/113613439): Use nodedef string utilities to properly handle all
+      // types.
+      LOG(FATAL) << "Type not supported";
+      break;
+  }
+
+  if (op == kUnpack) {
+    string attributes =
+        type_attribute + attr("num", "i: 2") + attr("axis", "i: 0");
+    AddTfOp("FlexUnpack", "Unpack", attributes, inputs, outputs);
+  } else if (op == kIdentity) {
+    string attributes = type_attribute;
+    AddTfOp("FlexIdentity", "Identity", attributes, inputs, outputs);
+  } else if (op == kAdd) {
+    string attributes = type_attribute;
+    AddTfOp("FlexAdd", "Add", attributes, inputs, outputs);
+  } else if (op == kMul) {
+    string attributes = type_attribute;
+    AddTfOp("FlexMul", "Mul", attributes, inputs, outputs);
+  } else if (op == kRfft) {
+    AddTfOp("FlexRFFT", "RFFT", "", inputs, outputs);
+  } else if (op == kImag) {
+    AddTfOp("FlexImag", "Imag", "", inputs, outputs);
+  } else if (op == kLoopCond) {
+    string attributes = type_attribute;
+    AddTfOp("FlexLoopCond", "LoopCond", attributes, inputs, outputs);
+  } else if (op == kNonExistent) {
+    AddTfOp("NonExistentOp", "NonExistentOp", "", inputs, outputs);
+  } else if (op == kIncompatibleNodeDef) {
+    // "Cast" op is created without attributes - making it incompatible.
+    AddTfOp("FlexCast", "Cast", "", inputs, outputs);
+  }
+}
+
+void FlexModelTest::AddTfOp(const char* tflite_name, const string& tf_name,
+                            const string& nodedef_str,
+                            const std::vector<int>& inputs,
+                            const std::vector<int>& outputs) {
+  static TfLiteRegistration reg = {nullptr, nullptr, nullptr, nullptr};
+  reg.builtin_code = BuiltinOperator_CUSTOM;
+  reg.custom_name = tflite_name;
+
+  tensorflow::NodeDef nodedef;
+  CHECK(tensorflow::protobuf::TextFormat::ParseFromString(
+      nodedef_str + " op: '" + tf_name + "'", &nodedef));
+  string serialized_nodedef;
+  CHECK(nodedef.SerializeToString(&serialized_nodedef));
+  flexbuffers::Builder fbb;
+  fbb.Vector([&]() {
+    fbb.String(nodedef.op());
+    fbb.String(serialized_nodedef);
+  });
+  fbb.Finish();
+
+  flexbuffers_.push_back(fbb.GetBuffer());
+  auto& buffer = flexbuffers_.back();
+  CHECK_EQ(interpreter_->AddNodeWithParameters(
+               inputs, outputs, reinterpret_cast<const char*>(buffer.data()),
+               buffer.size(), nullptr, &reg),
+           kTfLiteOk);
+}
+
+}  // namespace testing
+}  // namespace flex
+}  // namespace tflite

tensorflow/lite/delegates/flex/test_util.h

@@ -0,0 +1,135 @@
+/* Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_LITE_DELEGATES_FLEX_TEST_UTIL_H_
+#define TENSORFLOW_LITE_DELEGATES_FLEX_TEST_UTIL_H_
+
+#include "tensorflow/c/c_api_internal.h"
+#include "tensorflow/lite/kernels/test_util.h"
+
+namespace tflite {
+namespace flex {
+namespace testing {
+
+enum TfOpType {
+  kUnpack,
+  kIdentity,
+  kAdd,
+  kMul,
+  kRfft,
+  kImag,
+  kLoopCond,
+  // Represents an op that does not exist in TensorFlow.
+  kNonExistent,
+  // Represents an valid TensorFlow op where the NodeDef is incompatible.
+  kIncompatibleNodeDef,
+};
+
+// This class creates models with TF and TFLite ops. In order to use this class
+// to test the Flex delegate, implement a function that calls
+// interpreter->ModifyGraphWithDelegate.
+class FlexModelTest : public ::testing::Test {
+ public:
+  FlexModelTest() {}
+  ~FlexModelTest() override {}
+
+  bool Invoke();
+
+  // Sets the (typed) tensor's values at the given index.
+  template <typename T>
+  void SetTypedValues(int tensor_index, const std::vector<T>& values) {
+    memcpy(interpreter_->typed_tensor<T>(tensor_index), values.data(),
+           values.size() * sizeof(T));
+  }
+
+  // Returns the (typed) tensor's values at the given index.
+  template <typename T>
+  std::vector<T> GetTypedValues(int tensor_index) {
+    const TfLiteTensor* t = interpreter_->tensor(tensor_index);
+    const T* tdata = interpreter_->typed_tensor<T>(tensor_index);
+    return std::vector<T>(tdata, tdata + t->bytes / sizeof(T));
+  }
+
+  // Sets the tensor's values at the given index.
+  void SetValues(int tensor_index, const std::vector<float>& values) {
+    SetTypedValues<float>(tensor_index, values);
+  }
+  void SetStringValues(int tensor_index, const std::vector<string>& values);
+
+  // Returns the tensor's values at the given index.
+  std::vector<float> GetValues(int tensor_index) {
+    return GetTypedValues<float>(tensor_index);
+  }
+  std::vector<string> GetStringValues(int tensor_index) const;
+
+  // Sets the tensor's shape at the given index.
+  void SetShape(int tensor_index, const std::vector<int>& values);
+
+  // Returns the tensor's shape at the given index.
+  std::vector<int> GetShape(int tensor_index);
+
+  // Returns the tensor's type at the given index.
+  TfLiteType GetType(int tensor_index);
+
+  // Returns if the tensor at the given index is dynamic.
+  bool IsDynamicTensor(int tensor_index);
+
+  const TestErrorReporter& error_reporter() const { return error_reporter_; }
+
+  // Adds `num_tensor` tensors to the model. `inputs` contains the indices of
+  // the input tensors and `outputs` contains the indices of the output
+  // tensors. All tensors are set to have `type` and `dims`.
+  void AddTensors(int num_tensors, const std::vector<int>& inputs,
+                  const std::vector<int>& outputs, TfLiteType type,
+                  const std::vector<int>& dims);
+
+  // Set a constant tensor of the given shape, type and buffer at the given
+  // index.
+  void SetConstTensor(int tensor_index, const std::vector<int>& values,
+                      TfLiteType type, const char* buffer, size_t bytes);
+
+  // Adds a TFLite Mul op. `inputs` contains the indices of the input tensors
+  // and `outputs` contains the indices of the output tensors.
+  void AddTfLiteMulOp(const std::vector<int>& inputs,
+                      const std::vector<int>& outputs);
+
+  // Adds a TensorFlow op. `inputs` contains the indices of the
+  // input tensors and `outputs` contains the indices of the output tensors.
+  // This function is limited to the set of ops defined in TfOpType.
+  void AddTfOp(TfOpType op, const std::vector<int>& inputs,
+               const std::vector<int>& outputs);
+
+ protected:
+  std::unique_ptr<Interpreter> interpreter_;
+  TestErrorReporter error_reporter_;
+  std::vector<int> tf_ops_;
+
+ private:
+  // Helper method to add a TensorFlow op. tflite_names needs to start with
+  // "Flex" in order to work with the Flex delegate.
+  void AddTfOp(const char* tflite_name, const string& tf_name,
+               const string& nodedef_str, const std::vector<int>& inputs,
+               const std::vector<int>& outputs);
+
+  std::vector<std::vector<uint8_t>> flexbuffers_;
+
+  int next_op_index_ = 0;
+};
+
+}  // namespace testing
+}  // namespace flex
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_DELEGATES_FLEX_TEST_UTIL_H_

tensorflow/lite/delegates/flex/tflite_subgraph_execute.cc

@@ -0,0 +1,413 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <cstddef>
+#include <cstdlib>
+#include <cstring>
+#include <string>
+#include <unordered_map>
+#include <utility>
+#include <vector>
+
+#include "tensorflow/c/tf_tensor_internal.h"
+#include "tensorflow/core/framework/common_shape_fns.h"
+#include "tensorflow/core/framework/op.h"
+#include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/framework/op_requires.h"
+#include "tensorflow/core/framework/resource_handle.h"
+#include "tensorflow/core/framework/resource_mgr.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/platform/errors.h"
+#include "tensorflow/core/platform/mutex.h"
+#include "tensorflow/core/platform/tstring.h"
+#include "tensorflow/lite/core/c/c_api_types.h"
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/core/subgraph.h"
+#include "tensorflow/lite/delegates/flex/buffer_map_util.h"
+#include "tensorflow/lite/delegates/flex/subgraph_resource.h"
+#include "tensorflow/lite/delegates/flex/util.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
+#include "tensorflow/lite/kernels/kernel_util.h"
+#include "tensorflow/lite/string_util.h"
+
+namespace tensorflow {
+
+namespace {
+constexpr int kTfLiteSubgraphResource = 0;
+}
+
+REGISTER_OP("TfLiteSubgraphExecute")
+    .Input("subgraph_key: string")
+    .Input("args: Tin")
+    .Output("output: Tout")
+    .Attr("Tin: list(type) >= 0")
+    .Attr("Tout: list(type) >= 0")
+    .SetShapeFn(shape_inference::UnknownShape);
+
+// The `TfLiteSubgraphExecute` executes a tflite subgraph with the designated
+// inputs. This op will first look up the tflite subgraph from TF resource
+// manager based on the resource name stored on the first input, and then it
+// will call that specific subgraph with the remaining arguments. The first
+// input of this op is always a scalar string, which denotes the name of the
+// subgraph resource. The remaining inputs will be fed to the subgraph as
+// inputs, so the caller needs to ensure that the remaining inputs match with
+// the subgraph's expected inputs. This is currently WIP/experimental and
+// subject to change.
+class TfLiteSubgraphExecute : public OpKernel {
+ public:
+  explicit TfLiteSubgraphExecute(OpKernelConstruction* ctx)
+      : OpKernel(ctx),
+        tfl_tensors_need_allocation_(true),
+        fill_output_to_input_map_(true),
+        first_run_(true),
+        output_tensors_can_be_shared_(true) {}
+
+  void Compute(OpKernelContext* ctx) override {
+    // Fetch the TF Lite subgraph to execute.
+    tflite::flex::TFLiteSubgraphResource* resource = nullptr;
+    OP_REQUIRES_OK(
+        ctx,
+        ctx->resource_manager()->Lookup<tflite::flex::TFLiteSubgraphResource>(
+            "flex", ctx->input(kTfLiteSubgraphResource).flat<tstring>()(0),
+            &resource));
+    tensorflow::core::ScopedUnref unref_resource(resource);
+
+    // Try to acquire a mutex lock from this resource. This is because tflite
+    // subgraph is not thread-safe and we need to guarantee exclusive access to
+    // it.
+    mutex_lock lock(resource->GetExclusiveLock());
+    tflite::Subgraph& subgraph_selected = resource->GetSubgraphResource();
+
+    OutputToInputMap(ctx, subgraph_selected);
+    OP_REQUIRES(
+        ctx, ctx->num_inputs() == subgraph_selected.inputs().size() + 1,
+        absl::InvalidArgumentError(absl::StrCat(
+            "TF Lite subgraph expects ", subgraph_selected.inputs().size(),
+            " inputs, but received ", ctx->num_inputs() - 1, ".")));
+
+    // Resize input tensors if necessary.
+    ResizeInputTensor(ctx, subgraph_selected);
+    SetCustomAllocatorsForInputTensors(ctx, subgraph_selected);
+
+    if (tfl_tensors_need_allocation_) {
+      OP_REQUIRES(ctx, subgraph_selected.AllocateTensors() == kTfLiteOk,
+                  absl::InternalError("Failed to call allocate tensors"));
+      tfl_tensors_need_allocation_ = false;
+      // TODO(b/274934361): re-enable once allocation issue is solved.
+      // output_tensors_can_be_shared_ =!subgraph_selected.HasDynamicTensors();
+      output_tensors_can_be_shared_ = false;
+    }
+
+    if (output_tensors_can_be_shared_) {
+      SetCustomAllocatorsForOutputTensors(ctx, subgraph_selected);
+      // The output tensors have been allocated by
+      // `SetCustomAllocatorsForOutputTensors`. They have also been arena
+      // allocated by `AllocateTensors`. Calling `ReleaseMemory` and
+      // `AllocateTensors` allows us to reduce peak memory usage by no longer
+      // allocating the output tensors in the arena. This only needs to be done
+      // once as subsequent calls to `AllocateTensors` will not allocate the
+      // output tensors as they have been set to
+      // `kTfLiteCustom`.
+      if (first_run_) {
+        subgraph_selected.ReleaseMemory();
+        TfLiteStatus status = subgraph_selected.AllocateTensors();
+        if (status != kTfLiteOk) {
+          CleanUpCustomOutputs();
+          ctx->CtxFailure(absl::InternalError(absl::StrCat(
+              "Failed to call allocate tensors", subgraph_selected.GetName())));
+          return;
+        }
+        first_run_ = false;
+      }
+    }
+    // Copy input tensors to subgraph.
+    SetSubgraphInput(ctx, subgraph_selected, resource->GetFlexDelegate());
+
+    TfLiteStatus status = subgraph_selected.Invoke();
+    if (status != kTfLiteOk) {
+      CleanUpCustomOutputs();
+      ctx->CtxFailure(absl::InternalError(absl::StrCat(
+          "Failed to invoke tflite subgraph", subgraph_selected.GetName())));
+      return;
+    }
+
+    // Copy tflite results.
+    CopyTFLiteSubgraphResult(ctx, subgraph_selected);
+  }
+
+ private:
+  void CleanUpCustomOutputs() {
+    for (void* ptr : custom_output_ptrs_) {
+      tensorflow::cpu_allocator()->DeallocateRaw(ptr);
+    }
+  }
+  // Identifies subgraph inputs which are also outputs and maps the output
+  // tensor id to the input number. This means that tensors which are both
+  // subgraph inputs and outputs can be handled with zero copies.
+  void OutputToInputMap(OpKernelContext* ctx,
+                        tflite::Subgraph& subgraph_selected) {
+    if (!fill_output_to_input_map_) {
+      return;
+    }
+    for (int i = 0; i < subgraph_selected.inputs().size(); ++i) {
+      int input_idx = subgraph_selected.inputs()[i];
+      for (int j = 0; j < subgraph_selected.outputs().size(); ++j) {
+        int output_idx = subgraph_selected.outputs()[j];
+        if (input_idx == output_idx) {
+          output_to_input_map_[output_idx] = i;
+          break;
+        }
+      }
+    }
+    fill_output_to_input_map_ = false;
+  }
+
+  bool TensorCanBeShared(const TfLiteTensor* tensor) const {
+    if ((tensor->type == kTfLiteResource || tensor->type == kTfLiteVariant ||
+         tensor->type == kTfLiteString)) {
+      return false;
+    }
+    if (tensor->allocation_type != kTfLiteArenaRw &&
+        tensor->allocation_type != kTfLiteArenaRwPersistent &&
+        tensor->allocation_type != kTfLiteCustom) {
+      return false;
+    }
+    return true;
+  }
+
+  // Sets custom allocators for all output tensors which are not Resources,
+  // Variants or Strings. This means that these tensors can share the same
+  // memory as the TF tensors, reducing memcpys and memory usage.
+  void SetCustomAllocatorsForOutputTensors(
+      OpKernelContext* ctx, tflite::Subgraph& subgraph_selected) {
+    custom_output_ptrs_.clear();
+    for (int i = 0; i < subgraph_selected.outputs().size(); ++i) {
+      int tensor_idx = subgraph_selected.outputs()[i];
+      TfLiteTensor* subgraph_output = subgraph_selected.tensor(tensor_idx);
+      if (!TensorCanBeShared(subgraph_output)) {
+        continue;
+      }
+      if (output_to_input_map_.find(tensor_idx) != output_to_input_map_.end()) {
+        continue;
+      }
+      void* ptr = tensorflow::cpu_allocator()->AllocateRaw(
+          EIGEN_MAX_ALIGN_BYTES, subgraph_output->bytes);
+      custom_output_ptrs_.push_back(ptr);
+      TfLiteCustomAllocation allocation{ptr, subgraph_output->bytes};
+      OP_REQUIRES(
+          ctx,
+          subgraph_selected.SetCustomAllocationForTensor(
+              tensor_idx, allocation,
+              // Using kTfLiteCustomAllocationFlagsSkipAlignCheck is
+              // safe as the pointer comes from TensorFlow.
+              // TODO(b/257964109): Remove this flag when fixed.
+              kTfLiteCustomAllocationFlagsSkipAlignCheck) == kTfLiteOk,
+          absl::InternalError(absl::StrCat(
+              "Failed to set custom allocation for output tensor %d, name: %s",
+              tensor_idx, subgraph_output->name)));
+    }
+  }
+
+  // Sets custom allocators for all input tensors which are not Resources,
+  // Variants or Strings. This means that these tensors can share the same
+  // memory as the TF tensors, reducing memcpys and memory usage.
+  void SetCustomAllocatorsForInputTensors(OpKernelContext* ctx,
+                                          tflite::Subgraph& subgraph_selected) {
+    for (int i = 0; i < subgraph_selected.inputs().size(); ++i) {
+      int tensor_idx = subgraph_selected.inputs()[i];
+      TfLiteTensor* subgraph_input = subgraph_selected.tensor(tensor_idx);
+      if (!TensorCanBeShared(subgraph_input)) {
+        continue;
+      }
+      const Tensor& tf_tensor = ctx->input(i + 1);
+      TfLiteCustomAllocation allocation{tf_tensor.data(),
+                                        tf_tensor.AllocatedBytes()};
+      OP_REQUIRES(ctx,
+                  subgraph_selected.SetCustomAllocationForTensor(
+                      tensor_idx, allocation,
+                      // Using kTfLiteCustomAllocationFlagsSkipAlignCheck is
+                      // safe as the pointer comes from TensorFlow.
+                      // TODO(b/257964109): Remove this flag when fixed.
+                      kTfLiteCustomAllocationFlagsSkipAlignCheck) == kTfLiteOk,
+                  absl::InternalError(absl::StrCat(
+                      "Failed to set custom allocation for input tensor %d",
+                      tensor_idx)));
+    }
+  }
+
+  void ResizeInputTensor(OpKernelContext* ctx,
+                         tflite::Subgraph& subgraph_selected) {
+    for (int i = 0; i < subgraph_selected.inputs().size(); ++i) {
+      // Shift index by 1 since the first input is always the resource name.
+      const Tensor& tf_tensor = ctx->input(i + 1);
+      TfLiteTensor* subgraph_input =
+          subgraph_selected.tensor(subgraph_selected.inputs()[i]);
+
+      // Always resize for unranked tensors.
+      bool need_resize = (subgraph_input->dims->size == 0);
+      if (!need_resize) {
+        for (int dim = 0; dim < tf_tensor.shape().dims(); dim++) {
+          if (tf_tensor.shape().dim_size(dim) !=
+              subgraph_input->dims->data[dim]) {
+            need_resize = true;
+            break;
+          }
+        }
+      }
+      if (need_resize) {
+        std::vector<int> new_shape;
+        for (auto dim : tf_tensor.shape().dim_sizes()) {
+          new_shape.push_back(dim);
+        }
+        tfl_tensors_need_allocation_ = true;
+        OP_REQUIRES(ctx,
+                    subgraph_selected.ResizeInputTensor(
+                        subgraph_selected.inputs()[i], new_shape) == kTfLiteOk,
+                    absl::InternalError("Failed to resize tflite tensor"));
+      }
+    }
+  }
+
+  void SetSubgraphInput(OpKernelContext* ctx,
+                        tflite::Subgraph& subgraph_selected,
+                        TfLiteDelegate* flex_delegate) const {
+    auto InitializeVariantOrResource = [flex_delegate](
+                                           const Tensor& tf_tensor,
+                                           TfLiteTensor* subgraph_input) {
+      // The code here initializes the TfLiteTensor which points the data field
+      // to the original TF resource or variant tensor. This requires the TF
+      // tensor's lifetime must extend beyond the execution of callee subgraph.
+      // TODO(b/179094265): This is an experimental implementation, subject to
+      // change. This can be re-implemented with life cycle management
+      // mechanism like reference counting.
+      const size_t required_bytes =
+          tflite::flex::kTensorflowResourceTensorBytes;
+      const tensorflow::Tensor** tf_tensor_ptr =
+          reinterpret_cast<const tensorflow::Tensor**>(malloc(required_bytes));
+      *tf_tensor_ptr = &tf_tensor;
+
+      TfLiteTensorDataFree(subgraph_input);
+      subgraph_input->data.raw = reinterpret_cast<char*>(tf_tensor_ptr);
+      subgraph_input->bytes = required_bytes;
+      subgraph_input->data_is_stale = true;
+      subgraph_input->delegate = flex_delegate;
+    };
+
+    for (int i = 0; i < subgraph_selected.inputs().size(); ++i) {
+      const Tensor& tf_tensor = ctx->input(i + 1);
+      TfLiteTensor* subgraph_input =
+          subgraph_selected.tensor(subgraph_selected.inputs()[i]);
+
+      if (subgraph_input->type == kTfLiteString) {
+        OP_REQUIRES(
+            ctx, tf_tensor.dtype() == tensorflow::DT_STRING,
+            absl::InvalidArgumentError("Tensor doesn't have string type"));
+        tflite::DynamicBuffer dynamic_buffer;
+        auto tf_data = tf_tensor.flat<tensorflow::tstring>();
+        for (int i = 0; i < tf_tensor.NumElements(); ++i) {
+          dynamic_buffer.AddString(tf_data(i).data(), tf_data(i).size());
+        }
+
+        dynamic_buffer.WriteToTensor(subgraph_input, /*new_shape=*/nullptr);
+      } else if (subgraph_input->type == kTfLiteResource) {
+        // Here we will try to parse the input tensor handle to see if it
+        // contains a valid TF lite resource ID. If not, then we know that the
+        // input is a TF resource tensor.
+        tensorflow::ResourceHandle handle =
+            tf_tensor.flat<tensorflow::ResourceHandle>()(0);
+        if (!tflite::flex::GetTfLiteResourceTensorFromResourceHandle(
+                handle, subgraph_input)) {
+          InitializeVariantOrResource(tf_tensor, subgraph_input);
+        }
+      } else if (subgraph_input->type == kTfLiteVariant) {
+        InitializeVariantOrResource(tf_tensor, subgraph_input);
+      } else if (!TensorCanBeShared(subgraph_input)) {
+        absl::string_view tensor_data = tf_tensor.tensor_data();
+        OP_REQUIRES(ctx, subgraph_input->bytes == tensor_data.size(),
+                    absl::InternalError("tensor size doesn't match"));
+        // TODO(b/181352924): This could incur some overhead in memory copy.
+        // Optimize this away in the future.
+        std::memcpy(subgraph_input->data.raw, tensor_data.data(),
+                    tensor_data.size());
+      }
+    }
+  }
+
+  void CopyTFLiteSubgraphResult(OpKernelContext* ctx,
+                                tflite::Subgraph& subgraph_selected) const {
+    for (int i = 0; i < subgraph_selected.outputs().size(); ++i) {
+      OP_REQUIRES(
+          ctx,
+          subgraph_selected.EnsureTensorDataIsReadable(
+              subgraph_selected.outputs()[i]) == kTfLiteOk,
+          absl::InternalError("TF lite subgraph output is not readable"));
+      // Create an output tensor.
+      TfLiteTensor* subgraph_output =
+          subgraph_selected.tensor(subgraph_selected.outputs()[i]);
+
+      // If the output is also an input, create the output tensor using the same
+      // buffer as the input tensor so that no copy is needed and to also save
+      // memory.
+      int output_idx = subgraph_selected.outputs()[i];
+      auto it = output_to_input_map_.find(output_idx);
+      if (it != output_to_input_map_.end()) {
+        const Tensor& tf_tensor = ctx->input(it->second + 1);
+        Tensor tensor(tf_tensor);
+        ctx->set_output(i, std::move(tensor));
+        continue;
+      }
+      // Take ownership of the TFLite output ptr where possible, otherwise copy
+      // the output. This is because the callee subgraph might be invoked
+      // repeatedly for each item in the dataset, and the result TfLiteTensor's
+      // data should be immediately copied into tensorflow::Tensor.
+      Tensor tensor;
+      if (output_tensors_can_be_shared_ && TensorCanBeShared(subgraph_output)) {
+        tflite::flex::TfLiteTensorBuffer* buf =
+            new tflite::flex::TfLiteTensorBuffer(subgraph_output, true);
+        buf->TakeOwnershipOfBuffer();
+        tensorflow::TensorShape shape;
+        int num_dims = subgraph_output->dims->size;
+        for (int i = 0; i < num_dims; ++i) {
+          OP_REQUIRES_OK(
+              ctx, shape.AddDimWithStatus(subgraph_output->dims->data[i]));
+        }
+        tensor = tensorflow::TensorCApi::MakeTensor(
+            tflite::flex::GetTensorFlowDataType(subgraph_output->type), shape,
+            buf);
+        buf->Unref();
+      } else {
+        OP_REQUIRES_OK(
+            ctx, tflite::flex::SetTfTensorFromTfLite(subgraph_output, &tensor,
+                                                     /*allow_reusing=*/false));
+      }
+      ctx->set_output(i, std::move(tensor));
+    }
+  }
+
+  // Tells if the target subgraph needs to invoko AllocateTensors().
+  bool tfl_tensors_need_allocation_;
+  bool fill_output_to_input_map_;
+  // NOLINTNEXTLINE - absl::flat_hash_map increases binary size by 106kB.
+  std::unordered_map<int, int> output_to_input_map_;
+  bool first_run_;
+  bool output_tensors_can_be_shared_;
+  std::vector<void*> custom_output_ptrs_;
+  std::vector<TfLiteAllocationType> original_allocation_type_;
+};
+
+REGISTER_KERNEL_BUILDER(Name("TfLiteSubgraphExecute").Device(DEVICE_CPU),
+                        TfLiteSubgraphExecute);
+
+}  // namespace tensorflow

tensorflow/lite/delegates/flex/util.cc

@@ -0,0 +1,310 @@
+/* Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include "tensorflow/lite/delegates/flex/util.h"
+
+#include <cstdint>
+#include <cstring>
+#include <limits>
+#include <string>
+#include <vector>
+
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/numbers.h"
+#include "absl/strings/str_format.h"
+#include "absl/strings/str_split.h"
+#include "tensorflow/c/tf_datatype.h"
+#include "tensorflow/core/framework/resource_handle.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/framework/types.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/platform/status.h"
+#include "tensorflow/core/platform/tstring.h"
+#include "tensorflow/core/protobuf/error_codes.pb.h"
+#include "tensorflow/lite/c/common.h"
+#include "tensorflow/lite/core/c/c_api_types.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
+#include "tensorflow/lite/string_util.h"
+#include "tensorflow/lite/util.h"
+
+namespace tflite {
+namespace flex {
+
+static constexpr char kResourceVariablePrefix[] = "tflite_resource_variable";
+
+TfLiteStatus ConvertStatus(TfLiteContext* context, const absl::Status& status) {
+  if (!status.ok()) {
+    TF_LITE_KERNEL_LOG(context, "%s", absl::StatusMessageAsCStr(status));
+    return kTfLiteError;
+  }
+  return kTfLiteOk;
+}
+
+TfLiteStatus CopyShapeAndType(TfLiteContext* context,
+                              const tensorflow::Tensor& src,
+                              TfLiteTensor* tensor) {
+  tensor->type = GetTensorFlowLiteType(static_cast<TF_DataType>(src.dtype()));
+  if (tensor->type == kTfLiteNoType) {
+    TF_LITE_KERNEL_LOG(context,
+                       "TF Lite does not support TensorFlow data type: %s",
+                       DataTypeString(src.dtype()).c_str());
+    return kTfLiteError;
+  }
+
+  int num_dims = src.dims();
+  TfLiteIntArray* shape = TfLiteIntArrayCreate(num_dims);
+  for (int j = 0; j < num_dims; ++j) {
+    // We need to cast from TensorFlow's int64 to TF Lite's int32. Let's
+    // make sure there's no overflow.
+    if (src.dim_size(j) >= std::numeric_limits<int>::max()) {
+      TF_LITE_KERNEL_LOG(context,
+                         "Dimension value in TensorFlow shape is larger than "
+                         "supported by TF Lite");
+      TfLiteIntArrayFree(shape);
+      return kTfLiteError;
+    }
+    shape->data[j] = static_cast<int>(src.dim_size(j));
+  }
+  return context->ResizeTensor(context, tensor, shape);
+}
+
+TF_DataType GetTensorFlowDataType(TfLiteType type) {
+  switch (type) {
+    case kTfLiteNoType:
+      return TF_FLOAT;
+    case kTfLiteFloat32:
+      return TF_FLOAT;
+    case kTfLiteFloat16:
+      return TF_HALF;
+    case kTfLiteBFloat16:
+      return TF_BFLOAT16;
+    case kTfLiteFloat64:
+      return TF_DOUBLE;
+    case kTfLiteInt16:
+      return TF_INT16;
+    case kTfLiteUInt16:
+      return TF_UINT16;
+    case kTfLiteInt32:
+      return TF_INT32;
+    case kTfLiteUInt32:
+      return TF_UINT32;
+    case kTfLiteInt2:
+      // TODO(b/246806634): Tensorflow DT_INT2 type doesn't exist yet
+      return TF_INT8;
+    case kTfLiteUInt4:
+      // TODO(b/246806634): Tensorflow DT_UINT4 type doesn't exist yet
+      return TF_UINT8;
+    case kTfLiteInt4:
+      // TODO(b/246806634): Tensorflow DT_INT4 type doesn't exist yet
+      return TF_INT8;
+    case kTfLiteUInt8:
+      return TF_UINT8;
+    case kTfLiteInt8:
+      return TF_INT8;
+    case kTfLiteFloat8E4M3FN:
+      return TF_FLOAT8_E4M3FN;
+    case kTfLiteFloat8E5M2:
+      return TF_FLOAT8_E5M2;
+    case kTfLiteInt64:
+      return TF_INT64;
+    case kTfLiteUInt64:
+      return TF_UINT64;
+    case kTfLiteComplex64:
+      return TF_COMPLEX64;
+    case kTfLiteComplex128:
+      return TF_COMPLEX128;
+    case kTfLiteString:
+      return TF_STRING;
+    case kTfLiteBool:
+      return TF_BOOL;
+    case kTfLiteResource:
+      return TF_RESOURCE;
+    case kTfLiteVariant:
+      return TF_VARIANT;
+  }
+}
+
+TfLiteType GetTensorFlowLiteType(TF_DataType type) {
+  switch (type) {
+    case TF_FLOAT:
+      return kTfLiteFloat32;
+    case TF_HALF:
+      return kTfLiteFloat16;
+    case TF_BFLOAT16:
+      return kTfLiteBFloat16;
+    case TF_DOUBLE:
+      return kTfLiteFloat64;
+    case TF_INT16:
+      return kTfLiteInt16;
+    case TF_UINT16:
+      return kTfLiteUInt16;
+    case TF_INT32:
+      return kTfLiteInt32;
+    case TF_UINT32:
+      return kTfLiteUInt32;
+    case TF_UINT8:
+      return kTfLiteUInt8;
+    case TF_INT8:
+      return kTfLiteInt8;
+    case TF_INT64:
+      return kTfLiteInt64;
+    case TF_UINT64:
+      return kTfLiteUInt64;
+    case TF_COMPLEX64:
+      return kTfLiteComplex64;
+    case TF_COMPLEX128:
+      return kTfLiteComplex128;
+    case TF_STRING:
+      return kTfLiteString;
+    case TF_BOOL:
+      return kTfLiteBool;
+    case TF_RESOURCE:
+      return kTfLiteResource;
+    case TF_FLOAT8_E4M3FN:
+      return kTfLiteFloat8E4M3FN;
+    case TF_FLOAT8_E5M2:
+      return kTfLiteFloat8E5M2;
+    case TF_VARIANT:
+      return kTfLiteVariant;
+    default:
+      return kTfLiteNoType;
+  }
+}
+
+// Returns the TF data type name to be stored in the FunctionDef.
+const char* TfLiteTypeToTfTypeName(TfLiteType type) {
+  switch (type) {
+    case kTfLiteNoType:
+      return "invalid";
+    case kTfLiteFloat32:
+      return "float";
+    case kTfLiteInt16:
+      return "int16";
+    case kTfLiteUInt16:
+      return "uint16";
+    case kTfLiteInt32:
+      return "int32";
+    case kTfLiteUInt32:
+      return "uint32";
+    case kTfLiteInt2:
+      return "int2";
+    case kTfLiteUInt4:
+      return "uint4";
+    case kTfLiteInt4:
+      return "int4";
+    case kTfLiteUInt8:
+      return "uint8";
+    case kTfLiteInt8:
+      return "int8";
+    case kTfLiteFloat8E4M3FN:
+      return "float8_e4m3fn";
+    case kTfLiteFloat8E5M2:
+      return "float8_e5m2";
+    case kTfLiteInt64:
+      return "int64";
+    case kTfLiteUInt64:
+      return "uint64";
+    case kTfLiteBool:
+      return "bool";
+    case kTfLiteComplex64:
+      return "complex64";
+    case kTfLiteComplex128:
+      return "complex128";
+    case kTfLiteString:
+      return "string";
+    case kTfLiteFloat16:
+      return "float16";
+    case kTfLiteBFloat16:
+      return "bfloat16";
+    case kTfLiteFloat64:
+      return "float64";
+    case kTfLiteResource:
+      return "resource";
+    case kTfLiteVariant:
+      return "variant";
+  }
+  return "invalid";
+}
+
+std::string TfLiteResourceIdentifier(const TfLiteTensor* tensor) {
+  // TODO(b/199782192): Create a util function to get Resource ID from a TF Lite
+  // resource tensor.
+  const int resource_id = tensor->data.i32[0];
+  return absl::StrFormat("%s:%d", kResourceVariablePrefix, resource_id);
+}
+
+bool GetTfLiteResourceTensorFromResourceHandle(
+    const tensorflow::ResourceHandle& resource_handle, TfLiteTensor* tensor) {
+  std::vector<std::string> parts = absl::StrSplit(resource_handle.name(), ':');
+  if (parts.size() != 2) {
+    return false;
+  }
+  const int kBytesRequired = sizeof(int32_t);
+  TfLiteTensorRealloc(kBytesRequired, tensor);
+  int resource_id;
+  if (parts[0] == kResourceVariablePrefix &&
+      absl::SimpleAtoi<int32_t>(parts[1], &resource_id)) {
+    // TODO(b/199782192): Create a util function to set the Resource ID of
+    // a TF Lite resource tensor.
+    GetTensorData<int32_t>(tensor)[0] = resource_id;
+    return true;
+  }
+  return false;
+}
+
+absl::StatusOr<tensorflow::Tensor> CreateTfTensorFromTfLiteTensor(
+    const TfLiteTensor* tflite_tensor) {
+  if (IsResourceOrVariant(tflite_tensor)) {
+    // Returns error if the input tflite tensor has variant or resource type.
+    return absl::Status(absl::StatusCode::kInvalidArgument,
+                        "Input tensor has resource or variant type.");
+  }
+
+  tensorflow::TensorShape shape;
+  int num_dims = tflite_tensor->dims->size;
+  for (int i = 0; i < num_dims; ++i) {
+    shape.AddDim(tflite_tensor->dims->data[i]);
+  }
+
+  tensorflow::Tensor tf_tensor(
+      tensorflow::DataType(GetTensorFlowDataType(tflite_tensor->type)), shape);
+  if (tf_tensor.dtype() == tensorflow::DataType::DT_STRING &&
+      tf_tensor.data()) {
+    tensorflow::tstring* buf =
+        static_cast<tensorflow::tstring*>(tf_tensor.data());
+    for (int i = 0; i < tflite::GetStringCount(tflite_tensor); ++buf, ++i) {
+      auto ref = GetString(tflite_tensor, i);
+      buf->assign(ref.str, ref.len);
+    }
+  } else {
+    if (tf_tensor.tensor_data().size() != tflite_tensor->bytes) {
+      return absl::Status(
+          absl::StatusCode::kInternal,
+          "TfLiteTensor's size doesn't match the TF tensor's size.");
+    }
+    if (!tflite_tensor->data.raw) {
+      return absl::Status(absl::StatusCode::kInternal,
+                          "TfLiteTensor's data field is null.");
+    }
+    std::memcpy(tf_tensor.data(), tflite_tensor->data.raw,
+                tflite_tensor->bytes);
+  }
+
+  return tf_tensor;
+}
+
+}  // namespace flex
+}  // namespace tflite

tensorflow/lite/delegates/flex/util.h

@@ -0,0 +1,82 @@
+/* Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_DELEGATES_FLEX_UTIL_H_
+#define TENSORFLOW_LITE_DELEGATES_FLEX_UTIL_H_
+
+#include <string>
+
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "tensorflow/c/c_api_internal.h"
+#include "tensorflow/c/tf_datatype.h"
+#include "tensorflow/core/framework/resource_handle.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/lib/core/status.h"
+#include "tensorflow/core/platform/statusor.h"
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/util.h"
+
+namespace tflite {
+namespace flex {
+
+// Converts a tensorflow:Status into a TfLiteStatus. If the original status
+// represented an error, reports it using the given 'context'.
+TfLiteStatus ConvertStatus(TfLiteContext* context, const absl::Status& status);
+
+// Copies the given shape and type of the TensorFlow 'src' tensor into a TF Lite
+// 'tensor'. Logs an error and returns kTfLiteError if the shape or type can't
+// be converted.
+TfLiteStatus CopyShapeAndType(TfLiteContext* context,
+                              const tensorflow::Tensor& src,
+                              TfLiteTensor* tensor);
+
+// Returns the TF C API Data type that corresponds to the given TfLiteType.
+TF_DataType GetTensorFlowDataType(TfLiteType type);
+
+// Returns the TfLiteType that corresponds to the given TF C API Data type.
+TfLiteType GetTensorFlowLiteType(TF_DataType);
+
+// Returns the TF type name that corresponds to the given TfLiteType.
+const char* TfLiteTypeToTfTypeName(TfLiteType type);
+
+// Creates a `tensorflow::Tensor` from a TfLiteTensor for non-resource and
+// non-variant type. Returns error status if the conversion fails.
+absl::StatusOr<tensorflow::Tensor> CreateTfTensorFromTfLiteTensor(
+    const TfLiteTensor* tflite_tensor);
+
+// Returns the encoded string name for a TF Lite resource variable tensor.
+// This function will return a string in the format:
+// tflite_resource_variable:resource_id.
+std::string TfLiteResourceIdentifier(const TfLiteTensor* tensor);
+
+// Parses out the resource ID from the given `resource_handle` and sets it
+// to the corresponding TfLiteTensor. Returns true if succeed.
+bool GetTfLiteResourceTensorFromResourceHandle(
+    const tensorflow::ResourceHandle& resource_handle, TfLiteTensor* tensor);
+
+// We need a way to tell if we've stored a tensorflow::Tensor* in a resource
+// or if it's a standard kTfLiteResource tensor holding an integer. The proper
+// solution would be some way to set the TfLiteTensor::type field to something
+// unique for tensorflow::Tensor* resources. We don't want to do that, so we use
+// a hack instead: the `bytes` field of the tensor just needs to be big enough
+// to hold a pointer, but it can be larger. To disambiguate between a pointer on
+// a 32-bit machine and an int in a standard TFlite resource, we make the bytes
+// field a fixed constant big enough for a pointer on any platform.
+static constexpr int kTensorflowResourceTensorBytes = 8;
+
+}  // namespace flex
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_DELEGATES_FLEX_UTIL_H_

tensorflow/lite/delegates/flex/util_test.cc

@@ -0,0 +1,287 @@
+/* Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include "tensorflow/lite/delegates/flex/util.h"
+
+#include <cstdarg>
+#include <cstddef>
+#include <cstdint>
+#include <cstdio>
+#include <cstdlib>
+#include <cstring>
+#include <iterator>
+#include <string>
+#include <vector>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "absl/status/status.h"
+#include "tensorflow/c/tf_datatype.h"
+#include "tensorflow/core/framework/resource_handle.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/platform/errors.h"
+#include "tensorflow/core/platform/status.h"
+#include "tensorflow/core/platform/tstring.h"
+#include "tensorflow/core/protobuf/error_codes.pb.h"
+#include "tensorflow/lite/core/c/c_api_types.h"
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/string_type.h"
+#include "tensorflow/lite/string_util.h"
+#include "tensorflow/lite/util.h"
+
+namespace tflite {
+namespace flex {
+namespace {
+
+using tensorflow::DT_FLOAT;
+using tensorflow::DT_INT32;
+using tensorflow::Tensor;
+using ::testing::ElementsAre;
+
+struct TestContext : public TfLiteContext {
+  string error;
+  std::vector<int> new_size;
+};
+
+void ReportError(TfLiteContext* context, const char* format, ...) {
+  TestContext* c = static_cast<TestContext*>(context);
+  const size_t kBufferSize = 1024;
+  char temp_buffer[kBufferSize];
+
+  va_list args;
+  va_start(args, format);
+  vsnprintf(temp_buffer, kBufferSize, format, args);
+  va_end(args);
+
+  c->error = temp_buffer;
+}
+
+TfLiteStatus ResizeTensor(TfLiteContext* context, TfLiteTensor* tensor,
+                          TfLiteIntArray* new_size) {
+  TestContext* c = static_cast<TestContext*>(context);
+  c->new_size.clear();
+  for (int i = 0; i < new_size->size; ++i) {
+    c->new_size.push_back(new_size->data[i]);
+  }
+  TfLiteIntArrayFree(new_size);
+  return kTfLiteOk;
+}
+
+TEST(UtilTest, ConvertStatus) {
+  TestContext context;
+  context.ReportError = ReportError;
+
+  EXPECT_EQ(ConvertStatus(&context, absl::InternalError("Some Error")),
+            kTfLiteError);
+  EXPECT_EQ(context.error, "Some Error");
+
+  context.error.clear();
+  EXPECT_EQ(ConvertStatus(&context, absl::Status()), kTfLiteOk);
+  EXPECT_TRUE(context.error.empty());
+}
+
+TEST(UtilTest, CopyShapeAndType) {
+  TestContext context;
+  context.ReportError = ReportError;
+  context.ResizeTensor = ResizeTensor;
+
+  TfLiteTensor dst;
+
+  EXPECT_EQ(CopyShapeAndType(&context, Tensor(), &dst), kTfLiteOk);
+  EXPECT_THAT(context.new_size, ElementsAre(0));
+  EXPECT_EQ(dst.type, kTfLiteFloat32);
+
+  EXPECT_EQ(CopyShapeAndType(&context, Tensor(DT_FLOAT, {1, 2}), &dst),
+            kTfLiteOk);
+  EXPECT_THAT(context.new_size, ElementsAre(1, 2));
+  EXPECT_EQ(dst.type, kTfLiteFloat32);
+
+  EXPECT_EQ(CopyShapeAndType(&context, Tensor(DT_INT32, {1, 2}), &dst),
+            kTfLiteOk);
+  EXPECT_THAT(context.new_size, ElementsAre(1, 2));
+  EXPECT_EQ(dst.type, kTfLiteInt32);
+
+  EXPECT_EQ(CopyShapeAndType(&context, Tensor(DT_FLOAT, {1LL << 44, 2}), &dst),
+            kTfLiteError);
+  EXPECT_EQ(context.error,
+            "Dimension value in TensorFlow shape is larger than supported by "
+            "TF Lite");
+
+  EXPECT_EQ(
+      CopyShapeAndType(&context, Tensor(tensorflow::DT_HALF, {1, 2}), &dst),
+      kTfLiteOk);
+  EXPECT_THAT(context.new_size, ElementsAre(1, 2));
+  EXPECT_EQ(dst.type, kTfLiteFloat16);
+}
+
+TEST(UtilTest, TypeConversionsFromTFLite) {
+  EXPECT_EQ(TF_FLOAT, GetTensorFlowDataType(kTfLiteNoType));
+  EXPECT_EQ(TF_FLOAT, GetTensorFlowDataType(kTfLiteFloat32));
+  EXPECT_EQ(TF_HALF, GetTensorFlowDataType(kTfLiteFloat16));
+  EXPECT_EQ(TF_BFLOAT16, GetTensorFlowDataType(kTfLiteBFloat16));
+  EXPECT_EQ(TF_DOUBLE, GetTensorFlowDataType(kTfLiteFloat64));
+  EXPECT_EQ(TF_INT16, GetTensorFlowDataType(kTfLiteInt16));
+  EXPECT_EQ(TF_INT32, GetTensorFlowDataType(kTfLiteInt32));
+  EXPECT_EQ(TF_UINT8, GetTensorFlowDataType(kTfLiteUInt8));
+  EXPECT_EQ(TF_INT64, GetTensorFlowDataType(kTfLiteInt64));
+  EXPECT_EQ(TF_UINT64, GetTensorFlowDataType(kTfLiteUInt64));
+  EXPECT_EQ(TF_COMPLEX64, GetTensorFlowDataType(kTfLiteComplex64));
+  EXPECT_EQ(TF_COMPLEX128, GetTensorFlowDataType(kTfLiteComplex128));
+  EXPECT_EQ(TF_STRING, GetTensorFlowDataType(kTfLiteString));
+  EXPECT_EQ(TF_BOOL, GetTensorFlowDataType(kTfLiteBool));
+  EXPECT_EQ(TF_RESOURCE, GetTensorFlowDataType(kTfLiteResource));
+  EXPECT_EQ(TF_VARIANT, GetTensorFlowDataType(kTfLiteVariant));
+  // TODO(b/246806634): Tensorflow DT_INT4 type doesn't exist yet
+  EXPECT_EQ(TF_INT8, GetTensorFlowDataType(kTfLiteInt4));
+}
+
+TEST(UtilTest, TypeConversionsFromTensorFlow) {
+  EXPECT_EQ(kTfLiteFloat16, GetTensorFlowLiteType(TF_HALF));
+  EXPECT_EQ(kTfLiteBFloat16, GetTensorFlowLiteType(TF_BFLOAT16));
+  EXPECT_EQ(kTfLiteFloat32, GetTensorFlowLiteType(TF_FLOAT));
+  EXPECT_EQ(kTfLiteFloat64, GetTensorFlowLiteType(TF_DOUBLE));
+  EXPECT_EQ(kTfLiteInt16, GetTensorFlowLiteType(TF_INT16));
+  EXPECT_EQ(kTfLiteInt32, GetTensorFlowLiteType(TF_INT32));
+  EXPECT_EQ(kTfLiteUInt8, GetTensorFlowLiteType(TF_UINT8));
+  EXPECT_EQ(kTfLiteInt64, GetTensorFlowLiteType(TF_INT64));
+  EXPECT_EQ(kTfLiteUInt64, GetTensorFlowLiteType(TF_UINT64));
+  EXPECT_EQ(kTfLiteComplex64, GetTensorFlowLiteType(TF_COMPLEX64));
+  EXPECT_EQ(kTfLiteComplex128, GetTensorFlowLiteType(TF_COMPLEX128));
+  EXPECT_EQ(kTfLiteString, GetTensorFlowLiteType(TF_STRING));
+  EXPECT_EQ(kTfLiteBool, GetTensorFlowLiteType(TF_BOOL));
+  EXPECT_EQ(kTfLiteResource, GetTensorFlowLiteType(TF_RESOURCE));
+  EXPECT_EQ(kTfLiteVariant, GetTensorFlowLiteType(TF_VARIANT));
+}
+
+TEST(UtilTest, GetTfLiteResourceIdentifier) {
+  // Constructs a fake resource tensor.
+  TfLiteTensor tensor;
+  tensor.allocation_type = kTfLiteDynamic;
+  tensor.type = kTfLiteResource;
+  std::vector<int> dims = {1};
+  tensor.dims = ConvertVectorToTfLiteIntArray(dims);
+  tensor.data.raw = nullptr;
+  TfLiteTensorRealloc(sizeof(int32_t), &tensor);
+  tensor.delegate = nullptr;
+  tensor.data.i32[0] = 1;
+
+  EXPECT_EQ(TfLiteResourceIdentifier(&tensor), "tflite_resource_variable:1");
+  TfLiteIntArrayFree(tensor.dims);
+  TfLiteTensorDataFree(&tensor);
+}
+
+TEST(UtilTest, GetTfLiteResourceTensorFromResourceHandle) {
+  tensorflow::ResourceHandle handle;
+  handle.set_name("tflite_resource_variable:1");
+
+  TfLiteTensor tensor;
+  tensor.allocation_type = kTfLiteDynamic;
+  tensor.type = kTfLiteResource;
+  tensor.data.raw = nullptr;
+  std::vector<int> dims = {1};
+  tensor.dims = ConvertVectorToTfLiteIntArray(dims);
+  EXPECT_TRUE(GetTfLiteResourceTensorFromResourceHandle(handle, &tensor));
+  EXPECT_EQ(tensor.data.i32[0], 1);
+
+  TfLiteIntArrayFree(tensor.dims);
+  TfLiteTensorDataFree(&tensor);
+}
+
+TEST(UtilTest, CreateTfTensorFromTfLiteTensorResourceOrVariant) {
+  TfLiteTensor tensor;
+  tensor.type = kTfLiteResource;
+  EXPECT_EQ(CreateTfTensorFromTfLiteTensor(&tensor).status().code(),
+            absl::StatusCode::kInvalidArgument);
+  tensor.type = kTfLiteVariant;
+  EXPECT_EQ(CreateTfTensorFromTfLiteTensor(&tensor).status().code(),
+            absl::StatusCode::kInvalidArgument);
+}
+
+TEST(UtilTest, CreateTfTensorFromTfLiteTensorFloat) {
+  TfLiteTensor tflite_tensor;
+  tflite_tensor.type = kTfLiteFloat32;
+  tflite_tensor.allocation_type = kTfLiteDynamic;
+  tflite_tensor.sparsity = nullptr;
+  tflite_tensor.dims_signature = nullptr;
+
+  TfLiteQuantization quant;
+  quant.type = kTfLiteNoQuantization;
+  quant.params = nullptr;
+  tflite_tensor.quantization = quant;
+
+  TfLiteIntArray* dims = TfLiteIntArrayCreate(2);
+  dims->data[0] = 1;
+  dims->data[1] = 3;
+  tflite_tensor.dims = dims;
+  float data_arr[] = {1.1, 0.456, 0.322};
+  std::vector<float> data(std::begin(data_arr), std::end(data_arr));
+  size_t num_bytes = data.size() * sizeof(float);
+  tflite_tensor.data.raw = static_cast<char*>(malloc(num_bytes));
+  memcpy(tflite_tensor.data.raw, data.data(), num_bytes);
+  tflite_tensor.bytes = num_bytes;
+
+  auto tf_tensor_or = CreateTfTensorFromTfLiteTensor(&tflite_tensor);
+  EXPECT_TRUE(tf_tensor_or.ok());
+  tensorflow::Tensor tf_tensor = tf_tensor_or.value();
+  EXPECT_EQ(tf_tensor.NumElements(), 3);
+  auto* tf_data = static_cast<float*>(tf_tensor.data());
+  for (float weight : data_arr) {
+    EXPECT_EQ(*tf_data, weight);
+    tf_data++;
+  }
+
+  TfLiteTensorFree(&tflite_tensor);
+}
+
+TEST(UtilTest, CreateTfTensorFromTfLiteTensorString) {
+  TfLiteTensor tflite_tensor{};
+  tflite_tensor.type = kTfLiteString;
+  tflite_tensor.is_variable = false;
+  tflite_tensor.sparsity = nullptr;
+  tflite_tensor.data.raw = nullptr;
+  tflite_tensor.dims_signature = nullptr;
+  tflite_tensor.allocation_type = kTfLiteArenaRw;
+
+  TfLiteQuantization quant;
+  quant.type = kTfLiteNoQuantization;
+  quant.params = nullptr;
+  tflite_tensor.quantization = quant;
+
+  TfLiteIntArray* dims = TfLiteIntArrayCreate(2);
+  dims->data[0] = 1;
+  dims->data[1] = 2;
+  tflite_tensor.dims = dims;
+  std::string data_arr[] = {std::string("a_str\0ing", 9), "b_string"};
+  tflite::DynamicBuffer buf;
+  for (const auto& value : data_arr) {
+    ASSERT_EQ(buf.AddString(value.data(), value.length()), kTfLiteOk);
+  }
+  buf.WriteToTensor(&tflite_tensor, nullptr);
+
+  auto tf_tensor_or = CreateTfTensorFromTfLiteTensor(&tflite_tensor);
+  EXPECT_TRUE(tf_tensor_or.ok());
+  tensorflow::Tensor tf_tensor = tf_tensor_or.value();
+  EXPECT_EQ(tf_tensor.NumElements(), 2);
+  auto* tf_data = static_cast<tensorflow::tstring*>(tf_tensor.data());
+  for (const auto& str : data_arr) {
+    EXPECT_EQ(*tf_data, str);
+    tf_data++;
+  }
+  TfLiteTensorFree(&tflite_tensor);
+}
+
+}  // namespace
+}  // namespace flex
+}  // namespace tflite

tensorflow/lite/delegates/flex/version_script.lds

@@ -0,0 +1,6 @@
+VERS_1.0 {
+  global:
+    *AcquireFlexDelegate*;
+  local:
+    *;
+};

tensorflow/lite/delegates/gpu/BUILD

@@ -0,0 +1,388 @@
+load("@bazel_skylib//lib:selects.bzl", "selects")
+load("@build_bazel_rules_apple//apple:ios.bzl", "ios_static_framework")
+load("@build_bazel_rules_apple//apple:macos.bzl", "macos_dylib")
+load("@rules_cc//cc:cc_binary.bzl", "cc_binary")
+load("@rules_cc//cc:cc_library.bzl", "cc_library")
+load("@rules_cc//cc:cc_test.bzl", "cc_test")
+load("@rules_cc//cc:objc_library.bzl", "objc_library")
+load(
+    "//tensorflow/core/platform:build_config_root.bzl",
+    "tf_gpu_tests_tags",
+)
+load("//tensorflow/lite:build_def.bzl", "CXX17_BAZEL_ONLY_COPTS", "tflite_pagesize_linkopts")
+load("//tensorflow/lite:special_rules.bzl", "tflite_extra_gles_deps", "tflite_portable_test_suite")
+load("//tensorflow/lite/delegates/gpu:build_defs.bzl", "gpu_delegate_linkopts")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:LICENSE"],
+    default_visibility = ["//visibility:public"],
+    licenses = ["notice"],
+)
+
+exports_files([
+    "delegate.h",
+    "delegate_options.h",
+    "metal_delegate.h",
+])
+
+_DELEGATE_NO_GL_DEPS = select({
+    "//tensorflow:android": [
+        ":async_buffers",
+        "//tensorflow/lite/core/async/c:task",
+        "//tensorflow/lite/core/async/interop/c:attribute_map",
+        "//tensorflow/lite/core/async/interop/c:constants",
+        "//tensorflow/lite/delegates/gpu/gl:android_sync",
+        "//tensorflow/lite/delegates/gpu/gl:egl_environment",
+        "//tensorflow/lite/delegates/utils",
+        "//tensorflow/lite/delegates/utils:async_type_helpers",
+        "//tensorflow/lite/delegates/utils:ret_macros",
+        "//tensorflow/lite/delegates/utils:sync_fence",
+    ],
+    "//conditions:default": [],
+}) + [
+    ":android_hardware_buffer",
+    ":api",
+    ":delegate_options",
+    ":tflite_profile",
+    #"//third_party/GL:EGL_headers",
+    #"//third_party/GL:GLES3_headers",
+    # go/keep-sorted start
+    "@com_google_absl//absl/container:flat_hash_map",
+    "@com_google_absl//absl/container:flat_hash_set",
+    "@com_google_absl//absl/memory",
+    "@com_google_absl//absl/strings",
+    "@com_google_absl//absl/types:span",
+    "//tensorflow/lite/async:backend_async_kernel_interface",
+    "//tensorflow/lite/core/async/interop/c:types",
+    "//tensorflow/lite/core/c:common",
+    "//tensorflow/lite/delegates/gpu/cl:util",
+    "//tensorflow/lite/delegates/gpu/common:data_type",
+    "//tensorflow/lite/delegates/gpu/common:model",
+    "//tensorflow/lite/delegates/gpu/common:model_builder",
+    "//tensorflow/lite/delegates/gpu/common:model_builder_helper",
+    "//tensorflow/lite/delegates/gpu/common:quantization_util",
+    "//tensorflow/lite/delegates/gpu/common:status",
+    "//tensorflow/lite/delegates:serialization",
+    "//tensorflow/lite/kernels:kernel_util",
+    "//tensorflow/lite/profiling/telemetry",
+    "//tensorflow/lite/profiling/telemetry/c:telemetry_setting",
+    "//tensorflow/lite/profiling/telemetry/c:telemetry_setting_internal",
+    "//tensorflow/lite/profiling/telemetry:telemetry_status",
+    "//tensorflow/lite:kernel_api",
+    "//tensorflow/lite:minimal_logging",
+    # go/keep-sorted end
+]
+
+config_setting(
+    name = "tflite_gpu_binary_release",
+    values = {"copt": "-DTFLITE_GPU_BINARY_RELEASE"},
+)
+
+config_setting(
+    name = "tflite_gpu_extra_gles_deps",
+    constraint_values = [
+        "@platforms//cpu:x86_64",
+        "@platforms//os:linux",
+    ],
+    values = {
+        "copt": "-DTFLITE_GPU_EXTRA_GLES_DEPS",
+    },
+)
+
+cc_library(
+    name = "gl_delegate",
+    srcs = ["gl_delegate.cc"],
+    hdrs = ["gl_delegate.h"],
+    linkopts = gpu_delegate_linkopts(),
+    deps = [
+        "//tensorflow/lite:kernel_api",
+        "//tensorflow/lite:minimal_logging",
+        "//tensorflow/lite/core/c:common",
+        "//tensorflow/lite/delegates/gpu/common:convert",
+        "//tensorflow/lite/delegates/gpu/common:gpu_info",
+        "//tensorflow/lite/delegates/gpu/common:model",
+        "//tensorflow/lite/delegates/gpu/common:model_builder",
+        "//tensorflow/lite/delegates/gpu/common:model_transformer",
+        "//tensorflow/lite/delegates/gpu/common:shape",
+        "//tensorflow/lite/delegates/gpu/common:status",
+        "//tensorflow/lite/delegates/gpu/common:tensor",
+        "//tensorflow/lite/delegates/gpu/common/transformations:model_transformations",
+        "//tensorflow/lite/delegates/gpu/gl:api",
+        "//tensorflow/lite/delegates/gpu/gl:command_queue",
+        "//tensorflow/lite/delegates/gpu/gl:compiler",
+        "//tensorflow/lite/delegates/gpu/gl:compiler_options",
+        "//tensorflow/lite/delegates/gpu/gl:egl_environment",
+        "//tensorflow/lite/delegates/gpu/gl:gl_buffer",
+        "//tensorflow/lite/delegates/gpu/gl:gl_call",
+        "//tensorflow/lite/delegates/gpu/gl:object",
+        "//tensorflow/lite/delegates/gpu/gl:object_manager",
+        "//tensorflow/lite/delegates/gpu/gl:request_gpu_info",
+        "//tensorflow/lite/delegates/gpu/gl:runtime_options",
+        "//tensorflow/lite/delegates/gpu/gl/converters:bhwc_to_phwc4",
+        "//tensorflow/lite/delegates/gpu/gl/converters:phwc4_to_bhwc",
+        "//tensorflow/lite/delegates/gpu/gl/kernels:registry",
+        "//tensorflow/lite/delegates/gpu/gl/workgroups:best_effort_calculator",
+        "@com_google_absl//absl/base:core_headers",
+        "@com_google_absl//absl/types:span",
+    ] + select({
+        "//conditions:default": [
+            "//tensorflow/lite/delegates/gpu/gl:common_cc_fbs",
+            "//tensorflow/lite/delegates/gpu/gl:metadata_cc_fbs",
+            "//tensorflow/lite/delegates/gpu/gl:workgroups_cc_fbs",
+            "//tensorflow/lite/schema:schema_fbs",
+            "@flatbuffers",
+        ],
+        ":tflite_gpu_binary_release": [],
+    }) + tflite_extra_gles_deps(),
+)
+
+objc_library(
+    name = "metal_delegate",
+    srcs = ["metal_delegate.mm"],
+    hdrs = ["metal_delegate.h"],
+    copts = CXX17_BAZEL_ONLY_COPTS,
+    module_name = "TensorFlowLiteCMetal",
+    sdk_frameworks = ["Metal"],
+    deps = [
+        "//tensorflow/lite:kernel_api",
+        "//tensorflow/lite:minimal_logging",
+        "//tensorflow/lite/core/c:common",
+        "//tensorflow/lite/delegates/gpu/common:convert",
+        "//tensorflow/lite/delegates/gpu/common:gpu_info",
+        "//tensorflow/lite/delegates/gpu/common:model",
+        "//tensorflow/lite/delegates/gpu/common:model_builder",
+        "//tensorflow/lite/delegates/gpu/common:model_transformer",
+        "//tensorflow/lite/delegates/gpu/common:precision",
+        "//tensorflow/lite/delegates/gpu/common:quantization_util",
+        "//tensorflow/lite/delegates/gpu/common:shape",
+        "//tensorflow/lite/delegates/gpu/common:status",
+        "//tensorflow/lite/delegates/gpu/common:types",
+        "//tensorflow/lite/delegates/gpu/metal:buffer_convert",
+        "//tensorflow/lite/delegates/gpu/metal:common",
+        "//tensorflow/lite/delegates/gpu/metal:inference_context",
+        "//tensorflow/lite/delegates/gpu/metal:metal_spatial_tensor",
+        "//tensorflow/lite/kernels:kernel_util",
+        "@com_google_absl//absl/container:flat_hash_set",
+        "@com_google_absl//absl/types:span",
+    ],
+)
+
+objc_library(
+    name = "metal_delegate_internal",
+    hdrs = ["metal_delegate_internal.h"],
+    copts = CXX17_BAZEL_ONLY_COPTS,
+    sdk_frameworks = ["Metal"],
+    deps = ["//tensorflow/lite/delegates/gpu:metal_delegate"],
+)
+
+# build -c opt --config android_arm64 --copt -Os --copt -DTFLITE_GPU_BINARY_RELEASE --linkopt -s --strip always :libtensorflowlite_gpu_gl.so
+cc_binary(
+    name = "libtensorflowlite_gpu_gl.so",
+    linkopts = [
+        "-Wl,-soname=libtensorflowlite_gpu_gl.so",
+    ] + gpu_delegate_linkopts() + select({
+        "//tensorflow:windows": [],
+        "//conditions:default": [
+            "-fvisibility=hidden",
+        ],
+    }) + tflite_pagesize_linkopts(),
+    linkshared = 1,
+    linkstatic = 1,
+    tags = [
+        "nobuilder",
+        "notap",
+    ],
+    deps = [":gl_delegate"],
+)
+
+# build -c opt --config android_arm64 --copt -Os --copt -DTFLITE_GPU_BINARY_RELEASE --linkopt -s --strip always :libtensorflowlite_gpu_delegate.so
+cc_binary(
+    name = "libtensorflowlite_gpu_delegate.so",
+    linkopts = [
+        "-Wl,-soname=libtensorflowlite_gpu_delegate.so",
+    ] + gpu_delegate_linkopts() + select({
+        "//tensorflow:windows": [],
+        "//conditions:default": [
+            "-fvisibility=hidden",
+        ],
+    }),
+    linkshared = 1,
+    linkstatic = 1,
+    tags = [
+        "nobuilder",
+        "notap",
+    ],
+    deps = [":delegate"],
+)
+
+# bazel build -c opt --cpu ios_arm64 --copt -Os --copt -DTFLITE_GPU_BINARY_RELEASE --copt -fvisibility=hidden --linkopt -s --strip always --cxxopt=-std=c++14 :libtensorflowlite_gpu_metal --apple_platform_type=ios
+ios_static_framework(
+    name = "tensorflow_lite_gpu_framework",
+    hdrs = [
+        "metal_delegate.h",
+        "metal_delegate_internal.h",
+    ],
+    minimum_os_version = "12.0",
+    deps = [":metal_delegate"],
+)
+
+# Note: Support for MacOS is best-effort at the moment.
+# bazel build -c opt --copt -Os --copt -DTFLITE_GPU_BINARY_RELEASE --copt -fvisibility=hidden --linkopt -s --strip always --cxxopt=-std=c++14 :tensorflow_lite_gpu_dylib --apple_platform_type=macos
+macos_dylib(
+    name = "tensorflow_lite_gpu_dylib",
+    linkopts = [
+        "-all_load",
+        "-dead_strip",
+    ],
+    minimum_os_version = "12.0",
+    tags = [
+        "manual",
+        "nobuilder",
+        "notap",
+    ],
+    deps = [
+        ":metal_delegate",
+        ":metal_delegate_internal",
+    ],
+)
+
+cc_library(
+    name = "api",
+    srcs = ["api.cc"],
+    hdrs = ["api.h"],
+    deps = [
+        "//tensorflow/lite/delegates/gpu/common:data_type",
+        "//tensorflow/lite/delegates/gpu/common:status",
+        "//tensorflow/lite/delegates/gpu/common:util",
+        "//tensorflow/lite/delegates/gpu/gl:portable",
+        "@com_google_absl//absl/types:span",
+        "@com_google_absl//absl/types:variant",
+        "@opencl_headers",
+        "@vulkan_headers//:vulkan_headers_no_prototypes",
+    ],
+)
+
+cc_library(
+    name = "spi",
+    hdrs = ["spi.h"],
+    deps = [
+        ":api",
+        "//tensorflow/lite/delegates/gpu/common:access_type",
+        "//tensorflow/lite/delegates/gpu/common:status",
+    ],
+)
+
+# Currently the GPU delegate needs to be built on Android (due to EGL dependency),
+# or built with -DCL_DELEGATE_NO_GL (disabling OpenGL backend fallback), or both.
+selects.config_setting_group(
+    name = "supports_gpu_delegate",
+    match_any = [
+        "//tensorflow:android",
+        "//tensorflow/lite/delegates/gpu/cl:opencl_delegate_no_gl",
+    ],
+)
+
+cc_library(
+    name = "delegate_options",
+    srcs = ["delegate_options.cc"],
+    hdrs = ["delegate_options.h"],
+    deps = ["//tensorflow/lite/core/c:common"],
+)
+
+# copybara:uncomment_begin(google-only)
+# cc_library(
+#     name = "delegate_no_gl",
+#     srcs = [
+#         # copybara:comment_begin(oss-only)
+# "android_version.cc",
+#         # copybara:comment_end
+#         "delegate.cc",
+#     ],
+#     hdrs = ["delegate.h"],
+#     defines = ["CL_DELEGATE_NO_GL"],
+#     linkopts = gpu_delegate_linkopts(),
+#     deps = _DELEGATE_NO_GL_DEPS + [
+#         "//tensorflow/lite/delegates/gpu/cl:api_no_gl",
+#         "//tensorflow/lite/delegates/gpu/gl:api2",
+#     ],
+# )
+# copybara:uncomment_end
+
+cc_library(
+    name = "delegate",
+    srcs = [
+        # copybara:comment_begin(oss-only)
+        "android_version.cc",
+        # copybara:comment_end
+        "delegate.cc",
+    ],
+    hdrs = ["delegate.h"],
+    linkopts = gpu_delegate_linkopts(),
+    deps = select({
+        "//tensorflow/lite/delegates/gpu/cl:opencl_delegate_no_gl": [],
+        "//conditions:default": [
+            "//tensorflow/lite/delegates/gpu/gl:api2",
+        ],
+    }) + _DELEGATE_NO_GL_DEPS + ["//tensorflow/lite/delegates/gpu/cl:api"],
+)
+
+cc_library(
+    name = "tflite_profile",
+    srcs = ["tflite_profile.cc"],
+    hdrs = ["tflite_profile.h"],
+    deps = [
+        "//tensorflow/lite/core/api",
+        "//tensorflow/lite/delegates/gpu/common/task:profiling_info",
+        "@com_google_absl//absl/time",
+    ],
+)
+
+cc_library(
+    name = "android_hardware_buffer",
+    srcs = ["android_hardware_buffer.cc"],
+    hdrs = ["android_hardware_buffer.h"],
+)
+
+cc_test(
+    name = "android_hardware_buffer_test",
+    srcs = ["android_hardware_buffer_test.cc"],
+    deps = [
+        ":android_hardware_buffer",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
+
+cc_library(
+    name = "async_buffers",
+    srcs = ["async_buffers.cc"],
+    hdrs = ["async_buffers.h"],
+    deps = [
+        ":android_hardware_buffer",
+        ":api",
+        "//tensorflow/lite/delegates/gpu/common:data_type",
+        "//tensorflow/lite/delegates/gpu/gl:gl_errors",
+        "@com_google_absl//absl/status",
+    ],
+)
+
+cc_test(
+    name = "async_buffers_test",
+    srcs = ["async_buffers_test.cc"],
+    tags = tf_gpu_tests_tags() + [
+        "local",
+        "nobuilder",
+        "notap",
+        "tflite_not_portable_ios",
+    ],
+    deps = [
+        ":android_hardware_buffer",
+        ":api",
+        ":async_buffers",
+        ":delegate",
+        "//tensorflow/lite/delegates/gpu/common:data_type",
+        "//tensorflow/lite/delegates/gpu/gl:egl_environment",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
+
+tflite_portable_test_suite()

tensorflow/lite/delegates/gpu/README.md

@@ -0,0 +1,174 @@
+# TFLite on GPU
+
+TensorFlow Lite (TFLite) supports several hardware accelerators.  This document
+describes how to use the GPU backend using the TFLite delegate APIs on Android
+and iOS.
+
+GPUs are designed to have high throughput for massively parallelizable
+workloads.  Thus, they are well-suited for deep neural nets which consists of a
+huge number of operators, each working on some input tensor(s) that can be
+easily divided into smaller workloads and carried out in parallel, typically
+resulting in lower latency.  In the best scenario, inference on the GPU may now
+run fast enough and now become suitable for real-time applications if it was not
+before.
+
+GPUs do their computation with 16-bit or 32-bit floating point numbers and do
+not require quantization for optimal performance unlike the CPUs.  If
+quantization of your neural network was not an option due to lower accuracy
+caused by lost precision, such concern can be discarded when running deep neural
+net models on the GPU.
+
+Another benefit that comes with GPU inference is its power efficiency.  GPUs
+carry out the computations in a very efficient and optimized way, so that they
+consume less power and generate less heat than when the same task is run on the
+CPUs.
+
+TFLite on GPU supports the following ops in 16-bit and 32-bit float precision:
+
+* `ADD v1`
+* `AVERAGE_POOL_2D v1`
+* `CONCATENATION v1`
+* `CONV_2D v1`
+* `DEPTHWISE_CONV_2D v1-2`
+* `EXP v1`
+* `FULLY_CONNECTED v1`
+* `LOGISTIC v1`
+* `LSTM v2 (Basic LSTM only)`
+* `MAX_POOL_2D v1`
+* `MAXIMUM v1`
+* `MINIMUM v1`
+* `MUL v1`
+* `PAD v1`
+* `PRELU v1`
+* `RELU v1`
+* `RELU6 v1`
+* `RESHAPE v1`
+* `RESIZE_BILINEAR v1-3`
+* `SOFTMAX v1`
+* `STRIDED_SLICE v1`
+* `SUB v1`
+* `TRANSPOSE_CONV v1`
+
+## Basic Usage
+
+**Note:** Following section describes the example usage for Android GPU delegate
+with C++. For other languages and platforms, please see
+[the documentation](https://www.tensorflow.org/lite/performance/gpu).
+
+Using TFLite on GPU is as simple as getting the GPU delegate via
+`TfLiteGpuDelegateV2Create()` and then passing it to
+`InterpreterBuilder::AddDelegate()`:
+
+```c++
+////////
+// Set up InterpreterBuilder.
+auto model = FlatBufferModel::BuildFromFile(model_path);
+ops::builtin::BuiltinOpResolver op_resolver;
+InterpreterBuilder interpreter_builder(*model, op_resolver);
+
+////////
+// NEW: Prepare GPU delegate.
+auto* delegate = TfLiteGpuDelegateV2Create(/*default options=*/nullptr);
+interpreter_builder.AddDelegate(delegate);
+
+////////
+// Set up Interpreter.
+std::unique_ptr<Interpreter> interpreter;
+if (interpreter_builder(&interpreter) != kTfLiteOk) return;
+
+////////
+// IMPORTANT: AllocateTensors can be called only AFTER ModifyGraphWithDelegate
+
+////////
+// Run inference.
+WriteToInputTensor(interpreter->typed_input_tensor<float>(0));
+if (interpreter->Invoke() != kTfLiteOk) return;
+ReadFromOutputTensor(interpreter->typed_output_tensor<float>(0));
+
+////////
+// Clean up.
+TfLiteGpuDelegateV2Delete(delegate);
+```
+
+*IMPORTANT:* When calling `Interpreter::ModifyGraphWithDelegate()` or
+`InterpreterBuilder::operator()` or
+`Interpreter::Invoke()`, the caller must have a `EGLContext` in the current
+thread and `Interpreter::Invoke()` must be called from the same `EGLContext`.
+If such `EGLContext` does not exist, the delegate will internally create one,
+but then the developer must ensure that `Interpreter::Invoke()` is always called
+from the same thread `InterpreterBuilder::operator()` or
+`Interpreter::ModifyGraphWithDelegate()` was called.
+
+## Building and Runtime
+
+TFLite GPU backend uses OpenGL ES 3.1 compute shaders or OpenCL.
+
+```sh
+bazel build --config android_arm64 //path/to/your:project
+```
+
+Metal shaders are used for iOS, which were introduced with iOS 8.  Thus,
+compilation flags should look like:
+
+```sh
+bazel build --config ios_fat //path/to/your:project
+```
+
+## Advanced Usage: Delegate Options
+
+There are GPU options that can be set and passed on to
+`TfLiteGpuDelegateV2Create()`. When option is set to `nullptr` as shown in the
+Basic Usage, it translates to:
+
+```c++
+const TfLiteGpuDelegateOptionsV2 kDefaultOptions =
+    TfLiteGpuDelegateOptionsV2Default();
+```
+
+Similar for `TFLGpuDelegateCreate()`:
+
+```c++
+const TFLGpuDelegateOptions kDefaultOptions = {
+  .allow_precision_loss = false,
+  .wait_type = TFLGpuDelegateWaitTypePassive,
+  .enable_quantization = false,
+};
+```
+
+While it is convenient to just supply `nullptr`, it is recommended to explicitly
+set the options to avoid any unexpected artifacts in case default values are
+changed.
+
+*IMPORTANT:* Note that the default option may not be the fastest. For faster
+execution, you may want to set `allow_precision_loss` to `true` so that the GPU
+performs FP16 calculation internally, and set `wait_type` to
+`TFLGpuDelegateWaitTypeAggressive` to avoid GPU sleep mode.
+
+## Tips and Tricks
+
+* Some operations that are trivial on CPU side may be high cost in GPU land.
+  One class of such operation is various forms of reshape operations (including
+  `BATCH_TO_SPACE`, `SPACE_TO_BATCH`, `SPACE_TO_DEPTH`, etc.).  If those ops
+  are inserted into the network just for the network architect's logical
+  thinking, it is worth removing them for performance.
+
+* On GPU, tensor data is sliced into 4-channels.  Thus, a computation on a
+  tensor of shape `[B, H, W, 5]` will perform about the same on a tensor of
+  shape `[B, H, W, 8]`, but significantly worse than `[B, H, W, 4]`.
+
+* In that sense, if the camera hardware supports image frames in RGBA, feeding
+  that 4-channel input is significantly faster as a memory copy (from 3-channel
+  RGB to 4-channel RGBX) can be avoided.
+
+* For performance [best practices](https://www.tensorflow.org/lite/performance/best_practices), do not hesitate to re-train your classifier with
+  mobile-optimized network architecture.  That is a significant part of
+  optimization for on-device inference.
+
+## Publication
+
+*   [On-Device Neural Net Inference with Mobile GPUs](https://arxiv.org/abs/1907.01989)
+    *   Juhyun Lee, Nikolay Chirkov, Ekaterina Ignasheva, Yury Pisarchyk, Mogan
+        Shieh, Fabio Riccardi, Raman Sarokin, Andrei Kulik, and Matthias
+        Grundmann
+    *   CVPR Workshop
+        [Efficient Deep Learning for Computer Vision (ECV2019)](https://sites.google.com/corp/view/ecv2019)

tensorflow/lite/delegates/gpu/android_hardware_buffer.cc

@@ -0,0 +1,53 @@
+/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/lite/delegates/gpu/android_hardware_buffer.h"
+
+#include <dlfcn.h>
+
+namespace tflite::gpu {
+
+OptionalAndroidHardwareBuffer::OptionalAndroidHardwareBuffer() {
+#ifdef __ANDROID__
+  dlopen_handle_ = dlopen("libnativewindow.so", RTLD_NOW);
+  if (dlopen_handle_ == nullptr) {
+    supported_ = false;
+    return;
+  }
+  allocate_ = reinterpret_cast<decltype(allocate_)>(
+      dlsym(dlopen_handle_, "AHardwareBuffer_allocate"));
+  acquire_ = reinterpret_cast<decltype(acquire_)>(
+      dlsym(dlopen_handle_, "AHardwareBuffer_acquire"));
+  release_ = reinterpret_cast<decltype(release_)>(
+      dlsym(dlopen_handle_, "AHardwareBuffer_release"));
+  describe_ = reinterpret_cast<decltype(describe_)>(
+      dlsym(dlopen_handle_, "AHardwareBuffer_describe"));
+  is_supported_ = reinterpret_cast<decltype(is_supported_)>(
+      dlsym(dlopen_handle_, "AHardwareBuffer_isSupported"));
+  supported_ =
+      (allocate_ != nullptr && acquire_ != nullptr && release_ != nullptr &&
+       describe_ != nullptr && is_supported_ != nullptr);
+#else
+  dlopen_handle_ = nullptr;
+  allocate_ = nullptr;
+  acquire_ = nullptr;
+  release_ = nullptr;
+  describe_ = nullptr;
+  is_supported_ = nullptr;
+  supported_ = false;
+#endif
+}
+
+}  // namespace tflite::gpu

tensorflow/lite/delegates/gpu/android_hardware_buffer.h

@@ -0,0 +1,130 @@
+/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_LITE_DELEGATES_GPU_ANDROID_HARDWARE_BUFFER_H_
+#define TENSORFLOW_LITE_DELEGATES_GPU_ANDROID_HARDWARE_BUFFER_H_
+
+#include <stdint.h>
+
+#ifdef __ANDROID__
+#include <android/hardware_buffer.h>
+#else
+extern "C" {
+typedef struct AHardwareBuffer AHardwareBuffer;
+
+// struct is a copy of the Android NDK AHardwareBuffer_Desc struct in the link
+// below
+// https://developer.android.com/ndk/reference/struct/a-hardware-buffer-desc
+typedef struct AHardwareBuffer_Desc AHardwareBuffer_Desc;
+struct AHardwareBuffer_Desc {
+  uint32_t width;
+  uint32_t height;
+  uint32_t layers;
+  uint32_t format;
+  uint64_t usage;
+  uint32_t stride;
+  uint32_t rfu0;
+  uint64_t rfu1;
+};
+}  // extern "C"
+#endif  // __ANDROID__
+
+namespace tflite::gpu {
+
+// This header file and singleton class encapsulates the following Android NDK
+// features
+//   - header <android/hardware_buffer.h>
+//   - opaque struct type AHardwareBuffer
+//   - struct type AHardwareBuffer_Desc
+//   - function AHardwareBuffer_isSupported
+//   - function AHardwareBuffer_allocate
+//   - function AHardwareBuffer_acquire
+//   - function AHardwareBuffer_release
+//   - function AHardwareBuffer_describe
+//   - library libnativewindow.so (for the above features)
+//
+// For documentation on these features, see
+// <https://developer.android.com/ndk/reference/group/a-hardware-buffer>:
+//
+// Unlike using the native NDK functionality directly, this class only has a
+// run-time dependency on API level 26, not a build-time dependency.  So it can
+// be used even when building with NDK min SDK level < 26, as long as you are
+// very careful to check that Supported() returns true before calling any other
+// methods.
+class OptionalAndroidHardwareBuffer {
+ public:
+  static OptionalAndroidHardwareBuffer& Instance() {
+    static OptionalAndroidHardwareBuffer instance;
+    return instance;
+  }
+
+  // Returns true if the functionality in this class is supported.
+  bool Supported() { return supported_; }
+
+  // Like AHardwareBuffer_isSupported.
+  // Caller must check that Supported() returns true before calling this
+  // function.
+  int IsSupported(const AHardwareBuffer_Desc* description) {
+    return is_supported_(description);
+  }
+
+  // Like AHardwareBuffer_allocate.
+  // Caller must check that Supported() returns true before calling this
+  // function.
+  int Allocate(const AHardwareBuffer_Desc* description,
+               AHardwareBuffer** buffer) {
+    return allocate_(description, buffer);
+  }
+
+  // Like AHardwareBuffer_acquire.
+  // Caller must check that Supported() returns true before calling this
+  // function.
+  void Acquire(AHardwareBuffer* buffer) { return acquire_(buffer); }
+
+  // Like AHardwareBuffer_release.
+  // Caller must check that Supported() returns true before calling this
+  // function.
+  void Release(AHardwareBuffer* buffer) { return release_(buffer); }
+
+  // Like AHardwareBuffer_describe.
+  // Caller must check that Supported() returns true before calling this
+  // function.
+  void Describe(AHardwareBuffer* buffer, AHardwareBuffer_Desc* desc) {
+    return describe_(buffer, desc);
+  }
+
+ private:
+  void* dlopen_handle_;
+  int (*is_supported_)(const AHardwareBuffer_Desc* desc);
+  int (*allocate_)(const AHardwareBuffer_Desc* desc, AHardwareBuffer** buffer);
+  void (*acquire_)(AHardwareBuffer* buffer);
+  void (*release_)(AHardwareBuffer* buffer);
+  void (*describe_)(AHardwareBuffer* buffer, AHardwareBuffer_Desc* desc);
+  bool supported_;
+
+  OptionalAndroidHardwareBuffer();
+  OptionalAndroidHardwareBuffer(const OptionalAndroidHardwareBuffer&) = delete;
+  // Note that we deliberately do not call dlclose() in the destructor; doing
+  // so would complicate the code and would unnecessarily introduce additional
+  // failure scenarios. The object is a singleton and so is only destroyed when
+  // the process is about to exit, and the OS will automatically reclaim the
+  // resources on process exit anyway, so calling dlclose would only slow down
+  // process exit.
+  ~OptionalAndroidHardwareBuffer() = default;
+};
+
+}  // namespace tflite::gpu
+
+#endif  // TENSORFLOW_LITE_DELEGATES_GPU_ANDROID_HARDWARE_BUFFER_H_

tensorflow/lite/delegates/gpu/android_hardware_buffer_test.cc

@@ -0,0 +1,75 @@
+/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/lite/delegates/gpu/android_hardware_buffer.h"
+
+#include <gtest/gtest.h>
+
+using tflite::gpu::OptionalAndroidHardwareBuffer;
+auto Instance = OptionalAndroidHardwareBuffer::Instance;
+
+namespace {
+
+#ifndef __ANDROID__
+
+TEST(OptionalAndroidHardwareBufferTest, NotSupportedOnNonAndroid) {
+  EXPECT_EQ(Instance().Supported(), false);
+}
+
+#else  // defined(__ANDROID__)
+
+TEST(OptionalAndroidHardwareBufferTest, SupportedOnAndroid) {
+  EXPECT_EQ(Instance().Supported(), true);
+}
+
+TEST(OptionalAndroidHardwareBufferTest, CanAllocateAndReleaseOnAndroid) {
+  EXPECT_EQ(Instance().Supported(), true);
+  AHardwareBuffer* buffer;
+  AHardwareBuffer_Desc description{};
+  description.width = 1600;
+  description.height = 1;
+  description.layers = 1;
+  description.rfu0 = 0;
+  description.rfu1 = 0;
+  description.stride = 1;
+  description.format = AHARDWAREBUFFER_FORMAT_BLOB;
+  description.usage = AHARDWAREBUFFER_USAGE_CPU_READ_OFTEN;
+  EXPECT_TRUE(Instance().IsSupported(&description));
+  EXPECT_EQ(Instance().Allocate(&description, &buffer), 0);
+  Instance().Release(buffer);
+}
+
+TEST(OptionalAndroidHardwareBufferTest, CanAcquireAndReleaseOnAndroid) {
+  EXPECT_EQ(Instance().Supported(), true);
+  AHardwareBuffer* buffer;
+  AHardwareBuffer_Desc description{};
+  description.width = 1600;
+  description.height = 1;
+  description.layers = 1;
+  description.rfu0 = 0;
+  description.rfu1 = 0;
+  description.stride = 1;
+  description.format = AHARDWAREBUFFER_FORMAT_BLOB;
+  description.usage = AHARDWAREBUFFER_USAGE_CPU_READ_OFTEN;
+  EXPECT_TRUE(Instance().IsSupported(&description));
+  EXPECT_EQ(Instance().Allocate(&description, &buffer), 0);
+  Instance().Acquire(buffer);
+  Instance().Release(buffer);  // To match Acquire
+  Instance().Release(buffer);  // To match Allocate
+}
+
+#endif  // defined(__ANDROID__)
+
+}  // namespace

tensorflow/lite/delegates/gpu/android_version.cc

@@ -0,0 +1,58 @@
+/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#if defined(__ANDROID__)
+
+#include <pthread.h>
+#include <stdlib.h>
+#include <string.h>
+#include <sys/system_properties.h>
+
+extern "C" {
+
+static int SdkVersion;
+static int IsPreRelease;
+
+static void readSystemProperties(void) {
+  char buf[PROP_VALUE_MAX];
+
+  if (__system_property_get("ro.build.version.sdk", buf) == 0) {
+    // When the system property doesn't exist, defaults to future API level.
+    SdkVersion = __ANDROID_API_FUTURE__;
+  } else {
+    SdkVersion = atoi(buf);  // NOLINT(runtime/deprecated_fn)
+  }
+
+  if (__system_property_get("ro.build.version.codename", buf) == 0) {
+    IsPreRelease = 1;
+  } else {
+    IsPreRelease = strcmp(buf, "REL") != 0;
+  }
+  return;
+}
+
+int32_t __isOSVersionAtLeast(int32_t Major, int32_t Minor, int32_t Subminor) {
+  (int32_t) Minor;
+  (int32_t) Subminor;
+  static pthread_once_t once = PTHREAD_ONCE_INIT;
+  pthread_once(&once, readSystemProperties);
+
+  return SdkVersion >= Major ||
+         (IsPreRelease && Major == __ANDROID_API_FUTURE__);
+}
+
+}  // extern "C"
+
+#endif  // defined(__ANDROID__)

tensorflow/lite/delegates/gpu/api.cc

@@ -0,0 +1,213 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/lite/delegates/gpu/api.h"
+
+#include <cstdint>
+#include <limits>
+#include <variant>
+
+#include "tensorflow/lite/delegates/gpu/common/util.h"
+
+namespace tflite {
+namespace gpu {
+namespace {
+
+struct ObjectTypeGetter {
+  ObjectType operator()(std::monostate) const { return ObjectType::UNKNOWN; }
+  ObjectType operator()(OpenGlBuffer) const { return ObjectType::OPENGL_SSBO; }
+  ObjectType operator()(OpenGlTexture) const {
+    return ObjectType::OPENGL_TEXTURE;
+  }
+  ObjectType operator()(OpenClBuffer) const {
+    return ObjectType::OPENCL_BUFFER;
+  }
+  ObjectType operator()(OpenClTexture) const {
+    return ObjectType::OPENCL_TEXTURE;
+  }
+  ObjectType operator()(VulkanBuffer) const {
+    return ObjectType::VULKAN_BUFFER;
+  }
+  ObjectType operator()(VulkanTexture) const {
+    return ObjectType::VULKAN_TEXTURE;
+  }
+  ObjectType operator()(CpuMemory) const { return ObjectType::CPU_MEMORY; }
+};
+
+struct ObjectValidityChecker {
+  bool operator()(std::monostate) const { return false; }
+  bool operator()(OpenGlBuffer obj) const { return obj.id != GL_INVALID_INDEX; }
+  bool operator()(OpenGlTexture obj) const {
+    return obj.id != GL_INVALID_INDEX && obj.format != GL_INVALID_ENUM;
+  }
+  bool operator()(OpenClBuffer obj) const { return obj.memobj; }
+  bool operator()(OpenClTexture obj) const { return obj.memobj; }
+  bool operator()(VulkanBuffer obj) const { return obj.memory; }
+  bool operator()(VulkanTexture obj) const { return obj.memory; }
+  bool operator()(CpuMemory obj) const {
+    return obj.data != nullptr && obj.size_bytes > 0 &&
+           (data_type == DataType::UNKNOWN || data_type == DataType::BOOL ||
+            obj.size_bytes % SizeOf(data_type) == 0);
+  }
+  DataType data_type;
+};
+
+}  // namespace
+
+bool IsValid(const ObjectDef& def) {
+  return def.data_type != DataType::UNKNOWN &&
+         def.data_layout != DataLayout::UNKNOWN &&
+         def.object_type != ObjectType::UNKNOWN;
+}
+
+ObjectType GetType(const TensorObject& object) {
+  return std::visit(ObjectTypeGetter{}, object);
+}
+
+bool IsValid(const TensorObjectDef& def) {
+  return IsValid(def.object_def) &&
+         NumElements(def) <= std::numeric_limits<int32_t>::max();
+}
+
+bool IsValid(const TensorObjectDef& def, const TensorObject& object) {
+  return IsValid(def) && GetType(object) == def.object_def.object_type &&
+         std::visit(ObjectValidityChecker{def.object_def.data_type}, object);
+}
+
+bool IsObjectPresent(ObjectType type, const TensorObject& obj) {
+  switch (type) {
+    case ObjectType::CPU_MEMORY:
+      return std::holds_alternative<CpuMemory>(obj);
+    case ObjectType::OPENGL_SSBO:
+      return std::holds_alternative<OpenGlBuffer>(obj);
+    case ObjectType::OPENGL_TEXTURE:
+      return std::holds_alternative<OpenGlTexture>(obj);
+    case ObjectType::OPENCL_BUFFER:
+      return std::holds_alternative<OpenClBuffer>(obj);
+    case ObjectType::OPENCL_TEXTURE:
+      return std::holds_alternative<OpenClTexture>(obj);
+    case ObjectType::VULKAN_BUFFER:
+      return std::holds_alternative<VulkanBuffer>(obj);
+    case ObjectType::VULKAN_TEXTURE:
+      return std::holds_alternative<VulkanTexture>(obj);
+    case ObjectType::UNKNOWN:
+      return false;
+  }
+}
+
+bool IsObjectInitialized(const TensorObject& obj) {
+  return GetType(obj) != ObjectType::UNKNOWN;
+}
+
+int64_t NumElements(const TensorObjectDef& def) {
+  const auto& d = def.dimensions;
+  switch (def.object_def.data_layout) {
+    case DataLayout::BHWC:
+      return d.product();
+    case DataLayout::HWDC4:
+    case DataLayout::HDWC4:
+    case DataLayout::DHWC4:
+      return static_cast<int64_t>(d.b) * d.h * d.w * AlignByN(d.c, 4);
+    case DataLayout::UNKNOWN:
+      return 0;
+  }
+  return 0;
+}
+
+int GetPosition(const InferenceOptions& options, InferencePriority p) {
+  if (options.priority1 == p) return 1;
+  if (options.priority2 == p) return 2;
+  if (options.priority3 == p) return 3;
+  return 4;  // least important
+}
+
+PriorityImportance GetRelativeImportance(const InferenceOptions& options,
+                                         InferencePriority p1,
+                                         InferencePriority p2) {
+  int p1_position = GetPosition(options, p1);
+  int p2_position = GetPosition(options, p2);
+  if (p1_position == p2_position) return PriorityImportance::UNKNOWN;
+  return p1_position < p2_position ? PriorityImportance::HIGHER
+                                   : PriorityImportance::LOWER;
+}
+
+bool IsValid(const InferenceOptions& options) {
+  if (options.usage == InferenceUsage::UNKNOWN) {
+    return false;
+  }
+  if (options.priority1 == InferencePriority::UNKNOWN ||
+      options.priority2 == InferencePriority::UNKNOWN ||
+      options.priority3 == InferencePriority::UNKNOWN) {
+    return false;
+  }
+  if (options.priority1 == InferencePriority::AUTO) {
+    return false;
+  }
+  if (options.priority2 == InferencePriority::AUTO &&
+      options.priority3 != InferencePriority::AUTO) {
+    return false;
+  }
+  if (options.priority1 == options.priority2 ||
+      options.priority1 == options.priority3) {
+    return false;
+  }
+  if (options.priority2 == options.priority3 &&
+      options.priority2 != InferencePriority::AUTO) {
+    return false;
+  }
+  return true;
+}
+
+// Implementation note: this resolution logic is shared between GL and CL
+// backends, but they might have own logic. Thus, the function is defined
+// here just for code re-use purposes.
+void ResolveAutoPriority(InferenceOptions* options) {
+  // priority1 can not be AUTO as it would make options invalid.
+  if (options->priority2 == InferencePriority::AUTO) {
+    switch (options->priority1) {
+      case InferencePriority::MIN_LATENCY:
+        options->priority2 = InferencePriority::MIN_MEMORY_USAGE;
+        options->priority3 = InferencePriority::MAX_PRECISION;
+        return;
+      case InferencePriority::MIN_MEMORY_USAGE:
+        options->priority2 = InferencePriority::MAX_PRECISION;
+        options->priority3 = InferencePriority::MIN_LATENCY;
+        return;
+      case InferencePriority::MAX_PRECISION:
+        options->priority2 = InferencePriority::MIN_LATENCY;
+        options->priority3 = InferencePriority::MIN_MEMORY_USAGE;
+        return;
+      case InferencePriority::UNKNOWN:
+      case InferencePriority::AUTO:
+        // Invalid and unreachable option.
+        return;
+    }
+  }
+
+  if (options->priority3 == InferencePriority::AUTO) {
+    // Simply add missing priority
+    if (GetPosition(*options, InferencePriority::MIN_LATENCY) == 4) {
+      options->priority3 = InferencePriority::MIN_LATENCY;
+    } else if (GetPosition(*options, InferencePriority::MAX_PRECISION) == 4) {
+      options->priority3 = InferencePriority::MAX_PRECISION;
+    } else if (GetPosition(*options, InferencePriority::MIN_MEMORY_USAGE) ==
+               4) {
+      options->priority3 = InferencePriority::MIN_MEMORY_USAGE;
+    }
+  }
+}
+
+}  // namespace gpu
+}  // namespace tflite

tensorflow/lite/delegates/gpu/api.h

@@ -0,0 +1,416 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_LITE_DELEGATES_GPU_API_H_
+#define TENSORFLOW_LITE_DELEGATES_GPU_API_H_
+
+// Usage example:
+//
+//   // Builder is created from a model using GPU-specific parameters.
+//   std::unique_ptr<InferenceBuilder> builder = ...;
+//
+//   // input data is coming from a texture
+//   // output data goes to CPU
+//   builder->SetInputObjectDef(0, {DataType::FLOAT16, DataLayout::PHWC4,
+//                                  ObjectType::OPENGL_TEXTURE, true});
+//   builder->SetOutputObjectDef(0, {DataType::FLOAT32, DataLayout::BHWC,
+//                                  ObjectType::CPU_MEMORY, false});
+//   std::unique_ptr<InferenceRunner> runner;
+//   RETURN_IF_ERROR(builder->Build(&runner));  // may take significant time.
+//   RETURN_IF_ERROR(
+//       runner->SetInputObject(0, OpenGlTexture{texture_ud, texture_format}));
+//   RETURN_IF_ERROR(runner->Run());
+
+#include <cstdint>
+#include <memory>
+#include <variant>
+#include <vector>
+
+#include "absl/types/span.h"
+#include "absl/types/variant.h"
+#include <CL/cl.h>
+#include "tensorflow/lite/delegates/gpu/common/data_type.h"
+#include "tensorflow/lite/delegates/gpu/common/status.h"
+#include "tensorflow/lite/delegates/gpu/common/util.h"
+
+// The `absl::Status` conflicts with the macro definition in the X11/Xlib.h,
+// undefine the VK_USE_PLATFORM_XLIB_KHR to exclude the header file.
+#undef VK_USE_PLATFORM_XLIB_KHR
+#include "vulkan/vulkan.h"  // from @vulkan_headers
+
+#define GL_NO_PROTOTYPES
+#define EGL_NO_PROTOTYPES
+#include "tensorflow/lite/delegates/gpu/gl/portable_gl31.h"
+#undef GL_NO_PROTOTYPES
+#undef EGL_NO_PROTOTYPES
+
+namespace tflite {
+namespace gpu {
+
+// Common abbreviations:
+//   B  - batch
+//   H  - height
+//   W  - width
+//   C  - channels
+//   D  - depth := DivideRoundUp(C, 4)
+//   C4 - is the constant = 4.
+enum class DataLayout {
+  UNKNOWN,
+  BHWC,
+  DHWC4,
+  HWDC4,
+  HDWC4,
+};
+
+enum class ObjectType {
+  UNKNOWN,
+  OPENGL_SSBO,
+  OPENGL_TEXTURE,
+  CPU_MEMORY,
+  OPENCL_TEXTURE,
+  OPENCL_BUFFER,
+  VULKAN_BUFFER,
+  VULKAN_TEXTURE
+};
+
+struct OpenGlBuffer {
+  OpenGlBuffer() = default;
+  explicit OpenGlBuffer(GLuint new_id) : id(new_id) {}
+
+  GLuint id = GL_INVALID_INDEX;
+};
+
+struct OpenGlTexture {
+  OpenGlTexture() = default;
+  OpenGlTexture(GLuint new_id, GLenum new_format)
+      : id(new_id), format(new_format) {}
+
+  GLuint id = GL_INVALID_INDEX;
+  GLenum format = GL_INVALID_ENUM;
+};
+
+struct OpenClBuffer {
+  OpenClBuffer() = default;
+  explicit OpenClBuffer(cl_mem new_memobj) : memobj(new_memobj) {}
+
+  cl_mem memobj = nullptr;
+};
+
+struct OpenClTexture {
+  OpenClTexture() = default;
+  explicit OpenClTexture(cl_mem new_memobj) : memobj(new_memobj) {}
+
+  cl_mem memobj = nullptr;
+  // TODO(akulik): should it specify texture format?
+};
+
+struct VulkanBuffer {
+  VulkanBuffer() = default;
+  explicit VulkanBuffer(VkBuffer buffer_, VkDeviceSize size_,
+                        VkDeviceMemory memory_, VkDeviceSize offset_)
+      : buffer(buffer_), size(size_), memory(memory_), offset(offset_) {}
+
+  VkBuffer buffer;
+  VkDeviceSize size;
+  VkDeviceMemory memory;
+  VkDeviceSize offset;
+};
+
+struct VulkanTexture {
+  VulkanTexture() = default;
+  explicit VulkanTexture(VkDeviceMemory new_memory) : memory(new_memory) {}
+
+  VkImage image;
+  VkImageView image_view;
+  VkFormat format;
+  VkExtent3D extent;
+  VkDeviceMemory memory;
+  VkDeviceSize offset;
+};
+
+struct VulkanMemory {
+  VulkanMemory() = default;
+  explicit VulkanMemory(VkDeviceMemory new_memory) : memory(new_memory) {}
+
+  VkDeviceMemory memory;
+  VkDeviceSize size;
+  VkDeviceSize offset;
+};
+
+struct CpuMemory {
+  CpuMemory() = default;
+  CpuMemory(void* new_data, size_t new_size_bytes)
+      : data(new_data), size_bytes(new_size_bytes) {}
+
+  void* data = nullptr;
+  size_t size_bytes = 0;
+};
+
+template <typename T>
+inline CpuMemory MakeCpuMemory(absl::Span<T> t) {
+  CpuMemory m;
+  m.data = t.data();
+  m.size_bytes = t.size() * sizeof(T);
+  return m;
+}
+
+template <typename T>
+inline CpuMemory MakeReadableCpuMemory(absl::Span<const T> t) {
+  CpuMemory m;
+  m.data = const_cast<T*>(t.data());
+  m.size_bytes = t.size() * sizeof(T);
+  return m;
+}
+
+// Defines object representation.
+struct ObjectDef {
+  DataType data_type = DataType::UNKNOWN;
+  DataLayout data_layout = DataLayout::UNKNOWN;
+  ObjectType object_type = ObjectType::UNKNOWN;
+
+  // If true, then object is managed externally and needs to be provided to
+  // InferenceRunner by a user before running inference.
+  //
+  // User-provided objects will not be re-used internally for any purpose to
+  // lower overall memory usage.
+  bool user_provided = false;
+
+  bool operator==(const ObjectDef& other) const {
+    return data_type == other.data_type && data_layout == other.data_layout &&
+           object_type == other.object_type &&
+           user_provided == other.user_provided;
+  }
+};
+
+bool IsValid(const ObjectDef& def);
+
+struct Dimensions {
+  Dimensions() : b(1), h(1), w(1), c(1) {}
+
+  Dimensions(int32_t batch, int32_t height, int32_t width, int32_t channels)
+      : b(batch), h(height), w(width), c(channels) {}
+
+  int32_t d() const { return DivideRoundUp(c, 4); }
+
+  int64_t product() const { return static_cast<int64_t>(b) * h * w * c; }
+
+  bool operator==(const Dimensions& other) const {
+    return b == other.b && h == other.h && w == other.w && c == other.c;
+  }
+
+  int32_t b;
+  int32_t h;
+  int32_t w;
+  int32_t c;
+};
+
+// Connects tensor shape with corresponding object definition.
+struct TensorObjectDef {
+  // Dimensions semantic is defined by corresponding DataLayout.
+  Dimensions dimensions;
+  ObjectDef object_def;
+
+  bool operator==(const TensorObjectDef& other) const {
+    return dimensions == other.dimensions && object_def == other.object_def;
+  }
+};
+
+// @return true if tensor object def is defined.
+bool IsValid(const TensorObjectDef& def);
+
+// @return the number of elements in a tensor object.
+int64_t NumElements(const TensorObjectDef& def);
+
+using TensorObject =
+    std::variant<std::monostate, OpenGlBuffer, OpenGlTexture, CpuMemory,
+                 OpenClBuffer, OpenClTexture, VulkanBuffer, VulkanTexture>;
+
+// @return true if object is set and corresponding values are defined.
+bool IsValid(const TensorObjectDef& def, const TensorObject& object);
+
+ObjectType GetType(const TensorObject& object);
+
+// @return true if corresponding object is set for the given type
+bool IsObjectPresent(ObjectType type, const TensorObject& obj);
+
+// @return true if corresponding object has already been initialized and
+// assigned with a specific ObjectType.
+bool IsObjectInitialized(const TensorObject& obj);
+
+class InferenceRunner;
+
+// Allows to inspect and change input and output definitions before a graph is
+// prepared for the inference.
+class InferenceBuilder {
+ public:
+  virtual ~InferenceBuilder() {}
+
+  // Returns inference graph inputs and outputs definitions.
+  virtual std::vector<TensorObjectDef> inputs() const = 0;
+  virtual std::vector<TensorObjectDef> outputs() const = 0;
+
+  // Sets new shape for the input if underlying implementation and graph
+  // structure allows dynamic tensors.
+  virtual absl::Status SetInputShape(int index,
+                                     const Dimensions& dimensions) = 0;
+
+  // Updates object definitions for the given index. Implementation may allow
+  // to use different layouts and/or data type conversions between objects
+  // defined in a graph and given objects, for example:
+  //   input '0' is DataType::FLOAT32, DataLayout::BHWC.
+  //   A user, however, has an input in DataType::FLOAT16, DataLayout::PHWC4.
+  //   An implementation may allow this transformation to happen automatically
+  //   under the hood.
+  virtual absl::Status SetInputObjectDef(int index, ObjectDef def) = 0;
+  virtual absl::Status SetOutputObjectDef(int index, ObjectDef def) = 0;
+  virtual absl::Status SetAllInputObjectDefsTo(ObjectDef def) {
+    auto input_defs = inputs();
+    for (int i = 0; i < input_defs.size(); ++i) {
+      RETURN_IF_ERROR(SetInputObjectDef(i, def));
+    }
+    return absl::OkStatus();
+  }
+  virtual absl::Status SetAllOutputObjectDefsTo(ObjectDef def) {
+    auto output_defs = outputs();
+    for (int i = 0; i < output_defs.size(); ++i) {
+      RETURN_IF_ERROR(SetOutputObjectDef(i, def));
+    }
+    return absl::OkStatus();
+  }
+
+  // Creates new instance of the inference runner. InferenceBuilder stays valid
+  // and could be used to create another inference runner if needed.
+  //
+  // This method may take significant time to prepare new inference runner. For
+  // example, it may require to compile OpenGL shaders.
+  virtual absl::Status Build(std::unique_ptr<InferenceRunner>* runner) = 0;
+};
+
+// Runs prepared inference. Every object marked as external needs to be set
+// prior calling Run method.
+class InferenceRunner {
+ public:
+  virtual ~InferenceRunner() {}
+
+  // Returns inference graph inputs and outputs definitions.
+  virtual std::vector<TensorObjectDef> inputs() const = 0;
+  virtual std::vector<TensorObjectDef> outputs() const = 0;
+
+  // Getters provide access to underlying objects for the given index.
+  // Setters allow to set or change external object for the given index. Note,
+  // object need to match object definition set before in InferenceBuilder.
+
+  virtual absl::Status GetInputObject(int index, TensorObject* object) = 0;
+  virtual absl::Status GetOutputObject(int index, TensorObject* object) = 0;
+  virtual absl::Status SetInputObject(int index, TensorObject object) = 0;
+  virtual absl::Status SetOutputObject(int index, TensorObject object) = 0;
+
+  virtual absl::Status Run() = 0;
+};
+
+// Encapsulated compilation/runtime tradeoffs.
+enum class InferenceUsage {
+  UNKNOWN,
+
+  // InferenceRunner will be used only once. Therefore, it is important to
+  // minimize bootstrap time as well.
+  FAST_SINGLE_ANSWER,
+
+  // Prefer maximizing the throughput. Same inference runner will be used
+  // repeatedly on different inputs.
+  SUSTAINED_SPEED,
+
+  // Balance init latency and throughput. This option will result in slightly
+  // higher init latency than FAST_SINGLE_ANSWER but should have inference
+  // latency closer to SUSTAINED_SPEED.
+  BALANCED,
+};
+
+// Defines aspects to control while instantiating a runner.
+enum class InferencePriority {
+  UNKNOWN,
+
+  AUTO,
+
+  MIN_LATENCY,
+
+  MAX_PRECISION,
+
+  MIN_MEMORY_USAGE,
+};
+
+struct InferenceOptions {
+  InferenceUsage usage = InferenceUsage::SUSTAINED_SPEED;
+
+  // Ordered priorities provide better understanding of desired semantics,
+  // where priority(n) is more important than priority(n+1).
+  // AUTO priority is needed when a single priority is the most important
+  // factor. For example, priority1 = InferencePriority::MIN_LATENCY and leaving
+  // everything else to AUTO would result in configuration that achieves maximum
+  // performance.
+  //
+  // AUTO priority can only be used when higher priorities are fully specified.
+  // For example:
+  //   VALID:   priority1 = MIN_LATENCY, priority2 = AUTO, priority3 = AUTO
+  //   VALID:   priority1 = MIN_LATENCY, priority2 = MAX_PRECISION,
+  //            priority3 = AUTO
+  //   INVALID: priority1 = AUTO, priority2 = MIN_LATENCY, priority3 = AUTO
+  //   INVALID: priority1 = MIN_LATENCY, priority2 = AUTO,
+  //            priority3 = MAX_PRECISION
+  // Invalid priorities will result in error.
+  InferencePriority priority1 = InferencePriority::MAX_PRECISION;
+
+  InferencePriority priority2 = InferencePriority::AUTO;
+
+  InferencePriority priority3 = InferencePriority::AUTO;
+#ifdef TFLITE_GPU_ENABLE_INVOKE_LOOP
+  // Number of times to invoke the inference in GPU delegate, to collect more
+  // accurate latency result. Default as 1, which is the original behavior.
+  int gpu_invoke_loop_times = 1;
+#endif
+};
+
+// Returns a position number for the priority. If priority is missing,
+// then it would return 'max num priorities + 1'.
+int GetPosition(const InferenceOptions& options, InferencePriority p);
+
+// Return true if options are valid.
+bool IsValid(const InferenceOptions& options);
+
+// Resolves AUTO priorities and specifies them explicitly.
+// Note, no-one should assume that these mappings will not change.
+// Technically this function is declared here for code re-use purposes and
+// by no means it should be treated as canonical way to resolve AUTO.
+void ResolveAutoPriority(InferenceOptions* options);
+
+enum class PriorityImportance {
+  UNKNOWN,
+  HIGHER,
+  LOWER,
+};
+
+// If both p1 and p2 are not present in options, return UNKNOWN
+// If p1 is present, but p2 is not, return HIGHER
+// If p2 is present, but p1 is not, return LOWER
+// If both are present, and p1 is more important, return HIGHER, otherwise,
+// LOWER.
+PriorityImportance GetRelativeImportance(const InferenceOptions& options,
+                                         InferencePriority p1,
+                                         InferencePriority p2);
+
+}  // namespace gpu
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_DELEGATES_GPU_API_H_

tensorflow/lite/delegates/gpu/async_buffers.cc

@@ -0,0 +1,105 @@
+/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/lite/delegates/gpu/async_buffers.h"
+
+#include <EGL/egl.h>
+#include <EGL/eglext.h>
+#include <GLES2/gl2ext.h>
+#include <GLES3/gl31.h>
+
+#include "absl/status/status.h"
+#include "tensorflow/lite/delegates/gpu/android_hardware_buffer.h"
+#include "tensorflow/lite/delegates/gpu/gl/gl_errors.h"
+
+namespace {
+PFNGLBUFFERSTORAGEEXTERNALEXTPROC glBufferStorageExternalEXT;
+PFNEGLGETNATIVECLIENTBUFFERANDROIDPROC eglGetNativeClientBufferANDROID;
+
+bool IsGlSupported() {
+  static const bool extensions_allowed = [] {
+    eglGetNativeClientBufferANDROID =
+        reinterpret_cast<PFNEGLGETNATIVECLIENTBUFFERANDROIDPROC>(
+            eglGetProcAddress("eglGetNativeClientBufferANDROID"));
+    glBufferStorageExternalEXT =
+        reinterpret_cast<PFNGLBUFFERSTORAGEEXTERNALEXTPROC>(
+            eglGetProcAddress("glBufferStorageExternalEXT"));
+    return eglGetNativeClientBufferANDROID && glBufferStorageExternalEXT;
+  }();
+  return extensions_allowed;
+}
+}  // namespace
+
+namespace tflite {
+namespace gpu {
+
+// Where the AHWB<->SSBO mapping occurs
+absl::Status AsyncBuffer::MapAHardwareBufferToGlBuffer() {
+  if (!IsGlSupported()) {
+    return absl::UnknownError(
+        "No GL extension functions found to bind AHardwareBuffer and "
+        "OpenGL buffer");
+  }
+  EGLClientBuffer native_buffer = eglGetNativeClientBufferANDROID(ahwb_);
+  if (!native_buffer) {
+    return absl::UnknownError("Can't get native buffer");
+  }
+  // If an error is traced back to below fcn, check your ahwb usage flags.
+  // An example may be found in async_buffers_test.cc
+  glBufferStorageExternalEXT(GL_SHADER_STORAGE_BUFFER, 0, bytes_, native_buffer,
+                             GL_MAP_READ_BIT | GL_MAP_WRITE_BIT |
+                                 GL_MAP_COHERENT_BIT_EXT |
+                                 GL_MAP_PERSISTENT_BIT_EXT);
+  return gl::GetOpenGlErrors();
+}
+
+// Allocate SSBO, call the AHWB<->SSBO mapping, and fail gracefully if needed.
+absl::Status AsyncBuffer::AllocateOpenGlBuffer() {
+  if (opengl_buffer_ == GL_INVALID_INDEX) {
+    // Generate and bind SSBO
+    glGenBuffers(1, &opengl_buffer_);
+    glBindBuffer(GL_SHADER_STORAGE_BUFFER, opengl_buffer_);
+    absl::Status status = MapAHardwareBufferToGlBuffer();
+    if (!status.ok()) {
+      // If we can't map to SSBO, clear AHWB & SSBO
+      if (ahwb_ != nullptr) {
+        if (OptionalAndroidHardwareBuffer::Instance().Supported()) {
+          OptionalAndroidHardwareBuffer::Instance().Release(ahwb_);
+        }
+        ahwb_ = nullptr;
+      }
+      glBufferData(GL_SHADER_STORAGE_BUFFER, bytes_, nullptr, GL_STREAM_COPY);
+    }
+    glBindBuffer(GL_SHADER_STORAGE_BUFFER, 0);
+  }
+  return absl::OkStatus();
+}
+
+// Public function which will map the AHWB (from class constructor) to a SSBO
+// and return the associated the id by reference
+absl::Status AsyncBuffer::GetOpenGlBuffer(GLuint& buffer_ref) {
+  if (!valid_) {
+    absl::Status status = AllocateOpenGlBuffer();
+    if (!status.ok()) {
+      return status;
+    }
+  }
+  valid_ = true;
+  buffer_ref = opengl_buffer_;
+  return absl::OkStatus();
+}
+
+}  // namespace gpu
+}  // namespace tflite

tensorflow/lite/delegates/gpu/async_buffers.h

@@ -0,0 +1,58 @@
+/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_LITE_DELEGATES_GPU_ASYNC_BUFFERS_H_
+#define TENSORFLOW_LITE_DELEGATES_GPU_ASYNC_BUFFERS_H_
+
+#if defined(__ANDROID__)
+#include <android/hardware_buffer.h>
+#endif  // __ANDROID__
+
+#include <GLES3/gl31.h>
+
+#include "absl/status/status.h"
+#include "tensorflow/lite/delegates/gpu/api.h"
+#include "tensorflow/lite/delegates/gpu/common/data_type.h"
+
+extern "C" typedef struct AHardwareBuffer AHardwareBuffer;
+
+namespace tflite {
+namespace gpu {
+
+class AsyncBuffer {
+ private:
+  int bytes_;                                // Number of bytes in the buffer
+  bool valid_ = false;                       // Have we mapped to SSBO already
+  GLuint opengl_buffer_ = GL_INVALID_INDEX;  // SSBO buffer id
+  AHardwareBuffer* ahwb_ = nullptr;
+
+  // Where the AHWB<->SSBO mapping occurs
+  absl::Status MapAHardwareBufferToGlBuffer();
+  // Allocate SSBO, call the AHWB<->SSBO mapping; fail gracefully if needed.
+  absl::Status AllocateOpenGlBuffer();
+
+ public:
+  explicit AsyncBuffer(TensorObjectDef tensor_def, AHardwareBuffer* ahwb) {
+    bytes_ = NumElements(tensor_def) * SizeOf(tensor_def.object_def.data_type);
+    ahwb_ = ahwb;
+  }
+  // Map the AHWB (from class constructor) to an SSBO id
+  absl::Status GetOpenGlBuffer(GLuint& buffer_ref);
+};
+
+}  // namespace gpu
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_DELEGATES_GPU_ASYNC_BUFFERS_H_

tensorflow/lite/delegates/gpu/async_buffers_test.cc

@@ -0,0 +1,74 @@
+/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/lite/delegates/gpu/async_buffers.h"
+
+#include <memory>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "tensorflow/lite/delegates/gpu/android_hardware_buffer.h"
+#include "tensorflow/lite/delegates/gpu/api.h"
+#include "tensorflow/lite/delegates/gpu/common/data_type.h"
+#include "tensorflow/lite/delegates/gpu/gl/egl_environment.h"
+
+namespace tflite {
+namespace gpu {
+namespace {
+
+TEST(AsyncBufferTest, DuplicateTest) {
+  if (__builtin_available(android 26, *)) {
+    auto Instance = OptionalAndroidHardwareBuffer::Instance;
+    // Create tie
+    TensorObjectDef* tie = new TensorObjectDef();
+    tie->object_def.data_type = DataType::FLOAT32;
+    tie->object_def.data_layout = DataLayout::BHWC;
+    tie->dimensions = Dimensions(2, 2, 2, 2);
+
+    // Create AHWB
+    AHardwareBuffer_Desc buffDesc = {};
+    buffDesc.width = 1000;
+    buffDesc.height = 1;
+    buffDesc.layers = 1;
+    buffDesc.format = AHARDWAREBUFFER_FORMAT_BLOB;
+    buffDesc.usage = AHARDWAREBUFFER_USAGE_CPU_WRITE_OFTEN |
+                     AHARDWAREBUFFER_USAGE_CPU_READ_OFTEN |
+                     AHARDWAREBUFFER_USAGE_GPU_DATA_BUFFER;
+    AHardwareBuffer* ahwb;
+    EXPECT_TRUE(Instance().IsSupported(&buffDesc));
+    EXPECT_EQ(Instance().Allocate(&buffDesc, &ahwb), 0);
+
+    // Init GL Env to properly use gl fcns
+    std::unique_ptr<gl::EglEnvironment> env;
+    EXPECT_OK(gl::EglEnvironment::NewEglEnvironment(&env));
+    AsyncBuffer async_buffer1 = AsyncBuffer(*tie, ahwb);
+    GLuint buffer1, buffer2;
+    EXPECT_OK(async_buffer1.GetOpenGlBuffer(buffer1));
+    EXPECT_GE(buffer1, 0);
+    EXPECT_OK(async_buffer1.GetOpenGlBuffer(buffer2));
+    // Check that each instance of AsyncBuffer class has only one id
+    EXPECT_EQ(buffer1, buffer2);
+    AsyncBuffer async_buffer2 = AsyncBuffer(*tie, ahwb);
+    EXPECT_OK(async_buffer2.GetOpenGlBuffer(buffer2));
+    // Check that each different instance will produce unique id
+    EXPECT_NE(buffer1, buffer2);
+  } else {
+    GTEST_SKIP();
+  }
+}
+
+}  // namespace
+}  // namespace gpu
+}  // namespace tflite

tensorflow/lite/delegates/gpu/build_defs.bzl

@@ -0,0 +1,38 @@
+"""Additional build options needed for the GPU Delegate."""
+
+def gpu_delegate_linkopts():
+    """Additional link options needed when linking in the GPU Delegate."""
+    return select({
+        "//tensorflow:android": [
+            "-lEGL",
+            # We don't need to link libGLESv3, because if it exists,
+            # it is a symlink to libGLESv2.
+            # See Compatibility Definition Document:
+            # https://source.android.com/compatibility/10/android-10-cdd#7_1_4_1_opengl_es
+            "-lGLESv2",
+        ],
+        "//conditions:default": [],
+    })
+
+def tflite_angle_heapcheck_deps():
+    # copybara:uncomment_begin(google-only)
+    # return select({
+    # "//tensorflow/lite/delegates/gpu:tflite_gpu_angle": [
+    # "@com_google_googletest//:gtest_main_no_heapcheck",
+    # ],
+    # "//conditions:default": [
+    # "@com_google_googletest//:gtest_main",
+    # ],
+    # })
+    # copybara:uncomment_end
+    # copybara:comment_begin(oss-only)
+    return ["@com_google_googletest//:gtest_main"]
+    # copybara:comment_end
+
+def gtest_main_no_heapcheck_deps():
+    # copybara:uncomment_begin(google-only)
+    # return ["@com_google_googletest//:gtest_main_no_heapcheck"]
+    # copybara:uncomment_end
+    # copybara:comment_begin(oss-only)
+    return ["@com_google_googletest//:gtest_main"]
+    # copybara:comment_end

tensorflow/lite/delegates/gpu/cl/BUILD

@@ -0,0 +1,671 @@
+load("@flatbuffers//:build_defs.bzl", "flatbuffer_cc_library")
+load("@rules_cc//cc:cc_library.bzl", "cc_library")
+load("@rules_cc//cc:cc_test.bzl", "cc_test")
+load("//tensorflow:tensorflow.bzl", "if_google", "workspace_root")
+load(
+    "//tensorflow/core/platform:build_config_root.bzl",
+    "tf_gpu_tests_tags",
+)
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:LICENSE"],
+    default_visibility = ["//visibility:public"],
+    licenses = ["notice"],
+)
+
+_API_NO_GL_DEPS = [
+    ":cl_command_queue",
+    ":cl_errors",
+    ":cl_event",
+    ":environment",
+    ":inference_context",
+    ":opencl_wrapper",
+    ":tensor",
+    ":tensor_type_util",
+    "@com_google_absl//absl/memory",
+    "@com_google_absl//absl/types:span",
+    "//tensorflow/lite/delegates/gpu:api",
+    "//tensorflow/lite/delegates/gpu:tflite_profile",
+    "//tensorflow/lite/delegates/gpu/cl/kernels:converter",
+    "//tensorflow/lite/delegates/gpu/common:data_type",
+    "//tensorflow/lite/delegates/gpu/common:model",
+    "//tensorflow/lite/delegates/gpu/common:precision",
+    "//tensorflow/lite/delegates/gpu/common:shape",
+    "//tensorflow/lite/delegates/gpu/common:status",
+    "//tensorflow/lite/delegates/gpu/common:tensor",
+    "//tensorflow/lite/delegates/gpu/common/task:tensor_desc",
+]
+
+_GPU_API_DELEGATE_NO_GL_DEPS = [
+    ":opencl_wrapper",
+    ":tensor_type_util",
+    "@com_google_absl//absl/types:span",
+    "//tensorflow/lite:kernel_api",
+    "//tensorflow/lite/core/c:common",
+    "//tensorflow/lite/delegates/gpu:api",
+    "//tensorflow/lite/delegates/gpu/common:model",
+    "//tensorflow/lite/delegates/gpu/common:model_builder",
+    "//tensorflow/lite/delegates/gpu/common:model_transformer",
+    "//tensorflow/lite/delegates/gpu/common:status",
+    "//tensorflow/lite/delegates/gpu/common/transformations:model_transformations",
+]
+
+config_setting(
+    name = "opencl_delegate_no_gl",
+    values = {"copt": "-DCL_DELEGATE_NO_GL"},
+)
+
+# copybara:uncomment_begin(google-only)
+# cc_library(
+#     name = "api_no_gl",
+#     srcs = ["api.cc"],
+#     hdrs = ["api.h"],
+#     defines = ["CL_DELEGATE_NO_GL"],
+#     deps = _API_NO_GL_DEPS + ["//third_party/GL:EGL_headers"],
+# )
+# copybara:uncomment_end
+
+cc_library(
+    name = "api",
+    srcs = ["api.cc"],
+    hdrs = ["api.h"],
+    deps = select({
+        ":opencl_delegate_no_gl": [],
+        "//conditions:default": [
+            ":egl_sync",
+            ":gl_interop",
+        ],
+    }) + _API_NO_GL_DEPS,
+)
+
+cc_library(
+    name = "buffer",
+    srcs = ["buffer.cc"],
+    hdrs = ["buffer.h"],
+    deps = [
+        ":cl_command_queue",
+        ":cl_context",
+        ":gpu_object",
+        ":opencl_wrapper",
+        ":util",
+        "//tensorflow/lite/delegates/gpu/common:data_type",
+        "//tensorflow/lite/delegates/gpu/common:status",
+        "//tensorflow/lite/delegates/gpu/common/task:buffer_desc",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/types:span",
+    ],
+)
+
+cc_test(
+    name = "buffer_test",
+    srcs = ["buffer_test.cc"],
+    linkstatic = True,
+    tags = tf_gpu_tests_tags() + [
+        "linux",
+        "local",
+    ],
+    deps = [
+        ":buffer",
+        ":cl_test",
+        "//tensorflow/lite/delegates/gpu/common:status",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
+
+cc_library(
+    name = "cl_test",
+    testonly = 1,
+    hdrs = ["cl_test.h"],
+    deps = [
+        ":environment",
+        ":opencl_wrapper",
+        "@com_google_googletest//:gtest",
+    ],
+)
+
+cc_library(
+    name = "cl_arguments",
+    srcs = ["cl_arguments.cc"],
+    hdrs = ["cl_arguments.h"],
+    deps = [
+        ":buffer",
+        ":cl_context",
+        ":gpu_object",
+        ":qcom_thin_filter",
+        ":tensor",
+        "//tensorflow/lite/delegates/gpu/common:gpu_info",
+        "//tensorflow/lite/delegates/gpu/common:status",
+        "//tensorflow/lite/delegates/gpu/common:util",
+        "//tensorflow/lite/delegates/gpu/common/task:arguments",
+        "//tensorflow/lite/delegates/gpu/common/task:util",
+        "@com_google_absl//absl/strings",
+    ],
+)
+
+cc_test(
+    name = "cl_arguments_test",
+    srcs = ["cl_arguments_test.cc"],
+    linkstatic = True,
+    tags = tf_gpu_tests_tags() + [
+        "linux",
+        "local",
+    ],
+    deps = [
+        ":buffer",
+        ":cl_arguments",
+        ":cl_test",
+        ":gpu_object",
+        "//tensorflow/lite/delegates/gpu/common:gpu_info",
+        "@com_google_absl//absl/strings",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
+
+cc_library(
+    name = "cl_command_buffer",
+    srcs = ["cl_command_buffer.cc"],
+    hdrs = ["cl_command_buffer.h"],
+    deps = [
+        ":cl_command_queue",
+        ":cl_event",
+        ":opencl_wrapper",
+        ":util",
+        "//tensorflow/lite/delegates/gpu/common:status",
+        "@com_google_absl//absl/strings",
+    ],
+)
+
+cc_library(
+    name = "cl_command_queue",
+    srcs = ["cl_command_queue.cc"],
+    hdrs = ["cl_command_queue.h"],
+    deps = [
+        ":cl_context",
+        ":cl_device",
+        ":cl_event",
+        ":cl_kernel",
+        ":opencl_wrapper",
+        ":util",
+        "//tensorflow/lite/delegates/gpu/common:status",
+        "//tensorflow/lite/delegates/gpu/common:types",
+        "//tensorflow/lite/delegates/gpu/common/task:profiling_info",
+        "@com_google_absl//absl/strings",
+    ],
+)
+
+cc_library(
+    name = "cl_context",
+    srcs = ["cl_context.cc"],
+    hdrs = ["cl_context.h"],
+    deps = [
+        ":cl_device",
+        ":cl_image_format",
+        ":opencl_wrapper",
+        ":util",
+        "//tensorflow/lite/delegates/gpu/common:data_type",
+        "//tensorflow/lite/delegates/gpu/common:status",
+        "@com_google_absl//absl/strings",
+    ],
+)
+
+cc_library(
+    name = "cl_device",
+    srcs = ["cl_device.cc"],
+    hdrs = ["cl_device.h"],
+    deps = [
+        ":opencl_wrapper",
+        ":util",
+        "//tensorflow/lite/delegates/gpu/common:gpu_info",
+        "//tensorflow/lite/delegates/gpu/common:status",
+        "//tensorflow/lite/delegates/gpu/common:types",
+        "//tensorflow/lite/experimental/acceleration/compatibility:android_info",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/strings:str_format",
+    ],
+)
+
+cc_test(
+    name = "cl_device_test",
+    srcs = ["cl_device_test.cc"],
+    deps = [
+        ":cl_device",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
+
+cc_library(
+    name = "cl_errors",
+    hdrs = ["cl_errors.h"],
+    deps = [
+        ":util",
+        "//tensorflow/lite/delegates/gpu/common:status",
+    ],
+)
+
+cc_library(
+    name = "cl_event",
+    srcs = ["cl_event.cc"],
+    hdrs = ["cl_event.h"],
+    deps = [
+        ":opencl_wrapper",
+    ],
+)
+
+cc_library(
+    name = "cl_image_format",
+    srcs = ["cl_image_format.cc"],
+    hdrs = ["cl_image_format.h"],
+    deps = [
+        ":opencl_wrapper",
+        "//tensorflow/lite/delegates/gpu/common:data_type",
+    ],
+)
+
+cc_library(
+    name = "cl_kernel",
+    srcs = ["cl_kernel.cc"],
+    hdrs = ["cl_kernel.h"],
+    deps = [
+        ":cl_context",
+        ":cl_device",
+        ":cl_program",
+        ":opencl_wrapper",
+        ":util",
+        "//tensorflow/lite/delegates/gpu/common:kernel_info",
+        "//tensorflow/lite/delegates/gpu/common:status",
+        "@com_google_absl//absl/strings",
+    ],
+)
+
+cc_library(
+    name = "cl_memory",
+    srcs = ["cl_memory.cc"],
+    hdrs = ["cl_memory.h"],
+    deps = [
+        ":opencl_wrapper",
+        "//tensorflow/lite/delegates/gpu/common:access_type",
+        "//tensorflow/lite/delegates/gpu/common:status",
+    ],
+)
+
+cc_library(
+    name = "cl_operation",
+    srcs = ["cl_operation.cc"],
+    hdrs = ["cl_operation.h"],
+    deps = [
+        ":cl_arguments",
+        ":cl_command_queue",
+        ":cl_context",
+        ":cl_device",
+        ":cl_kernel",
+        ":program_cache",
+        ":tensor",
+        "//tensorflow/lite/delegates/gpu/common/task:compiler_options",
+        "//tensorflow/lite/delegates/gpu/common/task:gpu_operation",
+        "@com_google_absl//absl/strings",
+    ],
+)
+
+cc_library(
+    name = "cl_program",
+    srcs = ["cl_program.cc"],
+    hdrs = ["cl_program.h"],
+    deps = [
+        ":cl_context",
+        ":cl_device",
+        ":opencl_wrapper",
+        ":util",
+        "//tensorflow/lite/delegates/gpu/common:status",
+        "//tensorflow/lite/delegates/gpu/common/task:compiler_options",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/types:span",
+    ],
+)
+
+flatbuffer_cc_library(
+    name = "compiled_program_cache_cc_fbs",
+    srcs = ["compiled_program_cache.fbs"],
+    flatc_args = [
+        "--scoped-enums",
+    ],
+)
+
+cc_library(
+    name = "egl_sync",
+    srcs = ["egl_sync.cc"],
+    hdrs = ["egl_sync.h"],
+    defines = [
+        "EGL_EGLEXT_PROTOTYPES",
+    ],
+    deps = [
+        "//tensorflow/lite/delegates/gpu/common:status",
+        "//tensorflow/lite/delegates/gpu/gl:gl_call",
+    ],
+)
+
+cc_library(
+    name = "environment",
+    srcs = ["environment.cc"],
+    hdrs = ["environment.h"],
+    deps = [
+        ":cl_command_queue",
+        ":cl_context",
+        ":cl_device",
+        ":program_cache",
+        "//tensorflow/lite/delegates/gpu/common:data_type",
+        "//tensorflow/lite/delegates/gpu/common:gpu_info",
+        "//tensorflow/lite/delegates/gpu/common:precision",
+        "//tensorflow/lite/delegates/gpu/common:status",
+        "//tensorflow/lite/delegates/gpu/common:tensor",
+        "//tensorflow/lite/delegates/gpu/common/task:tensor_desc",
+    ],
+)
+
+cc_library(
+    name = "gl_interop",
+    srcs = ["gl_interop.cc"],
+    hdrs = ["gl_interop.h"],
+    deps = [
+        ":cl_command_queue",
+        ":cl_context",
+        ":cl_device",
+        ":cl_errors",
+        ":cl_event",
+        ":cl_memory",
+        ":egl_sync",
+        ":environment",
+        ":opencl_wrapper",
+        "//tensorflow/lite/delegates/gpu:spi",
+        "//tensorflow/lite/delegates/gpu/common:access_type",
+        "//tensorflow/lite/delegates/gpu/common:status",
+        "//tensorflow/lite/delegates/gpu/gl:gl_call",
+        "//tensorflow/lite/delegates/gpu/gl:gl_sync",
+        "//tensorflow/lite/delegates/gpu/gl:portable",
+        "@com_google_absl//absl/strings",
+    ],
+)
+
+# copybara:uncomment_begin(google-only)
+# cc_library(
+#     name = "gl_interop_no_gl",
+#     srcs = ["gl_interop.cc"],
+#     hdrs = ["gl_interop.h"],
+#     defines = ["CL_DELEGATE_NO_GL"],
+#     deps = [
+#         ":cl_command_queue",
+#         ":cl_device",
+#         ":environment",
+#         "//third_party/GL:EGL_headers",
+#         "//tensorflow/lite/delegates/gpu:spi",
+#         "//tensorflow/lite/delegates/gpu/common:access_type",
+#         "//tensorflow/lite/delegates/gpu/common:status",
+#     ],
+# )
+#
+# cc_test(
+#     name = "gl_interop_no_gl_test",
+#     srcs = ["gl_interop_no_gl_test.cc"],
+#     tags = ["requires-gpu-nvidia"],
+#     deps = [
+#         ":gl_interop_no_gl",
+#         "@com_google_googletest//:gtest_main",
+#     ],
+# )
+# copybara:uncomment_end
+
+# copybara:uncomment_begin(google-only)
+# cc_library(
+#     name = "gpu_api_delegate_no_gl",
+#     srcs = ["gpu_api_delegate.cc"],
+#     hdrs = ["gpu_api_delegate.h"],
+#     defines = ["CL_DELEGATE_NO_GL"],
+#     linkopts = select({
+#         "//tensorflow:android": [
+#             "-lEGL",
+#             "-lGLESv3",
+#         ],
+#         "//conditions:default": [],
+#     }),
+#     deps = _GPU_API_DELEGATE_NO_GL_DEPS + [
+#         ":api_no_gl",
+#         "//third_party/GL:EGL_headers",
+#         "//third_party/GL:GLES3_headers",
+#         "//tensorflow/lite/delegates/gpu:delegate_no_gl",
+#     ],
+# )
+# copybara:uncomment_end
+
+cc_library(
+    name = "gpu_api_delegate",
+    srcs = ["gpu_api_delegate.cc"],
+    hdrs = ["gpu_api_delegate.h"],
+    linkopts = select({
+        "//tensorflow:android": [
+            "-lEGL",
+            "-lGLESv3",
+        ],
+        "//conditions:default": [],
+    }),
+    deps = _GPU_API_DELEGATE_NO_GL_DEPS + [
+        ":api",
+        "//tensorflow/lite/delegates/gpu:delegate",
+    ],
+)
+
+cc_library(
+    name = "gpu_object",
+    hdrs = ["gpu_object.h"],
+    deps = [
+        ":opencl_wrapper",
+        "//tensorflow/lite/delegates/gpu/common:access_type",
+        "//tensorflow/lite/delegates/gpu/common:data_type",
+        "//tensorflow/lite/delegates/gpu/common:status",
+        "//tensorflow/lite/delegates/gpu/common/task:gpu_object_desc",
+    ],
+)
+
+cc_library(
+    name = "inference_context",
+    srcs = [
+        "inference_context.cc",
+    ],
+    hdrs = [
+        "inference_context.h",
+    ],
+    deps = [
+        ":buffer",
+        ":cl_command_buffer",
+        ":cl_command_queue",
+        ":cl_device",
+        ":cl_event",
+        ":cl_operation",
+        ":environment",
+        ":gpu_object",
+        ":opencl_wrapper",
+        ":recordable_queue_builder",
+        ":serialization_cc_fbs",
+        ":tensor",
+        "//tensorflow/lite/delegates/gpu/common:data_type",
+        "//tensorflow/lite/delegates/gpu/common:gpu_model",
+        "//tensorflow/lite/delegates/gpu/common:gpu_model_cc_fbs",
+        "//tensorflow/lite/delegates/gpu/common:memory_management",
+        "//tensorflow/lite/delegates/gpu/common:model",
+        "//tensorflow/lite/delegates/gpu/common:model_hints",
+        "//tensorflow/lite/delegates/gpu/common:precision",
+        "//tensorflow/lite/delegates/gpu/common:shape",
+        "//tensorflow/lite/delegates/gpu/common:status",
+        "//tensorflow/lite/delegates/gpu/common:tensor",
+        "//tensorflow/lite/delegates/gpu/common:types",
+        "//tensorflow/lite/delegates/gpu/common:util",
+        "//tensorflow/lite/delegates/gpu/common/task:gpu_operation",
+        "//tensorflow/lite/delegates/gpu/common/task:serialization_base",
+        "//tensorflow/lite/delegates/gpu/common/task:tensor_desc",
+        "@com_google_absl//absl/container:flat_hash_map",
+        "@com_google_absl//absl/container:flat_hash_set",
+    ],
+)
+
+cc_library(
+    name = "opencl_wrapper",
+    srcs = ["opencl_wrapper.cc"],
+    hdrs = ["opencl_wrapper.h"],
+    linkopts = select({
+        "//tensorflow:android": ["-lm"],
+        "//tensorflow:windows": [],  # Windows does not support dlopen().
+        "//conditions:default": ["-ldl"],  # opencl_wrapper calls dlopen()
+    }),
+    local_defines = select({
+        # copybara:uncomment_begin(google-only)
+        # "//tools/cc_target_os:linux-google": ["__LINUX_GOOGLE__"],
+        # copybara:uncomment_end
+        "//conditions:default": [],
+    }),
+    deps = [
+        "//tensorflow/lite/delegates/gpu/common:status",
+        "@com_google_absl//absl/strings",
+        "@opencl_headers",
+        "//tensorflow/lite/delegates/gpu/cl/" + if_google("google", "default") + ":qcom_wrapper",
+        "//tensorflow/lite/tools:logging",
+    ] + select({
+        "//conditions:default": [],
+    }),
+)
+
+cc_library(
+    name = "program_cache",
+    srcs = ["program_cache.cc"],
+    hdrs = ["program_cache.h"],
+    deps = [
+        ":cl_context",
+        ":cl_device",
+        ":cl_kernel",
+        ":cl_program",
+        ":compiled_program_cache_cc_fbs",
+        ":util",
+        "//tensorflow/lite/delegates/gpu/common:status",
+        "@com_google_absl//absl/container:flat_hash_map",
+        "@com_google_absl//absl/types:span",
+        "@farmhash_archive//:farmhash",
+        "@flatbuffers",
+    ],
+)
+
+cc_library(
+    name = "qcom_thin_filter",
+    srcs = ["qcom_thin_filter.cc"],
+    hdrs = ["qcom_thin_filter.h"],
+    deps = [
+        ":cl_context",
+        ":gpu_object",
+        ":util",
+        "//tensorflow/lite/delegates/gpu/common/task:qcom_thin_filter_desc",
+    ],
+)
+
+cc_library(
+    name = "recordable_queue",
+    hdrs = ["recordable_queue.h"],
+    deps = [
+        ":cl_command_queue",
+        ":cl_context",
+        ":cl_device",
+        ":cl_operation",
+        ":opencl_wrapper",
+        "//tensorflow/lite/delegates/gpu/common:status",
+    ],
+)
+
+cc_library(
+    name = "recordable_queue_builder",
+    hdrs = ["recordable_queue_builder.h"],
+    deps = [
+        ":cl_context",
+        ":cl_device",
+        ":cl_operation",
+        ":recordable_queue",
+        "//tensorflow/lite/delegates/gpu/cl/" + if_google("google", "default") + ":recordable_queue",
+    ],
+)
+
+flatbuffer_cc_library(
+    name = "serialization_cc_fbs",
+    srcs = ["serialization.fbs"],
+    flatc_args = [
+        "--scoped-enums",
+        "-I " + workspace_root,
+    ],
+    includes = [
+        "//tensorflow/lite/delegates/gpu/common:gpu_model_cc_fbs_includes",
+        "//tensorflow/lite/delegates/gpu/common/task:serialization_base_cc_fbs_includes",
+    ],
+)
+
+cc_library(
+    name = "tensor",
+    srcs = ["tensor.cc"],
+    hdrs = ["tensor.h"],
+    deps = [
+        ":buffer",
+        ":cl_command_queue",
+        ":cl_context",
+        ":cl_device",
+        ":cl_image_format",
+        ":cl_memory",
+        ":gpu_object",
+        ":util",
+        "//tensorflow/lite/delegates/gpu/common:data_type",
+        "//tensorflow/lite/delegates/gpu/common:shape",
+        "//tensorflow/lite/delegates/gpu/common:status",
+        "//tensorflow/lite/delegates/gpu/common:tensor",
+        "//tensorflow/lite/delegates/gpu/common:types",
+        "//tensorflow/lite/delegates/gpu/common/task:gpu_tensor",
+        "//tensorflow/lite/delegates/gpu/common/task:tensor_desc",
+        "@com_google_absl//absl/strings",
+    ],
+)
+
+cc_test(
+    name = "tensor_test",
+    srcs = ["tensor_test.cc"],
+    linkstatic = True,
+    tags = tf_gpu_tests_tags() + [
+        "linux",
+        "local",
+    ],
+    deps = [
+        ":cl_test",
+        ":tensor",
+        "//tensorflow/lite/delegates/gpu/common:data_type",
+        "//tensorflow/lite/delegates/gpu/common:shape",
+        "//tensorflow/lite/delegates/gpu/common:status",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
+
+cc_library(
+    name = "tensor_type_util",
+    srcs = ["tensor_type_util.cc"],
+    hdrs = ["tensor_type_util.h"],
+    deps = [
+        "//tensorflow/lite/delegates/gpu:api",
+        "//tensorflow/lite/delegates/gpu/common/task:tensor_desc",
+    ],
+)
+
+cc_library(
+    name = "util",
+    srcs = ["util.cc"],
+    hdrs = ["util.h"],
+    deps = [
+        ":opencl_wrapper",
+        "//tensorflow/lite/delegates/gpu/common:data_type",
+        "//tensorflow/lite/delegates/gpu/common:status",
+        "//tensorflow/lite/delegates/gpu/common:tensor",
+        "//tensorflow/lite/delegates/gpu/common:util",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/types:span",
+        "//tensorflow/lite/delegates/gpu/cl/" + if_google("google", "default") + ":util",
+    ],
+)

tensorflow/lite/delegates/gpu/cl/api.h

@@ -0,0 +1,175 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_LITE_DELEGATES_GPU_CL_API_H_
+#define TENSORFLOW_LITE_DELEGATES_GPU_CL_API_H_
+
+#ifdef CL_DELEGATE_NO_GL
+#define EGL_NO_PROTOTYPES
+#endif
+
+#include <EGL/egl.h>
+
+#include <cstdint>
+#include <memory>
+
+#include "absl/types/span.h"
+#include "tensorflow/lite/delegates/gpu/api.h"
+#include "tensorflow/lite/delegates/gpu/common/model.h"
+#include "tensorflow/lite/delegates/gpu/common/status.h"
+
+// Usage example:
+//
+//   std::unique_ptr<InferenceEnvironment> env;
+//   RETURN_IF_ERROR(NewInferenceEnvironment(option, &env));
+//
+//   InferenceOptions options;
+//
+//   std::unique_ptr<InferenceBuilder> builder;
+//   RETURN_IF_ERROR(env->NewInferenceBuilder(options, model, &builder));
+//   // now builder is ready to prepare inference runner.
+//
+// -----------------
+// Supported formats
+// -----------------
+//
+// OpenCL implementation uses 2D textures as the primary format.
+// Tensor in HWDC4 layout is {TEXTURE_2D, RGBA, width := W*D, height := H}.
+//
+
+namespace tflite {
+namespace gpu {
+namespace cl {
+
+struct InferenceOptions : public tflite::gpu::InferenceOptions {};
+
+// Indicates environment
+struct InferenceEnvironmentProperties {
+  bool is_opencl_available = false;
+
+  // GL objects (buffers and textures) could be shared with CL context.
+  bool is_gl_sharing_supported = false;
+
+  // Indicates whether fast GL->CL synchronization is supported.
+  bool is_gl_to_cl_fast_sync_supported = false;
+
+  // Indicates whether fast CL->GL synchronization is supported.
+  bool is_cl_to_gl_fast_sync_supported = false;
+};
+
+// Environment manages all resources that need to stay until any inference is
+// running using OpenCL backend.
+class InferenceEnvironment {
+ public:
+  virtual ~InferenceEnvironment() {}
+
+  // Converts GraphFloat32 into intermediate, device-specific representation.
+  // This serialized_model specific for device and InferenceOptions.
+  // serialized_model cannot be used with another device or InferenceOptions.
+  // Loading serialized_model is much faster than loading GraphFloat32.
+  // serialized_model must be used with appropriate NewInferenceBuilder
+  // method (see below).
+  // Normally BuildSerializedModel method need to be called whenever a model or
+  // OS GPU driver is updated.
+  virtual absl::Status BuildSerializedModel(
+      const InferenceOptions& options, GraphFloat32 model,
+      std::vector<uint8_t>* serialized_model) = 0;
+
+  // Serialized model can became invalid when environment changes. In this case
+  // this call will fail and model must be regenerated(with
+  // BuildSerializedModel).
+  virtual absl::Status NewInferenceBuilder(
+      const absl::Span<const uint8_t> serialized_model,
+      std::unique_ptr<InferenceBuilder>* builder) = 0;
+
+  virtual absl::Status NewInferenceBuilder(
+      const InferenceOptions& options, GraphFloat32 model,
+      std::unique_ptr<InferenceBuilder>* builder) = 0;
+
+  // Returns opaque binary blob that contains a collection of already compiled
+  // OpenCL kernels present in a cache. Returned data could be re-used later
+  // to speed up compilation time when new environment is created for the same
+  // set of models.
+  // Returned data is valid only if used on the same device, otherwise it will
+  // not be compatible and will be discarded.
+  virtual std::vector<uint8_t> GetSerializedBinaryCache() const = 0;
+};
+
+struct InferenceEnvironmentOptions {
+  // If any of these objects are set, created environment will use them instead
+  // of creating/choosing own instances.
+  cl_device_id device = nullptr;
+  cl_context context = nullptr;
+  cl_command_queue command_queue = nullptr;
+
+  // Whenever input and/or output is GL object, EGL display and context must be
+  // set to create GL aware OpenCL context. Do not set these variables whenever
+  // GL interoperability is not needed.
+  // It is the error to set egl_display, egl_context AND context at the same
+  // time. If egl_display and egl_context are set, they will be used to create
+  // GL-aware CL context.
+  EGLDisplay egl_display = EGL_NO_DISPLAY;
+  EGLContext egl_context = EGL_NO_CONTEXT;
+
+  // Should contain data returned from
+  // InferenceEnvironment::GetSerializedBinaryCache method.
+  // Invalid or incompatible data will be discarded. Compiled binary may become
+  // incompatible when GPU driver is updated.
+  absl::Span<const uint8_t> serialized_binary_cache;
+
+  bool IsGlAware() const {
+    return egl_context != EGL_NO_CONTEXT && egl_display != EGL_NO_DISPLAY;
+  }
+};
+
+// Creates new OpenCL environment that needs to stay around until all inference
+// runners are destroyed.
+absl::Status NewInferenceEnvironment(
+    const InferenceEnvironmentOptions& options,
+    std::unique_ptr<InferenceEnvironment>* environment,
+    InferenceEnvironmentProperties* properties /* optional */);
+
+class CLInferenceRunner : public ::tflite::gpu::InferenceRunner {
+ public:
+  // The RunWithoutExternalBufferCopy provides a contract where the user of this
+  // interface does not need
+  //    a. Inputs to be copied to the internal GPU buffer from the external CPU
+  //       input buffer
+  //    b. Outputs to be copied from the internal GPU buffer to the
+  //       external CPU buffer
+  //
+  // The user of this interface is responsible for copying the inputs prior to
+  // running the GPU kernels and outputs post running with the other interfaces
+  // provided here.
+  virtual absl::Status RunWithoutExternalBufferCopy() = 0;
+
+  // Copies from the external input tensor (normally CPU buffer) to the internal
+  // OpenCL buffer.  The call only guarantees a queueing of the command. The
+  // caller is expected to hold a copy of the queue and wait for completion if
+  // the external buffer is a CPU buffer.
+  virtual absl::Status CopyFromExternalInput(int index) = 0;
+
+  // Copies from the internal output OpenCL buffer to the external output
+  // tensor.  The call only guarantees a queueing of the command. The caller
+  // is expected to hold a copy of the queue and wait for completion if the
+  // external buffer is a CPU buffer.
+  virtual absl::Status CopyToExternalOutput(int index) = 0;
+};
+
+}  // namespace cl
+}  // namespace gpu
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_DELEGATES_GPU_CL_API_H_