wehub/tensorflow--tensorflow
1
0
Fork 0
Code Issues Pull Requests Actions 5 Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution
cffconvert / validate (push) Has been skipped
Details
License Check / license-check (push) Failing after 2s
Details

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 12:14:16 +08:00
commit 8a852e4b4e
36502 changed files with 9277225 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tensorflow/lite/delegates/flex/BUILD
+438
View File
@@ -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
+67
View File
@@ -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
+57
View File
@@ -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
+72
View File
@@ -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
+318
View File
@@ -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
+275
View File
@@ -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
+110
View File
@@ -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
+367
View File
@@ -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
+188
View File
@@ -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
+125
View File
@@ -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
+241
View File
@@ -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
+107
View File
@@ -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
+243
View File
@@ -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
+42
View File
@@ -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
+491
View File
@@ -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
+1
View File
@@ -0,0 +1 @@
*AcquireFlexDelegate*
tensorflow/lite/delegates/flex/java/src/main/java/org/tensorflow/lite/flex/BUILD
+9
View File
@@ -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
+69
View File
@@ -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
+35
View File
@@ -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
+43
View File
@@ -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
+857
View File
@@ -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
+72
View File
@@ -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
+516
View File
@@ -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
+65
View File
@@ -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
+115
View File
@@ -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
+216
View File
@@ -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
+135
View File
@@ -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
+413
View File
@@ -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
+310
View File
@@ -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
+82
View File
@@ -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
+287
View File
@@ -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
+6
View File
@@ -0,0 +1,6 @@
VERS_1.0 {
global:
*AcquireFlexDelegate*;
local:
*;
};