wehub/tensorflow--tensorflow
1
0
Fork 0
Code Issues Pull Requests Actions 3 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/compiler/mlir/BUILD
+277
View File
@@ -0,0 +1,277 @@
# Description:
# TensorFlow/TensorFlow Lite/XLA MLIR dialects and tools.
load(
"//tensorflow:tensorflow.bzl",
"tf_cc_binary",
"tf_cc_test",
)
load("//tensorflow:tensorflow.default.bzl", "filegroup")
load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
default_visibility = ["//visibility:public"],
licenses = ["notice"],
)
package_group(
name = "subpackages",
packages = ["//tensorflow/compiler/mlir/..."],
)
exports_files(glob(["g3doc/*.md"] + ["g3doc/_includes/tf_passes.md"]))
# To reference all tablegen files here when checking for updates to them.
filegroup(
name = "td_files",
srcs = glob(["**/*.td"]),
)
cc_library(
name = "op_or_arg_name_mapper",
srcs = ["op_or_arg_name_mapper.cc"],
hdrs = ["op_or_arg_name_mapper.h"],
deps = [
"//tensorflow/compiler/mlir/utils:name_utils",
"@com_google_absl//absl/strings",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Support",
],
)
tf_cc_test(
name = "op_or_arg_name_mapper_test",
srcs = ["op_or_arg_name_mapper_test.cc"],
deps = [
":op_or_arg_name_mapper",
"@com_google_absl//absl/strings",
"@com_google_googletest//:gtest_main",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Support",
],
)
cc_library(
name = "tf_mlir_opt_main",
testonly = True,
srcs = ["tf_mlir_opt_main.cc"],
deps = [
":init_mlir",
":passes",
":register_common_dialects",
"//tensorflow/compiler/mlir/quantization/stablehlo:bridge_passes",
"//tensorflow/compiler/mlir/tensorflow:mlprogram_util",
"//tensorflow/compiler/mlir/tensorflow/transforms:tensorflow_passes",
"//tensorflow/compiler/mlir/tensorflow/transforms:tensorflow_test_passes",
"//tensorflow/compiler/mlir/tensorflow/transforms:tf_graph_optimization_pass",
"//tensorflow/compiler/mlir/tensorflow/transforms:tf_saved_model_passes", # buildcleaner: keep
"//tensorflow/compiler/mlir/tensorflow/transforms/host_runtime:lower_cluster_to_runtime_ops",
"//tensorflow/compiler/mlir/tensorflow/transforms/host_runtime:runtime_passes",
"//tensorflow/compiler/mlir/tensorflow/transforms/sparsecore:sparsecore_passes",
"//tensorflow/compiler/mlir/tf2xla:compile_mlir_util",
"//tensorflow/compiler/mlir/tf2xla/internal/passes:clustering_passes",
"//tensorflow/compiler/mlir/tf2xla/internal/passes:mlir_to_graph_passes",
"//tensorflow/compiler/mlir/tf2xla/transforms:tf_xla_passes",
"//tensorflow/compiler/mlir/tf2xla/transforms:xla_legalize_tf",
"@llvm-project//mlir:AllPassesAndDialects",
"@llvm-project//mlir:MlirOptLib",
"@llvm-project//mlir:RegisterAllDialects", # buildcleaner: keep
"@llvm-project//mlir:RegisterAllPasses", # buildcleaner: keep
"@llvm-project//mlir:Support",
"@llvm-project//mlir:Transforms",
"@xla//xla/mlir_hlo:all_passes",
],
)
cc_library(
name = "passes",
visibility = ["//visibility:public"],
deps = [
"@llvm-project//mlir:AffineDialect",
"@llvm-project//mlir:QuantOps",
# Link jit lib to link JIT devices required to run
# xla-legalize-tf-with-tf2xla pass.
"//tensorflow/compiler/jit",
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/compiler/mlir/tensorflow/transforms:tensorflow_passes",
"//tensorflow/compiler/mlir/tensorflow/transforms:tf_dialect_passes",
"//tensorflow/compiler/mlir/tf2xla:compile_mlir_util",
],
)
cc_library(
name = "init_mlir",
srcs = ["init_mlir.cc"],
hdrs = ["init_mlir.h"],
visibility = ["//visibility:public"],
deps = [
"//tensorflow/core:lib",
"@llvm-project//llvm:Support",
],
)
cc_library(
name = "mlir_graph_optimization_pass",
srcs = ["mlir_graph_optimization_pass.cc"],
hdrs = ["mlir_graph_optimization_pass.h"],
deps = [
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/compiler/mlir/tensorflow:attribute_utils",
"//tensorflow/compiler/mlir/tensorflow:device_util",
"//tensorflow/compiler/mlir/tensorflow:dump_mlir_util",
"//tensorflow/compiler/mlir/tensorflow:mlir_roundtrip_flags",
"//tensorflow/compiler/mlir/tf2xla:mlir_bridge_rollout_policy",
"//tensorflow/compiler/mlir/tf2xla/api/v2:graph_to_tf_executor",
"//tensorflow/compiler/mlir/tf2xla/api/v2:tf_executor_to_graph",
"//tensorflow/core:core_cpu",
"//tensorflow/core:framework",
"//tensorflow/core:lib",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core/common_runtime:device_set",
"//tensorflow/core/protobuf:for_core_protos_cc",
"@com_google_absl//absl/container:flat_hash_set",
"@com_google_absl//absl/log",
"@com_google_absl//absl/log:check",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings:string_view",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:ArithDialect",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:FuncExtensions",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:ShapeDialect",
],
alwayslink = 1,
)
cc_library(
name = "mlir_graph_optimization_pass_registration",
srcs = [
"mlir_graph_optimization_pass_registration.cc",
],
deps = [
":mlir_graph_optimization_pass",
"//tensorflow/core:core_cpu",
],
alwayslink = 1,
)
# This should just be a wrapper around tf_mlir_opt_main. Don't add
# direct dependencies to this binary.
tf_cc_binary(
name = "tf-opt",
testonly = True,
deps = [
":tf_mlir_opt_main",
],
)
tf_cc_binary(
name = "tf-reduce",
testonly = True,
srcs = ["tf_mlir_reduce_main.cc"],
deps = [
":init_mlir",
":register_common_dialects",
"//tensorflow/compiler/mlir/tensorflow",
"@llvm-project//mlir:AllPassesAndDialects",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:MlirReduceLib",
"@llvm-project//mlir:Reducer",
"@llvm-project//mlir:Support",
],
)
cc_library(
name = "register_common_dialects",
testonly = True,
srcs = ["register_common_dialects.cc"],
hdrs = ["register_common_dialects.h"],
deps = [
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/compiler/mlir/tensorflow:mlprogram_util",
"//tensorflow/compiler/mlir/tools/kernel_gen/ir:tf_framework_ops",
"//tensorflow/core/ir/types:Dialect",
"@llvm-project//mlir:AllExtensions",
"@llvm-project//mlir:AllPassesAndDialects",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:QuantOps",
"@llvm-project//mlir:RegisterAllDialects", # buildcleaner: keep
"@llvm-project//mlir:RegisterAllExtensions", # buildcleaner: keep
"@llvm-project//mlir:RegisterAllPasses", # buildcleaner: keep
"@llvm-project//mlir:ShapeDialect",
"@llvm-project//mlir:TensorDialect",
"@llvm-project//mlir:TosaDialect",
"@stablehlo//:register",
"@xla//xla/mlir_hlo:hlo_dialect_registration",
],
)
tf_cc_test(
name = "register_common_dialects_test",
srcs = ["register_common_dialects_test.cc"],
deps = [
":register_common_dialects",
"@com_google_googletest//:gtest_main",
"@llvm-project//mlir:IR",
],
)
tf_cc_binary(
name = "tf-mlir-translate",
testonly = True,
srcs = ["tf_mlir_translate_main.cc"],
deps = [
":init_mlir",
"//tensorflow/compiler/mlir/tensorflow:tf_xla_mlir_translate",
"//tensorflow/compiler/mlir/tensorflow:translate_lib",
"//tensorflow/compiler/mlir/tf2xla/tests/registration:graph_to_tf_executor_registration",
"//tensorflow/compiler/mlir/tools:translate_cl_options",
"//tensorflow/compiler/mlir/tools:translate_registration",
"//tensorflow/core:lib",
"//tensorflow/core:tensorflow",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/types:span",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Support",
"@llvm-project//mlir:TranslateLib",
"@xla//xla/hlo/translate/hlo_to_mhlo:translate_registration",
"@xla//xla/hlo/translate/mhlo_to_hlo:translate_registration",
],
)
tf_cc_test(
name = "mlir_graph_optimization_pass_test",
srcs = ["mlir_graph_optimization_pass_test.cc"],
deps = [
":mlir_graph_optimization_pass",
"//tensorflow/core:core_cpu",
"//tensorflow/core:framework",
"//tensorflow/core:ops",
"//tensorflow/core:portable_gif_internal",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"//tensorflow/core/common_runtime:device_set",
"//tensorflow/core/common_runtime:optimization_registry",
"//tensorflow/core/framework:tensor_testutil",
"//tensorflow/core/lib/monitoring:cell_reader",
"//tensorflow/core/platform:status",
"//tensorflow/core/protobuf:for_core_protos_cc",
"@com_google_absl//absl/status",
"@com_google_googletest//:gtest_main",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:IR",
],
)
filegroup(
name = "litfiles",
srcs = glob(["runlit*py"]),
visibility = ["//tensorflow:__subpackages__"],
)
exports_files(["run_lit.sh"])
tensorflow/compiler/mlir/README.md
+35
View File
@@ -0,0 +1,35 @@
# MLIR dialects and utilities for TensorFlow, TensorFlow Lite and XLA.
This module contains the MLIR
([Multi-Level Intermediate Representation](https://mlir.llvm.org))
dialects and utilities for
1. TensorFlow
2. XLA
3. TF Lite
See [MLIR's website](https://mlir.llvm.org) for complete documentation.
## Getting started
Building dialects and utilities here follow the standard approach using
`bazel` as the rest of TensorFlow.
### Using local LLVM repo
To develop across MLIR core and TensorFlow, it is useful to override the repo to
use a local version instead of fetching from head. This can be achieved by
setting up your local repository for Bazel build. For this you will need to
create bazel workspace and build files:
```sh
LLVM_SRC=... # this the path to the LLVM local source directory you intend to use.
touch ${LLVM_SRC}/BUILD.bazel ${LLVM_SRC}/WORKSPACE
```
You can then use this overlay to build TensorFlow:
```
bazel build --override_repository="llvm-raw=${LLVM_SRC}" \
-c opt tensorflow/compiler/mlir:tf-opt
```
tensorflow/compiler/mlir/g3doc/README.md
+3
View File
@@ -0,0 +1,3 @@
# TensorFlow MLIR
These are the docs for: https://www.tensorflow.org/mlir
tensorflow/compiler/mlir/g3doc/_book.yaml
+44
View File
@@ -0,0 +1,44 @@
# Copyright 2026 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.
# ==============================================================================
upper_tabs:
# Tabs left of dropdown menu
- include: /_upper_tabs_left.yaml
- include: /api_docs/_upper_tabs_api.yaml
# Dropdown menu
- name: Resources
path: /resources
is_default: true
menu:
- include: /resources/_menu_toc.yaml
lower_tabs:
# Subsite tabs
other:
- name: Guide
contents:
- title: Overview
path: /mlir/overview
- heading: Dialects
- title: Overview
path: /mlir/dialects
- title: TensorFlow
path: /mlir/tf_ops
- title: TensorFlow Lite
path: /mlir/tfl_ops
- heading: Passes
- title: TF dialect
path: /mlir/tf_passes
- include: /_upper_tabs_right.yaml
tensorflow/compiler/mlir/g3doc/_includes/tf_passes.md
+2366
View File
File diff suppressed because it is too large Load Diff
tensorflow/compiler/mlir/g3doc/_index.yaml
+70
View File
@@ -0,0 +1,70 @@
# Copyright 2026 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.
# ==============================================================================
book_path: /mlir/_book.yaml
project_path: /mlir/_project.yaml
description: An intermediate representation and compiler framework, MLIR unifies the
infrastructure for high-performance ML models in TensorFlow.
landing_page:
custom_css_path: /site-assets/css/style.css
rows:
- heading: MLIR unifies the infrastructure for high-performance ML models in TensorFlow.
items:
- description: >
The <a href="https://mlir.llvm.org/" class="external">MLIR</a> project defines a common
intermediate representation (IR) that unifies the infrastructure required to execute high
performance machine learning models in TensorFlow and similar ML frameworks. This project
will include the application of HPC techniques, along with integration of
search algorithms like reinforcement learning. MLIR aims to reduce the
cost to bring up new hardware, and improve usability for existing
TensorFlow users.
- code_block: |
<pre class = "prettyprint">
// Syntactically similar to LLVM:
func @testFunction(%arg0: i32) {
%x = call @thingToCall(%arg0) : (i32) -> i32
br ^bb1
^bb1:
%y = arith.addi %x, %x : i32
return %y : i32
}
</pre>
- classname: devsite-landing-row-cards
items:
- heading: "Multi-Level Intermediate Representation for Compiler Infrastructure"
youtube_id: qzljG6DKgic
buttons:
- label: Watch the video
path: https://www.youtube.com/watch?v=qzljG6DKgic
- heading: "A new intermediate representation and compiler framework"
image_path: /resources/images/tf-logo-card-16x9.png
path: https://blog.tensorflow.org/2019/04/mlir-new-intermediate-representation.html
buttons:
- label: Read on TensorFlow blog
path: https://blog.tensorflow.org/2019/04/mlir-new-intermediate-representation.html
- heading: MLIR on GitHub
image_path: /resources/images/github-card-16x9.png
path: https://github.com/llvm/llvm-project/tree/main/mlir
buttons:
- label: View on GitHub
path: https://github.com/llvm/llvm-project/tree/main/mlir
- heading: TensorFlow MLIR on GitHub
image_path: /resources/images/github-card-16x9.png
path: https://github.com/tensorflow/tensorflow/tree/master/tensorflow/compiler/mlir
buttons:
- label: View on GitHub
path: https://github.com/tensorflow/tensorflow/tree/master/tensorflow/compiler/mlir
tensorflow/compiler/mlir/g3doc/dialects.md
+37
View File
@@ -0,0 +1,37 @@
# MLIR dialects
## Overview
To separate different hardware and software targets, MLIR has “dialects”,
including:
* TensorFlow IR, which represents all things possible in TensorFlow graphs.
* XLA HLO IR, which is designed to take advantage of XLA’s compilation
abilities (with output to, among other things, TPUs).
* An experimental affine dialect, which focuses on
[polyhedral representations](https://en.wikipedia.org/wiki/Polytope_model)
and optimizations.
* LLVM IR, which has a 1:1 mapping between it and LLVM’s own representation,
allowing MLIR to emit GPU and CPU code through LLVM.
* TensorFlow Lite, which will translate to running code on mobile platforms.
Each dialect consists of a set of defined operations which have invariants
placed on them, like: “This is a binary operator, and the inputs and outputs
have the same types.”
## Adding to MLIR
MLIR has no fixed/built-in list of globally known operations (no “intrinsics”).
Dialects can define entirely custom types, which is how MLIR can model things
like the LLVM IR type system (which has first class aggregates), domain
abstractions important for ML-optimized accelerators like quantized types, and
even the Swift or Clang type systems (which are built around Swift/Clang
declaration nodes) in the future.
If you want to connect a new low-level compiler, you would create a new dialect
and the lowerings between the TensorFlow Graph dialect and your dialect.
This smooths the path for hardware and compiler makers. You can even target
dialects at different levels in the same model; the higher-level optimizers
will respect the unfamiliar parts of the IR and wait for a lower level to handle
it.
tensorflow/compiler/mlir/g3doc/images/mlir-infra.svg
+1
View File
File diff suppressed because one or more lines are too long
Side by Side

After

Width:  |  Height:  |  Size: 148 KiB

tensorflow/compiler/mlir/g3doc/overview.md
+36
View File
@@ -0,0 +1,36 @@
# MLIR
## Overview
MLIR, or Multi-Level Intermediate Representation, is a representation format
and library of compiler utilities that sits between the model representation
and low-level compilers/executors that generate hardware-specific code.
MLIR is, at its heart, a flexible infrastructure for modern optimizing
compilers. This means it consists of a specification for intermediate
representations (IR) and a code toolkit to perform transformations on that
representation. (In compiler parlance, as you move from higher-level
representations to lower-level representations, these transformations can be
called “lowerings”)
MLIR is highly influenced by [LLVM](https://llvm.org/) and unabashedly reuses
many great ideas from it. It has a flexible type system, and allows
representing, analyzing and transforming graphs combining multiple levels of
abstraction in the same compilation unit. These abstractions include TensorFlow
operations, nested polyhedral loop regions, and even LLVM instructions and fixed
hardware operations and types.
We expect MLIR to be of interest to many groups, including:
* Compiler researchers and implementers looking to optimize performance and
memory consumption of machine learning models
* Hardware makers looking for a way to connect their hardware to TensorFlow,
such as TPUs, portable neural hardware in phones, and other custom ASICs
* People writing language bindings that want to take advantage of optimizing
compilers and hardware acceleration.
The TensorFlow ecosystem contains a number of compilers and optimizers that
operate at multiple levels of the software and hardware stack. We expect the
gradual adoption of MLIR to simplify every aspect of this stack.
<img alt="MLIR overview diagram" src="./images/mlir-infra.svg"/>
tensorflow/compiler/mlir/g3doc/tf_passes.md
+7
View File
@@ -0,0 +1,7 @@
# TensorFlow passes
[TOC]
## TF dialect passes
<<_includes/tf_passes.md>>
tensorflow/compiler/mlir/glob_lit_test.bzl
+163
View File
@@ -0,0 +1,163 @@
# Test definitions for Lit, the LLVM test runner.
#
# This is reusing the LLVM Lit test runner in the interim until the new build
# rules are upstreamed.
# TODO(b/136126535): remove this custom rule.
"""Lit runner globbing test
"""
load("@bazel_skylib//lib:paths.bzl", "paths")
load("@rules_python//python:py_test.bzl", "py_test")
load(
"@xla//xla:lit.bzl",
"lit_script_with_xla_gpu_cuda_data_dir",
)
# Default values used by the test runner.
_default_test_file_exts = ["mlir", ".pbtxt", ".td"]
_default_driver = "@llvm-project//mlir:run_lit.sh"
_default_size = "small"
_default_tags = []
# These are patterns which we should never match, for tests, subdirectories, or
# test input data files.
_ALWAYS_EXCLUDE = [
"**/LICENSE.txt",
"**/README.txt",
"**/lit.local.cfg",
# Exclude input files that have spaces in their names, since bazel
# cannot cope with such "targets" in the srcs list.
"**/* *",
"**/* */**",
]
def get_canonical_repo_name(apparent_repo_name):
"""Returns the canonical repo name for the given apparent repo name seen by the module this bzl file belongs to."""
if not apparent_repo_name.startswith("@"):
apparent_repo_name = "@" + apparent_repo_name
return Label(apparent_repo_name).workspace_name
def _run_lit_test(name, data, size, tags, driver, features, exec_properties):
"""Runs lit on all tests it can find in `data` under tensorflow/compiler/mlir.
Note that, due to Bazel's hermetic builds, lit only sees the tests that
are included in the `data` parameter, regardless of what other tests might
exist in the directory searched.
Args:
name: str, the name of the test, including extension.
data: [str], the data input to the test.
size: str, the size of the test.
tags: [str], tags to attach to the test.
driver: str, label of the driver shell script.
Note: use of a custom driver is not currently supported
and specifying a default driver will abort the tests.
features: [str], list of extra features to enable.
"""
# Disable tests on windows for now, to enable testing rest of all xla and mlir.
py_test(
name = name,
srcs = ["@llvm-project//llvm:lit"],
tags = tags + ["no_pip", "no_windows"],
args = [
"tensorflow/compiler/mlir/" + paths.basename(data[-1]) + " --config-prefix=runlit -v",
] + features,
data = data + [
"//tensorflow/compiler/mlir:litfiles",
"@llvm-project//llvm:FileCheck",
"@llvm-project//llvm:count",
"@llvm-project//llvm:not",
],
deps = ["@pypi//lit"],
size = size,
main = "lit.py",
env = {
"LLVM_CANONICAL_REPO_NAME": get_canonical_repo_name("@llvm-project"),
"XLA_CANONICAL_REPO_NAME": get_canonical_repo_name("@xla"),
},
exec_properties = exec_properties,
)
def glob_lit_tests(
name = None,
exclude = [],
test_file_exts = _default_test_file_exts,
default_size = _default_size,
size_override = {},
data = [],
per_test_extra_data = {},
default_tags = _default_tags,
tags_override = {},
driver = _default_driver,
features = [],
exec_properties = {},
use_lit_test_suite = None, # @unused
hermetic_cuda_data_dir = None):
"""Creates all plausible Lit tests (and their inputs) under this directory.
Args:
name: str, name of the test_suite rule to generate for running all tests.
exclude: [str], paths to exclude (for tests and inputs).
test_file_exts: [str], extensions for files that are tests.
default_size: str, the test size for targets not in "size_override".
size_override: {str: str}, sizes to use for specific tests.
data: [str], additional input data to the test.
per_test_extra_data: {str: [str]}, extra data to attach to a given file.
default_tags: [str], additional tags to attach to the test.
tags_override: {str: str}, tags to add to specific tests.
driver: str, label of the driver shell script.
Note: use of a custom driver is not currently supported
and specifying a default driver will abort the tests.
features: [str], list of extra features to enable.
exec_properties: a dictionary of properties to pass on.
hermetic_cuda_data_dir: string. If set, the tests will be run with a
`--xla_gpu_cuda_data_dir` flag set to the hermetic CUDA data directory.
use_lit_test_suite: unused. For compatibility.
"""
# Ignore some patterns by default for tests and input data.
exclude = _ALWAYS_EXCLUDE + exclude
tests = native.glob(
["*." + ext for ext in test_file_exts],
exclude = exclude,
)
# Run tests individually such that errors can be attributed to a specific
# failure.
all_tests = []
for curr_test in tests:
final_test_name = curr_test
if hermetic_cuda_data_dir:
output_file = "with_xla_gpu_cuda_data_dir_{}".format(curr_test)
rule_name = "script_{}".format(output_file)
lit_script_with_xla_gpu_cuda_data_dir(
rule_name,
curr_test,
output_file,
hermetic_cuda_data_dir,
)
final_test_name = output_file
all_tests.append(final_test_name + ".test")
# Instantiate this test with updated parameters.
_run_lit_test(
name = final_test_name + ".test",
data = data + [final_test_name] +
per_test_extra_data.get(curr_test, []),
size = size_override.get(curr_test, default_size),
tags = default_tags + tags_override.get(curr_test, []),
driver = driver,
features = features,
exec_properties = exec_properties,
)
# TODO: remove this check after making it a required param.
if name:
native.test_suite(
name = name,
tests = all_tests,
tags = ["manual"],
)
tensorflow/compiler/mlir/init_mlir.cc
+64
View File
@@ -0,0 +1,64 @@
/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/compiler/mlir/init_mlir.h"
#include "llvm/ADT/StringRef.h"
#include "llvm/Support/CommandLine.h"
#include "llvm/Support/PrettyStackTrace.h"
#include "tensorflow/core/platform/init_main.h"
static llvm::cl::extrahelp FlagSplittingHelp(R"(
The command line parsing is split between the two flag parsing libraries used by
TensorFlow and LLVM:
* Flags before the first '--' are parsed by tensorflow::InitMain while those
post are parsed by LLVM's command line parser.
* If there is no separator, then no flags are parsed by InitMain and only
LLVM command line parser used.e
The above help options reported are for LLVM's parser, run with `--help --` for
TensorFlow's help.
)");
namespace tensorflow {
InitMlir::InitMlir(int *argc, char ***argv) {
llvm::setBugReportMsg(
"TensorFlow crashed, please file a bug on "
"https://github.com/tensorflow/tensorflow/issues with the trace "
"below.\n");
constexpr char kSeparator[] = "--";
// Find index of separator between two sets of flags.
int pass_remainder = 1;
bool split = false;
for (int i = 0; i < *argc; ++i) {
if (llvm::StringRef((*argv)[i]) == kSeparator) {
pass_remainder = i;
*argc -= (i + 1);
split = true;
break;
}
}
tensorflow::port::InitMain((*argv)[0], &pass_remainder, argv);
if (split) {
*argc += pass_remainder;
(*argv)[1] = (*argv)[0];
++*argv;
}
}
} // namespace tensorflow
tensorflow/compiler/mlir/init_mlir.h
+33
View File
@@ -0,0 +1,33 @@
/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_COMPILER_MLIR_INIT_MLIR_H_
#define TENSORFLOW_COMPILER_MLIR_INIT_MLIR_H_
namespace tensorflow {
// Initializer to perform TF's InitMain initialization.
// InitMain also performs flag parsing and '--' is used to separate flags passed
// to it: Flags before the first '--' are parsed by InitMain and argc and argv
// progressed to the flags post. If there is no separator, then no flags are
// parsed by InitMain and argc/argv left unadjusted.
class InitMlir {
public:
InitMlir(int *argc, char ***argv);
};
} // namespace tensorflow
#endif // TENSORFLOW_COMPILER_MLIR_INIT_MLIR_H_
tensorflow/compiler/mlir/kernels/BUILD
+22
View File
@@ -0,0 +1,22 @@
load("@rules_cc//cc:cc_library.bzl", "cc_library")
load("//tensorflow:tensorflow.default.bzl", "get_compatible_with_portable")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
default_visibility = [
"//visibility:public",
],
licenses = ["notice"],
)
# LINT.IfChange(cc_library_padding)
cc_library(
name = "padding",
srcs = [],
hdrs = ["padding.h"],
compatible_with = get_compatible_with_portable(),
visibility = [
"//tensorflow/compiler/mlir/quantization/tensorflow/utils:__pkg__",
],
)
# LINT.ThenChange(//tensorflow/lite/kernels/BUILD)
tensorflow/compiler/mlir/kernels/padding.h
+63
View File
@@ -0,0 +1,63 @@
/* Copyright 2017 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_COMPILER_MLIR_KERNELS_PADDING_H_
#define TENSORFLOW_COMPILER_MLIR_KERNELS_PADDING_H_
typedef enum {
kTfLitePaddingUnknown = 0,
kTfLitePaddingSame,
kTfLitePaddingValid,
} TfLitePadding;
// LINT.IfChange
namespace tflite_migration {
// Matching GetWindowedOutputSize in TensorFlow.
inline int ComputeOutSize(TfLitePadding padding, int image_size,
int filter_size, int stride, int dilation_rate = 1) {
int effective_filter_size = (filter_size - 1) * dilation_rate + 1;
// TODO(b/186448822): This uses 0 since the function has no other way to
// report error case
if (stride == 0) return 0;
switch (padding) {
case kTfLitePaddingSame:
return (image_size + stride - 1) / stride;
case kTfLitePaddingValid:
return (image_size + stride - effective_filter_size) / stride;
default:
return 0;
}
}
// It's not guaranteed that padding is symmetric. It's important to keep
// offset for algorithms need all paddings.
inline int ComputePaddingWithOffset(int stride, int dilation_rate, int in_size,
int filter_size, int out_size,
int* offset) {
int effective_filter_size = (filter_size - 1) * dilation_rate + 1;
int total_padding =
((out_size - 1) * stride + effective_filter_size - in_size);
total_padding = total_padding > 0 ? total_padding : 0;
*offset = total_padding % 2;
return total_padding / 2;
}
} // namespace tflite_migration
// LINT.ThenChange(//tensorflow/lite/kernels/padding.h)
#endif // TENSORFLOW_COMPILER_MLIR_KERNELS_PADDING_H_
tensorflow/compiler/mlir/lite/BUILD
+2341
View File
File diff suppressed because it is too large Load Diff
tensorflow/compiler/mlir/lite/README.md
+35
View File
@@ -0,0 +1,35 @@
# The new MLIR based TensorFlow to TensorFlow Lite converter
This directory contains:
1. [MLIR](https://github.com/llvm/llvm-project/tree/main/mlir) dialects,
transformation passes and utilities for TensorFlow Lite.
## API:
The API for converting TensorFlow models to TensorFlow Lite will be through
`tf.lite.TFLiteConverter`. All the conversion code is open sourced, and
the API will be integrated soon.
### The conversion process from TensorFlow to TensorFlow Lite includes the
following major passes:
- Import from GraphDef, in .pb or .pbtxt format, into MLIR.
- Raise to Control-flow-graph. Converts TF Control Flow dialect to TF dialect.
- The Canonicalization pass iteratively applies canonicalization
transformations in a greedy way until no further changes occur.
Canonicalization includes constant folding.
- The Legalize pass converts TensorFlow operations to TensorFlow Lite
ones. The operations that cannot be mapped to TensorFlow Lite dialect
are left as TensorFlow operations. Unsupported op handling follows the
proposed TFLite mechanism.
- Optimizations are performed in both the TF & TFLite dialect; aiming for small
size and high performance (among the core value proposition of
TensorFlow Lite models).
- The Export pass writes out TensorFlow Lite FlatBuffer format. This pass
operates on MLIR TensorFlow Lite dialect and is simple/direct translation.
See
https://github.com/tensorflow/tensorflow/blob/master/tensorflow/compiler/mlir/lite/tf_tfl_passes.cc
for the full list of MLIR passes for conversion from TensorFlow to TensorFlow
Lite.
tensorflow/compiler/mlir/lite/allocation.cc
+163
View File
@@ -0,0 +1,163 @@
/* Copyright 2017 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/allocation.h"
#include <stddef.h>
#include <sys/stat.h>
#include <sys/types.h>
#include <cstdint>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <memory>
#include "tensorflow/compiler/mlir/lite/core/api/error_reporter.h"
namespace tflite {
#ifndef TFLITE_MCU
FileCopyAllocation::FileCopyAllocation(const char* filename,
ErrorReporter* error_reporter)
: Allocation(error_reporter, Allocation::Type::kFileCopy) {
// Obtain the file size using fstat, or report an error if that fails.
std::unique_ptr<FILE, decltype(&fclose)> file(fopen(filename, "rb"), fclose);
if (!file) {
error_reporter_->Report("Could not open '%s'.", filename);
return;
}
// Use 64-bit stat on Windows so that files larger than 2GB report a correct
// size. MSVC's `struct stat`/`fstat` use a 32-bit `st_size` and fail for such
// files; `_stat64`/`_fstat64` provide a 64-bit `st_size`.
#ifdef _WIN32
struct _stat64 sb;
if (_fstat64(_fileno(file.get()), &sb) != 0) {
#else
struct stat sb;
if (fstat(fileno(file.get()), &sb) != 0) {
#endif
error_reporter_->Report("Failed to get file size of '%s'.", filename);
return;
}
buffer_size_bytes_ = sb.st_size;
std::unique_ptr<char[]> buffer(new char[buffer_size_bytes_]);
if (!buffer) {
error_reporter_->Report("Malloc of buffer to hold copy of '%s' failed.",
filename);
return;
}
size_t bytes_read =
fread(buffer.get(), sizeof(char), buffer_size_bytes_, file.get());
if (bytes_read != buffer_size_bytes_) {
error_reporter_->Report("Read of '%s' failed (too few bytes read).",
filename);
return;
}
// Versions of GCC before 6.2.0 don't support std::move from non-const
// char[] to const char[] unique_ptrs.
copied_buffer_.reset(const_cast<char const*>(buffer.release()));
}
FileCopyAllocation::~FileCopyAllocation() = default;
const void* FileCopyAllocation::base() const { return copied_buffer_.get(); }
size_t FileCopyAllocation::bytes() const { return buffer_size_bytes_; }
bool FileCopyAllocation::valid() const { return copied_buffer_ != nullptr; }
#endif
MemoryAllocation::MemoryAllocation(const void* ptr, size_t num_bytes,
ErrorReporter* error_reporter)
: Allocation(error_reporter, Allocation::Type::kMemory) {
#ifdef __arm__
if ((reinterpret_cast<uintptr_t>(ptr) & 0x3) != 0) {
// The flatbuffer schema has alignment requirements of up to 16 bytes to
// guarantee that data can be correctly accesses by various backends.
// Therefore, model pointer should also be 16-bytes aligned to preserve this
// requirement. But this condition only checks 4-bytes alignment which is
// the mininum requirement to prevent SIGBUS fault on 32bit ARM. Some models
// could require 8 or 16 bytes alignment which is not checked yet.
//
// Note that 64-bit ARM may also suffer a performance impact, but no crash -
// that case is not checked.
TF_LITE_REPORT_ERROR(error_reporter,
"The supplied buffer is not 4-bytes aligned");
buffer_ = nullptr;
buffer_size_bytes_ = 0;
return;
}
#endif // __arm__
// `android_local_test` doesn't support zipalign b/356640509 so we need this
// workaround to keep our clients working.
// TODO: b/356413060 - Remove the workaround once b/356640509 is fixed.
#if defined(__x86_64__) && defined(UNDEFINED_BEHAVIOR_SANITIZER)
if ((reinterpret_cast<uintptr_t>(ptr) & 0x3) != 0) {
#if defined(_WIN32)
// Windows / MSVC
aligned_ptr_ = _aligned_malloc(num_bytes, 4);
#elif defined(__ANDROID__) && __ANDROID_API__ < 28
// Older Android (API < 28)
if (posix_memalign(&aligned_ptr_, 4, num_bytes) != 0) {
aligned_ptr_ = nullptr;
}
#elif defined(__APPLE__)
// macOS/iOS: aligned_alloc is technically 10.15+,
// posix_memalign is safer for backwards compatibility.
if (posix_memalign(&aligned_ptr_, 4, num_bytes) != 0) {
aligned_ptr_ = nullptr;
}
#else
// Standard C11 (Modern Linux, Android API 28+)
aligned_ptr_ = ::aligned_alloc(4, num_bytes);
#endif
if (aligned_ptr_ == nullptr) {
TF_LITE_REPORT_ERROR(error_reporter, "Failed to allocate aligned buffer");
buffer_ = nullptr;
buffer_size_bytes_ = 0;
return;
}
memcpy(aligned_ptr_, ptr, num_bytes);
buffer_ = aligned_ptr_;
} else {
buffer_ = ptr;
}
#else // defined(__x86_64__) && defined(UNDEFINED_BEHAVIOR_SANITIZER)
buffer_ = ptr;
#endif // defined(__x86_64__) && defined(UNDEFINED_BEHAVIOR_SANITIZER)
buffer_size_bytes_ = num_bytes;
}
MemoryAllocation::~MemoryAllocation() {
#if defined(__x86_64__) && defined(UNDEFINED_BEHAVIOR_SANITIZER)
if (aligned_ptr_) {
free(aligned_ptr_);
}
#endif
}
const void* MemoryAllocation::base() const { return buffer_; }
size_t MemoryAllocation::bytes() const { return buffer_size_bytes_; }
bool MemoryAllocation::valid() const { return buffer_ != nullptr; }
} // namespace tflite
tensorflow/compiler/mlir/lite/allocation.h
+173
View File
@@ -0,0 +1,173 @@
/* Copyright 2017 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.
==============================================================================*/
/// \file
///
/// Memory management for TF Lite.
#ifndef TENSORFLOW_COMPILER_MLIR_LITE_ALLOCATION_H_
#define TENSORFLOW_COMPILER_MLIR_LITE_ALLOCATION_H_
#include <stddef.h>
#include <cstdio>
#include <cstdlib>
#include <memory>
#include "tensorflow/compiler/mlir/lite/core/api/error_reporter.h"
namespace tflite {
/// A memory allocation handle. This could be a mmap or shared memory.
class Allocation {
public:
using Ptr = std::unique_ptr<Allocation>;
virtual ~Allocation() = default;
enum class Type {
kMMap,
kFileCopy,
kMemory,
};
/// Base pointer of this allocation
virtual const void* base() const = 0;
/// Size in bytes of the allocation
virtual size_t bytes() const = 0;
/// Whether the allocation is valid
virtual bool valid() const = 0;
/// Return the type of the Allocation.
Type type() const { return type_; }
protected:
Allocation(ErrorReporter* error_reporter, Type type)
: error_reporter_(error_reporter), type_(type) {}
ErrorReporter* error_reporter_;
private:
const Type type_;
};
/// Note that not all platforms support MMAP-based allocation.
/// Use `IsSupported()` to check.
class MMAPAllocation : public Allocation {
public:
/// Loads and maps the provided file to a memory region.
/// If map_private is true, the mapping is private and writeable. Otherwise,
/// the mapping is shared and read-only.
MMAPAllocation(const char* filename, ErrorReporter* error_reporter,
bool map_private = false);
/// Loads and maps the provided file to a memory region at the given
/// offset and length (both in bytes).
/// If map_private is true, the mapping is private and writeable. Otherwise,
/// the mapping is shared and read-only.
MMAPAllocation(const char* filename, size_t offset, size_t length,
ErrorReporter* error_reporter, bool map_private = false);
/// Maps the provided file descriptor to a memory region.
/// If map_private is true, the mapping is private and writeable. Otherwise,
/// the mapping is shared and read-only.
/// Note: The provided file descriptor will be dup'ed for usage; the caller
/// retains ownership of the provided descriptor and should close accordingly.
MMAPAllocation(int fd, ErrorReporter* error_reporter,
bool map_private = false);
/// Maps the provided file descriptor, with the given offset and length (both
/// in bytes), to a memory region.
/// If map_private is true, the mapping is private and writeable. Otherwise,
/// the mapping is shared and read-only.
/// Note: The provided file descriptor will be dup'ed for usage; the caller
/// retains ownership of the provided descriptor and should close accordingly.
MMAPAllocation(int fd, size_t offset, size_t length,
ErrorReporter* error_reporter, bool map_private = false);
~MMAPAllocation() override;
const void* base() const override;
size_t bytes() const override;
bool valid() const override;
int fd() const { return mmap_fd_; }
// The start address of the mmapped buffer.
// This will be base() rounded down to the nearest page boundary.
const void* mmapped_buffer() const { return mmapped_buffer_; }
// The size of the mmapped buffer.
size_t mmapped_buffer_size() const { return bytes() + offset_in_buffer_; }
// Offset of mmapped_buffer() in the file referenced by the file descriptor.
size_t mmapped_buffer_offset_in_file() const {
return offset_of_buffer_in_file_;
}
static bool IsSupported();
protected:
// Data required for mmap.
int mmap_fd_ = -1; // mmap file descriptor
const void* mmapped_buffer_;
size_t buffer_size_bytes_ = 0;
// Used when the address to mmap is not page-aligned.
size_t offset_in_buffer_ = 0;
size_t offset_of_buffer_in_file_ = 0;
private:
// Assumes ownership of the provided `owned_fd` instance.
MMAPAllocation(ErrorReporter* error_reporter, int owned_fd, bool map_private);
// Assumes ownership of the provided `owned_fd` instance, and uses the given
// offset and length (both in bytes) for memory mapping.
MMAPAllocation(ErrorReporter* error_reporter, int owned_fd, size_t offset,
size_t length, bool map_private);
};
class FileCopyAllocation : public Allocation {
public:
/// Loads the provided file into a heap memory region.
FileCopyAllocation(const char* filename, ErrorReporter* error_reporter);
~FileCopyAllocation() override;
const void* base() const override;
size_t bytes() const override;
bool valid() const override;
private:
std::unique_ptr<const char[]> copied_buffer_;
size_t buffer_size_bytes_ = 0;
};
class MemoryAllocation : public Allocation {
public:
/// Provides a (read-only) view of the provided buffer region as an
/// allocation.
/// Note: The caller retains ownership of `ptr`, and must ensure it remains
/// valid for the lifetime of the class instance.
MemoryAllocation(const void* ptr, size_t num_bytes,
ErrorReporter* error_reporter);
~MemoryAllocation() override;
const void* base() const override;
size_t bytes() const override;
bool valid() const override;
private:
const void* buffer_;
#if defined(__x86_64__) && defined(UNDEFINED_BEHAVIOR_SANITIZER)
void* aligned_ptr_ = nullptr;
#endif
size_t buffer_size_bytes_ = 0;
};
} // namespace tflite
#endif // TENSORFLOW_COMPILER_MLIR_LITE_ALLOCATION_H_
tensorflow/compiler/mlir/lite/build_def.bzl
+89
View File
@@ -0,0 +1,89 @@
"""Build macros for TF Lite."""
load("//tensorflow/compiler/mlir/lite:special_rules.bzl", "tflite_copts_extra")
load("//tensorflow/lite:build_def.bzl", "clean_dep")
# LINT.IfChange(tflite_copts)
def tflite_copts():
"""Defines common compile time flags for TFLite libraries."""
copts = [
"-DFARMHASH_NO_CXX_STRING",
"-DEIGEN_ALLOW_UNALIGNED_SCALARS", # TODO(b/296071640): Remove when underlying bugs are fixed.
] + select({
clean_dep("//tensorflow:android_arm"): [
"-mfpu=neon",
],
# copybara:uncomment_begin(google-only)
# clean_dep("//tensorflow:chromiumos_x86_64"): [],
# copybara:uncomment_end
clean_dep("//tensorflow:ios_x86_64"): [
"-msse4.1",
],
clean_dep("//tensorflow:linux_x86_64"): [
"-msse4.2",
],
clean_dep("//tensorflow:linux_x86_64_no_sse"): [],
clean_dep("//tensorflow:windows"): [
# copybara:uncomment_begin(no MSVC flags in google)
# "-DTFL_COMPILE_LIBRARY",
# "-Wno-sign-compare",
# copybara:uncomment_end_and_comment_begin
"/DTFL_COMPILE_LIBRARY",
"/wd4018", # -Wno-sign-compare
# copybara:comment_end
],
"//conditions:default": [
"-Wno-sign-compare",
],
}) + select({
clean_dep("//tensorflow:optimized"): ["-O3"],
"//conditions:default": [],
}) + select({
clean_dep("//tensorflow:android"): [
"-ffunction-sections", # Helps trim binary size.
"-fdata-sections", # Helps trim binary size.
],
"//conditions:default": [],
}) + select({
clean_dep("//tensorflow:windows"): [],
"//conditions:default": [
"-fno-exceptions", # Exceptions are unused in TFLite.
],
}) + select({
clean_dep("//tensorflow/compiler/mlir/lite:tflite_with_xnnpack_explicit_false"): ["-DTFLITE_WITHOUT_XNNPACK"],
"//conditions:default": [],
}) + select({
clean_dep("//tensorflow/compiler/mlir/lite:tensorflow_profiler_config"): ["-DTF_LITE_TENSORFLOW_PROFILER"],
"//conditions:default": [],
}) + select({
clean_dep("//tensorflow/compiler/mlir/lite/delegates:tflite_debug_delegate"): ["-DTFLITE_DEBUG_DELEGATE"],
"//conditions:default": [],
}) + select({
clean_dep("//tensorflow/compiler/mlir/lite:tflite_mmap_disabled"): ["-DTFLITE_MMAP_DISABLED"],
"//conditions:default": [],
})
return copts + tflite_copts_extra()
# LINT.ThenChange(//tensorflow/lite/build_def.bzl:tflite_copts)
# LINT.IfChange(tflite_copts_warnings)
def tflite_copts_warnings():
"""Defines common warning flags used primarily by internal TFLite libraries."""
# TODO(b/155906820): Include with `tflite_copts()` after validating clients.
return select({
clean_dep("//tensorflow:windows"): [
# We run into trouble on Windows toolchains with warning flags,
# as mentioned in the comments below on each flag.
# We could be more aggressive in enabling supported warnings on each
# Windows toolchain, but we compromise with keeping BUILD files simple
# by limiting the number of config_setting's.
],
"//conditions:default": [
"-Wall",
],
})
# LINT.ThenChange(//tensorflow/lite/build_def.bzl:tflite_copts_warnings)
tensorflow/compiler/mlir/lite/common/tfl_pass_config.h
+159
View File
@@ -0,0 +1,159 @@
/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_COMPILER_MLIR_LITE_COMMON_TFL_PASS_CONFIG_H_
#define TENSORFLOW_COMPILER_MLIR_LITE_COMMON_TFL_PASS_CONFIG_H_
#include <string>
#include <utility>
#include "absl/strings/str_join.h"
#include "llvm/ADT/ArrayRef.h"
#include "llvm/Support/raw_ostream.h"
#include "tensorflow/compiler/mlir/lite/converter_flags.pb.h"
#include "tensorflow/compiler/mlir/lite/quantization/common/quantization_lib/quantization_config.h"
namespace mlir {
namespace TFL {
// A config that controls which passes get run as part TFLite converter.
struct PassConfig {
explicit PassConfig(QuantizationSpecs specs)
: quant_specs(std::move(specs)) {}
// If `emit_builtin_tflite_ops` is true, TF Lite legalization passes will be
// added, which produces TF Lite ops.
bool emit_builtin_tflite_ops = true;
// If `lower_tensor_list_ops` is true, tensorlist ops will be lowered to basic
// TF ops before legalization to TF Lite dialect.
bool lower_tensor_list_ops = false;
// The allowlist of functions that would be preserved after trimming.
llvm::ArrayRef<std::string> trim_functions_allowlist;
// All information about quantization.
QuantizationSpecs quant_specs;
// If `form_clusters` is true , clusters are formed by grouping consecutive
// ops of the same device, under a `tf_device.launch` op.
bool form_clusters = false;
// If `unfold_batch_matmul` is true, the tf.BatchMatMul is unfolded to a set
// of tfl.fully_connected ops.
bool unfold_batch_matmul = true;
// Whether to outline WhileOp at the end of the pipeline.
bool outline_tf_while = false;
// Whether to do shape inference.
bool shape_inference = true;
// Whether to do TFLite runtime verification.
bool runtime_verification = true;
// Whether to enable TFLite variables or not, this will allow
// mutable variables and produce ReadVariable/AssignVariable ops in TFLite.
bool enable_tflite_variables = false;
// Whether to unfold large splat constant tensors and replace them with
// fill operation.
bool unfold_large_splat_constant = false;
// Whether to run the `GuaranteeAllFuncsOneUsePass` to ensure each function
// has a single use.
bool guarantee_all_funcs_one_use = false;
// Whether to enable the hlo/stablehlo to tf conversion. This also supports
// the case where a saved model contains both TF module and serialized
// StableHLO module.
bool enable_hlo_to_tf_conversion = false;
// Whether to disable the direct hlo/stablehlo to Tensorflow Lite conversion.
//
// This prevents from directly converting from HLO to TFLite without going
// through TF for some of the ops. Some conversions are only supported through
// this path.
bool disable_hlo_to_tfl_conversion = false;
// Whether to enable to use DynamicUpdateSlice op.
bool enable_dynamic_update_slice = false;
// Whether to preserve AssertOp during legalization.
bool preserve_assert_op = false;
// Whether to enable TF->stablehlo passes.
bool enable_stablehlo_conversion = false;
// Whether to convert `tf.TensorList*` to `tfl.custom_op` if they can all
// be supported.
bool legalize_custom_tensor_list_ops = false;
// Whether to convert some tensor types to a lower precision if all values
// within that tensor are within the range of the lower precision. This could
// have side effects e.g. reduced flatbuffer size. Only certain type
// conversions are supported.
bool reduce_type_precision = false;
// Whether to consider this model a quantized model with quantize/dequantize
// ops and to convert kernels to quantized kernels wherever appropriate.
QDQConversionMode qdq_conversion_mode = QDQConversionMode::kQDQNone;
// When set to true, StableHLO Quantizer is run. The full configuration for
// the quantizer is at `ConverterFlags::quantization_config`.
bool enable_stablehlo_quantizer = false;
// Enables the attempt to directly lower composites into tflite ops.
bool enable_composite_direct_lowering = true;
// Specifies the framework of the original model.
tflite::ConverterFlags::ModelOriginFramework model_origin_framework =
tflite::ConverterFlags::UNSET;
// When set to true, convert +Inf/-Inf to MIN/MAX float value and output of
// convert only contains finite values.
bool canonicalizing_inf_as_min_max_float = true;
// When set to true, allows fusion of dynamic shaped broadcast ops. It helps
// fusing implicit broadcasting ops when output shape has dynamic dimensions,
// but it may cause incorrect results when broadcasting ops are introduced by
// explicit broadcasting in the source model.
bool unsafe_fuse_dynamic_shaped_broadcast = false;
// When set to true, enable unsafe single batch rank reduction.
bool unsafe_single_batch_rank_reduction = false;
};
inline llvm::raw_ostream& operator<<(llvm::raw_ostream& os,
const PassConfig& pass_config) {
return os << "emit_builtin_tflite_ops: "
<< pass_config.emit_builtin_tflite_ops
<< "\nlower_tensor_list_ops: " << pass_config.lower_tensor_list_ops
<< "\ntrim_functions_allowlist: "
<< absl::StrJoin(pass_config.trim_functions_allowlist.vec(), ",")
<< "\nform_clusters: " << pass_config.form_clusters
<< "\nunfold_batch_matmul: " << pass_config.unfold_batch_matmul
<< "\noutline_tf_while: " << pass_config.outline_tf_while
<< "\nshape_inference: " << pass_config.shape_inference
<< "\nruntime_verification: " << pass_config.runtime_verification
<< "\nenable_tflite_variables: "
<< pass_config.enable_tflite_variables
<< "\nunfold_large_splat_constant: "
<< pass_config.unfold_large_splat_constant
<< "\nguarantee_all_funcs_one_use: "
<< pass_config.guarantee_all_funcs_one_use
<< "\nenable_hlo_to_tf_conversion: "
<< pass_config.enable_hlo_to_tf_conversion
<< "\nenable_stablehlo_conversion: "
<< pass_config.enable_stablehlo_conversion
<< "\nlegalize_custom_tensor_list_ops: "
<< pass_config.legalize_custom_tensor_list_ops
<< "\nunsafe_fuse_dynamic_shaped_broadcast: "
<< pass_config.unsafe_fuse_dynamic_shaped_broadcast
<< "\nreduce_type_precision: " << pass_config.reduce_type_precision
<< "\nconvert_qdq_format: "
<< GetQDQQuantModeString(pass_config.qdq_conversion_mode)
<< "\nmodel_origin_framework: "
<< tflite::ConverterFlags::ModelOriginFramework_Name(
pass_config.model_origin_framework)
<< "\nunsafe_single_batch_rank_reduction: "
<< pass_config.unsafe_single_batch_rank_reduction << "\n";
}
} // namespace TFL
} // namespace mlir
#endif // TENSORFLOW_COMPILER_MLIR_LITE_COMMON_TFL_PASS_CONFIG_H_
tensorflow/compiler/mlir/lite/converter_flags.proto
+389
View File
@@ -0,0 +1,389 @@
// Copyright 2023 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
syntax = "proto2";
package tflite;
import "tensorflow/compiler/mlir/lite/debug/debug_options.proto";
import "tensorflow/compiler/mlir/lite/types.proto";
// Supported I/O file formats. Some formats may be input-only or output-only.
enum FileFormat {
FILE_FORMAT_UNKNOWN = 0;
// GraphDef, tensorflow/core/framework/graph.proto
TENSORFLOW_GRAPHDEF = 1;
// Tensorflow's mobile inference model.
// tensorflow/lite/schema/schema.fbs
TFLITE = 2;
// GraphViz
// Export-only.
GRAPHVIZ_DOT = 3;
}
// TocoFlags encodes extra parameters that drive tooling operations, that
// are not normally encoded in model files and in general may not be thought
// of as properties of models, instead describing how models are to be
// processed in the context of the present tooling job.
//
// Next ID to use: 71.
message ConverterFlags {
reserved 54, 61;
// Input file format
optional FileFormat input_format = 1;
// Output file format
optional FileFormat output_format = 2;
// Similar to inference_type, but allows to control specifically the
// quantization of input arrays, separately from other arrays.
//
// If not set, then the value of inference_type is implicitly used, i.e.
// by default input arrays are quantized like other arrays.
//
// Like inference_type, this only affects real-number arrays. By "real-number"
// we mean float arrays, and quantized arrays. This excludes plain
// integer arrays, strings arrays, and every other data type.
//
// The typical use for this flag is for vision models taking a bitmap
// as input, typically with uint8 channels, yet still requiring floating-point
// inference. For such image models, the uint8 input is quantized, i.e.
// the uint8 values are interpreted as real numbers, and the quantization
// parameters used for such input arrays are their mean_value, std_value
// parameters.
optional IODataType inference_input_type = 11;
// Sets the type of real-number arrays in the output file, that is, controls
// the representation (quantization) of real numbers in the output file,
// except for input arrays, which are controlled by inference_input_type.
//
// NOTE: this flag only impacts real-number arrays. By "real-number"
// we mean float arrays, and quantized arrays. This excludes plain
// integer arrays, strings arrays, and every other data type.
//
// For real-number arrays, the impact of this flag is to allow the output
// file to choose a different real-numbers representation (quantization)
// from what the input file used. For any other types of arrays, changing
// the data type would not make sense.
//
// Specifically:
// - If FLOAT, then real-numbers arrays will be of type float in
// the output file. If they were quantized in the input file, then
// they get dequantized.
// - If QUANTIZED_UINT8, then real-numbers arrays will be quantized
// as uint8 in the output file. If they were float in the input file,
// then they get quantized.
// - If not set, then all real-numbers arrays retain the same type in the
// output file as they have in the input file.
//
optional IODataType inference_type = 4;
// default_ranges_min and default_ranges_max are helpers to experiment
// with quantization of models. Normally, quantization requires the input
// model to have (min, max) range information for every activations array.
// This is needed in order to know how to quantize arrays and still achieve
// satisfactory accuracy. However, in some circumstances one would just like
// to estimate the performance of quantized inference, without caring about
// accuracy. That is what default_ranges_min and default_ranges_max are for:
// when specified, they will be used as default (min, max) range boundaries
// for all activation arrays that lack (min, max) range information, thus
// allowing for quantization to proceed.
//
// It should be clear from the above explanation that these parameters are
// for experimentation purposes only and should not be used in production:
// they make it easy to quantize models, but the resulting quantized model
// will be inaccurate.
//
// These values only apply to arrays quantized with the kUint8 data type.
optional float default_ranges_min = 5;
optional float default_ranges_max = 6;
// Equivalent versions of default_ranges_min/_max for arrays quantized with
// the kInt16 data type.
optional float default_int16_ranges_min = 15;
optional float default_int16_ranges_max = 16;
// Ignore and discard FakeQuant nodes. For instance, that can be used to
// generate plain float code without fake-quantization from a quantized
// graph.
optional bool drop_fake_quant = 7;
// Normally, FakeQuant nodes must be strict boundaries for graph
// transformations, in order to ensure that quantized inference has the
// exact same arithmetic behavior as quantized training --- which is the
// whole point of quantized training and of FakeQuant nodes in the first
// place. However, that entails subtle requirements on where exactly
// FakeQuant nodes must be placed in the graph. Some quantized graphs
// have FakeQuant nodes at unexpected locations, that prevent graph
// transformations that are necessary in order to generate inference
// code for these graphs. Such graphs should be fixed, but as a
// temporary work-around, setting this reorder_across_fake_quant flag
// allows toco to perform necessary graph transformations on them,
// at the cost of no longer faithfully matching inference and training
// arithmetic.
optional bool reorder_across_fake_quant = 8;
// If true, allow TOCO to create TF Lite Custom operators for all the
// unsupported Tensorflow ops.
optional bool allow_custom_ops = 10;
// Applies only to the case when the input format is TENSORFLOW_GRAPHDEF.
// If true, then control dependencies will be immediately dropped during
// import.
// If not set, the default behavior is as follows:
// - Default to false if the output format is TENSORFLOW_GRAPHDEF.
// - Default to true in all other cases.
optional bool drop_control_dependency = 12;
// Disables transformations that fuse subgraphs such as known LSTMs (not all
// LSTMs are identified).
optional bool debug_disable_recurrent_cell_fusion = 13;
// Uses the FakeQuantWithMinMaxArgs.num_bits attribute to adjust quantized
// array data types throughout the graph. The graph must be properly annotated
// with FakeQuant* ops on at least the edges and may contain additional ops on
// the interior of the graph to widen/narrow as desired.
//
// Input and output array data types may change because of this propagation
// and users must be sure to query the final data_type values.
optional bool propagate_fake_quant_num_bits = 14;
// Some fast uint8 GEMM kernels require uint8 weights to avoid the value 0.
// This flag allows nudging them to 1 to allow proceeding, with moderate
// inaccuracy.
optional bool allow_nudging_weights_to_use_fast_gemm_kernel = 17;
// Minimum size of constant arrays to deduplicate; arrays smaller will not be
// deduplicated.
optional int64 dedupe_array_min_size_bytes = 18 [default = 64];
// Split the LSTM inputs from 5 tensors to 18 tensors for TFLite.
// Ignored if the output format is not TFLite.
optional bool split_tflite_lstm_inputs = 19 [default = true];
// Store weights as quantized weights followed by dequantize operations.
// Computation is still done in float, but reduces model size (at the cost of
// accuracy and latency).
// DEPRECATED: Please use post_training_quantize instead.
optional bool quantize_weights = 20 [default = false];
// Full filepath of folder to dump the graphs at various stages of processing
// GraphViz .dot files. Preferred over --output_format=GRAPHVIZ_DOT in order
// to keep the requirements of the output file.
optional string dump_graphviz_dir = 24;
// Boolean indicating whether to dump the graph after every graph
// transformation.
optional bool dump_graphviz_include_video = 25;
// Boolean indicating whether to quantize the weights of the converted float
// model. Model size will be reduced and there will be latency improvements
// (at the cost of accuracy).
optional bool post_training_quantize = 26 [default = false];
// This flag only works when converting to TensorFlow Lite format.
// When enabled, unsupported ops will be converted to select TensorFlow ops.
// `enable_select_tf_ops` should always be used with `allow_custom_ops`.
// WARNING: Experimental interface, subject to change
optional bool enable_select_tf_ops = 27 [default = false];
// This flag only works when converting to TensorFlow Lite format.
// When enabled, all TensorFlow ops will be converted to select TensorFlow
// ops.
// This will force `enable_select_tf_ops` to true.
// `force_select_tf_ops` should always be used with `enable_select_tf_ops`.
// WARNING: Experimental interface, subject to change
optional bool force_select_tf_ops = 28 [default = false];
// Boolean indicating whether to convert float32 constant buffers to
// float16. This is typically done to reduce model size. Delegates may also
// wish to implement kernels on reduced precision floats for performance
// gains.
optional bool quantize_to_float16 = 29 [default = false];
// Boolean flag indicating whether the converter should allow models with
// dynamic Tensor shape. When set to False, the converter will generate
// runtime memory offsets for activation Tensors (with 128 bits alignment)
// and error out on models with undetermined Tensor shape. (Default: True)
optional bool allow_dynamic_tensors = 30 [default = true];
// Full filepath of the folder to dump conversion logs. This includes a global
// view of the conversion process, and user can choose to submit those logs.
optional string conversion_summary_dir = 31;
// String representing the custom ops OpDefs that are included in the
// GraphDef.
// Deprecated do not use.
repeated string custom_opdefs = 32 [deprecated = true];
// Name of user's defined Tensorflow ops required in the TensorFlow Lite
// runtime. These ops will be supported as select TensorFlow ops.
repeated string select_user_tf_ops = 33;
// Whether to enable tflite resource variables during conversion or not.
// Note: This is an experimental feature.
optional bool enable_tflite_resource_variables = 34 [default = true];
// Whether to unfold tf.BatchMatMul to a set of tfl.fully_connected ops. If
// not, translate to tfl.batch_matmul.
// WARNING: Experimental interface, subject to change.
optional bool unfold_batchmatmul = 35 [default = false];
// Whether to lower static Tensor List ops to builtin ops. If not, use Flex
// tensor list ops.
// WARNING: Experimental interface, subject to change.
optional bool lower_tensor_list_ops = 36 [default = true];
// The accumulation type to use when quantize_to_float16 is true. Typical
// choices would be either float16 or float32.
optional IODataType accumulation_type = 37;
// Whether this model supports inference in bfloat16.
// Note: This is an experimental feature.
optional bool allow_bfloat16 = 38 [default = false];
// If true, automatically adds all tf ops into the model as select Tensorflow
// ops.
optional bool allow_all_select_tf_ops = 39;
// Whether to unfold large splat constant tensors in the flatbuffer to reduce
// model size.
optional bool unfold_large_splat_constant = 40 [default = false];
// Name of TFLite backends which are needed to check compatibility.
// WARNING: Experimental interface, subject to change.
repeated string supported_backends = 41;
// Whether to force to use batch size one when the batch size is None during
// lowering tensor list ops.
optional bool default_to_single_batch_in_tensor_list_ops = 42
[default = false];
// Disable per_channel quantization for dynamic range quantization.
// Note: This is an experimental feature
optional bool disable_per_channel_quantization = 43 [default = false];
// If false, the old TOCO dynamic range quantization is used.
// Note: This is an experimental feature
optional bool enable_mlir_dynamic_range_quantizer = 44 [default = false];
// When the output model is used for TF Quantization, this flag indicates the
// mode of TF Quantization. Ex: DEFAULT, LEGACY_INTEGER,...
optional string tf_quantization_mode = 45;
// Disable inferring tensor range for quantization.
// Note: This is an experimental feature
optional bool disable_infer_tensor_range = 46 [default = false];
// Enable using num bits set in fake quant attributes for quantization.
// Note: This is an experimental feature
optional bool use_fake_quant_num_bits = 47 [default = false];
// Enable converting to DynamicUpdateSlice op (for ops like TensorListSetItem)
// Note: This is an experimental feature
optional bool enable_dynamic_update_slice = 48 [default = false];
// Whether to preserve `TF::AssertOp`.
optional bool preserve_assert_op = 49 [default = false];
// Whether to ensure each function has a single use.
optional bool guarantee_all_funcs_one_use = 50 [default = false];
// Whether to convert model to stablehlo.
optional bool convert_to_stablehlo = 51 [default = false];
// If false, skip the variable quantization passes.
// Note: This is an experimental feature
optional bool enable_mlir_variable_quantization = 52 [default = false];
// If true, disable folding mul->fc as in layer norm during optimize pass.
optional bool disable_fuse_mul_and_fc = 53 [default = false];
// Flag to enable hlo to tf conversion.
// This is useful to exercise StableHLO -> HLO -> TF -> TFLite path.
optional bool enable_hlo_to_tf_conversion = 55
[default = false, deprecated = true];
// Additional parameters for controlling debug facilities.
optional tensorflow.converter.DebugOptions debug_options = 56;
// If true, TFlite will use offsets instead of flatbuffers array to store
// buffers and custom options Note: This is an experimental feature
optional bool use_buffer_offset = 57 [default = false];
// Whether to legalize "tf.TensorList*" ops to custom tflite if they
// can all be supported.
optional bool legalize_custom_tensor_list_ops = 58 [default = false];
// Whether to convert some tensor types to a lower precision if all values
// within that tensor are within the range of the lower precision. This could
// have side effects e.g. reduced flatbuffer size. Only certain type
// conversions are supported.
// WARNING: Experimental interface, subject to change.
optional bool reduce_type_precision = 59 [default = false];
// Whether to consider this model a quantized model with quantize/dequantize
// ops and to convert kernels to quantized kernels wherever appropriate.
// WARNING: Experimental interface, subject to change.
optional string qdq_conversion_mode = 60 [default = "NONE"];
// Disables per channel weights quantization for Dense layers and enables
// legacy per tensor quantization. The legacy quantization for Dense layers is
// inconsistent with Conv 1x1 which always performs per channel quantization.
optional bool disable_per_channel_quantization_for_dense_layers = 62
[default = false];
// Enables the attempt to directly lower composites into tflite ops.
// WARNING: Experimental interface, subject to change.
optional bool enable_composite_direct_lowering = 63 [default = false];
// The source model framework.
enum ModelOriginFramework {
UNSET = 0;
TENSORFLOW = 1;
KERAS = 2;
JAX = 3;
PYTORCH = 4;
}
// The source model type.
optional ModelOriginFramework model_origin_framework = 64 [default = UNSET];
// When set to true, convert +Inf/-Inf to MIN/MAX float value and output of
// convert only contains finite values.
optional bool canonicalizing_inf_as_min_max_float = 65 [default = false];
// When set to true, debug metadata will be generated and attached to
// serialized TFLite flatbuffer.
optional bool serialize_debug_metadata = 66 [default = false];
// When set, adheres to the QDQ annotations added by the framework when
// possible rather than quantizing any op that is possible to quantize.
// WARNING: Experimental interface, subject to change.
optional bool strict_qdq_mode = 67 [default = false];
// When set to true, allows fusion of dynamic shaped broadcast ops. It helps
// fusing implicit broadcasting ops when output shape has dynamic dimensions,
// but it may cause incorrect results when broadcasting ops are introduced by
// explicit broadcasting in the source model.
optional bool unsafe_fuse_dynamic_shaped_broadcast = 68 [default = false];
// If false, downcast x64 tensors and inputs to x32.
optional bool enable_x64 = 69 [default = true];
// If true, enable unsafe single batch rank reduction.
optional bool unsafe_single_batch_rank_reduction = 70 [default = false];
}
tensorflow/compiler/mlir/lite/converter_gen.cc
+652
View File
@@ -0,0 +1,652 @@
/* 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 <assert.h>
#include <sstream>
#include <string>
#include <vector>
#include "llvm/ADT/ArrayRef.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/Support/Casting.h"
#include "llvm/Support/CommandLine.h"
#include "llvm/Support/FormatVariadic.h"
#include "llvm/Support/InitLLVM.h"
#include "llvm/Support/Signals.h"
#include "llvm/TableGen/Error.h"
#include "llvm/TableGen/Main.h"
#include "llvm/TableGen/Record.h"
#include "llvm/TableGen/TableGenBackend.h"
#include "mlir/TableGen/Attribute.h" // from @llvm-project
#include "mlir/TableGen/Format.h" // from @llvm-project
#include "mlir/TableGen/Operator.h" // from @llvm-project
#include "mlir/TableGen/Predicate.h" // from @llvm-project
using llvm::DefInit;
using llvm::dyn_cast;
using llvm::formatv;
using llvm::LessRecord;
using llvm::raw_ostream;
using llvm::Record;
using llvm::RecordKeeper;
using llvm::RecordRecTy;
using llvm::SmallVector;
using llvm::StringInit;
using llvm::StringRef;
enum ActionType {
OpConv,
RuntimeVerify,
};
// NOLINTNEXTLINE
llvm::cl::opt<ActionType> action(
llvm::cl::desc("Action to perform:"),
llvm::cl::values(clEnumValN(OpConv, "gen-operator-converters",
"Generate operator converters"),
clEnumValN(RuntimeVerify, "gen-runtime-verifiers",
"Generate TFLite runtime verifiers")));
// Returns the associated option name for the given op definition.
static inline std::string GetOperatorOptionName(const Record &def) {
assert(def.getName().starts_with("TFL_") && "unexpected op prefix");
assert(def.getName().ends_with("Op") && "unexpected op suffix");
auto *custom_option = dyn_cast<StringInit>(def.getValueInit("customOption"));
std::ostringstream oss;
if (custom_option)
oss << custom_option->getValue().str();
else
oss << def.getName().drop_front(4).drop_back(2).str() << "Options";
return oss.str();
}
// Returns the builder function name for the given op definition.
static inline std::string GetOperatorBuilderName(StringRef op_name) {
assert(op_name.starts_with("TFL_") && "unexpected op prefix");
assert(op_name.ends_with("Op") && "unexpected op suffix");
// E.g., AddOp -> CreateAddOperator
std::ostringstream oss;
oss << "Create" << op_name.drop_front(4).str() << "erator";
return oss.str();
}
static inline bool IsLstmOp(const StringRef op_name) {
return op_name.take_back(6) == "LSTMOp";
}
static int HasOptions(const Record &def) {
if (def.getValueAsBit("hasOptions")) {
return 1;
}
if (def.getValueAsBit("hasOptions2")) {
return 2;
}
return 0;
}
static void EmitOptionBuilders(const RecordKeeper &record_keeper,
const std::vector<const Record *> &defs,
raw_ostream *ostream) {
raw_ostream &os = *ostream;
const auto attr_type = record_keeper.getClass("Attr");
for (const auto *def : defs) {
const int has_options = HasOptions(*def);
// TFLite ops without options are skipped over.
if (!has_options) {
continue;
}
StringRef op_name = def->getName().drop_front(4); // Strip 'TFL_' prefix
std::string option_name = GetOperatorOptionName(*def);
std::string tflite_option_name =
option_name == "BasicLSTMOptions" ? "LSTMOptions" : option_name;
os << "flatbuffers::Offset<tflite::" << tflite_option_name << "> Create"
<< option_name << "(mlir::TFL::" << op_name
<< " op, flatbuffers::FlatBufferBuilder *fbb) {\n";
// Construct all the builder option needed.
SmallVector<std::string, 8> options;
// Add options due to attributes (not-derived).
auto *arg_values = def->getValueAsDag("arguments");
mlir::tblgen::Operator op(*def);
for (unsigned i = 0, e = arg_values->getNumArgs(); i != e; ++i) {
auto arg = arg_values->getArg(i);
const auto *arg_def = dyn_cast<DefInit>(arg);
if (!arg_def) continue;
if (arg_def->getDef()->isSubClassOf(attr_type)) {
// This binds the name of the attribute in the TD file with the name
// of the add function of the builder and also with the conversion
// function to convert from the internal representation to the format
// expected by the flatbuffer builder. While this constrains the
// naming of the ops/attributes in the TD file, this also removes the
// need for specifying indirection. This tool is specific to TFLite
// conversion generation and so the simplicity was chosen over the
// flexibility.
StringRef arg_name = arg_values->getArgNameStr(i);
// Skip any "intermiadiateXXX" attribute as they are specially handled
// in the exporter. They are special because though they are attributes
// in the MLIR they are expressed as tensors in the flatbuffer instead
// of option.
if (IsLstmOp(op_name) && arg_name.take_back(12) == "intermediate")
continue;
os << formatv(
" auto {0} = Convert{1}ForOptionWriter(op.{2}(), fbb);\n",
arg_name, mlir::tblgen::Attribute(arg_def).getAttrDefName(),
op.getGetterName(arg_name));
options.push_back(arg_name.str());
}
}
// Add options due to derived attributes.
for (const auto &val : def->getValues()) {
if (auto *record = dyn_cast<RecordRecTy>(val.getType())) {
if (record->isSubClassOf(attr_type)) {
if (record->getClasses().size() != 1) {
PrintFatalError(
def->getLoc(),
"unsupported attribute modelling, only single class expected");
}
os << formatv(
" auto {0} = Convert{1}ForOptionWriter(op.{2}(), fbb);\n",
val.getName(), record->getClasses()[0]->getName(),
op.getGetterName(val.getName()));
options.push_back(std::string(val.getName()));
}
}
}
os << " tflite::" << tflite_option_name << "Builder b(*fbb);\n";
for (const auto &option : options)
os << formatv(" b.add_{0}(std::move({0}));\n", option);
os << " return b.Finish();\n}\n";
}
}
// For each TFLite op, emits a builder function that packs the TFLite op into
// the corresponding FlatBuffer object.
//
// TODO(hinsu): Revisit if only builtin_options and mutating_variable_inputs
// arguments that depend on op definitions should be auto-generated and then
// operator should be built by the caller because it does not require
// auto-generation.
static void EmitOperatorBuilders(const std::vector<const Record *> &defs,
raw_ostream *ostream) {
raw_ostream &os = *ostream;
for (const auto *def : defs) {
StringRef op_name = def->getName().drop_front(4);
const bool has_intermediates = op_name.take_back(6) == "LSTMOp";
// Signature
os << "static flatbuffers::Offset<tflite::Operator> "
<< GetOperatorBuilderName(def->getName()) << "(mlir::TFL::" << op_name
<< " tflOp, uint32_t opcode_index, "
<< "const std::vector<int32_t>& operands,"
<< "const std::vector<int32_t>& results,"
<< (has_intermediates ? "const std::vector<int32_t>& intermediate_index,"
: "")
<< "flatbuffers::FlatBufferBuilder *fbb,"
<< "int debug_metadata_index) {\n";
// Inputs & outputs
os << " auto inputs = fbb->CreateVector(operands);\n"
" auto outputs = fbb->CreateVector(results);\n\n";
// Intermediates for LSTM.
if (has_intermediates) {
os << " auto intermediates = fbb->CreateVector(intermediate_index);\n";
}
// Build the FlatBuffer operator
os << " return tflite::CreateOperator(\n"
" *fbb, opcode_index, inputs, outputs,\n";
const int has_options = HasOptions(*def);
if (has_options == 1) {
auto option_name = GetOperatorOptionName(*def);
std::string tflite_option_name =
option_name == "BasicLSTMOptions" ? "LSTMOptions" : option_name;
os << " tflite::BuiltinOptions_" << tflite_option_name << ", "
<< "Create" << option_name << "(tflOp, fbb).Union(),\n";
} else {
os << " tflite::BuiltinOptions_NONE, /*builtin_options=*/0,\n";
}
// Only built-in ops' builders are auto-generated. custom_options are only
// used by custom or flex ops and those ops are handled manually.
os << " /*custom_options=*/0, "
<< "tflite::CustomOptionsFormat_FLEXBUFFERS,\n"
<< " /*mutating_variable_inputs=*/0,"
<< (has_intermediates ? "intermediates" : "/*intermediates=*/0");
if (has_options == 2) {
os << ",\n"
<< " /*large_custom_options_offset=*/0,\n"
<< " /*large_custom_options_size=*/0";
os << ",\n";
const std::string option_name = GetOperatorOptionName(*def);
os << " tflite::BuiltinOptions2_" << option_name << ", "
<< "Create" << option_name << "(tflOp, fbb).Union()";
} else {
os << ",\n"
<< " /*large_custom_options_offset=*/0,\n"
<< " /*large_custom_options_size=*/0";
os << ",\n";
os << " tflite::BuiltinOptions2_NONE, /*builtin_options2=*/0";
}
os << ",\n"
<< " /*debug_metadata_index=*/debug_metadata_index";
os << ");\n}\n\n";
}
}
static inline std::string GetOperatorName(const Record &def) {
auto name = def.getValueAsString("opName");
// Special case for basic_lstm.
if (name == "basic_lstm") {
return "LSTM";
}
return name.upper();
}
// Emits a function that returns built-in operator code for each TFLite op.
//
// The signature of the function is:
//
// std::optional<tflite::BuiltinOperator>
// mlir::GetBuiltinOpCode(mlir::Operation* op);
//
// TODO(hinsu): Consider converting this to a static constant associative
// container instead of a series of if conditions, if required.
static void EmitGetBuiltinOpCode(const std::vector<const Record *> &defs,
raw_ostream *ostream) {
raw_ostream &os = *ostream;
os << "std::optional<tflite::BuiltinOperator> "
"mlir::GetBuiltinOpCode(mlir::Operation* op) {\n";
for (const auto *def : defs) {
StringRef op_name = def->getName().drop_front(4);
auto operator_name = GetOperatorName(*def);
os << " if (isa<mlir::TFL::" << op_name << ">(op))\n"
<< " return tflite::BuiltinOperator_" << operator_name << ";\n";
}
os << " return std::nullopt;\n"
"}\n";
}
// Emits functions that return the min/max operand numbers for a given tflite op
// name.
//
// Signature:
// llvm::MinMax mlir::OperandNumbersMinMax(llvm::StringRef op_name) {
// if(const auto *op = op_union.AsOptions()) {
// return {min, max};
// }
// ...
// return {0, 0};
// }
static void EmitOperandNumbers(const RecordKeeper &record_keeper,
const std::vector<const Record *> &defs,
raw_ostream *ostream) {
raw_ostream &os = *ostream;
const auto attr_type = record_keeper.getClass("Attr");
const auto optional_tensor = record_keeper.getClass("TFL_TensorOfOrNone");
os << "llvm::MinMax mlir::OperandNumbersMinMax(llvm::StringRef op_name) {\n";
for (const auto *def : defs) {
auto op_name = def->getValueAsString("opName");
int tail_optional_tensor = 0, tensor_number_max = 0;
auto *arg_values = def->getValueAsDag("arguments");
for (int i = 0, e = arg_values->getNumArgs(); i < e; ++i) {
auto arg = arg_values->getArg(i);
auto *arg_def = dyn_cast<DefInit>(arg);
if (!arg_def) continue;
if (!arg_def->getDef()->isSubClassOf(attr_type)) {
tensor_number_max++;
if (arg_def->getDef()->isSubClassOf(optional_tensor)) {
tail_optional_tensor++;
} else {
tail_optional_tensor = 0;
}
}
}
const int tensor_number_min = tensor_number_max - tail_optional_tensor;
os << formatv(" if (op_name == \"tfl.{0}\") {{\n", op_name)
<< " return {" << tensor_number_min << ", " << tensor_number_max
<< "};\n }\n";
}
os << " return {0, 0};\n}\n";
}
// Emits a builder function that returns the packed FlatBuffer object given
// a general mlir::Operation.
//
// The signature of the function is:
//
// std::optional<Flatbuffers::Offset<tflite::Operator>>
// mlir::CreateFlatBufferOperator(
// mlir::Operation* op,
// uint32_t opcode_index,
// const std::vector<int32_t>& operands,
// const std::vector<int32_t>& results,
// const std::vector<int32_t>& intermediates,
// flatbuffers::FlatBufferBuilder *fbb,
// std::optional<int> debug_metadata_index);
static void EmitBuildOperator(const std::vector<const Record *> &defs,
raw_ostream *ostream) {
raw_ostream &os = *ostream;
// Signature
os << "std::optional<flatbuffers::Offset<tflite::Operator>>\n"
"mlir::CreateFlatBufferOperator(mlir::Operation* op, "
"uint32_t opcode_index, "
"const std::vector<int32_t>& operands,"
"const std::vector<int32_t>& results,"
"const std::vector<int32_t>& intermediates,"
"flatbuffers::FlatBufferBuilder *fbb,"
"std::optional<int> debug_metadata_index) {\n";
for (const auto *def : defs) {
StringRef op_name = def->getName().drop_front(4);
// Try to cast to each op case and call the corresponding op builder
os << " if (auto tflOp = llvm::dyn_cast<mlir::TFL::" << op_name
<< ">(op))\n"
<< " return " << GetOperatorBuilderName(def->getName())
<< "(tflOp, opcode_index, operands, results, "
<< (op_name.take_back(6) == "LSTMOp" ? "intermediates, " : "")
<< "fbb, debug_metadata_index.value_or(-1));\n";
}
os << " return std::nullopt;\n"
"}\n";
}
// Emit a function that converts a BuiltinOptionsUnion to a vector of attributes
// Signature:
// void mlir::BuiltinOptions{id}ToAttributes(
// tflite::BuiltinOptions{id}Union op_union,
// mlir::Builder builder,
// llvm::SmallVectorImpl<mlir::NamedAttribute> &attributes);
//
// where id is an empty string if builtin_options_id is 1, or builtin_options_id
// otherwise.
static void EmitBuiltinOptionsToAttributes(
const RecordKeeper &record_keeper, const std::vector<const Record *> &defs,
raw_ostream *ostream, const int builtin_options_id) {
raw_ostream &os = *ostream;
const std::string builtin_options_suffix = [&] {
switch (builtin_options_id) {
case 1:
return "";
case 2:
return "2";
}
return "UnknownId";
}();
// Signature
os << "void mlir::BuiltinOptions" << builtin_options_suffix
<< "ToAttributes("
"tflite::BuiltinOptions"
<< builtin_options_suffix
<< "Union op_union, "
"mlir::Builder builder, "
"llvm::SmallVectorImpl<mlir::NamedAttribute> &attributes) {\n";
const auto attr_type = record_keeper.getClass("Attr");
for (const auto *def : defs) {
const int has_options = HasOptions(*def);
if (has_options != builtin_options_id) {
continue;
}
auto option_name = GetOperatorOptionName(*def);
// Basic LSTM and LSTM ops share the same option to attribute converter.
if (option_name == "BasicLSTMOptions") {
continue;
}
os << formatv(" if(const auto *op = op_union.As{0}()) {{\n", option_name);
// We only care about options that are in arguments
auto *arg_values = def->getValueAsDag("arguments");
for (unsigned i = 0, e = arg_values->getNumArgs(); i != e; ++i) {
auto arg = arg_values->getArg(i);
const auto *arg_def = dyn_cast<DefInit>(arg);
if (!arg_def) continue;
if (arg_def->getDef()->isSubClassOf(attr_type)) {
StringRef arg_name = arg_values->getArgNameStr(i);
// Already handle this case in flatbuffer_import.cc.
if ((option_name == "LSTMOptions" ||
option_name == "UnidirectionalSequenceLSTMOptions") &&
arg_name.take_back(12) == "intermediate")
continue;
StringRef attr_type = mlir::tblgen::Attribute(arg_def).getAttrDefName();
os << formatv(
" attributes.emplace_back(builder.getNamedAttr(\"{0}\","
" Build{1}(op->{0}, builder)));\n",
arg_name, attr_type);
}
}
os << " return;\n";
os << " }\n";
}
if (builtin_options_id == 2) {
os << " BuiltinOptions2ToAttributesManual(op_union, builder, "
"attributes);\n";
}
// Fallthrough case is no attributes
os << "}";
}
// The function below has a non-constant reference as that is required by LLVM's
// TableGenMain.
// NOLINTNEXTLINE
static bool OperatorWritersMain(raw_ostream &os, const RecordKeeper &records) {
emitSourceFileHeader("MLIR TFLite FlatBuffer Builders", os);
// Retrieve all the definitions derived from TFL_Op and sort by record name.
std::vector<const Record *> defs = records.getAllDerivedDefinitions("TFL_Op");
llvm::sort(defs, LessRecord());
for (const auto *def : defs) {
// TFLite ops in the .td file are expected to follow the naming convention:
// TFL_<OpName>Op.
// The generated TFLite op C++ class should be TFL::<OpName>Op.
// The generated operator's options should be tflite::<OpName>Options.
// The option builder should be Create<OpName>Options.
if (!def->getName().starts_with("TFL_"))
PrintFatalError(def->getLoc(),
"unexpected op name format: 'TFL_' prefix missing");
if (!def->getName().ends_with("Op"))
PrintFatalError(def->getLoc(),
"unexpected op name format: 'Op' suffix missing");
}
EmitOptionBuilders(records, defs, &os);
os << "\n\n";
EmitOperatorBuilders(defs, &os);
os << "\n\n";
EmitGetBuiltinOpCode(defs, &os);
os << "\n\n";
EmitBuildOperator(defs, &os);
os << "\n\n";
EmitBuiltinOptionsToAttributes(records, defs, &os, /*builtin_options_id=*/1);
os << "\n\n";
EmitBuiltinOptionsToAttributes(records, defs, &os, /*builtin_options_id=*/2);
os << "\n\n";
EmitOperandNumbers(records, defs, &os);
return false;
}
static void GenOperandResultVerifier(raw_ostream &os,
llvm::ArrayRef<const llvm::Init *> values,
StringRef valueKind) {
mlir::tblgen::FmtContext fctx;
bool first = true;
for (const auto &static_value : llvm::enumerate(values)) {
auto *definit = llvm::cast<llvm::DefInit>(static_value.value());
auto *val = definit->getDef()->getValue("tflRuntimeTypePredicate");
if (!val) continue;
// Create code block on first type to verify.
if (first) {
os << " {\n";
os << " unsigned index = " << static_value.index() << ";\n";
first = false;
}
mlir::tblgen::Pred pred(dyn_cast<llvm::DefInit>(val->getValue()));
auto desc =
definit->getDef()->getValueAsString("tflRuntimeTypeDescription");
// Emit a loop to check all operands.
os << formatv(" for (Value v : top.getODS{0}{1}s({2})) {{\n",
// Capitalize the first letter to match the function name
valueKind.substr(0, 1).upper(), valueKind.substr(1),
static_value.index());
os << " (void)v;\n"
<< " if (!("
<< tgfmt(pred.getCondition(), &fctx.withSelf("v.getType()")) << ")) {\n"
<< " if (emit_error_on_verify_fail) {\n"
<< formatv(
" return op->emitOpError(\"{0} #\") << index "
"<< \" must be {1}, but got \" << v.getType();\n",
valueKind, desc)
<< " } else {\n"
<< " return failure();\n"
<< " }\n"
<< " }\n" // if
<< " ++index;\n"
<< " }\n"; // for
}
// Emit closing brace if needed.
if (!first) os << " }\n";
}
// NOLINTNEXTLINE
static bool RuntimeVerifierWriterMain(raw_ostream &os,
const RecordKeeper &records) {
emitSourceFileHeader("MLIR TFLite Runtime Verifiers", os);
// Retrieve all the definitions derived from TFL_Op and sort by record name.
std::vector<const Record *> defs = records.getAllDerivedDefinitions("Op");
llvm::sort(defs, LessRecord());
// Iterate through all the ops defined.
for (const auto *def : defs) {
mlir::tblgen::Operator op(*def);
if (!op.getTrait("TflRuntimeVerifyOpInterface::Trait")) continue;
mlir::tblgen::FmtContext verify_ctx;
os << "::mlir::LogicalResult " << op.getCppClassName()
<< "::VerifyTflRuntimeConstraints(::mlir::Operation *op, bool "
"emit_error_on_verify_fail) {\n";
os << " auto top = cast<" << op.getCppClassName() << ">(op); (void)top;\n";
verify_ctx.addSubst("_op", "(*op)");
for (int i = 0, e = op.getNumOperands(); i < e; ++i) {
auto &value = op.getOperand(i);
// Skip from first variadic operands for now. Else getOperand index used
// below doesn't match.
if (value.isVariableLength()) break;
if (!value.name.empty())
verify_ctx.addSubst(value.name, formatv("op->getOperand({0})", i));
}
for (int i = 0, e = op.getNumResults(); i < e; ++i) {
auto &value = op.getResult(i);
// Skip from first variadic results for now. Else getResult index used
// below doesn't match.
if (value.isVariableLength()) break;
if (!value.name.empty())
verify_ctx.addSubst(value.name, formatv("op->getResult({0})", i));
}
GenOperandResultVerifier(os, def->getValueAsDag("arguments")->getArgs(),
"operand");
GenOperandResultVerifier(os, def->getValueAsDag("results")->getArgs(),
"result");
for (auto &trait : op.getTraits()) {
if (!trait.getDef().isSubClassOf("GenInternalOpTrait")) {
continue;
}
if (trait.getDef().getValueAsString("trait") !=
"::mlir::OpTrait::TFLRuntimeOpTrait") {
continue;
}
auto *val = trait.getDef().getValue("tflRuntimePredicate");
if (!val) continue;
auto desc = trait.getDef().getValueAsString("tflRuntimeDescription");
mlir::tblgen::Pred pred(dyn_cast<llvm::DefInit>(val->getValue()));
os << tgfmt(
" if (!($0)) {\n "
" if (emit_error_on_verify_fail) {\n"
" return top.emitOpError(\"failed to verify that $1\");\n"
" } else {\n"
" return failure();\n }\n }\n",
&verify_ctx, tgfmt(pred.getCondition(), &verify_ctx), desc);
}
os << " if (!emit_error_on_verify_fail) {\n";
os << "// Ignore transient errors by registering an no-op handler.\n"
"// Applying legalization patterns will emit unwanted, transient \n"
"// errors when the replaced TFLite ops do not meet the sanity \n"
"// checks. \n"
"// In order to ignore the transient errors, the following lines \n"
"// override a diagnostic handler with an no-op handler only\n"
"// while this pass runs.\n"
"uint64_t current_thread_id = llvm::get_threadid();\n"
"ScopedDiagnosticHandler scoped_diag_handler(\n"
" top.getContext(), [&current_thread_id](Diagnostic&) -> "
"LogicalResult "
"{\n"
" // Consume only errors that are coming from the same thread "
"in order not\n"
" // to ignore errors from other passes that are running. Things\n"
" // running\n"
" // in the pass manager can be multi-threaded.\n"
" return success(current_thread_id == llvm::get_threadid());\n"
" });\n";
os << " return top.verifyInvariants();\n";
os << " } else {\n";
os << " return top.verifyInvariants();\n}\n";
os << "}\n";
}
return false;
}
int main(int argc, char **argv) {
llvm::InitLLVM y(argc, argv);
llvm::cl::ParseCommandLineOptions(argc, argv);
if (action == ActionType::OpConv)
return TableGenMain(argv[0], &OperatorWritersMain);
return TableGenMain(argv[0], &RuntimeVerifierWriterMain);
}
tensorflow/compiler/mlir/lite/core/BUILD
+55
View File
@@ -0,0 +1,55 @@
load("@rules_cc//cc:cc_library.bzl", "cc_library")
load("//tensorflow:tensorflow.default.bzl", "get_compatible_with_portable")
load("//tensorflow/compiler/mlir/lite:build_def.bzl", "tflite_copts_warnings")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:LICENSE"],
licenses = ["notice"],
)
exports_files(
[
"model_builder_base.h",
],
visibility = ["//tensorflow/lite/core:__pkg__"],
)
cc_library(
name = "macros",
hdrs = ["macros.h"],
compatible_with = get_compatible_with_portable(),
visibility = ["//visibility:public"],
)
cc_library(
name = "model_builder_base",
srcs = ["model_builder_base.cc"],
hdrs = ["model_builder_base.h"],
compatible_with = get_compatible_with_portable(),
copts = tflite_copts_warnings(),
visibility = ["//visibility:public"],
deps = [
":macros",
"//tensorflow/compiler/mlir/lite:allocation",
"//tensorflow/compiler/mlir/lite/core/api:error_reporter",
"//tensorflow/compiler/mlir/lite/core/api:verifier",
"//tensorflow/compiler/mlir/lite/schema:schema_fbs",
"@com_google_absl//absl/strings",
"@flatbuffers",
],
alwayslink = 1,
)
cc_library(
name = "absl_error_model_builder",
srcs = ["absl_error_model_builder.cc"],
hdrs = ["absl_error_model_builder.h"],
compatible_with = get_compatible_with_portable(),
copts = tflite_copts_warnings(),
visibility = ["//visibility:public"],
deps = [
":model_builder_base",
"//tensorflow/compiler/mlir/lite/core/api:error_reporter",
"@com_google_absl//absl/log",
],
)
tensorflow/compiler/mlir/lite/core/absl_error_model_builder.cc
+40
View File
@@ -0,0 +1,40 @@
/* Copyright 2024 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/core/absl_error_model_builder.h"
#include <cstdarg>
#include <cstdio>
#include "absl/log/log.h"
#include "tensorflow/compiler/mlir/lite/core/api/error_reporter.h"
namespace mlir::TFL {
int AbslErrorReporter::Report(const char* format, va_list args) {
char buffer[1024];
#pragma clang diagnostic push
#pragma clang diagnostic ignored "-Wformat-nonliteral"
vsprintf(buffer, format, args);
#pragma clang diagnostic pop
LOG(ERROR) << buffer;
return 0;
}
tflite::ErrorReporter* GetAbslErrorReporter() {
static AbslErrorReporter* error_reporter = new AbslErrorReporter;
return error_reporter;
}
} // namespace mlir::TFL
tensorflow/compiler/mlir/lite/core/absl_error_model_builder.h
+47
View File
@@ -0,0 +1,47 @@
/* Copyright 2024 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_COMPILER_MLIR_LITE_CORE_ABSL_ERROR_MODEL_BUILDER_H_
#define TENSORFLOW_COMPILER_MLIR_LITE_CORE_ABSL_ERROR_MODEL_BUILDER_H_
#include <cstdarg>
#include "tensorflow/compiler/mlir/lite/core/api/error_reporter.h"
#include "tensorflow/compiler/mlir/lite/core/model_builder_base.h"
namespace mlir::TFL {
// An error reporter that uses absl logging.
class AbslErrorReporter : public tflite::ErrorReporter {
int Report(const char* format, va_list args) override;
};
tflite::ErrorReporter* GetAbslErrorReporter();
class FlatBufferModelAbslError
: public tflite::impl::FlatBufferModelBase<FlatBufferModelAbslError> {
public:
// Use stderr_reporter as the default error reporter.
static tflite::ErrorReporter* GetDefaultErrorReporter() {
return GetAbslErrorReporter();
}
// Inherit all constructors from FlatBufferModelBase since inherited factory
// methods refer to them.
using FlatBufferModelBase<FlatBufferModelAbslError>::FlatBufferModelBase;
};
} // namespace mlir::TFL
#endif // TENSORFLOW_COMPILER_MLIR_LITE_CORE_ABSL_ERROR_MODEL_BUILDER_H_
tensorflow/compiler/mlir/lite/core/api/BUILD
+84
View File
@@ -0,0 +1,84 @@
load("@rules_cc//cc:cc_library.bzl", "cc_library")
load("//tensorflow:tensorflow.bzl", "tf_cc_test")
load("//tensorflow:tensorflow.default.bzl", "get_compatible_with_portable")
load("//tensorflow/compiler/mlir/lite:build_def.bzl", "tflite_copts")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:LICENSE"],
default_visibility = ["//visibility:public"],
licenses = ["notice"],
)
exports_files(["error_reporter.h"])
filegroup(
name = "tflite_internal_cc_3p_api_deps_src",
srcs = [
"error_reporter.cc",
"error_reporter.h",
"verifier.h",
],
visibility = ["//tensorflow/lite:__pkg__"],
)
cc_library(
name = "error_reporter",
srcs = ["error_reporter.cc"],
hdrs = ["error_reporter.h"],
compatible_with = get_compatible_with_portable(),
copts = tflite_copts(),
deps = [],
)
exports_files(["verifier.h"])
cc_library(
name = "verifier",
hdrs = ["verifier.h"],
compatible_with = get_compatible_with_portable(),
copts = tflite_copts(),
visibility = ["//visibility:public"],
deps = [":error_reporter"],
)
tf_cc_test(
name = "error_reporter_test",
size = "small",
srcs = ["error_reporter_test.cc"],
deps = [
":error_reporter",
"@com_google_googletest//:gtest_main",
],
)
cc_library(
name = "flatbuffer_conversions",
srcs = ["flatbuffer_conversions.cc"],
hdrs = [
"flatbuffer_conversions.h",
],
compatible_with = get_compatible_with_portable(),
copts = tflite_copts(),
deps = [
"//tensorflow/compiler/mlir/lite/core/c:tflite_common",
"//tensorflow/compiler/mlir/lite/kernels/internal:compatibility_macros",
"//tensorflow/compiler/mlir/lite/schema:schema_fbs",
"@com_google_absl//absl/log:absl_log",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings:str_format",
"@flatbuffers//:runtime_cc",
],
)
tf_cc_test(
name = "flatbuffer_conversions_test",
size = "small",
srcs = ["flatbuffer_conversions_test.cc"],
deps = [
":flatbuffer_conversions",
"//tensorflow/compiler/mlir/lite/core/c:tflite_common",
"//tensorflow/compiler/mlir/lite/schema:schema_fbs",
"@com_google_googletest//:gtest_main",
"@flatbuffers//:runtime_cc",
],
)
tensorflow/compiler/mlir/lite/core/api/error_reporter.cc
+39
View File
@@ -0,0 +1,39 @@
/* Copyright 2017 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/core/api/error_reporter.h"
#include <cstdarg>
namespace tflite {
int ErrorReporter::Report(const char* format, ...) {
va_list args;
va_start(args, format);
int code = Report(format, args);
va_end(args);
return code;
}
// TODO(aselle): Make the name of ReportError on context the same, so
// we can use the ensure functions w/o a context and w/ a reporter.
int ErrorReporter::ReportError(void*, const char* format, ...) {
va_list args;
va_start(args, format);
int code = Report(format, args);
va_end(args);
return code;
}
} // namespace tflite
tensorflow/compiler/mlir/lite/core/api/error_reporter.h
+72
View File
@@ -0,0 +1,72 @@
/* Copyright 2017 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_COMPILER_MLIR_LITE_CORE_API_ERROR_REPORTER_H_
#define TENSORFLOW_COMPILER_MLIR_LITE_CORE_API_ERROR_REPORTER_H_
#include <cstdarg>
namespace tflite {
/// A functor that reports error to supporting system. Invoked similar to
/// printf.
///
/// Usage:
/// ErrorReporter foo;
/// foo.Report("test %d", 5);
/// or
/// va_list args;
/// foo.Report("test %d", args); // where args is va_list
///
/// Subclass ErrorReporter to provide another reporting destination.
/// For example, if you have a GUI program, you might redirect to a buffer
/// that drives a GUI error log box.
class ErrorReporter {
public:
virtual ~ErrorReporter() = default;
/// Converts `args` to character equivalents according to `format` string,
/// constructs the error string and report it.
/// Returns number of characters written or zero on success, and negative
/// number on error.
virtual int Report(const char* format, va_list args) = 0;
/// Converts arguments to character equivalents according to `format` string,
/// constructs the error string and report it.
/// Returns number of characters written or zero on success, and negative
/// number on error.
int Report(const char* format, ...);
/// Equivalent to `Report` above. The additional `void*` parameter is unused.
/// This method is for compatibility with macros that takes `TfLiteContext`,
/// like TF_LITE_ENSURE and related macros.
int ReportError(void*, const char* format, ...);
};
} // namespace tflite
// You should not make bare calls to the error reporter, instead use the
// TF_LITE_REPORT_ERROR macro, since this allows message strings to be
// stripped when the binary size has to be optimized. If you are looking to
// reduce binary size, define TF_LITE_STRIP_ERROR_STRINGS when compiling and
// every call will be stubbed out, taking no memory.
#ifndef TF_LITE_STRIP_ERROR_STRINGS
#define TF_LITE_REPORT_ERROR(reporter, ...) \
do { \
static_cast<::tflite::ErrorReporter*>(reporter)->Report(__VA_ARGS__); \
} while (false)
#else // TF_LITE_STRIP_ERROR_STRINGS
#define TF_LITE_REPORT_ERROR(reporter, ...)
#endif // TF_LITE_STRIP_ERROR_STRINGS
#endif // TENSORFLOW_COMPILER_MLIR_LITE_CORE_API_ERROR_REPORTER_H_
tensorflow/compiler/mlir/lite/core/api/error_reporter_test.cc
+61
View File
@@ -0,0 +1,61 @@
/* 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/compiler/mlir/lite/core/api/error_reporter.h"
#include <cstdio>
#include <gtest/gtest.h>
namespace tflite {
class MockErrorReporter : public ErrorReporter {
public:
MockErrorReporter() { buffer_[0] = 0; }
int Report(const char* format, va_list args) override {
vsnprintf(buffer_, kBufferSize, format, args);
return 0;
}
char* GetBuffer() { return buffer_; }
private:
static constexpr int kBufferSize = 256;
char buffer_[kBufferSize];
};
TEST(ErrorReporter, TestReport) {
MockErrorReporter mock_reporter;
ErrorReporter* reporter = &mock_reporter;
reporter->Report("Error: %d", 23);
EXPECT_EQ(0, strcmp(mock_reporter.GetBuffer(), "Error: 23"));
}
TEST(ErrorReporter, TestReportMacro) {
MockErrorReporter mock_reporter;
// Only define the reporter if it's used, to avoid warnings.
#ifndef TF_LITE_STRIP_ERROR_STRINGS
ErrorReporter* reporter = &mock_reporter;
#endif // TFLITE_STRIP_ERROR_STRINGS
TF_LITE_REPORT_ERROR(reporter, "Error: %d", 23);
#ifndef TF_LITE_STRIP_ERROR_STRINGS
EXPECT_EQ(0, strcmp(mock_reporter.GetBuffer(), "Error: 23"));
#else // TF_LITE_STRIP_ERROR_STRINGS
EXPECT_EQ(0, strcmp(mock_reporter.GetBuffer(), ""));
#endif // TF_LITE_STRIP_ERROR_STRINGS
}
} // namespace tflite
tensorflow/compiler/mlir/lite/core/api/flatbuffer_conversions.cc
+3243
View File
File diff suppressed because it is too large Load Diff
tensorflow/compiler/mlir/lite/core/api/flatbuffer_conversions.h
+427
View File
@@ -0,0 +1,427 @@
/* Copyright 2024 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_COMPILER_MLIR_LITE_CORE_API_FLATBUFFER_CONVERSIONS_H_
#define TENSORFLOW_COMPILER_MLIR_LITE_CORE_API_FLATBUFFER_CONVERSIONS_H_
#include <cstddef>
#include <type_traits>
#include "absl/status/status.h"
#include "tensorflow/compiler/mlir/lite/core/c/tflite_types.h"
#include "tensorflow/compiler/mlir/lite/schema/schema_generated.h"
// The namespace tflite_file is for the data structures that define the .tflite
// file format, and code that is tightly coupled with those data structures.
// The .tflite file format is the serialized flatbuffer representation of
// computations on tensors that TF Lite uses for distribution of compiled ML
// models.
namespace tflite_file {
// This namespace contains functions that transform code and data structures
// that are defined in the flatbuffer serialization format into
// in-memory values that are used by the runtime API, interpreter and compiler.
namespace flatbuffer_conversions {
using tflite::Operator;
// Interface class for builtin data allocations.
class BuiltinDataAllocator {
public:
virtual void* Allocate(size_t size, size_t alignment_hint) = 0;
virtual void Deallocate(void* data) = 0;
// Allocate a structure, but make sure it is a POD structure that doesn't
// require constructors to run. The reason we do this, is that Interpreter's C
// extension part will take ownership so destructors will not be run during
// deallocation.
template <typename T>
T* AllocatePOD() {
static_assert(std::is_trivially_destructible<T>::value,
"Builtin data structure must be POD.");
void* allocated_memory = this->Allocate(sizeof(T), alignof(T));
return new (allocated_memory) T();
}
virtual ~BuiltinDataAllocator() = default;
};
// Parse the appropriate data out of the op.
//
// This handles builtin data explicitly as there are flatbuffer schemas.
// If it returns kTfLiteOk, it passes the data out with `builtin_data`. The
// calling function has to pass in an allocator object, and this allocator
// will be called to reserve space for the output data. If the calling
// function's allocator reserves memory on the heap, then it's the calling
// function's responsibility to free it.
// If it returns kTfLiteError, `builtin_data` will be `nullptr`.
absl::Status ParseOpData(const tflite::Operator* op,
tflite::BuiltinOperator op_type,
BuiltinDataAllocator* allocator, void** builtin_data);
// Converts the tensor data type used in the flat buffer to the representation
// used by the runtime.
absl::Status ConvertTensorType(tflite::TensorType tensor_type,
TfLiteType* type);
absl::Status ParseAbs(const Operator* op, BuiltinDataAllocator* allocator,
void** builtin_data);
absl::Status ParseAdd(const Operator* op, BuiltinDataAllocator* allocator,
void** builtin_data);
absl::Status ParseAddN(const Operator* op, BuiltinDataAllocator* allocator,
void** builtin_data);
absl::Status ParseArgMax(const Operator* op, BuiltinDataAllocator* allocator,
void** builtin_data);
absl::Status ParseArgMin(const Operator* op, BuiltinDataAllocator* allocator,
void** builtin_data);
absl::Status ParseAssignVariable(const Operator* op,
BuiltinDataAllocator* allocator,
void** builtin_data);
absl::Status ParseBatchMatMul(const Operator* op,
BuiltinDataAllocator* allocator,
void** builtin_data);
absl::Status ParseBatchToSpaceNd(const Operator* op,
BuiltinDataAllocator* allocator,
void** builtin_data);
absl::Status ParseBroadcastArgs(const Operator* op,
BuiltinDataAllocator* allocator,
void** builtin_data);
absl::Status ParseBroadcastTo(const Operator* op,
BuiltinDataAllocator* allocator,
void** builtin_data);
absl::Status ParseCallOnce(const Operator* op, BuiltinDataAllocator* allocator,
void** builtin_data);
absl::Status ParseCeil(const Operator* op, BuiltinDataAllocator* allocator,
void** builtin_data);
absl::Status ParseCast(const Operator* op, BuiltinDataAllocator* allocator,
void** builtin_data);
absl::Status ParseConcatenation(const Operator* op,
BuiltinDataAllocator* allocator,
void** builtin_data);
absl::Status ParseConv2D(const Operator* op, BuiltinDataAllocator* allocator,
void** builtin_data);
absl::Status ParseCos(const Operator* op, BuiltinDataAllocator* allocator,
void** builtin_data);
absl::Status ParseCumsum(const Operator* op, BuiltinDataAllocator* allocator,
void** builtin_data);
absl::Status ParseDepthToSpace(const Operator* op,
BuiltinDataAllocator* allocator,
void** builtin_data);
absl::Status ParseDepthwiseConv2D(const Operator* op,
BuiltinDataAllocator* allocator,
void** builtin_data);
absl::Status ParseDequantize(const Operator* op,
BuiltinDataAllocator* allocator,
void** builtin_data);
absl::Status ParseDiv(const Operator* op, BuiltinDataAllocator* allocator,
void** builtin_data);
absl::Status ParseElu(const Operator* op, BuiltinDataAllocator* allocator,
void** builtin_data);
absl::Status ParseEmbeddingLookup(const Operator* op,
BuiltinDataAllocator* allocator,
void** builtin_data);
absl::Status ParseEqual(const Operator* op, BuiltinDataAllocator* allocator,
void** builtin_data);
absl::Status ParseExp(const Operator* op, BuiltinDataAllocator* allocator,
void** builtin_data);
absl::Status ParseExpandDims(const Operator* op,
BuiltinDataAllocator* allocator,
void** builtin_data);
absl::Status ParseFill(const Operator* op, BuiltinDataAllocator* allocator,
void** builtin_data);
absl::Status ParseFloor(const Operator* op, BuiltinDataAllocator* allocator,
void** builtin_data);
absl::Status ParseFloorDiv(const Operator* op, BuiltinDataAllocator* allocator,
void** builtin_data);
absl::Status ParseFloorMod(const Operator* op, BuiltinDataAllocator* allocator,
void** builtin_data);
absl::Status ParseFullyConnected(const Operator* op,
BuiltinDataAllocator* allocator,
void** builtin_data);
absl::Status ParseGather(const Operator* op, BuiltinDataAllocator* allocator,
void** builtin_data);
absl::Status ParseGatherNd(const Operator* op, BuiltinDataAllocator* allocator,
void** builtin_data);
absl::Status ParseGreater(const Operator* op, BuiltinDataAllocator* allocator,
void** builtin_data);
absl::Status ParseGreaterEqual(const Operator* op,
BuiltinDataAllocator* allocator,
void** builtin_data);
absl::Status ParseHardSwish(const Operator* op, BuiltinDataAllocator* allocator,
void** builtin_data);
absl::Status ParseIf(const Operator* op, BuiltinDataAllocator* allocator,
void** builtin_data);
absl::Status ParseL2Normalization(const Operator* op,
BuiltinDataAllocator* allocator,
void** builtin_data);
absl::Status ParseLeakyRelu(const Operator* op, BuiltinDataAllocator* allocator,
void** builtin_data);
absl::Status ParseLess(const Operator* op, BuiltinDataAllocator* allocator,
void** builtin_data);
absl::Status ParseLessEqual(const Operator* op, BuiltinDataAllocator* allocator,
void** builtin_data);
absl::Status ParseLog(const Operator* op, BuiltinDataAllocator* allocator,
void** builtin_data);
absl::Status ParseLogicalAnd(const Operator* op,
BuiltinDataAllocator* allocator,
void** builtin_data);
absl::Status ParseLogicalNot(const Operator* op,
BuiltinDataAllocator* allocator,
void** builtin_data);
absl::Status ParseLogicalOr(const Operator* op, BuiltinDataAllocator* allocator,
void** builtin_data);
absl::Status ParseLogistic(const Operator* op, BuiltinDataAllocator* allocator,
void** builtin_data);
absl::Status ParseLogSoftmax(const Operator* op,
BuiltinDataAllocator* allocator,
void** builtin_data);
absl::Status ParseLSTM(const Operator* op, BuiltinDataAllocator* allocator,
void** builtin_data);
absl::Status ParseMaximum(const Operator* op, BuiltinDataAllocator* allocator,
void** builtin_data);
absl::Status ParseMinimum(const Operator* op, BuiltinDataAllocator* allocator,
void** builtin_data);
absl::Status ParseMirrorPad(const Operator* op, BuiltinDataAllocator* allocator,
void** builtin_data);
absl::Status ParseMul(const Operator* op, BuiltinDataAllocator* allocator,
void** builtin_data);
absl::Status ParseNeg(const Operator* op, BuiltinDataAllocator* allocator,
void** builtin_data);
absl::Status ParseNotEqual(const Operator* op, BuiltinDataAllocator* allocator,
void** builtin_data);
absl::Status ParsePack(const Operator* op, BuiltinDataAllocator* allocator,
void** builtin_data);
absl::Status ParsePad(const Operator* op, BuiltinDataAllocator* allocator,
void** builtin_data);
absl::Status ParsePadV2(const Operator* op, BuiltinDataAllocator* allocator,
void** builtin_data);
absl::Status ParsePool(const Operator* op, BuiltinDataAllocator* allocator,
void** builtin_data);
absl::Status ParsePow(const Operator* op, BuiltinDataAllocator* allocator,
void** builtin_data);
absl::Status ParsePrelu(const Operator* op, BuiltinDataAllocator* allocator,
void** builtin_data);
absl::Status ParseQuantize(const Operator* op, BuiltinDataAllocator* allocator,
void** builtin_data);
absl::Status ParseReadVariable(const Operator* op,
BuiltinDataAllocator* allocator,
void** builtin_data);
absl::Status ParseReducer(const Operator* op, BuiltinDataAllocator* allocator,
void** builtin_data);
absl::Status ParseRelu(const Operator* op, BuiltinDataAllocator* allocator,
void** builtin_data);
absl::Status ParseRelu6(const Operator* op, BuiltinDataAllocator* allocator,
void** builtin_data);
absl::Status ParseReshape(const Operator* op, BuiltinDataAllocator* allocator,
void** builtin_data);
absl::Status ParseResizeBilinear(const Operator* op,
BuiltinDataAllocator* allocator,
void** builtin_data);
absl::Status ParseResizeNearestNeighbor(const Operator* op,
BuiltinDataAllocator* allocator,
void** builtin_data);
absl::Status ParseRound(const Operator* op, BuiltinDataAllocator* allocator,
void** builtin_data);
absl::Status ParseRsqrt(const Operator* op, BuiltinDataAllocator* allocator,
void** builtin_data);
absl::Status ParseSelectV2(const Operator* op, BuiltinDataAllocator* allocator,
void** builtin_data);
absl::Status ParseShape(const Operator* op, BuiltinDataAllocator* allocator,
void** builtin_data);
absl::Status ParseSin(const Operator* op, BuiltinDataAllocator* allocator,
void** builtin_data);
absl::Status ParseSlice(const Operator* op, BuiltinDataAllocator* allocator,
void** builtin_data);
absl::Status ParseSoftmax(const Operator* op, BuiltinDataAllocator* allocator,
void** builtin_data);
absl::Status ParseSpaceToBatchNd(const Operator* op,
BuiltinDataAllocator* allocator,
void** builtin_data);
absl::Status ParseSpaceToDepth(const Operator* op,
BuiltinDataAllocator* allocator,
void** builtin_data);
absl::Status ParseSplit(const Operator* op, BuiltinDataAllocator* allocator,
void** builtin_data);
absl::Status ParseSplitV(const Operator* op, BuiltinDataAllocator* allocator,
void** builtin_data);
absl::Status ParseSqueeze(const Operator* op, BuiltinDataAllocator* allocator,
void** builtin_data);
absl::Status ParseSqrt(const Operator* op, BuiltinDataAllocator* allocator,
void** builtin_data);
absl::Status ParseSquare(const Operator* op, BuiltinDataAllocator* allocator,
void** builtin_data);
absl::Status ParseSquaredDifference(const Operator* op,
BuiltinDataAllocator* allocator,
void** builtin_data);
absl::Status ParseStridedSlice(const Operator* op,
BuiltinDataAllocator* allocator,
void** builtin_data);
absl::Status ParseSub(const Operator* op, BuiltinDataAllocator* allocator,
void** builtin_data);
absl::Status ParseSvdf(const Operator* op, BuiltinDataAllocator* allocator,
void** builtin_data);
absl::Status ParseTanh(const Operator* op, BuiltinDataAllocator* allocator,
void** builtin_data);
absl::Status ParseTranspose(const Operator* op, BuiltinDataAllocator* allocator,
void** builtin_data);
absl::Status ParseTransposeConv(const Operator* op,
BuiltinDataAllocator* allocator,
void** builtin_data);
absl::Status ParseUnpack(const Operator* op, BuiltinDataAllocator* allocator,
void** builtin_data);
absl::Status ParseUnidirectionalSequenceLSTM(const Operator* op,
BuiltinDataAllocator* allocator,
void** builtin_data);
absl::Status ParseVarHandle(const Operator* op, BuiltinDataAllocator* allocator,
void** builtin_data);
absl::Status ParseWhile(const Operator* op, BuiltinDataAllocator* allocator,
void** builtin_data);
absl::Status ParseZerosLike(const Operator* op, BuiltinDataAllocator* allocator,
void** builtin_data);
absl::Status ParseBitwiseXor(const Operator* op,
BuiltinDataAllocator* allocator,
void** builtin_data);
absl::Status ParseRightShift(const Operator* op,
BuiltinDataAllocator* allocator,
void** builtin_data);
absl::Status ParseStablehloScatter(const Operator* op,
BuiltinDataAllocator* allocator,
void** builtin_data);
absl::Status ParseStablehloRngBitGenerator(const Operator* op,
BuiltinDataAllocator* allocator,
void** builtin_data);
absl::Status ParseStablehloGather(const Operator* op,
BuiltinDataAllocator* allocator,
void** builtin_data);
absl::Status ParseStablehloReduceWindow(const Operator* op,
BuiltinDataAllocator* allocator,
void** builtin_data);
absl::Status ParseStablehloPad(const Operator* op,
BuiltinDataAllocator* allocator,
void** builtin_data);
absl::Status ParseStablehloComposite(const Operator* op,
BuiltinDataAllocator* allocator,
void** builtin_data);
absl::Status ParseStablehloShiftLeft(const Operator* op,
BuiltinDataAllocator* allocator,
void** builtin_data);
absl::Status ParseStablehloCase(const Operator* op,
BuiltinDataAllocator* allocator,
void** builtin_data);
} // namespace flatbuffer_conversions
} // namespace tflite_file
#endif // TENSORFLOW_COMPILER_MLIR_LITE_CORE_API_FLATBUFFER_CONVERSIONS_H_
tensorflow/compiler/mlir/lite/core/api/flatbuffer_conversions_test.cc
+875
View File
@@ -0,0 +1,875 @@
/* 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/compiler/mlir/lite/core/api/flatbuffer_conversions.h"
#include <cstdint>
#include <cstdio>
#include <cstring>
#include <string>
#include <tuple>
#include <vector>
#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include "flatbuffers/buffer.h" // from @flatbuffers
#include "flatbuffers/flatbuffer_builder.h" // from @flatbuffers
#include "tensorflow/compiler/mlir/lite/core/c/builtin_op_data.h"
#include "tensorflow/compiler/mlir/lite/core/c/tflite_types.h"
#include "tensorflow/compiler/mlir/lite/schema/schema_generated.h"
using testing::AllOf;
using testing::Each;
using testing::ElementsAre;
using testing::Eq;
using testing::HasSubstr;
using testing::StrEq;
using tflite::BuiltinOptions;
using tflite::BuiltinOptions2;
using tflite::BuiltinOptions_SqueezeOptions;
using tflite::CustomOptionsFormat_FLEXBUFFERS;
namespace tflite_file {
namespace flatbuffer_conversions {
using tflite::ActivationFunctionType_RELU;
using tflite::BuiltinOperator_CONV_2D;
using tflite::BuiltinOperator_CUSTOM;
using tflite::BuiltinOperator_FULLY_CONNECTED;
using tflite::BuiltinOperator_RESHAPE;
using tflite::BuiltinOperator_SQUEEZE;
using tflite::BuiltinOperator_STABLEHLO_PAD;
using tflite::BuiltinOperator_STABLEHLO_REDUCE_WINDOW;
using tflite::BuiltinOptions2_StablehloPadOptions;
using tflite::BuiltinOptions2_StablehloReduceWindowOptions;
using tflite::BuiltinOptions_Conv2DOptions;
using tflite::BuiltinOptions_FullyConnectedOptions;
using tflite::BuiltinOptions_NONE;
using tflite::BuiltinOptions_ReshapeOptions;
using tflite::CreateReshapeOptions;
using tflite::CreateSqueezeOptions;
using tflite::CreateStablehloPadOptions;
using tflite::CreateStablehloReduceWindowOptions;
using tflite::FullyConnectedOptionsWeightsFormat;
using tflite::Padding_SAME;
using tflite::TensorType_BFLOAT16;
using tflite::TensorType_FLOAT16;
using tflite::TensorType_FLOAT32;
using tflite::TensorType_INT4;
namespace {
using std::string;
// Used to determine how the op data parsing function creates its working space.
class MockDataAllocator : public BuiltinDataAllocator {
public:
MockDataAllocator() : is_allocated_(false) {}
void* Allocate(size_t size, size_t alignment_hint) override {
EXPECT_FALSE(is_allocated_);
const int max_size = kBufferSize;
EXPECT_LE(size, max_size);
is_allocated_ = true;
return buffer_;
}
void Deallocate(void* data) override { is_allocated_ = false; }
private:
static constexpr int kBufferSize = 1024;
char buffer_[kBufferSize];
bool is_allocated_;
};
} // namespace
class FlatbufferConversionsTest : public ::testing::Test {
public:
const Operator* BuildTestOperator(BuiltinOptions op_type,
flatbuffers::Offset<void> options) {
flatbuffers::Offset<Operator> offset =
CreateOperatorDirect(builder_, 0, nullptr, nullptr, op_type, options,
nullptr, CustomOptionsFormat_FLEXBUFFERS, nullptr);
builder_.Finish(offset);
void* pointer = builder_.GetBufferPointer();
return flatbuffers::GetRoot<Operator>(pointer);
}
const Operator* BuildTestOperator(BuiltinOptions2 op_type,
flatbuffers::Offset<void> options) {
flatbuffers::Offset<Operator> offset = CreateOperatorDirect(
builder_, /*opcode_index=*/0, /*inputs=*/nullptr, /*outputs=*/nullptr,
/*builtin_options_type=*/tflite::BuiltinOptions_NONE,
/*builtin_options=*/0, /*custom_options=*/nullptr,
/*custom_options_format=*/tflite::CustomOptionsFormat_FLEXBUFFERS,
/*mutating_variable_inputs=*/nullptr, /*intermediates=*/nullptr,
/*large_custom_options_offset=*/0, /*large_custom_options_size=*/0,
/*builtin_options_2_type=*/op_type,
/*builtin_options_2=*/options);
builder_.Finish(offset);
void* pointer = builder_.GetBufferPointer();
return flatbuffers::GetRoot<Operator>(pointer);
}
protected:
MockDataAllocator mock_allocator_;
flatbuffers::FlatBufferBuilder builder_;
};
TEST_F(FlatbufferConversionsTest, ParseSqueezeAll) {
const Operator* op = BuildTestOperator(
BuiltinOptions_SqueezeOptions, CreateSqueezeOptions(builder_).Union());
void* output_data = nullptr;
EXPECT_TRUE(
ParseOpData(op, BuiltinOperator_SQUEEZE, &mock_allocator_, &output_data)
.ok());
}
TEST_F(FlatbufferConversionsTest, ParseDynamicReshape) {
const Operator* op = BuildTestOperator(
BuiltinOptions_ReshapeOptions, CreateReshapeOptions(builder_).Union());
void* output_data = nullptr;
EXPECT_TRUE(
ParseOpData(op, BuiltinOperator_RESHAPE, &mock_allocator_, &output_data)
.ok());
}
TEST_F(FlatbufferConversionsTest, TestParseOpDataConv) {
const Operator* conv_op =
BuildTestOperator(BuiltinOptions_Conv2DOptions,
CreateConv2DOptions(builder_, Padding_SAME, 1, 2,
ActivationFunctionType_RELU, 3, 4)
.Union());
void* output_data = nullptr;
EXPECT_TRUE(ParseOpData(conv_op, BuiltinOperator_CONV_2D, &mock_allocator_,
&output_data)
.ok());
EXPECT_NE(nullptr, output_data);
TfLiteConvParams* params = reinterpret_cast<TfLiteConvParams*>(output_data);
EXPECT_EQ(kTfLitePaddingSame, params->padding);
EXPECT_EQ(1, params->stride_width);
EXPECT_EQ(2, params->stride_height);
EXPECT_EQ(kTfLiteActRelu, params->activation);
EXPECT_EQ(3, params->dilation_width_factor);
EXPECT_EQ(4, params->dilation_height_factor);
}
TEST_F(FlatbufferConversionsTest, ParseBadFullyConnected) {
const Operator* conv_op = BuildTestOperator(
BuiltinOptions_FullyConnectedOptions,
CreateFullyConnectedOptions(
builder_, ActivationFunctionType_RELU,
static_cast<FullyConnectedOptionsWeightsFormat>(-1), true)
.Union());
void* output_data = nullptr;
EXPECT_FALSE(ParseOpData(conv_op, BuiltinOperator_FULLY_CONNECTED,
&mock_allocator_, &output_data)
.ok());
}
TEST_F(FlatbufferConversionsTest, TestParseOpDataCustom) {
const Operator* custom_op =
BuildTestOperator(BuiltinOptions_NONE, flatbuffers::Offset<void>());
void* output_data = nullptr;
EXPECT_TRUE(ParseOpData(custom_op, BuiltinOperator_CUSTOM, &mock_allocator_,
&output_data)
.ok());
EXPECT_EQ(nullptr, output_data);
}
TEST_F(FlatbufferConversionsTest, TestConvertTensorType) {
TfLiteType type;
EXPECT_TRUE(ConvertTensorType(TensorType_FLOAT32, &type).ok());
EXPECT_EQ(kTfLiteFloat32, type);
}
TEST_F(FlatbufferConversionsTest, TestConvertTensorTypeFloat16) {
TfLiteType type;
EXPECT_TRUE(ConvertTensorType(TensorType_FLOAT16, &type).ok());
EXPECT_EQ(kTfLiteFloat16, type);
}
TEST_F(FlatbufferConversionsTest, TestConvertTensorTypeBFloat16) {
TfLiteType type;
EXPECT_TRUE(ConvertTensorType(TensorType_BFLOAT16, &type).ok());
EXPECT_EQ(kTfLiteBFloat16, type);
}
TEST_F(FlatbufferConversionsTest, TestConvertTensorTypeInt4) {
TfLiteType type;
EXPECT_TRUE(ConvertTensorType(TensorType_INT4, &type).ok());
EXPECT_EQ(kTfLiteInt4, type);
}
class StablehloReduceWindowFlatbufferConversionsTest
: public FlatbufferConversionsTest {
public:
static constexpr int kMaxDims =
TFLITE_STABLEHLO_REDUCE_WINDOW_PARAMS_MAX_DIMENSION_COUNT;
static constexpr int64_t kValidValue = 5;
auto ValidAttr() {
return builder_.CreateVector(std::vector<int64_t>(kMaxDims, kValidValue));
}
auto InvalidAttr() {
return builder_.CreateVector(
std::vector<int64_t>(kMaxDims + 1, kValidValue));
}
auto ValidPaddingAttr() {
return builder_.CreateVector(
std::vector<int64_t>(2 * kMaxDims, kValidValue));
}
auto InvalidPaddingAttr() {
return builder_.CreateVector(
std::vector<int64_t>(2 * kMaxDims + 1, kValidValue));
}
auto EmptyAttr() { return builder_.CreateVector<int64_t>({}); }
};
TEST_F(StablehloReduceWindowFlatbufferConversionsTest, Succeeds) {
const Operator* stablehlo_reduce_window_op = BuildTestOperator(
BuiltinOptions2_StablehloReduceWindowOptions,
CreateStablehloReduceWindowOptions(
builder_,
/*window_dimensions=*/builder_.CreateVector<int64_t>({1, 2}),
/*window_strides=*/builder_.CreateVector<int64_t>({3, 4}),
/*base_dilations=*/builder_.CreateVector<int64_t>({5, 6}),
/*window_dilations=*/builder_.CreateVector<int64_t>({7, 8}),
/*padding=*/builder_.CreateVector<int64_t>({9, 10, 11, 12}),
/*body_subgraph_index=*/13)
.Union());
TfLiteStablehloReduceWindowParams* output_data = nullptr;
EXPECT_TRUE(ParseOpData(stablehlo_reduce_window_op,
BuiltinOperator_STABLEHLO_REDUCE_WINDOW,
&mock_allocator_, (void**)&output_data)
.ok());
EXPECT_THAT(std::make_tuple(output_data->window_dimensions, 2),
ElementsAre(1, 2));
EXPECT_THAT(std::make_tuple(output_data->window_strides, 2),
ElementsAre(3, 4));
EXPECT_THAT(std::make_tuple(output_data->base_dilations, 2),
ElementsAre(5, 6));
EXPECT_THAT(std::make_tuple(output_data->window_dilations, 2),
ElementsAre(7, 8));
EXPECT_THAT(std::make_tuple(output_data->padding, 4),
ElementsAre(9, 10, 11, 12));
EXPECT_THAT(output_data->body_subgraph_index, Eq(13));
}
TEST_F(StablehloReduceWindowFlatbufferConversionsTest,
FailsWithNoWindowDimensions) {
TfLiteStablehloReduceWindowParams* output_data = nullptr;
auto status = ParseOpData(
BuildTestOperator(
BuiltinOptions2_StablehloReduceWindowOptions,
CreateStablehloReduceWindowOptions(builder_,
/*window_dimensions=*/0,
/*window_strides=*/ValidAttr(),
/*base_dilations=*/ValidAttr(),
/*window_dilations=*/ValidAttr(),
/*padding=*/ValidPaddingAttr(),
/*body_subgraph_index=*/13)
.Union()),
BuiltinOperator_STABLEHLO_REDUCE_WINDOW, &mock_allocator_,
(void**)&output_data);
EXPECT_FALSE(status.ok());
EXPECT_THAT(status.message(),
HasSubstr("'window_dimensions' attribute is not optional for "
"'stablehlo.reduce_window' and cannot be empty."));
}
TEST_F(StablehloReduceWindowFlatbufferConversionsTest,
SucceedsWithNoWindowStrides) {
TfLiteStablehloReduceWindowParams* output_data = nullptr;
auto status = ParseOpData(
BuildTestOperator(
BuiltinOptions2_StablehloReduceWindowOptions,
CreateStablehloReduceWindowOptions(builder_,
/*window_dimensions=*/ValidAttr(),
/*window_strides=*/0,
/*base_dilations=*/ValidAttr(),
/*window_dilations=*/ValidAttr(),
/*padding=*/ValidPaddingAttr(),
/*body_subgraph_index=*/13)
.Union()),
BuiltinOperator_STABLEHLO_REDUCE_WINDOW, &mock_allocator_,
(void**)&output_data);
EXPECT_TRUE(status.ok());
EXPECT_THAT(status.message(), StrEq(""));
EXPECT_THAT(std::make_tuple(output_data->window_dimensions, kMaxDims),
Each(kValidValue));
EXPECT_THAT(std::make_tuple(output_data->window_strides, kMaxDims), Each(1));
EXPECT_THAT(std::make_tuple(output_data->base_dilations, kMaxDims),
Each(kValidValue));
EXPECT_THAT(std::make_tuple(output_data->window_dilations, kMaxDims),
Each(kValidValue));
EXPECT_THAT(std::make_tuple(output_data->padding, 2 * kMaxDims),
Each(kValidValue));
EXPECT_THAT(output_data->body_subgraph_index, Eq(13));
}
TEST_F(StablehloReduceWindowFlatbufferConversionsTest,
SucceedsWithNoBaseDilations) {
TfLiteStablehloReduceWindowParams* output_data = nullptr;
auto status = ParseOpData(
BuildTestOperator(
BuiltinOptions2_StablehloReduceWindowOptions,
CreateStablehloReduceWindowOptions(builder_,
/*window_dimensions=*/ValidAttr(),
/*window_strides=*/ValidAttr(),
/*base_dilations=*/0,
/*window_dilations=*/ValidAttr(),
/*padding=*/ValidPaddingAttr(),
/*body_subgraph_index=*/13)
.Union()),
BuiltinOperator_STABLEHLO_REDUCE_WINDOW, &mock_allocator_,
(void**)&output_data);
EXPECT_TRUE(status.ok());
EXPECT_THAT(status.message(), StrEq(""));
EXPECT_THAT(std::make_tuple(output_data->window_dimensions, kMaxDims),
Each(kValidValue));
EXPECT_THAT(std::make_tuple(output_data->window_strides, kMaxDims),
Each(kValidValue));
EXPECT_THAT(std::make_tuple(output_data->base_dilations, kMaxDims), Each(1));
EXPECT_THAT(std::make_tuple(output_data->window_dilations, kMaxDims),
Each(kValidValue));
EXPECT_THAT(std::make_tuple(output_data->padding, 2 * kMaxDims),
Each(kValidValue));
EXPECT_THAT(output_data->body_subgraph_index, Eq(13));
}
TEST_F(StablehloReduceWindowFlatbufferConversionsTest,
SucceedsWithNoWindowDilations) {
TfLiteStablehloReduceWindowParams* output_data = nullptr;
auto status = ParseOpData(
BuildTestOperator(
BuiltinOptions2_StablehloReduceWindowOptions,
CreateStablehloReduceWindowOptions(builder_,
/*window_dimensions=*/ValidAttr(),
/*window_strides=*/ValidAttr(),
/*base_dilations=*/ValidAttr(),
/*window_dilations=*/0,
/*padding=*/ValidPaddingAttr(),
/*body_subgraph_index=*/13)
.Union()),
BuiltinOperator_STABLEHLO_REDUCE_WINDOW, &mock_allocator_,
(void**)&output_data);
EXPECT_TRUE(status.ok());
EXPECT_THAT(status.message(), StrEq(""));
EXPECT_THAT(std::make_tuple(output_data->window_dimensions, kMaxDims),
Each(kValidValue));
EXPECT_THAT(std::make_tuple(output_data->window_strides, kMaxDims),
Each(kValidValue));
EXPECT_THAT(std::make_tuple(output_data->base_dilations, kMaxDims),
Each(kValidValue));
EXPECT_THAT(std::make_tuple(output_data->window_dilations, kMaxDims),
Each(1));
EXPECT_THAT(std::make_tuple(output_data->padding, 2 * kMaxDims),
Each(kValidValue));
EXPECT_THAT(output_data->body_subgraph_index, Eq(13));
}
TEST_F(StablehloReduceWindowFlatbufferConversionsTest, SucceedsWithNoPadding) {
TfLiteStablehloReduceWindowParams* output_data = nullptr;
auto status = ParseOpData(
BuildTestOperator(
BuiltinOptions2_StablehloReduceWindowOptions,
CreateStablehloReduceWindowOptions(builder_,
/*window_dimensions=*/ValidAttr(),
/*window_strides=*/ValidAttr(),
/*base_dilations=*/ValidAttr(),
/*window_dilations=*/ValidAttr(),
/*padding=*/0,
/*body_subgraph_index=*/13)
.Union()),
BuiltinOperator_STABLEHLO_REDUCE_WINDOW, &mock_allocator_,
(void**)&output_data);
EXPECT_TRUE(status.ok());
EXPECT_THAT(status.message(), StrEq(""));
EXPECT_THAT(std::make_tuple(output_data->window_dimensions, kMaxDims),
Each(kValidValue));
EXPECT_THAT(std::make_tuple(output_data->window_strides, kMaxDims),
Each(kValidValue));
EXPECT_THAT(std::make_tuple(output_data->base_dilations, kMaxDims),
Each(kValidValue));
EXPECT_THAT(std::make_tuple(output_data->window_dilations, kMaxDims),
Each(kValidValue));
EXPECT_THAT(std::make_tuple(output_data->padding, 2 * kMaxDims), Each(0));
EXPECT_THAT(output_data->body_subgraph_index, Eq(13));
}
TEST_F(StablehloReduceWindowFlatbufferConversionsTest,
FailsWithEmptyWindowDimensions) {
TfLiteStablehloReduceWindowParams* output_data = nullptr;
auto status = ParseOpData(
BuildTestOperator(
BuiltinOptions2_StablehloReduceWindowOptions,
CreateStablehloReduceWindowOptions(builder_,
/*window_dimensions=*/EmptyAttr(),
/*window_strides=*/ValidAttr(),
/*base_dilations=*/ValidAttr(),
/*window_dilations=*/ValidAttr(),
/*padding=*/ValidPaddingAttr(),
/*body_subgraph_index=*/13)
.Union()),
BuiltinOperator_STABLEHLO_REDUCE_WINDOW, &mock_allocator_,
(void**)&output_data);
EXPECT_FALSE(status.ok());
EXPECT_THAT(status.message(),
HasSubstr("'window_dimensions' attribute is not optional for "
"'stablehlo.reduce_window' and cannot be empty."));
}
TEST_F(StablehloReduceWindowFlatbufferConversionsTest,
SucceedsWithEmptyWindowStrides) {
TfLiteStablehloReduceWindowParams* output_data = nullptr;
auto status = ParseOpData(
BuildTestOperator(
BuiltinOptions2_StablehloReduceWindowOptions,
CreateStablehloReduceWindowOptions(builder_,
/*window_dimensions=*/ValidAttr(),
/*window_strides=*/EmptyAttr(),
/*base_dilations=*/ValidAttr(),
/*window_dilations=*/ValidAttr(),
/*padding=*/ValidPaddingAttr(),
/*body_subgraph_index=*/13)
.Union()),
BuiltinOperator_STABLEHLO_REDUCE_WINDOW, &mock_allocator_,
(void**)&output_data);
EXPECT_TRUE(status.ok());
EXPECT_THAT(status.message(), StrEq(""));
EXPECT_THAT(std::make_tuple(output_data->window_dimensions, kMaxDims),
Each(kValidValue));
EXPECT_THAT(std::make_tuple(output_data->window_strides, kMaxDims), Each(1));
EXPECT_THAT(std::make_tuple(output_data->base_dilations, kMaxDims),
Each(kValidValue));
EXPECT_THAT(std::make_tuple(output_data->window_dilations, kMaxDims),
Each(kValidValue));
EXPECT_THAT(std::make_tuple(output_data->padding, 2 * kMaxDims),
Each(kValidValue));
EXPECT_THAT(output_data->body_subgraph_index, Eq(13));
}
TEST_F(StablehloReduceWindowFlatbufferConversionsTest,
SucceedsWithEmptyBaseDilations) {
TfLiteStablehloReduceWindowParams* output_data = nullptr;
auto status = ParseOpData(
BuildTestOperator(
BuiltinOptions2_StablehloReduceWindowOptions,
CreateStablehloReduceWindowOptions(builder_,
/*window_dimensions=*/ValidAttr(),
/*window_strides=*/ValidAttr(),
/*base_dilations=*/EmptyAttr(),
/*window_dilations=*/ValidAttr(),
/*padding=*/ValidPaddingAttr(),
/*body_subgraph_index=*/13)
.Union()),
BuiltinOperator_STABLEHLO_REDUCE_WINDOW, &mock_allocator_,
(void**)&output_data);
EXPECT_TRUE(status.ok());
EXPECT_THAT(status.message(), StrEq(""));
EXPECT_THAT(std::make_tuple(output_data->window_dimensions, kMaxDims),
Each(kValidValue));
EXPECT_THAT(std::make_tuple(output_data->window_strides, kMaxDims),
Each(kValidValue));
EXPECT_THAT(std::make_tuple(output_data->base_dilations, kMaxDims), Each(1));
EXPECT_THAT(std::make_tuple(output_data->window_dilations, kMaxDims),
Each(kValidValue));
EXPECT_THAT(std::make_tuple(output_data->padding, 2 * kMaxDims),
Each(kValidValue));
EXPECT_THAT(output_data->body_subgraph_index, Eq(13));
}
TEST_F(StablehloReduceWindowFlatbufferConversionsTest,
SucceedsWithEmptyWindowDilations) {
TfLiteStablehloReduceWindowParams* output_data = nullptr;
auto status = ParseOpData(
BuildTestOperator(
BuiltinOptions2_StablehloReduceWindowOptions,
CreateStablehloReduceWindowOptions(builder_,
/*window_dimensions=*/ValidAttr(),
/*window_strides=*/ValidAttr(),
/*base_dilations=*/ValidAttr(),
/*window_dilations=*/EmptyAttr(),
/*padding=*/ValidPaddingAttr(),
/*body_subgraph_index=*/13)
.Union()),
BuiltinOperator_STABLEHLO_REDUCE_WINDOW, &mock_allocator_,
(void**)&output_data);
EXPECT_TRUE(status.ok());
EXPECT_THAT(status.message(), StrEq(""));
EXPECT_THAT(std::make_tuple(output_data->window_dimensions, kMaxDims),
Each(kValidValue));
EXPECT_THAT(std::make_tuple(output_data->window_strides, kMaxDims),
Each(kValidValue));
EXPECT_THAT(std::make_tuple(output_data->base_dilations, kMaxDims),
Each(kValidValue));
EXPECT_THAT(std::make_tuple(output_data->window_dilations, kMaxDims),
Each(1));
EXPECT_THAT(std::make_tuple(output_data->padding, 2 * kMaxDims),
Each(kValidValue));
EXPECT_THAT(output_data->body_subgraph_index, Eq(13));
}
TEST_F(StablehloReduceWindowFlatbufferConversionsTest,
SucceedsWithEmptyPadding) {
TfLiteStablehloReduceWindowParams* output_data = nullptr;
auto status = ParseOpData(
BuildTestOperator(
BuiltinOptions2_StablehloReduceWindowOptions,
CreateStablehloReduceWindowOptions(builder_,
/*window_dimensions=*/ValidAttr(),
/*window_strides=*/ValidAttr(),
/*base_dilations=*/ValidAttr(),
/*window_dilations=*/ValidAttr(),
/*padding=*/EmptyAttr(),
/*body_subgraph_index=*/13)
.Union()),
BuiltinOperator_STABLEHLO_REDUCE_WINDOW, &mock_allocator_,
(void**)&output_data);
EXPECT_TRUE(status.ok());
EXPECT_THAT(status.message(), StrEq(""));
EXPECT_THAT(std::make_tuple(output_data->window_dimensions, kMaxDims),
Each(kValidValue));
EXPECT_THAT(std::make_tuple(output_data->window_strides, kMaxDims),
Each(kValidValue));
EXPECT_THAT(std::make_tuple(output_data->base_dilations, kMaxDims),
Each(kValidValue));
EXPECT_THAT(std::make_tuple(output_data->window_dilations, kMaxDims),
Each(kValidValue));
EXPECT_THAT(std::make_tuple(output_data->padding, 2 * kMaxDims), Each(0));
EXPECT_THAT(output_data->body_subgraph_index, Eq(13));
}
TEST_F(StablehloReduceWindowFlatbufferConversionsTest,
SucceedsWithParamsAtMaxDims) {
TfLiteStablehloReduceWindowParams* output_data = nullptr;
auto status = ParseOpData(
BuildTestOperator(
BuiltinOptions2_StablehloReduceWindowOptions,
CreateStablehloReduceWindowOptions(builder_,
/*window_dimensions=*/ValidAttr(),
/*window_strides=*/ValidAttr(),
/*base_dilations=*/ValidAttr(),
/*window_dilations=*/ValidAttr(),
/*padding=*/ValidPaddingAttr(),
/*body_subgraph_index=*/13)
.Union()),
BuiltinOperator_STABLEHLO_REDUCE_WINDOW, &mock_allocator_,
(void**)&output_data);
EXPECT_TRUE(status.ok());
EXPECT_THAT(status.message(), StrEq(""));
}
TEST_F(StablehloReduceWindowFlatbufferConversionsTest,
FailsWhenWindowDimensionsHasMoreThanMaxDims) {
TfLiteStablehloReduceWindowParams* output_data = nullptr;
auto status = ParseOpData(
BuildTestOperator(BuiltinOptions2_StablehloReduceWindowOptions,
CreateStablehloReduceWindowOptions(
builder_,
/*window_dimensions=*/InvalidAttr(),
/*window_strides=*/ValidAttr(),
/*base_dilations=*/ValidAttr(),
/*window_dilations=*/ValidAttr(),
/*padding=*/ValidPaddingAttr(),
/*body_subgraph_index=*/13)
.Union()),
BuiltinOperator_STABLEHLO_REDUCE_WINDOW, &mock_allocator_,
(void**)&output_data);
EXPECT_FALSE(status.ok());
EXPECT_THAT(status.message(),
AllOf(HasSubstr("Found too many dimensions in the input array of "
"operation 'stablehlo.reduce_window'."),
HasSubstr("Check the 'window_dimensions' attribute.")));
}
TEST_F(StablehloReduceWindowFlatbufferConversionsTest,
FailsWhenWindowStridesHasWrongDimCount) {
TfLiteStablehloReduceWindowParams* output_data = nullptr;
auto status = ParseOpData(
BuildTestOperator(
BuiltinOptions2_StablehloReduceWindowOptions,
CreateStablehloReduceWindowOptions(builder_,
/*window_dimensions=*/ValidAttr(),
/*window_strides=*/InvalidAttr(),
/*base_dilations=*/ValidAttr(),
/*window_dilations=*/ValidAttr(),
/*padding=*/ValidPaddingAttr(),
/*body_subgraph_index=*/13)
.Union()),
BuiltinOperator_STABLEHLO_REDUCE_WINDOW, &mock_allocator_,
(void**)&output_data);
EXPECT_FALSE(status.ok());
EXPECT_THAT(
status.message(),
HasSubstr("'window_strides' attribute of 'stablehlo.reduce_window' does "
"not have the expected size"));
}
TEST_F(StablehloReduceWindowFlatbufferConversionsTest,
FailsWhenBaseDilationsHasWrongDimCount) {
TfLiteStablehloReduceWindowParams* output_data = nullptr;
auto status = ParseOpData(
BuildTestOperator(
BuiltinOptions2_StablehloReduceWindowOptions,
CreateStablehloReduceWindowOptions(builder_,
/*window_dimensions=*/ValidAttr(),
/*window_strides=*/ValidAttr(),
/*base_dilations=*/InvalidAttr(),
/*window_dilations=*/ValidAttr(),
/*padding=*/ValidPaddingAttr(),
/*body_subgraph_index=*/13)
.Union()),
BuiltinOperator_STABLEHLO_REDUCE_WINDOW, &mock_allocator_,
(void**)&output_data);
EXPECT_FALSE(status.ok());
EXPECT_THAT(
status.message(),
HasSubstr("'base_dilations' attribute of 'stablehlo.reduce_window' does "
"not have the expected size"));
}
TEST_F(StablehloReduceWindowFlatbufferConversionsTest,
FailsWhenWindowDilationsHasWrongDimCount) {
TfLiteStablehloReduceWindowParams* output_data = nullptr;
auto status = ParseOpData(
BuildTestOperator(
BuiltinOptions2_StablehloReduceWindowOptions,
CreateStablehloReduceWindowOptions(builder_,
/*window_dimensions=*/ValidAttr(),
/*window_strides=*/ValidAttr(),
/*base_dilations=*/ValidAttr(),
/*window_dilations=*/InvalidAttr(),
/*padding=*/ValidPaddingAttr(),
/*body_subgraph_index=*/13)
.Union()),
BuiltinOperator_STABLEHLO_REDUCE_WINDOW, &mock_allocator_,
(void**)&output_data);
EXPECT_FALSE(status.ok());
EXPECT_THAT(
status.message(),
HasSubstr(
"'window_dilations' attribute of 'stablehlo.reduce_window' does "
"not have the expected size"));
}
TEST_F(StablehloReduceWindowFlatbufferConversionsTest,
FailsWhenPaddingHasWrongDimCount) {
TfLiteStablehloReduceWindowParams* output_data = nullptr;
auto status = ParseOpData(
BuildTestOperator(
BuiltinOptions2_StablehloReduceWindowOptions,
CreateStablehloReduceWindowOptions(builder_,
/*window_dimensions=*/ValidAttr(),
/*window_strides=*/ValidAttr(),
/*base_dilations=*/ValidAttr(),
/*window_dilations=*/ValidAttr(),
/*padding=*/InvalidPaddingAttr(),
/*body_subgraph_index=*/13)
.Union()),
BuiltinOperator_STABLEHLO_REDUCE_WINDOW, &mock_allocator_,
(void**)&output_data);
EXPECT_FALSE(status.ok());
EXPECT_THAT(status.message(),
HasSubstr("'padding' attribute of 'stablehlo.reduce_window' does "
"not have the expected size"));
}
TEST_F(StablehloReduceWindowFlatbufferConversionsTest, FailsWithWrongOptions) {
const Operator* stablehlo_reduce_window_op =
BuildTestOperator(BuiltinOptions2_StablehloReduceWindowOptions, 0);
TfLiteStablehloReduceWindowParams* output_data = nullptr;
auto status = ParseOpData(stablehlo_reduce_window_op,
BuiltinOperator_STABLEHLO_REDUCE_WINDOW,
&mock_allocator_, (void**)&output_data);
EXPECT_FALSE(status.ok());
EXPECT_THAT(
status.message(),
HasSubstr(
"Could not get 'stablehlo.reduce_window' operation parameters."));
}
TEST_F(StablehloReduceWindowFlatbufferConversionsTest, DeathTests) {
const Operator* stablehlo_reduce_window_op = BuildTestOperator(
BuiltinOptions2_StablehloReduceWindowOptions,
CreateStablehloReduceWindowOptions(
builder_, /*window_dimensions=*/ValidAttr(),
/*window_strides=*/ValidAttr(),
/*base_dilations=*/ValidAttr(),
/*window_dilations=*/ValidAttr(),
/*padding=*/ValidPaddingAttr(), /*body_subgraph_index=*/13)
.Union());
TfLiteStablehloReduceWindowParams* output_data = nullptr;
#ifdef NDEBUG
GTEST_SKIP();
#endif
EXPECT_DEATH(ParseOpData(nullptr, BuiltinOperator_STABLEHLO_REDUCE_WINDOW,
&mock_allocator_, (void**)&output_data)
.IgnoreError(),
"");
EXPECT_DEATH(ParseOpData(stablehlo_reduce_window_op,
BuiltinOperator_STABLEHLO_REDUCE_WINDOW, nullptr,
(void**)&output_data)
.IgnoreError(),
"");
EXPECT_DEATH(ParseOpData(stablehlo_reduce_window_op,
BuiltinOperator_STABLEHLO_REDUCE_WINDOW,
&mock_allocator_, nullptr)
.IgnoreError(),
"");
}
class StablehloPadFlatbufferConversionsTest : public FlatbufferConversionsTest {
public:
static constexpr int kMaxDims =
TFLITE_STABLEHLO_PAD_PARAMS_MAX_DIMENSION_COUNT;
static constexpr int64_t kValidValue = 5;
};
TEST_F(StablehloPadFlatbufferConversionsTest, Succeeds) {
const Operator* stablehlo_pad_op = BuildTestOperator(
BuiltinOptions2_StablehloPadOptions,
CreateStablehloPadOptions(
builder_,
/*edge_padding_low=*/builder_.CreateVector<int64_t>({1, 0, -1}),
/*edge_padding_high=*/builder_.CreateVector<int64_t>({2, 0, -2}),
/*interior_padding=*/builder_.CreateVector<int64_t>({3, 0, 3}))
.Union());
TfLiteStablehloPadParams* output_data = nullptr;
EXPECT_TRUE(ParseOpData(stablehlo_pad_op, BuiltinOperator_STABLEHLO_PAD,
&mock_allocator_, (void**)&output_data)
.ok());
EXPECT_THAT(std::make_tuple(output_data->edge_padding_low, 3),
ElementsAre(1, 0, -1));
EXPECT_THAT(std::make_tuple(output_data->edge_padding_high, 3),
ElementsAre(2, 0, -2));
EXPECT_THAT(std::make_tuple(output_data->interior_padding, 3),
ElementsAre(3, 0, 3));
}
TEST_F(StablehloPadFlatbufferConversionsTest, FailsWithMissingLowPadding) {
const Operator* stablehlo_pad_op = BuildTestOperator(
BuiltinOptions2_StablehloPadOptions,
CreateStablehloPadOptions(
builder_,
/*edge_padding_low=*/0,
/*edge_padding_high=*/builder_.CreateVector<int64_t>({2, 0, -2}),
/*interior_padding=*/builder_.CreateVector<int64_t>({3, 0, 3}))
.Union());
TfLiteStablehloPadParams* output_data = nullptr;
auto status = ParseOpData(stablehlo_pad_op, BuiltinOperator_STABLEHLO_PAD,
&mock_allocator_, (void**)&output_data);
EXPECT_FALSE(status.ok());
EXPECT_THAT(
status.message(),
AllOf(
HasSubstr("Input array not provided for operation 'stablehlo.pad'."),
HasSubstr("Check the 'edge_padding_low' attribute.")));
}
TEST_F(StablehloPadFlatbufferConversionsTest, FailsWithMissingHighPadding) {
const Operator* stablehlo_pad_op = BuildTestOperator(
BuiltinOptions2_StablehloPadOptions,
CreateStablehloPadOptions(
builder_,
/*edge_padding_low=*/builder_.CreateVector<int64_t>({1, 0, -1}),
/*edge_padding_high=*/0,
/*interior_padding=*/builder_.CreateVector<int64_t>({3, 0, 3}))
.Union());
TfLiteStablehloPadParams* output_data = nullptr;
auto status = ParseOpData(stablehlo_pad_op, BuiltinOperator_STABLEHLO_PAD,
&mock_allocator_, (void**)&output_data);
EXPECT_FALSE(status.ok());
EXPECT_THAT(
status.message(),
AllOf(
HasSubstr("Input array not provided for operation 'stablehlo.pad'."),
HasSubstr("Check the 'edge_padding_high' attribute.")));
}
TEST_F(StablehloPadFlatbufferConversionsTest, FailsWithMissingInteriorPadding) {
const Operator* stablehlo_pad_op = BuildTestOperator(
BuiltinOptions2_StablehloPadOptions,
CreateStablehloPadOptions(
builder_,
/*edge_padding_low=*/builder_.CreateVector<int64_t>({1, 0, -1}),
/*edge_padding_high=*/builder_.CreateVector<int64_t>({2, 0, -2}),
/*interior_padding=*/0)
.Union());
TfLiteStablehloPadParams* output_data = nullptr;
auto status = ParseOpData(stablehlo_pad_op, BuiltinOperator_STABLEHLO_PAD,
&mock_allocator_, (void**)&output_data);
EXPECT_FALSE(status.ok());
EXPECT_THAT(
status.message(),
AllOf(
HasSubstr("Input array not provided for operation 'stablehlo.pad'."),
HasSubstr("Check the 'interior_padding' attribute.")));
}
TEST_F(StablehloPadFlatbufferConversionsTest, FailsInconsistentSizes) {
const Operator* stablehlo_pad_op = BuildTestOperator(
BuiltinOptions2_StablehloPadOptions,
CreateStablehloPadOptions(
builder_,
/*edge_padding_low=*/builder_.CreateVector<int64_t>({1, 0, -1}),
/*edge_padding_high=*/builder_.CreateVector<int64_t>({2, 0, -2}),
/*interior_padding=*/builder_.CreateVector<int64_t>({3, 0, -3, 5}))
.Union());
TfLiteStablehloPadParams* output_data = nullptr;
auto status = ParseOpData(stablehlo_pad_op, BuiltinOperator_STABLEHLO_PAD,
&mock_allocator_, (void**)&output_data);
EXPECT_FALSE(status.ok());
EXPECT_THAT(status.message(),
HasSubstr("'stablehlo.pad' operation parameter array sizes are "
"not consistent."));
}
TEST_F(StablehloPadFlatbufferConversionsTest, FailsWithWrongOptions) {
const Operator* stablehlo_pad_op = BuildTestOperator(BuiltinOptions_NONE, 0);
TfLiteStablehloPadParams* output_data = nullptr;
auto status = ParseOpData(stablehlo_pad_op, BuiltinOperator_STABLEHLO_PAD,
&mock_allocator_, (void**)&output_data);
EXPECT_FALSE(status.ok());
EXPECT_THAT(status.message(),
HasSubstr("Could not get 'stablehlo.pad' operation parameters."));
}
TEST_F(StablehloPadFlatbufferConversionsTest, DeathTests) {
const Operator* stablehlo_pad_op = BuildTestOperator(BuiltinOptions_NONE, 0);
TfLiteStablehloPadParams* output_data = nullptr;
#ifdef NDEBUG
GTEST_SKIP();
#endif
EXPECT_DEATH(ParseOpData(nullptr, BuiltinOperator_STABLEHLO_PAD,
&mock_allocator_, (void**)&output_data)
.IgnoreError(),
"");
EXPECT_DEATH(ParseOpData(stablehlo_pad_op, BuiltinOperator_STABLEHLO_PAD,
nullptr, (void**)&output_data)
.IgnoreError(),
"");
EXPECT_DEATH(ParseOpData(stablehlo_pad_op, BuiltinOperator_STABLEHLO_PAD,
&mock_allocator_, nullptr)
.IgnoreError(),
"");
}
} // namespace flatbuffer_conversions
} // namespace tflite_file
tensorflow/compiler/mlir/lite/core/api/verifier.h
+39
View File
@@ -0,0 +1,39 @@
/* 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.
==============================================================================*/
/// \file
///
/// Abstract interface for verifying a model.
#ifndef TENSORFLOW_COMPILER_MLIR_LITE_CORE_API_VERIFIER_H_
#define TENSORFLOW_COMPILER_MLIR_LITE_CORE_API_VERIFIER_H_
#include "tensorflow/compiler/mlir/lite/core/api/error_reporter.h"
namespace tflite {
/// Abstract interface that verifies whether a given model is legit.
/// It facilitates the use-case to verify and build a model without loading it
/// twice.
/// (See also "tensorflow/lite/tools/verifier.h".)
class TfLiteVerifier {
public:
/// Returns true if the model is legit.
virtual bool Verify(const char* data, int length,
ErrorReporter* reporter) = 0;
virtual ~TfLiteVerifier() = default;
};
} // namespace tflite
#endif // TENSORFLOW_COMPILER_MLIR_LITE_CORE_API_VERIFIER_H_
tensorflow/compiler/mlir/lite/core/c/BUILD
+67
View File
@@ -0,0 +1,67 @@
load("@rules_cc//cc:cc_library.bzl", "cc_library")
load("//tensorflow:tensorflow.default.bzl", "get_compatible_with_portable")
load("//tensorflow/compiler/mlir/lite:build_def.bzl", "tflite_copts")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:LICENSE"],
default_visibility = [
"//visibility:public",
],
licenses = ["notice"],
)
filegroup(
name = "tflite_internal_cc_3p_api_deps_src",
srcs = [
":tflite_types.h",
],
visibility = ["//tensorflow/lite:__pkg__"],
)
filegroup(
name = "lite_headers_filegroup",
srcs = [
"builtin_op_data.h",
"tflite_types.h",
],
visibility = ["//visibility:public"],
)
exports_files(
srcs = [
"builtin_op_data.h",
"tflite_types.h",
],
visibility = [
# copybara:uncomment_begin(google-only)
# "//tensorflow/lite:__subpackages__",
# copybara:uncomment_end_and_comment_begin
"//visibility:public",
# copybara:comment_end
],
)
cc_library(
name = "tflite_types",
hdrs = [
"tflite_types.h",
],
compatible_with = get_compatible_with_portable(),
copts = tflite_copts(),
visibility = ["//visibility:public"],
)
# LINT.IfChange(common)
cc_library(
name = "tflite_common",
srcs = [],
hdrs = [
"builtin_op_data.h",
"tflite_types.h",
],
compatible_with = get_compatible_with_portable(),
copts = tflite_copts(),
visibility = ["//visibility:public"],
alwayslink = 1, # Why?? TODO(b/161243354): eliminate this.
)
# LINT.ThenChange(//tensorflow/lite/core/c:common)
tensorflow/compiler/mlir/lite/core/c/builtin_op_data.h
+671
View File
@@ -0,0 +1,671 @@
/* Copyright 2024 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.
==============================================================================*/
/// WARNING: Users of TensorFlow Lite should not include this file directly,
/// only the TensorFlow Lite implementation itself should.
// IWYU pragma: private, include "third_party/tensorflow/lite/c/builtin_op_data.h"
#ifndef TENSORFLOW_COMPILER_MLIR_LITE_CORE_C_BUILTIN_OP_DATA_H_
#define TENSORFLOW_COMPILER_MLIR_LITE_CORE_C_BUILTIN_OP_DATA_H_
#include <stdbool.h> // IWYU pragma: keep
#include <stddef.h>
#include <stdint.h>
#include "tensorflow/compiler/mlir/lite/core/c/tflite_types.h"
#ifdef __cplusplus
extern "C" {
#endif // __cplusplus
// TfLiteReshapeParams can't have dynamic data so we fix the maximum possible
// number of dimensions.
#define TFLITE_RESHAPE_PARAMS_MAX_DIMENSION_COUNT 8
#define TFLITE_STABLEHLO_SCATTER_PARAMS_MAX_DIMENSION_COUNT 8
#define TFLITE_STABLEHLO_GATHER_PARAMS_MAX_DIMENSION_COUNT 8
#define TFLITE_STABLEHLO_REDUCE_WINDOW_PARAMS_MAX_DIMENSION_COUNT 8
#define TFLITE_STABLEHLO_PAD_PARAMS_MAX_DIMENSION_COUNT 8
#define TFLITE_STABLEHLO_CASE_PARAMS_MAX_BRANCHES_COUNT 20
// TODO(aselle): Consider using "if this then that" for testing.
// Useful placeholder to put in otherwise empty structs to avoid size warnings.
typedef struct {
char dummy;
} EmptyStructPlaceholder;
// IMPORTANT: All new members of structs must be added at the end to ensure
// backwards compatibility.
// Possible padding types (for convolutions)
typedef enum {
kTfLitePaddingUnknown = 0,
kTfLitePaddingSame,
kTfLitePaddingValid,
} TfLitePadding;
typedef enum {
kTfLiteMirrorPaddingUnknown = 0,
kTfLiteMirrorPaddingReflect,
kTfLiteMirrorPaddingSymmetric,
} TfLiteMirrorPaddingMode;
// TODO(b/130259536): We should move this out of builtin_op_data.
typedef struct {
int width;
int height;
int width_offset;
int height_offset;
} TfLitePaddingValues;
typedef struct {
TfLiteMirrorPaddingMode mode;
} TfLiteMirrorPaddingParams;
// Possible fused activation functions.
typedef enum {
kTfLiteActNone = 0,
kTfLiteActRelu,
kTfLiteActReluN1To1, // min(max(-1, x), 1)
kTfLiteActRelu6, // min(max(0, x), 6)
kTfLiteActTanh,
kTfLiteActSignBit,
kTfLiteActSigmoid,
} TfLiteFusedActivation;
typedef struct {
// Parameters for CONV_2D version 1.
TfLitePadding padding;
int stride_width;
int stride_height;
TfLiteFusedActivation activation;
// Parameters for CONV_2D version 2.
// Note: Version 2 supports dilation values not equal to 1.
int dilation_width_factor;
int dilation_height_factor;
// Parameters for CONV_2D version 7 or above.
// Used to determine the default value for the quantized bias.
TfLiteType quantized_bias_type;
} TfLiteConvParams;
typedef struct {
TfLitePadding padding;
int stride_width;
int stride_height;
int stride_depth;
int dilation_width_factor;
int dilation_height_factor;
int dilation_depth_factor;
TfLiteFusedActivation activation;
} TfLiteConv3DParams;
typedef TfLiteConv3DParams TfLiteConv3DTransposeParams;
typedef struct {
TfLitePadding padding;
int stride_width;
int stride_height;
int filter_width;
int filter_height;
TfLiteFusedActivation activation;
struct {
TfLitePaddingValues padding;
} computed;
} TfLitePoolParams;
typedef struct {
// Parameters for DepthwiseConv version 1 or above.
TfLitePadding padding;
int stride_width;
int stride_height;
// `depth_multiplier` is redundant. It's used by CPU kernels in
// TensorFlow 2.0 or below, but ignored in versions above.
//
// The information can be deduced from the shape of input and the shape of
// weights. Since the TFLiteConverter toolchain doesn't support partially
// specified shapes, relying on `depth_multiplier` stops us from supporting
// graphs with dynamic shape tensors.
//
// Note: Some of the delegates (e.g. NNAPI, GPU) are still relying on this
// field.
int depth_multiplier;
TfLiteFusedActivation activation;
// Parameters for DepthwiseConv version 2 or above.
int dilation_width_factor;
int dilation_height_factor;
} TfLiteDepthwiseConvParams;
typedef struct {
int rank;
TfLiteFusedActivation activation;
// Parameter for SVDF version 4.
bool asymmetric_quantize_inputs;
} TfLiteSVDFParams;
typedef struct {
TfLiteFusedActivation activation;
// Parameter for RNN version 3.
bool asymmetric_quantize_inputs;
} TfLiteRNNParams;
typedef struct {
bool time_major;
TfLiteFusedActivation activation;
// Parameter for Sequence RNN version 3.
bool asymmetric_quantize_inputs;
} TfLiteSequenceRNNParams;
typedef struct {
bool time_major;
TfLiteFusedActivation activation;
bool merge_outputs;
// Parameter for Bidirectional RNN version 3.
bool asymmetric_quantize_inputs;
} TfLiteBidirectionalSequenceRNNParams;
typedef enum {
kTfLiteFullyConnectedWeightsFormatDefault = 0,
kTfLiteFullyConnectedWeightsFormatShuffled4x16Int8 = 1,
} TfLiteFullyConnectedWeightsFormat;
typedef struct {
// Parameters for FullyConnected version 1 or above.
TfLiteFusedActivation activation;
// Parameters for FullyConnected version 2 or above.
TfLiteFullyConnectedWeightsFormat weights_format;
// Parameters for FullyConnected version 5 or above.
// If set to true, then the number of dimensions in the input and the output
// tensors are the same. Furthermore, all but the last dimension of the input
// and output shapes will be equal.
bool keep_num_dims;
// Parameters for FullyConnected version 7 or above.
// If set to true and the weights are quantized, then non constant inputs
// are quantized at evaluation time with asymmetric quantization.
bool asymmetric_quantize_inputs;
// Parameters for FullyConnected version 10 or above.
// Used to determine the default value for the quantized bias.
TfLiteType quantized_bias_type;
} TfLiteFullyConnectedParams;
typedef enum {
kTfLiteLshProjectionUnknown = 0,
kTfLiteLshProjectionSparse = 1,
kTfLiteLshProjectionDense = 2,
} TfLiteLSHProjectionType;
typedef struct {
TfLiteLSHProjectionType type;
} TfLiteLSHProjectionParams;
typedef struct {
float beta;
} TfLiteSoftmaxParams;
typedef struct {
int axis;
TfLiteFusedActivation activation;
} TfLiteConcatenationParams;
typedef struct {
TfLiteFusedActivation activation;
// Parameter added for the version 4.
bool pot_scale_int16;
} TfLiteAddParams;
typedef struct {
EmptyStructPlaceholder placeholder;
} TfLiteSpaceToBatchNDParams;
typedef struct {
EmptyStructPlaceholder placeholder;
} TfLiteBatchToSpaceNDParams;
typedef struct {
bool adj_x;
bool adj_y;
// Parameters for BatchMatMul version 4 or above.
// If set to true and the weights are quantized, then non constant inputs
// are quantized at evaluation time with asymmetric quantization.
bool asymmetric_quantize_inputs;
} TfLiteBatchMatMulParams;
typedef struct {
TfLiteFusedActivation activation;
} TfLiteMulParams;
typedef struct {
TfLiteFusedActivation activation;
// Parameter added for the version 5.
bool pot_scale_int16;
} TfLiteSubParams;
typedef struct {
TfLiteFusedActivation activation;
} TfLiteDivParams;
typedef struct {
TfLiteFusedActivation activation;
} TfLiteL2NormParams;
typedef struct {
int radius;
float bias;
float alpha;
float beta;
} TfLiteLocalResponseNormParams;
typedef enum {
kTfLiteLSTMFullKernel = 0,
kTfLiteLSTMBasicKernel
} TfLiteLSTMKernelType;
typedef struct {
// Parameters for LSTM version 1.
TfLiteFusedActivation activation;
float cell_clip;
float proj_clip;
// Parameters for LSTM version 2.
// kTfLiteLSTMBasicKernel is only supported in version 2 or above.
TfLiteLSTMKernelType kernel_type;
// Parameters for LSTM version 4.
bool asymmetric_quantize_inputs;
} TfLiteLSTMParams;
typedef struct {
// Parameters needed for the underlying LSTM.
TfLiteFusedActivation activation;
float cell_clip;
float proj_clip;
// If set to true then the first dimension is time, otherwise batch.
bool time_major;
// Parameter for unidirectional sequence RNN version 3.
bool asymmetric_quantize_inputs;
// Parameter for unidirectional sequence RNN version 4.
bool diagonal_recurrent_tensors;
} TfLiteUnidirectionalSequenceLSTMParams;
typedef struct {
// Parameters supported by version 1:
// Parameters inherited for the LSTM kernel.
TfLiteFusedActivation activation;
float cell_clip;
float proj_clip;
// If true, store the outputs of both directions in the first output.
bool merge_outputs;
// Parameters supported by version 2:
// If set to true then the first dimension is time, otherwise batch.
bool time_major;
// Parameters supported by version 3:
// If set to true, then hybrid ops use asymmetric quantization for inputs.
bool asymmetric_quantize_inputs;
} TfLiteBidirectionalSequenceLSTMParams;
typedef struct {
bool align_corners;
// half_pixel_centers assumes pixels are of half the actual dimensions, and
// yields more accurate resizes. Corresponds to the same argument for the
// original TensorFlow op in TF2.0.
bool half_pixel_centers;
} TfLiteResizeBilinearParams;
typedef struct {
bool align_corners;
bool half_pixel_centers;
} TfLiteResizeNearestNeighborParams;
typedef struct {
EmptyStructPlaceholder placeholder;
} TfLitePadParams;
typedef struct {
EmptyStructPlaceholder placeholder;
} TfLitePadV2Params;
typedef struct {
// These fields are only used in old models for backward compatibility.
// In the current implementation, we use the 2nd input of the op as the shape,
// and these fields are unused.
int32_t shape[TFLITE_RESHAPE_PARAMS_MAX_DIMENSION_COUNT];
int num_dimensions;
} TfLiteReshapeParams;
typedef struct {
int ngram_size;
int max_skip_size;
bool include_all_ngrams;
} TfLiteSkipGramParams;
typedef struct {
int block_size;
} TfLiteSpaceToDepthParams;
typedef struct {
int block_size;
} TfLiteDepthToSpaceParams;
typedef struct {
TfLiteType in_data_type;
TfLiteType out_data_type;
} TfLiteCastParams;
typedef enum {
kTfLiteCombinerTypeSum = 0,
kTfLiteCombinerTypeMean = 1,
kTfLiteCombinerTypeSqrtn = 2,
} TfLiteCombinerType;
typedef struct {
TfLiteCombinerType combiner;
} TfLiteEmbeddingLookupSparseParams;
typedef struct {
int axis;
int batch_dims;
} TfLiteGatherParams;
typedef struct {
EmptyStructPlaceholder placeholder;
} TfLiteTransposeParams;
typedef struct {
bool keep_dims;
} TfLiteReducerParams;
typedef struct {
int num_splits;
} TfLiteSplitParams;
typedef struct {
int num_splits;
} TfLiteSplitVParams;
typedef struct {
// TODO(ahentz): We can't have dynamic data in this struct, at least not yet.
// For now we will fix the maximum possible number of dimensions.
int32_t squeeze_dims[8];
int num_squeeze_dims;
} TfLiteSqueezeParams;
typedef struct {
int begin_mask;
int end_mask;
int ellipsis_mask;
int new_axis_mask;
int shrink_axis_mask;
// Parameters supported by version 8:
// If true, then the end tensor is an offset of the begin tensor.
bool offset;
} TfLiteStridedSliceParams;
typedef struct {
TfLiteType output_type;
} TfLiteArgMaxParams;
typedef struct {
TfLiteType output_type;
} TfLiteArgMinParams;
typedef struct {
// Parameters supported by version 1:
TfLitePadding padding;
int stride_width;
int stride_height;
// Parameters supported by version 4:
TfLiteFusedActivation activation;
// Parameters for TransposeConv version 5 or above.
// Used to determine the default value for the quantized bias.
TfLiteType quantized_bias_type;
} TfLiteTransposeConvParams;
typedef struct {
bool validate_indices;
} TfLiteSparseToDenseParams;
typedef struct {
TfLiteType out_type;
} TfLiteShapeParams;
typedef struct {
EmptyStructPlaceholder placeholder;
} TfLiteRankParams;
typedef struct {
// Parameters supported by version 1:
float min;
float max;
int num_bits;
// Parameters supported by version 2:
bool narrow_range;
} TfLiteFakeQuantParams;
typedef struct {
int values_count;
int axis;
} TfLitePackParams;
typedef struct {
int axis;
} TfLiteOneHotParams;
typedef struct {
int num;
int axis;
} TfLiteUnpackParams;
typedef struct {
float alpha;
} TfLiteLeakyReluParams;
typedef struct {
TfLiteType index_out_type;
} TfLiteUniqueParams;
typedef struct {
int seq_dim;
int batch_dim;
} TfLiteReverseSequenceParams;
typedef struct {
EmptyStructPlaceholder placeholder;
} TfLiteMatrixDiagParams;
typedef struct {
EmptyStructPlaceholder placeholder;
} TfLiteMatrixSetDiagParams;
typedef struct {
int then_subgraph_index;
int else_subgraph_index;
} TfLiteIfParams;
typedef struct {
int cond_subgraph_index;
int body_subgraph_index;
} TfLiteWhileParams;
typedef struct {
bool exclusive;
bool reverse;
} TfLiteCumsumParams;
typedef struct {
int init_subgraph_index;
} TfLiteCallOnceParams;
typedef struct {
int table_id;
TfLiteType key_dtype;
TfLiteType value_dtype;
} TfLiteHashtableParams;
typedef struct {
const char* container;
const char* shared_name;
} TfLiteVarHandleParams;
typedef struct {
int seed;
int seed2;
} TfLiteRandomParams;
typedef struct {
int num_boundaries;
// This points to the memory stored in the model (flatbuffer),
// and is not owned.
const float* boundaries;
} TfLiteBucketizeParams;
typedef struct {
bool approximate;
} TfLiteGeluParams;
typedef struct {
int64_t dimension;
} TfLiteStablehloConcatenateParams;
typedef struct {
// See the stablehlo spec for the explanation of the attributes:
// https://github.com/openxla/stablehlo/blob/main/docs/spec.md#scatter
bool indices_are_sorted;
int64_t
update_window_dims[TFLITE_STABLEHLO_SCATTER_PARAMS_MAX_DIMENSION_COUNT];
int num_update_window_dims;
int64_t
inserted_window_dims[TFLITE_STABLEHLO_SCATTER_PARAMS_MAX_DIMENSION_COUNT];
int num_inserted_window_dims;
int64_t scatter_dims_to_operand_dims
[TFLITE_STABLEHLO_SCATTER_PARAMS_MAX_DIMENSION_COUNT];
int num_scatter_dims_to_operand_dims;
int64_t index_vector_dim;
bool unique_indices;
int update_computation_subgraph_index;
} TfLiteStablehloScatterParams;
typedef enum {
kTfLiteRngAlgorithmUnknown = 0,
// An algorithm auto-selected by the system according to device type.
kTfLiteRngAlgorithmDefault,
// The Philox algorithm, as described in paper
// ['Parallel Random Numbers: As Easy as 1, 2, 3']
// (https://www.thesalmons.org/john/random123/papers/random123sc11.pdf)
kTfLiteRngAlgorithmPhilox,
// The ThreeFry algorithm, as described in paper
// ['Parallel Random Numbers: As Easy as 1, 2, 3']
// (https://www.thesalmons.org/john/random123/papers/random123sc11.pdf)
kTfLiteRngAlgorithmThreefry,
} TfLiteRngAlgorithm;
typedef struct {
TfLiteRngAlgorithm algorithm;
} TfLiteStablehloRngBitGeneratorParams;
typedef struct {
// See the stablehlo spec for the explanation of the attributes:
// https://github.com/openxla/stablehlo/blob/main/docs/spec.md#gather
int64_t offset_dims[TFLITE_STABLEHLO_GATHER_PARAMS_MAX_DIMENSION_COUNT];
int num_offset_dims;
int64_t
collapsed_slice_dims[TFLITE_STABLEHLO_GATHER_PARAMS_MAX_DIMENSION_COUNT];
int num_collapsed_slice_dims;
int64_t start_index_map[TFLITE_STABLEHLO_GATHER_PARAMS_MAX_DIMENSION_COUNT];
int num_start_index_map;
int64_t index_vector_dim;
int64_t slice_sizes[TFLITE_STABLEHLO_GATHER_PARAMS_MAX_DIMENSION_COUNT];
int num_slice_sizes;
bool indices_are_sorted;
} TfLiteStablehloGatherParams;
typedef struct {
// See the stablehlo spec for the explanation of the attributes:
// https://github.com/openxla/stablehlo/blob/main/docs/spec.md#reduce_window
int64_t window_dimensions
[TFLITE_STABLEHLO_REDUCE_WINDOW_PARAMS_MAX_DIMENSION_COUNT];
int64_t
window_strides[TFLITE_STABLEHLO_REDUCE_WINDOW_PARAMS_MAX_DIMENSION_COUNT];
int64_t
base_dilations[TFLITE_STABLEHLO_REDUCE_WINDOW_PARAMS_MAX_DIMENSION_COUNT];
int64_t window_dilations
[TFLITE_STABLEHLO_REDUCE_WINDOW_PARAMS_MAX_DIMENSION_COUNT];
int64_t
padding[2 * TFLITE_STABLEHLO_REDUCE_WINDOW_PARAMS_MAX_DIMENSION_COUNT];
int body_subgraph_index;
} TfLiteStablehloReduceWindowParams;
enum TfLiteReduceWindowFunction {
TfLiteReduceWindowFunctionUnsupported,
TfLiteReduceWindowFunctionAdd,
TfLiteReduceWindowFunctionMul,
TfLiteReduceWindowFunctionMin,
TfLiteReduceWindowFunctionMax,
TfLiteReduceWindowFunctionAll,
TfLiteReduceWindowFunctionAny
};
typedef struct {
enum TfLiteReduceWindowFunction reduce_function;
} TfLiteReduceWindowParams;
typedef struct {
// See the stablehlo spec for the explanation of the attributes:
// https://github.com/openxla/stablehlo/blob/main/docs/spec.md#pad
int64_t edge_padding_low[TFLITE_STABLEHLO_PAD_PARAMS_MAX_DIMENSION_COUNT];
int64_t edge_padding_high[TFLITE_STABLEHLO_PAD_PARAMS_MAX_DIMENSION_COUNT];
int64_t interior_padding[TFLITE_STABLEHLO_PAD_PARAMS_MAX_DIMENSION_COUNT];
} TfLiteStablehloPadParams;
typedef struct {
const char* name;
int32_t subgraph_index;
int32_t version;
const uint8_t* attributes;
size_t attributes_size;
} TfLiteStablehloCompositeParams;
typedef struct {
// See the stablehlo spec for the explanation of the attributes:
// https://github.com/openxla/stablehlo/blob/main/docs/spec.md#case
int32_t
branch_subgraph_indices[TFLITE_STABLEHLO_CASE_PARAMS_MAX_BRANCHES_COUNT];
uint32_t num_branches;
} TfLiteStablehloCaseParams;
#ifdef __cplusplus
} // extern "C"
#endif // __cplusplus
#endif // TENSORFLOW_COMPILER_MLIR_LITE_CORE_C_BUILTIN_OP_DATA_H_
tensorflow/compiler/mlir/lite/core/c/tflite_types.h
+94
View File
@@ -0,0 +1,94 @@
/* Copyright 2024 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.
==============================================================================*/
// This file hosts data structures that are needed both for LiteRT and
// Compiler.
// WARNING: Users of TensorFlow Lite should not include this file directly, but
// should instead include "third_party/tensorflow/lite/c/c_api_types.h".
// Only the TensorFlow Lite implementation itself should include this file
// directly.
// clang-format off
// NOLINTBEGIN(whitespace/line_length)
/// \note Users of TensorFlow Lite should use
/// \code
/// #include "tensorflow/lite/c/c_api_types.h"
/// \endcode
/// to access the APIs documented on this page.
// NOLINTEND(whitespace/line_length)
// clang-format on
// IWYU pragma: private, include "third_party/tensorflow/lite/c/c_api_types.h"
#ifndef TENSORFLOW_COMPILER_MLIR_LITE_CORE_C_TFLITE_TYPES_H_
#define TENSORFLOW_COMPILER_MLIR_LITE_CORE_C_TFLITE_TYPES_H_
#include <stdint.h>
#ifdef __cplusplus
extern "C" {
#endif
/// Types supported by tensor
// LINT.IfChange
typedef enum {
kTfLiteNoType = 0,
kTfLiteFloat32 = 1,
kTfLiteInt32 = 2,
kTfLiteUInt8 = 3,
kTfLiteInt64 = 4,
kTfLiteString = 5,
kTfLiteBool = 6,
kTfLiteInt16 = 7,
kTfLiteComplex64 = 8,
kTfLiteInt8 = 9,
kTfLiteFloat16 = 10,
kTfLiteFloat64 = 11,
kTfLiteComplex128 = 12,
kTfLiteUInt64 = 13,
kTfLiteResource = 14,
kTfLiteVariant = 15,
kTfLiteUInt32 = 16,
kTfLiteUInt16 = 17,
kTfLiteInt4 = 18,
kTfLiteBFloat16 = 19,
kTfLiteInt2 = 20,
kTfLiteUInt4 = 21,
kTfLiteFloat8E4M3FN = 22,
kTfLiteFloat8E5M2 = 23,
} TfLiteType;
// LINT.ThenChange(//tensorflow/lite/profiling/proto/model_runtime_info.proto:EdgeDataType)
/// Legacy. Will be deprecated in favor of `TfLiteAffineQuantization`.
/// If per-layer quantization is specified this field will still be populated in
/// addition to `TfLiteAffineQuantization`.
/// Parameters for asymmetric quantization. Quantized values can be converted
/// back to float using: `real_value = scale * (quantized_value - zero_point)`
typedef struct TfLiteQuantizationParams {
float scale;
int32_t zero_point;
} TfLiteQuantizationParams;
/// Storage format of each dimension in a sparse tensor.
typedef enum TfLiteDimensionType {
kTfLiteDimDense = 0,
kTfLiteDimSparseCSR,
} TfLiteDimensionType;
#ifdef __cplusplus
} // extern C
#endif
#endif // TENSORFLOW_COMPILER_MLIR_LITE_CORE_C_TFLITE_TYPES_H_
tensorflow/compiler/mlir/lite/core/macros.h
+50
View File
@@ -0,0 +1,50 @@
/* Copyright 2024 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.
==============================================================================*/
// This provides utility macros and functions that are inherently platform
// specific or shared across runtime & converter.
#ifndef TENSORFLOW_COMPILER_MLIR_LITE_CORE_MACROS_H_
#define TENSORFLOW_COMPILER_MLIR_LITE_CORE_MACROS_H_
#ifndef TF_LITE_STATIC_MEMORY
// maximum size of a valid flatbuffer
inline constexpr unsigned int flatbuffer_size_max = 2147483648;
// If none zero then the buffer is stored outside of the flatbuffers, string
inline constexpr char tflite_metadata_buffer_location[] = "buffer_location";
// field for minimum runtime version, string
inline constexpr char tflite_metadata_min_runtime_version[] =
"min_runtime_version";
// the stablehlo op version is supported by the tflite runtime
inline constexpr char tflite_supported_stablehlo_version[] = "1.0.0";
#endif
// LINT.IfChange(TFLITE_NOINLINE)
#ifdef _WIN32
#define TFLITE_NOINLINE __declspec(noinline)
#else
#ifdef __has_attribute
#if __has_attribute(noinline)
#define TFLITE_NOINLINE __attribute__((noinline))
#else
#define TFLITE_NOINLINE
#endif // __has_attribute(noinline)
#else
#define TFLITE_NOINLINE
#endif // __has_attribute
#endif // _WIN32
// LINT.ThenChange(//tensorflow/lite/core/macros.h)
#endif // TENSORFLOW_COMPILER_MLIR_LITE_CORE_MACROS_H_
tensorflow/compiler/mlir/lite/core/model_builder_base.cc
+62
View File
@@ -0,0 +1,62 @@
/* Copyright 2017 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/core/model_builder_base.h"
#include <stddef.h>
#include <memory>
#include <string>
#include "absl/strings/str_cat.h"
#include "tensorflow/compiler/mlir/lite/allocation.h"
#include "tensorflow/compiler/mlir/lite/core/api/error_reporter.h"
namespace tflite {
#ifndef TFLITE_MCU
// Loads a model from `filename`. If `mmap_file` is true then use mmap,
// otherwise make a copy of the model in a buffer.
std::unique_ptr<Allocation> GetAllocationFromFile(const char* filename,
ErrorReporter* error_reporter,
bool allow_modifications) {
std::unique_ptr<Allocation> allocation;
if (MMAPAllocation::IsSupported()) {
allocation = std::make_unique<MMAPAllocation>(filename, error_reporter,
allow_modifications);
} else {
allocation = std::make_unique<FileCopyAllocation>(filename, error_reporter);
}
return allocation;
}
// Loads a model from `fd`. If `mmap_file` is true then use mmap,
// otherwise make a copy of the model in a buffer.
std::unique_ptr<Allocation> GetAllocationFromFile(int fd,
ErrorReporter* error_reporter,
bool allow_modifications) {
std::unique_ptr<Allocation> allocation;
if (MMAPAllocation::IsSupported()) {
allocation = std::make_unique<MMAPAllocation>(fd, error_reporter,
allow_modifications);
} else {
allocation = std::make_unique<FileCopyAllocation>(
absl::StrCat("/proc/self/fd/", fd).c_str(), error_reporter);
}
return allocation;
}
#endif
} // namespace tflite
tensorflow/compiler/mlir/lite/core/model_builder_base.h
+627
View File
@@ -0,0 +1,627 @@
/* Copyright 2017 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.
==============================================================================*/
/// \file
///
/// Deserialization infrastructure for tflite. Provides functionality
/// to go from a serialized tflite model in flatbuffer format to an
/// in-memory representation of the model.
///
/// WARNING: Users of TensorFlow Lite should not include this file directly,
/// but should instead include "third_party/tensorflow/lite/model_builder.h".
/// Only the TensorFlow Lite implementation itself should include this
/// file directly.
// IWYU pragma: private, include "third_party/tensorflow/lite/model_builder.h"
#ifndef TENSORFLOW_COMPILER_MLIR_LITE_CORE_MODEL_BUILDER_BASE_H_
#define TENSORFLOW_COMPILER_MLIR_LITE_CORE_MODEL_BUILDER_BASE_H_
#include <stddef.h>
#include <algorithm>
#include <cstdint>
#include <map>
#include <memory>
#include <string>
#include <utility>
#include "flatbuffers/base.h" // from @flatbuffers
#include "flatbuffers/buffer.h" // from @flatbuffers
#include "flatbuffers/vector.h" // from @flatbuffers
#include "flatbuffers/verifier.h" // from @flatbuffers
#include "tensorflow/compiler/mlir/lite/allocation.h"
#include "tensorflow/compiler/mlir/lite/core/api/error_reporter.h"
#include "tensorflow/compiler/mlir/lite/core/api/verifier.h"
#include "tensorflow/compiler/mlir/lite/core/macros.h"
#include "tensorflow/compiler/mlir/lite/schema/schema_generated.h"
namespace tflite {
std::unique_ptr<Allocation> GetAllocationFromFile(
const char* filename, ErrorReporter* error_reporter,
bool allow_modifications = false);
std::unique_ptr<Allocation> GetAllocationFromFile(
int fd, ErrorReporter* error_reporter, bool allow_modifications = false);
namespace impl {
/// An RAII object that represents a read-only tflite model, copied from disk,
/// or mmapped. This uses flatbuffers as the serialization format.
///
/// NOTE: The current API requires that a FlatBufferModelBase instance be kept
/// alive by the client as long as it is in use by any dependent Interpreter
/// instances. As the FlatBufferModelBase instance is effectively immutable
/// after creation, the client may safely use a single model with multiple
/// dependent Interpreter instances, even across multiple threads (though note
/// that each Interpreter instance is *not* thread-safe).
///
/// <pre><code>
/// using namespace tflite;
/// StderrReporter error_reporter;
/// auto model = FlatBufferModelBase::BuildFromFile("interesting_model.tflite",
/// &error_reporter);
/// MyOpResolver resolver; // You need to subclass OpResolver to provide
/// // implementations.
/// InterpreterBuilder builder(*model, resolver);
/// std::unique_ptr<Interpreter> interpreter;
/// if(builder(&interpreter) == kTfLiteOk) {
/// .. run model inference with interpreter
/// }
/// </code></pre>
///
/// OpResolver must be defined to provide your kernel implementations to the
/// interpreter. This is environment specific and may consist of just the
/// builtin ops, or some custom operators you defined to extend tflite.
template <typename T>
class FlatBufferModelBase {
public:
/// Builds a model based on a file.
/// Caller retains ownership of `error_reporter` and must ensure its lifetime
/// is longer than the FlatBufferModelBase instance.
/// Returns a nullptr in case of failure.
static std::unique_ptr<T> BuildFromFile(
const char* filename,
ErrorReporter* error_reporter = T::GetDefaultErrorReporter()) {
error_reporter = ValidateErrorReporter(error_reporter);
std::unique_ptr<T> model = BuildFromAllocation(
GetAllocationFromFile(filename, error_reporter), error_reporter);
#if FLATBUFFERS_LITTLEENDIAN == 1
return model;
#else
return ByteConvertModel(std::move(model), error_reporter);
#endif
}
/// Verifies whether the content of the file is legit, then builds a model
/// based on the file.
/// The extra_verifier argument is an additional optional verifier for the
/// file contents. By default, we always check with tflite::VerifyModelBuffer.
/// If extra_verifier is supplied, the file contents is also checked against
/// the extra_verifier after the check against tflite::VerifyModelBuilder.
/// Caller retains ownership of `error_reporter` and must ensure its lifetime
/// is longer than the FlatBufferModelBase instance.
/// Returns a nullptr in case of failure.
static std::unique_ptr<T> VerifyAndBuildFromFile(
const char* filename, TfLiteVerifier* extra_verifier = nullptr,
ErrorReporter* error_reporter = T::GetDefaultErrorReporter()) {
error_reporter = ValidateErrorReporter(error_reporter);
std::unique_ptr<T> model = VerifyAndBuildFromAllocation(
GetAllocationFromFile(filename, error_reporter), extra_verifier,
error_reporter);
#if FLATBUFFERS_LITTLEENDIAN == 1
return model;
#else
return ByteConvertModel(std::move(model), error_reporter);
#endif
}
/// Builds a model based on a file descriptor.
/// Caller retains ownership of `error_reporter` and must ensure its lifetime
/// is longer than the FlatBufferModelBase instance. Caller retains ownership
/// of `fd` and must ensure it is closed after BuildFromFile returns. Returns
/// a nullptr in case of failure.
static std::unique_ptr<T> BuildFromFileDescriptor(
int fd, ErrorReporter* error_reporter = T::GetDefaultErrorReporter()) {
error_reporter = ValidateErrorReporter(error_reporter);
std::unique_ptr<T> model = BuildFromAllocation(
GetAllocationFromFile(fd, error_reporter), error_reporter);
#if FLATBUFFERS_LITTLEENDIAN == 1
return model;
#else
return ByteConvertModel(std::move(model), error_reporter);
#endif
}
/// Verifies whether the content of the file descriptor is legit, then builds
/// a model based on the file.
/// The extra_verifier argument is an additional optional verifier for the
/// file contents. By default, we always check with tflite::VerifyModelBuffer.
/// If extra_verifier is supplied, the file contents is also checked against
/// the extra_verifier after the check against tflite::VerifyModelBuilder.
/// Caller retains ownership of `error_reporter` and must ensure its lifetime
/// is longer than the FlatBufferModelBase instance.
/// Returns a nullptr in case of failure.
static std::unique_ptr<T> VerifyAndBuildFromFileDescriptor(
int fd, TfLiteVerifier* extra_verifier = nullptr,
ErrorReporter* error_reporter = T::GetDefaultErrorReporter()) {
error_reporter = ValidateErrorReporter(error_reporter);
std::unique_ptr<FlatBufferModelBase<T>> model =
VerifyAndBuildFromAllocation(GetAllocationFromFile(fd, error_reporter),
extra_verifier, error_reporter);
#if FLATBUFFERS_LITTLEENDIAN == 1
return model;
#else
return ByteConvertModel(std::move(model), error_reporter);
#endif
}
/// Builds a model based on a pre-loaded flatbuffer.
/// Caller retains ownership of the buffer and should keep it alive until
/// the returned object is destroyed. Caller also retains ownership of
/// `error_reporter` and must ensure its lifetime is longer than the
/// FlatBufferModelBase instance.
/// Returns a nullptr in case of failure.
/// NOTE: this validates the base Flatbuffer buffer structure. Use
/// VerifyAndBuildFromBuffer if an additional TfLiteVerifier check is
/// required.
static std::unique_ptr<T> BuildFromBuffer(
const char* caller_owned_buffer, size_t buffer_size,
ErrorReporter* error_reporter = T::GetDefaultErrorReporter()) {
error_reporter = ValidateErrorReporter(error_reporter);
std::unique_ptr<Allocation> allocation(
new MemoryAllocation(caller_owned_buffer, buffer_size, error_reporter));
return BuildFromAllocation(std::move(allocation), error_reporter);
}
/// Verifies whether the content of the buffer is legit, then builds a model
/// based on the pre-loaded flatbuffer.
/// The extra_verifier argument is an additional optional verifier for the
/// buffer. By default, we always check with tflite::VerifyModelBuffer. If
/// extra_verifier is supplied, the buffer is checked against the
/// extra_verifier after the check against tflite::VerifyModelBuilder. The
/// caller retains ownership of the buffer and should keep it alive until the
/// returned object is destroyed. Caller retains ownership of `error_reporter`
/// and must ensure its lifetime is longer than the FlatBufferModelBase
/// instance. Returns a nullptr in case of failure.
static std::unique_ptr<T> VerifyAndBuildFromBuffer(
const char* caller_owned_buffer, size_t buffer_size,
TfLiteVerifier* extra_verifier = nullptr,
ErrorReporter* error_reporter = T::GetDefaultErrorReporter()) {
error_reporter = ValidateErrorReporter(error_reporter);
std::unique_ptr<Allocation> allocation(
new MemoryAllocation(caller_owned_buffer, buffer_size, error_reporter));
return VerifyAndBuildFromAllocation(std::move(allocation), extra_verifier,
error_reporter);
}
#if FLATBUFFERS_LITTLEENDIAN == 0
static void ByteSwapSerializedModel(std::string* serialized_model,
bool from_big_endian) {
const uint8_t* buffer =
reinterpret_cast<const uint8_t*>(serialized_model->c_str());
const tflite::Model* input_model = tflite::GetModel(buffer);
ByteSwapTFLiteModel(input_model, from_big_endian);
}
static void ByteSwapBuffer(int8_t tensor_type, size_t buffer_size,
uint8_t* buffer, bool from_big_endian) {
switch (tensor_type) {
case tflite::TensorType_STRING: {
auto bp = reinterpret_cast<int32_t*>(buffer);
int num_of_strings =
from_big_endian ? bp[0] : flatbuffers::EndianSwap(bp[0]);
for (int i = 0; i < num_of_strings + 2; i++)
bp[i] = flatbuffers::EndianSwap(bp[i]);
break;
}
// 16-bit types
case tflite::TensorType_FLOAT16:
case tflite::TensorType_INT16:
case tflite::TensorType_UINT16: {
auto bp = reinterpret_cast<uint16_t*>(buffer);
for (int i = 0; i < buffer_size / 2; i++)
bp[i] = flatbuffers::EndianSwap(bp[i]);
break;
}
// 32-bit types
case tflite::TensorType_FLOAT32:
case tflite::TensorType_INT32:
case tflite::TensorType_UINT32:
case tflite::TensorType_COMPLEX64: {
auto bp = reinterpret_cast<uint32_t*>(buffer);
for (int i = 0; i < buffer_size / 4; i++)
bp[i] = flatbuffers::EndianSwap(bp[i]);
break;
}
// 64-bit types
case tflite::TensorType_INT64:
case tflite::TensorType_FLOAT64:
case tflite::TensorType_UINT64:
case tflite::TensorType_COMPLEX128: {
auto bp = reinterpret_cast<uint64_t*>(buffer);
for (int i = 0; i < buffer_size / 8; i++)
bp[i] = flatbuffers::EndianSwap(bp[i]);
break;
}
default:
break;
}
}
static void ByteSwapTFLiteModel(const tflite::Model* tfl_model,
bool from_big_endian) {
std::vector<bool> buffer_swapped(tfl_model->buffers()->size(), false);
for (size_t subgraph_idx = 0; subgraph_idx < tfl_model->subgraphs()->size();
subgraph_idx++) {
const tflite::SubGraph* subgraph =
tfl_model->subgraphs()->Get(subgraph_idx);
for (size_t ts_idx = 0; ts_idx < subgraph->tensors()->size(); ts_idx++) {
const tflite::Tensor* tensor = subgraph->tensors()->Get(ts_idx);
if (tensor->buffer() > 0 &&
tensor->buffer() < tfl_model->buffers()->size() &&
!buffer_swapped[tensor->buffer()]) {
const tflite::Buffer* buffer_ =
(*tfl_model->buffers())[tensor->buffer()];
if (!buffer_ || !buffer_->data()) continue;
auto* buffer = buffer_->data();
uint8_t* buff_ = const_cast<uint8_t*>(buffer->data());
ByteSwapBuffer(tensor->type(), buffer->size(), buff_,
from_big_endian);
buffer_swapped[tensor->buffer()] = true;
}
}
}
}
static std::unique_ptr<T> ByteConvertModel(std::unique_ptr<T> model,
ErrorReporter* error_reporter,
bool from_big_endian = false) {
if (model == nullptr) return model;
auto tfl_model = model->GetModel();
if (tfl_model->subgraphs()->size() == 0) return model;
if (tfl_model->subgraphs()->Get(0)->tensors()->size() == 0) return model;
if (tfl_model->buffers()->size() < 2) return model;
return ByteSwapFlatBufferModelBase(std::move(model), error_reporter,
from_big_endian);
}
static std::unique_ptr<T> ByteSwapFlatBufferModelBase(
std::unique_ptr<T> model, ErrorReporter* error_reporter,
bool from_big_endian) {
FlatBufferModelBase<T>* modelp = model.release();
auto tflite_model = modelp->GetModel();
auto copied_model = std::make_unique<tflite::ModelT>();
tflite_model->UnPackTo(copied_model.get(), nullptr);
ByteSwapTFLiteModelT(copied_model.get(), from_big_endian);
std::unique_ptr<flatbuffers::FlatBufferBuilder> builder(
new flatbuffers::FlatBufferBuilder());
auto packed_model = tflite::Model::Pack(*builder, copied_model.get());
tflite::FinishModelBuffer(*builder, packed_model);
flatbuffers::FlatBufferBuilder* builder_ = builder.release();
return BuildFromBuffer(
reinterpret_cast<const char*>(builder_->GetBufferPointer()),
builder_->GetSize(), error_reporter);
}
static void ByteSwapTFLiteModelT(tflite::ModelT* tfl_modelt,
bool from_big_endian) {
size_t bytes_per_elem = 0;
std::vector<bool> buffer_swapped(tfl_modelt->buffers.size(), false);
for (size_t subgraph_idx = 0; subgraph_idx < tfl_modelt->subgraphs.size();
subgraph_idx++) {
tflite::SubGraphT* subgraph =
tfl_modelt->subgraphs.at(subgraph_idx).get();
for (size_t ts_idx = 0; ts_idx < subgraph->tensors.size(); ts_idx++) {
tflite::TensorT* tensor = subgraph->tensors[ts_idx].get();
if (tensor->buffer > 0 && tensor->buffer < tfl_modelt->buffers.size() &&
!buffer_swapped[tensor->buffer]) {
const auto* buffer =
&(tfl_modelt->buffers[tensor->buffer].get()->data);
if (buffer && buffer->data()) {
uint8_t* buff_ = const_cast<uint8_t*>(buffer->data());
ByteSwapBuffer(tensor->type, buffer->size(), buff_,
from_big_endian);
buffer_swapped[tensor->buffer] = true;
}
}
}
}
}
#endif
private:
// Internal helper to build a model from an allocation without running base
// verification.
static std::unique_ptr<T> BuildFromAllocationInternal(
std::unique_ptr<Allocation> allocation, ErrorReporter* error_reporter) {
std::unique_ptr<T> model(new T(std::move(allocation), error_reporter));
if (!model->initialized()) {
model.reset();
} else {
model->ValidateModelBuffers(model->error_reporter());
}
return model;
}
public:
/// Builds a model directly from an allocation.
/// Ownership of the allocation is passed to the model, but the caller
/// retains ownership of `error_reporter` and must ensure its lifetime is
/// longer than the FlatBufferModelBase instance.
/// Returns a nullptr in case of failure (e.g., the allocation is invalid).
static std::unique_ptr<T> BuildFromAllocation(
std::unique_ptr<Allocation> allocation,
ErrorReporter* error_reporter = T::GetDefaultErrorReporter()) {
return VerifyAndBuildFromAllocation(std::move(allocation),
/*extra_verifier=*/nullptr,
error_reporter);
}
/// Verifies whether the content of the allocation is legit, then builds a
/// model based on the provided allocation.
/// The extra_verifier argument is an additional optional verifier for the
/// buffer. By default, we always check with tflite::VerifyModelBuffer. If
/// extra_verifier is supplied, the buffer is checked against the
/// extra_verifier after the check against tflite::VerifyModelBuilder.
/// Ownership of the allocation is passed to the model, but the caller
/// retains ownership of `error_reporter` and must ensure its lifetime is
/// longer than the FlatBufferModelBase instance.
/// Returns a nullptr in case of failure.
static std::unique_ptr<T> VerifyAndBuildFromAllocation(
std::unique_ptr<Allocation> allocation,
TfLiteVerifier* extra_verifier = nullptr,
ErrorReporter* error_reporter = T::GetDefaultErrorReporter()) {
error_reporter = ValidateErrorReporter(error_reporter);
if (!allocation || !allocation->valid()) {
TF_LITE_REPORT_ERROR(error_reporter,
"The model allocation is null/empty");
return nullptr;
}
{
// Flatbuffers can only be smaller than 2GB. The file format appends some
// data after the actual flabuffer. We truncate the allocation size to 2GB
// so that the verifier doesn't early exit on us.
size_t allocation_size =
std::min(allocation->bytes(),
static_cast<size_t>(FLATBUFFERS_MAX_BUFFER_SIZE - 1));
flatbuffers::Verifier base_verifier(
reinterpret_cast<const uint8_t*>(allocation->base()), allocation_size,
flatbuffers::Verifier::Options());
if (!VerifyModelBuffer(base_verifier)) {
TF_LITE_REPORT_ERROR(error_reporter,
"The model is not a valid Flatbuffer buffer");
return nullptr;
}
if (extra_verifier &&
!extra_verifier->Verify(static_cast<const char*>(allocation->base()),
allocation_size, error_reporter)) {
// The verifier will have already logged an appropriate error message.
return nullptr;
}
}
return BuildFromAllocationInternal(std::move(allocation), error_reporter);
}
/// Builds a model directly from a flatbuffer pointer
/// Caller retains ownership of the buffer and should keep it alive until the
/// returned object is destroyed. Caller retains ownership of `error_reporter`
/// and must ensure its lifetime is longer than the FlatBufferModelBase
/// instance. Returns a nullptr in case of failure.
static std::unique_ptr<T> BuildFromModel(
const tflite::Model* caller_owned_model_spec,
ErrorReporter* error_reporter = T::GetDefaultErrorReporter()) {
error_reporter = ValidateErrorReporter(error_reporter);
if (CheckBufferOutsideModel(caller_owned_model_spec)) {
TF_LITE_REPORT_ERROR(error_reporter,
"The model contains weights not accessible from "
"tflite::Model *, please use other api");
return nullptr;
}
std::unique_ptr<T> model(new T(caller_owned_model_spec, error_reporter));
if (!model->initialized()) {
model.reset();
} else {
model->ValidateModelBuffers(error_reporter);
}
return model;
}
// Releases memory or unmaps mmaped memory.
~FlatBufferModelBase() = default;
// Copying or assignment is disallowed to simplify ownership semantics.
FlatBufferModelBase(const FlatBufferModelBase&) = delete;
FlatBufferModelBase& operator=(const FlatBufferModelBase&) = delete;
bool initialized() const { return model_ != nullptr; }
const tflite::Model* operator->() const { return model_; }
const tflite::Model* GetModel() const { return model_; }
ErrorReporter* error_reporter() const { return error_reporter_; }
const Allocation* allocation() const { return allocation_.get(); }
// Returns the minimum runtime version from the flatbuffer. This runtime
// version encodes the minimum required interpreter version to run the
// flatbuffer model. If the minimum version can't be determined, an empty
// string will be returned.
// Note that the returned minimum version is a lower-bound but not a strict
// lower-bound; ops in the graph may not have an associated runtime version,
// in which case the actual required runtime might be greater than the
// reported minimum.
std::string GetMinimumRuntime() const {
if (!model_ || !model_->metadata()) return "";
for (int i = 0; i < model_->metadata()->size(); ++i) {
auto metadata = model_->metadata()->Get(i);
if (metadata->name()->str() == tflite_metadata_min_runtime_version) {
auto buf = metadata->buffer();
auto* buffer = (*model_->buffers())[buf];
auto* array = buffer->data();
// Get the real length of the runtime string, since there might be
// trailing
// '\0's in the buffer.
for (int len = 0; len < array->size(); ++len) {
if (array->data()[len] == '\0') {
return std::string(reinterpret_cast<const char*>(array->data()),
len);
}
}
// If there is no '\0' in the buffer, this indicates that the flatbuffer
// is malformed.
TF_LITE_REPORT_ERROR(
error_reporter_,
"Min_runtime_version in model metadata is malformed");
break;
}
}
return "";
}
// Return model metadata as a mapping of name & buffer strings.
// See Metadata table in TFLite schema.
std::map<std::string, std::string> ReadAllMetadata() const {
return ReadAllMetadata(model_);
}
// // Return model metadata as a mapping of name & buffer strings.
// // See Metadata table in TFLite schema.
static std::map<std::string, std::string> ReadAllMetadata(
const ::tflite::Model* model) {
std::map<std::string, std::string> keys_values;
if (!model || !model->metadata() || !model->buffers()) return keys_values;
for (size_t i = 0; i < model->metadata()->size(); ++i) {
auto metadata = model->metadata()->Get(i);
auto buf = metadata->buffer();
if (buf >= model->buffers()->size()) continue;
const tflite::Buffer* buffer = (*model->buffers())[buf];
if (!buffer || !buffer->data()) continue;
const flatbuffers::Vector<uint8_t>* array = buffer->data();
if (!array) continue;
std::string val = std::string(
reinterpret_cast<const char*>(array->data()), array->size());
// Skip if key or value of metadata is empty.
if (!metadata->name() || val.empty()) continue;
keys_values[metadata->name()->str()] = val;
}
return keys_values;
}
// Validates if the FlatBufferModelBase's buffer is well-formed. Specifically,
// it checks if the 0th entry of the model buffers is an empty buffer
// (sentinel). This is a convention so that tensors without a buffer can
// provide 0 as their buffer. NOTE: The function doesn't explicitly fail for
// backward compatibility reasons; it just provides a warning in case of
// failures.
void ValidateModelBuffers(ErrorReporter* error_reporter) {
auto buffers = model_->buffers();
if (buffers && !buffers->empty()) {
auto first_buffer = buffers->Get(0);
if (first_buffer && first_buffer->size() != 0) {
// Note the 0th entry of this array must be an empty buffer (sentinel).
// This is a convention so that tensors without a buffer can provide 0
// as their buffer.
TF_LITE_REPORT_ERROR(
error_reporter,
"The 0th entry of the model buffer must be an empty buffer.");
}
}
}
/// Returns true if the model identifier is correct (otherwise false and
/// reports an error).
bool CheckModelIdentifier() const {
if (allocation_->bytes() < 7) {
TF_LITE_REPORT_ERROR(
error_reporter_,
"Model provided must have at least 7 bytes to hold identifier.\n");
return false;
}
if (!tflite::ModelBufferHasIdentifier(allocation_->base())) {
const char* ident = flatbuffers::GetBufferIdentifier(allocation_->base());
// Suppress unused variable warning.
(void)ident;
TF_LITE_REPORT_ERROR(
error_reporter_,
"Model provided has model identifier '%c%c%c%c', should be '%s'\n",
ident[0], ident[1], ident[2], ident[3], tflite::ModelIdentifier());
return false;
}
return true;
}
/// Check If the buffer is stored as part of the Flatbuffer or outside
/// Return false if the buffers are part of the Flatbuffer
static bool CheckBufferOutsideModel(const tflite::Model* model) {
if (!model || !model->metadata()) return false;
for (int i = 0; i < model->metadata()->size(); ++i) {
auto metadata = model->metadata()->Get(i);
if (metadata->name()->str() == tflite_metadata_buffer_location) {
return true;
}
}
return false;
}
protected:
/// Loads a model from a given allocation. FlatBufferModelBase will take over
/// the ownership of `allocation`, and delete it in destructor. The ownership
/// of `error_reporter`remains with the caller and must have lifetime at least
/// as much as FlatBufferModelBase. This is to allow multiple models to use
/// the same ErrorReporter instance.
explicit FlatBufferModelBase(
std::unique_ptr<Allocation> allocation,
ErrorReporter* error_reporter = T::GetDefaultErrorReporter())
: error_reporter_(ValidateErrorReporter(error_reporter)),
allocation_(std::move(allocation)) {
if (!allocation_ || !allocation_->valid() || !CheckModelIdentifier()) {
return;
}
model_ = ::tflite::GetModel(allocation_->base());
}
/// Loads a model from Model flatbuffer. The `model` has to remain alive and
/// unchanged until the end of this flatbuffer model's lifetime.
FlatBufferModelBase(const Model* model, ErrorReporter* error_reporter)
: model_(model), error_reporter_(ValidateErrorReporter(error_reporter)) {}
static ErrorReporter* ValidateErrorReporter(ErrorReporter* error_reporter) {
return error_reporter ? error_reporter : T::GetDefaultErrorReporter();
}
/// Flatbuffer traverser pointer. (Model* is a pointer that is within the
/// allocated memory of the data allocated by allocation's internals.
const tflite::Model* model_ = nullptr;
/// The error reporter to use for model errors and subsequent errors when
/// the interpreter is created
ErrorReporter* error_reporter_;
/// The allocator used for holding memory of the model. Note that this will
/// be null if the client provides a tflite::Model directly.
std::unique_ptr<Allocation> allocation_;
};
} // namespace impl
} // namespace tflite
#endif // TENSORFLOW_COMPILER_MLIR_LITE_CORE_MODEL_BUILDER_BASE_H_
tensorflow/compiler/mlir/lite/core/special_rules.bzl
+19
View File
@@ -0,0 +1,19 @@
# Copyright 2022 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Build macros for C++ headers."""
def macros_visibility_allowlist():
"""Returns a list of packages that can depend on macros.h."""
return []
tensorflow/compiler/mlir/lite/debug/BUILD
+69
View File
@@ -0,0 +1,69 @@
load("@rules_cc//cc:cc_library.bzl", "cc_library")
load("//tensorflow:tensorflow.bzl", "tf_cc_test")
load(
"//tensorflow/core/platform:build_config.bzl",
"tf_proto_library",
)
# copybara:uncomment package(default_applicable_licenses = ["//tensorflow:LICENSE"])
# OSS only: This target is header-only. Link `debug_options_proto_cc_impl` only to
# `libtensorflow_framework.so` via `lib_internal_impl`. Do NOT link `debug_options_proto_cc_impl`
# directly unless the target does not link `libtensorflow_framework.so`.
tf_proto_library(
name = "debug_options_proto",
srcs = ["debug_options.proto"],
make_default_target_header_only = True,
visibility = ["//visibility:public"],
)
cc_library(
name = "debug",
srcs = ["debug.cc"],
hdrs = ["debug.h"],
visibility = ["//tensorflow/compiler/mlir/lite:__subpackages__"],
deps = [
":debug_options_proto_cc",
"//tensorflow/compiler/mlir/lite/metrics:error_collector_inst",
"//tensorflow/core:portable_gif_internal",
"@com_google_absl//absl/container:flat_hash_set",
"@com_google_absl//absl/log",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/time",
"@com_googlesource_code_re2//:re2",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Pass",
"@llvm-project//mlir:Support",
"@llvm-project//mlir:Transforms",
"@tsl//tsl/platform:path",
"@tsl//tsl/platform:stringpiece",
"@xla//xla/tsl/lib/io:buffered_file",
"@xla//xla/tsl/platform:env",
],
)
tf_cc_test(
name = "debug_test",
srcs = ["debug_test.cc"],
deps = [
":debug",
":debug_options_proto_cc",
"//tensorflow/compiler/mlir/lite:tensorflow_lite",
"//tensorflow/core:portable_gif_internal",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings",
"@com_google_googletest//:gtest_main",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:ArithDialect",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Pass",
"@llvm-project//mlir:Support",
"@tsl//tsl/platform:path",
"@xla//xla/tsl/lib/core:status_test_util",
"@xla//xla/tsl/platform:env",
],
)
tensorflow/compiler/mlir/lite/debug/debug.cc
+358
View File
@@ -0,0 +1,358 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/compiler/mlir/lite/debug/debug.h"
#include <stddef.h>
#include <stdint.h>
#include <functional>
#include <memory>
#include <string>
#include <utility>
#include <vector>
#include "absl/container/flat_hash_set.h"
#include "absl/log/log.h"
#include "absl/status/status.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/str_join.h"
#include "absl/time/clock.h"
#include "absl/time/time.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/ADT/StringRef.h"
#include "llvm/Support/raw_ostream.h"
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/BuiltinAttributes.h" // from @llvm-project
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "mlir/IR/Operation.h" // from @llvm-project
#include "mlir/IR/OperationSupport.h" // from @llvm-project
#include "mlir/Pass/Pass.h" // from @llvm-project
#include "mlir/Pass/PassInstrumentation.h" // from @llvm-project
#include "mlir/Pass/PassManager.h" // from @llvm-project
#include "mlir/Support/FileUtilities.h" // from @llvm-project
#include "mlir/Transforms/ViewOpGraph.h" // from @llvm-project
#include "re2/re2.h" // IWYU pragma: keep
#include "tensorflow/compiler/mlir/lite/debug/debug_options.pb.h"
#include "tensorflow/compiler/mlir/lite/metrics/error_collector_inst.h"
#include "xla/tsl/lib/io/buffered_file.h"
#include "xla/tsl/platform/env.h"
#include "xla/tsl/platform/file_system.h"
#include "tensorflow/core/platform/logging.h"
#include "tsl/platform/path.h"
#include "tsl/platform/stringpiece.h"
// IWYU pragma: no_include "util/regexp/re2/re2.h"
namespace tensorflow {
namespace {
// Simple raw_ostream that prints to a file.
struct WritableFileRawStream : public llvm::raw_ostream {
explicit WritableFileRawStream(std::unique_ptr<tsl::WritableFile> file)
: file(std::move(file)) {
SetUnbuffered();
}
~WritableFileRawStream() override = default;
uint64_t current_pos() const override {
int64_t position;
if (file->Tell(&position).ok()) {
return position;
} else {
// MLIR uses os.tell() to determine whether something was written by
// a subroutine or not, so it's important we have a working current_pos().
LOG(WARNING)
<< "Couldn't query file position. Stream might be malformed.\n";
return -1;
}
}
void write_impl(const char* ptr, size_t size) override {
// Write the file if it is still valid. If the write fails, null out the
// file to avoid encountering another error.
if (file && !file->Append(absl::string_view(ptr, size)).ok()) {
file = nullptr;
}
}
// The file being written to.
std::unique_ptr<tsl::WritableFile> file;
};
// Reproducer stream that emits a reproducer to the given `llvm::raw_ostream`.
class ReproducerStream : public mlir::ReproducerStream {
public:
ReproducerStream(std::string name, std::unique_ptr<llvm::raw_ostream> os)
: name_(std::move(name)), os_(std::move(os)) {}
llvm::StringRef description() override { return name_; }
llvm::raw_ostream& os() override { return *os_; }
private:
std::string name_;
std::unique_ptr<llvm::raw_ostream> os_;
};
// Returns a function that builds a reproducer stream, or nullptr if the MLIR
// reproducer will not be enabled.
mlir::ReproducerStreamFactory GetReproducerStreamFactory(
absl::string_view dump_dir) {
std::string path = tsl::io::JoinPath(dump_dir, "tfl_mlir_crash_repro.mlir");
return [path = std::move(path)](
std::string& error) -> std::unique_ptr<mlir::ReproducerStream> {
std::unique_ptr<tsl::WritableFile> file;
if (auto status = tsl::Env::Default()->NewWritableFile(path, &file);
!status.ok()) {
error = status.ToString();
absl::StrAppend(&error, "; failed to open '", path,
"' for writing an MLIR reproducer");
return nullptr;
}
file = std::make_unique<tsl::BufferedWritableFile>(std::move(file));
return std::make_unique<ReproducerStream>(
path, std::make_unique<WritableFileRawStream>(std::move(file)));
};
}
// Removes unwanted characters for readability and to eliminate issues when
// saving a file.
std::string Sanitize(absl::string_view string) {
static const auto& kUnwantedChars = *new absl::flat_hash_set<char>{
'<', '>', ':', '\"', '/', '\\', '|', '?', '*', ' ', '(', ')'};
std::string sanitized;
sanitized.reserve(string.size());
bool skip = false;
for (const char& c : string) {
if (auto it = kUnwantedChars.find(c); it != kUnwantedChars.end()) {
skip = true;
continue;
}
if (skip) {
skip = false;
sanitized.push_back('_');
}
sanitized.push_back(c);
}
return sanitized;
}
// Pass instrumentation that dumps MLIR based on the criteria specified by
// `ir_dump_*` debug options.
//
// While `mlir::PassManager::enableIRPrinting` provides a similar functionality,
// it is cumbersome to manually copy printed IRs and run them with `litert-opt`.
// Also, long MLIR dumps are often truncated during printing. Instead, this
// instrumentation dumps MLIR to external directories for convenience.
class DumpInstrumentation : public mlir::PassInstrumentation {
public:
explicit DumpInstrumentation(absl::string_view dump_dir,
absl::string_view dump_pass_regex,
absl::string_view dump_func_regex)
: dump_dir_(dump_dir),
dump_pass_re_(std::make_unique<RE2>(dump_pass_regex)),
dump_func_re_(std::make_unique<RE2>(dump_func_regex)) {}
DumpInstrumentation(const DumpInstrumentation& other) = delete;
DumpInstrumentation& operator=(const DumpInstrumentation& other) = delete;
void runBeforePass(mlir::Pass* pass, mlir::Operation* op) override {
// Always print before the first pass.
if (!printed_) {
Dump("before_all", op);
printed_ = true;
}
if (RE2::FullMatch(pass->getName(), *dump_pass_re_)) {
Dump(absl::StrCat(absl::string_view(pass->getName()), "_before"), op,
absl::StrCat(absl::Hex(pass_counter_, absl::kZeroPad8)));
}
}
void runAfterPass(mlir::Pass* pass, mlir::Operation* op) override {
if (RE2::FullMatch(pass->getName(), *dump_pass_re_)) {
Dump(absl::StrCat(absl::string_view(pass->getName()), "_after"), op,
absl::StrCat(absl::Hex(pass_counter_++, absl::kZeroPad8)));
}
}
private:
// Dumps the given op. `name` is used as part of the filename to help
// distinguish dumps at different passes.
void Dump(absl::string_view name, mlir::Operation* op,
std::string prefix = "") {
static constexpr char kFiletypeSuffix[] = "mlir";
// Find names of all func ops with public visibility and check whether at
// least one of them matches `dump_func_re_`.
llvm::SmallVector<absl::string_view> func_names;
bool match = false;
op->walk([&](mlir::func::FuncOp func) {
if (func.isPublic()) {
const absl::string_view name = func.getSymName();
if (name.empty()) {
return;
}
func_names.push_back(name);
if (RE2::FullMatch(name, *dump_func_re_)) {
match = true;
}
}
});
if (!func_names.empty() && !match) {
return;
}
// Sort the function names for determinism.
llvm::sort(func_names);
std::string joined_func_names = Sanitize(absl::StrJoin(func_names, "-"));
std::string sanitized_name = Sanitize(name);
std::vector<absl::string_view> name_parts;
if (!prefix.empty()) {
name_parts.emplace_back(prefix);
}
if (!joined_func_names.empty()) {
name_parts.emplace_back(joined_func_names);
}
name_parts.emplace_back(sanitized_name);
name_parts.emplace_back(kFiletypeSuffix);
// Build a filename such that it contains function names and pass names for
// easy disambiguation.
const std::string filename = tsl::io::JoinPath(
dump_dir_, absl::StrJoin(name_parts.begin(), name_parts.end(), "."));
// Open the file for dumping. Failures are logged instead of being
// propagated to the client because they are non-fatal.
std::unique_ptr<tsl::WritableFile> file;
if (auto status = tsl::Env::Default()->NewWritableFile(filename, &file);
!status.ok()) {
LOG(ERROR) << "Unable to open '" << filename
<< "' for dumping TFLite MLIR output: " << status;
return;
}
file = std::make_unique<tsl::BufferedWritableFile>(std::move(file));
WritableFileRawStream os(std::move(file));
op->print(os);
}
const std::string dump_dir_;
const std::unique_ptr<RE2> dump_pass_re_;
const std::unique_ptr<RE2> dump_func_re_;
// Counter used for pass name prefix to signify sequence
int pass_counter_ = 0;
bool printed_ = false;
};
std::function<bool(mlir::Pass*, mlir::Operation*)> CreatePrintIRFun(
const std::string& pass_regex) {
std::function<bool(mlir::Pass*, mlir::Operation*)> fun;
if (pass_regex.empty()) {
return fun;
}
return [pr = pass_regex](mlir::Pass* p, mlir::Operation*) {
static const RE2* const re = new RE2(pr);
if (RE2::FullMatch(p->getName(), *re)) {
return true;
}
return false;
};
}
} // namespace
void InitPassManager(mlir::PassManager& pm,
const converter::DebugOptions& options,
llvm::raw_ostream& out) {
std::string dump_dir = options.ir_dump_dir();
bool dump_to_dir = !dump_dir.empty();
bool print_to_stdout =
!options.print_ir_before().empty() || !options.print_ir_after().empty();
if (dump_to_dir || print_to_stdout) {
// Necessary for maintaining sequence of passes when dumping MLIR to files
// or stdout.
pm.getContext()->disableMultithreading();
}
if (dump_to_dir) {
dump_dir = tsl::io::JoinPath(
dump_dir, absl::FormatTime("%E4Y%m%d_%H%M%E6S", absl::Now(),
absl::LocalTimeZone()));
// Get a valid file path to dump with.
tsl::Env* env = tsl::Env::Default();
if (auto status = env->RecursivelyCreateDir(dump_dir); !status.ok()) {
LOG(WARNING) << "Failed to create '" << dump_dir
<< "' directory for dumping: " << status;
return;
}
// Set a default crash reproducer for easier debugging.
if (auto reproducer_stream_factory = GetReproducerStreamFactory(dump_dir)) {
pm.enableCrashReproducerGeneration(std::move(reproducer_stream_factory));
}
pm.addInstrumentation(std::make_unique<DumpInstrumentation>(
dump_dir, options.ir_dump_pass_regex(), options.ir_dump_func_regex()));
}
if (print_to_stdout) {
std::function<bool(mlir::Pass*, mlir::Operation*)>
should_print_ir_before_pass(
CreatePrintIRFun(options.print_ir_before()));
std::function<bool(mlir::Pass*, mlir::Operation*)>
should_print_ir_after_pass(CreatePrintIRFun(options.print_ir_after()));
mlir::OpPrintingFlags opPrintingFlags = mlir::OpPrintingFlags();
if (options.has_elide_elementsattrs_if_larger()) {
opPrintingFlags.elideLargeElementsAttrs(
options.elide_elementsattrs_if_larger());
}
pm.enableIRPrinting(should_print_ir_before_pass, should_print_ir_after_pass,
options.print_ir_module_scope(),
/*printAfterOnlyOnChange=*/true,
/*printAfterOnlyOnFailure=*/false, out,
opPrintingFlags);
}
// Enable pass timing. Note: MLIR expects `mlir::PassManager::enableTiming` to
// be called after all instrumentations are added.
if (options.enable_timing()) {
pm.enableTiming();
}
pm.addInstrumentation(
std::make_unique<mlir::TFL::ErrorCollectorInstrumentation>(
pm.getContext()));
}
} // namespace tensorflow
tensorflow/compiler/mlir/lite/debug/debug.h
+34
View File
@@ -0,0 +1,34 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_COMPILER_MLIR_LITE_DEBUG_DEBUG_H_
#define TENSORFLOW_COMPILER_MLIR_LITE_DEBUG_DEBUG_H_
#include "llvm/Support/raw_ostream.h"
#include "mlir/Pass/PassManager.h" // from @llvm-project
#include "tensorflow/compiler/mlir/lite/debug/debug_options.pb.h"
namespace tensorflow {
// Initializes the pass manager with default options that make debugging easier.
// The `out` method parameter is exposed for testing purposes and not intended
// to be specified by client code.
void InitPassManager(mlir::PassManager& pm,
const converter::DebugOptions& options,
llvm::raw_ostream& out = llvm::outs());
} // namespace tensorflow
#endif // TENSORFLOW_COMPILER_MLIR_LITE_DEBUG_DEBUG_H_
tensorflow/compiler/mlir/lite/debug/debug_options.proto
+61
View File
@@ -0,0 +1,61 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
syntax = "proto2";
package tensorflow.converter;
// Additional parameters that control the debug behavior of the Converter.
//
// Next ID: 9
message DebugOptions {
// If not empty, dumps MLIR to the specified directory. The initial state of
// the MLIR after import will be dumped at the beginning of each pass manager
// run. Additionally, MLIR will be dumped before and after each pass
// depending on pass names and functions matched using the
// `ir_dump_pass_regex` and `ir_dump_func_regex` values.
optional string ir_dump_dir = 1 [default = ""];
// Regular expression that matches the names of passes in pascal case (e.g.,
// FooPass) before/after which MLIR will be dumped. Effective only if
// ir_dump_dir is not empty.
optional string ir_dump_pass_regex = 2 [default = ".*"];
// Regular expression that matches the names of functions to be dumped. MLIR
// modules are dumped only if there's at least one public function in the
// module whose name matches the pattern. Effective only if ir_dump_dir is
// not empty.
optional string ir_dump_func_regex = 3 [default = ".*"];
// If true, report the execution time of each MLIR pass.
optional bool enable_timing = 4 [default = false];
// Prints MLIR before specified passes. Supports regular expressions for
// matching against the names of the desired passes.
optional string print_ir_before = 5 [default = ""];
// Prints MLIR after specified passes. Supports regular expressions for
// matching against the names of the desired passes. Currently only prints
// after a pass if the MLIR is mutated.
optional string print_ir_after = 6 [default = ""];
// If true, always print the top-level operation when printing IR for
// print_ir_[before|after].
optional bool print_ir_module_scope = 7 [default = true];
// Elide ElementsAttrs with \"...\" that have more elements than the given
// upper limit.
optional int64 elide_elementsattrs_if_larger = 8;
}
tensorflow/compiler/mlir/lite/debug/debug_test.cc
+288
View File
@@ -0,0 +1,288 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/compiler/mlir/lite/debug/debug.h"
#include <stdint.h>
#include <cstdlib>
#include <memory>
#include <string>
#include <vector>
#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include "absl/status/status.h"
#include "absl/strings/str_cat.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/Support/raw_ostream.h"
#include "mlir/Dialect/Arith/IR/Arith.h" // from @llvm-project
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/Builders.h" // from @llvm-project
#include "mlir/IR/BuiltinAttributes.h" // from @llvm-project
#include "mlir/IR/BuiltinDialect.h" // from @llvm-project
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/IR/BuiltinTypes.h" // from @llvm-project
#include "mlir/IR/Dialect.h" // from @llvm-project
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "mlir/IR/OpDefinition.h" // from @llvm-project
#include "mlir/IR/Operation.h" // from @llvm-project
#include "mlir/IR/OwningOpRef.h" // from @llvm-project
#include "mlir/Pass/Pass.h" // from @llvm-project
#include "mlir/Pass/PassManager.h" // from @llvm-project
#include "mlir/Support/LogicalResult.h" // from @llvm-project
#include "mlir/Support/TypeID.h" // from @llvm-project
#include "tensorflow/compiler/mlir/lite/debug/debug_options.pb.h"
#include "tensorflow/compiler/mlir/lite/ir/tfl_ops.h"
#include "xla/tsl/lib/core/status_test_util.h"
#include "xla/tsl/platform/env.h"
#include "tensorflow/core/platform/types.h"
#include "tsl/platform/path.h"
namespace tensorflow {
namespace debug_test {
class NopPass : public mlir::PassWrapper<NopPass, mlir::OperationPass<>> {
public:
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(NopPass)
void runOnOperation() override {}
};
class MutatePass : public mlir::PassWrapper<MutatePass, mlir::OperationPass<>> {
public:
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(MutatePass)
void runOnOperation() override {
mlir::OpBuilder builder(&getContext());
getOperation()->setAttr("tfl.random_attr", builder.getUnitAttr());
}
};
class AlwaysFailPass
: public mlir::PassWrapper<AlwaysFailPass, mlir::OperationPass<>> {
public:
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(AlwaysFailPass)
void runOnOperation() override { signalPassFailure(); }
};
} // namespace debug_test
namespace {
using ::testing::HasSubstr;
using ::testing::IsEmpty;
using ::testing::Not;
using namespace tensorflow::debug_test;
class InitPassManagerTest : public testing::Test {
protected:
InitPassManagerTest()
: path_(GetOutputPath()), context_([]() {
mlir::registerPassManagerCLOptions();
mlir::DialectRegistry registry;
registry.insert<mlir::BuiltinDialect>();
registry.insert<mlir::arith::ArithDialect>();
registry.insert<mlir::func::FuncDialect>();
registry.insert<mlir::TFL::TensorFlowLiteDialect>();
return registry;
}()) {
context_.loadAllAvailableDialects();
mlir::OpBuilder builder(&context_);
module_ = mlir::ModuleOp::create(builder, builder.getUnknownLoc());
builder.setInsertionPointToStart(module_->getBody());
auto func = mlir::func::FuncOp::create(builder, //
builder.getUnknownLoc(), "main",
builder.getFunctionType({}, {}));
func->setAttr("tfl.func", builder.getUnitAttr());
builder.setInsertionPointToStart(func.addEntryBlock());
llvm::SmallVector<int> shape{1, 2, 3, 4, 5, 6, 7, 8, 9, 10};
mlir::arith::ConstantOp::create(
builder, builder.getUnknownLoc(),
mlir::DenseIntElementsAttr::get(
mlir::RankedTensorType::get(shape.size(), builder.getI32Type()),
shape));
mlir::func::ReturnOp::create(builder, builder.getUnknownLoc());
}
absl::Status GetDumpDir(std::string* dump_dir) {
std::vector<std::string> files;
if (auto status = tsl::Env::Default()->GetChildren(path_, &files);
!status.ok()) {
return status;
}
if (files.size() != 1) {
return absl::FailedPreconditionError(
"Expecting directory to have one child.");
}
*dump_dir = tsl::io::JoinPath(path_, files[0]);
return absl::OkStatus();
}
std::string path_;
mlir::MLIRContext context_;
mlir::OwningOpRef<mlir::ModuleOp> module_;
private:
std::string GetOutputPath() {
const auto* const test_info =
testing::UnitTest::GetInstance()->current_test_info();
return tsl::io::JoinPath(
getenv("TEST_UNDECLARED_OUTPUTS_DIR"),
absl::StrCat(test_info->test_suite_name(), ".", test_info->name()));
}
};
TEST_F(InitPassManagerTest, CrashReproducer) {
converter::DebugOptions debug_options;
*debug_options.mutable_ir_dump_dir() = path_;
mlir::PassManager pm(&context_);
InitPassManager(pm, debug_options);
pm.addPass(std::make_unique<AlwaysFailPass>());
ASSERT_TRUE(mlir::failed(pm.run(*module_)));
std::string dump_dir;
TF_ASSERT_OK(GetDumpDir(&dump_dir));
std::string mlir_dump;
TF_ASSERT_OK(tsl::ReadFileToString(
tsl::Env::Default(),
tsl::io::JoinPath(dump_dir, "tfl_mlir_crash_repro.mlir"), &mlir_dump));
EXPECT_THAT(mlir_dump, Not(IsEmpty()));
}
TEST_F(InitPassManagerTest, DumpToDir) {
converter::DebugOptions debug_options;
*debug_options.mutable_ir_dump_dir() = path_;
*debug_options.mutable_ir_dump_pass_regex() = R"(.*NopPass)";
mlir::PassManager pm(&context_);
InitPassManager(pm, debug_options);
pm.addPass(std::make_unique<NopPass>());
ASSERT_TRUE(mlir::succeeded(pm.run(*module_)));
std::string dump_dir;
TF_ASSERT_OK(GetDumpDir(&dump_dir));
{
std::string mlir_dump;
TF_ASSERT_OK(tsl::ReadFileToString(
tsl::Env::Default(),
tsl::io::JoinPath(
dump_dir, "00000000.main.tensorflow_debug_test_NopPass_after.mlir"),
&mlir_dump));
EXPECT_THAT(mlir_dump, Not(IsEmpty()));
}
{
std::string mlir_dump;
TF_ASSERT_OK(tsl::ReadFileToString(
tsl::Env::Default(),
tsl::io::JoinPath(
dump_dir,
"00000000.main.tensorflow_debug_test_NopPass_before.mlir"),
&mlir_dump));
EXPECT_THAT(mlir_dump, Not(IsEmpty()));
}
}
TEST_F(InitPassManagerTest, PrintIRBeforeEverything) {
converter::DebugOptions debug_options;
*debug_options.mutable_print_ir_before() = R"(.*)";
std::string captured_out;
llvm::raw_string_ostream out(captured_out);
mlir::PassManager pm(&context_);
InitPassManager(pm, debug_options, out);
pm.addPass(std::make_unique<NopPass>());
ASSERT_TRUE(mlir::succeeded(pm.run(*module_)));
EXPECT_THAT(captured_out,
HasSubstr("IR Dump Before tensorflow::debug_test::NopPass"));
EXPECT_THAT(captured_out,
Not(HasSubstr("IR Dump After tensorflow::debug_test::NopPass")));
}
TEST_F(InitPassManagerTest, PrintIRAfterEverything) {
converter::DebugOptions debug_options;
*debug_options.mutable_print_ir_after() = R"(.*)";
std::string captured_out;
llvm::raw_string_ostream out(captured_out);
mlir::PassManager pm(&context_);
InitPassManager(pm, debug_options, out);
pm.addPass(std::make_unique<MutatePass>());
ASSERT_TRUE(mlir::succeeded(pm.run(*module_)));
EXPECT_THAT(captured_out,
HasSubstr("IR Dump After tensorflow::debug_test::MutatePass"));
EXPECT_THAT(
captured_out,
Not(HasSubstr("IR Dump Before tensorflow::debug_test::MutatePass")));
}
TEST_F(InitPassManagerTest, PrintIRBeforeAndAfterEverything) {
converter::DebugOptions debug_options;
*debug_options.mutable_print_ir_before() = R"(.*)";
*debug_options.mutable_print_ir_after() = R"(.*)";
std::string captured_out;
llvm::raw_string_ostream out(captured_out);
mlir::PassManager pm(&context_);
InitPassManager(pm, debug_options, out);
pm.addPass(std::make_unique<MutatePass>());
ASSERT_TRUE(mlir::succeeded(pm.run(*module_)));
EXPECT_THAT(captured_out,
HasSubstr("IR Dump After tensorflow::debug_test::MutatePass"));
EXPECT_THAT(captured_out,
HasSubstr("IR Dump Before tensorflow::debug_test::MutatePass"));
}
TEST_F(InitPassManagerTest, ElideLargeElementAttrs) {
converter::DebugOptions debug_options;
*debug_options.mutable_print_ir_before() = R"(.*)";
debug_options.set_elide_elementsattrs_if_larger(5);
std::string captured_out;
llvm::raw_string_ostream out(captured_out);
mlir::PassManager pm(&context_);
InitPassManager(pm, debug_options, out);
pm.addPass(std::make_unique<MutatePass>());
ASSERT_TRUE(mlir::succeeded(pm.run(*module_)));
EXPECT_THAT(captured_out, HasSubstr("dense_resource<__elided__>"));
}
TEST_F(InitPassManagerTest, DontElideSmallerElementAttrs) {
converter::DebugOptions debug_options;
*debug_options.mutable_print_ir_before() = R"(.*)";
debug_options.set_elide_elementsattrs_if_larger(11);
std::string captured_out;
llvm::raw_string_ostream out(captured_out);
mlir::PassManager pm(&context_);
InitPassManager(pm, debug_options, out);
pm.addPass(std::make_unique<MutatePass>());
ASSERT_TRUE(mlir::succeeded(pm.run(*module_)));
EXPECT_THAT(captured_out,
HasSubstr("dense<[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]>"));
}
} // namespace
} // namespace tensorflow
tensorflow/compiler/mlir/lite/delegates/BUILD
+27
View File
@@ -0,0 +1,27 @@
# 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(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:LICENSE"],
default_visibility = ["//visibility:public"],
licenses = ["notice"],
)
# LINT.IfChange(config_setting)
config_setting(
name = "tflite_debug_delegate",
define_values = {"tflite_debug_delegate": "true"},
)
# LINT.ThenChange(//tensorflow/lite/delegates/BUILD)
tensorflow/compiler/mlir/lite/delegates/flex/BUILD
+45
View File
@@ -0,0 +1,45 @@
load("@rules_cc//cc:cc_library.bzl", "cc_library")
load(
"//tensorflow:tensorflow.bzl",
"if_mobile",
"if_not_mobile",
)
load("//tensorflow:tensorflow.default.bzl", "get_compatible_with_portable")
load("//tensorflow/lite:special_rules.bzl", "internal_visibility_allowlist")
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"],
)
cc_library(
name = "allowlisted_flex_ops_lib",
srcs = [
"allowlisted_flex_ops.cc",
],
hdrs = [
"allowlisted_flex_ops.h",
"allowlisted_flex_ops_internal.h",
],
compatible_with = get_compatible_with_portable(),
copts = select({
"//tensorflow:android": ["-Wno-private-header"],
"//conditions:default": [],
}),
visibility = internal_visibility_allowlist(),
deps = if_mobile([
"//tensorflow/core:portable_tensorflow_lib_lite",
]) + if_not_mobile([
"//tensorflow/core:framework",
]),
)
tensorflow/compiler/mlir/lite/delegates/flex/allowlisted_flex_ops.cc
+852
View File
@@ -0,0 +1,852 @@
/* 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/compiler/mlir/lite/delegates/flex/allowlisted_flex_ops.h"
#include <set>
#include <string>
#include "tensorflow/compiler/mlir/lite/delegates/flex/allowlisted_flex_ops_internal.h"
#include "tensorflow/core/framework/op.h"
namespace tflite {
namespace flex {
const std::set<std::string>& GetFlexAllowlist() {
// LINT.IfChange
static const std::set<std::string>* allowlisted_flex_ops =
new std::set<std::string>({
// go/keep-sorted start
"Abort",
"Abs",
"Add",
"AddN",
"AddV2",
"AdjustContrast",
"AdjustContrastv2",
"AdjustHue",
"AdjustSaturation",
"All",
"Angle",
"Any",
"ApplyAdaMax",
"ApplyAdadelta",
"ApplyAdagrad",
"ApplyAdagradDA",
"ApplyAdagradV2",
"ApplyAdam",
"ApplyAddSign",
"ApplyCenteredRMSProp",
"ApplyFtrl",
"ApplyFtrlV2",
"ApplyGradientDescent",
"ApplyMomentum",
"ApplyPowerSign",
"ApplyProximalAdagrad",
"ApplyProximalGradientDescent",
"ApplyRMSProp",
"ApproximateEqual",
"ArgMax",
"ArgMin",
"AsString",
"Assert",
"Assign",
"AssignAdd",
"AssignAddVariableOp",
"AssignSub",
"AssignSubVariableOp",
"AssignVariableOp",
"Atan",
"Atan2",
"AudioSpectrogram",
"AvgPool",
"AvgPool3D",
"AvgPool3DGrad",
"AvgPoolGrad",
"BatchCholesky",
"BatchDatasetV2",
"BatchMatMul",
"BatchMatMulV2",
"BatchMatrixBandPart",
"BatchMatrixDeterminant",
"BatchMatrixDiag",
"BatchMatrixDiagPart",
"BatchMatrixInverse",
"BatchMatrixSetDiag",
"BatchMatrixTriangularSolve",
"BatchNormWithGlobalNormalization",
"BatchNormWithGlobalNormalizationGrad",
"BatchToSpace",
"BatchToSpaceND",
"BiasAdd",
"BiasAddGrad",
"BiasAddV1",
"Bincount",
"Bitcast",
"BitwiseAnd",
"BitwiseOr",
"BitwiseXor",
"BroadcastArgs",
"BroadcastGradientArgs",
"BroadcastTo",
"Bucketize",
"CTCBeamSearchDecoder",
"CTCGreedyDecoder",
"Case",
"Cast",
"Ceil",
"CheckNumerics",
"CheckNumericsV2",
"Cholesky",
"ClipByValue",
"CombinedNonMaxSuppression",
"Complex",
"ComplexAbs",
"Concat",
"ConcatOffset",
"ConcatV2",
"Conj",
"ConjugateTranspose",
"Const",
"ControlTrigger",
"Conv2D",
"Conv2DBackpropFilter",
"Conv2DBackpropInput",
"Conv3D",
"Conv3DBackpropFilter",
"Conv3DBackpropFilterV2",
"Conv3DBackpropInput",
"Conv3DBackpropInputV2",
"Cos",
"Cosh",
"CropAndResize",
"CropAndResizeGradBoxes",
"CropAndResizeGradImage",
"Cumprod",
"Cumsum",
"CumulativeLogsumexp",
"DataFormatDimMap",
"DataFormatVecPermute",
"DebugGradientIdentity",
"DebugGradientRefIdentity",
"DecodeAndCropJpeg",
"DecodeBase64",
"DecodeBmp",
"DecodeGif",
"DecodeImage",
"DecodeJpeg",
"DecodePaddedRaw",
"DecodePng",
"DecodeRaw",
"DecodeWav",
"DeepCopy",
"DeleteSessionTensor",
"DenseBincount",
"DenseToDenseSetOperation",
"DenseToSparseSetOperation",
"DepthToSpace",
"DepthwiseConv2dNative",
"DepthwiseConv2dNativeBackpropFilter",
"DepthwiseConv2dNativeBackpropInput",
"Dequantize",
"DestroyResourceOp",
"DestroyTemporaryVariable",
"Diag",
"DiagPart",
"Dilation2D",
"Dilation2DBackpropFilter",
"Dilation2DBackpropInput",
"Div",
"DivNoNan",
"DynamicPartition",
"DynamicStitch",
"Einsum",
"Elu",
"EluGrad",
"Empty",
"EmptyTensorList",
"EmptyTensorMap",
"EncodeBase64",
"EncodeJpeg",
"EncodeJpegVariableQuality",
"EncodePng",
"EncodeWav",
"EnsureShape",
"Enter",
"Equal",
"Erf",
"Exit",
"Exp",
"ExpandDims",
"ExtractImagePatches",
"FFT",
"FFT2D",
"FFT3D",
"FIFOQueue",
"FIFOQueueV2",
"FakeQuantWithMinMaxArgs",
"FakeQuantWithMinMaxArgsGradient",
"FakeQuantWithMinMaxVars",
"FakeQuantWithMinMaxVarsGradient",
"FakeQuantWithMinMaxVarsPerChannel",
"FakeQuantWithMinMaxVarsPerChannelGradient",
"FakeQueue",
"Fill",
"FilterDataset",
"FinalizeDataset",
"Fingerprint",
"FlatMapDataset",
"Floor",
"FloorDiv",
"FloorMod",
"FusedBatchNorm",
"FusedBatchNormGrad",
"FusedBatchNormGradV2",
"FusedBatchNormGradV3",
"FusedBatchNormV2",
"FusedBatchNormV3",
"FusedPadConv2D",
"FusedResizeAndPadConv2D",
"Gather",
"GatherNd",
"GatherV2",
"GetSessionHandle",
"GetSessionHandleV2",
"GetSessionTensor",
"Greater",
"GreaterEqual",
"HSVToRGB",
"HashTable",
"HashTableV2",
"HistogramSummary",
"IFFT",
"IFFT2D",
"IFFT3D",
"IRFFT",
"IRFFT2D",
"IRFFT3D",
"Identity",
"IdentityN",
"Imag",
"ImageProjectiveTransformV2",
"ImageProjectiveTransformV3",
"ImmutableConst",
"InTopK",
"InTopKV2",
"InitializeTable",
"InitializeTableFromDataset",
"InitializeTableFromTextFile",
"InitializeTableFromTextFileV2",
"InitializeTableV2",
"InplaceAdd",
"InplaceSub",
"InplaceUpdate",
"Inv",
"InvGrad",
"Invert",
"InvertPermutation",
"IsFinite",
"IsNan",
"IsVariableInitialized",
"LRN",
"LeakyRelu",
"LeakyReluGrad",
"LeftShift",
"Less",
"LessEqual",
"LinSpace",
"ListDiff",
"Log",
"LogMatrixDeterminant",
"LogSoftmax",
"LogicalAnd",
"LogicalNot",
"LogicalOr",
"LookupTableExport",
"LookupTableExportV2",
"LookupTableFind",
"LookupTableFindV2",
"LookupTableImport",
"LookupTableImportV2",
"LookupTableInsert",
"LookupTableInsertV2",
"LookupTableRemoveV2",
"LookupTableSize",
"LookupTableSizeV2",
"LoopCond",
"MapDataset",
"MatMul",
"MatrixBandPart",
"MatrixDeterminant",
"MatrixDiag",
"MatrixDiagPart",
"MatrixDiagPartV2",
"MatrixDiagPartV3",
"MatrixDiagV2",
"MatrixDiagV3",
"MatrixInverse",
"MatrixSetDiag",
"MatrixSetDiagV2",
"MatrixSetDiagV3",
"MatrixTriangularSolve",
"Max",
"MaxPool",
"MaxPool3D",
"MaxPool3DGrad",
"MaxPool3DGradGrad",
"MaxPoolGrad",
"MaxPoolGradGrad",
"MaxPoolGradGradV2",
"MaxPoolGradV2",
"MaxPoolGradWithArgmax",
"MaxPoolV2",
"MaxPoolWithArgmax",
"Maximum",
"Mean",
"Merge",
"MergeSummary",
"MergeV2Checkpoints",
"Mfcc",
"Min",
"Minimum",
"MirrorPad",
"MirrorPadGrad",
"ModelDataset",
"Mul",
"MulNoNan",
"Multinomial",
"MutableDenseHashTable",
"MutableDenseHashTableV2",
"MutableHashTable",
"MutableHashTableOfTensors",
"MutableHashTableOfTensorsV2",
"MutableHashTableV2",
"Neg",
"NextIteration",
"NoOp",
"NonMaxSuppression",
"NonMaxSuppressionV2",
"NonMaxSuppressionV3",
"NonMaxSuppressionV4",
"NonMaxSuppressionV5",
"NonMaxSuppressionWithOverlaps",
"NotEqual",
"OneHot",
"OnesLike",
"OptimizeDatasetV2",
"OptionalFromValue",
"OptionalGetValue",
"OptionalHasValue",
"OptionalNone",
"Pack",
"Pad",
"PadV2",
"PaddingFIFOQueue",
"PaddingFIFOQueueV2",
"ParallelConcat",
"ParallelDynamicStitch",
"ParseExample",
"ParseExampleV2",
"ParseSequenceExample",
"ParseSequenceExampleV2",
"ParseSingleExample",
"ParseSingleSequenceExample",
"Placeholder",
"PlaceholderV2",
"PlaceholderWithDefault",
"PopulationCount",
"Pow",
"PreventGradient",
"Print",
"PrintV2",
"Prod",
"Qr",
"QuantizeDownAndShrinkRange",
"QuantizeV2",
"QuantizedAdd",
"QuantizedAvgPool",
"QuantizedBatchNormWithGlobalNormalization",
"QuantizedBiasAdd",
"QuantizedConcat",
"QuantizedConv2D",
"QuantizedInstanceNorm",
"QuantizedMatMul",
"QuantizedMaxPool",
"QuantizedMul",
"QuantizedRelu",
"QuantizedRelu6",
"QuantizedReshape",
"QuantizedResizeBilinear",
"QueueClose",
"QueueCloseV2",
"QueueDequeue",
"QueueDequeueMany",
"QueueDequeueManyV2",
"QueueDequeueUpTo",
"QueueDequeueUpToV2",
"QueueDequeueV2",
"QueueEnqueue",
"QueueEnqueueMany",
"QueueEnqueueManyV2",
"QueueEnqueueV2",
"QueueIsClosed",
"QueueIsClosedV2",
"QueueSize",
"QueueSizeV2",
"RFFT",
"RFFT2D",
"RFFT3D",
"RGBToHSV",
"RaggedBincount",
"RaggedGather",
"RaggedRange",
"RaggedTensorFromVariant",
"RaggedTensorToSparse",
"RaggedTensorToTensor",
"RaggedTensorToVariant",
"RaggedTensorToVariantGradient",
"RandomGamma",
"RandomPoisson",
"RandomPoissonV2",
"RandomShuffle",
"RandomStandardNormal",
"RandomUniform",
"RandomUniformInt",
"Range",
"Rank",
"ReadFile",
"ReadVariableOp",
"Real",
"RealDiv",
"Reciprocal",
"ReciprocalGrad",
"Recv",
"ReduceDataset",
"ReduceJoin",
"RefEnter",
"RefExit",
"RefIdentity",
"RefMerge",
"RefNextIteration",
"RefSelect",
"RefSwitch",
"RegexFullMatch",
"RegexReplace",
"Relu",
"Relu6",
"Relu6Grad",
"ReluGrad",
"RemoteCall",
"RepeatDataset",
"RequantizationRange",
"Requantize",
"Reshape",
"ResizeBicubic",
"ResizeBicubicGrad",
"ResizeBilinear",
"ResizeBilinearGrad",
"ResizeNearestNeighbor",
"ResizeNearestNeighborGrad",
"ResourceApplyAdaMax",
"ResourceApplyAdadelta",
"ResourceApplyAdagrad",
"ResourceApplyAdagradDA",
"ResourceApplyAdagradV2",
"ResourceApplyAdam",
"ResourceApplyAdamWithAmsgrad",
"ResourceApplyAddSign",
"ResourceApplyCenteredRMSProp",
"ResourceApplyFtrl",
"ResourceApplyFtrlV2",
"ResourceApplyGradientDescent",
"ResourceApplyKerasMomentum",
"ResourceApplyMomentum",
"ResourceApplyPowerSign",
"ResourceApplyProximalAdagrad",
"ResourceApplyProximalGradientDescent",
"ResourceApplyRMSProp",
"ResourceGather",
"ResourceGatherNd",
"ResourceScatterAdd",
"ResourceScatterDiv",
"ResourceScatterMax",
"ResourceScatterMin",
"ResourceScatterMul",
"ResourceScatterNdAdd",
"ResourceScatterNdMax",
"ResourceScatterNdMin",
"ResourceScatterNdSub",
"ResourceScatterNdUpdate",
"ResourceScatterSub",
"ResourceScatterUpdate",
"ResourceSparseApplyAdadelta",
"ResourceSparseApplyAdagrad",
"ResourceSparseApplyAdagradDA",
"ResourceSparseApplyAdagradV2",
"ResourceSparseApplyCenteredRMSProp",
"ResourceSparseApplyFtrl",
"ResourceSparseApplyFtrlV2",
"ResourceSparseApplyKerasMomentum",
"ResourceSparseApplyMomentum",
"ResourceSparseApplyProximalAdagrad",
"ResourceSparseApplyProximalGradientDescent",
"ResourceSparseApplyRMSProp",
"ResourceStridedSliceAssign",
"Restore",
"RestoreSlice",
"RestoreV2",
"Reverse",
"ReverseSequence",
"ReverseV2",
"RightShift",
"Roll",
"Round",
"Rsqrt",
"RsqrtGrad",
"SampleDistortedBoundingBox",
"SampleDistortedBoundingBoxV2",
"Save",
"SaveSlices",
"SaveV2",
"ScalarSummary",
"ScatterNd",
"ScatterNdAdd",
"ScatterNdMax",
"ScatterNdMin",
"ScatterNdNonAliasingAdd",
"ScatterNdSub",
"ScatterNdUpdate",
"SegmentMax",
"SegmentMean",
"SegmentMin",
"SegmentProd",
"SegmentSum",
"Select",
"SelectV2",
"Selu",
"SeluGrad",
"Send",
"SerializeTensor",
"Shape",
"ShapeN",
"ShardedFilename",
"ShardedFilespec",
"Sigmoid",
"SigmoidGrad",
"Sign",
"Sin",
"Sinh",
"Size",
"Slice",
"Softmax",
"SoftmaxCrossEntropyWithLogits",
"Softplus",
"SoftplusGrad",
"Softsign",
"SoftsignGrad",
"SpaceToBatch",
"SpaceToBatchND",
"SpaceToDepth",
"SparseAdd",
"SparseApplyAdadelta",
"SparseApplyAdagrad",
"SparseApplyAdagradDA",
"SparseApplyAdagradV2",
"SparseApplyCenteredRMSProp",
"SparseApplyFtrl",
"SparseApplyFtrlV2",
"SparseApplyMomentum",
"SparseApplyProximalAdagrad",
"SparseApplyProximalGradientDescent",
"SparseApplyRMSProp",
"SparseBincount",
"SparseCross",
"SparseCrossHashed",
"SparseCrossV2",
"SparseFillEmptyRows",
"SparseFillEmptyRowsGrad",
"SparseReduceSum",
"SparseReorder",
"SparseReshape",
"SparseSegmentMean",
"SparseSegmentMeanGrad",
"SparseSegmentMeanWithNumSegments",
"SparseSegmentSqrtN",
"SparseSegmentSqrtNGrad",
"SparseSegmentSqrtNWithNumSegments",
"SparseSegmentSum",
"SparseSegmentSumGrad",
"SparseSegmentSumWithNumSegments",
"SparseSlice",
"SparseSoftmaxCrossEntropyWithLogits",
"SparseTensorDenseMatMul",
"SparseToDense",
"SparseToSparseSetOperation",
"Split",
"SplitV",
"Sqrt",
"SqrtGrad",
"Square",
"SquaredDifference",
"Squeeze",
"Stack",
"StackClose",
"StackCloseV2",
"StackPop",
"StackPopV2",
"StackPush",
"StackPushV2",
"StackV2",
"StatelessMultinomial",
"StatelessRandomGammaV2",
"StatelessRandomGammaV3",
"StatelessRandomGetAlg",
"StatelessRandomGetKeyCounter",
"StatelessRandomGetKeyCounterAlg",
"StatelessRandomNormal",
"StatelessRandomNormalV2",
"StatelessRandomPoisson",
"StatelessRandomUniform",
"StatelessRandomUniformFullInt",
"StatelessRandomUniformFullIntV2",
"StatelessRandomUniformInt",
"StatelessRandomUniformIntV2",
"StatelessRandomUniformV2",
"StatelessSampleDistortedBoundingBox",
"StatelessTruncatedNormal",
"StatelessTruncatedNormalV2",
"StaticRegexFullMatch",
"StaticRegexReplace",
"StopGradient",
"StridedSlice",
"StridedSliceAssign",
"StridedSliceGrad",
"StringFormat",
"StringJoin",
"StringLength",
"StringLower",
"StringSplit",
"StringSplitV2",
"StringStrip",
"StringToHashBucket",
"StringToHashBucketFast",
"StringToHashBucketStrong",
"StringToNumber",
"Sub",
"Substr",
"Sum",
"Switch",
"SymbolicGradient",
"TakeDataset",
"TakeWhileDataset",
"Tan",
"Tanh",
"TanhGrad",
"TemporaryVariable",
"TensorArray",
"TensorArrayClose",
"TensorArrayCloseV2",
"TensorArrayCloseV3",
"TensorArrayConcat",
"TensorArrayConcatV2",
"TensorArrayConcatV3",
"TensorArrayGather",
"TensorArrayGatherV2",
"TensorArrayGatherV3",
"TensorArrayGrad",
"TensorArrayGradV2",
"TensorArrayGradV3",
"TensorArrayGradWithShape",
"TensorArrayPack",
"TensorArrayRead",
"TensorArrayReadV2",
"TensorArrayReadV3",
"TensorArrayScatter",
"TensorArrayScatterV2",
"TensorArrayScatterV3",
"TensorArraySize",
"TensorArraySizeV2",
"TensorArraySizeV3",
"TensorArraySplit",
"TensorArraySplitV2",
"TensorArraySplitV3",
"TensorArrayUnpack",
"TensorArrayV2",
"TensorArrayV3",
"TensorArrayWrite",
"TensorArrayWriteV2",
"TensorArrayWriteV3",
"TensorListConcat",
"TensorListConcatLists",
"TensorListConcatV2",
"TensorListElementShape",
"TensorListFromTensor",
"TensorListGather",
"TensorListGetItem",
"TensorListLength",
"TensorListPopBack",
"TensorListPushBack",
"TensorListPushBackBatch",
"TensorListReserve",
"TensorListResize",
"TensorListScatter",
"TensorListScatterIntoExistingList",
"TensorListScatterV2",
"TensorListSetItem",
"TensorListSplit",
"TensorListStack",
"TensorMapErase",
"TensorMapHasKey",
"TensorMapInsert",
"TensorMapLookup",
"TensorMapSize",
"TensorMapStackKeys",
"TensorScatterAdd",
"TensorScatterMax",
"TensorScatterMin",
"TensorScatterSub",
"TensorScatterUpdate",
"TensorSliceDataset",
"TensorStridedSliceUpdate",
"Tile",
"TileGrad",
"Timestamp",
"TopK",
"TopKV2",
"Transpose",
"TruncateDiv",
"TruncatedNormal",
"UnicodeDecode",
"UnicodeDecodeWithOffsets",
"UnicodeEncode",
"UnicodeTranscode",
"Unique",
"UniqueV2",
"UniqueWithCounts",
"UniqueWithCountsV2",
"Unpack",
"UnsortedSegmentJoin",
"UnsortedSegmentMax",
"UnsortedSegmentMin",
"UnsortedSegmentProd",
"UnsortedSegmentSum",
"UnwrapDatasetVariant",
"UpperBound",
"VarHandleOp",
"VarIsInitializedOp",
"Variable",
"VariableShape",
"VariableV2",
"Where",
"WrapDatasetVariant",
"WriteFile",
"Xdivy",
"Xlog1py",
"Xlogy",
"ZerosLike",
"_Arg",
"_ArrayToList",
"_DeviceArg",
"_DeviceRetval",
"_FusedConv2D",
"_HostCast",
"_HostRecv",
"_HostSend",
"_ListToArray",
"_ParallelConcatStart",
"_ParallelConcatUpdate",
"_ReadVariablesOp",
"_Recv",
"_Retval",
"_Send",
"_SwitchN",
"_VarHandlesOp",
// go/keep-sorted end
});
// LINT.ThenChange(//tensorflow/lite/g3doc/guide/op_select_allowlist.md)
return *allowlisted_flex_ops;
// Prevent lint error about this function being too long. This function
// is a set of ops, and making it shorter won't help readbility.
// NOLINTNEXTLINE
}
const std::set<std::string>& GetTFTextFlexAllowlist() {
// LINT.IfChange
static const std::set<std::string>* tftext_flex_ops =
new std::set<std::string>({
"CaseFoldUTF8",
"ConstrainedSequence",
"MaxSpanningTree",
"NormalizeUTF8",
"NormalizeUTF8WithOffsetsMap",
"RegexSplitWithOffsets",
"RougeL",
"SentenceFragments",
"SentencepieceOp",
"SentencepieceTokenizeOp",
"SentencepieceTokenizeWithOffsetsOp",
"SentencepieceDetokenizeOp",
"SentencepieceVocabSizeOp",
"SplitMergeTokenizeWithOffsets",
"TFText>NgramsStringJoin",
"TFText>WhitespaceTokenizeWithOffsetsV2",
"TokenizerFromLogits",
"UnicodeScriptTokenizeWithOffsets",
"WhitespaceTokenizeWithOffsets",
"WordpieceTokenizeWithOffsets",
});
// LINT.ThenChange(//tensorflow/lite/g3doc/guide/op_select_allowlist.md)
return *tftext_flex_ops;
}
// Allow the tf.text ops if they are registered in the global op registry.
bool IsAllowedTFTextOpForFlex(const std::string& op_name) {
if (GetTFTextFlexAllowlist().count(op_name) == 0) return false;
return tensorflow::OpRegistry::Global()->LookUp(op_name) != nullptr;
}
const std::set<std::string>& GetSentencePieceFlexAllowlist() {
// LINT.IfChange
static const std::set<std::string>* sentencepiece_flex_ops =
new std::set<std::string>({
"SentencepieceGetPieceSize",
"SentencepiecePieceToId",
"SentencepieceIdToPiece",
"SentencepieceEncodeDense",
"SentencepieceEncodeSparse",
"SentencepieceDecode",
});
// LINT.ThenChange(//tensorflow/lite/g3doc/guide/op_select_allowlist.md)
return *sentencepiece_flex_ops;
}
// Allow the sentencepiece ops if they are registered in the global op registry.
bool IsAllowedSentencePieceOpForFlex(const std::string& op_name) {
if (GetSentencePieceFlexAllowlist().count(op_name) == 0) return false;
return tensorflow::OpRegistry::Global()->LookUp(op_name) != nullptr;
}
bool IsAllowlistedFlexOp(const std::string& tensorflow_op_name) {
if (GetFlexAllowlist().count(tensorflow_op_name) != 0) return true;
// Check if the op is an allowlisted tf.text or sentencepiece op.
return IsAllowedTFTextOpForFlex(tensorflow_op_name) ||
IsAllowedSentencePieceOpForFlex(tensorflow_op_name);
}
} // namespace flex
} // namespace tflite
tensorflow/compiler/mlir/lite/delegates/flex/allowlisted_flex_ops.h
+45
View File
@@ -0,0 +1,45 @@
/* 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_COMPILER_MLIR_LITE_DELEGATES_FLEX_ALLOWLISTED_FLEX_OPS_H_
#define TENSORFLOW_COMPILER_MLIR_LITE_DELEGATES_FLEX_ALLOWLISTED_FLEX_OPS_H_
#include <set>
#include <string>
namespace tflite {
namespace flex {
// Whether the given op has been statically allowlisted for flex export.
//
// This static allowlist is formed by the intersection of ops supported by
// TensorFlowMobile on both iOS and Android. As the converter is likely running
// on a host that has the full suite of TensorFlow ops available, we use this
// static allowlist to ensure compatibility when deploying to a mobile device.
// TODO(b/118389105): Automate generation of the allowlisted flex ops.
bool IsAllowlistedFlexOp(const std::string& tensorflow_op_name);
// Return the list of allowlisted flex ops.
const std::set<std::string>& GetFlexAllowlist();
// Return the list of TF.Text flex ops.
const std::set<std::string>& GetTFTextFlexAllowlist();
// Return the list of SentencePiece flex ops.
const std::set<std::string>& GetSentencePieceFlexAllowlist();
} // namespace flex
} // namespace tflite
#endif // TENSORFLOW_COMPILER_MLIR_LITE_DELEGATES_FLEX_ALLOWLISTED_FLEX_OPS_H_
tensorflow/compiler/mlir/lite/delegates/flex/allowlisted_flex_ops_internal.h
+29
View File
@@ -0,0 +1,29 @@
/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_COMPILER_MLIR_LITE_DELEGATES_FLEX_ALLOWLISTED_FLEX_OPS_INTERNAL_H_
#define TENSORFLOW_COMPILER_MLIR_LITE_DELEGATES_FLEX_ALLOWLISTED_FLEX_OPS_INTERNAL_H_
#include <string>
namespace tflite {
namespace flex {
// Return true if op_name is a tf.text op need to be supported by flex delegate.
bool IsAllowedTFTextOpForFlex(const std::string& op_name);
} // namespace flex
} // namespace tflite
#endif // TENSORFLOW_COMPILER_MLIR_LITE_DELEGATES_FLEX_ALLOWLISTED_FLEX_OPS_INTERNAL_H_
tensorflow/compiler/mlir/lite/delegates/flex/build_def.bzl
+190
View File
@@ -0,0 +1,190 @@
"""Generate custom flex delegate library."""
load("@rules_cc//cc:cc_library.bzl", "cc_library")
load(
"//tensorflow:tensorflow.bzl",
"clean_dep",
"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/compiler/mlir/lite:special_rules.bzl", "flex_portable_tensorflow_deps")
# LINT.IfChange
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:portable_op_registrations_and_gradients"),
clean_dep("//tensorflow/core/kernels:portable_core_ops"),
clean_dep("//tensorflow/core/kernels:portable_extended_ops"),
]) + [CUSTOM_KERNEL_HEADER.header],
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(),
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:android"): [
portable_tensorflow_lib,
],
clean_dep("//tensorflow:ios"): [
portable_tensorflow_lib,
],
clean_dep("//tensorflow:chromiumos"): [
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,
)
# LINT.ThenChange(//tensorflow/lite/delegates/flex/build_def.bzl)
tensorflow/compiler/mlir/lite/experimental/common/BUILD
+22
View File
@@ -0,0 +1,22 @@
load("@rules_cc//cc:cc_library.bzl", "cc_library")
load("//tensorflow:tensorflow.default.bzl", "get_compatible_with_portable")
# copybara:uncomment package(default_applicable_licenses = ["//tensorflow:LICENSE"])
cc_library(
name = "outline_operations",
srcs = ["outline_operations.cc"],
hdrs = ["outline_operations.h"],
compatible_with = get_compatible_with_portable(),
visibility = ["//visibility:public"],
deps = [
"//tensorflow/compiler/mlir/lite:tensorflow_lite",
"//tensorflow/compiler/mlir/tensorflow:cluster_util",
"@com_google_absl//absl/strings",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Pass",
"@llvm-project//mlir:Support",
],
)
tensorflow/compiler/mlir/lite/experimental/common/outline_operations.cc
+212
View File
@@ -0,0 +1,212 @@
/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/compiler/mlir/lite/experimental/common/outline_operations.h"
#include <cassert>
#include <string>
#include "absl/strings/str_cat.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SetVector.h"
#include "llvm/ADT/SmallVector.h"
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/Block.h" // from @llvm-project
#include "mlir/IR/Builders.h" // from @llvm-project
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/IR/Matchers.h" // from @llvm-project
#include "mlir/IR/OpDefinition.h" // from @llvm-project
#include "mlir/IR/Operation.h" // from @llvm-project
#include "mlir/IR/Value.h" // from @llvm-project
#include "mlir/IR/Visitors.h" // from @llvm-project
#include "mlir/Support/LLVM.h" // from @llvm-project
#include "tensorflow/compiler/mlir/lite/ir/tfl_ops.h"
#include "tensorflow/compiler/mlir/lite/utils/utils.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/cluster_util.h"
namespace mlir {
namespace TFL {
namespace common {
bool IsConstantOrNone(Operation* op) {
return (op->getNumResults() == 1 &&
mlir::isa<NoneType>(op->getResult(0).getType())) ||
matchPattern(op, m_Constant()) || isa<QConstOp>(op);
}
// Pre-order traverse, adding results and BlockArgs to `been_defined` and
// collecting operands not contained within `been_defined`. If we encounter an
// operand that references a Value that has been defined (and added to
// `been_defined`) it is garuanteed that the Value definition is not contained
// in descedant node of reference, and given that the input DAG is valid, the
// definition is self-contained within `op` so it is not depended upon.
// Otherwise, the operand must have been defined somewhere above the Subgraph,
// so union with other operand dependencies.
llvm::SmallVector<Value> AccumulateOperandsDefinedAbove(
const llvm::SetVector<Operation*>& partition_ops) {
// Assuming that all are topologically sorted.
llvm::SetVector<Value> been_defined;
llvm::SetVector<Value> results;
auto update_from_op = [&](Operation* op) {
been_defined.insert(op->getResults().begin(), op->getResults().end());
for (Value input : op->getOperands()) {
if (been_defined.contains(input)) {
continue;
}
results.insert(input);
}
};
for (Operation* op : partition_ops) {
update_from_op(op);
op->walk<WalkOrder::PreOrder>([&](Block* nested_block) {
been_defined.insert(nested_block->getArguments().begin(),
nested_block->getArguments().end());
for (Operation& op : nested_block->getOperations()) update_from_op(&op);
});
}
return SmallVector<Value>(results.getArrayRef());
}
llvm::SmallVector<Value> AccumulateResultsDefinedWithin(
const llvm::SetVector<Operation*>& partition_ops) {
llvm::SmallVector<Value> values_for_results;
for (Operation* op : partition_ops) {
if (IsConstantOrNone(op)) {
continue;
}
for (Value output : op->getResults()) {
bool output_consumed_outside_subgraph = false;
for (Operation* consumer : output.getUsers()) {
if (llvm::all_of(partition_ops, [&](Operation* op) {
return !op->isAncestor(consumer);
})) {
output_consumed_outside_subgraph = true;
}
}
if (output_consumed_outside_subgraph) {
values_for_results.push_back(output);
}
}
}
return values_for_results;
}
// Compute signature for raised func from arugments and outputs of
// Operation partition.
llvm::SmallVector<Type> TypesFromValues(
const llvm::SmallVector<Value>& values) {
llvm::SmallVector<Type> types;
for (auto value : values) {
types.push_back(value.getType());
}
return types;
}
func::FuncOp BuildFuncOp(const Subgraph& subgraph, OpBuilder& builder,
ModuleOp& module, OpsAdded& ops_added) {
// The parameters of the new MLIR function are taken to be the union
// of all operands referenced by Operations within the subraph.
// Likewise the results of the function are any Value(s) that are defined
// within the subgraph and are referenced outside the subgraph.
llvm::SmallVector<Type> input_types =
TypesFromValues(subgraph.FuncArguments());
llvm::SmallVector<Type> return_types =
TypesFromValues(subgraph.FuncOutputs());
FunctionType function_type =
builder.getFunctionType(input_types, return_types);
std::string function_name = absl::StrCat("func_", subgraph.subgraph_id_);
func::FuncOp new_func = func::FuncOp::create(builder.getUnknownLoc(),
function_name, function_type);
new_func.setVisibility(func::FuncOp::Visibility::Private);
new_func.addEntryBlock();
// To form the body of the new function we need to clone each
// Operation along with its respective operands and result Values(s).
// The semantic of `Operation::clone` is copying given entity *into* this
// entity. The new FuncOp body is populated by cloning partitioned ops into
// it. Cloning Operation(s) will create cloned Value(s) for the results of a
// cloned op, but it needs a reference to the new operand Value(s) which are
// the result of the cloned ops. The approach is to traverse the subgraph in
// order, accumulating clones of defined Values into a `IRMapping`
// and pass that map to calls to clone ops.
OpBuilder function_builder(new_func.getBody());
// Preferred data structure for mapping MLIR values.
IRMapping values_in_scope;
// Function arguments can appear as operands, so they clone should
// be aware of them.
assert(subgraph.FuncArguments().size() == new_func.getNumArguments());
for (int i = 0; i < subgraph.FuncArguments().size(); ++i) {
Value original_value = subgraph.FuncArguments()[i];
Value new_func_arg = new_func.getArgument(i);
values_in_scope.map(original_value, new_func_arg);
}
for (Operation* op : subgraph.partition_ops_) {
function_builder.clone(*op, values_in_scope);
}
SmallVector<Value> return_operands;
for (Value result : subgraph.FuncOutputs()) {
Value cloned_output = values_in_scope.lookup(result);
return_operands.push_back(cloned_output);
}
mlir::func::ReturnOp::create(function_builder, new_func.getLoc(),
return_operands);
ops_added.func_op = new_func;
module.push_back(new_func);
return new_func;
}
void ExtractSubgraphToFunc(const Subgraph& subgraph, OpBuilder& builder,
ModuleOp& module, OpsAdded& ops_added) {
func::FuncOp func = BuildFuncOp(subgraph, builder, module, ops_added);
// We just use the location of the last ops in the subgraph as the location
// for the call_op.
Operation* last_output = subgraph.partition_ops_.back();
builder.setInsertionPoint(last_output);
auto call_op = func::CallOp::create(builder, last_output->getLoc(), func,
subgraph.FuncArguments());
ops_added.call_op = call_op;
// FuncOutputs refer to the original `Values` in input module which are now
// invalid after pulling out the defining ops. The values in
// `call_ops.getResult` refer to the clones of original `Values` which are now
// returned by the new `FuncOp`. We can replace each in `FuncOutputs` with
// clone in `call_op` to fix up.
for (int i = 0; i < subgraph.FuncOutputs().size(); ++i) {
Value output = subgraph.FuncOutputs()[i];
output.replaceAllUsesWith(call_op.getResult(i));
}
// Clear the subgraph.
// Those ops should be removed.
for (auto* op : subgraph.partition_ops_) {
if (IsConstantOrNone(op)) {
continue;
}
op->dropAllDefinedValueUses();
op->dropAllReferences();
op->erase();
}
// Ensure that users of the call op's results appear after the launch op in
// order to preserve the dominance property.
TF::ReorderOpResultUses(call_op);
}
} // namespace common
} // namespace TFL
} // namespace mlir
tensorflow/compiler/mlir/lite/experimental/common/outline_operations.h
+133
View File
@@ -0,0 +1,133 @@
/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_COMMON_OUTLINE_OPERATIONS_H_
#define TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_COMMON_OUTLINE_OPERATIONS_H_
#include <memory>
#include <string>
#include <utility>
#include "llvm/ADT/SetVector.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/ADT/StringRef.h"
#include "llvm/Support/Casting.h"
#include "llvm/Support/raw_os_ostream.h"
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/Attributes.h" // from @llvm-project
#include "mlir/IR/Block.h" // from @llvm-project
#include "mlir/IR/Builders.h" // from @llvm-project
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "mlir/IR/OpDefinition.h" // from @llvm-project
#include "mlir/IR/Operation.h" // from @llvm-project
#include "mlir/IR/UseDefLists.h" // from @llvm-project
#include "mlir/IR/Value.h" // from @llvm-project
#include "mlir/IR/ValueRange.h" // from @llvm-project
#include "mlir/IR/Visitors.h" // from @llvm-project
#include "mlir/Pass/Pass.h" // from @llvm-project
#include "mlir/Pass/PassRegistry.h" // from @llvm-project
#include "mlir/Support/LLVM.h" // from @llvm-project
#include "tensorflow/compiler/mlir/lite/ir/tfl_ops.h"
#include "tensorflow/compiler/mlir/lite/utils/utils.h"
namespace mlir {
namespace TFL {
namespace common {
// Returns true if the `op` is a constant-like op or produces none type.
bool IsConstantOrNone(Operation* op);
// Computes the list of Value(s) referenced by Subgraph Operations that are
// not defined within the Subgraph. Any such Value(s)
// are validly in-scope for the initial Operation. They must be either
// defined above the subgraph or appear as an argument to the containing func.
// These Value(s) are taken to be the arguments of the new raised func.
// An operand dependency is a Value referenced anywhere in an Op
// that is defined above the Op. All SSA Values are assigned/defined in a
// BlockArg or as a result of an Operation.
llvm::SmallVector<Value> AccumulateOperandsDefinedAbove(
const llvm::SetVector<Operation*>& partition_ops);
// Similar to `AccumulateOperandsDefinedAbove()`, computes the Value(s) that are
// defined within a Subgraph and referenced in a descendant Operation. These
// Values(s) are to be returned by the new raised function.
llvm::SmallVector<Value> AccumulateResultsDefinedWithin(
const llvm::SetVector<Operation*>& partition_ops);
// Represents a view of a set of mlir Operations that form a subgraph of the
// entire Module's DAG. `Subgraph` can be thought of as segment of sequential
// Operations within a func definition. Additional facts:
// 1. Subgraphs are restricted to a single Block. They do not span
// branching instructions. Thus the subgraph is a simple 1-degree path.
// 2. All Operations in a subgraph belong to the same block in a
// funtion body.
// 3. Function bodies are assumed to have only one block in some places.
class Subgraph {
// Set vector preserves insertion order, must insert Ops in topological order.
public:
const llvm::SetVector<Operation*> partition_ops_;
// Subgraphs are given a unique incremented integer id based on when
// they were encountered in this pass.
const int subgraph_id_;
const llvm::StringRef dialect_namespace_;
Subgraph(const llvm::SetVector<Operation*> partition_ops, int num_subgraphs)
: partition_ops_(partition_ops),
subgraph_id_(num_subgraphs),
func_arguments_(AccumulateOperandsDefinedAbove(partition_ops)),
func_outputs_(AccumulateResultsDefinedWithin(partition_ops)) {}
const llvm::SmallVector<Value>& FuncArguments() const {
// `Value`s in MLIR library are implemented as having "value semantics"
// see "llvm/llvm-project/mlir/include/mlir/IR/Value.h" so copying is fine.
return func_arguments_;
}
const llvm::SmallVector<Value>& FuncOutputs() const { return func_outputs_; }
private:
// Compute once at construction and save as field.
const llvm::SmallVector<Value> func_arguments_;
const llvm::SmallVector<Value> func_outputs_;
};
// Helper data structure for output parameters to `ExtractSubgraphToFunc`.
// `ExtractSubgraphToFunc` adds exactly two "new" `Operations`, a FuncOp and
// a CallOp. Pass these back to the caller for setting more specific attributes
// after graph mutation has taken place.
struct OpsAdded {
mlir::func::FuncOp func_op;
mlir::func::CallOp call_op;
};
// Given a `Subgraph` containing a sequence of adjacent `Operations` from
// the `module`, raise these `Operations` (and any ops contained nested within)
// to the body of a new seperate root level function. Replace in their current
// location with a `CallOp` which invokes said `FuncOp`. The inputs to
// this new functions are taken to be the `Values` that appear as operands
// to ops in the subgraph, which are not self-contained within the subgraph.
// The outputs of this function are taken to be the results of ops in the
// subgraph which are referenced as operands outside of the subgraph.
// Also refer to documention of `AccumulateOperandsDefinedAbove` &
// `AccumulateResultsDefinedWithin`.
void ExtractSubgraphToFunc(const Subgraph& subgraph, OpBuilder& builder,
ModuleOp& module, OpsAdded& ops_added);
} // namespace common
} // namespace TFL
} // namespace mlir
#endif // TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_COMMON_OUTLINE_OPERATIONS_H_
tensorflow/compiler/mlir/lite/experimental/remat/BUILD
+44
View File
@@ -0,0 +1,44 @@
load("@rules_cc//cc:cc_library.bzl", "cc_library")
load("//tensorflow:tensorflow.bzl", "tf_cc_test")
load("//tensorflow:tensorflow.default.bzl", "get_compatible_with_portable")
# copybara:uncomment package(default_applicable_licenses = ["//tensorflow:LICENSE"])
cc_library(
name = "rematerializer",
srcs = ["rematerializer.cc"],
hdrs = ["rematerializer.h"],
deps = [
],
)
tf_cc_test(
name = "rematerializer_test",
size = "small",
srcs = ["rematerializer_test.cc"],
deps = [
":rematerializer",
"@com_google_googletest//:gtest_main",
],
)
cc_library(
name = "metadata_util",
srcs = ["metadata_util.cc"],
hdrs = ["metadata_util.h"],
compatible_with = get_compatible_with_portable(),
visibility = ["//visibility:public"],
deps = [
"//tensorflow/compiler/mlir/lite:control_edges",
],
)
tf_cc_test(
name = "metadata_util_test",
size = "small",
srcs = ["metadata_util_test.cc"],
deps = [
":metadata_util",
"@com_google_googletest//:gtest_main",
],
)
tensorflow/compiler/mlir/lite/experimental/remat/metadata_util.cc
+124
View File
@@ -0,0 +1,124 @@
/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/compiler/mlir/lite/experimental/remat/metadata_util.h"
#include <cstddef>
#include <cstdint>
#include <string>
#include <utility>
#include <vector>
namespace {
// We serialize unsigneds as protobuf varints, i.e., in chunks of 7 bits each.
constexpr int kMod = (1 << 7);
void Serialize(std::string* out, uint32_t value) {
for (; value >= kMod; value /= kMod) {
out->push_back(value % kMod + kMod);
}
out->push_back(value);
}
bool Parse(const char** data, size_t* size, uint32_t* out) {
*out = 0;
uint32_t mul = 1;
for (bool done = false; !done;
mul *= kMod, done = !(**data & kMod), ++*data, --*size) {
if (*size == 0) {
return false;
}
*out += static_cast<unsigned char>(**data) % kMod * mul;
}
return true;
}
// Signed ints are zigzag-encoded as unsigned varints, [..., -2, -1, 0, 1, 2,
// ...] ->
// [..., 3, 1, 0, 2, 4, ...].
void Serialize(std::string* out, int32_t value) {
Serialize(out, static_cast<uint32_t>(
value < 0 ? static_cast<uint32_t>(-(value + 1)) * 2 + 1
: static_cast<uint32_t>(value) * 2));
}
bool Parse(const char** data, size_t* size, int32_t* out) {
uint32_t value = 0;
if (!Parse(data, size, &value)) {
return false;
}
const int32_t magnitude = value / 2;
*out = (value % 2) ? (-magnitude - 1) : magnitude;
return true;
}
// Pairs are serialized as the concatenation of their elements' serialization.
template <class First, class Second>
void Serialize(std::string* out, const std::pair<First, Second>& in) {
Serialize(out, in.first);
Serialize(out, in.second);
}
template <class First, class Second>
bool Parse(const char** data, size_t* size, std::pair<First, Second>* out) {
return Parse(data, size, &(out->first)) && Parse(data, size, &(out->second));
}
// Vectors are serialized as the concetation of the serialization of their size
// and the the serializations of their elements.
template <class Value>
void Serialize(std::string* out, const std::vector<Value>& in) {
Serialize(out, static_cast<uint32_t>(in.size()));
for (const auto& val : in) {
Serialize(out, val);
}
}
template <class T>
bool Parse(const char** data, size_t* size, std::vector<T>* out) {
uint32_t num_elems = 0;
if (!Parse(data, size, &num_elems)) {
return false;
}
out->assign(num_elems, T{});
for (auto& elem : *out) {
if (!Parse(data, size, &elem)) {
return false;
}
}
return true;
}
} // namespace
namespace tflite {
std::string SerializeModelControlDependencies(
const ModelControlDependencies& in) {
std::string out;
Serialize(&out, kModelControlDependenciesMetadataVersion);
Serialize(&out, in);
return out;
}
bool ParseModelControlDependencies(const char* data, size_t size,
ModelControlDependencies* out) {
out->clear();
uint32_t version = 0;
return Parse(&data, &size, &version) &&
(version == kModelControlDependenciesMetadataVersion) &&
Parse(&data, &size, out) && (size == 0);
}
} // namespace tflite
tensorflow/compiler/mlir/lite/experimental/remat/metadata_util.h
+62
View File
@@ -0,0 +1,62 @@
/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
/// \file
///
/// Functions for serializiation/deserialization of control dependency
/// information to/from model metadata.
///
#ifndef TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_REMAT_METADATA_UTIL_H_
#define TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_REMAT_METADATA_UTIL_H_
#include <cstddef>
#include <cstdint>
#include <string>
#include <vector>
#include "tensorflow/compiler/mlir/lite/utils/control_edges.h"
namespace tflite {
/// Control dependencies for the model is the collection of control dependencies
/// for its subgraphs.
using ModelControlDependencies = std::vector<ControlEdges>;
/// Serializes `in` into the returned string. The result is parseable with
/// ParseModelControlDependencies.
std::string SerializeModelControlDependencies(
const ModelControlDependencies& in);
/// Deserializes `*out` from a character buffer of size `size` at `data`.
/// Returns true iff successful. `*out` needn't be empty before invocation.
/// When returning false, `*out`'s state is undefined.
bool ParseModelControlDependencies(const char* data, size_t size,
ModelControlDependencies* out);
/// The key under which to store the serialized control dependencies in the
/// model's metadata.
constexpr char kModelControlDependenciesMetadataKey[] =
"model_control_dependencies";
/// To allow future changes to the format, serialized control dependency data
/// will contain a version; this constant is the version that will be used for
/// serialization. For deserialization, past versions should remain parseable.
constexpr uint32_t kModelControlDependenciesMetadataVersion = 1;
inline constexpr char kModelUseStablehloTensorKey[] = "keep_stablehlo_constant";
} // namespace tflite
#endif // TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_REMAT_METADATA_UTIL_H_
tensorflow/compiler/mlir/lite/experimental/remat/metadata_util_test.cc
+55
View File
@@ -0,0 +1,55 @@
/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/compiler/mlir/lite/experimental/remat/metadata_util.h"
#include <cstdint>
#include <limits>
#include <string>
#include <gmock/gmock.h>
#include <gtest/gtest.h>
namespace tflite {
namespace {
class MetadataSerializerTest : public ::testing::Test {
protected:
static constexpr auto kHuge = std::numeric_limits<int32_t>::max();
static constexpr auto kTiny = std::numeric_limits<int32_t>::min();
std::string RoundTrip(const ModelControlDependencies &in) const {
ModelControlDependencies out = {{{-1, -1}}};
const std::string serialized =
tflite::SerializeModelControlDependencies(in);
return tflite::ParseModelControlDependencies(serialized.data(),
serialized.size(), &out)
? (out == in) ? "ok" : "mismatch"
: "malformed";
}
};
TEST_F(MetadataSerializerTest, nothing) { EXPECT_THAT(RoundTrip({}), "ok"); }
TEST_F(MetadataSerializerTest, something) {
EXPECT_THAT(
RoundTrip({{{1, 2}, {2, 3}, {4, 5}},
{},
{{kHuge, kTiny}, {kTiny, kHuge}, {kHuge - 1, kTiny + 1}},
{{1, 0}}}),
"ok");
}
} // namespace
} // namespace tflite
tensorflow/compiler/mlir/lite/experimental/remat/rematerializer.cc
+381
View File
@@ -0,0 +1,381 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/compiler/mlir/lite/experimental/remat/rematerializer.h"
#include <algorithm>
#include <map>
#include <tuple>
#include <utility>
#include <vector>
namespace mlir {
namespace TFL {
namespace {
// Helper functions for sorted + deduped int vectors.
std::tuple<std::vector<int>::iterator, bool> Find(const int item,
std::vector<int>& items) {
const auto iter = std::lower_bound(items.begin(), items.end(), item);
return std::make_tuple(iter, iter != items.end() && *iter == item);
}
void Insert(const int item, std::vector<int>& items) {
const auto [iter, found] = Find(item, items);
if (!found) items.insert(iter, item);
}
void Erase(const int item, std::vector<int>& items) {
const auto [iter, found] = Find(item, items);
if (found) items.erase(iter);
}
} // namespace
int Rematerializer::AddOperation(const bool is_stateful) {
operations_.emplace_back();
operations_.back().is_stateful = is_stateful;
return operations_.size() - 1;
}
int Rematerializer::AddTensor(const SizeT size) {
tensors_.emplace_back();
tensors_.back().size = size;
return tensors_.size() - 1;
}
void Rematerializer::DelUse(const int ioperation, const int itensor) {
auto& tensor = tensors_[itensor];
auto& operation = operations_[ioperation];
const auto& size = tensor.size;
// Was the dependence to be deleted the first/last (or both) use of this
// tensor?
const bool was_first_use =
(!tensor.operations.empty() && ioperation == tensor.first_use());
const bool was_last_use =
(!tensor.operations.empty() && ioperation == tensor.last_use());
Erase(ioperation, tensor.operations);
Erase(itensor, operation.tensors);
if (was_first_use) {
operation.alloc -= size;
if (!was_last_use) {
operations_[tensor.first_use()].alloc += size;
}
}
if (was_last_use) {
operation.dealloc -= size;
if (!was_first_use) {
operations_[tensor.last_use()].dealloc += size;
}
}
}
void Rematerializer::AddUse(const int ioperation, const int itensor) {
auto& tensor = tensors_[itensor];
auto& operation = operations_[ioperation];
const auto& size = tensor.size;
const bool will_be_first_use =
tensor.operations.empty() || ioperation < tensor.first_use();
const bool will_be_last_use =
tensor.operations.empty() || ioperation > tensor.last_use();
if (will_be_first_use) {
operation.alloc += size;
if (!will_be_last_use) {
operations_[tensor.first_use()].alloc -= size;
}
}
if (will_be_last_use) {
operation.dealloc += size;
if (!will_be_first_use) {
operations_[tensor.last_use()].dealloc -= size;
}
}
Insert(ioperation, tensor.operations);
Insert(itensor, operation.tensors);
}
Rematerializer::SizeT Rematerializer::MaxSavings(const int begin, const int end,
const int peak_loc) const {
SizeT max_savings = 0;
// We're looking at the outputs of all operators in [`begin`, `end`). If an
// output is alive after `peak_loc`, rematerializing the operator anywhere
// after `peak_loc` would lower the memory profile at `peak_loc`.
for (int ioperation = begin; ioperation != end; ++ioperation) {
for (const int itensor : operations_[ioperation].tensors) {
if (const Tensor& tensor = tensors_[itensor];
tensor.first_use() == ioperation /* output */ &&
tensor.last_use() > peak_loc /* used later */) {
max_savings += tensor.size;
}
}
}
return max_savings;
}
std::tuple<Rematerializer::SizeT, Rematerializer::RematSpec>
Rematerializer::FindBestRemat(const SizeT min_savings, const int begin_len,
const int end_len) const {
const auto peak = GetPeakMemory();
SizeT best_peak_mem = peak.size;
RematSpec best_remat = {};
for (int len = begin_len; len < end_len; ++len) {
std::vector<std::tuple<SizeT, int, int>> pre_screen;
for (int begin = 0, end = begin + len; end <= peak.op_index;
++begin, ++end) {
if (!std::any_of(operations_.begin() + begin, operations_.begin() + end,
[](const Operation& s) { return s.is_stateful; })) {
if (const auto max_savings = MaxSavings(begin, end, peak.op_index);
max_savings >= min_savings) {
pre_screen.emplace_back(max_savings, begin, end);
}
}
}
std::sort(pre_screen.begin(), pre_screen.end());
for (; !pre_screen.empty(); pre_screen.pop_back()) {
const auto& [max_savings, begin, end] = pre_screen.back();
const auto insert_before = FindBestRematPoint(begin, end, peak.op_index);
if (insert_before == operations_.size()) {
continue;
}
const RematSpec this_remat = {begin, end, insert_before};
if (const auto new_peak = GetPeakMemory(this_remat);
new_peak.size < best_peak_mem &&
peak.size >= new_peak.size + min_savings) {
best_peak_mem = new_peak.size;
best_remat = this_remat;
}
// If the actual savings achieved is bigger than the maximal savings that
// can be possibly achieved, leave early.
if (peak.size >= max_savings + best_peak_mem) {
break;
}
}
// We already found one savings for this length.
if (peak.size >= min_savings + best_peak_mem) {
break;
}
}
return std::make_tuple(best_peak_mem, best_remat);
}
std::vector<Rematerializer::MemSpec> Rematerializer::GetDeltas(
const RematSpec& remat) const {
// We're cloning operations [remat.begin, remat.end) at position
// remat.insert. We store changes to the alloc/dealloc sizes due to the
// insertion in a vector `delta`: A change `c_alloc` of `operations_[i].alloc`
// as `delta[i] += c_alloc`, and a change `c_dealloc` of
// `operations_[i].dealloc` as `delta[i+1] -= c_dealloc`.
std::vector<MemSpec> deltas;
if (remat.begin == remat.end) {
return deltas;
}
// Helper lambda: converts an operation index in the original range to its
// equivalent in the cloned range.
const auto source_to_target = [&](int i) {
return i + (remat.insert - remat.begin);
};
// Helper struct: bundles first and last use of a tensor within
// a contiguous range of operations.
struct TensorUse {
int first_use;
int last_use;
};
// For all tensors in the operation range to be cloned, store their first and
// last use within that range. This will be needed below to decide whether a
// tensor's life will be extended, or if it is to be rematerialized.
std::map<int, TensorUse> source_uses;
for (int ioperation = remat.begin; ioperation < remat.end; ++ioperation) {
const auto& operation = operations_[ioperation];
for (const int itensor : operation.tensors) {
// insertion will only happen for the first use.
const auto [iter, inserted] = source_uses.emplace(
itensor,
TensorUse{/*first_use=*/ioperation, /*last_use=*/ioperation});
if (!inserted) {
// Otherwise, update last_use.
iter->second.last_use = ioperation;
}
}
}
deltas.reserve(2 * source_uses.size());
for (const auto& [itensor, source] : source_uses) {
auto& tensor = tensors_[itensor];
const TensorUse global = {tensor.first_use(), tensor.last_use()};
auto add_alloc = [&](int pos) { deltas.emplace_back(pos, tensor.size); };
auto add_dealloc = [&](int pos) {
deltas.emplace_back(pos + 1, -tensor.size);
};
auto del_dealloc = [&](int pos) {
deltas.emplace_back(pos + 1, tensor.size);
};
if (global.first_use < remat.begin) {
// The tensor is created before the source range.
if (global.last_use < remat.insert) {
// It currently gets deallocated before the newly inserted range, so we
// need to extend its lifetime: It will now be deallocated at its last
// use in the inserted range.
del_dealloc(global.last_use);
add_dealloc(source_to_target(source.last_use));
}
} else {
// Tensor is created in the source range. It will be rematerialized in the
// cloned range.
add_alloc(source_to_target(source.first_use));
if (global.last_use < remat.insert) {
// The last use of the original tensor is before the target range, so
// the lifetime of the rematerialized tensor ends with the target range.
add_dealloc(source_to_target(source.last_use));
} else {
// There are uses of the original tensor after the cloned range. They
// can be replaced with uses of the rematerialized tensor, and the
// original tensor can be deallocated after its last use before the
// rematerialization.
add_dealloc(*std::partition_point(
tensor.operations.rbegin(), tensor.operations.rend(),
[&](int i) { return i >= remat.insert; }));
}
}
}
std::sort(deltas.begin(), deltas.end(), ByOpIndex);
return deltas;
}
Rematerializer::MemProfile Rematerializer::GetMemProfile(
const RematSpec& remat) const {
const auto num_inserted = remat.end - remat.begin;
std::vector<SizeT> profile(operations_.size() + num_inserted);
MapMem([&](const MemSpec& m) { profile[m.op_index] = m.size; }, remat);
return profile;
}
Rematerializer::MemSpec Rematerializer::GetPeakMemory(
const RematSpec& remat) const {
MemSpec peak;
MapMem([&](const MemSpec& m) { peak = std::max(m, peak, BySize); }, remat);
return peak;
}
int Rematerializer::FindBestRematPoint(const int begin, const int end,
const int peak_loc) const {
// All tensors necessarily have their last use before or at the last operator
// (which is the yield operation of a function), so an unchanged best ==
// operations_.size() value means that there is no valid configuration.
int best = operations_.size();
// All tensors whose first use is in [begin, end) and that are not destroyed
// until after peak_loc.
for (int ioperation = begin; ioperation < end; ++ioperation) {
for (const int itensor : operations_[ioperation].tensors) {
if (const auto& tensor = tensors_[itensor];
tensor.first_use() >= begin && tensor.first_use() < end &&
tensor.last_use() > peak_loc) {
for (const int ioperation : tensor.operations) {
// We return the earliest dependence on any output tensor.
if (ioperation > peak_loc && ioperation < best) {
best = ioperation;
break;
}
}
}
}
}
return best;
}
void Rematerializer::Remat(const RematSpec& remat) {
const int num_inserted = remat.end - remat.begin;
// Rewrite all operation indices.
for (auto& tensor : tensors_) {
std::for_each(std::lower_bound(tensor.operations.begin(),
tensor.operations.end(), remat.insert),
tensor.operations.end(),
[&](int& iop) { iop += num_inserted; });
}
operations_.insert(operations_.begin() + remat.insert, num_inserted, {});
// Copy all tensors dependencies of the old region to the new region.
// For any output tensor, a new tensor will be created.
std::vector<std::pair<int, int>> new_tensors;
for (int iop_old = remat.begin, iop_new = remat.insert; iop_old < remat.end;
++iop_old, ++iop_new) {
for (const auto itensor : operations_[iop_old].tensors) {
if (tensors_[itensor].first_use() == iop_old) {
new_tensors.emplace_back(itensor, AddTensor(tensors_[itensor].size));
}
AddUse(iop_new, itensor);
}
}
std::sort(new_tensors.begin(), new_tensors.end());
// Let all inserted + downstream operations refer to the new tensors.
for (int iop = remat.insert; iop < operations_.size(); ++iop) {
// We have to copy, otherwise we're modifying the set during the iteration.
for (const int old_tensor : std::vector<int>(operations_[iop].tensors)) {
const auto new_tensor =
std::lower_bound(new_tensors.begin(), new_tensors.end(),
std::make_pair(old_tensor, 0));
if (new_tensor != new_tensors.end() && new_tensor->first == old_tensor) {
DelUse(iop, old_tensor);
AddUse(iop, new_tensor->second);
}
}
}
}
void Rematerializer::RunGreedyAlgorithm(const int max_cost,
const int max_block_length,
const SizeT min_savings) {
const bool unlimited_cost = (max_cost < 0);
for (int min_block_length = 1, cost = 0;
min_block_length <= max_block_length &&
(unlimited_cost || cost <= max_cost);
min_block_length *= 2) {
while (unlimited_cost || cost <= max_cost) {
const auto [peak, remat] = FindBestRemat(
/*min_savings*/ min_savings,
/*begin_len=*/min_block_length,
/*end_len=*/
std::min(1 + (unlimited_cost
? max_block_length
: std::min(max_block_length, max_cost - cost)),
2 * min_block_length));
if (remat.begin == remat.end) break;
Remat(remat);
ApplyRemat(remat);
cost += (remat.end - remat.begin);
}
}
}
} // namespace TFL
} // namespace mlir
tensorflow/compiler/mlir/lite/experimental/remat/rematerializer.h
+262
View File
@@ -0,0 +1,262 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_REMAT_REMATERIALIZER_H_
#define TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_REMAT_REMATERIALIZER_H_
// This file declares the Rematerializer class, which is used by an MLIR-based
// set of transformations for TFLite IR that lower memory usage by redoing
// operations with small inputs and large outputs instead of keeping the result
// in memory. This class allows us to compactly and efficiently represent the
// (idealized) memory profile of a TFLite graph and simulate the effect of
// re-inserting operations on that memory profile.
#include <algorithm>
#include <cinttypes>
#include <cstdint>
#include <tuple>
#include <vector>
namespace mlir {
namespace TFL {
// A class that
// (1) Encodes in concise form the memory requirements of a computational graph
// (2) Allows for the fast simulation of changes to the peak memory requirement
// under rematerialization of intermediate results in the graph
// (3) Implements a greedy algorithm for finding rematerializations of
// intermediate results in that graph to lower peak memory requirements.
class Rematerializer {
public:
Rematerializer() = default;
virtual ~Rematerializer() = default;
// The type used for memory sizes (in bytes) and differences thereof.
using SizeT = int64_t;
// The memory profile: The i-th element gives the amount of memory
// that is needed when performing the i-th operation. This is the
// sum of the sizes of
//
// (1) input tensors of that operation,
// (2) output tensors of that operation,
// (3) output tensors of preceding operations that are input tensors
// of subsequent operations.
using MemProfile = std::vector<SizeT>;
// Used for specifying memory consumption at a certain operation in the
// computational graph.
struct MemSpec {
int op_index; // The index of the operation
SizeT size; // The amount of memory needed in order to execute this
// operation, i.e., the sum of input and output sizes and the
// sizes of outputs of previous operations that are needed as
// inputs of subsequent operations.
explicit MemSpec(int op_index = 0, SizeT size = 0)
: op_index(op_index), size(size) {}
};
static bool BySize(const MemSpec& a, const MemSpec& b) {
return std::tie(a.size, a.op_index) < std::tie(b.size, b.op_index);
}
static bool ByOpIndex(const MemSpec& a, const MemSpec& b) {
return std::tie(a.op_index, a.size) < std::tie(b.op_index, b.size);
}
// Specifies an elementary rematerialization operation: The operations in
// operations [`begin`, `end`) will be rescheduled before operation `insert`.
// A valid `RematSpec` requires begin <= end <= insert <= number of
// operations. Note that (1) `end` is exclusive -- begin == end signifies a
// trivial RematSpec (no operation will be rescheduled), (2) the
// zero-initialized RematSpec {} is trivial and always valid.
struct RematSpec {
int begin;
int end;
int insert;
};
// Gives the peak memory location and size after inserting operations
// according to `remat` (but doesn't actually insert them.) Ties are broken
// towards later locations. `remat` must be valid (see above).
MemSpec GetPeakMemory(const RematSpec& remat = {}) const;
// Gives memory profile after inserting operations according to `remat` (but
// doesn't actually insert them). `remat` must be valid (see above).
MemProfile GetMemProfile(const RematSpec& remat = {}) const;
// Runs the greedy incremental block algorithm: Finds a sequence of
// rematerializations of block size up to max_block_length, each reducing peak
// memory by at least min_savings. If max_cost >= 0, at most max_cost
// operations will be re-inserted. For each rematerialization found,
// ApplyRemat is invoked (which can be used to apply the rematerialization to
// the higher- level representation, e.g., MLIR, flatbuffer, ...)
void RunGreedyAlgorithm(int max_cost, int max_block_length,
SizeT min_savings);
virtual void ApplyRemat(const RematSpec& remat) {}
protected:
// Rematerializes the outputs of the operations [`remat.begin`, `remat.end`)
// before operation remat.insert by copying that operation range before
// remat.insert and updating tensor references so that any operation that can
// will make use of the rematerialized outputs rather than the original ones.
// `remat` must be valid (see above).
void Remat(const RematSpec& remat);
// The protected methods below are to be used by derived classes to create the
// low-level (this class) representation from a high-level one.
// Creates a new tensor-like object that takes `size` bytes. Returns a
// contiguous increasing index for each new object, starting at 0.
int AddTensor(SizeT size);
// Creates an operation. If `is_stateful`, the operation (and any block of
// operations containing it) will never be considered for rematerialization.
// Returns a contiguous increasing index for each new object, starting at 0.
int AddOperation(bool is_stateful);
// The operator with index `ioperation` will be assumed to produce and/or
// consume the tensor with index `itensor`. NoOp if that's already the case.
// The arguments must be valid indices (i.e., obtained with
// `AddOperation`/`AddTensor`).
void AddUse(int ioperation, int itensor);
// Undoes an AddUse(ioperation, itensor). NoOp if there was no prior `AddUse`.
// The arguments must be valid indices (i.e., obtained with
// `AddOperation`/`AddTensor`).
void DelUse(int ioperation, int itensor);
private:
// Find the best remat operation that saves at least `min_savings` bytes for a
// block of operators with a length is between [`begin_len`, `end_len`).
// 'Best' means with the highest savings, ties are broken towards shorter
// blocks.
std::tuple<SizeT, RematSpec> FindBestRemat(SizeT min_savings, int begin_len,
int end_len) const;
// Optimization: Estimate (from above) the remat savings of instruction block
// [begin, end) after operation `peak_location`
SizeT MaxSavings(int begin, int end, int peak_loc) const;
// If I want to remat ops [begin, end) after the op at operation `peak_loc`,
// find the latest point at which to reinsert them (the op before which to
// insert.)
int FindBestRematPoint(int begin, int end, int peak_loc) const;
// The memory objects.
struct Tensor {
SizeT size; // The size of the object (in bytes.)
std::vector<int> operations; // The operations it is used in. This vector
// is kept sorted + unique.
// The operation that makes the first use of this tensor.
int first_use() const { return *operations.begin(); }
// The operation that makes the last use of this tensor.
int last_use() const { return *operations.rbegin(); }
};
// The operators.
struct Operation {
bool is_stateful = false; // Results of an Operation can be rematerialized
// only if `!is_stateful`. This probably should
// be replaced with a more-fine grained
// approach--for example, the results of a "read
// resource variable" operation can be
// rematerialized as long as this doesn't happen
// after the corresponding "write resource
// variable" operation.
std::vector<int> tensors; // The tensors that are used (input or output) by
// this operation. They needn't correspond to
// tensors in the TF graph -- we may add fake
// tensors to model memory consumed in addition
// to input and output tensors. This vector is
// kept sorted + unique.
SizeT alloc = 0; // The number of bytes that need to be allocated before
// this operation.
SizeT dealloc = 0; // The number of bytes that can be deallocated after
// this operation.
};
// Given the current state of `operations_` and `tensors_`, return a vector of
// corrections that transform the current memory profile into the one that we
// would get after applying `remat`.
//
// The memory profile of a sequence of operations is the partial sum of the
// sizes of the allocations that are necessary before an operation and the
// negative sizes of the deallocations that are possible after the previous
// operation.
//
// If we modify the operation sequence by cloning an operation range, that
// memory profile will change--cloning makes it necessary to extend the
// lifetime of some tensors, while other tensors can be deallocated early and
// rematerialized later.
//
// This method represents these changes in compact form: It returns a vector
// of (position of operation, delta) pairs in lexicographic order; one
// obtains the memory profile after `remat` by adding the deltas from any
// entries (i, delta) to the i-th entry of the partial sum.
//
// This allows us to efficiently compute the change to the peak of a memory
// profile due to cloning an operation range without having to actually clone
// that range and without having to build a profile vector.
//
// The returned vector has at most 2 entries for each tensor referenced in
// [remat.begin, remat.end). There may be multiple entries for a single
// operation position; operation positions refer to the sequence *after*
// cloning [`remat.begin`, `remat.end`) before `remat.insert`.
std::vector<MemSpec> GetDeltas(const RematSpec& remat) const;
// Helper template: Iterates through all `MemSpec`s (i.e., operation
// index/memory usage pairs) for the current graph in operation order and
// calls `mapper` on them. This is an optimization -- by instantiating with an
// appropriate `Mapper`, it allows us to e.g. compute the peak memory without
// having to instantiate an actual memory profile vector.
template <class Mapper>
void MapMem(const Mapper& mapper, const RematSpec& remat) const {
const auto deltas = GetDeltas(remat);
const auto len = (remat.end - remat.begin);
auto idelta = deltas.begin();
for (MemSpec m; m.op_index < operations_.size() + len; ++m.op_index) {
// Are we in the cloned portion of the new operation sequence?
// Then all alloc/dealloc information must come from deltas.
const bool patch =
(m.op_index >= remat.insert) && (m.op_index < remat.insert + len);
// Are we past the insertion portion of the new operation sequence?
// Then we need to convert indices back to the original sequence.
const int shift = (m.op_index >= remat.insert + len) ? len : 0;
m.size += patch ? 0 : operations_[m.op_index - shift].alloc;
// deltas is sorted by location; apply any corrections to the current
// operator.
for (; idelta != deltas.end() && idelta->op_index == m.op_index;
++idelta) {
m.size += idelta->size;
}
mapper(m);
m.size -= patch ? 0 : operations_[m.op_index - shift].dealloc;
}
}
std::vector<Operation> operations_;
std::vector<Tensor> tensors_;
};
} // namespace TFL
} // namespace mlir
#endif // TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_REMAT_REMATERIALIZER_H_
tensorflow/compiler/mlir/lite/experimental/remat/rematerializer_test.cc
+519
View File
@@ -0,0 +1,519 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/compiler/mlir/lite/experimental/remat/rematerializer.h"
#include <algorithm>
#include <array>
#include <cstdlib>
#include <initializer_list>
#include <random>
#include <vector>
#include <gmock/gmock.h>
#include <gtest/gtest.h>
namespace mlir {
namespace TFL {
namespace {
using ::testing::ElementsAre;
using ::testing::ElementsAreArray;
using ::testing::Eq;
using ::testing::FieldsAre;
using ::testing::StrictMock;
class RematTest : public ::testing::Test {
protected:
class TestableRematerializer : public Rematerializer {
public:
using Rematerializer::AddOperation;
using Rematerializer::AddTensor;
using Rematerializer::AddUse;
using Rematerializer::DelUse;
using Rematerializer::Remat;
};
TestableRematerializer r_;
};
TEST_F(RematTest, TensorUseSimple) {
for (int i = 0; i < 6; ++i) {
r_.AddOperation(/*is_stateful=*/false);
r_.AddTensor(/*size=*/1 << i);
}
r_.AddUse(/*ioperation=*/2, /*itensor=*/2);
EXPECT_THAT(r_.GetMemProfile(), ElementsAre(0, 0, 4, 0, 0, 0));
EXPECT_THAT(r_.GetPeakMemory(), FieldsAre(Eq(2), Eq(4)));
r_.AddUse(/*ioperation=*/2, /*itensor=*/2);
EXPECT_THAT(r_.GetMemProfile(), ElementsAre(0, 0, 4, 0, 0, 0));
EXPECT_THAT(r_.GetPeakMemory(), FieldsAre(Eq(2), Eq(4)));
r_.AddUse(/*ioperation=*/4, /*itensor=*/2);
EXPECT_THAT(r_.GetMemProfile(), ElementsAre(0, 0, 4, 4, 4, 0));
EXPECT_THAT(r_.GetPeakMemory(), FieldsAre(Eq(4), Eq(4)));
r_.DelUse(/*ioperation=*/2, /*itensor=*/2);
EXPECT_THAT(r_.GetMemProfile(), ElementsAre(0, 0, 0, 0, 4, 0));
EXPECT_THAT(r_.GetPeakMemory(), FieldsAre(Eq(4), Eq(4)));
r_.DelUse(/*ioperation=*/2, /*itensor=*/2);
EXPECT_THAT(r_.GetMemProfile(), ElementsAre(0, 0, 0, 0, 4, 0));
EXPECT_THAT(r_.GetPeakMemory(), FieldsAre(Eq(4), Eq(4)));
r_.DelUse(/*ioperation=*/4, /*itensor=*/2);
EXPECT_THAT(r_.GetMemProfile(), ElementsAre(0, 0, 0, 0, 0, 0));
EXPECT_THAT(r_.GetPeakMemory(), FieldsAre(Eq(5), Eq(0)));
}
TEST_F(RematTest, TensorUseMany) {
constexpr int n = 6;
for (int i = 0; i < n; ++i) {
r_.AddUse(/*ioperation=*/r_.AddOperation(/*is_stateful=*/false),
/*itensor=*/r_.AddTensor(1 << (n - i - 1)));
}
for (int i = 0; i < n; ++i) {
r_.AddUse(/*ioperation=*/r_.AddOperation(/*is_stateful=*/false),
/*itensor=*/n - 1 - i);
}
EXPECT_THAT(r_.GetMemProfile(), ElementsAreArray({32, 48, 56, 60, 62, 63, 63,
62, 60, 56, 48, 32}));
EXPECT_THAT(r_.GetPeakMemory(), FieldsAre(Eq(6), Eq(63)));
}
TEST_F(RematTest, PeakTiesAreBrokenInFavorOfLaterOperations) {
r_.AddUse(/*ioperation=*/r_.AddOperation(/*is_stateful=*/false),
/*itensor=*/r_.AddTensor(100));
r_.AddUse(/*ioperation=*/r_.AddOperation(/*is_stateful=*/false),
/*itensor=*/r_.AddTensor(1));
r_.AddUse(/*ioperation=*/r_.AddOperation(/*is_stateful=*/false),
/*itensor=*/r_.AddTensor(100));
ASSERT_THAT(r_.GetMemProfile(), ElementsAreArray({100, 1, 100}));
EXPECT_THAT(r_.GetPeakMemory(), FieldsAre(Eq(2), Eq(100)));
}
TEST_F(RematTest, RematRecreatesOutput) {
r_.AddUse(r_.AddOperation(/*is_stateful=*/false), r_.AddTensor(100));
r_.AddOperation(/*is_stateful=*/false);
// /* before: */
// %0 = f1()
// f2()
ASSERT_THAT(r_.GetMemProfile(), ElementsAre(100, 0));
EXPECT_THAT(r_.GetMemProfile({/*begin=*/0, /*end=*/1, /*insert=*/2}),
ElementsAre(100, 0, 100));
r_.Remat({/*begin=*/0, /*end=*/1, /*insert=*/2});
// /* after: */
// %0 = f1()
// f2()
// %1 = f1()
EXPECT_THAT(r_.GetMemProfile(), ElementsAre(100, 0, 100));
EXPECT_THAT(r_.AddTensor(0), 2);
}
TEST_F(RematTest, RematExtendsInputAndRecreatesOutput) {
r_.AddUse(/*ioperation=*/r_.AddOperation(/*is_stateful=*/false),
/*itensor=*/r_.AddTensor(1));
r_.AddUse(/*ioperation=*/r_.AddOperation(/*is_stateful=*/false),
/*itensor=*/r_.AddTensor(100));
r_.AddUse(/*ioperation=*/1, 0);
r_.AddOperation(/*is_stateful=*/false);
r_.AddOperation(/*is_stateful=*/false);
// /* before: */
// %0 = f1()
// %1 = f2(%0)
// f3()
// f4()
ASSERT_THAT(r_.GetMemProfile(), ElementsAre(1, 101, 0, 0));
EXPECT_THAT(r_.GetMemProfile({/*begin=*/1, /*end=*/2, /*insert=*/3}),
ElementsAre(1, 101, 1, 101, 0));
r_.Remat({/*begin=*/1, /*end=*/2, /*insert=*/3});
// /* after: */
// %0 = f1()
// %1 = f2(%0)
// f3() /* %0 is kept alive */
// %2 = f2(%0)
// f4()
EXPECT_THAT(r_.GetMemProfile(), ElementsAre(1, 101, 1, 101, 0));
EXPECT_THAT(r_.AddTensor(0), 3);
}
TEST_F(RematTest, BlockRematDuplicatesIntraBlockValues) {
r_.AddUse(/*ioperation=*/r_.AddOperation(/*is_stateful=*/false),
/*itensor=*/r_.AddTensor(1));
r_.AddUse(/*ioperation=*/r_.AddOperation(/*is_stateful=*/false),
/*itensor=*/r_.AddTensor(10));
r_.AddUse(/*ioperation=*/r_.AddOperation(/*is_stateful=*/false),
/*itensor=*/r_.AddTensor(100));
r_.AddUse(/*ioperation=*/r_.AddOperation(/*is_stateful=*/false),
/*itensor=*/r_.AddTensor(1000));
r_.AddOperation(/*is_stateful=*/false);
r_.AddUse(/*ioperation=*/1, /*itensor=*/0);
r_.AddUse(/*ioperation=*/2, /*itensor=*/0);
r_.AddUse(/*ioperation=*/2, /*itensor=*/1);
r_.AddUse(/*ioperation=*/3, /*itensor=*/0);
r_.AddUse(/*ioperation=*/3, /*itensor=*/1);
r_.AddUse(/*ioperation=*/3, /*itensor=*/2);
// /* before */
// %0 = f1()
// %1 = f2(%0)
// %2 = f3(%0,%1)
// %3 = f4(%0,%1,%2)
// f5()
ASSERT_THAT(r_.GetMemProfile(), ElementsAre(1, 11, 111, 1111, 0));
EXPECT_THAT(r_.GetMemProfile({/*begin=*/1, /*end=*/4, /*insert=*/5}),
ElementsAre(1, 11, 111, 1111, 1, 11, 111, 1111));
r_.Remat({/*begin=*/1, /*end=*/4, /*insert=*/5});
EXPECT_THAT(r_.GetMemProfile(),
ElementsAre(1, 11, 111, 1111, 1, 11, 111, 1111));
EXPECT_THAT(r_.AddTensor(0), 7);
// /* after */
// %0 = f1()
// %1 = f2(%0)
// %2 = f3(%0,%1)
// %3 = f4(%0,%1,%2)
// f5()
// %4 = f2(%0)
// %5 = f3(%0,%4)
// %6 = f4(%0,%4,%5)
}
class RematSimulationTest : public testing::Test {
protected:
class RandomRemat : public Rematerializer {
public:
using Rematerializer::Remat; // For testing.
RandomRemat(const int num_operations, const int num_tensors,
const int num_uses, std::mt19937& rng) {
std::uniform_int_distribution<int> some_size_log(0, 16);
std::uniform_int_distribution<int> some_tensor(0, num_tensors - 1);
std::uniform_int_distribution<int> some_operation(0, num_operations - 1);
for (int i = 0; i < num_tensors; ++i) {
AddTensor(SizeT{1} << some_size_log(rng));
}
for (int i = 0; i < num_operations; ++i) {
AddOperation(/*is_stateful=*/false);
}
for (int i = 0; i < num_uses; ++i) {
AddUse(some_operation(rng), some_tensor(rng));
}
}
};
};
TEST_F(RematSimulationTest, SimulationAgreesWithReality) {
constexpr int kNumOperations = 128;
constexpr int kNumTensors = 32;
constexpr int kNumUses = kNumOperations * kNumTensors / 4;
std::mt19937 rng;
for (int i = 0; i < 1024; ++i) {
RandomRemat remat(kNumOperations, kNumTensors, kNumUses, rng);
// Worst-case scenario: we might double the length of the computation
// schedule each time we remat, so only a few iterations...
std::array<int, 3> randos;
const auto& [begin, end, insert] = randos;
for (int i = 0, num_operations = kNumOperations; i < 4;
++i, num_operations += end - begin) {
std::uniform_int_distribution<int> some_op(0, num_operations - 1);
for (auto& rando : randos) {
rando = some_op(rng);
}
// We need begin <= end <= insert.
std::sort(randos.begin(), randos.end());
const Rematerializer::RematSpec spec{begin, end, insert};
const auto simulated_profile = remat.GetMemProfile(spec);
remat.Remat(spec);
const auto actual_profile = remat.GetMemProfile();
EXPECT_THAT(simulated_profile, ElementsAreArray(actual_profile));
}
}
}
class GreedyRematTest : public testing::Test {
protected:
// A Rematerializer for a graph whose memory profile is a sequence of
// "rainbows" due to nested ops (similar to what happens in backwards pass
// gradient calculations.)
class RainbowRemat : public Rematerializer {
public:
// `sizes` is a vector of vector of sizes. Each sizes vector of length n
// generates a `rainbow` profile of the SSA form
//
// %1 = f1()
// %2 = f2()
// ...
// %n = fn()
// gn(%n)
// ...
// g2(%2)
// g1(%1)
//
// with the sizes of intermediates %1,...%n given by the sizes entries.
//
// If extra_ops > 0, each assignment above becomes a sequence
// %i.1 = fi.1()
// %i.2 = fi.1(%i.1)
// %i.3 = fi.1(%i.2)
// ...
// %i = f1(%i.m)
// where the 'dotted' intermediates have size `extra_size`. This
// allows for testing the effect of window sizes.
explicit RainbowRemat(const std::vector<std::vector<int>>& sizes,
int extra_ops = 0, SizeT extra_size = 0) {
for (const auto& rainbow : sizes) {
int tensor = 0;
int op = 0;
for (const auto& size : rainbow) {
for (int i = 0; i < extra_ops; ++i) {
op = AddOperation(/*is_stateful=*/false);
if (i != 0) {
AddUse(op, tensor);
}
tensor = AddTensor(extra_size);
AddUse(op, tensor);
}
// using negative sizes to signal forbidden operations.
op = AddOperation(/*is_stateful=*/size < 0);
if (extra_ops > 0) {
AddUse(op, tensor);
}
tensor = AddTensor(std::abs(size));
AddUse(op, tensor);
}
for (int i = 0; i < rainbow.size(); ++i) {
op = AddOperation(/*is_stateful=*/false);
AddUse(op, tensor - i);
}
}
}
};
// Similar to above: A multilayer perceptron. The forward pass is
// f[n](f[n-1](...(f[1]())), with dimensions |f[1]|...|f[n]| handed in the
// constructor. The backward pass calculates the loss gradient for a single
// weight in the first layer.
class MlpRemat : public Rematerializer {
public:
explicit MlpRemat(const std::vector<int>& sizes) {
int forward_tensor = -1;
int backward_tensor = -1;
int op = -1;
// Forward pass:
for (const int size : sizes) {
op = AddOperation(/*is_stateful=*/false);
if (forward_tensor >= 0) AddUse(op, forward_tensor);
forward_tensor = AddTensor(size);
AddUse(op, forward_tensor);
}
// Backward pass: Right-multiply the jacobian from the outside in.
// dLoss/df[n] * df[n]/df[n-1] * ...
// The i-th term g[i] depends on g[i-1] and f[n-i] and has size |f[n-i]|.
for (; forward_tensor >= 0; --forward_tensor) {
op = AddOperation(/*is_stateful=*/false);
AddUse(op, forward_tensor);
if (backward_tensor >= 0) AddUse(op, backward_tensor);
backward_tensor = AddTensor(sizes[forward_tensor]);
AddUse(op, backward_tensor);
}
}
// We will also instrument the actual optimizations performed.
MOCK_METHOD(void, ApplyRemat, (const RematSpec&));
};
};
TEST_F(GreedyRematTest, MlpBasic) {
StrictMock<MlpRemat> remat(std::vector<int>({1, 1, 1}));
// (o)ut, (i)n, (l)ive: 0 1 2 3 4 5 Sum
// %0 = f0() o 1
// %1 = f1(%0) i o 2
// %2 = f2(%1) l i o 3
// %3 = g2( %2) l l i o 4
// %4 = g1(%3,%1) l i i o 4
// %5 = g0(%4,%0) i i o 3
ASSERT_THAT(remat.GetMemProfile(), ElementsAreArray({1, 2, 3, 4, 4, 3}));
// Little we can do here -- remat %0 before %5
EXPECT_CALL(remat, ApplyRemat(FieldsAre(/*begin=*/0,
/*end=*/1,
/*insert=*/5)));
// (o)ut, (i)n, (l)ive: 0 1 2 3 4 5 Sum of live sizes
// %0 = f0() o 1
// %1 = f1(%0) i o 2
// %2 = f2(%1) i o 2
// %3 = g2( %2) l i o 3
// %4 = g1(%3,%1) i i o 3
// %0' = f0() o l 2
// %5 = g0(%4,%0') i i o 3
remat.RunGreedyAlgorithm(/*max_cost=*/-1, /*max_block_length=*/1,
/*min_savings=*/1);
EXPECT_THAT(remat.GetMemProfile(), ElementsAreArray({1, 2, 2, 3, 3, 2, 3}));
}
TEST_F(GreedyRematTest, MlpBinary) {
StrictMock<MlpRemat> remat(std::vector<int>({1, 2, 4, 8}));
// (o)ut, (i)n, (l)ive: 0 1 2 3 4 5 6 7 Sum of live sizes
// %0 = f0() o 1
// %1 = f1(%0) i o 3
// %2 = f2(%1) l i o 7
// %3 = f3(%2) l l i o 15
// %4 = g3( %3) l l l i o 23
// %5 = g2(%4 %2) l l i i o 19
// %6 = g1(%5,%1) l i i o 9
// %7 = g0(%6,%0) i i o 4
ASSERT_THAT(remat.GetMemProfile(),
ElementsAreArray({1, 3, 7, 15, 23, 19, 9, 4}));
EXPECT_CALL(remat, ApplyRemat(FieldsAre(/*begin=*/2,
/*end=*/3,
/*insert=*/5)));
EXPECT_CALL(remat, ApplyRemat(FieldsAre(/*begin=*/0,
/*end=*/1,
/*insert=*/8)));
// (o)ut, (i)n, (l)ive: 0 1 2 3 4 5 6 7 Sum of live sizes
// %0 = f0() o 1
// %1 = f1(%0) i o 3
// %2 = f2(%1) i o 6
// %3 = f3(%2) l i o 14
// %4 = g3( %3) l i o 18
// %2'= f2(%1) l o l 14
// %5 = g2(%4 %2) l i i o 18
// %6 = g1(%5,%1) i i o 8
// %0'= f(0) o l 3
// %7 = g0(%6,%0) i i o 4
remat.RunGreedyAlgorithm(/*max_cost=*/-1, /*max_block_length=*/4,
/*min_savings=*/1);
EXPECT_THAT(remat.GetMemProfile(),
ElementsAreArray({1, 3, 6, 14, 18, 14, 18, 8, 3, 4}));
}
TEST_F(GreedyRematTest, SimpleMax) {
RainbowRemat remat({{1, 2, 4, 8, 16}});
ASSERT_THAT(remat.GetMemProfile(),
ElementsAreArray({1, 3, 7, 15, 31, 31, 15, 7, 3, 1}));
remat.RunGreedyAlgorithm(/*max_cost=*/-1, /*max_block_length=*/1,
/*min_savings=*/1);
// Profile is flattened to its minimum--16.
EXPECT_THAT(remat.GetMemProfile(),
ElementsAreArray({1, 2, 4, 8, 16, 16, 8, 8, 4, 4, 2, 2, 1, 1}));
}
TEST_F(GreedyRematTest, SimpleMaxLongWindow) {
RainbowRemat remat({{1, 2, 4, 8, 16}});
ASSERT_THAT(remat.GetMemProfile(),
ElementsAreArray({1, 3, 7, 15, 31, 31, 15, 7, 3, 1}));
remat.RunGreedyAlgorithm(/*max_cost=*/-1, /*max_block_length=*/4,
/*min_savings=*/1);
// Profile is flattened to its minimum--16.
EXPECT_THAT(remat.GetMemProfile(),
ElementsAreArray({1, 2, 4, 8, 16, 16, 8, 8, 4, 4, 2, 2, 1, 1}));
}
TEST_F(GreedyRematTest, SimpleSizeThreshold) {
RainbowRemat remat({{1, 2, 4, 8, 16}});
ASSERT_THAT(remat.GetMemProfile(),
ElementsAreArray({1, 3, 7, 15, 31, 31, 15, 7, 3, 1}));
// Only do remats of at least 4 bytes of savings -- this will lower the
// profile by 4 + 8 instead of 1 + 2 + 4 + 8 as before.
remat.RunGreedyAlgorithm(/*max_cost=*/-1, /*max_block_length=*/1,
/*min_savings=*/4);
EXPECT_THAT(remat.GetMemProfile(),
ElementsAreArray({1, 3, 7, 11, 19, 19, 11, 11, 7, 7, 3, 1}));
}
TEST_F(GreedyRematTest, SimpleCostThreshold) {
RainbowRemat remat({{1, 2, 4, 8, 16}});
ASSERT_THAT(remat.GetMemProfile(),
ElementsAreArray({1, 3, 7, 15, 31, 31, 15, 7, 3, 1}));
// Do at most 1 remat -- this will lower the profile by 8.
remat.RunGreedyAlgorithm(/*max_cost=*/1, /*max_block_length=*/1,
/*min_savings=*/1);
// Only as single remat is done -- this will be the best one possible,
// lowering the profile by 8 instead of the maximum 1 + 2 + 4 + 8.
EXPECT_THAT(remat.GetMemProfile(),
ElementsAreArray({1, 3, 7, 15, 23, 23, 15, 15, 7, 3, 1}));
}
TEST_F(GreedyRematTest, SimpleForbiddenOps) {
// Operator generating size-4 tensor is stateful, so it won't be materialized.
RainbowRemat remat({{1, 2, -4, 8, 16}});
ASSERT_THAT(remat.GetMemProfile(),
ElementsAreArray({1, 3, 7, 15, 31, 31, 15, 7, 3, 1}));
// Best we can do is lower the profile to 8 + 2 + 1.
remat.RunGreedyAlgorithm(/*max_cost=*/-1, /*max_block_length=*/1,
/*min_savings=*/1);
EXPECT_THAT(remat.GetMemProfile(),
ElementsAreArray({1, 2, 4, 12, 20, 20, 12, 12, 4, 2, 2, 1, 1}));
}
TEST_F(GreedyRematTest, DoubleMax) {
// Generate a profile with two local maxima of size 31 and 28.
RainbowRemat remat({{1, 2, 4, 8, 16}, {4, 8, 16}});
ASSERT_THAT(remat.GetMemProfile(),
ElementsAreArray(
{1, 3, 7, 15, 31, 31, 15, 7, 3, 1, 4, 12, 28, 28, 12, 4}));
remat.RunGreedyAlgorithm(/*max_cost=*/-1, /*max_block_length=*/1,
/*min_savings=*/1);
// Two rainbows; the first is lowered by 1 + 2 + 4 + 8 == 15, the second by 4
// + 8 == 12.
EXPECT_THAT(remat.GetMemProfile(),
ElementsAreArray({1, 2, 4, 8, 16, 16, 8, 8, 4, 4, 2,
2, 1, 1, 4, 8, 16, 16, 8, 8, 4, 4}));
}
TEST_F(GreedyRematTest, DoubleCostThreshold) {
// Generate a profile with two local maxima of size 31 and 28.
RainbowRemat remat({{1, 2, 4, 8, 16}, {4, 8, 16}});
ASSERT_THAT(remat.GetMemProfile(),
ElementsAreArray(
{1, 3, 7, 15, 31, 31, 15, 7, 3, 1, 4, 12, 28, 28, 12, 4}));
remat.RunGreedyAlgorithm(/*max_cost=*/2, /*max_block_length=*/1,
/*min_savings=*/1);
// Profile can be flattened twice--first, the global maximum of 31 is reduced
// by 8; then the newly-global maximum of 28 is reduced by 8.
EXPECT_THAT(remat.GetMemProfile(),
ElementsAreArray({1, 3, 7, 15, 23, 23, 15, 15, 7, 3, 1, 4, 12, 20,
20, 12, 12, 4}));
}
TEST_F(GreedyRematTest, SingleLongerBlocksByWindowSize) {
std::vector<Rematerializer::SizeT> best_for_window_size;
for (int window_size : {0, 1, 2, 3, 4, 5}) {
RainbowRemat remat({{1, 2, 4, 8}}, 2, 16);
remat.RunGreedyAlgorithm(/*max_cost=*/-1, /*max_block_length=*/window_size,
/*min_savings=*/1);
best_for_window_size.push_back(remat.GetPeakMemory().size);
}
EXPECT_THAT(best_for_window_size, ElementsAreArray({44, 36, 36, 32, 32, 32}));
}
} // namespace
} // namespace TFL
} // namespace mlir
tensorflow/compiler/mlir/lite/experimental/tac/BUILD
+399
View File
@@ -0,0 +1,399 @@
load("@flatbuffers//:build_defs.bzl", "flatbuffer_cc_library")
load(
"@llvm-project//mlir:tblgen.bzl",
"gentbl_cc_library",
)
load("@xla//third_party/rules_python/python:py_library.bzl", "py_library")
load("//tensorflow:tensorflow.bzl", "if_google", "tf_cc_binary", "tf_cc_test")
load("//tensorflow:tensorflow.default.bzl", "get_compatible_with_portable")
load(
"//tensorflow/core/platform:build_config.bzl",
"tf_proto_library",
)
load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:LICENSE"],
default_visibility = [
"//visibility:public",
],
licenses = ["notice"],
)
flatbuffer_cc_library(
name = "runtime_metadata_fbs",
srcs = ["runtime_metadata.fbs"],
compatible_with = get_compatible_with_portable(),
)
cc_library(
name = "execution_metadata_exporter",
srcs = [
"execution_metadata_exporter.cc",
],
hdrs = [
"execution_metadata_exporter.h",
],
deps = [
":common",
":runtime_metadata_fbs",
"//tensorflow/compiler/mlir/lite:tensorflow_lite",
"//tensorflow/compiler/mlir/lite/experimental/tac/hardwares:target_hardware",
"//tensorflow/compiler/mlir/lite/quantization/ir:QuantOps",
"//tensorflow/compiler/mlir/tensorflow",
"@flatbuffers",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:ArithDialect",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Support",
],
)
tf_cc_test(
name = "execution_metadata_exporter_test",
srcs = [
"execution_metadata_exporter_test.cc",
],
deps = [
":execution_metadata_exporter",
":runtime_metadata_fbs",
"//tensorflow/compiler/mlir/lite:tensorflow_lite",
"//tensorflow/compiler/mlir/lite/experimental/tac/hardwares:all-target-hardwares",
"@com_google_googletest//:gtest_main",
"@flatbuffers",
"@llvm-project//mlir:ArithDialect",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Parser",
],
)
cc_library(
name = "common",
srcs = [
"common/utils.cc",
],
hdrs = [
"common/cost.h",
"common/subgraph.h",
"common/targets.h",
"common/utils.h",
],
deps = [
"//tensorflow/compiler/mlir/lite:tensorflow_lite",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:ArithDialect",
"@llvm-project//mlir:BytecodeOpInterface",
"@llvm-project//mlir:CallOpInterfaces",
"@llvm-project//mlir:CastInterfaces",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:QuantOps",
"@llvm-project//mlir:SideEffectInterfaces",
"@llvm-project//mlir:Support",
],
)
gentbl_cc_library(
name = "transform_patterns_inc_gen",
compatible_with = get_compatible_with_portable(),
tbl_outs = {"transforms/generated_transform_patterns.inc": ["-gen-rewriters"]},
tblgen = "@llvm-project//mlir:mlir-tblgen",
td_file = "transforms/transform_patterns.td",
deps = [
"//tensorflow/compiler/mlir/lite:tensorflow_lite_ops_td_files",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_ops_td_files",
"@llvm-project//mlir:ArithOpsTdFiles",
"@llvm-project//mlir:FuncTdFiles",
],
)
cc_library(
name = "device_transform_patterns",
srcs = [
"transforms/device_transform_patterns.cc",
],
hdrs = [
"transforms/device_transform_patterns.h",
],
textual_hdrs = [
"transforms/generated_transform_patterns.inc",
],
deps = [
":common",
"//tensorflow/compiler/mlir/lite:tensorflow_lite",
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/compiler/mlir/tensorflow:verification_utils",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Pass",
"@llvm-project//mlir:QuantOps",
"@llvm-project//mlir:Support",
],
)
cc_library(
name = "device_transform_gpu",
srcs = [
"transforms/device_transform_gpu.cc",
],
hdrs = [
"transforms/device_transform_gpu.h",
],
deps = [
":common",
"//tensorflow/compiler/mlir/lite:tensorflow_lite",
"//tensorflow/compiler/mlir/lite/experimental/tac/hardwares:gpu_hardware",
"//tensorflow/compiler/mlir/tensorflow",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Pass",
"@llvm-project//mlir:Support",
"@llvm-project//mlir:TransformUtils",
],
alwayslink = 1,
)
cc_library(
name = "device_transform_nnapi",
srcs = [
"transforms/device_transform_nnapi.cc",
],
deps = [
":common",
"//tensorflow/compiler/mlir/lite:tensorflow_lite",
"//tensorflow/compiler/mlir/lite/experimental/tac/hardwares:nnapi_hardware",
"//tensorflow/compiler/mlir/tensorflow",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Pass",
"@llvm-project//mlir:Support",
"@llvm-project//mlir:TransformUtils",
],
alwayslink = 1,
)
cc_library(
name = "device_transform",
srcs = [
"transforms/device_transform.cc",
],
hdrs = [
"transforms/device_transform.h",
],
deps = [
":common",
":device_transform_gpu",
":device_transform_nnapi",
":device_transform_patterns",
"//tensorflow/compiler/mlir/lite:tensorflow_lite",
"//tensorflow/compiler/mlir/lite/experimental/tac/hardwares:target_hardware",
"//tensorflow/compiler/mlir/tensorflow",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:QuantOps",
"@llvm-project//mlir:Support",
"@llvm-project//mlir:TransformUtils",
"@llvm-project//mlir:Transforms",
],
)
cc_library(
name = "cost_model",
srcs = [
"transforms/cost_model.cc",
],
hdrs = [
"transforms/cost_model.h",
],
deps = [
":common",
"//tensorflow/compiler/mlir/lite:tensorflow_lite",
"//tensorflow/compiler/mlir/lite/experimental/tac/hardwares:target_hardware",
"@com_google_absl//absl/strings",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:CallOpInterfaces",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Pass",
"@llvm-project//mlir:Support",
],
alwayslink = 1,
)
# TODO(b/177376459): split tac_module and passes dependency to separate libraries.
cc_library(
name = "target_aware_conversion",
srcs = [
"tac_module.cc",
"transforms/compute_cost.cc",
"transforms/fold_constants_to_subgraph.cc",
"transforms/get_alternative_subgraph.cc",
"transforms/pick_subgraphs.cc",
"transforms/raise_target_subgraphs.cc",
"transforms/tac_filter.cc",
"transforms/target_annotation.cc",
],
hdrs = [
"tac_module.h",
"transforms/passes.h",
"transforms/tac_pass.h",
],
deps = [
":common",
":cost_model",
":device_transform",
":tac_filter_proto_cc",
":tac_importer_exporter",
"//tensorflow/compiler/mlir/lite:tensorflow_lite",
"//tensorflow/compiler/mlir/lite:tf_tfl_passes",
"//tensorflow/compiler/mlir/lite/experimental/common:outline_operations",
"//tensorflow/compiler/mlir/lite/experimental/tac/hardwares:target_hardware",
"//tensorflow/compiler/mlir/tensorflow:cluster_util",
"//tensorflow/compiler/mlir/tensorflow:error_util",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_analysis",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings",
"@com_google_protobuf//:protobuf_headers",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:ArithDialect",
"@llvm-project//mlir:CallOpInterfaces",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Pass",
"@llvm-project//mlir:SideEffectInterfaces",
"@llvm-project//mlir:Support",
"@llvm-project//mlir:TransformUtils",
"@llvm-project//mlir:Transforms",
"@stablehlo//:stablehlo_ops",
],
alwayslink = 1,
)
cc_library(
name = "tac-opt_lib",
testonly = True,
deps = [
":target_aware_conversion",
"//tensorflow/compiler/mlir/lite:litert_mlir_opt_main",
],
alwayslink = 1,
)
# Binary with no hardwares linked.
tf_cc_binary(
name = "tac-opt",
testonly = True,
deps = [
":tac-opt_lib",
],
)
# Binary with all backends linked.
tf_cc_binary(
name = "tac-opt-all-backends",
testonly = True,
deps = [
":tac-opt_lib",
"//tensorflow/compiler/mlir/lite/experimental/common:outline_operations",
"//tensorflow/compiler/mlir/lite/experimental/tac/hardwares:all-target-hardwares",
],
)
cc_library(
name = "tac-translate-lib",
srcs = [
"tac_translate.cc",
],
deps = [
":common",
":execution_metadata_exporter",
":target_aware_conversion",
":tflite_importer_exporter",
"//tensorflow/compiler/mlir:init_mlir",
"//tensorflow/compiler/mlir/lite:tensorflow_lite_legalize_tf",
"//tensorflow/compiler/mlir/lite:tensorflow_lite_optimize",
"//tensorflow/compiler/mlir/lite/experimental/tac/hardwares:target_hardware",
"//tensorflow/compiler/mlir/lite/experimental/tac/utils",
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/compiler/mlir/tensorflow:error_util",
"//tensorflow/compiler/mlir/tensorflow/transforms:tensorflow_passes",
"//tensorflow/core:lib",
"@com_google_absl//absl/log",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/strings:str_format",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Parser",
"@llvm-project//mlir:Pass",
"@llvm-project//mlir:Support",
"@llvm-project//mlir:Transforms",
"@llvm-project//mlir:TranslateLib",
],
)
tf_cc_binary(
name = "tac-translate",
deps = [
":tac-translate-lib",
"//tensorflow/compiler/mlir/lite/experimental/tac/hardwares:all-target-hardwares",
],
)
cc_library(
name = "tac_importer_exporter",
hdrs = ["tac_importer_exporter.h"],
deps = [
"@com_google_absl//absl/status:statusor",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
],
)
cc_library(
name = "tflite_importer_exporter",
srcs = ["tflite_import_export.cc"],
hdrs = ["tflite_import_export.h"],
deps = [
":common",
":execution_metadata_exporter",
":tac_importer_exporter",
"//tensorflow/compiler/mlir/lite/experimental/tac/hardwares:target_hardware",
"//tensorflow/compiler/mlir/lite/experimental/tac/utils",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:Support",
],
)
exports_files([
"run_lit.sh",
])
py_library(
name = "tac",
srcs = [
"tac.py",
],
strict_deps = True,
visibility = ["//visibility:public"],
deps = [
"//tensorflow/compiler/mlir/lite/experimental/tac/py_wrapper:_pywrap_tac_wrapper",
],
)
tf_proto_library(
name = "tac_filter_proto",
srcs = ["tac_filter.proto"],
protodeps = if_google(["//google/protobuf:any"]),
visibility = ["//visibility:public"],
)
tensorflow/compiler/mlir/lite/experimental/tac/README.md
+271
View File
@@ -0,0 +1,271 @@
# Target Aware Conversion (TAC)
Different hardwares have different capabilities and restrictions.
TAC is designed to leverage hardwares' capabilities to:
* Perform device-specific optimizations (such as unsupported ops lowering,
layout transformations, etc.)
* Graph partitioning based on the hardware costs modeling.
* It supports general import/export where you can hook your own
importer/exporter from any format to MLIR and export MLIR to anything.
For more details, please checkout the
[TAC workflow](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/compiler/mlir/lite/experimental/tac/README.md#tac-workflow)
section
## How to use
Once you have a converted TfLite model ready, you can use the following command
to use TAC to optimize for your model:
```
bazel run -c opt //tensorflow/compiler/mlir/lite/experimental/tac:tac-translate -- <PATH_TO_YOUR_MODEL> -o=<OUTPUT_PATH> -device-specs=<HARDWARE_BACKENDS>
```
The devices_specs is a list of the names of the desired hardware backends,
separated by comma, e.g., "GPU,CPU".
If you're interested in what are the subgraphs being explored for different
backends, you can pass in `-output-mlir -inline-subgraphs=false` and check out
the output mlir file.
## How to add a hardware backend
If you want to add a hardware backend for TAC, you can start with the
`SimpleHardware` interface.
For example:
```
class FooHardware : public SimpleHardware {
public:
static constexpr char kId[] = "FOO";
mlir::RewritePatternSet GetTransformations(
MLIRContext* context) const override {
mlir::RewritePatternSet patterns;
// Pick the transformations that we want to perform,
// We can add other transformations we like here.
patterns.add<LowerPackIntoConcatReshape, UnrollSplit, UnrollSplitV,
PadSlice>(context);
return patterns;
}
mlir::TypeID GetTypeId() const override {
return mlir::TypeID::get<FooHardware>();
}
// We can specify what ops are not supported here.
bool IsNotSupportedOp(mlir::Operation* op) const override { return false; }
// This is basically saying how fast are we comparing to CPU.
// The larger the value the better.
float AdvantageOverCPU() const override { return 5.0; }
};
```
Then we need to register our hardware like below:
```
std::unique_ptr<TargetHardware> CreateFooHardware() {
return std::make_unique<FooHardware>();
}
TargetHardwareRegistration<FooHardware> foo_hardware(
"Target device for FOO", CreateFooHardware);
```
### Advanced user
For advanced users (e.g., you may already have your own hardware dialect
defined), please just use `TargetHardware` directly. See the following code
snippet for reference.
```
class MyCustomHardware : public TargetHardware {
public:
static constexpr char kId[] = "MY_CUSTOM_HARDWARE";
mlir::TypeID GetTypeId() const override {
return mlir::TypeID::get<MyCustomHardware>();
}
bool IsOpSupported(mlir::Operation* op) const override {
// check whether the op is supported, if the user has they own dialect,
// this can be target dialect legalization process.
}
double GetHardwareSwitchingCost(const TargetHardware* from,
size_t buffer_size) const override {
// Get the hardware switching cost from the source hardware.
}
double GetOpCost(mlir::Operation* op) const override {
// call customized cost model.
}
mlir::RewritePatternSet GetTransformations(
MLIRContext* context) const override {
// customized transformations patterns: ops lowering/fusion, layout
// transformation, etc.
}
};
```
## TAC workflow
The workflow of target-aware-conversion is as followed:
1 Try to break down the whole graph into several subgraphs based on hardwares'
capabilities. See the diagram below, let's say our desired target backends are
"GPU" and "CPU", and currently "C" is not supported on "GPU", but the rest are
supported by "GPU". So we will end up with 3 subgraphs as shown in the diagram.
![Target Annotation](g3doc/images/target_annotation.png)
2 Perform ops-lowering & target-specific optimizations for
different hardware backends. As shown in the below diagram, the red & the
yellow subgraph will be duplicated as "alternative subgraph view" for "CPU".
"C" op can be lowered into "G" + "H" op which can be supported by "GPU".
![Target Optimization](g3doc/images/target_optimization.png)
3 Estimate the costs for each subgraph (and their alternative views)
based on the hardware cost model. See the following diagram.
![Estimate costs](g3doc/images/compute_cost.png)
4 Pick the proper subgraphs from the alternative views for execution based on
costs(computation costs, transfer costs, quant/dequant costs). As shown in the
diagram below, since cross-device data transferring cost is high, even "G" + "H"
running on GPU maybe less efficient than "C" running on "CPU", we will still
pick "G" + "H" subgraph.
![Pick subgraphs](g3doc/images/pick_subgraphs.png)
The final graph looks like below:
![Final graph](g3doc/images/final_graph.png)
## TAC components
### Hardwares
Hardwares are used to modeling target device capabilities & also ops cost for
the target devices.
We have already modeled `cpu_hardware` & `gpu_hardware` as well as the
`nnapi_hardware`.
### Passes
#### Target Annotation Pass
In this pass, every op will be targeted with the user specified targets based on
the device capabilites. For example, If the user specified the desired targets
are "GPU", "CPU", `conv2d` can run on both "GPU" and "CPU", we will annotate
the op `conv2d` with "GPU" since it's preferred; `pack` can only run on "CPU",
so we will annotate the op with "CPU" since "GPU" does not support this op.
#### Raise Target Subgraphs Pass
In this pass, ops will be broken down into subgraph. Those ops have the same
target annotation will be raised as subgraphs.
In this pass, subgraph is actually implemented with `FuncOp`.
Take the following code as an example:
```
func @simpleTest(%arg0: tensor<1xf32>, %arg1: tensor<1xf32>, %arg2: tensor<1xf32>, %arg3: tensor<1xf32>) -> tensor<2x1xf32> {
%0 = "tfl.add"(%arg0, %arg1) {tac.device = "GPU", fused_activation_function = "RELU6", tac.inference_type = "FLOAT"} : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32>
%1 = "tfl.mul"(%0, %arg2) {tac.device = "GPU", fused_activation_function = "RELU6", tac.inference_type = "FLOAT"} : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32>
%2 = "tfl.add"(%arg0, %arg3) {tac.device = "GPU", fused_activation_function = "RELU6", tac.inference_type = "FLOAT"} : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32>
%3 = "tfl.pack"(%1, %2) {tac.device = "CPU", tac.inference_type = "FLOAT", axis = 0 : i32, values_count = 2 : i32} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2x1xf32>
return %3 : tensor<2x1xf32>
}
```
In this code, `%3` is annotated with "CPU", while others are annotated with
"GPU", in this case, `%3` will be raised as a separate function like below:
```
func private @func_1_GPU_FLOAT(%arg0: tensor<1xf32>, %arg1: tensor<1xf32>) -> tensor<1xf32> attributes {tac.device = "GPU", tac.inference_type = "FLOAT", tac.interface_name = "func_1"} {
%0 = tfl.add %arg0, %arg1 {fused_activation_function = "RELU6", tac.device = "GPU", tac.inference_type = "FLOAT"} : tensor<1xf32>
return %0 : tensor<1xf32>
}
```
And the rest ops will be raised as below:
```
func private @func_2_CPU_FLOAT(%arg0: tensor<1xf32>, %arg1: tensor<1xf32>) -> tensor<2x1xf32> attributes {tac.device = "CPU", tac.inference_type = "FLOAT", tac.interface_name = "func_2"} {
%0 = "tfl.pack"(%arg0, %arg1) {axis = 0 : i32, tac.device = "CPU", tac.inference_type = "FLOAT", values_count = 2 : i32} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2x1xf32>
return %0 : tensor<2x1xf32>
}
func private @func_0_GPU_FLOAT(%arg0: tensor<1xf32>, %arg1: tensor<1xf32>, %arg2: tensor<1xf32>) -> tensor<1xf32> attributes {tac.device = "GPU", tac.inference_type = "FLOAT", tac.interface_name = "func_0"} {
%0 = tfl.add %arg0, %arg1 {fused_activation_function = "RELU6", tac.device = "GPU", tac.inference_type = "FLOAT"} : tensor<1xf32>
%1 = tfl.mul %0, %arg2 {fused_activation_function = "RELU6", tac.device = "GPU", tac.inference_type = "FLOAT"} : tensor<1xf32>
return %1 : tensor<1xf32>
}
```
And the original function will be replaced by `CallOps` to those `FuncOps`:
```
func @simpleTest(%arg0: tensor<1xf32>, %arg1: tensor<1xf32>, %arg2: tensor<1xf32>, %arg3: tensor<1xf32>) -> tensor<2x1xf32> {
%0 = call @func_0_GPU_FLOAT(%arg0, %arg1, %arg2) {tac.device = "GPU", tac.inference_type = "FLOAT", tac.interface_name = "func_0"} : (tensor<1xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32>
%1 = call @func_1_GPU_FLOAT(%arg0, %arg3) {tac.device = "GPU", tac.inference_type = "FLOAT", tac.interface_name = "func_1"} : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32>
%2 = call @func_2_CPU_FLOAT(%0, %1) {tac.device = "CPU", tac.inference_type = "FLOAT", tac.interface_name = "func_2"} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2x1xf32>
return %2 : tensor<2x1xf32>
}
```
Why we need to raise those ops into `FuncOps`? Please see the following section.
#### Get Alternative Subgraph View Pass
In the Get Alternative Subgraph View Pass, we will essentially duplicate those
`FuncOps` and perform unsupported ops lowering & target-specific optimization.
For example, `Pack` is not supported by "GPU", but it can be lowered into
`Concat` + `Reshape` which can be supported by "GPU".
So the original example:
```
func private @func_1_GPU_FLOAT(%arg0: tensor<1xf32>, %arg1: tensor<1xf32>) -> tensor<1xf32> attributes {tac.device = "GPU", tac.inference_type = "FLOAT", tac.interface_name = "func_1"} {
%0 = tfl.add %arg0, %arg1 {fused_activation_function = "RELU6", tac.device = "GPU", tac.inference_type = "FLOAT"} : tensor<1xf32>
return %0 : tensor<1xf32>
}
```
Will be transformed into:
```
func private @func_2_CPU_FLOAT(%arg0: tensor<1xf32>, %arg1: tensor<1xf32>) -> tensor<2x1xf32> attributes {tac.device = "CPU", tac.inference_type = "FLOAT", tac.interface_name = "func_2"} {
%0 = "tfl.pack"(%arg0, %arg1) {axis = 0 : i32, tac.device = "CPU", tac.inference_type = "FLOAT", values_count = 2 : i32} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2x1xf32>
return %0 : tensor<2x1xf32>
}
func private @func_2_GPU_FLOAT(%arg0: tensor<1xf32>, %arg1: tensor<1xf32>) -> tensor<2x1xf32> attributes {tac.device = "GPU", tac.inference_type = "FLOAT", tac.interface_name = "func_2"} {
%cst = arith.constant dense<1> : tensor<4xi32>
%cst_0 = arith.constant dense<2> : tensor<1xi32>
%cst_1 = arith.constant dense<[2, 1]> : tensor<2xi32>
%0 = "tfl.reshape"(%arg0, %cst) {tac.device = "GPU", tac.inference_type = "FLOAT"} : (tensor<1xf32>, tensor<4xi32>) -> tensor<1x1x1x1xf32>
%1 = "tfl.reshape"(%arg1, %cst) {tac.device = "GPU", tac.inference_type = "FLOAT"} : (tensor<1xf32>, tensor<4xi32>) -> tensor<1x1x1x1xf32>
%2 = "tfl.concatenation"(%0, %1) {axis = 3 : i32, fused_activation_function = "NONE", tac.device = "GPU", tac.inference_type = "FLOAT"} : (tensor<1x1x1x1xf32>, tensor<1x1x1x1xf32>) -> tensor<1x1x1x2xf32>
%3 = "tfl.reshape"(%2, %cst_0) {tac.device = "GPU", tac.inference_type = "FLOAT"} : (tensor<1x1x1x2xf32>, tensor<1xi32>) -> tensor<2xf32>
%4 = "tfl.reshape"(%3, %cst_1) {tac.device = "GPU", tac.inference_type = "FLOAT"} : (tensor<2xf32>, tensor<2xi32>) -> tensor<2x1xf32>
return %4 : tensor<2x1xf32>
}
```
#### Compute Costs Pass
In the compute cost pass, we will essentially compute the cost of each op within
the `FuncOp` based on the target-device cost model and sum them together.
#### Pick Subgraphs Pass
In the pick subgraphs pass, we will pick those subgraphs which can minimize the
global costs (we will take the tensor transferring costs as well).
tensorflow/compiler/mlir/lite/experimental/tac/common/cost.h
+50
View File
@@ -0,0 +1,50 @@
/* 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_COMPILER_MLIR_LITE_EXPERIMENTAL_TAC_COMMON_COST_H_
#define TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_TAC_COMMON_COST_H_
#include <string>
#include "mlir/IR/Builders.h" // from @llvm-project
#include "mlir/IR/Operation.h" // from @llvm-project
namespace mlir {
namespace TFL {
namespace tac {
// Cost attribute string on the TFL dialect.
constexpr char kCost[] = "tac.cost";
inline void UpdateCost(Operation* op, float cost, OpBuilder* builder) {
op->setAttr(kCost, builder->getF32FloatAttr(cost));
}
// Get the cost annotated with kCost.
inline bool GetCostOnOp(Operation* op, float* cost) {
auto cost_type = op->getAttrOfType<FloatAttr>(kCost);
if (cost_type == nullptr) {
return false;
}
*cost = cost_type.getValueAsDouble();
return true;
}
} // namespace tac
} // namespace TFL
} // namespace mlir
#endif // TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_TAC_COMMON_COST_H_
tensorflow/compiler/mlir/lite/experimental/tac/common/subgraph.h
+52
View File
@@ -0,0 +1,52 @@
/* 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_COMPILER_MLIR_LITE_EXPERIMENTAL_TAC_COMMON_SUBGRAPH_H_
#define TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_TAC_COMMON_SUBGRAPH_H_
#include <optional>
#include <string>
#include "llvm/ADT/StringRef.h"
#include "mlir/IR/Operation.h" // from @llvm-project
namespace mlir {
namespace TFL {
namespace tac {
// Interface name here is the "hook" between the CallOp and FuncOps.
// Take the following example:
//
// call @func_1_CPU {tac.interface_name = "func_1"}
//
// "func_1" is the interface name where "func_1_cpu" is the real implementation
// we can have multiple FuncOps like "func_1_cpu" and "func_1_gpu" and they
// both implement "func_1".
//
// The attribute on the FuncOp means what it actually implements while the
// attribute on the CallOp means what it actually looks for.
constexpr char kInterfaceNameAttr[] = "tac.interface_name";
inline std::optional<std::string> GetInterFaceName(Operation* op) {
auto name_attr = op->getAttrOfType<StringAttr>(kInterfaceNameAttr);
if (!name_attr) return std::nullopt;
return name_attr.getValue().str();
}
} // namespace tac
} // namespace TFL
} // namespace mlir
#endif // TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_TAC_COMMON_SUBGRAPH_H_
tensorflow/compiler/mlir/lite/experimental/tac/common/targets.h
+158
View File
@@ -0,0 +1,158 @@
/* 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_COMPILER_MLIR_LITE_EXPERIMENTAL_TAC_COMMON_TARGETS_H_
#define TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_TAC_COMMON_TARGETS_H_
#include <algorithm>
#include <cctype>
#include <cstddef>
#include <functional>
#include <optional>
#include <string>
#include <vector>
#include "llvm/ADT/ArrayRef.h"
#include "llvm/ADT/StringRef.h"
#include "mlir/IR/BuiltinAttributes.h" // from @llvm-project
#include "mlir/IR/Operation.h" // from @llvm-project
namespace mlir {
namespace TFL {
namespace tac {
// Device attribute string on the TFL dialect.
constexpr char kDevice[] = "tac.device";
// Inference type.
constexpr char kInferenceType[] = "tac.inference_type";
// Inference type.
constexpr char kSkipTargetAnnotation[] = "tac.skip_target_annotation";
// Custom options fingerprint to apply different options for different filters.
constexpr char kCustomOptionsFingerprint[] = "tac.custom_options_fingerprint";
// Delegate compiler stack version.
constexpr char kDelegateCompilerVersion[] = "tac.delegate_compiler_version";
// TODO(renjieliu): Add more inference types.
enum InferenceType {
UNKNOWN = 0,
FLOAT = 1,
QUANTIZED_INT8 = 2,
QUANTIZED_UINT8 = 3,
HYBRID = 4
};
inline InferenceType GetInferenceTypeEnum(llvm::StringRef inference_type_str) {
if (inference_type_str == "FLOAT") {
return FLOAT;
} else if (inference_type_str == "QUANTIZED_INT8") {
return QUANTIZED_INT8;
} else if (inference_type_str == "QUANTIZED_UINT8") {
return QUANTIZED_UINT8;
} else if (inference_type_str == "HYBRID") {
return HYBRID;
} else {
return UNKNOWN;
}
}
inline std::string GetInferenceString(InferenceType inference_type) {
if (inference_type == FLOAT) {
return "FLOAT";
} else if (inference_type == QUANTIZED_INT8) {
return "QUANTIZED_INT8";
} else if (inference_type == QUANTIZED_UINT8) {
return "QUANTIZED_UINT8";
} else if (inference_type == HYBRID) {
return "HYBRID";
} else {
return "UNKNOWN";
}
}
// Returns canonical representation for hardware name (All uppercase).
// TODO(b/177376459): Remove this in favor of the string defined by hardwares
// MyHardware::kId.
inline std::string GetCanonicalHardwareName(const std::string& hardware_name) {
std::string name = hardware_name;
std::transform(
name.begin(), name.end(), name.begin(),
[](unsigned char c) -> unsigned char { return std::toupper(c); });
return name;
}
// Get the target annotation form the op.
inline std::optional<std::string> GetTargetAnnotation(Operation* op) {
auto device = op->getAttrOfType<StringAttr>(kDevice);
if (device == nullptr || device.getValue().empty()) return std::nullopt;
return GetCanonicalHardwareName(device.getValue().str());
}
// Get inference type attribute from the operation if available.
inline std::optional<InferenceType> GetInferenceTypeAnnotation(Operation* op) {
auto inference_type = op->getAttrOfType<StringAttr>(kInferenceType);
if (inference_type == nullptr) return std::nullopt;
llvm::StringRef device_name_str = inference_type.getValue();
return GetInferenceTypeEnum(device_name_str);
}
// InferenceDeviceType is a combination of the hardware with inference type.
struct InferenceDeviceType {
std::string hardware;
InferenceType inference_type;
bool operator==(const InferenceDeviceType& other) const {
return (hardware == other.hardware) &&
(inference_type == other.inference_type);
}
bool operator!=(const InferenceDeviceType& other) const {
return !(*this == other);
}
struct inference_device_type_hash {
size_t operator()(const InferenceDeviceType& p) const {
auto hash1 = std::hash<std::string>{}(p.hardware);
auto hash2 = std::hash<InferenceType>{}(p.inference_type);
return hash1 ^ hash2;
}
};
};
// Get InferenceDeviceType attribute from the operation if available.
inline std::optional<InferenceDeviceType> GetInferenceDeviceTypeForOp(
Operation* op) {
auto hardware = GetTargetAnnotation(op);
if (!hardware.has_value()) return std::nullopt;
auto inference_type = GetInferenceTypeAnnotation(op);
if (!inference_type.has_value()) return std::nullopt;
InferenceDeviceType inference_device_type;
inference_device_type.hardware = hardware.value();
inference_device_type.inference_type = inference_type.value();
return inference_device_type;
}
} // namespace tac
} // namespace TFL
} // namespace mlir
#endif // TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_TAC_COMMON_TARGETS_H_
tensorflow/compiler/mlir/lite/experimental/tac/common/utils.cc
+80
View File
@@ -0,0 +1,80 @@
/* 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/compiler/mlir/lite/experimental/tac/common/utils.h"
#include "llvm/Support/Casting.h"
#include "mlir/IR/OpDefinition.h" // from @llvm-project
#include "tensorflow/compiler/mlir/lite/experimental/tac/common/targets.h"
#include "tensorflow/compiler/mlir/lite/ir/tfl_ops.h"
#include "tensorflow/compiler/mlir/lite/utils/utils.h"
namespace mlir {
namespace TFL {
namespace tac {
bool NotTFLQuantDequantizeOp(Operation* op) {
if (!op) return false;
if (llvm::isa<TFL::QuantizeOp, TFL::DequantizeOp>(op)) return false;
return true;
}
bool IsTerminatorOp(Operation* op) {
if (!op) return false;
return op->hasTrait<OpTrait::IsTerminator>();
}
// Try to guess the inference type of the op.
InferenceType GetInferenceType(Operation* op) {
bool float_type_observed = false;
bool int8_type_observed = false;
bool uint8_type_observed = false;
for (auto& input : op->getOpOperands()) {
auto input_type = input.get().getType();
if (IsF32ShapedType(input_type)) {
float_type_observed = true;
} else if (IsQI8Type(input_type)) {
int8_type_observed = true;
} else if (IsQUI8Type(input_type)) {
uint8_type_observed = true;
}
}
// We should not observe both uint8 & int8.
if (int8_type_observed && uint8_type_observed) return UNKNOWN;
if (float_type_observed) {
if (int8_type_observed || uint8_type_observed) {
return HYBRID;
} else {
return FLOAT;
}
}
if (int8_type_observed) {
return QUANTIZED_INT8;
}
if (uint8_type_observed) {
return QUANTIZED_UINT8;
}
// Default to float inference.
return FLOAT;
}
} // namespace tac
} // namespace TFL
} // namespace mlir
tensorflow/compiler/mlir/lite/experimental/tac/common/utils.h
+94
View File
@@ -0,0 +1,94 @@
/* 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_COMPILER_MLIR_LITE_EXPERIMENTAL_TAC_COMMON_UTILS_H_
#define TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_TAC_COMMON_UTILS_H_
#include "llvm/Support/Casting.h"
#include "mlir/Bytecode/BytecodeOpInterface.h" // from @llvm-project
#include "mlir/Dialect/Arith/IR/Arith.h" // from @llvm-project
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/Dialect/Quant/IR/QuantTypes.h" // from @llvm-project
#include "mlir/IR/BuiltinTypes.h" // from @llvm-project
#include "mlir/IR/OpDefinition.h" // from @llvm-project
#include "mlir/IR/Types.h" // from @llvm-project
#include "mlir/Interfaces/CallInterfaces.h" // from @llvm-project
#include "mlir/Interfaces/CastInterfaces.h" // from @llvm-project
#include "mlir/Support/LLVM.h" // from @llvm-project
#include "tensorflow/compiler/mlir/lite/experimental/tac/common/targets.h"
#include "tensorflow/compiler/mlir/lite/ir/tfl_ops.h"
#include "tensorflow/compiler/mlir/lite/utils/utils.h"
namespace mlir {
namespace TFL {
namespace tac {
// Returns true if 'op' is non const op. Returns false otherwise or if
// 'op' is null.
inline bool IsNonConstOp(Operation* op) {
if (!op) return false;
if (llvm::isa<arith::ConstantOp, mlir::func::ConstantOp>(op)) return false;
if (op->hasTrait<OpTrait::ConstantLike>()) return false;
if (llvm::isa<TFL::ConstOp, TFL::QConstOp>(op)) return false;
return true;
}
// Returns true if 'op' is a terminator op, otherwise false.
bool IsTerminatorOp(Operation* op);
// Returns true if 'op' is not TFL Quant / Dequant op. Returns False otherwise
// or if 'op' is null.
bool NotTFLQuantDequantizeOp(Operation* op);
// Returns true if it is a shaped type of f32 elements.
inline bool IsF32ShapedType(Type t) {
if (auto shaped_type = mlir::dyn_cast_or_null<ShapedType>(t)) {
return shaped_type.getElementType().isF32();
}
return false;
}
// Return true when the given element_type is QI8.
inline bool IsQI8Type(Type t) {
auto quantized_type = quant::QuantizedType::getQuantizedElementType(t);
return quantized_type != nullptr &&
quantized_type.getStorageTypeIntegralWidth() == 8 &&
quantized_type.isSigned();
}
// Return true when the given element_type is QUI8.
inline bool IsQUI8Type(Type t) {
auto quantized_type = quant::QuantizedType::getQuantizedElementType(t);
return quantized_type != nullptr &&
quantized_type.getStorageTypeIntegralWidth() == 8 &&
!quantized_type.isSigned();
}
// Return true when the given element_type is QI32.
inline bool IsQI32Type(Type t) {
auto quantized_type = quant::QuantizedType::getQuantizedElementType(t);
return quantized_type != nullptr &&
quantized_type.getStorageTypeIntegralWidth() == 32 &&
quantized_type.isSigned();
}
// Try to guess the inference type of the op.
InferenceType GetInferenceType(Operation* op);
} // namespace tac
} // namespace TFL
} // namespace mlir
#endif // TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_TAC_COMMON_UTILS_H_
tensorflow/compiler/mlir/lite/experimental/tac/examples/BUILD
+31
View File
@@ -0,0 +1,31 @@
load("//tensorflow:tensorflow.bzl", "tf_cc_binary")
load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:LICENSE"],
default_visibility = [
"//visibility:public",
],
licenses = ["notice"], # Apache 2.0
)
cc_library(
name = "example_hardware",
srcs = ["example_hardware.cc"],
hdrs = ["example_hardware.h"],
deps = [
"//tensorflow/compiler/mlir/lite:tensorflow_lite",
"//tensorflow/compiler/mlir/lite/experimental/tac:device_transform_patterns",
"//tensorflow/compiler/mlir/lite/experimental/tac/hardwares:simple_hardware",
],
alwayslink = 1,
)
tf_cc_binary(
name = "example-hardware-translate",
deps = [
":example_hardware",
"//tensorflow/compiler/mlir/lite/experimental/tac:tac-translate-lib",
"//tensorflow/compiler/mlir/lite/experimental/tac/hardwares:all-target-hardwares",
],
)
tensorflow/compiler/mlir/lite/experimental/tac/examples/example_hardware.cc
+47
View File
@@ -0,0 +1,47 @@
/* 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/compiler/mlir/lite/experimental/tac/examples/example_hardware.h"
#include <memory>
#include "tensorflow/compiler/mlir/lite/experimental/tac/transforms/device_transform_patterns.h"
#include "tensorflow/compiler/mlir/lite/ir/tfl_ops.h"
namespace mlir {
namespace TFL {
namespace tac {
constexpr char ExampleHardware::kId[]; // Define kId.
mlir::RewritePatternSet ExampleHardware::GetTransformations(
MLIRContext* context) const {
mlir::RewritePatternSet patterns(context);
patterns.add<LowerPackIntoConcatReshape, UnrollSplit, UnrollSplitV, PadSlice,
PadConcat>(context);
return patterns;
}
std::unique_ptr<TargetHardware> CreateExampleHardware() {
return std::make_unique<ExampleHardware>();
}
TargetHardwareRegistration<ExampleHardware> example_hardware(
"Example device", CreateExampleHardware);
} // namespace tac
} // namespace TFL
} // namespace mlir
tensorflow/compiler/mlir/lite/experimental/tac/examples/example_hardware.h
+45
View File
@@ -0,0 +1,45 @@
/* 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_COMPILER_MLIR_LITE_EXPERIMENTAL_TAC_EXAMPLES_EXAMPLE_HARDWARE_H_
#define TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_TAC_EXAMPLES_EXAMPLE_HARDWARE_H_
#include "tensorflow/compiler/mlir/lite/experimental/tac/hardwares/simple_hardware.h"
namespace mlir {
namespace TFL {
namespace tac {
class ExampleHardware : public SimpleHardware {
public:
static constexpr char kId[] = "ExampleHardware";
mlir::RewritePatternSet GetTransformations(
MLIRContext* context) const override;
mlir::TypeID GetTypeId() const override {
return mlir::TypeID::get<ExampleHardware>();
}
bool IsNotSupportedOp(mlir::Operation* op) const override { return false; }
float AdvantageOverCPU() const override { return 5.0; }
};
} // namespace tac
} // namespace TFL
} // namespace mlir
#endif // TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_TAC_EXAMPLES_EXAMPLE_HARDWARE_H_
tensorflow/compiler/mlir/lite/experimental/tac/execution_metadata_exporter.cc
+206
View File
@@ -0,0 +1,206 @@
// 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/experimental/tac/execution_metadata_exporter.h"
#include <cstdint>
#include <map>
#include <optional>
#include <string>
#include <utility>
#include <vector>
#include "flatbuffers/buffer.h" // from @flatbuffers
#include "flatbuffers/flatbuffer_builder.h" // from @flatbuffers
#include "flatbuffers/string.h" // from @flatbuffers
#include "flatbuffers/vector.h" // from @flatbuffers
#include "llvm/Support/Casting.h"
#include "mlir/Dialect/Arith/IR/Arith.h" // from @llvm-project
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/Attributes.h" // from @llvm-project
#include "mlir/IR/BuiltinAttributes.h" // from @llvm-project
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/IR/Operation.h" // from @llvm-project
#include "mlir/IR/Region.h" // from @llvm-project
#include "mlir/Support/LLVM.h" // from @llvm-project
#include "tensorflow/compiler/mlir/lite/experimental/tac/common/targets.h"
#include "tensorflow/compiler/mlir/lite/experimental/tac/hardwares/target_hardware.h"
#include "tensorflow/compiler/mlir/lite/experimental/tac/runtime_metadata_generated.h"
#include "tensorflow/compiler/mlir/lite/ir/tfl_ops.h"
#include "tensorflow/compiler/mlir/lite/quantization/ir/QuantOps.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
namespace tflite {
namespace {
bool IsConst(mlir::Operation* op) {
return llvm::isa<mlir::arith::ConstantOp, mlir::TF::ConstOp,
mlir::TFL::ConstOp, mlir::TFL::QConstOp>(op);
}
bool IsOpSupported(mlir::Operation* op, const std::string& hardware) {
auto* devce_hardware = mlir::TFL::tac::GetTargetHardware(hardware);
if (devce_hardware == nullptr) return {};
return devce_hardware->IsOpSupported(op);
}
bool HasValidHardwareTarget(mlir::Operation* op) {
// All TFLite ops has CPU interface, should be enough to check for cpu.
return IsOpSupported(op, "CPU");
}
std::optional<std::string> GetDeviceName(mlir::Operation* op) {
if (IsConst(op)) return std::nullopt;
// The model may contain quant stats op which is unrelevant to the
// execution.
if (llvm::isa<mlir::func::ReturnOp, mlir::quantfork::StatisticsOp>(op))
return std::nullopt;
if (!HasValidHardwareTarget(op)) return std::nullopt;
auto device = op->getAttrOfType<mlir::StringAttr>(mlir::TFL::tac::kDevice);
if (device == nullptr) return std::nullopt;
llvm::StringRef device_name_str = device.getValue();
return device_name_str.str();
}
std::optional<std::vector<float>> GetPerDeviceCosts(
const std::map<std::string, uint8_t>& hardware_map, mlir::Operation* op) {
auto device_costs_attr =
op->getAttrOfType<mlir::DictionaryAttr>("per_device_costs");
if (device_costs_attr == nullptr) return std::nullopt;
std::vector<float> device_costs(hardware_map.size(), -1.f);
for (const auto& kv : hardware_map) {
auto cost_attr = device_costs_attr.getNamed(kv.first);
if (!cost_attr.has_value()) return std::nullopt;
float cost = mlir::dyn_cast_or_null<mlir::FloatAttr>(cost_attr->getValue())
.getValueAsDouble();
device_costs[kv.second] = cost;
}
return device_costs;
}
flatbuffers::Offset<SubgraphMetadata> CreateSubgraphMetadata(
const std::map<std::string, uint8_t>& hardware_map, mlir::Region* Region,
flatbuffers::FlatBufferBuilder* builder) {
auto& block = Region->front();
int index = 0;
std::vector<flatbuffers::Offset<tflite::OpMetadata>> ops;
for (auto& inst : block) {
// Const nodes are mapped to const vectors in flatbuffer, so skip.
if (IsConst(&inst)) continue;
// The model may contain quant stats op which is unrelevant to the
// execution.
if (llvm::isa<mlir::func::ReturnOp, mlir::quantfork::StatisticsOp>(&inst))
continue;
// If an op doesn't implement any of the hardware interface we skip it.
// This can happen in cases like Flex when we have non TFLite ops.
auto device_name = GetDeviceName(&inst);
if (device_name.has_value()) {
// Add per device costs if present.
auto per_device_cost = GetPerDeviceCosts(hardware_map, &inst);
flatbuffers::Offset<flatbuffers::Vector<float>> per_device_cost_offset;
if (per_device_cost.has_value()) {
per_device_cost_offset = builder->CreateVector(*per_device_cost);
}
OpMetadataBuilder op_builder(*builder);
op_builder.add_index(index);
uint8_t hardware = hardware_map.at(*device_name);
op_builder.add_hardware(hardware);
if (per_device_cost.has_value()) {
op_builder.add_op_costs(per_device_cost_offset);
}
ops.push_back(op_builder.Finish());
}
index++;
}
return CreateSubgraphMetadata(*builder, builder->CreateVector(ops));
}
flatbuffers::Offset<tflite::HardwareMetadata>
CreateHardwareMetadataAndPopulateLookupTable(
std::vector<mlir::func::FuncOp>* funcs,
flatbuffers::FlatBufferBuilder* builder,
std::map<std::string, uint8_t>* hardware_names) {
uint8_t index = 0;
for (auto& func : *funcs) {
func.walk([&hardware_names, &index](mlir::Operation* op) {
auto device_name = GetDeviceName(op);
if (!device_name.has_value()) return;
auto iter = hardware_names->find(*device_name);
if (iter == hardware_names->end()) {
hardware_names->insert({*device_name, index++});
}
});
}
// Build the flatbuffer.
std::vector<flatbuffers::Offset<flatbuffers::String>> hardwares;
for (const auto& kv : *hardware_names) {
hardwares.push_back(builder->CreateString(kv.first));
}
return CreateHardwareMetadata(*builder, builder->CreateVector(hardwares));
}
} // namespace
std::optional<std::string> ExportRuntimeMetadata(mlir::ModuleOp module) {
mlir::func::FuncOp main_fn = module.lookupSymbol<mlir::func::FuncOp>("main");
if (!main_fn) return std::string("");
flatbuffers::FlatBufferBuilder fb_builder;
std::vector<mlir::func::FuncOp> funcs;
funcs.push_back(main_fn);
module.walk([&](mlir::func::FuncOp fn) {
if (fn != main_fn) {
funcs.push_back(fn);
}
});
// Populate the hardware metadata.
// And collect the hardwares used.
std::map<std::string, uint8_t> hardware_map;
flatbuffers::Offset<tflite::HardwareMetadata> hardware_metadata_offset =
CreateHardwareMetadataAndPopulateLookupTable(&funcs, &fb_builder,
&hardware_map);
// Populate the runtime metadata.
std::vector<flatbuffers::Offset<SubgraphMetadata>> subgraphs_metadata;
subgraphs_metadata.reserve(funcs.size());
for (auto& func : funcs) {
subgraphs_metadata.push_back(
CreateSubgraphMetadata(hardware_map, &func.getBody(), &fb_builder));
}
auto runtime_metadata =
CreateRuntimeMetadata(fb_builder, hardware_metadata_offset,
fb_builder.CreateVector(subgraphs_metadata));
fb_builder.Finish(runtime_metadata);
return std::string(
reinterpret_cast<const char*>(fb_builder.GetBufferPointer()),
fb_builder.GetSize());
}
} // namespace tflite
tensorflow/compiler/mlir/lite/experimental/tac/execution_metadata_exporter.h
+31
View File
@@ -0,0 +1,31 @@
// Copyright 2020 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#ifndef TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_TAC_EXECUTION_METADATA_EXPORTER_H_
#define TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_TAC_EXECUTION_METADATA_EXPORTER_H_
#include <optional>
#include <string>
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
namespace tflite {
// Returns serialized string for the generated flatbuffer.
std::optional<std::string> ExportRuntimeMetadata(mlir::ModuleOp module);
} // namespace tflite
#endif // TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_TAC_EXECUTION_METADATA_EXPORTER_H_
tensorflow/compiler/mlir/lite/experimental/tac/execution_metadata_exporter_test.cc
+128
View File
@@ -0,0 +1,128 @@
// 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/experimental/tac/execution_metadata_exporter.h"
#include <string>
#include <vector>
#include <gtest/gtest.h>
#include "flatbuffers/buffer.h" // from @flatbuffers
#include "flatbuffers/flatbuffer_builder.h" // from @flatbuffers
#include "flatbuffers/string.h" // from @flatbuffers
#include "mlir/Dialect/Arith/IR/Arith.h" // from @llvm-project
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/IR/DialectRegistry.h" // from @llvm-project
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "mlir/IR/OwningOpRef.h" // from @llvm-project
#include "mlir/Parser/Parser.h" // from @llvm-project
#include "tensorflow/compiler/mlir/lite/experimental/tac/runtime_metadata_generated.h"
#include "tensorflow/compiler/mlir/lite/ir/tfl_ops.h"
namespace tflite {
std::string CreateRuntimeMetadata() {
flatbuffers::FlatBufferBuilder fb_builder;
std::vector<flatbuffers::Offset<flatbuffers::String>> device_names = {
fb_builder.CreateString("GPU"), fb_builder.CreateString("CPU")};
const auto hardwares =
CreateHardwareMetadata(fb_builder, fb_builder.CreateVector(device_names));
const auto ops = {
CreateOpMetadata(fb_builder, 0, 0,
fb_builder.CreateVector(std::vector<float>({1.0, 5.0}))),
CreateOpMetadata(fb_builder, 1, 0,
fb_builder.CreateVector(std::vector<float>({1.0, 5.0}))),
CreateOpMetadata(fb_builder, 2, 0,
fb_builder.CreateVector(std::vector<float>({1.0, 5.0}))),
CreateOpMetadata(
fb_builder, 3, 1,
fb_builder.CreateVector(std::vector<float>({-1.0, 2.0}))),
};
const auto subgraphs = {CreateSubgraphMetadata(
fb_builder, fb_builder.CreateVector(ops.begin(), ops.size()))};
const auto metadata = CreateRuntimeMetadata(
fb_builder, hardwares,
fb_builder.CreateVector(subgraphs.begin(), subgraphs.size()));
fb_builder.Finish(metadata);
return std::string(
reinterpret_cast<const char*>(fb_builder.GetBufferPointer()),
fb_builder.GetSize());
}
void Verify(const RuntimeMetadata* result, const RuntimeMetadata* expected) {
EXPECT_EQ(result->subgraph_metadata()->size(),
expected->subgraph_metadata()->size());
for (int i = 0; i < result->subgraph_metadata()->size(); ++i) {
auto result_subgraph_metadata =
result->subgraph_metadata()->GetAs<SubgraphMetadata>(i);
auto expected_subgraph_metadata =
expected->subgraph_metadata()->GetAs<SubgraphMetadata>(i);
if (expected_subgraph_metadata->op_metadata() == nullptr &&
result_subgraph_metadata->op_metadata() == nullptr) {
return;
}
ASSERT_EQ(expected_subgraph_metadata->op_metadata()->size(),
result_subgraph_metadata->op_metadata()->size());
for (int j = 0; j < expected_subgraph_metadata->op_metadata()->size();
++j) {
auto result_op_metadata =
result_subgraph_metadata->op_metadata()->GetAs<OpMetadata>(j);
auto expected_op_metadata =
expected_subgraph_metadata->op_metadata()->GetAs<OpMetadata>(j);
EXPECT_EQ(result_op_metadata->index(), expected_op_metadata->index());
EXPECT_EQ(result_op_metadata->hardware(),
expected_op_metadata->hardware());
EXPECT_EQ(result_op_metadata->op_costs()->size(),
expected_op_metadata->op_costs()->size());
for (int i = 0; i < result_op_metadata->op_costs()->size(); ++i) {
EXPECT_FLOAT_EQ(result_op_metadata->op_costs()->Get(i),
expected_op_metadata->op_costs()->Get(i));
}
}
}
}
TEST(ExporterTest, Valid) {
const std::string kMLIR = R"(
func.func @main(%arg0: tensor<1xf32>, %arg1: tensor<1xf32>, %arg2: tensor<1xf32>, %arg3: tensor<1xf32>) -> tensor<2x1xf32> {
%0 = "tfl.add"(%arg0, %arg1) {fused_activation_function = "RELU6", per_device_costs = {CPU = 5.0 : f32, GPU = 1.0 : f32}, tac.device = "GPU"} : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32>
%1 = "tfl.mul"(%0, %arg2) {fused_activation_function = "RELU6", per_device_costs = {CPU = 5.0 : f32, GPU = 1.0 : f32}, tac.device = "GPU"} : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32>
%2 = "tfl.add"(%arg0, %arg3) {fused_activation_function = "RELU6", per_device_costs = {CPU = 5.0 : f32, GPU = 1.0 : f32}, tac.device = "GPU"} : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32>
%3 = "tfl.pack"(%1, %2) {axis = 0 : i32, per_device_costs = {CPU = 2.0 : f32, GPU = -1.0 : f32}, values_count = 2 : i32, tac.device = "CPU"} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2x1xf32>
func.return %3 : tensor<2x1xf32>
})";
const std::string kExpectedFB = CreateRuntimeMetadata();
mlir::DialectRegistry registry;
registry.insert<mlir::TFL::TensorFlowLiteDialect, mlir::arith::ArithDialect,
mlir::func::FuncDialect>();
mlir::MLIRContext context(registry);
auto module = mlir::OwningOpRef<mlir::ModuleOp>(
mlir::parseSourceString<mlir::ModuleOp>(kMLIR, &context));
auto module_op = module.get();
auto serialized_result_fb = ExportRuntimeMetadata(module_op);
const auto* result = GetRuntimeMetadata(serialized_result_fb.value().c_str());
const auto* expected = GetRuntimeMetadata(kExpectedFB.c_str());
ASSERT_TRUE(result != nullptr);
ASSERT_TRUE(result->subgraph_metadata() != nullptr);
ASSERT_TRUE(expected->subgraph_metadata() != nullptr);
Verify(result, expected);
}
} // namespace tflite
tensorflow/compiler/mlir/lite/experimental/tac/g3doc/images/compute_cost.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 38 KiB

tensorflow/compiler/mlir/lite/experimental/tac/g3doc/images/final_graph.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 29 KiB

tensorflow/compiler/mlir/lite/experimental/tac/g3doc/images/pick_subgraphs.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 41 KiB

tensorflow/compiler/mlir/lite/experimental/tac/g3doc/images/target_annotation.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 38 KiB

tensorflow/compiler/mlir/lite/experimental/tac/g3doc/images/target_optimization.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 38 KiB

tensorflow/compiler/mlir/lite/experimental/tac/hardwares/BUILD
+89
View File
@@ -0,0 +1,89 @@
load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:LICENSE"],
default_visibility = [
"//visibility:public",
],
licenses = ["notice"],
)
cc_library(
name = "gpu_hardware",
srcs = ["gpu_hardware.cc"],
hdrs = ["gpu_hardware.h"],
deps = [
":target_hardware",
"//tensorflow/compiler/mlir/lite:cost_estimators",
"//tensorflow/compiler/mlir/lite:tensorflow_lite",
"//tensorflow/compiler/mlir/lite/experimental/tac:common",
"//tensorflow/compiler/mlir/lite/experimental/tac:device_transform_patterns",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Support",
],
alwayslink = 1,
)
cc_library(
name = "cpu_hardware",
srcs = ["cpu_hardware.cc"],
deps = [
":target_hardware",
"//tensorflow/compiler/mlir/lite:cost_estimators",
"//tensorflow/compiler/mlir/lite:tensorflow_lite",
"//tensorflow/compiler/mlir/lite/experimental/tac:common",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Support",
],
alwayslink = 1,
)
cc_library(
name = "nnapi_hardware",
srcs = ["nnapi_hardware.cc"],
hdrs = ["nnapi_hardware.h"],
deps = [
":target_hardware",
"//tensorflow/compiler/mlir/lite:cost_estimators",
"//tensorflow/compiler/mlir/lite:tensorflow_lite",
"//tensorflow/compiler/mlir/lite/experimental/tac:device_transform_patterns",
"//tensorflow/compiler/mlir/lite/experimental/tac/hardwares:simple_hardware",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Support",
],
alwayslink = 1,
)
cc_library(
name = "target_hardware",
srcs = ["target_hardware.cc"],
hdrs = ["target_hardware.h"],
deps = [
"//tensorflow/compiler/mlir/lite:tensorflow_lite",
"//tensorflow/compiler/mlir/lite/experimental/tac:common",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Support",
],
)
cc_library(
name = "simple_hardware",
srcs = ["simple_hardware.cc"],
hdrs = ["simple_hardware.h"],
deps = [
":target_hardware",
"@llvm-project//mlir:IR",
],
)
cc_library(
name = "all-target-hardwares",
deps = [
":cpu_hardware",
":gpu_hardware",
":nnapi_hardware",
],
alwayslink = 1,
)
tensorflow/compiler/mlir/lite/experimental/tac/hardwares/cpu_hardware.cc
+181
View File
@@ -0,0 +1,181 @@
/* 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 <cstdint>
#include <memory>
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "mlir/IR/Operation.h" // from @llvm-project
#include "mlir/IR/PatternMatch.h" // from @llvm-project
#include "mlir/Support/TypeID.h" // from @llvm-project
#include "tensorflow/compiler/mlir/lite/experimental/tac/common/targets.h"
#include "tensorflow/compiler/mlir/lite/experimental/tac/common/utils.h"
#include "tensorflow/compiler/mlir/lite/experimental/tac/hardwares/target_hardware.h"
#include "tensorflow/compiler/mlir/lite/ir/tfl_ops.h"
#include "tensorflow/compiler/mlir/lite/utils/arithmetic_count_util.h"
namespace mlir {
namespace TFL {
namespace tac {
namespace {
// CPU
constexpr float kCPUArithmeticUnitCost = 1.0;
// This basically assumes pure load/store. This is just fake data.
constexpr float kCPUCopyUnitCost = 0.5;
// Default values.
constexpr float kCPUDefaultFixedValuedCost = 10000.0;
// Quantized inference cost efficiency.
// For CPU, quantized inference is ~3x faster than the float alternative, this
// is just an estimation.
constexpr float kQuantizedInferenceEfficiency = 0.3;
inline float InferenceTypeEfficiency(InferenceType inference_type) {
if (inference_type == QUANTIZED_INT8 || inference_type == QUANTIZED_UINT8) {
return kQuantizedInferenceEfficiency;
}
return 1.0;
}
// CPU hardware class which handles CPU capabilities in TFLite.
// This is used by TAC to get op supported/ op cost estimates on CPU.
class CpuHardware : public TargetHardware {
public:
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(CpuHardware)
// String Identifier for CPU hardware.
static constexpr char kId[] = "CPU";
double GetHardwareSwitchingCost(const TargetHardware* from,
size_t buffer_size) const override {
auto from_type = from->GetTypeId();
auto to_type = GetTypeId();
if (from_type == to_type) return 0.0f;
// TODO(renjieliu): Implement a better version for different hardware cases.
return buffer_size * kCrossHardwareTransferPerByteCost / 8.0 +
kCrossHardwareTransferFixedCost;
}
mlir::RewritePatternSet GetTransformations(
MLIRContext* context) const override {
return {context};
}
mlir::TypeID GetTypeId() const override {
return mlir::TypeID::get<CpuHardware>();
}
bool IsOpSupported(mlir::Operation* op) const override {
// All ops in TFL dialect are supported on CPU.
if (op->getDialect() == nullptr) return false;
if (op->getDialect()->getNamespace() != "tfl") return false;
return true;
}
};
constexpr char CpuHardware::kId[]; // Define kId.
std::unique_ptr<TargetHardware> CreateCpuHardware() {
return std::make_unique<CpuHardware>();
}
TargetHardwareRegistration<CpuHardware> cpu_hardware("Target device for CPU",
CreateCpuHardware);
#define TAC_REGISTER_CPU_OP(Op, Create) \
TargetHardwareOpRegistration<CpuHardware, Op> Op##_CpuHardware_hardware( \
Create);
// Operation costs on CPU
// Currently used for these ops:
// tfl.conv_2d / tfl.depthwise_conv_2d / tfl.fully_connected
class CpuConvOp : public TargetHardwareOperation {
double GetOpCost(mlir::Operation* op) const override {
float cost = 0.0;
int64_t arithmetic_count;
if (ArithmeticCountUtilHelper::GetArithmeticCountForConvAndFullyconnectedOp(
op, &arithmetic_count)) {
cost = arithmetic_count * kCPUArithmeticUnitCost;
} else {
cost = kCPUDefaultFixedValuedCost;
}
return cost * InferenceTypeEfficiency(GetInferenceType(op));
}
bool IsOpSupported(mlir::Operation* op) const override { return true; }
};
std::unique_ptr<TargetHardwareOperation> CreateConvOp() {
return std::make_unique<CpuConvOp>();
}
// Currently used for these ops:
// tfl.Add / tfl.mul
class CpuArithmeticOp : public TargetHardwareOperation {
double GetOpCost(mlir::Operation* op) const override {
float cost = 0.0;
int64_t count;
if (ArithmeticCountUtilHelper::GetFirstOutputCount(op, &count)) {
cost = kCPUArithmeticUnitCost * count;
} else {
cost = kCPUDefaultFixedValuedCost;
}
return cost * InferenceTypeEfficiency(GetInferenceType(op));
}
bool IsOpSupported(mlir::Operation* op) const override { return true; }
};
std::unique_ptr<TargetHardwareOperation> CreateArithmeticOp() {
return std::make_unique<CpuArithmeticOp>();
}
// Currently used for these ops:
// tfl.concatenation / tfl.reshape / tfl.pack
class CpuConcatOp : public TargetHardwareOperation {
double GetOpCost(mlir::Operation* op) const override {
float cost = 0.0;
int64_t count;
if (ArithmeticCountUtilHelper::GetInputTensorTotalSize(op, &count)) {
cost = kCPUCopyUnitCost * count;
} else {
cost = kCPUDefaultFixedValuedCost;
}
return cost * InferenceTypeEfficiency(GetInferenceType(op));
}
bool IsOpSupported(mlir::Operation* op) const override { return true; }
};
std::unique_ptr<TargetHardwareOperation> CreateConcatOp() {
return std::make_unique<CpuConcatOp>();
}
TAC_REGISTER_CPU_OP(Conv2DOp, CreateConvOp);
TAC_REGISTER_CPU_OP(DepthwiseConv2DOp, CreateConvOp);
TAC_REGISTER_CPU_OP(FullyConnectedOp, CreateConvOp);
TAC_REGISTER_CPU_OP(AddOp, CreateArithmeticOp);
TAC_REGISTER_CPU_OP(MulOp, CreateArithmeticOp);
TAC_REGISTER_CPU_OP(ConcatenationOp, CreateConcatOp);
TAC_REGISTER_CPU_OP(ReshapeOp, CreateConcatOp);
TAC_REGISTER_CPU_OP(PackOp, CreateConcatOp);
#undef TAC_REGISTER_CPU_OP
} // namespace
} // namespace tac
} // namespace TFL
} // namespace mlir
tensorflow/compiler/mlir/lite/experimental/tac/hardwares/gpu_hardware.cc
+229
View File
@@ -0,0 +1,229 @@
/* 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/compiler/mlir/lite/experimental/tac/hardwares/gpu_hardware.h"
#include <cstddef>
#include <cstdint>
#include <memory>
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "mlir/IR/Operation.h" // from @llvm-project
#include "mlir/IR/PatternMatch.h" // from @llvm-project
#include "tensorflow/compiler/mlir/lite/experimental/tac/common/targets.h"
#include "tensorflow/compiler/mlir/lite/experimental/tac/common/utils.h"
#include "tensorflow/compiler/mlir/lite/experimental/tac/hardwares/target_hardware.h"
#include "tensorflow/compiler/mlir/lite/experimental/tac/transforms/device_transform_patterns.h"
#include "tensorflow/compiler/mlir/lite/ir/tfl_ops.h"
#include "tensorflow/compiler/mlir/lite/utils/arithmetic_count_util.h"
namespace mlir {
namespace TFL {
namespace tac {
namespace {
#include "tensorflow/compiler/mlir/lite/experimental/tac/transforms/generated_transform_patterns.inc"
} // namespace
constexpr char GpuHardware::kId[]; // Define kId.
mlir::RewritePatternSet GpuHardware::GetTransformations(
MLIRContext* context) const {
mlir::RewritePatternSet patterns(context);
patterns.add<LowerPackIntoConcatReshape, UnrollSplit, UnrollSplitV, SubToAdd,
EnsureBiasForConv2d, PadSlice, FullyConnectedToConv, PadConcat,
SquaredDifference>(context);
return patterns;
}
double GpuHardware::GetHardwareSwitchingCost(const TargetHardware* from,
size_t buffer_size) const {
auto from_type = from->GetTypeId();
auto to_type = GetTypeId();
if (from_type == to_type) return 0.0f;
// TODO(renjieliu): Implement a better version for different hardware cases.
return buffer_size * kCrossHardwareTransferPerByteCost / 8.0 +
kCrossHardwareTransferFixedCost;
}
bool GpuHardware::IsOpSupported(mlir::Operation* op) const {
if (TargetHardware::IsOpSupported(op)) {
return true;
}
// We also support quantized ops.
return !NotTFLQuantDequantizeOp(op);
}
namespace {
// GPU
constexpr float kGPUArithmeticUnitCost = 0.2;
// The copy can be non-consectutive copy. This is just fake data.
constexpr float kGPUCopyUnitCost = 0.2;
// Default values.
constexpr float kGPUDefaultFixedValuedCost = 10000.0;
std::unique_ptr<TargetHardware> CreateGpuHardware() {
return std::make_unique<GpuHardware>();
}
TargetHardwareRegistration<GpuHardware> gpu_hardware("Target device for GPU",
CreateGpuHardware);
#define TAC_REGISTER_GPU_OP(Op, Create) \
TargetHardwareOpRegistration<GpuHardware, Op> Op##_GpuHardware_hardware( \
Create);
// Currently used for these ops:
// tfl.Abs / tfl.Average_pool_2d / tfl.Cos / tfl.div / tfl.exp / tfl.hardswish /
// tfl.log / tfl.logistic / tfl.max_pool_2d / tfl.mirror_pad / tfl.maximum /
// tfl.custom / tfl.mean / tfl.minimum / tfl.pad / tfl.pow / tfl.prelu /
// tfl.relu / tfl.relu6 / tfl.rsqrt / tfl.sin / tfl.slice / tfl.softmax /
// tfl.space_to_depth / tfl.sqrt / tfl.square / tfl.squared_difference /
// tfl.strided_slice / tfl.tanh / tfl.transpose / tfl.transpose_conv
class GpuBasicSupportedOpNoCost : public TargetHardwareOperation {
double GetOpCost(mlir::Operation* op) const override { return 0; }
bool IsOpSupported(mlir::Operation* op) const override {
InferenceType inference_type = GetInferenceType(op);
if (inference_type != FLOAT) {
return false;
}
return true;
}
};
std::unique_ptr<TargetHardwareOperation> CreateBasicOpNoCost() {
return std::make_unique<GpuBasicSupportedOpNoCost>();
}
// Currently used for these ops:
// tfl.Add / tfl.mul
class GpuArithmeticOp : public TargetHardwareOperation {
double GetOpCost(mlir::Operation* op) const override {
int64_t count;
if (ArithmeticCountUtilHelper::GetFirstOutputCount(op, &count))
return kGPUArithmeticUnitCost * count;
return kGPUDefaultFixedValuedCost;
}
bool IsOpSupported(mlir::Operation* op) const override {
InferenceType inference_type = GetInferenceType(op);
if (inference_type != FLOAT) {
return false;
}
return true;
}
};
std::unique_ptr<TargetHardwareOperation> CreateArithmeticOp() {
return std::make_unique<GpuArithmeticOp>();
}
// Currently used for these ops:
// tfl.concatenation / tfl.reshape
class GpuConcatOp : public TargetHardwareOperation {
double GetOpCost(mlir::Operation* op) const override {
int64_t count;
if (ArithmeticCountUtilHelper::GetInputTensorTotalSize(op, &count))
return kGPUCopyUnitCost * count;
return kGPUDefaultFixedValuedCost;
}
bool IsOpSupported(mlir::Operation* op) const override {
InferenceType inference_type = GetInferenceType(op);
if (inference_type != FLOAT) {
return false;
}
return true;
}
};
std::unique_ptr<TargetHardwareOperation> CreateConcatOp() {
return std::make_unique<GpuConcatOp>();
}
// Currently used for these ops:
// tfl.conv_2d / tfl.depthwise_conv_2d / tfl.fully_connected
class GpuConvOp : public TargetHardwareOperation {
double GetOpCost(mlir::Operation* op) const override {
int64_t arithmetic_count;
if (ArithmeticCountUtilHelper::GetArithmeticCountForConvAndFullyconnectedOp(
op, &arithmetic_count)) {
return arithmetic_count * kGPUArithmeticUnitCost;
}
return kGPUDefaultFixedValuedCost;
}
bool IsOpSupported(mlir::Operation* op) const override {
InferenceType inference_type = GetInferenceType(op);
if (inference_type != FLOAT) {
return false;
}
return true;
}
};
std::unique_ptr<TargetHardwareOperation> CreateConvOp() {
return std::make_unique<GpuConvOp>();
}
// Op registrations
TAC_REGISTER_GPU_OP(AbsOp, CreateBasicOpNoCost);
TAC_REGISTER_GPU_OP(AveragePool2DOp, CreateBasicOpNoCost);
TAC_REGISTER_GPU_OP(CosOp, CreateBasicOpNoCost);
TAC_REGISTER_GPU_OP(DivOp, CreateBasicOpNoCost);
TAC_REGISTER_GPU_OP(ExpOp, CreateBasicOpNoCost);
TAC_REGISTER_GPU_OP(HardSwishOp, CreateBasicOpNoCost);
TAC_REGISTER_GPU_OP(LogOp, CreateBasicOpNoCost);
TAC_REGISTER_GPU_OP(LogisticOp, CreateBasicOpNoCost);
TAC_REGISTER_GPU_OP(MaxPool2DOp, CreateBasicOpNoCost);
TAC_REGISTER_GPU_OP(MirrorPadOp, CreateBasicOpNoCost);
TAC_REGISTER_GPU_OP(MaximumOp, CreateBasicOpNoCost);
TAC_REGISTER_GPU_OP(MinimumOp, CreateBasicOpNoCost);
TAC_REGISTER_GPU_OP(MeanOp, CreateBasicOpNoCost);
TAC_REGISTER_GPU_OP(CustomOp, CreateBasicOpNoCost);
TAC_REGISTER_GPU_OP(PadOp, CreateBasicOpNoCost);
TAC_REGISTER_GPU_OP(PowOp, CreateBasicOpNoCost);
TAC_REGISTER_GPU_OP(PReluOp, CreateBasicOpNoCost);
TAC_REGISTER_GPU_OP(ReluOp, CreateBasicOpNoCost);
TAC_REGISTER_GPU_OP(Relu6Op, CreateBasicOpNoCost);
TAC_REGISTER_GPU_OP(RsqrtOp, CreateBasicOpNoCost);
TAC_REGISTER_GPU_OP(SinOp, CreateBasicOpNoCost);
TAC_REGISTER_GPU_OP(SliceOp, CreateBasicOpNoCost);
TAC_REGISTER_GPU_OP(SoftmaxOp, CreateBasicOpNoCost);
TAC_REGISTER_GPU_OP(SpaceToDepthOp, CreateBasicOpNoCost);
TAC_REGISTER_GPU_OP(SqrtOp, CreateBasicOpNoCost);
TAC_REGISTER_GPU_OP(SquareOp, CreateBasicOpNoCost);
TAC_REGISTER_GPU_OP(SquaredDifferenceOp, CreateBasicOpNoCost);
TAC_REGISTER_GPU_OP(StridedSliceOp, CreateBasicOpNoCost);
TAC_REGISTER_GPU_OP(TanhOp, CreateBasicOpNoCost);
TAC_REGISTER_GPU_OP(TransposeOp, CreateBasicOpNoCost);
TAC_REGISTER_GPU_OP(TransposeConvOp, CreateBasicOpNoCost);
TAC_REGISTER_GPU_OP(ConcatenationOp, CreateConcatOp);
TAC_REGISTER_GPU_OP(ReshapeOp, CreateConcatOp);
TAC_REGISTER_GPU_OP(Conv2DOp, CreateConvOp);
TAC_REGISTER_GPU_OP(DepthwiseConv2DOp, CreateConvOp);
TAC_REGISTER_GPU_OP(FullyConnectedOp, CreateConvOp);
TAC_REGISTER_GPU_OP(AddOp, CreateArithmeticOp);
TAC_REGISTER_GPU_OP(MulOp, CreateArithmeticOp);
#undef TAC_REGISTER_GPU_OP
} // namespace
} // namespace tac
} // namespace TFL
} // namespace mlir
tensorflow/compiler/mlir/lite/experimental/tac/hardwares/gpu_hardware.h
+51
View File
@@ -0,0 +1,51 @@
/* 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_COMPILER_MLIR_LITE_EXPERIMENTAL_TAC_HARDWARES_GPU_HARDWARE_H_
#define TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_TAC_HARDWARES_GPU_HARDWARE_H_
#include <cstddef>
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "mlir/IR/PatternMatch.h" // from @llvm-project
#include "mlir/Support/TypeID.h" // from @llvm-project
#include "tensorflow/compiler/mlir/lite/experimental/tac/hardwares/target_hardware.h"
#include "tensorflow/compiler/mlir/lite/ir/tfl_ops.h"
namespace mlir {
namespace TFL {
namespace tac {
// Gpu hardware class which handles GPU capabilities in TFLite.
// This is used by TAC to get op supported/ op cost estimates on GPU.
class GpuHardware : public TargetHardware {
public:
static constexpr char kId[] = "GPU";
mlir::RewritePatternSet GetTransformations(
MLIRContext* context) const override;
mlir::TypeID GetTypeId() const override {
return mlir::TypeID::get<GpuHardware>();
}
double GetHardwareSwitchingCost(const TargetHardware* from,
size_t buffer_size) const override;
bool IsOpSupported(mlir::Operation* op) const override;
};
} // namespace tac
} // namespace TFL
} // namespace mlir
#endif // TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_TAC_HARDWARES_GPU_HARDWARE_H_
tensorflow/compiler/mlir/lite/experimental/tac/hardwares/nnapi_hardware.cc
+100
View File
@@ -0,0 +1,100 @@
/* 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/compiler/mlir/lite/experimental/tac/hardwares/nnapi_hardware.h"
#include <cstdint>
#include <memory>
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "mlir/IR/Operation.h" // from @llvm-project
#include "mlir/IR/PatternMatch.h" // from @llvm-project
#include "tensorflow/compiler/mlir/lite/experimental/tac/hardwares/target_hardware.h"
#include "tensorflow/compiler/mlir/lite/experimental/tac/transforms/device_transform_patterns.h"
#include "tensorflow/compiler/mlir/lite/ir/tfl_ops.h"
#include "tensorflow/compiler/mlir/lite/utils/arithmetic_count_util.h"
namespace mlir {
namespace TFL {
namespace tac {
// The copy can be non-consectutive copy. This is just fake data.
constexpr float kNNAPICopyUnitCost = 0.2;
// Default values.
constexpr float kNNAPIDefaultFixedValuedCost = 10000.0;
constexpr char NNAPIHardware::kId[]; // Define kId.
mlir::RewritePatternSet NNAPIHardware::GetTransformations(
MLIRContext* context) const {
mlir::RewritePatternSet patterns(context);
patterns.add<SquaredDifference, LowerPackIntoConcatReshape,
ReduceMeanToAvgPool, InsertRequantForReduceMean>(context);
return patterns;
}
std::unique_ptr<TargetHardware> CreateNNAPIHardware() {
return std::make_unique<NNAPIHardware>();
}
TargetHardwareRegistration<NNAPIHardware> nnapi_hardware(
"Target device for NNAPI", CreateNNAPIHardware);
// Currently used for these ops:
// tfl.squared_difference
class NNAPIBasicSupportedOpNoCost : public TargetHardwareOperation {
double GetOpCost(mlir::Operation* op) const override { return 0; }
bool IsOpSupported(mlir::Operation* op) const override {
return true;
}
};
std::unique_ptr<TargetHardwareOperation> CreateBasicOpNoCost() {
return std::make_unique<NNAPIBasicSupportedOpNoCost>();
}
// Currently used for these ops:
// tfl.concatenation / tfl.reshape / tfl.pack
class NNAPIConcatOp : public TargetHardwareOperation {
double GetOpCost(mlir::Operation* op) const override {
int64_t count;
if (ArithmeticCountUtilHelper::GetInputTensorTotalSize(op, &count))
return kNNAPICopyUnitCost * count;
return kNNAPIDefaultFixedValuedCost;
}
bool IsOpSupported(mlir::Operation* op) const override { return true; }
};
std::unique_ptr<TargetHardwareOperation> CreateConcatOp() {
return std::make_unique<NNAPIConcatOp>();
}
#define TAC_REGISTER_NNAPI_OP(Op, Create) \
TargetHardwareOpRegistration<NNAPIHardware, Op> Op##_NNAPIHardware_hardware( \
Create);
// Op registeration
TAC_REGISTER_NNAPI_OP(SquaredDifferenceOp, CreateBasicOpNoCost);
TAC_REGISTER_NNAPI_OP(ConcatenationOp, CreateConcatOp);
TAC_REGISTER_NNAPI_OP(ReshapeOp, CreateConcatOp);
TAC_REGISTER_NNAPI_OP(PackOp, CreateConcatOp);
#undef TAC_REGISTER_NNAPI_OP
} // namespace tac
} // namespace TFL
} // namespace mlir
tensorflow/compiler/mlir/lite/experimental/tac/hardwares/nnapi_hardware.h
+50
View File
@@ -0,0 +1,50 @@
/* 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.
==============================================================================*/
/* NNAPI Hardware Implementation */
#ifndef TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_TAC_HARDWARES_NNAPI_HARDWARE_H_
#define TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_TAC_HARDWARES_NNAPI_HARDWARE_H_
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "mlir/IR/Operation.h" // from @llvm-project
#include "mlir/IR/PatternMatch.h" // from @llvm-project
#include "mlir/Support/TypeID.h" // from @llvm-project
#include "tensorflow/compiler/mlir/lite/experimental/tac/hardwares/simple_hardware.h"
namespace mlir {
namespace TFL {
namespace tac {
class NNAPIHardware : public SimpleHardware {
public:
static constexpr char kId[] = "NNAPI";
mlir::RewritePatternSet GetTransformations(
MLIRContext* context) const override;
mlir::TypeID GetTypeId() const override {
return mlir::TypeID::get<NNAPIHardware>();
}
bool IsNotSupportedOp(mlir::Operation* op) const override { return false; }
float AdvantageOverCPU() const override { return 5.0; }
};
} // namespace tac
} // namespace TFL
} // namespace mlir
#endif // TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_TAC_HARDWARES_NNAPI_HARDWARE_H_
tensorflow/compiler/mlir/lite/experimental/tac/hardwares/simple_hardware.cc
+53
View File
@@ -0,0 +1,53 @@
/* 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/compiler/mlir/lite/experimental/tac/hardwares/simple_hardware.h"
#include <cstddef>
#include "mlir/IR/Operation.h" // from @llvm-project
#include "tensorflow/compiler/mlir/lite/experimental/tac/hardwares/target_hardware.h"
namespace mlir {
namespace TFL {
namespace tac {
bool SimpleHardware::IsOpSupported(mlir::Operation* op) const {
if (IsNotSupportedOp(op)) {
return false;
}
const TargetHardware* cpu = GetTargetHardware("CPU");
return cpu->IsOpSupported(op);
}
double SimpleHardware::GetHardwareSwitchingCost(const TargetHardware* from,
size_t buffer_size) const {
auto from_type = from->GetTypeId();
auto to_type = GetTypeId();
if (from_type == to_type) return 0.0f;
// TODO(renjieliu): Implement a better version for different hardware cases.
return buffer_size * kCrossHardwareTransferPerByteCost / 8.0 +
kCrossHardwareTransferFixedCost;
}
double SimpleHardware::GetOpCost(mlir::Operation* op) const {
const TargetHardware* cpu = GetTargetHardware("CPU");
return cpu->GetOpCost(op) / AdvantageOverCPU();
}
} // namespace tac
} // namespace TFL
} // namespace mlir
tensorflow/compiler/mlir/lite/experimental/tac/hardwares/simple_hardware.h
+67
View File
@@ -0,0 +1,67 @@
/* 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_COMPILER_MLIR_LITE_EXPERIMENTAL_TAC_HARDWARES_SIMPLE_HARDWARE_H_
#define TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_TAC_HARDWARES_SIMPLE_HARDWARE_H_
#include <cstddef>
#include "mlir/IR/Operation.h" // from @llvm-project
#include "tensorflow/compiler/mlir/lite/experimental/tac/hardwares/target_hardware.h"
namespace mlir {
namespace TFL {
namespace tac {
// A simple hardware is an interface makes you add a target backend easily if
// you don't want too much customization.
//
// It allows you to easily specify the ops capabilities (by
// specifying the denylist), the rest ops will be considered supported. Also you
// can also specify the advantage over CPU.
//
// If you need more customization, e.g., if you have your own hardware dialect,
// please consider use TargetHardware directly.
class SimpleHardware : public TargetHardware {
public:
// This is essentially a denylist.
// TODO(renjieliu): Consider whether we want an allowlist for custom op as
// well.
virtual bool IsNotSupportedOp(mlir::Operation* op) const = 0;
// The larger the value is, the more preferrable over CPU.
// If the value > 1, means the hardware has advantage over CPU.
// If the value < 1, means CPU is more preferred.
// If we specify 10.0, meaning the hardware is 10x faster than CPU.
// The value should be > 0.
// TODO(renjieliu): Consider add an interface for more detailed customization,
// for example, users should be able to specify some ops are preferred and
// some are less preferred.
virtual float AdvantageOverCPU() const = 0;
private:
bool IsOpSupported(mlir::Operation* op) const override;
double GetHardwareSwitchingCost(const TargetHardware* from,
size_t buffer_size) const override;
double GetOpCost(mlir::Operation* op) const override;
};
} // namespace tac
} // namespace TFL
} // namespace mlir
#endif // TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_TAC_HARDWARES_SIMPLE_HARDWARE_H_
tensorflow/compiler/mlir/lite/experimental/tac/hardwares/target_hardware.cc
+284
View File
@@ -0,0 +1,284 @@
/* 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/experimental/tac/hardwares/target_hardware.h"
#include <functional>
#include <memory>
#include <string>
#include <vector>
#include "llvm/ADT/ArrayRef.h"
#include "llvm/ADT/DenseMap.h"
#include "llvm/Support/Casting.h"
#include "llvm/Support/raw_ostream.h"
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/OpDefinition.h" // from @llvm-project
#include "mlir/IR/Operation.h" // from @llvm-project
#include "mlir/Support/TypeID.h" // from @llvm-project
#include "tensorflow/compiler/mlir/lite/experimental/tac/common/targets.h"
#include "tensorflow/compiler/mlir/lite/experimental/tac/common/utils.h"
#include "tensorflow/compiler/mlir/lite/ir/tfl_ops.h"
#include "tensorflow/compiler/mlir/lite/utils/utils.h"
namespace mlir {
namespace TFL {
namespace tac {
namespace {
struct RegisteredTargetHardware {
RegisteredTargetHardware(
const std::string& name, const std::string& description,
mlir::TypeID type_id,
std::function<std::unique_ptr<TargetHardware>()> target_hardware_factory)
: unique_name(GetCanonicalHardwareName(name)),
description(description),
type_id(type_id),
target_hardware(target_hardware_factory()),
target_hardware_factory(target_hardware_factory) {}
std::string unique_name;
std::string description;
mlir::TypeID type_id;
std::unique_ptr<TargetHardware> target_hardware;
std::function<std::unique_ptr<TargetHardware>()> target_hardware_factory;
};
struct RegisteredTargetHardwareOps {
explicit RegisteredTargetHardwareOps(mlir::TypeID hardware_type)
: hardware_typeid(hardware_type) {}
// Key is the Operation TypeID
llvm::DenseMap<mlir::TypeID, std::unique_ptr<TargetHardwareOperation>>
target_hardware_ops;
// Key is the Operation TypeID
llvm::DenseMap<mlir::TypeID,
std::function<std::unique_ptr<TargetHardwareOperation>()>>
target_hardware_ops_factory;
mlir::TypeID hardware_typeid;
};
std::vector<std::unique_ptr<RegisteredTargetHardwareOps>>*
GetRegisteredTargetHardwareOps() {
static std::vector<std::unique_ptr<RegisteredTargetHardwareOps>>*
hardwares_ops =
[]() -> std::vector<std::unique_ptr<RegisteredTargetHardwareOps>>* {
return new std::vector<std::unique_ptr<RegisteredTargetHardwareOps>>();
}();
return hardwares_ops;
}
std::vector<RegisteredTargetHardware>* GetRegisteredHardwares() {
static std::vector<RegisteredTargetHardware>* hardwares =
[]() -> std::vector<RegisteredTargetHardware>* {
return new std::vector<RegisteredTargetHardware>();
}();
return hardwares;
}
llvm::DenseMap<mlir::TypeID, std::unique_ptr<TargetHardwareOperation>>*
getRegisteredOperationsForHardware(mlir::TypeID type_id) {
auto* hardwares = GetRegisteredTargetHardwareOps();
for (auto& hardware : *hardwares) {
if (hardware->hardware_typeid == type_id) {
return &hardware->target_hardware_ops;
}
}
return nullptr;
}
// A deny list for op cost computation since those ops are not arithemtic.
inline bool IsNonArithmeticOp(mlir::Operation* op) {
if (llvm::isa<func::ReturnOp, func::FuncOp>(op)) return true;
if (op->hasTrait<OpTrait::ConstantLike>()) return true;
if (llvm::isa<QConstOp, SparseQConstOp>(op)) return true;
if (!NotTFLQuantDequantizeOp(op)) return true;
return false;
}
} // namespace
bool TargetHardware::Init() {
auto* hardware_ops_factory = GetRegisteredTargetHardwareOps();
for (auto& hardware_ops : *hardware_ops_factory) {
if (hardware_ops->hardware_typeid != this->GetTypeId()) continue;
auto& op_factories = hardware_ops->target_hardware_ops_factory;
for (auto& op_factory : op_factories) {
hardware_ops_.emplace_back(op_factory.getSecond()());
}
break;
}
return true;
}
double TargetHardware::GetOpCost(mlir::Operation* op) const {
auto* registered_ops = getRegisteredOperationsForHardware(GetTypeId());
if (registered_ops == nullptr) {
return kDefaultFixedValuedCost;
}
auto abstract_op = op->getRegisteredInfo();
auto hardware_op = registered_ops->find(abstract_op->getTypeID());
if (hardware_op == registered_ops->end()) return kDefaultFixedValuedCost;
return hardware_op->second->GetOpCost(op);
}
bool TargetHardware::IsOpSupported(mlir::Operation* op) const {
auto* registered_ops = getRegisteredOperationsForHardware(GetTypeId());
if (registered_ops == nullptr) {
return false;
}
auto abstract_op = op->getRegisteredInfo();
auto hardware_op = registered_ops->find(abstract_op->getTypeID());
if (hardware_op == registered_ops->end()) return false;
return hardware_op->second->IsOpSupported(op);
}
double TargetHardware::GetFuncCost(func::FuncOp* func) const {
double total_cost = 0.0;
func->walk([&](Operation* op) {
if (IsNonArithmeticOp(op)) return;
// We will always defer to the hardware to decide the cost.
total_cost += GetOpCost(op);
});
return total_cost;
}
const TargetHardware* GetTargetHardware(const std::string& hardware_name) {
const std::string canonical_name = GetCanonicalHardwareName(hardware_name);
// Just loop for now, we don't expect number of hardwares to be huge.
// Revisit to have map if number of elements increased.
auto* registered_hardwares = GetRegisteredHardwares();
for (const auto& hardware : *registered_hardwares) {
if (hardware.unique_name == canonical_name) {
return hardware.target_hardware.get();
}
}
return nullptr;
}
std::function<std::unique_ptr<TargetHardware>()> GetTargetHardwareFactory(
const std::string& hardware_name) {
const std::string canonical_name = GetCanonicalHardwareName(hardware_name);
// Just loop for now, we don't expect number of hardwares to be huge.
// Revisit to have map if number of elements increased.
auto* registered_hardwares = GetRegisteredHardwares();
for (const auto& hardware : *registered_hardwares) {
if (hardware.unique_name == canonical_name) {
return hardware.target_hardware_factory;
}
}
return nullptr;
}
namespace internal {
void RegisterTargetHardwareFactory(
const std::string& unique_name, const std::string& description,
mlir::TypeID type_id,
std::function<std::unique_ptr<TargetHardware>()> target_hardware_factory) {
auto* registered_hardwares = GetRegisteredHardwares();
for (auto& hardware : *registered_hardwares) {
if (hardware.unique_name == unique_name) {
llvm::errs() << "Ignoring duplicate hardware. Hardware " << unique_name
<< " already registered\n";
hardware.target_hardware_factory = target_hardware_factory;
return;
}
}
registered_hardwares->push_back(RegisteredTargetHardware(
unique_name, description, type_id, target_hardware_factory));
}
void RegisterTargetHardwareOp(
mlir::TypeID hardware_type, mlir::TypeID op_type,
std::function<std::unique_ptr<TargetHardwareOperation>()>
target_hardware_op_factory) {
auto* registered_hardware_ops = GetRegisteredTargetHardwareOps();
for (auto& hardware : *registered_hardware_ops) {
if (hardware->hardware_typeid == hardware_type) {
if (hardware->target_hardware_ops.count(op_type)) {
llvm::errs() << "Trying to register duplicate Op";
return;
}
hardware->target_hardware_ops[op_type] = target_hardware_op_factory();
return;
}
}
registered_hardware_ops->push_back(
std::make_unique<RegisteredTargetHardwareOps>(
RegisteredTargetHardwareOps(hardware_type)));
registered_hardware_ops->back()->target_hardware_ops[op_type] =
target_hardware_op_factory();
}
void RegisterTargetHardwareOpFactory(
mlir::TypeID hardware_type, mlir::TypeID op_type,
std::function<std::unique_ptr<TargetHardwareOperation>()>
target_hardware_op_factory) {
auto* registered_hardware_ops = GetRegisteredTargetHardwareOps();
for (auto& hardware : *registered_hardware_ops) {
if (hardware->hardware_typeid == hardware_type) {
if (hardware->target_hardware_ops_factory.count(op_type)) {
llvm::errs() << "Trying to register duplicate Op";
return;
}
hardware->target_hardware_ops_factory[op_type] =
target_hardware_op_factory;
return;
}
}
registered_hardware_ops->push_back(
std::make_unique<RegisteredTargetHardwareOps>(
RegisteredTargetHardwareOps(hardware_type)));
registered_hardware_ops->back()->target_hardware_ops_factory[op_type] =
target_hardware_op_factory;
}
} // namespace internal
bool ProcessTargetDevices(llvm::ArrayRef<std::string> specified_device_specs,
std::vector<std::string>* device_specs) {
bool cpu_include = false;
for (auto& device_spec : specified_device_specs) {
auto device = GetCanonicalHardwareName(device_spec);
if (device == "CPU") cpu_include = true;
device_specs->push_back(device);
}
if (!cpu_include) {
device_specs->push_back("CPU");
}
// Make sure all the devices are registered.
for (const std::string& device : *device_specs) {
if (GetTargetHardware(device) == nullptr) {
llvm::errs() << "cannot get target hardware for device: " << device;
return false;
}
}
return true;
}
std::string GetHardwareName(const TargetHardware* hardware) {
const auto* registered_hardwares = GetRegisteredHardwares();
for (const auto& registered_hardware : *registered_hardwares) {
if (registered_hardware.type_id == hardware->GetTypeId())
return registered_hardware.unique_name;
}
return "";
}
} // namespace tac
} // namespace TFL
} // namespace mlir
tensorflow/compiler/mlir/lite/experimental/tac/hardwares/target_hardware.h
+196
View File
@@ -0,0 +1,196 @@
/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_TAC_HARDWARES_TARGET_HARDWARE_H_
#define TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_TAC_HARDWARES_TARGET_HARDWARE_H_
#include <cstddef>
#include <functional>
#include <memory>
#include <string>
#include <vector>
#include "llvm/ADT/ArrayRef.h"
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/DialectRegistry.h" // from @llvm-project
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "mlir/IR/Operation.h" // from @llvm-project
#include "mlir/IR/PatternMatch.h" // from @llvm-project
#include "mlir/Support/TypeID.h" // from @llvm-project
namespace mlir {
namespace TFL {
namespace tac {
// Default fixed values for ops.
constexpr static float kDefaultFixedValuedCost = 1000000.0;
// This is just fake data.
constexpr static float kCrossHardwareTransferPerByteCost = 5.0f;
// This is just fake data.
constexpr static float kCrossHardwareTransferFixedCost = 10.f;
// Interface for an Operation capabilities which should be tied to
// a specific hardware.
// Users should implement the interface and use TargetHardwareOpRegistration
// for registering the operation.
class TargetHardwareOperation {
public:
virtual ~TargetHardwareOperation() = default;
virtual double GetOpCost(mlir::Operation* op) const = 0;
virtual bool IsOpSupported(mlir::Operation* op) const = 0;
};
// Abstract base class for a hardware.
// To introduce new hardware
// users should implement the interface and use TargetHardwareRegistration
// for registering the hardware.
// Subclasses must implement the pure virtual function interface and
// define static member variable that retrieves string identifying the Target
// Hardware. Example,
// class MyType : public TargetHardware {
// public:
// static constexpr char kId[] = "MyHardware";
// };
class TargetHardware {
public:
virtual ~TargetHardware() = default;
// Initializes all TargetHardwareOperation registered for this hardware.
// Users overriding this function, should call the base class method to
// initialize the ops.
virtual bool Init();
// Returns the cost of running 'op' on this Hardware.
virtual double GetOpCost(mlir::Operation* op) const;
// Returns the cost of running the whole function on this hardware.
// By default this is the sum of the cost of individual cost for each op.
virtual double GetFuncCost(func::FuncOp* func) const;
// Returns true if 'op' can run on this Hardware.
virtual bool IsOpSupported(mlir::Operation* op) const;
// Switching cost between from hardware and this hardware.
// If both the hardwares are the same, the transfer cost is basically 0.
virtual double GetHardwareSwitchingCost(const TargetHardware* from,
size_t buffer_size) const = 0;
// Returns a list of all patterns to apply for this hardware.
virtual mlir::RewritePatternSet GetTransformations(
MLIRContext* context) const = 0;
// Returns TypeId for the provided hardware.
// Usually should be something like mlir::TypeID::get<MyType>()
virtual mlir::TypeID GetTypeId() const = 0;
virtual void GetDependentDialects(mlir::DialectRegistry& registry) const {}
protected:
// All registered hardware ops.
std::vector<std::unique_ptr<TargetHardwareOperation>> hardware_ops_;
};
// Returns pointer to the Hardware identified by 'hardware_name'.
// If not found nullptr is returned.
// DEPRECATED: Do not use, prefer GetTargetHardwareFactory instead.
const TargetHardware* GetTargetHardware(const std::string& hardware_name);
// Returns the factory method for the requested hardware if present.
std::function<std::unique_ptr<TargetHardware>()> GetTargetHardwareFactory(
const std::string& hardware_name);
namespace internal {
void RegisterTargetHardwareFactory(
const std::string& unique_name, const std::string& description,
mlir::TypeID type_id,
std::function<std::unique_ptr<TargetHardware>()> target_hardware_factory);
// Registers the provided target hardware factory.
template <typename T>
void RegisterTargetHardwareFactory(
const std::string& description,
std::function<std::unique_ptr<TargetHardware>()> target_hardware_factory) {
RegisterTargetHardwareFactory(T::kId, description, mlir::TypeID::get<T>(),
target_hardware_factory);
}
// DEPRECATED: Do not use, prefer RegisterTargetHardwareOpFactory intstead.
void RegisterTargetHardwareOp(
mlir::TypeID hardware_type, mlir::TypeID op_type,
std::function<std::unique_ptr<TargetHardwareOperation>()>
target_hardware_op_factory);
void RegisterTargetHardwareOpFactory(
mlir::TypeID hardware_type, mlir::TypeID op_type,
std::function<std::unique_ptr<TargetHardwareOperation>()>
target_hardware_op_factory);
} // namespace internal
// Register target hardware.
template <typename Hardware>
struct TargetHardwareRegistration {
TargetHardwareRegistration(const std::string& description,
std::function<std::unique_ptr<TargetHardware>()>
target_hardware_factory) {
internal::RegisterTargetHardwareFactory<Hardware>(description,
target_hardware_factory);
}
};
// Register Op capabilities for specific hardware.
template <typename Hardware, typename Op>
struct TargetHardwareOpRegistration {
explicit TargetHardwareOpRegistration(
std::function<std::unique_ptr<TargetHardwareOperation>()>
target_hardware_op_factory) {
// TODO(b/177376459): remove this.
internal::RegisterTargetHardwareOp(mlir::TypeID::get<Hardware>(),
mlir::TypeID::get<Op>(),
target_hardware_op_factory);
internal::RegisterTargetHardwareOpFactory(mlir::TypeID::get<Hardware>(),
mlir::TypeID::get<Op>(),
target_hardware_op_factory);
}
};
//======== util functions ==========
// Process user specified device specs, will always add CPU if it's not there.
// specified_device_specs: ',' separated, like "GPU,DSP,CPU".
// device_specs: processed device specs enum.
bool ProcessTargetDevices(llvm::ArrayRef<std::string> specified_device_specs,
std::vector<std::string>* device_specs);
// Check whether two hardwares are the same.
inline bool IsSameHardware(const TargetHardware* lhs,
const TargetHardware* rhs) {
return lhs->GetTypeId() == rhs->GetTypeId();
}
// Returns the ID identifying 'hardware'. This should match the ID defined
// in the hardware field ID.
// For example, if MyHardware is passed the value returned should match
// MyHardware::kId.
std::string GetHardwareName(const TargetHardware* hardware);
} // namespace tac
} // namespace TFL
} // namespace mlir
#endif // TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_TAC_HARDWARES_TARGET_HARDWARE_H_
tensorflow/compiler/mlir/lite/experimental/tac/py_wrapper/BUILD
+110
View File
@@ -0,0 +1,110 @@
load("@rules_cc//cc:cc_library.bzl", "cc_library")
load("//tensorflow:tensorflow.bzl", "VERSION")
load("//tensorflow:tensorflow.default.bzl", "pybind_extension")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:LICENSE"],
default_visibility = [
"//tensorflow/compiler/mlir/lite/experimental/tac:__subpackages__",
],
licenses = ["notice"],
)
cc_library(
name = "tac_wrapper_lib",
srcs = ["tac_wrapper.cc"],
hdrs = [
"tac_wrapper.h",
],
deps = [
"//tensorflow/compiler/mlir/lite/experimental/tac:tac_importer_exporter",
"//tensorflow/compiler/mlir/lite/experimental/tac:target_aware_conversion",
"//tensorflow/compiler/mlir/lite/experimental/tac:tflite_importer_exporter",
"//tensorflow/compiler/mlir/lite/experimental/tac/hardwares:all-target-hardwares",
"//tensorflow/compiler/mlir/tensorflow",
"@com_google_absl//absl/status",
"@llvm-project//mlir:IR",
"@xla//third_party/python_runtime:headers", # buildcleaner: keep
],
)
pybind_extension(
name = "_pywrap_tac_wrapper",
srcs = [
"tac_wrapper_pybind11.cc",
],
hdrs = ["tac_wrapper.h"],
dynamic_deps = select({
"//tensorflow:macos": ["//tensorflow:libtensorflow_framework.%s.dylib" % VERSION],
"//tensorflow:windows": [],
"//conditions:default": ["//tensorflow:libtensorflow_framework.so.%s" % VERSION],
}),
enable_stub_generation = True,
pytype_srcs = [
"_pywrap_tac_wrapper.pyi",
],
static_deps = [
"@arm_neon_2_x86_sse//:__subpackages__",
"@bazel_tools//:__subpackages__",
"@boringssl//:__subpackages__",
"@clog//:__subpackages__",
"@com_github_cares_cares//:__subpackages__",
"@com_github_googlecloudplatform_tensorflow_gcp_tools//:__subpackages__",
"@com_github_grpc_grpc//:__subpackages__",
"@com_google_absl//:__subpackages__",
"@com_google_googleapis//:__subpackages__",
"@com_google_protobuf//:__subpackages__",
"@com_googlesource_code_re2//:__subpackages__",
"@compute_library//:__subpackages__",
"@cpuinfo//:__subpackages__",
"@curl//:__subpackages__",
"@eigen_archive//:__subpackages__",
"@farmhash_archive//:__subpackages__",
"@farmhash_gpu_archive//:__subpackages__",
"@fft2d//:__subpackages__",
"@flatbuffers//:__subpackages__",
"@FP16//:__subpackages__",
"@FXdiv//:__subpackages__",
"@gemmlowp//:__subpackages__",
"@go_googleapis//google/api:annotations_proto",
"@gif//:__subpackages__",
"@highwayhash//:__subpackages__",
"@hwloc//:__subpackages__",
"@icu//:__subpackages__",
"@jsoncpp_git//:__subpackages__",
"@libjpeg_turbo//:__subpackages__",
"@llvm_openmp//:__subpackages__",
"@llvm-project//:__subpackages__",
"@llvm_terminfo//:__subpackages__",
"@llvm_zlib//:__subpackages__",
"@local_config_cuda//:__subpackages__",
"@local_config_git//:__subpackages__",
"@local_config_nccl//:__subpackages__",
"@local_config_python//:__subpackages__",
"@local_config_rocm//:__subpackages__",
"@local_config_tensorrt//:__subpackages__",
"@mkl_dnn_acl_compatible//:__subpackages__",
"@onednn//:__subpackages__",
"@org_sqlite//:__subpackages__",
"@platforms//:__subpackages__",
"@png//:__subpackages__",
"@pthreadpool//:__subpackages__",
"@pybind11//:__subpackages__",
"@ruy//:__subpackages__",
"@snappy//:__subpackages__",
"@sobol_data//:__subpackages__",
"@stablehlo//:__subpackages__",
"//:__subpackages__",
"@upb//:__subpackages__",
"@XNNPACK//:__subpackages__",
"@zlib//:__subpackages__",
"@tsl//tsl:__subpackages__",
"@xla//xla:__subpackages__",
],
deps = [
":tac_wrapper_lib",
"//tensorflow/python/lib/core:pybind11_lib",
"@pybind11",
"@xla//third_party/python_runtime:headers",
],
)
tensorflow/compiler/mlir/lite/experimental/tac/py_wrapper/_pywrap_tac_wrapper.pyi
+16
View File
@@ -0,0 +1,16 @@
# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
def run_tac(arg0: str, arg1: list[str], arg2: str) -> bool: ...
tensorflow/compiler/mlir/lite/experimental/tac/py_wrapper/tac_wrapper.cc
+72
View File
@@ -0,0 +1,72 @@
/* 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/compiler/mlir/lite/experimental/tac/py_wrapper/tac_wrapper.h"
#include <memory>
#include <string>
#include <vector>
#include "absl/status/status.h"
#include "mlir/IR/DialectRegistry.h" // from @llvm-project
#include "tensorflow/compiler/mlir/lite/experimental/tac/tac_importer_exporter.h"
#include "tensorflow/compiler/mlir/lite/experimental/tac/tac_module.h"
#include "tensorflow/compiler/mlir/lite/experimental/tac/tflite_import_export.h"
#include "tensorflow/compiler/mlir/lite/experimental/tac/transforms/passes.h"
#include "tensorflow/compiler/mlir/tensorflow/dialect_registration.h"
namespace tflite {
namespace {
std::unique_ptr<mlir::TFL::tac::TacImporter> CreateTfLiteImporter(
const std::string input_file_name) {
mlir::TFL::tac::TfLiteImporter::Options options;
options.file_name = input_file_name;
options.input_mlir = false;
return std::make_unique<mlir::TFL::tac::TfLiteImporter>(options);
}
std::unique_ptr<mlir::TFL::tac::TacExporter> CreateTfLiteExporter(
const std::string output_file_name,
const std::vector<std::string>& target_hardware_backends) {
mlir::TFL::tac::TfLiteExporter::Options options;
options.output_mlir = false;
options.output_file_name = output_file_name;
options.export_runtime_metadata = false;
options.target_hardware_backends = target_hardware_backends;
return std::make_unique<mlir::TFL::tac::TfLiteExporter>(options);
}
} // namespace
// Run target-aware-conversion for the given tflite model with the given device
// specs.
// Warning: The API is experimental and subject to changes.
bool run_tac(const std::string& model_file_path,
const std::vector<std::string>& device_specs,
const std::string& model_output_path) {
mlir::TFL::tac::TacModule::Options options;
options.hardware_backends = device_specs;
options.enable_inliner = true;
options.legalize_to_tflite_ops = true;
mlir::TFL::tac::TacModule tac_module(options);
mlir::DialectRegistry registry;
mlir::RegisterAllTensorFlowDialects(registry);
tac_module.RegisterExtraDialects(registry);
tac_module.SetImporter(CreateTfLiteImporter(model_file_path));
tac_module.SetExporter(
CreateTfLiteExporter(model_output_path, options.hardware_backends));
return tac_module.Run().ok();
}
} // namespace tflite
tensorflow/compiler/mlir/lite/experimental/tac/py_wrapper/tac_wrapper.h
+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.
==============================================================================*/
#ifndef TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_TAC_PY_WRAPPER_TAC_WRAPPER_H_
#define TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_TAC_PY_WRAPPER_TAC_WRAPPER_H_
#include <functional>
#include <memory>
#include <string>
#include <vector>
// Place `<locale>` before <Python.h> to avoid build failures in macOS.
#include <locale>
// The empty line above is on purpose as otherwise clang-format will
// automatically move <Python.h> before <locale>.
#include <Python.h>
namespace tflite {
// Run target-aware-conversion for the given tflite model with the given device
// specs.
// Warning: The API is experimental and subject to change.
bool run_tac(const std::string& model_file_path,
const std::vector<std::string>& device_specs,
const std::string& model_output_path);
} // namespace tflite
#endif // TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_TAC_PY_WRAPPER_TAC_WRAPPER_H_
tensorflow/compiler/mlir/lite/experimental/tac/py_wrapper/tac_wrapper_pybind11.cc
+37
View File
@@ -0,0 +1,37 @@
/* 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 <pybind11/stl.h>
#include <string>
#include <vector>
#include "pybind11/pybind11.h" // from @pybind11
#include "tensorflow/compiler/mlir/lite/experimental/tac/py_wrapper/tac_wrapper.h"
// Warning: The API is experimental and subject to change.
PYBIND11_MODULE(_pywrap_tac_wrapper, m) {
m.def(
"run_tac",
[](const std::string& model_file_path,
const std::vector<std::string>& device_specs,
const std::string& model_output_path) {
return ::tflite::run_tac(model_file_path, device_specs,
model_output_path);
},
R"pbdoc(
Run target-aware-conversion with the given device specs.
)pbdoc");
}
tensorflow/compiler/mlir/lite/experimental/tac/runtime_metadata.fbs
+89
View File
@@ -0,0 +1,89 @@
// 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.
namespace tflite;
// Here is an overview of the pipeline & runtime metadata looks like:
//
// ======== IRs after TAC ========
// op1 = .... {tac.device = "CPU"}
// op2 = .... {tac.device = "GPU"}
// op3 = .... {tac.device = "GPU"}
// op4 = .... {tac.device = "CPU"}
//
// We will run a separate cost estimation for each op with each different
// hardwares.
//
// ======== IRs after cost-estimation ========
// op1 = .... {tac.device = "CPU"} {"CPU": 10, "GPU": -1.0}
// op2 = .... {tac.device = "GPU"} {"CPU": 20, "GPU": 5.0}
// op3 = .... {tac.device = "GPU"} {"CPU": 30, "GPU": 10.0}
// op4 = .... {tac.device = "CPU"} {"CPU": 40, "GPU": -1.0}
//
// The exported runtime metadata will look like below:
// ======== Runtime Metadata ========
//
// RuntimeMetadata: {
// hardwares: {
// ["CPU", "GPU"]
// }
//
// subgraph_metadata: {
// [
// op1:
// hardware: 0
// op_costs: [10, -1],
// op2:
// hardware: 1
// op_costs: [20, 5],
// op3:
// hardware: 1
// op_costs: [30, 10],
// op4:
// hardware: 0
// op_costs: [40, -1],
// ]
// }
// }
table HardwareMetadata {
// List of different hardwares this model can use.
hardware_name:[string];
}
// Metadata for a single Op.
table OpMetadata {
// Node index in the corresponding subgraph.
index:int;
// The hardware suggested for inference.
// This will correspond to the index in the hardware_name field in the
// hardware metadata.
hardware:uint8;
// The estimated costs for the given op.
// The index matches the hardware array.
op_costs: [float];
}
table SubgraphMetadata {
// Metadata for all ops in the current subgraph.
op_metadata:[OpMetadata];
}
table RuntimeMetadata {
// Metadata about different hardwares used by the model.
hardwares: HardwareMetadata;
// Metadata for each subgraph in same order as in the TFLite model file.
subgraph_metadata:[SubgraphMetadata];
}
root_type RuntimeMetadata;
tensorflow/compiler/mlir/lite/experimental/tac/tac.py
+46
View File
@@ -0,0 +1,46 @@
# 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.
# ==============================================================================
"""Target aware conversion for TFLite model."""
from tensorflow.compiler.mlir.lite.experimental.tac.py_wrapper import _pywrap_tac_wrapper
def run_tac(model_path, targets, output_path):
"""Run target aware conversion for the given tflite model file.
Args:
model_path: Path to the tflite model file.
targets: A list of string of the desired targets. E.g., ['GPU', 'CPU'].
output_path: The output path.
Returns:
Whether the optimization succeeded.
Raises:
ValueError:
Invalid model_path.
Targets are not specified.
Invalid output_path.
"""
if not model_path:
raise ValueError("Invalid model_path.")
if not targets:
raise ValueError("Targets are not specified.")
if not output_path:
raise ValueError("Invalid output_path.")
return _pywrap_tac_wrapper.run_tac(model_path, targets, output_path)
tensorflow/compiler/mlir/lite/experimental/tac/tac_filter.proto
+87
View File
@@ -0,0 +1,87 @@
// Copyright 2026 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.
// ==============================================================================
syntax = "proto3";
package third_party.tensorflow.compiler.mlir.lite.experimental.tac;
import "google/protobuf/any.proto";
// A list of filters for TAC users to run ops/functions on ML hardwares. The
// intuition is that, for ops/functions that can be run on ML hardware (e.g.
// EdgeTPU) and TFLite CPU, TAC users give a hint that they're more performant
// to run on TFLite CPU. These filters give the TAC users freedom to specify the
// parts that they want to use other hardware to accelerate.
message TacFilters {
// A list of filters/rules to specify the parts that user wants to run on
// other hardware.
repeated TacFilter tac_filters = 1;
}
// A filter can be used for an op or function.
message TacFilter {
oneof filter {
OpFilter op_filter = 1;
FunctionFilter function_filter = 2;
}
}
// Function filter is to include/exclude a function in the target annotation
// pass in the TAC tool pipeline.
message FunctionFilter {
// Function filter types that are supported. If one function is matched for
// two rules with conflict, INCLUDE_TARGET_ANNOTATION has higher priority.
enum FunctionFilterType {
// To skip this function in the target annotation pass. This means all ops
// in this function run on TFLite CPU.
SKIP_TARGET_ANNOTATION = 0;
// To include this function in the target annotation pass. This has higher
// priority than `SKIP_TARGET_ANNOTATION`.
INCLUDE_TARGET_ANNOTATION = 1;
}
// This name corresponds to the TFLite subgraph name in the flatbuffer.
// `function_name_pattern` supports regex matching.
string function_name_pattern = 1;
FunctionFilterType filter_type = 2;
}
// Op filter is to filter out ops that user wants to run. Ops with this filter
// run on TFLite CPU.
message OpFilter {
// This name corresponds to the mlir::Location of the tensor.
// `op_name_pattern` supports regex matching.
string op_name_pattern = 1;
enum MatchType {
// To match the op name with the `op_name_pattern` directly.
MATCH = 0;
// To invert the match with the `op_name_pattern`, ie, the filter will
// select an op if its name does not match the pattern above.
INVERT_MATCH = 1;
}
MatchType match_type = 2;
enum DeviceType {
// To run the op on CPU.
CPU = 0;
// To run the op on an hardware unit other than the CPU.
OTHER = 1;
}
DeviceType device_type = 3;
// An arbitrary string that can be used to identify the filter or a set of
// filters.
string tag = 4;
// Custom options that can be used to pass additional information for the
// filter.
google.protobuf.Any custom_options = 5;
}
tensorflow/compiler/mlir/lite/experimental/tac/tac_importer_exporter.h
+52
View File
@@ -0,0 +1,52 @@
/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_TAC_TAC_IMPORTER_EXPORTER_H_
#define TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_TAC_TAC_IMPORTER_EXPORTER_H_
#include "absl/status/statusor.h"
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
namespace mlir {
namespace TFL {
namespace tac {
// Interface for Importing program to TAC (Target Aware Conversion) Module.
// This class is an interface for importing program in TAC.
// See TacModule in how to register it with the module and use it.
class TacImporter {
public:
virtual ~TacImporter() = default;
// Imports and returns the Module for the imported program.
virtual absl::StatusOr<mlir::OwningOpRef<mlir::ModuleOp>> Import() = 0;
};
// Interface for exporting a module.
// Users should implement the interface for exporting the result from TAC
// in their preferred way.
// See TacModule in how to register it with the module and use it.
class TacExporter {
public:
virtual ~TacExporter() = default;
// Imports and returns the Module for the imported program.
virtual absl::Status Export(mlir::ModuleOp module) = 0;
};
} // namespace tac
} // namespace TFL
} // namespace mlir
#endif // TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_TAC_TAC_IMPORTER_EXPORTER_H_
tensorflow/compiler/mlir/lite/experimental/tac/tac_module.cc
+149
View File
@@ -0,0 +1,149 @@
/* 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/compiler/mlir/lite/experimental/tac/tac_module.h"
#include <memory>
#include <string>
#include <vector>
#include "absl/status/status.h"
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/Transforms/Passes.h" // from @llvm-project
#include "tensorflow/compiler/mlir/lite/experimental/tac/hardwares/target_hardware.h"
#include "tensorflow/compiler/mlir/lite/experimental/tac/transforms/passes.h"
#include "tensorflow/compiler/mlir/lite/transforms/passes.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/error_util.h"
namespace mlir {
namespace TFL {
namespace tac {
namespace {
// TODO(b/177376459): We should make this configureable.
void AddExportTFLPass(mlir::OpPassManager* pass_manager, bool enable_inliner) {
if (enable_inliner) pass_manager->addPass(mlir::createInlinerPass());
pass_manager->addPass(mlir::createSymbolDCEPass());
pass_manager->addNestedPass<mlir::func::FuncOp>(
mlir::createCanonicalizerPass());
pass_manager->addNestedPass<mlir::func::FuncOp>(mlir::createCSEPass());
}
} // namespace
// TODO(b/177376459): We should make this configureable.
void TacModule::AddTACPass(mlir::OpPassManager* pass_manager,
llvm::ArrayRef<std::string> device_specs) {
pass_manager->addPass(mlir::TFL::tac::CreateTargetAnnotationPass(this));
pass_manager->addPass(mlir::TFL::tac::CreateRaiseTargetSubgraphsPass());
pass_manager->addPass(mlir::TFL::tac::CreateFoldConstantsToSubgraphPass(
/*fold_all_constants=*/false));
pass_manager->addPass(
mlir::TFL::tac::CreateAlternativeSubgraphPass(device_specs));
if (options_.legalize_to_tflite_ops) {
// After we creat the alternative subgraph, we can still do canonicalization
// legalization & other optimizations as long as we're not inlining the
// function.
// And in fact, we probably need to do the proper legalization, for the
// compute cost to work. (in case we added some TF ops)
pass_manager->addPass(mlir::TFL::CreatePrepareTFPass(
/*unfold_batch_matmul=*/true,
/*allow_bf16_and_f16_type_legalization=*/false));
pass_manager->addNestedPass<mlir::func::FuncOp>(
mlir::createCanonicalizerPass());
pass_manager->addPass(
mlir::TFL::CreateLegalizeTFPass(/*run_tfl_runtime_verification=*/true));
pass_manager->addPass(mlir::TFL::CreateOptimizePass());
}
pass_manager->addPass(mlir::TFL::tac::CreateComputeCostPass());
pass_manager->addPass(mlir::TFL::tac::CreatePickSubgraphsPass());
// After this pass, we may consider add a pass to merge small functions into
// large functions (and maybe other metadata as well).
}
const tac::TargetHardware* TacModule::GetTargetHardware(
const std::string& hardware_name) const {
for (auto& hardware : backends_) {
if (GetHardwareName(hardware.get()) == hardware_name) return hardware.get();
}
return nullptr;
}
absl::Status TacModule::RunTacPasses(mlir::ModuleOp* module, bool debug_mode) {
mlir::PassManager pm((*module)->getName(),
mlir::OpPassManager::Nesting::Implicit);
AddTACPass(&pm, options_.hardware_backends);
if (!debug_mode) {
AddExportTFLPass(&pm, options_.enable_inliner);
}
mlir::StatusScopedDiagnosticHandler statusHandler(module->getContext(),
/*propagate=*/true);
if (failed(pm.run(*module))) {
return absl::InternalError("conversion error");
}
return absl::OkStatus();
}
std::vector<std::unique_ptr<tac::TargetHardware>>
TacModule::InstantiateBackends() {
std::vector<std::unique_ptr<tac::TargetHardware>> backends;
for (const auto& hardware_name : options_.hardware_backends) {
auto factory = tac::GetTargetHardwareFactory(hardware_name);
backends.emplace_back(factory());
backends.back()->Init();
}
return backends;
}
absl::Status TacModule::Run() {
// Construct all backends.
backends_ = InstantiateBackends();
const_backends_.resize(backends_.size());
for (const auto& backend : backends_)
const_backends_.emplace_back(backend.get());
if (!importer_) {
return absl::Status(absl::StatusCode::kFailedPrecondition,
"Null Importer provided");
}
if (!exporter_) {
return absl::Status(absl::StatusCode::kFailedPrecondition,
"Null Exporter provided");
}
auto module_status = importer_->Import();
if (!module_status.ok()) {
return module_status.status();
}
auto module = module_status->get();
auto* context = module->getContext();
context->appendDialectRegistry(registry_);
context->loadAllAvailableDialects();
// Run TAC passes.
auto status = RunTacPasses(&module, options_.debug_mode);
if (!status.ok()) {
return status;
}
return exporter_->Export(module);
}
void TacModule::RegisterExtraDialects(mlir::DialectRegistry& registry) {
registry.appendTo(registry_);
}
} // namespace tac
} // namespace TFL
} // namespace mlir

Some files were not shown because too many files have changed in this diff Show More