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/tfr/BUILD
+381
View File
@@ -0,0 +1,381 @@
load(
"@llvm-project//mlir:tblgen.bzl",
"gentbl_cc_library",
"td_library",
)
load("@xla//third_party/rules_python/python:py_library.bzl", "py_library")
load(
"//tensorflow:tensorflow.bzl",
"tf_cc_binary",
"tf_cc_test",
)
load("//tensorflow:tensorflow.default.bzl", "get_compatible_with_portable", "tf_py_strict_test", "tf_python_pybind_extension")
load("//tensorflow/compiler/mlir:glob_lit_test.bzl", "glob_lit_tests")
load("//tensorflow/compiler/mlir/tfr:build_defs.bzl", "gen_op_libraries")
load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
default_visibility = [
":friends",
],
licenses = ["notice"],
)
package_group(
name = "friends",
packages = [
"//tensorflow/c/...",
"//tensorflow/compiler/...",
# Allow visibility from the mlir language server.
"//learning/brain/mlir/mlir_lsp_server/...",
],
)
td_library(
name = "tfr_ops_td_files",
srcs = [
"ir/tfr_ops.td",
],
compatible_with = get_compatible_with_portable(),
deps = [
"//tensorflow/compiler/mlir/tensorflow:tensorflow_ops_td_files",
"@llvm-project//mlir:CallInterfacesTdFiles",
"@llvm-project//mlir:ControlFlowInterfacesTdFiles",
"@llvm-project//mlir:FunctionInterfacesTdFiles",
"@llvm-project//mlir:OpBaseTdFiles",
"@llvm-project//mlir:QuantizationOpsTdFiles",
"@llvm-project//mlir:ShapeOpsTdFiles",
"@llvm-project//mlir:SideEffectInterfacesTdFiles",
],
)
gentbl_cc_library(
name = "tfr_ops_inc_gen",
compatible_with = get_compatible_with_portable(),
tbl_outs = {
"ir/tfr_ops.h.inc": ["-gen-op-decls"],
"ir/tfr_ops.cc.inc": ["-gen-op-defs"],
},
tblgen = "@llvm-project//mlir:mlir-tblgen",
td_file = "ir/tfr_ops.td",
deps = [
":tfr_ops_td_files",
],
)
gentbl_cc_library(
name = "tfr_decompose_inc_gen",
compatible_with = get_compatible_with_portable(),
tbl_outs = {"passes/generated_decompose.inc": ["-gen-rewriters"]},
tblgen = "@llvm-project//mlir:mlir-tblgen",
td_file = "passes/decompose_patterns.td",
deps = [
":tfr_ops_td_files",
"@llvm-project//mlir:ArithOpsTdFiles",
"@llvm-project//mlir:FuncTdFiles",
],
)
cc_library(
name = "tfr",
srcs = [
"ir/tfr_ops.cc",
"ir/tfr_ops.cc.inc",
],
hdrs = [
"ir/tfr_ops.h",
"ir/tfr_ops.h.inc",
"ir/tfr_types.h",
],
deps = [
":tfr_ops_inc_gen",
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_attributes",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_ops",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_types",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:ArithDialect",
"@llvm-project//mlir:ControlFlowInterfaces",
"@llvm-project//mlir:Dialect",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:FunctionInterfaces",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:InferTypeOpInterface",
"@llvm-project//mlir:InliningUtils",
"@llvm-project//mlir:QuantOps",
"@llvm-project//mlir:ShapeDialect",
"@llvm-project//mlir:SideEffectInterfaces",
"@llvm-project//mlir:Support",
"@llvm-project//mlir:TransformUtils",
],
)
cc_library(
name = "utils",
srcs = [
"utils/utils.cc",
],
hdrs = [
"utils/utils.h",
],
deps = [
":tfr",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Support",
],
)
cc_library(
name = "passes",
srcs = [
"passes/canonicalize.cc",
"passes/decompose.cc",
"passes/generated_decompose.inc",
"passes/raise_to_tf.cc",
"passes/rewrite_quantized_io.cc",
],
hdrs = [
"passes/passes.h",
],
deps = [
":tfr",
":utils",
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/core:lib",
"@com_google_absl//absl/memory",
"@com_google_absl//absl/strings",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:AffineUtils",
"@llvm-project//mlir:ArithDialect",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:InliningUtils",
"@llvm-project//mlir:Pass",
"@llvm-project//mlir:QuantOps",
"@llvm-project//mlir:SCFDialect",
"@llvm-project//mlir:SCFToControlFlow",
"@llvm-project//mlir:Support",
"@llvm-project//mlir:TransformUtils",
],
alwayslink = 1,
)
tf_cc_binary(
name = "tfr-opt",
srcs = ["passes/tfr_opt.cc"],
deps = [
":passes",
":tfr",
"//tensorflow/compiler/mlir:init_mlir",
"//tensorflow/compiler/mlir:passes",
"//tensorflow/compiler/mlir/quantization/common/ir:QuantOps",
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/compiler/mlir/tensorflow/transforms:tensorflow_passes",
"@llvm-project//mlir:AllPassesAndDialects",
"@llvm-project//mlir:ArithDialect",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:FuncExtensions",
"@llvm-project//mlir:MlirOptLib",
"@llvm-project//mlir:QuantOps",
"@llvm-project//mlir:RegisterAllPasses",
"@llvm-project//mlir:SCFDialect",
"@llvm-project//mlir:ShapeDialect",
],
)
glob_lit_tests(
name = "all_tests",
data = [":test_utilities"],
driver = "//tensorflow/compiler/mlir:run_lit.sh",
test_file_exts = ["mlir"],
)
# Bundle together all of the test utilities that are used by tests.
filegroup(
name = "test_utilities",
testonly = True,
data = [
":tfr-opt",
"@llvm-project//llvm:FileCheck",
"@llvm-project//llvm:not",
"@llvm-project//mlir:run_lit.sh",
],
)
cc_library(
name = "tfr_decompose_ctx",
srcs = ["integration/tfr_decompose_ctx.cc"],
hdrs = ["integration/tfr_decompose_ctx.h"],
deps = [
":passes",
":tfr",
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/compiler/mlir/tensorflow:convert_attr",
"//tensorflow/compiler/mlir/tensorflow:convert_type",
"//tensorflow/compiler/mlir/tensorflow:mlir_roundtrip_flags",
"//tensorflow/compiler/mlir/tensorflow/transforms:tensorflow_passes",
"//tensorflow/compiler/mlir/tf2xla/api/v2:tf_executor_to_graph",
"//tensorflow/core:framework",
"//tensorflow/core:framework_types_hdr",
"//tensorflow/core:lib",
"//tensorflow/core:lib_internal",
"//tensorflow/core:protos_all_cc",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings",
"@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:Parser",
"@llvm-project//mlir:Pass",
"@llvm-project//mlir:SCFDialect",
"@llvm-project//mlir:ShapeDialect",
],
)
tf_cc_test(
name = "tfr_decompose_ctx_test",
srcs = ["integration/tfr_decompose_ctx_test.cc"],
deps = [
":tfr_decompose_ctx",
"//tensorflow/core:framework",
"//tensorflow/core:ops",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"@com_google_absl//absl/types:span",
"@llvm-project//mlir:IR",
"@xla//xla/hlo/testlib:test",
],
)
tf_python_pybind_extension(
name = "tfr_wrapper",
srcs = ["python/tfr_wrapper.cc"],
enable_stub_generation = True,
pytype_srcs = [
"tfr_wrapper.pyi",
],
starlark_only = True,
visibility = [
"//tensorflow/python:__pkg__",
],
deps = [
":tfr",
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/python/lib/core:pybind11_lib",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:ArithDialect",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Parser",
"@llvm-project//mlir:SCFDialect",
"@llvm-project//mlir:ShapeDialect",
"@pybind11",
],
)
py_library(
name = "composite",
srcs = ["python/composite.py"],
strict_deps = True,
)
py_library(
name = "tfr_gen",
srcs = ["python/tfr_gen.py"],
strict_deps = True,
deps = [
":tfr_wrapper",
"//tensorflow:tensorflow_py", # buildcleaner: keep
"//tensorflow/core:protos_all_py",
"//tensorflow/python/autograph/converters:control_flow",
"//tensorflow/python/autograph/converters:return_statements",
"//tensorflow/python/autograph/impl:api",
"//tensorflow/python/autograph/pyct:anno",
"//tensorflow/python/autograph/pyct:cfg",
"//tensorflow/python/autograph/pyct:qual_names",
"//tensorflow/python/autograph/pyct:transformer",
"//tensorflow/python/autograph/pyct:transpiler",
"//tensorflow/python/autograph/pyct/static_analysis:activity",
"//tensorflow/python/autograph/pyct/static_analysis:reaching_definitions",
"//tensorflow/python/autograph/pyct/static_analysis:reaching_fndefs",
"//tensorflow/python/autograph/pyct/static_analysis:type_inference",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:load_library",
"//tensorflow/python/framework:op_def_registry",
"//tensorflow/python/platform:tf_logging",
"//tensorflow/python/util:tf_inspect",
"@pypi//gast",
],
)
tf_py_strict_test(
name = "tfr_gen_test",
size = "medium",
srcs = ["python/tfr_gen_test.py"],
deps = [
":composite",
":tfr_gen",
"//tensorflow/compiler/mlir/python/mlir_wrapper:filecheck_wrapper",
"//tensorflow/compiler/mlir/tfr/resources:test_ops",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:array_ops_gen",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:math_ops_gen",
"//tensorflow/python/platform:client_testlib",
],
)
py_library(
name = "op_reg_gen",
srcs = ["python/op_reg_gen.py"],
strict_deps = True,
deps = [
"//tensorflow:tensorflow_py",
"//tensorflow/python/autograph/pyct:transformer",
"//tensorflow/python/autograph/pyct:transpiler",
"//tensorflow/python/framework:op_def_registry",
"//tensorflow/python/util:tf_inspect",
"@pypi//gast",
],
)
tf_py_strict_test(
name = "op_reg_gen_test",
size = "small",
srcs = ["python/op_reg_gen_test.py"],
deps = [
":composite",
":op_reg_gen",
"//tensorflow/compiler/mlir/python/mlir_wrapper:filecheck_wrapper",
"//tensorflow/python/platform:client_testlib",
],
)
py_library(
name = "test_utils",
srcs = ["python/test_utils.py"],
strict_deps = True,
deps = [
"//tensorflow:tensorflow_py",
"//tensorflow/python/eager:backprop",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/platform:client_testlib",
],
)
gen_op_libraries(
name = "one_op",
src = "define_op_template.py",
deps = [
"//tensorflow/python/platform:flags",
"@absl_py//absl:app",
],
)
tensorflow/compiler/mlir/tfr/README.md
+168
View File
@@ -0,0 +1,168 @@
# Composable Tensorflow
## Composable Tensorflow
Composable TensorFlow (TF) is the framework for defining portable TF ops with
composition in the authoring language.
The set of standard TF ops is currently open. New ops are defined for special
purposes but it is hard to make them work end-to-end: The op
needs to be handled separately by several backends (tf2xla bridge, tflite
converter, CPU kernels, etc.). Writing shape functions and gradients for these
ops is extremely difficult. `tf.function` makes some parts of the implementation
simpler, but it introduces runtime overhead and it cannot easily be used to
apply dedicated optimizations to op kernels.
The composable TF framework allows the user to define portable TF ops as
compositions of other TF ops. It translates a Python function used to define the
composition directly into a portable IR at build time, and uses it to expand the
composite op in the TF program during compilation / execution. By using this
expansion mechanism, new op are readily available on different platforms without
extra work. Moreover, since the expansion is optional, the backend can easily
treat it as a monolithic op when needed, for instance to apply optimizations or
associate it with a custom kernel.
### Benefits
Using the Composable TF API to define a new op and its composition can bring the
following benefits:
* *Automatic backend support*: As long as it is composed of ops supported by
the backend, the new op is automatically supported (as a `tf.function`
alternative);
* *Reduced tracing overhead*: Unlike `tf.function`, the composition function
is compiled at build time, hence TF only needs to trace a single op to build
the `graph`;
* *Easy fused op/kernel optimization*: Even if it has complex semantics, the
new op is presented as a single node in the graph, thus optimization passes
and kernels can easily be specialized to this op for better performance.
* *Automatic shape/type inference support*: No shape functions are required
for the new op;
* *Automatic gradient support (WIP)*: The user doesn't need to author gradient
a function of the op for training.
### Use Cases
* (Portability) User wants to add a new op and run this op on different
platforms (CPU, TPU, TFLite, etc.) to be portable.
* *Solution*: The user should define the new op as a composition. The ops
used inside the composition should have support for these platforms.
These ops can also be composite ops.
* (Performance) User defines a custom kernel for a regular structure (i.e.
LSTM), but it is hard to add the logic to fuse the individual ops to target
this kernel in the inference graph.
* *Solution*: The user should define a new TF op, which corresponds to the
fused kernel, with composition, and use this op to build the model for
both training and inference. For the platforms where a fused kernel is
not available, the execution will use the composition instead.
## Gradient
(TODO)
## Authoring Op Composition in Python
The composable TF provides a single API to define a new op with its composition
at the same time. For example, the following code defines a new
`FusedFullyConnected` op, which have `MatMul`, `Add` and some
`activation function` (specified by an op attribute) fused.
```python
import tensorflow as tf
@Composite(
'FusedFullyConnected',
inputs=['input_: T', 'filter_: T', 'bias: T'],
attrs=['act: {"", "RELU", "RELU6", "TANH"} = ""'],
derived_attrs=['T: {float, int8}'],
outputs=['o: T'])
def _composite_fully_connected(input_, filter_, bias, act):
res = tf.raw_ops.MatMul(
a=input_, b=filter_, transpose_a=False, transpose_b=True)
res = tf.raw_ops.Add(x=res, y=bias)
if act == 'RELU':
return tf.raw_ops.Relu(features=res)
elif act == 'RELU6':
return tf.raw_ops.Relu6(features=res)
elif act == 'TANH':
return tf.raw_ops.Tanh(x=res)
else:
return res
```
Besides defining new ops, composition can be specified for an existing op
for portability. The following code defines the semantics of `AddNOp`:
```python
@Composite('AddNOp')
def _my_op_c(ins):
N = len(ins)
if N == 1:
return ins[0]
sum = ins[0]
for i in range(1, N):
sum += ins[i]
return sum
```
Utilities have been built to compile the Python composition functions down to
the backend IR. The project also provides a set of graph optimization passes to
expand the composite ops in the graph by using the input backend IR. These
passes have been added to the TF
[common runtime](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/common_runtime)
for graph execution and
[eager runtime](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/common_runtime/eager)
for eager execution.
## Compiling Op Composition
### Ahead-Of-Time (AOT) mode
Like the op kernels, the op composition can be pre-compiled to the backend IR
so the decomposition can be invoked at runtime. A Python [define_op_template.py](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/compiler/mlir/tfr/define_op_template.py)
file is provided as an example to build composite ops in the users project
directory. All the targets required to build the new ops are created by the
following target:
```BUILD
load("//tensorflow/compiler/mlir/tfr:build_defs.bzl", "gen_op_libraries")
gen_op_libraries(
name = "test_ops",
src = "define_op_template.py",
deps = [
"//third_party/py/tensorflow",
],
)
```
More composite op definitions and usages are here included in the
[examples](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/compiler/mlir/tfr/examples)
directory.
### Just-In-Time (JIT) mode
(TODO)
## Known Limitations
* `while` statement
* condition of `if` statement couldn't be a tensor
## RFC
This project is an alternative implementation of
[RFC:Standardizing composite ops in tensorflow to support efficient inference](https://github.com/tensorflow/community/blob/master/rfcs/20190610-standardizing-composite_ops.md).
This project doesn't rely on the tracing functionality provided by `tf.function`
to avoid all its pitfalls and it helps to build more general transformations in
the backends.
## Team
* Feng Liu
* Dan Moldovan
tensorflow/compiler/mlir/tfr/build_defs.bzl
+165
View File
@@ -0,0 +1,165 @@
"""BUILD extension for TF composition project."""
load("@rules_cc//cc:cc_library.bzl", "cc_library")
load("@rules_ml_toolchain//py/rules_pywrap:pywrap.default.bzl", "use_pywrap_rules")
load("@xla//third_party/rules_python/python:py_binary.bzl", "py_binary")
load("@xla//third_party/rules_python/python:py_library.bzl", "py_library")
load("//tensorflow:tensorflow.bzl", "tf_custom_op_library", "tf_gen_op_wrapper_py")
load("//tensorflow:tensorflow.default.bzl", "tf_custom_op_py_library")
# TODO(b/356020232): cleanup use_pywrap_rules once migration is done
def gen_op_libraries(
name,
src,
deps = [],
tags = [],
test = False):
"""gen_op_libraries() generates all cc and py libraries for composite op source.
Args:
name: used as the name component of all the generated libraries.
src: File contains the composite ops.
deps: Libraries the 'src' depends on.
tags:
test:
"""
if not src.endswith(".py") or name == src[:-3]:
fail("'src' %s conflicts with op Python wrapper. Rename it to be different from 'name'." % src)
py_deps = []
if use_pywrap_rules():
py_deps = ["//tensorflow/python:_pywrap_tensorflow"]
py_deps += [
"//tensorflow/compiler/mlir/tfr:op_reg_gen",
"//tensorflow/compiler/mlir/tfr:tfr_gen",
"//tensorflow/compiler/mlir/tfr:composite",
] + deps
gen_op_lib_exec = src[:-3] # Strip off the .py
py_binary(
name = gen_op_lib_exec,
srcs = [src],
srcs_version = "PY3",
python_version = "PY3",
deps = py_deps,
)
registered_op = "registered_" + name
native.genrule(
name = registered_op,
srcs = [],
outs = [name + ".inc.cc"],
cmd =
"$(location %s) --output=$@ --gen_register_op=true" % gen_op_lib_exec,
tools = [":" + gen_op_lib_exec],
tags = tags,
)
cc_library(
name = name + "_cc",
testonly = test,
srcs = [":" + registered_op],
deps = [
"//tensorflow/core:framework",
"//tensorflow/core:lib",
"//tensorflow/core:protos_all_cc",
],
alwayslink = 1,
)
tf_custom_op_library(
name = name + ".so",
srcs = [":" + registered_op],
)
tf_gen_op_wrapper_py(
name = "gen_" + name,
out = "gen_" + name + ".py",
py_lib_rule = py_library,
deps = [
":%s_cc" % name,
],
extra_py_deps = [
"//tensorflow/python:pywrap_tfe",
"//tensorflow/python/util:dispatch",
"//tensorflow/python/util:deprecation",
"//tensorflow/python/util:tf_export",
] + py_deps,
)
tf_custom_op_py_library(
name = name,
dso = [":%s.so" % name],
kernels = [":%s_cc" % name],
srcs_version = "PY3",
# copybara:uncomment(OSS version passes this to py_library) lib_rule = py_library,
deps = [
":gen_%s" % name,
],
)
# Link the register op and rebuild the binary
gen_tfr_lib_exec = gen_op_lib_exec + "_with_op_library"
py_binary(
name = gen_tfr_lib_exec,
main = src,
srcs = [src],
python_version = "PY3",
srcs_version = "PY3",
deps = py_deps + [":%s" % name],
)
native.genrule(
name = name + "_mlir",
srcs = [],
outs = [name + ".mlir"],
cmd =
"$(location %s) --output=$@ --gen_register_op=false" % gen_tfr_lib_exec,
tools = [":" + gen_tfr_lib_exec],
tags = tags,
)
py_library(
name = name + "_py",
srcs = [src],
srcs_version = "PY3",
deps = py_deps,
)
def gen_op_bindings(name):
cc_library(
name = name + "_ops_cc",
srcs = [name + "_ops.cc"],
deps = [
"//tensorflow/core:framework",
"//tensorflow/core:lib",
"//tensorflow/core:protos_all_cc",
],
alwayslink = 1,
)
tf_custom_op_library(
name = name + "_ops.so",
srcs = [name + "_ops.cc"],
)
tf_gen_op_wrapper_py(
name = "gen_" + name + "_ops",
out = "gen_" + name + "_ops.py",
py_lib_rule = py_library,
deps = [":" + name + "_ops_cc"],
extra_py_deps = [
"//tensorflow/python:pywrap_tfe",
"//tensorflow/python/util:dispatch",
"//tensorflow/python/util:deprecation",
"//tensorflow/python/util:tf_export",
],
)
tf_custom_op_py_library(
name = name + "_ops",
dso = [":" + name + "_ops.so"],
kernels = [":" + name + "_ops_cc"],
# copybara:uncomment(OSS version passes this to py_library) lib_rule = py_library,
deps = [":gen_" + name + "_ops"],
)
tensorflow/compiler/mlir/tfr/define_op_template.py
+60
View File
@@ -0,0 +1,60 @@
# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""A template to define composite ops."""
# pylint: disable=g-direct-tensorflow-import
import os
import sys
from absl import app
from tensorflow.compiler.mlir.tfr.python.composite import Composite
from tensorflow.compiler.mlir.tfr.python.op_reg_gen import gen_register_op
from tensorflow.compiler.mlir.tfr.python.tfr_gen import tfr_gen_from_module
from tensorflow.python.platform import flags
FLAGS = flags.FLAGS
flags.DEFINE_string(
'output', None,
'Path to write the genereated register op file and MLIR file.')
flags.DEFINE_bool('gen_register_op', True,
'Generate register op cc file or tfr mlir file.')
flags.mark_flag_as_required('output')
@Composite('TestRandom', derived_attrs=['T: numbertype'], outputs=['o: T'])
def _composite_random_op():
pass
def main(_):
if FLAGS.gen_register_op:
assert FLAGS.output.endswith('.cc')
generated_code = gen_register_op(sys.modules[__name__], '_composite_')
else:
assert FLAGS.output.endswith('.mlir')
generated_code = tfr_gen_from_module(sys.modules[__name__], '_composite_')
dirname = os.path.dirname(FLAGS.output)
if not os.path.exists(dirname):
os.makedirs(dirname)
with open(FLAGS.output, 'w') as f:
f.write(generated_code)
if __name__ == '__main__':
app.run(main=main)
tensorflow/compiler/mlir/tfr/integration/tfr_decompose_ctx.cc
+247
View File
@@ -0,0 +1,247 @@
/* 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/tfr/integration/tfr_decompose_ctx.h"
#include <memory>
#include <string>
#include <utility>
#include <vector>
#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/string_view.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/ADT/StringRef.h"
#include "llvm/Support/LogicalResult.h"
#include "llvm/Support/MemoryBuffer.h"
#include "llvm/Support/SMLoc.h"
#include "llvm/Support/SourceMgr.h"
#include "mlir/Dialect/Arith/IR/Arith.h" // from @llvm-project
#include "mlir/Dialect/Func/Extensions/AllExtensions.h" // from @llvm-project
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/Dialect/SCF/IR/SCF.h" // from @llvm-project
#include "mlir/Dialect/Shape/IR/Shape.h" // from @llvm-project
#include "mlir/IR/Attributes.h" // from @llvm-project
#include "mlir/IR/Builders.h" // from @llvm-project
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/IR/BuiltinTypes.h" // from @llvm-project
#include "mlir/IR/Location.h" // from @llvm-project
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "mlir/IR/OperationSupport.h" // from @llvm-project
#include "mlir/IR/OwningOpRef.h" // from @llvm-project
#include "mlir/IR/Types.h" // from @llvm-project
#include "mlir/IR/Verifier.h" // from @llvm-project
#include "mlir/Parser/Parser.h" // from @llvm-project
#include "mlir/Pass/PassManager.h" // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_device.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_dialect.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_executor.h"
#include "tensorflow/compiler/mlir/tensorflow/transforms/passes.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/convert_attr.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/convert_type.h"
#include "tensorflow/compiler/mlir/tf2xla/api/v2/tf_executor_to_graph.h"
#include "tensorflow/compiler/mlir/tfr/ir/tfr_ops.h"
#include "tensorflow/compiler/mlir/tfr/passes/passes.h"
#include "xla/tsl/platform/errors.h"
#include "xla/tsl/platform/statusor.h"
#include "tensorflow/core/framework/function.pb.h"
#include "tensorflow/core/framework/node_def.pb.h"
#include "tensorflow/core/framework/node_def_util.h"
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/op_def.pb.h"
#include "tensorflow/core/framework/types.h"
#include "tensorflow/core/platform/env.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/platform/path.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/platform/stringpiece.h"
#include "tensorflow/core/platform/types.h"
#include "tensorflow/core/util/env_var.h"
namespace tensorflow {
namespace tfr {
const char* const kTFRLibEnv = "TF_MLIR_TFR_LIB_DIR";
absl::StatusOr<std::unique_ptr<TFRDecomposeContext>> TFRDecomposeContext::Get(
mlir::MLIRContext* mlir_ctx) {
Env* env = Env::Default();
std::string tfr_lib_dir;
TF_RETURN_IF_ERROR(ReadStringFromEnvVar(
kTFRLibEnv, "tensorflow/compiler/mlir/tfr/resources", &tfr_lib_dir));
std::string composite_mlir_dir =
io::JoinPath(env->GetRunfilesDir(), tfr_lib_dir);
std::vector<std::string> files;
TF_RETURN_IF_ERROR(env->GetChildren(composite_mlir_dir, &files));
if (files.empty()) {
return absl::InternalError(absl::StrCat(
"Failed to find the decomposition lib from path ", composite_mlir_dir));
}
std::string tfr_raw_text;
for (const auto& file : files) {
std::string fullpath = io::JoinPath(composite_mlir_dir, file);
if (env->MatchPath(fullpath, io::JoinPath(composite_mlir_dir, "*.mlir"))) {
std::string text;
TF_RETURN_IF_ERROR(ReadFileToString(env, fullpath, &text));
tfr_raw_text.append(text);
}
}
auto ctx = TFRDecomposeContext::GetFromText(tfr_raw_text, mlir_ctx);
if (!ctx) {
return absl::InternalError(absl::StrCat(
"Failed to load the imported decomposition lib: ", tfr_raw_text));
}
return ctx;
}
std::unique_ptr<TFRDecomposeContext> TFRDecomposeContext::GetFromText(
absl::string_view tfr_raw_text, mlir::MLIRContext* mlir_ctx) {
mlir_ctx->allowUnregisteredDialects(/*allow=*/true);
// Load dialects involved in the conversion
mlir::DialectRegistry registry;
// clang-format off
registry.insert<mlir::arith::ArithDialect,
mlir::func::FuncDialect,
mlir::scf::SCFDialect,
mlir::shape::ShapeDialect,
mlir::TF::TensorFlowDialect,
mlir::tf_device::TensorFlowDeviceDialect,
mlir::tf_executor::TensorFlowExecutorDialect,
mlir::TFR::TFRDialect>();
// clang-format on
mlir::func::registerAllExtensions(registry);
mlir_ctx->appendDialectRegistry(registry);
mlir_ctx->loadAllAvailableDialects();
// Load the TFR functions in a mlir::ModuleOp
auto memory_buffer = llvm::MemoryBuffer::getMemBuffer(
llvm::StringRef(tfr_raw_text.data(), tfr_raw_text.size()));
llvm::SourceMgr source_mgr;
source_mgr.AddNewSourceBuffer(std::move(memory_buffer), llvm::SMLoc());
mlir::OwningOpRef<mlir::ModuleOp> module =
mlir::parseSourceFile<mlir::ModuleOp>(source_mgr, mlir_ctx);
// The MLIRContext owns the module
auto module_op = module.release();
// Create the context
return std::make_unique<TFRDecomposeContext>(module_op);
}
absl::StatusOr<FunctionDef> TFRDecomposeContext::ExpandNode(
const NodeDef& node_def, absl::string_view func_name) {
const OpDef* op_def;
TF_RETURN_IF_ERROR(OpRegistry::Global()->LookUpOpDef(node_def.op(), &op_def));
DataTypeVector input_dtys, output_dtys;
TF_RETURN_IF_ERROR(InputTypesForNode(node_def, *op_def, &input_dtys));
TF_RETURN_IF_ERROR(OutputTypesForNode(node_def, *op_def, &output_dtys));
mlir::MLIRContext* context = tfr_module_.getContext();
llvm::SmallVector<mlir::Type, 4> input_tys, output_tys;
mlir::Builder builder(context);
for (auto ty : input_dtys) {
mlir::Type elt_ty;
TF_RETURN_IF_ERROR(ConvertDataType(ty, builder, &elt_ty));
mlir::TensorType mlir_ty = mlir::UnrankedTensorType::get(elt_ty);
input_tys.push_back(mlir_ty);
}
for (auto ty : output_dtys) {
mlir::Type elt_ty;
TF_RETURN_IF_ERROR(ConvertDataType(ty, builder, &elt_ty));
mlir::TensorType mlir_ty = mlir::UnrankedTensorType::get(elt_ty);
output_tys.push_back(mlir_ty);
}
llvm::SmallVector<mlir::NamedAttribute, 4> attrs;
for (const auto& attr : node_def.attr()) {
TF_ASSIGN_OR_RETURN(auto mlir_attr,
ConvertAttributeValue(attr.second, &builder));
attrs.push_back({mlir::StringAttr::get(context, attr.first), mlir_attr});
}
mlir::Location loc = mlir::UnknownLoc::get(context);
mlir::ModuleOp module = mlir::ModuleOp::create(loc);
mlir::FunctionType func_type =
mlir::FunctionType::get(context, input_tys, output_tys);
llvm::StringRef func_name_str(func_name.data(), func_name.size());
auto func = mlir::func::FuncOp::create(loc, func_name_str, func_type, {});
module.push_back(func);
func.addEntryBlock();
mlir::OpBuilder op_builder(func.getBody());
// Create the TF op
const std::string tf_op_full_name = absl::StrCat("tf.", node_def.op());
mlir::OperationState op_state(loc, tf_op_full_name);
op_state.addOperands(func.getArguments());
op_state.addTypes(output_tys);
op_state.addAttributes(attrs);
mlir::Operation* tf_op = op_builder.create(op_state);
mlir::func::ReturnOp::create(op_builder, loc, tf_op->getResults());
// Run the decompose passes on the module
TF_RETURN_IF_ERROR(DecomposeGraph(module));
// Export the result as a FunctionDef.
FunctionDef func_def;
TF_RETURN_IF_ERROR(
tensorflow::tf2xla::v2::ConvertMlirFunctionToFunctionLibraryDef(
func, export_confs_, &func_def));
module.erase();
return func_def;
}
absl::Status TFRDecomposeContext::DecomposeGraph(mlir::ModuleOp user_module) {
// Call the decompose passes by using the external symbol table.
if (failed(pm_.run(user_module))) {
return absl::InternalError("Failed to run the decompose passes.");
}
return absl::OkStatus();
}
// Constructor of the decompose context.
TFRDecomposeContext::TFRDecomposeContext(mlir::ModuleOp tfr_module)
: tfr_module_(tfr_module), pm_(tfr_module_.getContext()) {
mlir::OpPassManager& func_pm = pm_.nest<mlir::func::FuncOp>();
// Prepare the imported graph.
func_pm.addPass(mlir::CreateExecutorDialectToFunctionalConversionPass());
// Run TFR lowering, inlining and raising to tf.
func_pm.addPass(mlir::TFR::CreateDecomposeTFOpsPass(tfr_module_));
func_pm.addPass(mlir::TFR::CreateRaiseToTFOpsPass(
tfr_module_, /*materialize_derived_attrs=*/true));
// Prepare to be exported.
func_pm.addPass(mlir::CreateFunctionalToExecutorDialectConversionPass());
pm_.addPass(mlir::CreateBreakUpIslandsPass());
}
void TFRDecomposeContext::Destroy() { tfr_module_.erase(); }
absl::StatusOr<FunctionDef> ExpandNode(const NodeDef& node_def,
absl::string_view func_name) {
mlir::MLIRContext mlir_ctx;
TF_ASSIGN_OR_RETURN(auto ctx, TFRDecomposeContext::Get(&mlir_ctx));
return ctx->ExpandNode(node_def, func_name);
}
absl::Status DecomposeGraph(mlir::ModuleOp user_module) {
mlir::MLIRContext* mlir_ctx = user_module.getContext();
TF_ASSIGN_OR_RETURN(auto ctx, TFRDecomposeContext::Get(mlir_ctx));
return ctx->DecomposeGraph(user_module);
}
} // namespace tfr
} // namespace tensorflow
tensorflow/compiler/mlir/tfr/integration/tfr_decompose_ctx.h
+87
View File
@@ -0,0 +1,87 @@
/* 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_TFR_INTEGRATION_TFR_DECOMPOSE_CTX_H_
#define TENSORFLOW_COMPILER_MLIR_TFR_INTEGRATION_TFR_DECOMPOSE_CTX_H_
#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "absl/strings/string_view.h"
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "mlir/Pass/PassManager.h" // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/translate/mlir_roundtrip_flags.h"
#include "xla/tsl/platform/statusor.h"
#include "tensorflow/core/framework/function.pb.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/platform/stringpiece.h"
namespace tensorflow {
namespace tfr {
extern const char* const kTFRLibEnv;
using tsl::StatusOr;
// An wrapper for all the objects used to decompose a module (graph mode) and
// node_def (eager mode). Note that this class owns the decomposition library.
class TFRDecomposeContext {
public:
// The entry function to get a decompose context. All the required passes have
// been initialized.
static absl::StatusOr<std::unique_ptr<TFRDecomposeContext>> Get(
mlir::MLIRContext* mlir_ctx);
// Constructor of the decompose context. To share the decompose library, the
// whole decompose TFR function library is loaded.
explicit TFRDecomposeContext(mlir::ModuleOp tfr_module);
// Constructs the decompose context from the tfr text module and the mlir
// context. The tfr text module is added to the mlir context.
static std::unique_ptr<TFRDecomposeContext> GetFromText(
absl::string_view tfr_raw_text, mlir::MLIRContext* mlir_ctx);
// Decomposes the op in the NodeDef to a set of primitive ops according to the
// decompose library in the context. Wrap the decomposed result in a
// FunctionDef.
absl::StatusOr<FunctionDef> ExpandNode(const NodeDef& node_def,
absl::string_view func_name);
// Runs the decompose passes on the user_module.
absl::Status DecomposeGraph(mlir::ModuleOp user_module);
// Erases the tfr_module created.
void Destroy();
private:
mlir::ModuleOp tfr_module_;
mlir::PassManager pm_;
GraphExportConfig export_confs_;
};
// Decomposes the NodeDef to a set of primitive ops according to the decompose
// library loaded. Wrap the decomposed result in a FunctionDef.
absl::StatusOr<FunctionDef> ExpandNode(const NodeDef& node_def,
absl::string_view func_name);
// Decomposes the ops in the ModuleOp to a set of primitive ops according to
// decompose library in the context.
absl::Status DecomposeGraph(mlir::ModuleOp user_module);
} // namespace tfr
} // namespace tensorflow
#endif // TENSORFLOW_COMPILER_MLIR_TFR_INTEGRATION_TFR_DECOMPOSE_CTX_H_
tensorflow/compiler/mlir/tfr/integration/tfr_decompose_ctx_test.cc
+159
View File
@@ -0,0 +1,159 @@
/* 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/tfr/integration/tfr_decompose_ctx.h"
#include <memory>
#include <string>
#include <utility>
#include <vector>
#include "absl/types/span.h"
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "xla/hlo/testlib/test.h"
#include "tensorflow/core/framework/attr_value.pb.h"
#include "tensorflow/core/framework/common_shape_fns.h"
#include "tensorflow/core/framework/function.pb.h"
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/framework/node_def_builder.h"
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/platform/test.h"
using testing::ElementsAreArray;
using testing::Test;
using NodeAndType = std::pair<std::string, tensorflow::DataType>;
namespace tensorflow {
namespace {
REGISTER_OP("MyAddN")
.Input("inputs: N * T")
.Output("sum: T")
.Attr("N: int >= 1")
.Attr("T: {numbertype, variant}")
.SetIsCommutative()
.SetIsAggregate()
.SetShapeFn(shape_inference::UnchangedShape);
REGISTER_OP("RiscAddDummy")
.Input("x: T")
.Input("y: T")
.Output("z: T")
.Attr(
"T: {bfloat16, half, float, double, uint8, int8, int16, int32, int64, "
"complex64, complex128, string}")
.SetShapeFn(shape_inference::UnchangedShape);
constexpr char tfr_raw_text[] = R"(
tfr.func @tf__my_add_n(%values: !tfr.tensor_list,
%n: i64 {tfr.name="N"}) -> !tfr.tensor {
%index = arith.constant 0 : index
%cst = arith.constant 1 : i64
%eq = arith.cmpi "eq", %n, %cst : i64
%v1 = tfr.get_element %values[%index] : (!tfr.tensor_list, index) -> !tfr.tensor
%res = scf.if %eq -> !tfr.tensor {
scf.yield %v1 : !tfr.tensor
} else {
%step = arith.index_cast %cst : i64 to index
%end = arith.index_cast %n : i64 to index
%reduce = scf.for %i = %step to %end step %step iter_args(%reduce_iter=%v1) -> !tfr.tensor {
%v = tfr.get_element %values[%i] : (!tfr.tensor_list, index) -> !tfr.tensor
%reduce_next = tfr.call @tf__risc_add_dummy(%reduce_iter, %v) : (!tfr.tensor, !tfr.tensor) -> !tfr.tensor
scf.yield %reduce_next : !tfr.tensor
}
scf.yield %reduce : !tfr.tensor
}
tfr.return %res : !tfr.tensor
}
tfr.func @tf__my_add_n_(!tfr.tensor_list<N,T>, i64 {tfr.name="N"}) -> !tfr.tensor attributes{N,T}
tfr.func @tf__risc_add_dummy_(!tfr.tensor<T>, !tfr.tensor<T>) -> !tfr.tensor<T> attributes{T}
)";
class TFRDecomposeContextTest : public Test {
protected:
void SetUp() override {
test_ctx_ = tfr::TFRDecomposeContext::GetFromText(tfr_raw_text, &ctx_);
}
void TearDown() override { test_ctx_->Destroy(); }
mlir::MLIRContext ctx_;
std::unique_ptr<tfr::TFRDecomposeContext> test_ctx_;
};
std::vector<NodeAndType> NodesSequenceOf(const FunctionDef& graph) {
std::vector<NodeAndType> nodes;
for (auto& node : graph.node_def()) {
nodes.push_back({node.op(), node.attr().at("T").type()});
}
return nodes;
}
TEST_F(TFRDecomposeContextTest, FLOAT_1_ins) {
std::vector<NodeDefBuilder::NodeOut> src_list;
src_list.emplace_back("input", 0, DT_FLOAT);
NodeDef test_node;
auto status = NodeDefBuilder("float_add", "MyAddN")
.Input(src_list)
.Finalize(&test_node);
EXPECT_TRUE(status.ok());
auto decomposed = test_ctx_->ExpandNode(test_node, "test");
EXPECT_TRUE(decomposed.ok());
std::vector<NodeAndType> expected_results{{"Identity", DT_FLOAT}};
EXPECT_THAT(NodesSequenceOf(decomposed.value()),
ElementsAreArray(expected_results));
}
TEST_F(TFRDecomposeContextTest, FLOAT_3_ins) {
std::vector<NodeDefBuilder::NodeOut> src_list;
src_list.emplace_back("in0", 0, DT_FLOAT);
src_list.emplace_back("in1", 0, DT_FLOAT);
src_list.emplace_back("in2", 0, DT_FLOAT);
NodeDef test_node;
auto status = NodeDefBuilder("float_add_3", "MyAddN")
.Input(src_list)
.Finalize(&test_node);
EXPECT_TRUE(status.ok());
auto decomposed = test_ctx_->ExpandNode(test_node, "test");
EXPECT_TRUE(decomposed.ok());
std::vector<NodeAndType> expected_results{{"RiscAddDummy", DT_FLOAT},
{"RiscAddDummy", DT_FLOAT}};
EXPECT_THAT(NodesSequenceOf(decomposed.value()),
ElementsAreArray(expected_results));
}
TEST_F(TFRDecomposeContextTest, INT32_3_ins) {
std::vector<NodeDefBuilder::NodeOut> src_list;
src_list.emplace_back("in0", 0, DT_INT32);
src_list.emplace_back("in1", 0, DT_INT32);
src_list.emplace_back("in2", 0, DT_INT32);
NodeDef test_node;
auto status =
NodeDefBuilder("int_add", "MyAddN").Input(src_list).Finalize(&test_node);
EXPECT_TRUE(status.ok());
auto decomposed = test_ctx_->ExpandNode(test_node, "test");
EXPECT_TRUE(decomposed.ok());
std::vector<NodeAndType> expected_results{{"RiscAddDummy", DT_INT32},
{"RiscAddDummy", DT_INT32}};
EXPECT_THAT(NodesSequenceOf(decomposed.value()),
ElementsAreArray(expected_results));
}
} // namespace
} // namespace tensorflow
tensorflow/compiler/mlir/tfr/ir/tfr_ops.cc
+1037
View File
File diff suppressed because it is too large Load Diff
tensorflow/compiler/mlir/tfr/ir/tfr_ops.h
+66
View File
@@ -0,0 +1,66 @@
/* 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_TFR_IR_TFR_OPS_H_
#define TENSORFLOW_COMPILER_MLIR_TFR_IR_TFR_OPS_H_
#include "llvm/ADT/StringSet.h"
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/Dialect/Shape/IR/Shape.h" // from @llvm-project
#include "mlir/IR/Attributes.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/DialectImplementation.h" // from @llvm-project
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "mlir/IR/Operation.h" // from @llvm-project
#include "mlir/IR/Types.h" // from @llvm-project
#include "mlir/Interfaces/CallInterfaces.h" // from @llvm-project
#include "mlir/Interfaces/ControlFlowInterfaces.h" // from @llvm-project
#include "mlir/Interfaces/FunctionInterfaces.h" // from @llvm-project
#include "mlir/Interfaces/InferTypeOpInterface.h" // from @llvm-project
#include "mlir/Interfaces/SideEffectInterfaces.h" // from @llvm-project
#include "mlir/Support/LLVM.h" // from @llvm-project
namespace mlir {
namespace TFR {
constexpr char kAttrArgumentNameAttr[] = "tfr.name";
constexpr char kAttrArgumentDefaultAttr[] = "tfr.default";
constexpr char kAttrArgumentTypeAttr[] = "tfr.type";
class TFRDialect : public Dialect {
public:
explicit TFRDialect(MLIRContext *context);
static StringRef getDialectNamespace() { return "tfr"; }
Operation *materializeConstant(OpBuilder &builder, Attribute value, Type type,
Location loc) override;
// Parse a type registered to this dialect.
Type parseType(DialectAsmParser &parser) const override;
// Prints a type registered to this dialect.
void printType(Type ty, DialectAsmPrinter &os) const override;
};
} // namespace TFR
} // namespace mlir
#define GET_OP_CLASSES
#include "tensorflow/compiler/mlir/tfr/ir/tfr_ops.h.inc"
#endif // TENSORFLOW_COMPILER_MLIR_TFR_IR_TFR_OPS_H_
tensorflow/compiler/mlir/tfr/ir/tfr_ops.td
+615
View File
@@ -0,0 +1,615 @@
/* 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.
==============================================================================*/
// This is the operation definition file for TFR
#ifndef DIALECT_TFR_OPS_
#define DIALECT_TFR_OPS_
include "mlir/Dialect/Shape/IR/ShapeBase.td"
include "mlir/IR/OpBase.td"
include "mlir/Interfaces/FunctionInterfaces.td"
include "mlir/Dialect/Quant/IR/QuantBase.td"
include "mlir/IR/SymbolInterfaces.td"
include "mlir/Interfaces/CallInterfaces.td"
include "mlir/Interfaces/ControlFlowInterfaces.td"
include "mlir/Interfaces/InferTypeOpInterface.td"
include "mlir/Interfaces/SideEffectInterfaces.td"
include "tensorflow/compiler/mlir/tensorflow/ir/tf_op_base.td"
//===----------------------------------------------------------------------===//
// Dialect
//===----------------------------------------------------------------------===//
def TFR_Dialect : Dialect {
let name = "tfr";
let description = [{
The TensorFlow Composition dialect.
}];
let cppNamespace = "::mlir::TFR";
}
//===----------------------------------------------------------------------===//
// Type classes
//===----------------------------------------------------------------------===//
// tensor argument types
class TFR_Type<string name> : DialectType<TFR_Dialect,
CPred<"llvm::isa<mlir::TFR::" # name # "Type>($_self)">,
"TFR " # name #" type">,
BuildableType<"$_builder.getType<mlir::TFR::" # name # "Type>()">;
def TFR_TensorType : TFR_Type<"TFRTensor">;
def TFR_TensorListType : TFR_Type<"TFRTensorList">;
def TFR_AllTensorTypes : Type<Or<[
TFR_TensorType.predicate,
TFR_TensorListType.predicate]>, "all tensor related types">;
// attribute argument types
def TFR_AttrType : TFR_Type<"TFRAttr">;
def TFR_AttrScalarType: TypeAlias<TF_ElementType, "scalar attribute">;
def TFR_AttrVectorType : VectorOfNonZeroRankOf<[TF_ElementType, TFR_AttrType]>;
def TFR_AllAttrTypes : Type<Or<[
TFR_AttrType.predicate,
Index.predicate,
TFR_AttrScalarType.predicate,
TFR_AttrVectorType.predicate]>, "all attribute related types">;
// all allowed arguments types
def TFR_allowedArgType : Type<Or<[
TFR_AllTensorTypes.predicate,
TFR_AllAttrTypes.predicate]>, "allowed tfr.call operand types">;
def TFR_allowedConstValues : Attr<Or<[
FlatSymbolRefAttr.predicate,
TypeAttr.predicate,
StrAttr.predicate,
ArrayAttr.predicate]>, "allowed tfr.constant value"> {
let storageType = "Attribute";
let returnType = "Attribute";
let convertFromStorage = "$_self";
let constBuilderCall = "$0";
}
// all allowed result types
def TFR_allowedResultType : TypeAlias<TFR_AllTensorTypes,
"allowed tfr.call result types">;
// standard tensor type and tfr.tensor types can be casted to each other.
def TFR_singleTensorType : Type<Or<[
TFR_TensorType.predicate,
TF_Tensor.predicate,
TensorOf<[quant_QuantizedType]>.predicate]>, "single tensor or tfr.tensor type">;
// all allowed build list input types
def TFR_allowedBuiltListType : Type<Or<[
TFR_TensorType.predicate,
TF_ElementType.predicate,
TFR_AttrType.predicate]>, "single tfr.tensor or tensor element type">;
// all allowed build list result types
def TFR_allowedListResultType : Type<Or<[
TFR_TensorListType.predicate,
TFR_AttrType.predicate]>, "tfr.tensor_list or tfr.attr type">;
//===----------------------------------------------------------------------===//
// Op classes
//===----------------------------------------------------------------------===//
class TFR_Op<string mnemonic, list<Trait> traits> :
Op<TFR_Dialect, mnemonic, traits>;
def TFR_CallOp : TFR_Op<"call", [CallOpInterface]> {
let description = [{
The `call` operation represents a direct call to a function that is within
the same symbol scope as the callee. The operands and result types of the
call must match the specified function type. The callee is encoded as a
symbol reference attribute named "callee".
Example:
```mlir
%2 = tfr.call @my_add(%0, %1) : (tfr.tensor, f32) -> tfr.tensor_list
```
Note that the operands of the `call` operation can only be with tfr.tensor,
tfr.tensor_list, tfr.attr and mlir float and integer types. The results of
the `call` operation can only be with tfr.tensor and tfr.tensor_list types.
}];
let arguments = (ins
FlatSymbolRefAttr:$callee,
Variadic<TFR_allowedArgType>:$args,
OptionalAttr<DictArrayAttr>:$arg_attrs,
OptionalAttr<DictArrayAttr>:$res_attrs);
let results = (outs
Variadic<TFR_allowedResultType>:$outs);
let extraClassDeclaration = [{
// Get the argument operands to the called function.
operand_range getArgOperands() { return getArgs(); }
MutableOperandRange getArgOperandsMutable() {
return getArgsMutable();
}
// Return the callee of this operation.
CallInterfaceCallable getCallableForCallee() { return getCalleeAttr(); }
// Sets the callee from the callable
void setCalleeFromCallable(CallInterfaceCallable callee);
}];
let assemblyFormat = [{
$callee `(` $args `)` attr-dict `:` functional-type($args, results)
}];
}
def TFR_CastOp : TFR_Op<"cast", [Pure]> {
let description = [{
The `cast` operation converts the operand with built-in tensor type to
tfr.tensor type, or vice versa.
Example:
```mlir
%1 = tfr.cast(%0) : tensor<f32> -> !tfr.tensor
%3 = tfr.cast(%1) : !tfr.tensor -> tensor<f32>
```
}];
let arguments = (ins TFR_singleTensorType:$arg);
let results = (outs TFR_singleTensorType:$out);
let extraClassDeclaration = [{
// Return element type of the input tensor type. Only available when the
// input is a MLIR built-in tensor type.
Attribute getInputElementType() {
if (auto ty = llvm::dyn_cast<TensorType>(getArg().getType())) {
return TypeAttr::get(ty.getElementType());
}
return {};
}
}];
let hasCanonicalizer = 1;
}
def TFR_GetShapeOp : TFR_Op<"get_shape", [Pure]> {
let description = [{
The `get_shape` operation gets the shape of a tfr.tensor and returns
!shape.shape type.
Example:
```mlir
%1 = "tfr.get_shape"(%0) : !tfr.tensor -> !shape.shape
%1 = tfr.get_shape %0 -> !shape.shape
```
}];
let arguments = (ins TFR_TensorType:$arg);
let results = (outs Shape_ShapeType:$out);
let assemblyFormat = "$arg attr-dict `->` type($out)";
let hasCanonicalizer = 1;
}
def TFR_GetElementTypeOp : TFR_Op<"get_element_type", [Pure]> {
let description = [{
The `get_element_type` operation gets the element type of a tfr.tensor and
returns !tfr.attr.
Example:
```mlir
%1 = "tfr.get_element_type"(%0) : !tfr.tensor -> !tfr.attr
%1 = tfr.get_element_type %0 -> !tfr.attr
```
}];
let arguments = (ins TFR_TensorType:$arg);
let results = (outs TFR_AttrType:$out);
let assemblyFormat = "$arg attr-dict `->` type($out)";
}
def TFR_EqualOp : TFR_Op<"equal", [Pure, SameTypeOperands]> {
let description = [{
The `equal` operation compares the values of the tfr.attr type arguments.
The operation returns an i1 boolean indicating if the two values are the
same.
Example:
```mlir
%x = tfr.equal %lhs, %rhs -> i1
%x = "tfr.equal"(%lhs, %rhs) : (!tfr.attr, !tfr.attr) -> i1
```
}];
let arguments = (ins
TFR_AttrType:$lhs,
TFR_AttrType:$rhs
);
let results = (outs BoolLike:$result);
let hasFolder = 1;
let assemblyFormat = "$lhs `,` $rhs attr-dict `->` type($result)";
}
def TFR_ConstOp : TFR_Op<"constant", [ConstantLike, Pure]> {
let description = [{
The `attr` operation stores TF op's attribute, which doesn't support
arithmetic operations.
Example:
```mlir
%1 = "tfr.constant"() { value: i32 } : () -> !tfr.attr
%2 = "tfr.constant"() { value: [i32, f32] } : () -> !tfr.attr
%3 = tfr.constant [i32, f32] -> !tfr.attr
%4 = tfr.constant f32 -> !tfr.attr
```
}];
let arguments = (ins TFR_allowedConstValues:$value);
let results = (outs TFR_AttrType:$out);
let hasFolder = 1;
let builders = [
OpBuilder<(ins "Attribute":$value),
[{
auto* ctx = value.getContext();
$_state.addAttribute("value", value);
$_state.addTypes(TFRAttrType::get(ctx));
}]>
];
let assemblyFormat = [{
$value attr-dict `->` type($out)
}];
}
def TFR_ConstantTensorOp : TFR_Op<"constant_tensor", [Pure]> {
let description = [{
The `constant_tensor` operation converts the operand with non-built-in
tensor type to built-in tensor type or tfr.tensor type. If it is built-in
tensor type, the shape shouldn't be changed during the conversion.
Example:
```mlir
%1 = tfr.constant_tensor(%0) : f32 -> tensor<f32>
%3 = tfr.constant_tensor(%2) : vector<1xf32> -> tensor<1xf32>
```
}];
let arguments = (ins TFR_AllAttrTypes:$arg);
let results = (outs TFR_singleTensorType:$out);
let hasCanonicalizer = 1;
let hasVerifier = 1;
}
def TFR_GetElementOp : TFR_Op<"get_element", [Pure]> {
let description = [{
The `get_element` operation extracts one tfr.tensor element from a
tfr.tensor_list.
Example:
```mlir
%2 = tfr.get_element %1[%0] : (tfr.tensor, index) -> tfr.tensor
```
}];
let arguments = (ins
TFR_TensorListType:$tensor_list,
Index:$index);
let results = (outs TFR_TensorType:$out);
let hasCanonicalizer = 1;
let assemblyFormat = [{
$tensor_list `[` $index `]` attr-dict `:`
`(` type($tensor_list) `,` type($index) `)` `->` type($out)
}];
}
def TFR_BuildListOp : TFR_Op<"build_list", [Pure]> {
let description = [{
The `build_list` operation builds a tensor list from a list of tensors, or
an tfr.attr from a list of scalars.
Example:
```mlir
%3 = tfr.build_list(%2, %1, %0) :
(tfr.tensor, tfr.tensor, tfr.tensor) -> tfr.tensor_list
%3 = tfr.build_list(%2, %1, %0) : (i32, i32, i32) -> tfr.attr
```
}];
let arguments = (ins Variadic<TFR_allowedBuiltListType>:$tensors);
let results = (outs TFR_allowedListResultType:$out);
let hasCanonicalizer = 1;
}
def TFR_GetLengthOp : TFR_Op<"get_length", [Pure]> {
let description = [{
The `get_length` operation returns the number of tensors for a
tfr.tensor_list.
Example:
```mlir
%2 = tfr.get_length(%1) : tfr.tensor -> index
%2 = tfr.get_length %1 -> index
```
}];
let arguments = (ins TFR_TensorListType:$tensor_list);
let results = (outs Index:$out);
let hasCanonicalizer = 1;
let assemblyFormat = [{
$tensor_list attr-dict `->` type($out)
}];
}
//===----------------------------------------------------------------------===//
// Function related classes
//===----------------------------------------------------------------------===//
def TFR_TFRFuncOp : TFR_Op<"func", [HasParent<"ModuleOp">,
FunctionOpInterface,
IsolatedFromAbove, Symbol]> {
let summary = "TFR Function defines a composition of other ops";
let description = [{
Defines a function that can be used to decompose an TF function call to
the invocation of a set of other TF ops.
Syntax:
```
op ::= `tfr.func` visibility? symbol-ref-id `(` argument-list `)` (`->`
function-result-list)? function-attributes? region
```
Example:
```mlir
tfr.func @foo(%arg0: !tfr.tensor, %arg1: !tfr.tensor_list<T>,
%arg2: int {tfr.name="T", tfr.default=1})
attributes {qux: "quux"} {
tfr.return
}
```
Note the arguments are ordered by the following rule:
tfr.tensor > tfr.tensor_list > tfr.attr/i32/...,
and only one trfr.tensor_list argument is allowed.
}];
let arguments = (ins
TypeAttrOf<FunctionType>:$function_type,
StrAttr:$sym_name,
OptionalAttr<DictArrayAttr>:$arg_attrs,
OptionalAttr<DictArrayAttr>:$res_attrs
);
let results = (outs);
// When the regions is empty, the tfr.func is an external function and used
// to model the element type constraints of the tf op. Otherwise, there is one
// region containing the composition.
let regions = (region VariadicRegion<AnyRegion>:$body);
let extraClassDeclaration = [{
LogicalResult verifyType() { return success(); }
mlir::Region *getCallableRegion();
/// Returns the argument types of this function.
ArrayRef<Type> getArgumentTypes() { return getFunctionType().getInputs(); }
/// Returns the result types of this function.
ArrayRef<Type> getResultTypes() { return getFunctionType().getResults(); }
// Get the names of all defined attributes, including both derived and
// non-derived ones.
llvm::StringSet<> getDefinedAttributeNames() {
llvm::StringSet<> all_attrs;
for (auto& attr : (*this)->getAttrs()) {
all_attrs.insert(attr.getName().strref());
}
for (const auto& operand : llvm::enumerate(getArgumentTypes())) {
if (auto attr_name = getArgAttrOfType<StringAttr>(
operand.index(), kAttrArgumentNameAttr)) {
all_attrs.insert(attr_name.getValue());
}
}
return all_attrs;
}
}];
let hasVerifier = 1;
let hasCustomAssemblyFormat = 1;
}
def TFR_TFRReturnOp : TFR_Op<"return", [HasParent<"TFRFuncOp">, Pure,
ReturnLike, Terminator]> {
let description = [{
A terminator operation for regions that appear in the body of `tfr.func`
functions. The operands to the `tfr.return` are the result values returned
by an invocation of the `tfr.func`.
Note that only the tfr.tensor and tfr.tensor_list can be returned.
}];
let arguments = (ins Variadic<TFR_allowedResultType>:$arguments);
let assemblyFormat = "attr-dict ($arguments^ `:` type($arguments))?";
}
//===----------------------------------------------------------------------===//
// Quantization related operations
//===----------------------------------------------------------------------===//
def TFR_TFRQuantActRangeOp : TFR_Op<"quant_act_range", [Pure]> {
let description = [{
The `quant_act_range` returns the a pair of integers to indicate the fixed
range for the fused activation `act` with the quantization defined by the
`scale` and `zero point`. Currently, the allowed activations are
`NONE`, `RELU`, `RELU6` and `RELU_N1_TO_1`.
Example:
```mlir
%3, %4 = tfr.quant_act_range(%2, %1, %0) :
(tfr.attr, float, i64) -> (tfr.tensor, tfr.tensor)
```
}];
let arguments = (ins
TFR_AttrType:$act,
F32:$scale,
I64:$zp);
let results = (outs TFR_TensorType:$min, TFR_TensorType:$max);
let assemblyFormat = [{
`(` $act `,` $scale `,` $zp `)` attr-dict `:` functional-type(operands, results)
}];
}
def TFR_TFRQuantRescaleOp : TFR_Op<"quant_rescale", [Pure]> {
let description = [{
The `quant_rescale` rescales the elements of the integer tensor by the
floating-point rescale factor. This op needs to be legalized to the preferred
operations of the backends.
Example:
```mlir
%3 = tfr.quant_rescale(%2, %1, %0) :
(tfr.tensor, tfr.tensor, i64) -> (tfr.tensor)
```
}];
let arguments = (ins
TFR_TensorType:$input,
TFR_TensorType:$scale,
I64:$zp);
let results = (outs TFR_TensorType:$output);
let assemblyFormat = [{
`(` $input `,` $scale `,` $zp `)` attr-dict `:` functional-type(operands, results)
}];
let hasCanonicalizer = 1;
}
def TFR_TFRQuantRawDataOp : TFR_Op<"quant_raw_data", [
Pure,
SameOperandsAndResultType]> {
let description = [{
The `quant_raw_data` removes the quantization parameter from the intput
tensor(s).
Example:
```mlir
%3 = tfr.quant_raw_data(%0) : (tfr.tensor) -> (tfr.tensor)
```
}];
let arguments = (ins TFR_AllTensorTypes:$input);
let results = (outs TFR_AllTensorTypes:$output);
let assemblyFormat = [{
`(` $input `)` attr-dict `:` functional-type($input, results)
}];
let hasCanonicalizer = 1;
}
def TFR_TFRQuantQParamsOp : TFR_Op<"quant_qparam", [Pure]> {
let description = [{
The `quant_qparam` returns the quantization parameter of the input
tensors.
Example:
```mlir
%3 = tfr.quant_qparam(%0) : (tfr.tensor) -> (float, tfr.tensor)
```
}];
let arguments = (ins TFR_TensorType:$input);
let results = (outs TFR_TensorType:$scale, TFR_TensorType:$zp);
let assemblyFormat = [{
`(` $input `)` attr-dict `:` functional-type($input, results)
}];
let hasCanonicalizer = 1;
}
def TFR_TFRQuantScaleFactorOp : TFR_Op<"quant_scale_factor", [Pure]> {
let description = [{
The `quant_scale_factor` computes the effective scale factor according to the
output scale and input scales.
Example:
```mlir
%3 = tfr.quant_scale_factor(%0) : (f32, tfr.tensor_list) -> (tfr.tensor)
```
}];
let arguments = (ins
F32:$out_scale,
TFR_TensorListType:$in_scales);
let results = (outs TFR_TensorType:$scale_factor);
let assemblyFormat = [{
`(` $out_scale `,` $in_scales `)` attr-dict `:` functional-type(operands, results)
}];
let hasCanonicalizer = 1;
}
#endif // DIALECT_TFR_OPS_
tensorflow/compiler/mlir/tfr/ir/tfr_types.h
+124
View File
@@ -0,0 +1,124 @@
/* 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_TFR_IR_TFR_TYPES_H_
#define TENSORFLOW_COMPILER_MLIR_TFR_IR_TFR_TYPES_H_
#include <memory>
#include <string>
#include "mlir/IR/Attributes.h" // from @llvm-project
#include "mlir/IR/BuiltinTypes.h" // from @llvm-project
#include "mlir/IR/Diagnostics.h" // from @llvm-project
#include "mlir/IR/Location.h" // from @llvm-project
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "mlir/IR/Operation.h" // from @llvm-project
#include "mlir/IR/TypeSupport.h" // from @llvm-project
#include "mlir/IR/Types.h" // from @llvm-project
namespace mlir {
namespace TFR {
class TFRType : public Type {
public:
using Type::Type;
static bool classof(Type type);
};
namespace detail {
struct TFRTypeStorage final
: public TypeStorage,
public llvm::TrailingObjects<TFRTypeStorage, StringAttr> {
using KeyTy = ArrayRef<StringAttr>;
explicit TFRTypeStorage(unsigned num_attrs) : num_attrs(num_attrs) {}
static TFRTypeStorage* construct(TypeStorageAllocator& allocator, KeyTy key) {
// Allocate a new storage instance.
auto byteSize = TFRTypeStorage::totalSizeToAlloc<StringAttr>(key.size());
auto rawMem = allocator.allocate(byteSize, alignof(TFRTypeStorage));
auto result = ::new (rawMem) TFRTypeStorage(key.size());
// Copy in the string attributes into the trailing storage.
std::uninitialized_copy(key.begin(), key.end(),
result->getTrailingObjects());
return result;
}
bool operator==(const KeyTy& attrs) const { return attrs == GetAttrs(); }
KeyTy GetAttrs() const { return {getTrailingObjects(), num_attrs}; }
unsigned num_attrs;
};
template <typename Derived>
class TFRTypeImpl : public Type::TypeBase<Derived, TFRType, TFRTypeStorage> {
public:
using Base = Type::TypeBase<Derived, TFRType, TFRTypeStorage>;
using TFRBase = TFRTypeImpl<Derived>;
using Base::Base;
static Derived get(ArrayRef<StringAttr> attrs, MLIRContext* context) {
return Base::get(context, attrs);
}
static Derived getChecked(ArrayRef<StringAttr> attrs, Location loc) {
return Base::getChecked(loc, loc.getContext(), attrs);
}
static Derived getChecked(function_ref<InFlightDiagnostic()> emitError,
MLIRContext* context, ArrayRef<StringAttr> attrs) {
return Base::getChecked(emitError, context, attrs);
}
static Derived get(MLIRContext* context) { return get({}, context); }
// TODO(fengliuai): fix the implementation
static LogicalResult verify(function_ref<InFlightDiagnostic()> emitError,
ArrayRef<StringAttr> attrs) {
return success();
}
ArrayRef<StringAttr> getAttrKeys() { return Base::getImpl()->GetAttrs(); }
};
} // namespace detail
class TFRTensorType : public detail::TFRTypeImpl<TFRTensorType> {
public:
using TFRBase::TFRBase;
static constexpr StringLiteral name = "tfr.tensor";
static std::string getTypeName() { return "TFRTensorType"; }
};
class TFRTensorListType : public detail::TFRTypeImpl<TFRTensorListType> {
public:
using TFRBase::TFRBase;
static constexpr StringLiteral name = "tfr.tensor_list";
static std::string getTypeName() { return "TFRTensorListType"; }
};
class TFRAttrType : public Type::TypeBase<TFRAttrType, TFRType, TypeStorage> {
public:
using Base::Base;
static constexpr StringLiteral name = "tfr.attr";
static std::string getTypeName() { return "TFRAttrType"; }
};
} // namespace TFR
} // namespace mlir
#endif // TENSORFLOW_COMPILER_MLIR_TFR_IR_TFR_TYPES_H_
tensorflow/compiler/mlir/tfr/passes/canonicalize.cc
+177
View File
@@ -0,0 +1,177 @@
/* 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 <cstdint>
#include "llvm/Support/raw_ostream.h"
#include "mlir/Conversion/SCFToControlFlow/SCFToControlFlow.h" // from @llvm-project
#include "mlir/Dialect/Affine/LoopUtils.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/SCF/IR/SCF.h" // from @llvm-project
#include "mlir/IR/Attributes.h" // from @llvm-project
#include "mlir/IR/IRMapping.h" // from @llvm-project
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "mlir/IR/Matchers.h" // from @llvm-project
#include "mlir/IR/PatternMatch.h" // from @llvm-project
#include "mlir/IR/Region.h" // from @llvm-project
#include "mlir/Support/LLVM.h" // from @llvm-project
#include "mlir/Support/LogicalResult.h" // from @llvm-project
#include "mlir/Transforms/Inliner.h" // from @llvm-project
#include "mlir/Transforms/InliningUtils.h" // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/compiler/mlir/tfr/ir/tfr_ops.h"
#include "tensorflow/compiler/mlir/tfr/passes/passes.h"
//===----------------------------------------------------------------------===//
// Canonicalization patterns for the scf.for and scf.if ops. They are used to
// optimize the control flow in the tfr function. Technically, both patterns
// should be upstreamed to be part of the op definition.
// TODO(fengliuai): sync with the llvm upstream for both patterns.
//
namespace mlir {
namespace TFR {
namespace {
class UnrollSCFForOp : public OpRewritePattern<scf::ForOp> {
using OpRewritePattern<scf::ForOp>::OpRewritePattern;
public:
LogicalResult matchAndRewrite(scf::ForOp for_op,
PatternRewriter &rewriter) const override {
Location loc = for_op.getLoc();
APInt lower_bound, upper_bound, step;
if (!matchPattern(for_op.getLowerBound(), m_ConstantInt(&lower_bound)) ||
!matchPattern(for_op.getUpperBound(), m_ConstantInt(&upper_bound)) ||
!matchPattern(for_op.getStep(), m_ConstantInt(&step))) {
return failure();
}
uint64_t trip_count = (upper_bound - lower_bound).sdiv(step).getZExtValue();
if (trip_count <= 0) return failure();
// TODO(fengliuai): use loopUnrollByFactor once the iter_arg is supported
Block *single_block = for_op.getBody();
IRMapping mapping;
Value iv = for_op.getInductionVar();
for (auto iter_op :
llvm::zip(for_op.getRegionIterArgs(), for_op.getInitArgs())) {
mapping.map(std::get<0>(iter_op), std::get<1>(iter_op));
}
mapping.map(iv, for_op.getLowerBound());
for (auto i = 0; i < trip_count; ++i) {
if (!iv.use_empty()) {
// iv' = iv + step * i;
Value iter = arith::ConstantIndexOp::create(rewriter, loc, i);
Value step_cst =
arith::ConstantIndexOp::create(rewriter, loc, step.getSExtValue());
Value stride = arith::MulIOp::create(rewriter, loc, step_cst, iter);
Value iv_unroll =
arith::AddIOp::create(rewriter, loc, mapping.lookup(iv), stride);
mapping.map(iv, iv_unroll);
}
Operation *terminator_op;
for (auto it = single_block->begin(); it != single_block->end(); ++it) {
terminator_op = rewriter.clone(*it, mapping);
}
// Map the block arguments to the yield results.
for (auto iter_op : llvm::zip(for_op.getRegionIterArgs(),
terminator_op->getOperands())) {
mapping.map(std::get<0>(iter_op), std::get<1>(iter_op));
}
rewriter.eraseOp(terminator_op);
}
SmallVector<Value, 4> returned;
for (Value arg : for_op.getRegionIterArgs()) {
returned.push_back(mapping.lookup(arg));
}
rewriter.replaceOp(for_op, returned);
return success();
}
};
// TODO(fengliuai): up stream this pattern.
class SimplifySCFIfOp : public OpRewritePattern<scf::IfOp> {
using OpRewritePattern<scf::IfOp>::OpRewritePattern;
public:
LogicalResult matchAndRewrite(scf::IfOp if_op,
PatternRewriter &rewriter) const override {
// Then branch
if (matchPattern(if_op.getCondition(), m_NonZero())) {
return InlineRegion(if_op.getLoc(), rewriter, if_op,
&if_op.getThenRegion());
}
// Else branch
if (matchPattern(if_op.getCondition(), m_Zero())) {
if (if_op.getElseRegion().empty()) {
// Remove the op
rewriter.eraseOp(if_op);
return success();
} else {
return InlineRegion(if_op.getLoc(), rewriter, if_op,
&if_op.getElseRegion());
}
}
// Not a constant condition
return failure();
}
private:
LogicalResult InlineRegion(Location loc, PatternRewriter &rewriter,
Operation *inline_point, Region *region) const;
};
LogicalResult SimplifySCFIfOp::InlineRegion(Location loc,
PatternRewriter &rewriter,
Operation *inline_point,
Region *region) const {
InlinerInterface interface(loc.getContext());
InlinerConfig config;
if (failed(inlineRegion(interface, config.getCloneCallback(), region,
inline_point, {}, inline_point->getResults(), loc,
/*shouldCloneInlinedRegion=*/true))) {
return failure();
}
// If the inlining was successful then erase the scf.if op.
rewriter.eraseOp(inline_point);
return success();
}
} // namespace
void populateCanonicalizationPatterns(func::FuncOp func,
RewritePatternSet &patterns) {
MLIRContext *context = func.getContext();
mlir::Dialect *tf = context->getLoadedDialect<mlir::TF::TensorFlowDialect>();
// Load all official canonicalization patterns. Here we skip the
// canonicalization of the ops in the tf dialect, because they couldn't
// propagate the attributes correctly. These optimization will be played by
// bridge.
func->walk([&](Operation *op) {
if (op->getDialect() != tf) {
op->getRegisteredInfo()->getCanonicalizationPatterns(patterns, context);
}
});
patterns.add<UnrollSCFForOp, SimplifySCFIfOp>(context);
}
} // namespace TFR
} // namespace mlir
tensorflow/compiler/mlir/tfr/passes/decompose.cc
+370
View File
@@ -0,0 +1,370 @@
/* 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 <algorithm>
#include <cmath>
#include <cstdint>
#include <iterator>
#include <limits>
#include <memory>
#include <optional>
#include <string>
#include <utility>
#include "llvm/ADT/ArrayRef.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/ADT/StringRef.h"
#include "llvm/Support/Casting.h"
#include "llvm/Support/CommandLine.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/Dialect/SCF/IR/SCF.h" // from @llvm-project
#include "mlir/IR/Attributes.h" // from @llvm-project
#include "mlir/IR/Builders.h" // from @llvm-project
#include "mlir/IR/BuiltinAttributes.h" // from @llvm-project
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/IR/BuiltinTypes.h" // from @llvm-project
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "mlir/IR/SymbolTable.h" // from @llvm-project
#include "mlir/IR/Value.h" // from @llvm-project
#include "mlir/IR/Visitors.h" // from @llvm-project
#include "mlir/Pass/Pass.h" // from @llvm-project
#include "mlir/Support/LLVM.h" // from @llvm-project
#include "mlir/Support/LogicalResult.h" // from @llvm-project
#include "mlir/Transforms/DialectConversion.h" // from @llvm-project
#include "mlir/Transforms/GreedyPatternRewriteDriver.h" // from @llvm-project
#include "mlir/Transforms/Inliner.h" // from @llvm-project
#include "mlir/Transforms/InliningUtils.h" // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/compiler/mlir/tfr/ir/tfr_ops.h"
#include "tensorflow/compiler/mlir/tfr/ir/tfr_types.h"
#include "tensorflow/compiler/mlir/tfr/passes/passes.h"
#include "tensorflow/compiler/mlir/tfr/utils/utils.h"
#include "tensorflow/core/lib/monitoring/counter.h"
namespace tensorflow {
namespace {
auto* tf_core_op_expansion_op_counter =
monitoring::Counter<1>::New("/tensorflow/core/op_expansion/op_counter",
"The number of composite op expanded.", "name");
}
void IncreaseOpExpansionExecuteCounterByOne(const std::string& op_name) {
tf_core_op_expansion_op_counter->GetCell(op_name)->IncrementBy(1);
}
} // namespace tensorflow
//===----------------------------------------------------------------------===//
// The pass to decompose unregistered TF ops with the TFR compose function.
//
namespace mlir {
namespace TFR {
namespace {
// Quantize the float value based on given scale and zero point attributes.
IntegerAttr Quantize(float value, Attribute scale_attr, Attribute zp_attr,
OpBuilder builder) {
double scale = mlir::cast<FloatAttr>(scale_attr).getValueAsDouble();
int64_t zp = mlir::cast<IntegerAttr>(zp_attr).getInt();
int quantized = static_cast<int>(std::round(value / scale) + zp);
quantized =
std::min(quantized, static_cast<int>(std::numeric_limits<int8_t>::max()));
quantized =
std::max(quantized, static_cast<int>(std::numeric_limits<int8_t>::min()));
return builder.getI32IntegerAttr(quantized);
}
// Decompose the TF ops with the registered composition library.
class DecomposeTFOpsPass
: public PassWrapper<DecomposeTFOpsPass, OperationPass<func::FuncOp>> {
public:
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(DecomposeTFOpsPass)
explicit DecomposeTFOpsPass(std::optional<ModuleOp> external_tfr_module)
: external_tfr_module_(external_tfr_module) {}
StringRef getArgument() const final { return "tfr-decompose"; }
StringRef getDescription() const final {
return "Decompose TF ops with the registered composition library.";
}
void runOnOperation() override;
private:
// Apply canonicalization, mainly constant folding, on the function.
void ApplyCanonicalization();
// Rewrite unregistered TF ops to TFR func call ops. Return failure if all the
// ops are registered or the compose function doesn't exist.
LogicalResult RewriteUnregisteredTFOps();
// Inline the TFR func call ops.
LogicalResult InlineTFRFuncCalls();
// Optional external symbol table to look up the TFR function.
std::optional<ModuleOp> external_tfr_module_;
};
#include "tensorflow/compiler/mlir/tfr/passes/generated_decompose.inc"
void DecomposeTFOpsPass::ApplyCanonicalization() {
func::FuncOp func = getOperation();
RewritePatternSet patterns(&getContext());
populateWithGenerated(patterns);
populateCanonicalizationPatterns(func, patterns);
(void)applyPatternsGreedily(func, std::move(patterns));
}
LogicalResult DecomposeTFOpsPass::RewriteUnregisteredTFOps() {
func::FuncOp func = getOperation();
SymbolTable table(external_tfr_module_.has_value()
? *external_tfr_module_
: func->getParentOfType<ModuleOp>());
OpBuilder builder(func);
bool changed = false;
func.walk([&table, &builder, &changed](Operation* op) {
// Only the un-registered ops requires decomposition. The remaining ones
// either will be constant folded or lowered by the rules defined in the
// bridge.
if (op->isRegistered()) {
return WalkResult::advance();
}
// Find out the compose function
auto compose_func_name = GetComposeFuncName(op->getName().getStringRef());
auto compose_func = table.lookup<TFRFuncOp>(compose_func_name);
if (!compose_func || compose_func.isExternal()) {
// There are no decomposition methods defined for this op, skip.
return WalkResult::advance();
}
// Make sure all the attributes are valid. An attribute is valid when it is
// in the signature or it is allowed explicitly.
auto compose_func_signature =
table.lookup<TFRFuncOp>(compose_func_name + "_");
if (!compose_func_signature) compose_func_signature = compose_func;
auto defined_attrs = compose_func_signature.getDefinedAttributeNames();
if (failed(ValidateAttrs(op, defined_attrs))) {
return WalkResult::interrupt();
}
tensorflow::IncreaseOpExpansionExecuteCounterByOne(
op->getName().getStringRef().str());
auto compose_func_type = compose_func.getFunctionType();
builder.setInsertionPoint(op);
TFRTensorType unconstrainted_tensor_type = builder.getType<TFRTensorType>();
// Create the new operands. This is mapping the operands from the target
// TF ops to the TFR function arguments. If the TFR function argument is
// a tensor_list, a "tfr.build_list" op is used to concat the available
// TF op operands. If the TFR function argument isn't a tensor/tensor_list,
// a constant is created by using the attribute stored in the TF op or the
// default value in the argument attribute.
llvm::SmallVector<Value, 4> new_operands;
for (auto arg : llvm::enumerate(compose_func_type.getInputs())) {
if (auto tensor_type = mlir::dyn_cast<TFRTensorType>(arg.value())) {
auto casted = CastOp::create(builder, op->getLoc(), tensor_type,
op->getOperand(arg.index()));
new_operands.push_back(casted);
} else if (auto list_type =
mlir::dyn_cast<TFRTensorListType>(arg.value())) {
llvm::SmallVector<Value, 4> variadic_operands;
for (int i = arg.index(); i < op->getNumOperands(); i++) {
auto casted =
CastOp::create(builder, op->getLoc(), unconstrainted_tensor_type,
op->getOperand(i));
variadic_operands.push_back(casted);
}
auto build_list_op = BuildListOp::create(builder, op->getLoc(),
list_type, variadic_operands);
new_operands.push_back(build_list_op.getOut());
} else {
auto attr_name = compose_func.getArgAttrOfType<StringAttr>(
arg.index(), kAttrArgumentNameAttr);
auto attribute = op->getAttr(attr_name.getValue());
if (!attribute) {
attribute =
compose_func.getArgAttr(arg.index(), kAttrArgumentDefaultAttr);
}
if (!attribute && attr_name.getValue() == "out_type") {
auto type = op->getResult(0).getType();
if (mlir::isa<TensorType>(type)) {
type = mlir::cast<TensorType>(type).getElementType();
}
attribute = TypeAttr::get(type);
}
Value attr_cst;
// Wrap these special attributes as a special TFR constant, so the SSA
// value has a valid type to be used as TFR function argument. These
// attributes are not expected to be manipulated by the lowering passes.
if (mlir::isa<TypeAttr>(attribute) || mlir::isa<ArrayAttr>(attribute) ||
mlir::isa<StringAttr>(attribute) ||
mlir::isa<FlatSymbolRefAttr>(attribute)) {
TFRAttrType output_type = TFRAttrType::get(builder.getContext());
attr_cst =
ConstOp::create(builder, op->getLoc(), output_type, attribute);
} else {
attr_cst = mlir::arith::ConstantOp::create(
builder, op->getLoc(), cast<TypedAttr>(attribute));
}
new_operands.push_back(attr_cst);
}
}
// Create the TFR call op
auto new_op = CallOp::create(
builder, op->getLoc(), compose_func_type.getResults(),
SymbolRefAttr::get(builder.getContext(), compose_func.getName()),
new_operands, /*args_attrs=*/nullptr, /*res_attrs=*/nullptr);
// Replace the use of the old op. This is mapping the results from the
// target TF ops to the TFR function returns. If the TFR function return is
// a tensor_list, "tfr.get_element" op is used to extract the required TF
// op result.
llvm::SmallVector<Value, 4> new_results;
for (auto res : llvm::enumerate(compose_func_type.getResults())) {
if (mlir::dyn_cast<TFRTensorType>(res.value())) {
new_results.push_back(new_op.getResult(res.index()));
} else if (auto list_type =
mlir::dyn_cast<TFRTensorListType>(res.value())) {
for (int i = res.index(), j = 0; i < op->getNumResults(); i++, j++) {
auto index = mlir::arith::ConstantOp::create(builder, op->getLoc(),
builder.getIndexAttr(j));
auto element_op = GetElementOp::create(
builder, op->getLoc(), unconstrainted_tensor_type,
new_op.getResult(res.index()), index.getResult());
new_results.push_back(element_op.getOut());
}
}
}
for (auto res : llvm::zip(op->getResults(), new_results)) {
auto casted = CastOp::create(
builder, op->getLoc(), std::get<0>(res).getType(), std::get<1>(res));
std::get<0>(res).replaceAllUsesWith(casted.getOut());
}
// Copy all the unregisted attributes to the new op.
if (failed(CopyAllowedUnregisteredAttrs(op, new_op, defined_attrs))) {
return WalkResult::interrupt();
}
op->erase();
changed |= true;
return WalkResult::advance();
});
// If `changed` is false, it is considered as a failure, so the recursive
// rewrite will stop.
return success(changed);
}
LogicalResult DecomposeTFOpsPass::InlineTFRFuncCalls() {
// The Inliner will automatically use the registered dialect inliner.
InlinerInterface inliner(&getContext());
InlinerConfig config;
func::FuncOp func = getOperation();
SymbolTable table(external_tfr_module_.has_value()
? *external_tfr_module_
: func->getParentOfType<ModuleOp>());
// The inliner only inlines the TFR call op.
bool changed = false;
auto walk_result = func.walk([&](CallOp call_op) {
auto callee = table.lookup<TFRFuncOp>(call_op.getCallee());
if (!callee || callee.isExternal()) return WalkResult::advance();
// Record the boundary of the inlined operations. The inlined operation will
// be inserted between these two operations.
Operation* inlined_point = call_op.getOperation();
Operation* after_inlined_point =
&*std::next(Block::iterator(call_op.getOperation()));
// Use the inliner to replace all the uses of the call_op by its
// composition.
if (failed(inlineCall(inliner, config.getCloneCallback(),
cast<CallOpInterface>(call_op.getOperation()),
cast<CallableOpInterface>(callee.getOperation()),
callee.getCallableRegion(),
/**shouldCloneInLinedRegion=*/true))) {
// This failure is usually because the decompose function is not defined.
// This call will be raised to TF ops.
return WalkResult::interrupt();
}
// Propagate all the attributes to the inlined operations, which are defined
// by the two boundary operations.
PropagateAttrsToOperations(call_op, Block::iterator(inlined_point),
Block::iterator(after_inlined_point));
// Remove the call_op to finish the op expansion.
call_op.erase();
changed |= true;
return WalkResult::advance();
});
if (walk_result.wasInterrupted()) {
signalPassFailure();
return failure();
}
// If `changed` is false, it is considered as a failure, so the recursive
// rewrite will stop.
return success(changed);
}
void DecomposeTFOpsPass::runOnOperation() {
// Set a maximum iteration threshold in case there are infinite loops in the
// call stack.
int max_iterators = 10;
do {
// canonicalization
ApplyCanonicalization();
// rewrite unregistered tf ops. Failed either because no ops can be
// decomposed or the compose function isn't defined.
auto rewrite_status = RewriteUnregisteredTFOps();
// inline the tfr call op until there are no tfr.call op can be inlined.
auto inline_status = InlineTFRFuncCalls();
if (failed(rewrite_status) && failed(inline_status)) {
break;
}
} while (max_iterators-- >= 0);
}
} // namespace
// Creates an instance of the pass to decompose the TF ops.
std::unique_ptr<OperationPass<func::FuncOp>> CreateDecomposeTFOpsPass(
std::optional<ModuleOp> tfr_module) {
return std::make_unique<DecomposeTFOpsPass>(tfr_module);
}
static PassRegistration<DecomposeTFOpsPass> pass([] {
return CreateDecomposeTFOpsPass();
});
} // namespace TFR
} // namespace mlir
tensorflow/compiler/mlir/tfr/passes/decompose_patterns.td
+59
View File
@@ -0,0 +1,59 @@
/* 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 "mlir/IR/OpBase.td"
include "mlir/IR/PatternBase.td"
include "mlir/Dialect/Arith/IR/ArithOps.td"
include "mlir/Dialect/Func/IR/FuncOps.td"
include "tensorflow/compiler/mlir/tfr/ir/tfr_ops.td"
class Quantize<string value> : NativeCodeCall<"TFR::Quantize(" # value # ", $0, $1, $_builder)">;
class HasStringAttr<string value> : AttrConstraint<
CPred<"llvm::cast<StringAttr>($_self).getValue() == \"" # value # "\"">>;
def QuantActRangeNonePattern :
Pattern<
(TFR_TFRQuantActRangeOp
(TFR_ConstOp HasStringAttr<"NONE">:$act), $scale, $zp),
[(TFR_ConstantTensorOp (Arith_ConstantOp ConstantAttr<I32Attr, "-128">)),
(TFR_ConstantTensorOp (Arith_ConstantOp ConstantAttr<I32Attr, "127">))]>;
def QuantActRangeReluPattern :
Pattern<
(TFR_TFRQuantActRangeOp
(TFR_ConstOp HasStringAttr<"RELU">:$act),
(ConstantLikeMatcher F32Attr:$scale),
(ConstantLikeMatcher I64Attr:$zp)),
[(TFR_ConstantTensorOp (Arith_ConstantOp (Quantize<"0.0f"> $scale, $zp))),
(TFR_ConstantTensorOp (Arith_ConstantOp ConstantAttr<I32Attr, "127">))]>;
def QuantActRangeRelu6Pattern :
Pattern<
(TFR_TFRQuantActRangeOp
(TFR_ConstOp HasStringAttr<"RELU6">:$act),
(ConstantLikeMatcher F32Attr:$scale),
(ConstantLikeMatcher I64Attr:$zp)),
[(TFR_ConstantTensorOp (Arith_ConstantOp (Quantize<"0.0f"> $scale, $zp))),
(TFR_ConstantTensorOp (Arith_ConstantOp (Quantize<"6.0f"> $scale, $zp)))]>;
def QuantActRangeReluN1To1Pattern :
Pattern<
(TFR_TFRQuantActRangeOp
(TFR_ConstOp HasStringAttr<"RELU_N1_TO_1">:$act),
(ConstantLikeMatcher F32Attr:$scale),
(ConstantLikeMatcher I64Attr:$zp)),
[(TFR_ConstantTensorOp (Arith_ConstantOp (Quantize<"-1.0f"> $scale, $zp))),
(TFR_ConstantTensorOp (Arith_ConstantOp (Quantize<"1.0f"> $scale, $zp)))]>;
tensorflow/compiler/mlir/tfr/passes/passes.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_TFR_PASSES_PASSES_H_
#define TENSORFLOW_COMPILER_MLIR_TFR_PASSES_PASSES_H_
#include <memory>
#include <optional>
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/Pass/Pass.h" // from @llvm-project
#include "mlir/Support/LogicalResult.h" // from @llvm-project
namespace mlir {
namespace TFR {
// Scans the func op and adds all the canonicalization patterns of the ops
// except the tf ops, inside the function.
void populateCanonicalizationPatterns(func::FuncOp func,
RewritePatternSet &patterns);
// Decompose ops.
std::unique_ptr<OperationPass<func::FuncOp>> CreateDecomposeTFOpsPass(
std::optional<ModuleOp> tfr_module = std::nullopt);
// Rewrites quantized operands and results with their storage types.
// This pass should be run at module level after decomposition, if there are
// quantized operands or results.
std::unique_ptr<OperationPass<ModuleOp>> CreateRewriteQuantizedIOPass();
// Raise to TF ops.
std::unique_ptr<OperationPass<func::FuncOp>> CreateRaiseToTFOpsPass(
std::optional<ModuleOp> tfr_module = std::nullopt,
bool materialize_derived_attrs = false);
} // namespace TFR
} // namespace mlir
#endif // TENSORFLOW_COMPILER_MLIR_TFR_PASSES_PASSES_H_
tensorflow/compiler/mlir/tfr/passes/raise_to_tf.cc
+522
View File
@@ -0,0 +1,522 @@
/* 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 <memory>
#include <optional>
#include <string>
#include <utility>
#include "llvm/ADT/ArrayRef.h"
#include "llvm/ADT/DenseSet.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/ADT/StringMap.h"
#include "llvm/ADT/StringRef.h"
#include "llvm/Support/Casting.h"
#include "llvm/Support/CommandLine.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/Dialect/SCF/IR/SCF.h" // from @llvm-project
#include "mlir/IR/Attributes.h" // from @llvm-project
#include "mlir/IR/Builders.h" // from @llvm-project
#include "mlir/IR/BuiltinAttributes.h" // from @llvm-project
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/IR/BuiltinTypes.h" // from @llvm-project
#include "mlir/IR/Diagnostics.h" // from @llvm-project
#include "mlir/IR/Location.h" // from @llvm-project
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "mlir/IR/Matchers.h" // from @llvm-project
#include "mlir/IR/OperationSupport.h" // from @llvm-project
#include "mlir/IR/SymbolTable.h" // from @llvm-project
#include "mlir/IR/Types.h" // from @llvm-project
#include "mlir/IR/Value.h" // from @llvm-project
#include "mlir/IR/Visitors.h" // from @llvm-project
#include "mlir/Pass/Pass.h" // from @llvm-project
#include "mlir/Support/LLVM.h" // from @llvm-project
#include "mlir/Support/LogicalResult.h" // from @llvm-project
#include "mlir/Transforms/DialectConversion.h" // from @llvm-project
#include "mlir/Transforms/GreedyPatternRewriteDriver.h" // from @llvm-project
#include "mlir/Transforms/InliningUtils.h" // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/compiler/mlir/tfr/ir/tfr_ops.h"
#include "tensorflow/compiler/mlir/tfr/ir/tfr_types.h"
#include "tensorflow/compiler/mlir/tfr/passes/passes.h"
#include "tensorflow/compiler/mlir/tfr/utils/utils.h"
//===----------------------------------------------------------------------===//
// The pass to rewrite the TFR function call ops by TF ops. The callee of the
// TFR function call defines the signatures of the TF ops.
//
namespace mlir {
namespace TFR {
namespace {
// This pattern is to rewrite the "tfr.call" op and the "tfr.cast" ops on the
// operands by a TF op with "tfr.cast" ops on the results. The result type of
// the new TF op is an unranked tensor with element type derived.
class RewriteTFRCallOp : public OpRewritePattern<CallOp> {
using OpRewritePattern<CallOp>::OpRewritePattern;
public:
explicit RewriteTFRCallOp(MLIRContext* context, const SymbolTable& table,
bool materialize_derived_attrs)
: OpRewritePattern<CallOp>(context),
symbol_table_(table),
materialize_derived_attrs_(materialize_derived_attrs) {}
LogicalResult matchAndRewrite(CallOp call_op,
PatternRewriter& rewriter) const override;
private:
// Derives the attribute values for the attributes attached to the
// `input_tfr_type`. These attributes are only for the element type of the
// inputs, and these type information has been collected in the `input_types`.
// The result is stored in `derived_attrs` as the named attributes. Returns
// failure if the attributes stored in the `input_tfr_type` violates the
// assumptions.
LogicalResult AddDerivedAttrs(
PatternRewriter& rewriter, Type input_tfr_type,
ArrayRef<Attribute> input_types,
llvm::StringMap<Attribute>* derived_attrs) const;
// Collects the operands and attributes for the TF op. At the same time, it
// collects all the derived attribute values to derive the output types of the
// TF op.
LogicalResult CollectInputsAndAttributes(
PatternRewriter& rewriter, TFRFuncOp signature, CallOp call_op,
SmallVectorImpl<Value>* inputs, NamedAttrList* arg_attrs,
llvm::StringMap<Attribute>* derived_attrs) const;
// Uses the collected attribute values to derive all the output types.
LogicalResult DeriveOutputTypes(Location loc, FunctionType signature,
const llvm::StringMap<Attribute>& attrs,
SmallVectorImpl<Type>* output_types) const;
// Creates the TF op and also the necessary tfr.cast ops to replace the
// original TFR call op.
LogicalResult CreateAndReplaceOp(
PatternRewriter& rewriter, CallOp call_op,
const SmallVectorImpl<Type>& output_types,
const SmallVectorImpl<Value>& inputs, const NamedAttrList& attr_list,
const llvm::StringMap<Attribute>& derived_attrs) const;
// Converts the attribute to the specific type.
Attribute ProcessAttributeValue(Attribute attr, StringAttr attr_type) const;
Type GetFixedElementType(StringRef element_type, Builder& builder) const {
if (element_type == "i32_") return builder.getI32Type();
if (element_type == "i64_") return builder.getI64Type();
if (element_type == "f32_") return builder.getF32Type();
if (element_type == "i1_") return builder.getI1Type();
return {};
}
// Adds a tf.Cast op if the tfr.tensor attribute indicated a fixed element
// type.
// TODO(fengliuai): This method is required when the operand types are not set
// by the frontend correctly.
Value CastToNonDerivedType(PatternRewriter& rewriter, Location loc,
CastOp cast_op, Type input_tfr_type) const {
auto tensor_type = mlir::dyn_cast<TFRTensorType>(input_tfr_type);
if (!tensor_type) return cast_op.getArg();
auto attr_names = tensor_type.getAttrKeys();
if (attr_names.empty() || attr_names.size() > 1) return cast_op.getArg();
StringRef tfr_type_attr = attr_names[0].getValue();
if (!fixed_elt_type_attrs_.contains(tfr_type_attr)) return cast_op.getArg();
Type result_elt_type = GetFixedElementType(tfr_type_attr, rewriter);
if (!result_elt_type) {
return cast_op.getArg();
}
Type original_input_type =
mlir::cast<TypeAttr>(cast_op.getInputElementType()).getValue();
if (result_elt_type != original_input_type) {
UnrankedTensorType result_type = UnrankedTensorType::get(result_elt_type);
return TF::CastOp::create(rewriter, loc, result_type, cast_op.getArg());
}
return cast_op.getArg();
}
// For variadic operands, we have to enforce them to use the same types.
// TODO(fengliuai): This method is required when the operand types are not set
// by the frontend correctly.
void CastValuesToSameType(PatternRewriter& rewriter, Location loc,
const llvm::SmallVectorImpl<Attribute>& input_types,
llvm::SmallVectorImpl<Value>& input_values) const {
if (input_types.size() <= 1) return;
Type target_input_type = mlir::cast<TypeAttr>(input_types[0]).getValue();
auto result_type = UnrankedTensorType::get(target_input_type);
for (auto i = 1; i < input_types.size(); ++i) {
Type current_input_type = mlir::cast<TypeAttr>(input_types[i]).getValue();
if (current_input_type != target_input_type) {
input_values[i] =
TF::CastOp::create(rewriter, loc, result_type, input_values[i]);
}
}
}
const SymbolTable& symbol_table_;
const bool materialize_derived_attrs_;
const llvm::SmallDenseSet<StringRef, 4> fixed_elt_type_attrs_{"i32_", "i64_",
"f32_", "i1_"};
};
LogicalResult RewriteTFRCallOp::AddDerivedAttrs(
PatternRewriter& rewriter, Type input_tfr_type,
ArrayRef<Attribute> input_types,
llvm::StringMap<Attribute>* derived_attrs) const {
// If there is an attribute associated to the input in the signature, we
// store it as an derived attribute.
if (auto tensor_type = mlir::dyn_cast<TFRTensorType>(input_tfr_type)) {
auto attr_names = tensor_type.getAttrKeys();
if (attr_names.empty()) return success();
if (attr_names.size() == 1) {
derived_attrs->insert({attr_names[0].getValue(), input_types[0]});
return success();
}
}
// If there is an attribute associated to the input in the signature,
// we store it as an derived attribute.
if (auto list_type = mlir::dyn_cast<TFRTensorListType>(input_tfr_type)) {
auto attr_names = list_type.getAttrKeys();
if (attr_names.empty()) return success();
// N*T case
if (attr_names.size() == 2) {
derived_attrs->insert({attr_names[0].getValue(),
rewriter.getI32IntegerAttr(input_types.size())});
// Note that this uses the first element of the list to infer the T value.
// A tf.Cast is required to cast the other inputs to the same type.
derived_attrs->insert({attr_names[1].getValue(), input_types[0]});
return success();
}
// list(dtype) case
if (attr_names.size() == 1) {
derived_attrs->insert(
{attr_names[0].getValue(), rewriter.getArrayAttr(input_types)});
return success();
}
}
return failure();
}
LogicalResult RewriteTFRCallOp::CollectInputsAndAttributes(
PatternRewriter& rewriter, TFRFuncOp signature, CallOp call_op,
SmallVectorImpl<Value>* inputs, NamedAttrList* arg_attrs,
llvm::StringMap<Attribute>* derived_attrs) const {
for (const auto& operand :
llvm::enumerate(signature.getFunctionType().getInputs())) {
// If the index is larger than the operand number of the call_op, the
// default value of the operand needs to be used.
if (operand.index() >= call_op.getNumOperands()) {
auto attr_name = signature.getArgAttrOfType<StringAttr>(
operand.index(), kAttrArgumentNameAttr);
auto attr_value =
signature.getArgAttr(operand.index(), kAttrArgumentDefaultAttr);
arg_attrs->push_back(
rewriter.getNamedAttr(attr_name.getValue(), attr_value));
continue;
}
// The index is valid for the call_op.
Value input = call_op.getOperand(operand.index());
Operation* input_op = input.getDefiningOp();
auto input_tfr_type =
signature.getFunctionType().getInputs()[operand.index()];
// There are three cases for the preceding input_op:
// 1. The preceding op can be a tfr.cast op, which will be fused to the
// current op, so the result op has input with tensor type.
if (auto cast_op = dyn_cast_or_null<CastOp>(input_op)) {
Value input_to_cast = CastToNonDerivedType(rewriter, call_op.getLoc(),
cast_op, input_tfr_type);
inputs->push_back(input_to_cast);
if (failed(AddDerivedAttrs(rewriter, input_tfr_type,
{cast_op.getInputElementType()},
derived_attrs))) {
return failure();
}
continue;
}
// 2. The preceding op is a tfr.build_list op, which collects multiple
// values with tensor types via the tfr.cast ops. These ops will be fused
// to the current op as well, so all the tfr.cast op inputs will be inputs
// to the result op.
if (auto list_op = dyn_cast_or_null<BuildListOp>(input_op)) {
// Find out all the inputs to the build list op
// TODO(fengliuai): make build_list op only take tensor argument
llvm::SmallVector<Attribute, 4> list_input_types;
llvm::SmallVector<Value, 4> list_inputs;
for (auto list_input : list_op.getOperands()) {
auto cast_op = dyn_cast_or_null<CastOp>(list_input.getDefiningOp());
if (!cast_op) return failure();
list_inputs.push_back(cast_op.getArg());
list_input_types.push_back(cast_op.getInputElementType());
}
CastValuesToSameType(rewriter, call_op.getLoc(), list_input_types,
list_inputs);
inputs->append(list_inputs.begin(), list_inputs.end());
if (failed(AddDerivedAttrs(rewriter, input_tfr_type, list_input_types,
derived_attrs))) {
return failure();
}
continue;
}
// 3. The preceding op is a constant, thus the value of this constant is
// used to create an attribute of the result op, according to the signature.
Attribute arg_value;
// A failure indicates the argument isn't a constant value, so we should
// not use it as an attribute.
if (!matchPattern(input, m_Constant(&arg_value))) {
return failure();
}
auto attr_name = signature.getArgAttrOfType<StringAttr>(
operand.index(), kAttrArgumentNameAttr);
auto attr_type = signature.getArgAttrOfType<StringAttr>(
operand.index(), kAttrArgumentTypeAttr);
auto value = ProcessAttributeValue(arg_value, attr_type);
arg_attrs->push_back(rewriter.getNamedAttr(attr_name.getValue(), value));
}
return success();
}
Attribute RewriteTFRCallOp::ProcessAttributeValue(Attribute attr,
StringAttr attr_type) const {
if (!attr_type) return attr;
if (attr_type.getValue() == "tensor") {
if (auto f = mlir::dyn_cast<FloatAttr>(attr)) {
RankedTensorType type = RankedTensorType::get({}, f.getType());
return DenseFPElementsAttr::get(type, attr);
}
// TODO(fengliuai): handles ArrayAttr. Note that it can be nested ArrayAttr.
}
return attr;
}
// For each output, uses the attribute name associated to the tfr types to find
// out the attribute value from the collected `attrs` and create the output type
// of the result op by using the attribute value as the element type.
LogicalResult RewriteTFRCallOp::DeriveOutputTypes(
Location loc, FunctionType signature,
const llvm::StringMap<Attribute>& attrs,
SmallVectorImpl<Type>* output_types) const {
for (auto res : llvm::enumerate(signature.getResults())) {
if (auto tensor_type = mlir::dyn_cast<TFRTensorType>(res.value())) {
// tfr.tensor should only have one attribute attached.
auto attr_key = tensor_type.getAttrKeys().front();
Builder builder(signature.getContext());
if (auto attr = attrs.lookup(attr_key.getValue())) {
output_types->push_back(
UnrankedTensorType::get(mlir::cast<TypeAttr>(attr).getValue()));
} else if (Type element_type =
GetFixedElementType(attr_key.getValue(), builder)) {
output_types->push_back(UnrankedTensorType::get(element_type));
} else {
emitError(loc) << "type " << attr_key.getValue()
<< " can't be resolved for the signature of the op";
return failure();
}
continue;
}
if (auto list_type = mlir::dyn_cast<TFRTensorListType>(res.value())) {
// There are two cases: N*T or list(dtype)
auto attr_keys = list_type.getAttrKeys();
// N*T case
if (attr_keys.size() == 2) {
// The first one is N, and the second one is T
int list_size =
mlir::cast<IntegerAttr>(attrs.lookup(attr_keys[0].getValue()))
.getInt();
Type list_type =
mlir::cast<TypeAttr>(attrs.lookup(attr_keys[1].getValue()))
.getValue();
for (int i = 0; i < list_size; ++i) {
output_types->push_back(UnrankedTensorType::get(list_type));
}
continue;
}
// TODO(fengliuai): list(dtype) case
}
return failure();
}
return success();
}
LogicalResult RewriteTFRCallOp::CreateAndReplaceOp(
PatternRewriter& rewriter, CallOp call_op,
const SmallVectorImpl<Type>& output_types,
const SmallVectorImpl<Value>& inputs, const NamedAttrList& attr_list,
const llvm::StringMap<Attribute>& derived_attrs) const {
// Create the new op
Location loc = call_op.getLoc();
rewriter.setInsertionPointAfter(call_op);
std::string tf_op_name = GetTFOpName(call_op.getCallee());
OperationState new_state(loc, tf_op_name, inputs, output_types, attr_list);
Operation* new_op = rewriter.create(new_state);
if (materialize_derived_attrs_) {
for (const auto& attr : derived_attrs) {
// Add or update the derived attribute with the value. Skip the fixed
// element type attributes, in case they are present in the NodeDef.
if (!fixed_elt_type_attrs_.contains(attr.first())) {
new_op->setAttr(attr.first(), attr.second);
}
}
}
// Create the tfr.cast ops on the results and replace the uses of the
// original call op.
TFRTensorType unconstrainted_type = rewriter.getType<TFRTensorType>();
SmallVector<Value, 4> new_results;
for (auto res : llvm::enumerate(call_op.getResultTypes())) {
Type res_type = res.value();
if (mlir::dyn_cast<TFRTensorType>(res_type)) {
Value new_res = new_op->getResult(res.index());
auto casted = CastOp::create(rewriter, loc, res_type, new_res);
new_results.push_back(casted.getOut());
} else if (auto list_type =
mlir::dyn_cast<TFRTensorListType>(res.value())) {
SmallVector<Value, 4> tensor_list;
for (int i = res.index(); i < new_op->getNumResults(); i++) {
Value new_res = new_op->getResult(i);
auto casted =
CastOp::create(rewriter, loc, unconstrainted_type, new_res);
tensor_list.push_back(casted.getOut());
}
auto list_op = BuildListOp::create(rewriter, loc, res_type, tensor_list);
new_results.push_back(list_op.getOut());
}
}
// Copy all the allowed attributes to the new op.
if (failed(CopyNonSymbolRefAttrs(call_op, new_op))) return failure();
rewriter.replaceOp(call_op, new_results);
return success();
}
LogicalResult RewriteTFRCallOp::matchAndRewrite(
CallOp call_op, PatternRewriter& rewriter) const {
// Get the func op and verify that it is external. The type of this external
// func op is used as the signature of the corresponding TF ops. All the
// external func ops have the trailing underscore.
std::string external_callee_name = call_op.getCallee().str().append("_");
TFRFuncOp func = symbol_table_.lookup<TFRFuncOp>(external_callee_name);
if (!func || !func.isExternal()) return failure();
// Get the inputs and attributes. The attributes include these from the
// argument list and also these derived from the inputs.
SmallVector<Value, 4> inputs;
NamedAttrList argument_attrs;
llvm::StringMap<Attribute> derived_attrs;
if (failed(CollectInputsAndAttributes(rewriter, func, call_op, &inputs,
&argument_attrs, &derived_attrs))) {
return failure();
}
// Derive the output types. The result type is derived by using the
// attributes attched to the result type of the signature. The attribute
// value should be either in the attribute argument list or the derived
// attribute from the input tensors. All the result type
// are unranked, and shape inference should be applied afterwards.
SmallVector<Type, 4> output_types;
// Merge the attributes from the argument list to the derived ones.
for (auto& attr : argument_attrs) {
derived_attrs.insert({attr.getName(), attr.getValue()});
}
// Derive the output types by using the attributes attached to the tfr
// types.
if (failed(DeriveOutputTypes(call_op->getLoc(), func.getFunctionType(),
derived_attrs, &output_types))) {
return failure();
}
// Create the new op and replace the old TFR call op.
return CreateAndReplaceOp(rewriter, call_op, output_types, inputs,
argument_attrs, derived_attrs);
}
// Raise TFR call ops to the TF ops.
class RaiseToTFOpsPass
: public PassWrapper<RaiseToTFOpsPass, OperationPass<func::FuncOp>> {
public:
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(RaiseToTFOpsPass)
void getDependentDialects(DialectRegistry& registry) const override {
registry.insert<TFRDialect, TF::TensorFlowDialect, scf::SCFDialect,
arith::ArithDialect, func::FuncDialect>();
}
explicit RaiseToTFOpsPass(std::optional<ModuleOp> tfr_module,
bool materialize_derived_attrs)
: external_tfr_module_(tfr_module),
materialize_derived_attrs_(materialize_derived_attrs) {}
StringRef getArgument() const final { return "tfr-raise-to-tf"; }
StringRef getDescription() const final {
return "Raise all the TFR call ops to TF ops.";
}
void runOnOperation() override;
private:
std::optional<ModuleOp> external_tfr_module_;
const bool materialize_derived_attrs_;
};
void RaiseToTFOpsPass::runOnOperation() {
func::FuncOp func = getOperation();
MLIRContext* ctx = &getContext();
SymbolTable table(external_tfr_module_.has_value()
? *external_tfr_module_
: func->getParentOfType<ModuleOp>());
RewritePatternSet patterns(&getContext());
patterns.add<RewriteTFRCallOp>(ctx, table, materialize_derived_attrs_);
populateCanonicalizationPatterns(func, patterns);
(void)applyPatternsGreedily(func, std::move(patterns));
}
} // namespace
// Creates an instance of the pass to raise TFR call ops to the TF ops.
std::unique_ptr<OperationPass<func::FuncOp>> CreateRaiseToTFOpsPass(
std::optional<ModuleOp> tfr_module, bool materialize_derived_attrs) {
return std::make_unique<RaiseToTFOpsPass>(tfr_module,
materialize_derived_attrs);
}
static PassRegistration<RaiseToTFOpsPass> pass([] {
return CreateRaiseToTFOpsPass();
});
} // namespace TFR
} // namespace mlir
tensorflow/compiler/mlir/tfr/passes/rewrite_quantized_io.cc
+120
View File
@@ -0,0 +1,120 @@
/* 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 <memory>
#include <string>
#include "llvm/ADT/StringRef.h"
#include "llvm/Support/Casting.h"
#include "mlir/Dialect/Quant/IR/QuantTypes.h" // from @llvm-project
#include "mlir/IR/Block.h" // from @llvm-project
#include "mlir/IR/Builders.h" // from @llvm-project
#include "mlir/IR/BuiltinTypes.h" // from @llvm-project
#include "mlir/IR/Operation.h" // from @llvm-project
#include "mlir/IR/Value.h" // from @llvm-project
#include "mlir/Pass/Pass.h" // from @llvm-project
#include "mlir/Pass/PassRegistry.h" // from @llvm-project
#include "tensorflow/compiler/mlir/tfr/ir/tfr_ops.h"
#include "tensorflow/compiler/mlir/tfr/passes/passes.h"
namespace mlir {
namespace TFR {
class RewriteQuantizedIOPass
: public PassWrapper<RewriteQuantizedIOPass, OperationPass<ModuleOp>> {
public:
StringRef getArgument() const final { return "tfr-rewrite-quantized-io"; }
StringRef getDescription() const final {
return "Replaces operands and results that has quantized type with their "
"storage types.";
}
void runOnOperation() override;
};
void RewriteQuantizedIOPass::runOnOperation() {
ModuleOp module = getOperation();
OpBuilder builder(module);
module.walk([&](func::FuncOp func) {
Block& block = func.front();
Operation* terminator = block.getTerminator();
// Replace input_arg(tensor<quant_type>) -> tfr.cast
// with input_arg(tensor<storage_type>) -> tfr.cast
for (BlockArgument arg : block.getArguments()) {
Type arg_type = arg.getType();
if (auto quant_type = llvm::dyn_cast<quant::QuantizedType>(
llvm::cast<TensorType>(arg_type).getElementType())) {
if (arg.hasOneUse() && llvm::isa<TFR::CastOp>(*arg.user_begin())) {
arg.setType(llvm::cast<TensorType>(arg_type).clone(
quant_type.getStorageType()));
} else {
std::string error_message;
llvm::raw_string_ostream os{error_message};
os << "The argument with type ";
arg.getType().print(os);
os << " should have one user, which should be tfr.cast.";
func->emitError(error_message);
return;
}
}
}
builder.setInsertionPoint(terminator);
// Replace tfr.cast(tensor<quant_type>) -> output
// with tfr.cast(tensor<storage_type>) -> output
for (OpOperand& returned_value : terminator->getOpOperands()) {
auto returned_type =
llvm::dyn_cast<TensorType>(returned_value.get().getType());
if (!returned_type ||
!llvm::isa<quant::QuantizedType>(returned_type.getElementType())) {
continue;
}
if (auto returned_op =
returned_value.get().getDefiningOp<TFR::CastOp>()) {
auto new_type = returned_type.clone(
llvm::cast<quant::QuantizedType>(returned_type.getElementType())
.getStorageType());
auto new_op = TFR::CastOp::create(builder, returned_op->getLoc(),
new_type, returned_op.getArg());
returned_value.set(new_op.getResult());
if (returned_op.use_empty()) {
returned_op.erase();
}
} else {
returned_value.get().getDefiningOp()->emitError(
"The producer of quantized type result should be a tfr.cast op.");
return;
}
}
auto new_func_type = builder.getFunctionType(block.getArgumentTypes(),
terminator->getOperandTypes());
func.setType(new_func_type);
});
}
// Creates an instance of the pass to decompose the TF ops.
std::unique_ptr<OperationPass<ModuleOp>> CreateRewriteQuantizedIOPass() {
return std::make_unique<RewriteQuantizedIOPass>();
}
static PassRegistration<RewriteQuantizedIOPass> pass([] {
return CreateRewriteQuantizedIOPass();
});
} // namespace TFR
} // namespace mlir
tensorflow/compiler/mlir/tfr/passes/tfr_opt.cc
+44
View File
@@ -0,0 +1,44 @@
/* 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 "mlir/Dialect/Arith/IR/Arith.h" // from @llvm-project
#include "mlir/Dialect/Func/Extensions/AllExtensions.h" // from @llvm-project
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/Dialect/Quant/IR/Quant.h" // from @llvm-project
#include "mlir/Dialect/Quant/IR/QuantTypes.h" // from @llvm-project
#include "mlir/Dialect/SCF/IR/SCF.h" // from @llvm-project
#include "mlir/Dialect/Shape/IR/Shape.h" // from @llvm-project
#include "mlir/InitAllDialects.h" // from @llvm-project
#include "mlir/InitAllPasses.h" // from @llvm-project
#include "mlir/Tools/mlir-opt/MlirOptMain.h" // from @llvm-project
#include "tensorflow/compiler/mlir/init_mlir.h"
#include "tensorflow/compiler/mlir/quantization/common/ir/QuantOps.h"
#include "tensorflow/compiler/mlir/tensorflow/dialect_registration.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/compiler/mlir/tfr/ir/tfr_ops.h"
int main(int argc, char **argv) {
tensorflow::InitMlir y(&argc, &argv);
mlir::registerAllPasses();
mlir::DialectRegistry registry;
registry.insert<mlir::scf::SCFDialect, mlir::TF::TensorFlowDialect,
mlir::arith::ArithDialect, mlir::func::FuncDialect,
mlir::shape::ShapeDialect, mlir::quant::QuantDialect,
mlir::quant::ir::TFQuantDialect, mlir::TFR::TFRDialect>();
mlir::func::registerAllExtensions(registry);
return failed(mlir::MlirOptMain(argc, argv, "TFR Pass Driver\n", registry));
}
tensorflow/compiler/mlir/tfr/python/composite.py
+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.
# ==============================================================================
"""Op composition registration."""
# TODO(fengliuai): add the tf_export decrator
class Composite(object):
"""A decorator to register a function as a composition for an TF operator.
The argument to the decorator must be the name of a TF raw operator the
function composites for. Decorated function must take positional arguments
which corresponds to the input and attributes in OpDef of the TF operation.
# TODO(fengliuai): more documents here.
Example:
@composite.Composite('AddN')
def _compose_add_n(inputs, N):
if N == 1:
....
"""
# TODO(fengliuai): support input_binding and output_binding so the arguments
# are not positional.
def __init__(self,
op_name,
inputs=None,
attrs=None,
derived_attrs=None,
outputs=None):
self._op_name = op_name
self._inputs = inputs
self._attrs = attrs
self._derived_attrs = derived_attrs
self._outputs = outputs
def __call__(self, compose_fn):
# TODO(fengliuai): more sanity check of the input function and make sure
# the bounded arguments of the function matches the 'inputs' and 'attrs'.
setattr(compose_fn, '_tfr_op_name', self._op_name)
return compose_fn
tensorflow/compiler/mlir/tfr/python/op_reg_gen.py
+145
View File
@@ -0,0 +1,145 @@
# 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.
# ==============================================================================
"""op_reg_gen: Generate op registration code from composite op code."""
# pylint: disable=invalid-name
# pylint: disable=missing-function-docstring
# pylint: disable=g-direct-tensorflow-import
import gast as ast
from tensorflow.python.autograph.pyct import transformer
from tensorflow.python.autograph.pyct import transpiler
from tensorflow.python.framework import op_def_registry
from tensorflow.python.util import tf_inspect
_COMPOSITE_ARG_LIST = ['op_name', 'inputs', 'attrs', 'derived_attrs', 'outputs']
class OpRegGenImpl(transformer.CodeGenerator):
"""Visit the AST and generate C++ op registration functions."""
def __init__(self, ctx):
super(OpRegGenImpl, self).__init__(ctx)
self.ctx = ctx
def visit_Name(self, node):
return node.id
def visit_Constant(self, node):
return node.value
def visit_keyword(self, node):
return node.arg, self.visit(node.value)
def visit_List(self, node):
return [self.visit(cst) for cst in node.elts]
def visit_arguments(self, node):
return [self.visit(arg) for arg in node.args]
def visit_FunctionDef(self, node):
# TODO(fengliuai): create one utility method to match different apis and
# shared it with the tfr_gen.py module.
compose_dec = []
for dec in node.decorator_list:
if isinstance(dec, ast.Call):
if isinstance(dec.func, ast.Attribute) and dec.func.attr == 'Composite':
compose_dec.append(dec)
if isinstance(dec.func, ast.Name) and dec.func.id == 'Composite':
compose_dec.append(dec)
if not compose_dec:
# skip a non-composition function
return
elif len(compose_dec) > 1:
raise KeyError('More than one TF ops decomposes for.')
all_dec_args = {}
for arg_name, arg_value in zip(_COMPOSITE_ARG_LIST, compose_dec[0].args):
all_dec_args[arg_name] = self.visit(arg_value)
kw_dec_args = dict([self.visit(kw) for kw in compose_dec[0].keywords])
if all_dec_args.keys() & kw_dec_args.keys():
raise KeyError('More arguments than expected.')
all_dec_args.update(kw_dec_args)
op_name = all_dec_args['op_name']
op_def = op_def_registry.get(op_name)
if op_def:
if len(all_dec_args) > 1:
# Op has been registered, so it is a user error to specify op def.
raise ValueError('Op has been registered: ' + op_name)
else:
# Op has been registered, then we don't need to generate register code.
return
# Validates the function inputs match what are in the decorator.
inputs = all_dec_args.get('inputs', [])
attrs = all_dec_args.get('attrs', [])
expected_args = [arg.split(':')[0] for arg in inputs + attrs]
all_func_args = self.visit(node.args)
if len(expected_args) != len(all_func_args):
raise KeyError(
'Composition arguments for {} do not match the registration. {} vs {}'
.format(op_name, expected_args, all_func_args))
cxx_reg_code = ['\nREGISTER_OP("{}")'.format(op_name)]
for input_ in inputs:
cxx_reg_code.append('.Input("{}")'.format(input_))
for attr in attrs:
py_str = attr.replace('"', "'")
cxx_reg_code.append('.Attr("{}")'.format(py_str))
for attr in all_dec_args.get('derived_attrs', []):
py_str = attr.replace('"', "'")
cxx_reg_code.append('.Attr("{}")'.format(py_str))
for output_ in all_dec_args.get('outputs', []):
cxx_reg_code.append('.Output("{}")'.format(output_))
cxx_reg_code[-1] += ';\n'
self.emit('\n '.join(cxx_reg_code))
class OpRegGen(transpiler.GenericTranspiler):
"""Transforms Python objects into TFR MLIR source code."""
def transform_ast(self, node, ctx):
gen = OpRegGenImpl(ctx)
gen.visit(node)
return gen.code_buffer
def op_reg_gen(func):
"""Parse a function and emit the TFR functions."""
op_reg_code, _ = OpRegGen().transform(func, None)
return op_reg_code
def gen_register_op(source, method_prefix=None):
"""Parse a python code and emit the TFR functions from a target class."""
mlir_funcs = [
op_reg_gen(func)
for name, func in tf_inspect.getmembers(source, tf_inspect.isfunction)
if not method_prefix or name.startswith(method_prefix)
]
headers = r"""
#include "tensorflow/core/framework/op.h"
namespace tensorflow {
"""
code = '\n'.join(mlir_funcs)
return headers + code + '} // namespace tensorflow\n'
tensorflow/compiler/mlir/tfr/python/op_reg_gen_test.py
+77
View File
@@ -0,0 +1,77 @@
# 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.
# ==============================================================================
"""Tests for `op_reg_gen` module."""
# pylint: disable=missing-function-docstring
# pylint: disable=invalid-name
# pylint: disable=g-direct-tensorflow-import
import sys
from tensorflow.compiler.mlir.python.mlir_wrapper import filecheck_wrapper as fw
from tensorflow.compiler.mlir.tfr.python import composite
from tensorflow.compiler.mlir.tfr.python.op_reg_gen import gen_register_op
from tensorflow.python.platform import test
Composite = composite.Composite
@composite.Composite(
'TestNoOp', derived_attrs=['T: numbertype'], outputs=['o1: T'])
def _composite_no_op():
pass
@Composite(
'TestCompositeOp',
inputs=['x: T', 'y: T'],
attrs=['act: {"", "relu"}', 'trans: bool = true'],
derived_attrs=['T: numbertype'],
outputs=['o1: T', 'o2: T'])
def _composite_op(x, y, act, trans):
return x + act, y + trans
class TFRGenTensorTest(test.TestCase):
"""MLIR Generation Tests for MLIR TFR Program."""
def test_op_reg_gen(self):
cxx_code = gen_register_op(sys.modules[__name__])
cxx_code_exp = r"""
CHECK: #include "tensorflow/core/framework/op.h"
CHECK-EMPTY
CHECK: namespace tensorflow {
CHECK-EMPTY
CHECK-LABEL: REGISTER_OP("TestNoOp")
CHECK-NEXT: .Attr("T: numbertype")
CHECK-NEXT: .Output("o1: T");
CHECK-EMPTY
CHECK-LABEL: REGISTER_OP("TestCompositeOp")
CHECK-NEXT: .Input("x: T")
CHECK-NEXT: .Input("y: T")
CHECK-NEXT: .Attr("act: {'', 'relu'}")
CHECK-NEXT: .Attr("trans: bool = true")
CHECK-NEXT: .Attr("T: numbertype")
CHECK-NEXT: .Output("o1: T")
CHECK-NEXT: .Output("o2: T");
CHECK-EMPTY
CHECK: } // namespace tensorflow
"""
self.assertTrue(fw.check(str(cxx_code), cxx_code_exp), str(cxx_code))
if __name__ == '__main__':
test.main()
tensorflow/compiler/mlir/tfr/python/test_utils.py
+47
View File
@@ -0,0 +1,47 @@
# 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.
"""Test utils for composite op definition."""
from tensorflow.python.eager import backprop
from tensorflow.python.framework import test_util
from tensorflow.python.platform import test
class OpsDefsTest(test.TestCase):
"""Test utils."""
def _assertOpAndComposite(self, vars_, compute_op, compute_composite, kwargs,
op_kwargs=None):
if op_kwargs is None:
op_kwargs = kwargs
if test_util.IsMklEnabled():
self.skipTest("Not compatible with oneDNN custom ops.")
# compute with op.
with backprop.GradientTape() as gt:
for var_ in vars_:
gt.watch(var_)
y = compute_op(**op_kwargs) # uses op and decomposites by the graph pass.
grads = gt.gradient(y, vars_) # uses registered gradient function.
# compute with composition
with backprop.GradientTape() as gt:
for var_ in vars_:
gt.watch(var_)
re_y = compute_composite(**kwargs) # uses composite function.
re_grads = gt.gradient(re_y, vars_) # uses gradients compposite function.
for v, re_v in zip(y, re_y):
self.assertAllClose(v, re_v)
for g, re_g in zip(grads, re_grads):
self.assertAllClose(g, re_g)
tensorflow/compiler/mlir/tfr/python/tfr_gen.py
+1573
View File
File diff suppressed because it is too large Load Diff
tensorflow/compiler/mlir/tfr/python/tfr_gen_test.py
+645
View File
@@ -0,0 +1,645 @@
# 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.
# ==============================================================================
"""Tests for `tfr_gen` module."""
# pylint: disable=missing-function-docstring
import sys
from tensorflow.compiler.mlir.python.mlir_wrapper import filecheck_wrapper as fw
from tensorflow.compiler.mlir.tfr.python import composite
from tensorflow.compiler.mlir.tfr.python.tfr_gen import tfr_gen_from_module as tfr_gen
from tensorflow.compiler.mlir.tfr.resources import gen_test_ops as test_ops
from tensorflow.python.framework import dtypes
from tensorflow.python.ops import gen_array_ops as array_ops
from tensorflow.python.ops import gen_math_ops as math_ops
from tensorflow.python.platform import test
Composite = composite.Composite
#--- test fn for mlir location ---
@Composite('TestInputNOp')
def _tfr_loc_test(x):
n = 10
x_sum = x[0]
for i in range(1, n):
x_sum = math_ops.Add(x_sum, x[i])
return x_sum
#--- test fn for tfr tensors ---
@composite.Composite('TestNoOp')
def _tfr_tensor_empty_arg():
pass
@composite.Composite('TestIdentityOp')
def _tfr_tensor_tensor(x):
return x
@composite.Composite('TestIdentityNOp')
def _tfr_tensor_tensor_list(x):
return x
@composite.Composite('TestInputNOp')
def _tfr_tensor_tensor_list_get_elt(x):
return x[1]
@composite.Composite('TestOutputNOp')
def _tfr_tensor_tensor_list_output(x):
return [x, x]
@composite.Composite('TestTwoInputsOp')
def _tfr_tensor_tensor_list_split(x, y, pred):
z, _ = array_ops.Split(axis=0, value=x, num_split=2)
(y, pred) # pylint: disable=pointless-statement
return z
@composite.Composite('TestTwoOutputsOp')
def _tfr_tensor_two_output(x):
z = array_ops.Split(axis=0, value=x, num_split=2)
return z[0], z[1]
@composite.Composite('TestNumAttrsOp')
def _tfr_tensor_tensor_with_cst(x1, y1, x2, y2):
x = array_ops.OneHot(
indices=[0, 2, -1, x1], depth=y1, on_value=True, off_value=False)
(x, x2, y2) # pylint: disable=pointless-statement
return
#--- test fn for scf control flow ---
@composite.Composite('TestTwoInputsOp')
def _tfr_control_flow_if(x, y, pred):
if pred:
return x
else:
return y
@composite.Composite('TestThreeInputsOp')
def _tfr_control_flow_nested_if(x, y, z, select):
if select == 'x':
return x
elif select == 'y':
return y
else:
return z
@composite.Composite('TestInputNOp')
def _tfr_control_flow_range_for(x):
# TODO(fengliuai): use len(x) instead
n = 10
x_sum = x[0]
for i in range(1, n):
x_sum = math_ops.Add(x_sum, x[i])
return x_sum
@composite.Composite('TestInputNOp')
def _tfr_control_flow_tensor_list_size(ins):
n = len(ins)
if n == 0:
return array_ops.Const(value=[[0, 1], [2, 3]], dtype=dtypes.int64)
else:
return math_ops.AddN(ins)
#--- test fn for tf ops ---
@composite.Composite('TestComplexTFOp')
def _tfr_tf_ops_complex(lhs, rhs):
left_padding, _ = array_ops.SplitV(
value=lhs, size_splits=[rhs, -1], axis=0, num_split=2)
_, right_padding = array_ops.SplitV(
value=lhs, size_splits=[rhs, rhs], axis=1, num_split=2)
return [left_padding, right_padding]
@composite.Composite('TestIdentityOp')
def _tfr_tf_ops_tensor(x):
return array_ops.Identity(x)
@composite.Composite('TestTwoInputsOp')
def _tfr_tf_ops_tensors(x, y, pred):
if pred:
return math_ops.Add(x, y)
else:
return array_ops.Concat(0, [x, y])
@composite.Composite('TestInputNOp')
def _tfr_tf_ops_with_defaults(ins):
return test_ops.TestTwoInputsOp(ins[0], ins[1])
#--- test fn for tfr attributes ---
@composite.Composite('TestNumAttrsOp')
def _tfr_attrs_num_type(x, y, x1, y1):
# int
z0 = [x, y]
z1 = x == y
z2 = x < y
z3 = x <= y
z4 = x > y
z5 = x >= y
z6 = x != y
z7 = x + y
z8 = x - y
z8 += x
z8 += 1
(z0, z1, z2, z3, z4, z5, z6, z7, z8) # pylint: disable=pointless-statement
# float
z9 = x1 > y1
z10 = x1 + y1
z11 = [x1, y1]
(z9, z10, z11) # pylint: disable=pointless-statement
return
@composite.Composite('TestNonNumAttrsOp')
def _tfr_attrs_tfr_type(x, y, z):
z1 = x == y
z2 = x == 'test'
z3 = y == z
(z1, z2, z3) # pylint: disable=pointless-statement
return
#--- test fn for shapes ---
@composite.Composite('TestIdentityOp')
def _tfr_shapes(x):
s1 = x.shape
s3 = x.shape.as_list()
for i in range(len(s3)):
s3[i] # pylint: disable=pointless-statement
for i in range(1, len(s3), 2):
s3[i] # pylint: disable=pointless-statement
s5 = array_ops.Shape(x)
(s1, s3, s5) # pylint: disable=pointless-statement
return x
#--- test fn for nested functions ---
@composite.Composite('TestIdentityNOp')
def _tfr_temp_op(x):
return x
@composite.Composite('TestIdentityOp')
def _tfr_temp_use_op(x):
y = _tfr_temp_op([x])
return y[0]
#--- test fn for quant built-ins ---
# pylint: disable=undefined-variable
@composite.Composite('TestIdentityOp')
def _tfr_quant_test(x):
y = _tfr_quant_raw_data(x)
s, z = _tfr_quant_qparam(x)
s = _tfr_quant_scale_factor(1.0, [s, s])
s = _tfr_quant_scale_factor(1.0, [s])
y = math_ops.Sub(y, z)
qmin, qmax = _tfr_quant_act_range('RELU', 1.0, 0)
(qmin, qmax) # pylint: disable=pointless-statement
d = _tfr_quant_rescale(y, s, 0)
e = math_ops.Cast(x=d, DstT=dtypes.int16)
f = math_ops.Cast(x=e, DstT=dtypes.int8)
return f
@composite.Composite('TestIdentityNOp')
def _tfr_quant_test_n(x):
y = _tfr_quant_raw_data(x)
return y
class TFRGenTestBase(test.TestCase):
def _check_code(self, tfr_code, exp_tfr_code):
return self.assertTrue(fw.check(str(tfr_code), exp_tfr_code), str(tfr_code))
class TFRGenTensorTest(TFRGenTestBase):
"""MLIR Generation Tests for MLIR TFR Program."""
def test_tfr_loc(self):
mlir_code = tfr_gen(sys.modules[__name__], '_tfr_loc', [test_ops])
mlir_code_exp = r"""
CHECK-LABEL: tfr.func @tf__test_input_n_op(%x: !tfr.tensor_list) -> (!tfr.tensor) {
CHECK-NEXT: %[[n:.*]] = arith.constant 10 : i64
CHECK-SAME loc("tfr_gen_test.py":%{{.*}}:6)
CHECK-NEXT: %[[cst:.*]] = arith.constant 0 : index
CHECK-SAME loc("tfr_gen_test.py":%[[sum_line:.*]]:10)
CHECK-NEXT: %[[elt:.*]] = tfr.get_element %x[%[[cst]]] : (!tfr.tensor_list, index) -> !tfr.tensor
CHECK-SAME loc("tfr_gen_test.py":%[[sum_line]]:10)
CHECK-NEXT: %[[cst_1:.*]] = arith.constant 1 : i64
CHECK-SAME loc("tfr_gen_test.py":%[[for_line:.*]]:2)
CHECK-NEXT: %[[begin:.*]] = arith.index_cast %[[cst_1]] : i64 to index
CHECK-SAME loc("tfr_gen_test.py":%[[for_line]]:2)
CHECK-NEXT: %[[end:.*]] = arith.index_cast %[[n]] : i64 to index
CHECK-SAME loc("tfr_gen_test.py":%[[for_line]]:2)
CHECK-NEXT: %[[step:.*]] = arith.constant 1 : index
CHECK-SAME loc("tfr_gen_test.py":%[[for_line]]:2)
CHECK-NEXT: %[[for_stmt:.*]] = scf.for %[[itr_1:.*]] = %[[begin]] to %[[end]] step %[[step]]
CHECK-SAME: iter_args(%[[it_arg:.*]] = %[[elt]]) -> (!tfr.tensor) {
CHECK-NEXT: %[[elt_1:.*]] = tfr.get_element %x[%itr_1] : (!tfr.tensor_list, index) -> !tfr.tensor
CHECK-SAME loc("tfr_gen_test.py":%[[add_line:.*]]:34)
CHECK-NEXT: %[[Add:.*]] = tfr.call @tf__add(%[[it_arg]], %[[elt_1]]) : (!tfr.tensor, !tfr.tensor) -> (!tfr.tensor)
CHECK-SAME loc("tfr_gen_test.py":%[[add_line]]:12)
CHECK-NEXT: scf.yield %[[Add]] : !tfr.tensor
CHECK-SAME loc(unknown)
CHECK-NEXT: }
CHECK-SAME loc("tfr_gen_test.py":%[[for_line]]:2)
CHECK-NEXT: %{{.*}} = arith.constant true
CHECK-SAME loc(unknown)
CHECK-NEXT: tfr.return %[[for_stmt]] : !tfr.tensor
CHECK-SAME loc(unknown)
CHECK-NEXT: }
CHECK-SAME loc("tfr_gen_test.py":%{{def_line:.*}}:0)
"""
self._check_code(mlir_code, mlir_code_exp)
def test_tfr_tensors(self):
mlir_code = tfr_gen(sys.modules[__name__], '_tfr_tensor', [test_ops])
mlir_code_exp = r"""
CHECK-LABEL: tfr.func @tf__test_no_op() -> () {
CHECK-NEXT: tfr.return
CHECK-NEXT: }
CHECK-LABEL: tfr.func @tf__test_identity_op(%x: !tfr.tensor) -> (!tfr.tensor) {
CHECK-NEXT: constant true
CHECK-NEXT: tfr.return %x : !tfr.tensor
CHECK-NEXT: }
CHECK-LABEL: tfr.func @tf__test_identity_n_op(%x: !tfr.tensor_list) -> (!tfr.tensor_list) {
CHECK-NEXT: constant true
CHECK-NEXT: tfr.return %x : !tfr.tensor_list
CHECK-NEXT: }
CHECK-LABEL: tfr.func @tf__test_input_n_op(%x: !tfr.tensor_list) -> (!tfr.tensor) {
CHECK-NEXT: constant true
CHECK-NEXT: %[[index:.*]] = arith.constant 1 : index
CHECK-NEXT: %[[sub:.*]] = tfr.get_element %x[%cst_1] : (!tfr.tensor_list, index) -> !tfr.tensor
CHECK-NEXT: tfr.return %[[sub]] : !tfr.tensor
CHECK-NEXT: }
CHECK-LABEL: tfr.func @tf__test_output_n_op(%x: !tfr.tensor) -> (!tfr.tensor_list) {
CHECK-NEXT: constant true
CHECK-NEXT: %[[list:.*]] = "tfr.build_list"(%x, %x) : (!tfr.tensor, !tfr.tensor) -> !tfr.tensor_list
CHECK-NEXT: tfr.return %[[list]] : !tfr.tensor_list
CHECK-NEXT: }
CHECK-LABEL: tfr.func @tf__test_two_inputs_op(%x: !tfr.tensor, %y: !tfr.tensor, %pred: i1{tfr.name="pred",tfr.default=false}) -> (!tfr.tensor) {
CHECK-NEXT: %[[cst:.*]] = arith.constant 0 : i64
CHECK-NEXT: %[[cst_1:.*]] = arith.constant 2 : i64
CHECK-NEXT: %[[cst_2:.*]] = "tfr.constant_tensor"(%[[cst]]) : (i64) -> !tfr.tensor
CHECK-NEXT: %[[Split:.*]] = tfr.call @tf__split(%[[cst_2]], %x, %[[cst_1]]) : (!tfr.tensor, !tfr.tensor, i64) -> (!tfr.tensor_list)
CHECK-NEXT: %[[cst_4:.*]] = arith.constant 0 : index
CHECK-NEXT: %[[elt:.*]] = tfr.get_element %[[Split]][%idx] : (!tfr.tensor_list, index) -> !tfr.tensor
CHECK-NEXT: %[[cst_5:.*]] = arith.constant 1 : index
CHECK-NEXT: %[[elt_1:.*]] = tfr.get_element %[[Split]][%idx_1] : (!tfr.tensor_list, index) -> !tfr.tensor
CHECK-NEXT: constant true
CHECK-NEXT: tfr.return %[[elt]] : !tfr.tensor
CHECK-NEXT: }
CHECK-LABEL: tfr.func @tf__test_two_outputs_op(%x: !tfr.tensor) -> (!tfr.tensor, !tfr.tensor) {
CHECK-NEXT: %[[cst:.*]] = arith.constant 0 : i64
CHECK-NEXT: %[[cst_1:.*]] = arith.constant 2 : i64
CHECK-NEXT: %[[cst_2:.*]] = "tfr.constant_tensor"(%[[cst]]) : (i64) -> !tfr.tensor
CHECK-NEXT: %[[Split:.*]] = tfr.call @tf__split(%[[cst_2]], %x, %[[cst_1]]) : (!tfr.tensor, !tfr.tensor, i64) -> (!tfr.tensor_list)
CHECK-NEXT: constant true
CHECK-NEXT: %[[cst_4:.*]] = arith.constant 0 : index
CHECK-NEXT: %[[elt:.*]] = tfr.get_element %[[Split]][%cst_4] : (!tfr.tensor_list, index) -> !tfr.tensor
CHECK-NEXT: %[[cst_5:.*]] = arith.constant 1 : index
CHECK-NEXT: %[[elt_1:.*]] = tfr.get_element %[[Split]][%cst_5] : (!tfr.tensor_list, index) -> !tfr.tensor
CHECK-NEXT: tfr.return %[[elt]], %[[elt_1]] : !tfr.tensor, !tfr.tensor
CHECK-NEXT: }
CHECK-LABEL: tfr.func @tf__test_num_attrs_op(%x1: i64{tfr.name="x1",tfr.default=-10}, %y1: i64{tfr.name="y1",tfr.default=1}, %x2: f32{tfr.name="x2",tfr.default=0.0}, %y2: f32{tfr.name="y2",tfr.default=-3.0}) -> () {
CHECK-NEXT: %[[cst:.*]] = arith.constant 0 : i64
CHECK-NEXT: %[[cst_1:.*]] = arith.constant 2 : i64
CHECK-NEXT: %[[cst_2:.*]] = arith.constant 1 : i64
CHECK-NEXT: %[[zero:.*]] = arith.constant 0 : i64
CHECK-NEXT: %[[cst_3:.*]] = arith.subi %zero, %cst_2 : i64
CHECK-NEXT: %[[list:.*]] = "tfr.build_list"(%[[cst]], %[[cst_1]], %[[cst_3]], %x1) : (i64, i64, i64, i64) -> !tfr.attr
CHECK-NEXT: %[[cst_4:.*]] = arith.constant true
CHECK-NEXT: %[[cst_5:.*]] = arith.constant false
CHECK-NEXT: %[[cst_6:.*]] = "tfr.constant_tensor"(%[[list]]) : (!tfr.attr) -> !tfr.tensor
CHECK-NEXT: %[[cst_7:.*]] = "tfr.constant_tensor"(%y1) : (i64) -> !tfr.tensor
CHECK-NEXT: %[[cst_8:.*]] = "tfr.constant_tensor"(%[[cst_4]]) : (i1) -> !tfr.tensor
CHECK-NEXT: %[[cst_9:.*]] = "tfr.constant_tensor"(%[[cst_5]]) : (i1) -> !tfr.tensor
CHECK-NEXT: %[[cst_10:.*]] = arith.constant -1 : i64
CHECK-NEXT: %[[OneHot:.*]] = tfr.call @tf__one_hot(%[[cst_6]], %[[cst_7]], %[[cst_8]], %[[cst_9]], %[[cst_10]])
CHECK-SAME: (!tfr.tensor, !tfr.tensor, !tfr.tensor, !tfr.tensor, i64) -> (!tfr.tensor)
CHECK-NEXT: constant true
CHECK-NEXT: tfr.return
CHECK-NEXT: }
"""
self._check_code(mlir_code, mlir_code_exp)
def test_tfr_control_flow(self):
mlir_code = tfr_gen(sys.modules[__name__], '_tfr_control_flow', [test_ops])
mlir_code_exp = r"""
CHECK-LABEL: tfr.func @tf__test_two_inputs_op(%x: !tfr.tensor, %y: !tfr.tensor,
CHECK-SAME: %pred: i1{tfr.name="pred",tfr.default=false}) -> (!tfr.tensor) {
CHECK-NEXT: %[[if:.*]] = scf.if %pred -> (!tfr.tensor) {
CHECK-NEXT: arith.constant true
CHECK-NEXT: scf.yield %x : !tfr.tensor
CHECK-NEXT: } else {
CHECK-NEXT: arith.constant true
CHECK-NEXT: scf.yield %y : !tfr.tensor
CHECK-NEXT: }
CHECK-NEXT: tfr.return %if_stmt : !tfr.tensor
CHECK-NEXT: }
CHECK-LABEL: tfr.func @tf__test_three_inputs_op(%x: !tfr.tensor, %y: !tfr.tensor, %z: !tfr.tensor,
CHECK-SAME: %select: !tfr.attr{tfr.name="act",tfr.default="z"}) -> (!tfr.tensor) {
CHECK-NEXT: %[[cst:.*]] = tfr.constant "x" -> !tfr.attr
CHECK-NEXT: %[[eq:.*]] = tfr.equal %select, %[[cst]] -> i1
CHECK-NEXT: %[[if_stmt:.*]] = scf.if %[[eq]] -> (!tfr.tensor) {
CHECK-NEXT: %[[cst_1:.*]] = arith.constant true
CHECK-NEXT: scf.yield %x : !tfr.tensor
CHECK-NEXT: } else {
CHECK-NEXT: %[[cst_2:.*]] = tfr.constant "y" -> !tfr.attr
CHECK-NEXT: %[[eq_1:.*]] = tfr.equal %select, %[[cst_2]] -> i1
CHECK-NEXT: %[[if_stmt1:.*]] = scf.if %[[eq_1]] -> (!tfr.tensor) {
CHECK-NEXT: %[[cst_3:.*]] = arith.constant true
CHECK-NEXT: scf.yield %y : !tfr.tensor
CHECK-NEXT: } else {
CHECK-NEXT: %[[cst_4:.*]] = arith.constant true
CHECK-NEXT: scf.yield %z : !tfr.tensor
CHECK-NEXT: }
CHECK-NEXT: scf.yield %[[if_stmt1]] : !tfr.tensor
CHECK-NEXT: }
CHECK-NEXT: tfr.return %[[if_stmt]] : !tfr.tensor
CHECK-NEXT: }
CHECK-LABEL: tfr.func @tf__test_input_n_op(%x: !tfr.tensor_list) -> (!tfr.tensor) {
CHECK-NEXT: %[[n:.*]] = arith.constant 10 : i64
CHECK-NEXT: %[[cst:.*]] = arith.constant 0 : index
CHECK-NEXT: %[[elt:.*]] = tfr.get_element %x[%[[cst]]] : (!tfr.tensor_list, index) -> !tfr.tensor
CHECK-NEXT: %[[cst_1:.*]] = arith.constant 1 : i64
CHECK-NEXT: %[[begin:.*]] = arith.index_cast %[[cst_1]] : i64 to index
CHECK-NEXT: %[[end:.*]] = arith.index_cast %[[n]] : i64 to index
CHECK-NEXT: %[[step:.*]] = arith.constant 1 : index
CHECK-NEXT: %[[for_stmt:.*]] = scf.for %[[itr_1:.*]] = %[[begin]] to %[[end]] step %[[step]]
CHECK-SAME: iter_args(%[[it_arg:.*]] = %[[elt]]) -> (!tfr.tensor) {
CHECK-NEXT: %[[elt_1:.*]] = tfr.get_element %x[%itr_1] : (!tfr.tensor_list, index) -> !tfr.tensor
CHECK-NEXT: %[[Add:.*]] = tfr.call @tf__add(%[[it_arg]], %[[elt_1]]) : (!tfr.tensor, !tfr.tensor) -> (!tfr.tensor)
CHECK-NEXT: scf.yield %[[Add]] : !tfr.tensor
CHECK-NEXT: }
CHECK-NEXT: %{{.*}} = arith.constant true
CHECK-NEXT: tfr.return %[[for_stmt]] : !tfr.tensor
CHECK-NEXT: }
CHECK-LABEL: tfr.func @tf__test_input_n_op(%ins: !tfr.tensor_list) -> (!tfr.tensor) {
CHECK: %[[attr:.*]] = tfr.constant i64 -> !tfr.attr
CHECK: %Const = tfr.call @tf__const(%{{.*}}, %[[attr]]) : (!tfr.attr, !tfr.attr) -> (!tfr.tensor)
"""
self._check_code(mlir_code, mlir_code_exp)
def test_tfr_tf_ops(self):
mlir_code = tfr_gen(sys.modules[__name__], '_tfr_tf_ops', [test_ops])
mlir_code_exp = r"""
CHECK-LABEL: tfr.func @tf__test_complex_tf_op(%lhs: !tfr.tensor, %rhs: !tfr.tensor) -> (!tfr.tensor_list) {
CHECK-NEXT: %[[cst:.*]] = arith.constant 1 : i64
CHECK-NEXT: %[[zero:.*]] = arith.constant 0 : i64
CHECK-NEXT: %[[cst_1:.*]] = arith.subi %[[zero]], %cst : i64
CHECK-NEXT: %[[cst_2:.*]] = "tfr.constant_tensor"(%[[cst_1]]) : (i64) -> !tfr.tensor
CHECK-NEXT: %[[list:.*]] = "tfr.build_list"(%rhs, %[[cst_2]]) : (!tfr.tensor, !tfr.tensor) -> !tfr.tensor_list
CHECK-NEXT: %[[cst_3:.*]] = arith.constant 0 : i64
CHECK-NEXT: %[[cst_4:.*]] = arith.constant 2 : i64
CHECK-NEXT: %[[zero_1:.*]] = arith.constant 0 : i64
CHECK-NEXT: %[[pack:.*]] = tfr.call @tf__pack(%[[list]], %[[zero_1]]) : (!tfr.tensor_list, i64) -> !tfr.tensor
CHECK-NEXT: %[[cst_5:.*]] = "tfr.constant_tensor"(%[[cst_3]]) : (i64) -> !tfr.tensor
CHECK-NEXT: %[[SplitV:.*]] = tfr.call @tf__split_v(%lhs, %[[pack]], %[[cst_5]], %[[cst_4]])
CHECK-NEXT: %[[idx:.*]] = arith.constant 0 : index
CHECK-NEXT: %[[elt:.*]] = tfr.get_element %SplitV[%idx] : (!tfr.tensor_list, index) -> !tfr.tensor
CHECK-NEXT: %[[idx_1:.*]] = arith.constant 1 : index
CHECK-NEXT: %[[elt_1:.*]] = tfr.get_element %SplitV[%idx_1] : (!tfr.tensor_list, index) -> !tfr.tensor
CHECK-NEXT: %[[list_1:.*]] = "tfr.build_list"(%rhs, %rhs) : (!tfr.tensor, !tfr.tensor) -> !tfr.tensor_list
CHECK-NEXT: %[[cst_6:.*]] = arith.constant 1 : i64
CHECK-NEXT: %[[cst_7:.*]] = arith.constant 2 : i64
CHECK-NEXT: %[[zero_2:.*]] = arith.constant 0 : i64
CHECK-NEXT: %[[pack_1:.*]] = tfr.call @tf__pack(%[[list_1]], %[[zero_2]]) : (!tfr.tensor_list, i64) -> !tfr.tensor
CHECK-NEXT: %[[cst_8:.*]] = "tfr.constant_tensor"(%[[cst_6]]) : (i64) -> !tfr.tensor
CHECK-NEXT: %[[SplitV_1:.*]] = tfr.call @tf__split_v(%lhs, %[[pack_1]], %[[cst_8]], %[[cst_7]])
CHECK-NEXT: %[[idx_2:.*]] = arith.constant 0 : index
CHECK-NEXT: %[[elt_2:.*]] = tfr.get_element %SplitV_1[%idx_2] : (!tfr.tensor_list, index) -> !tfr.tensor
CHECK-NEXT: %[[idx_3:.*]] = arith.constant 1 : index
CHECK-NEXT: %[[elt_3:.*]] = tfr.get_element %SplitV_1[%idx_3] : (!tfr.tensor_list, index) -> !tfr.tensor
CHECK-NEXT: %[[cst_9:.*]] = arith.constant true
CHECK-NEXT: %[[list_2:.*]] = "tfr.build_list"(%[[elt]], %[[elt_3]]) : (!tfr.tensor, !tfr.tensor) -> !tfr.tensor_list
CHECK-NEXT: tfr.return %[[list_2]] : !tfr.tensor_list
CHECK-NEXT: }
CHECK-LABEL: tfr.func @tf__test_identity_op(%x: !tfr.tensor) -> (!tfr.tensor) {
CHECK-NEXT: %cst = arith.constant true
CHECK-NEXT: %[[Id:.*]] = tfr.call @tf__identity(%x) : (!tfr.tensor) -> (!tfr.tensor)
CHECK-NEXT: tfr.return %[[Id]] : !tfr.tensor
CHECK-NEXT: }
CHECK-LABEL: tfr.func @tf__test_two_inputs_op(%x: !tfr.tensor, %y: !tfr.tensor,
CHECK-SAME: %pred: i1{tfr.name="pred",tfr.default=false}) -> (!tfr.tensor) {
CHECK-NEXT: %[[if_stmt:.*]] = scf.if %pred -> (!tfr.tensor) {
CHECK-NEXT: %cst = arith.constant true
CHECK-NEXT: %[[Add:.*]] = tfr.call @tf__add(%x, %y) : (!tfr.tensor, !tfr.tensor) -> (!tfr.tensor)
CHECK-NEXT: scf.yield %[[Add]] : !tfr.tensor
CHECK-NEXT: } else {
CHECK-NEXT: %cst_1 = arith.constant true
CHECK-NEXT: %[[cst_2:.*]] = arith.constant 0 : i64
CHECK-NEXT: %[[list:.*]] = "tfr.build_list"(%x, %y) : (!tfr.tensor, !tfr.tensor) -> !tfr.tensor_list
CHECK-NEXT: %[[Concat:.*]] = tfr.call @tf__concat(%[[cst_2]], %[[list]]) : (i64, !tfr.tensor_list) -> (!tfr.tensor)
CHECK-NEXT: scf.yield %[[Concat]] : !tfr.tensor
CHECK-NEXT: }
CHECK-NEXT: tfr.return %[[if_stmt]] : !tfr.tensor
CHECK-NEXT: }
CHECK-LABEL: tfr.func @tf__test_input_n_op(%ins: !tfr.tensor_list) -> (!tfr.tensor) {
CHECK-NEXT: %cst = arith.constant true
CHECK-NEXT: %[[cst_1:.*]] = arith.constant 0 : index
CHECK-NEXT: %[[elt:.*]] = tfr.get_element %ins[%cst_1] : (!tfr.tensor_list, index) -> !tfr.tensor
CHECK-NEXT: %[[cst_2:.*]] = arith.constant 1 : index
CHECK-NEXT: %[[elt_1:.*]] = tfr.get_element %ins[%cst_2] : (!tfr.tensor_list, index) -> !tfr.tensor
CHECK-NEXT: %[[cst_3:.*]] = arith.constant false
CHECK-NEXT: %[[call:.*]] = tfr.call @tf__test_two_inputs_op(
CHECK-SAME: %[[elt]], %[[elt_1]], %[[cst_3]]) : (!tfr.tensor, !tfr.tensor, i1) -> (!tfr.tensor)
CHECK-NEXT: tfr.return %[[call]] : !tfr.tensor
CHECK-NEXT: }
CHECK-LABEL: tfr.func @tf__add_(!tfr.tensor<T>,!tfr.tensor<T>) -> (!tfr.tensor<T>) attributes {T,f32_,i1_,i32_,i64_}
CHECK-LABEL: tfr.func @tf__concat_(!tfr.tensor<i32_>,!tfr.tensor_list<N,T>) -> (!tfr.tensor<T>) attributes {N,T,f32_,i1_,i32_,i64_}
CHECK-LABEL: tfr.func @tf__identity_(!tfr.tensor<T>) -> (!tfr.tensor<T>) attributes {T,f32_,i1_,i32_,i64_}
CHECK-LABEL: tfr.func @tf__pack_(!tfr.tensor_list<N,T>,i64{tfr.name="axis",tfr.type="int"}) -> (!tfr.tensor<T>) attributes {N,T,axis,f32_,i1_,i32_,i64_}
CHECK-LABEL: tfr.func @tf__split_v_(!tfr.tensor<T>,!tfr.tensor<Tlen>,!tfr.tensor<i32_>,i64{tfr.name="num_split",tfr.type="int"}) -> (!tfr.tensor_list<num_split,T>) attributes {T,Tlen,f32_,i1_,i32_,i64_,num_split}
CHECK-LABEL: tfr.func @tf__test_complex_tf_op_(!tfr.tensor<T>,!tfr.tensor<Tlen>,i64{tfr.name="N",tfr.type="int"}) -> (!tfr.tensor_list<N,T>) attributes {N,T,Tlen,f32_,i1_,i32_,i64_}
CHECK-LABEL: tfr.func @tf__test_identity_op_(!tfr.tensor<T>) -> (!tfr.tensor<T>) attributes {T,f32_,i1_,i32_,i64_}
CHECK-LABEL: tfr.func @tf__test_input_n_op_(!tfr.tensor_list<N,T>) -> (!tfr.tensor<T>) attributes {N,T,f32_,i1_,i32_,i64_}
CHECK-LABEL: tfr.func @tf__test_two_inputs_op_(!tfr.tensor<T>,!tfr.tensor<T>,i1{tfr.name="pred",tfr.type="bool"}) -> (!tfr.tensor<T>) attributes {T,f32_,i1_,i32_,i64_,pred}
CHECK-LABEL: tfr.func @tf__test_two_outputs_op_(!tfr.tensor<T>) -> (!tfr.tensor<T>,!tfr.tensor<T>) attributes {T,f32_,i1_,i32_,i64_}
"""
self._check_code(mlir_code, mlir_code_exp)
def test_tfr_attrs(self):
mlir_code = tfr_gen(sys.modules[__name__], '_tfr_attrs', [test_ops])
mlir_code_exp = r"""
CHECK-LABEL: tfr.func @tf__test_num_attrs_op(
CHECK-SAME: %x: i64{tfr.name="x1",tfr.default=-10},
CHECK-SAME: %y: i64{tfr.name="y1",tfr.default=1},
CHECK-SAME: %x1: f32{tfr.name="x2",tfr.default=0.0},
CHECK-SAME: %y1: f32{tfr.name="y2",tfr.default=-3.0}) -> () {
CHECK-NEXT: %{{.*}} = "tfr.build_list"(%x, %y) : (i64, i64) -> !tfr.attr
CHECK-NEXT: %{{.*}} = arith.cmpi "eq", %x, %y : i64
CHECK-NEXT: %{{.*}} = arith.cmpi "ult", %x, %y : i64
CHECK-NEXT: %{{.*}} = arith.cmpi "ule", %x, %y : i64
CHECK-NEXT: %{{.*}} = arith.cmpi "ugt", %x, %y : i64
CHECK-NEXT: %{{.*}} = arith.cmpi "uge", %x, %y : i64
CHECK-NEXT: %{{.*}} = arith.cmpi "ne", %x, %y : i64
CHECK-NEXT: %{{.*}} = arith.addi %x, %y : i64
CHECK-NEXT: %[[sub_1:.*]] = arith.subi %x, %y : i64
CHECK-NEXT: %[[add_1:.*]] = arith.addi %[[sub_1]], %x : i64
CHECK-NEXT: %[[cst:.*]] = arith.constant 1 : i64
CHECK-NEXT: %{{.*}} = arith.addi %[[add_1]], %[[cst]] : i64
CHECK-NEXT: %{{.*}} = arith.cmpf "ugt", %x1, %y1 : f32
CHECK-NEXT: %{{.*}} = arith.addf %x1, %y1 : f32
CHECK-NEXT: %{{.*}} = "tfr.build_list"(%x1, %y1) : (f32, f32) -> !tfr.attr
CHECK-NEXT: %{{.*}} = arith.constant true
CHECK-NEXT: tfr.return
CHECK-NEXT: }
CHECK-LABEL: tfr.func @tf__test_non_num_attrs_op(
CHECK-SAME: %x: !tfr.attr{tfr.name="z"},
CHECK-SAME: %y: !tfr.attr{tfr.name="x",tfr.default="hello"},
CHECK-SAME: %z: !tfr.attr{tfr.name="y",tfr.default=f32}) -> () {
CHECK-NEXT: %{{.*}} = tfr.equal %x, %y -> i1
CHECK-NEXT: %[[cst:.*]] = tfr.constant "test" -> !tfr.attr
CHECK-NEXT: %{{.*}} = tfr.equal %x, %[[cst]] -> i1
CHECK-NEXT: %{{.*}} = tfr.equal %y, %z -> i1
CHECK-NEXT: %{{.*}} = arith.constant true
CHECK-NEXT: tfr.return
CHECK-NEXT: }
"""
self._check_code(mlir_code, mlir_code_exp)
def test_tf_tensor_shape(self):
mlir_code = tfr_gen(sys.modules[__name__], '_tfr_shapes', [test_ops])
mlir_code_exp = r"""
CHECK-LABEL: tfr.func @tf__test_identity_op(%x: !tfr.tensor) -> (!tfr.tensor) {
CHECK-NEXT: %[[shape:.*]] = tfr.get_shape %x -> !shape.shape
CHECK-NEXT: %[[shape_1:.*]] = tfr.get_shape %x -> !shape.shape
CHECK-NEXT: %[[len:.*]] = shape.rank %[[shape_1]] : !shape.shape -> !shape.size
CHECK-NEXT: %[[index:.*]] = shape.size_to_index %[[len]] : !shape.size
CHECK-NEXT: %[[begin:.*]] = arith.constant 0 : index
CHECK-NEXT: %[[step:.*]] = arith.constant 1 : index
CHECK-NEXT: scf.for %[[itr_1:.*]] = %[[begin]] to %[[index]] step %[[step]] {
CHECK-NEXT: %[[size:.*]] = shape.get_extent %[[shape_1]], %[[itr_1]]: !shape.shape, index -> !shape.size
CHECK-NEXT: %[[elt:.*]] = shape.size_to_index %[[size]] : !shape.size
CHECK-NEXT: scf.yield
CHECK-NEXT: }
CHECK-NEXT: %[[cst:.*]] = arith.constant 1 : i64
CHECK-NEXT: %[[len_1:.*]] = shape.rank %shape_1 : !shape.shape -> !shape.size
CHECK-NEXT: %[[len_size_1:.*]] = shape.size_to_index %[[len_1]] : !shape.size
CHECK-NEXT: %[[cst_1:.*]] = arith.constant 2 : i64
CHECK-NEXT: %[[begin_1:.*]] = arith.index_cast %[[cst]] : i64 to index
CHECK-NEXT: %[[step_1:.*]] = arith.index_cast %[[cst_1]] : i64 to index
CHECK-NEXT: scf.for %[[itr_3:.*]] = %[[begin_1]] to %[[len_size_1]] step %[[step_1]]
CHECK: %[[cst:.*]] = tfr.constant i32 -> !tfr.attr
CHECK-NEXT: %[[Shape:.*]] = tfr.call @tf__shape(%x, %[[cst]]) : (!tfr.tensor, !tfr.attr) -> (!tfr.tensor)
CHECK-NEXT: %{{.*}} = arith.constant true
CHECK-NEXT: tfr.return %x : !tfr.tensor
CHECK-NEXT: }
"""
self._check_code(mlir_code, mlir_code_exp)
def test_temp_function(self):
mlir_code = tfr_gen(sys.modules[__name__], '_tfr_temp', [test_ops])
mlir_code_exp = r"""
CHECK-LABEL: tfr.func @tf__test_identity_n_op(%x: !tfr.tensor_list) -> (!tfr.tensor_list)
CHECK-LABEL: tfr.func @tf__test_identity_op(%x: !tfr.tensor) -> (!tfr.tensor) {
CHECK-NEXT: %[[list:.*]] = "tfr.build_list"(%x) : (!tfr.tensor) -> !tfr.tensor_list
CHECK-NEXT: %[[call:.*]] = tfr.call @tf__test_identity_n_op(%[[list]]) : (!tfr.tensor_list)
"""
self._check_code(mlir_code, mlir_code_exp)
def test_quant_builtins(self):
mlir_code = tfr_gen(sys.modules[__name__], '_tfr_quant', [test_ops])
mlir_code_exp = r"""
CHECK-LABEL: tfr.func @tf__test_identity_op(%x: !tfr.tensor) -> (!tfr.tensor) {
CHECK-NEXT: %[[raw_data:.*]] = tfr.quant_raw_data(%x) : (!tfr.tensor) -> (!tfr.tensor)
CHECK-NEXT: %[[qparam:.*]]:2 = tfr.quant_qparam(%x) : (!tfr.tensor) -> (!tfr.tensor, !tfr.tensor)
CHECK: %[[list:.*]] = "tfr.build_list"(%[[qparam]]#0, %[[qparam]]#0) : (!tfr.tensor, !tfr.tensor) -> !tfr.tensor_list
CHECK: %[[factor:.*]] = tfr.quant_scale_factor(%{{.*}}, %[[list]]) : (f32, !tfr.tensor_list) -> (!tfr.tensor)
CHECK: %[[list1:.*]] = "tfr.build_list"(%[[factor]]) : (!tfr.tensor) -> !tfr.tensor_list
CHECK: %[[factor1:.*]] = tfr.quant_scale_factor(%{{.*}}, %[[list1]]) : (f32, !tfr.tensor_list) -> (!tfr.tensor)
CHECK-NEXT: %[[Sub:.*]] = tfr.call @tf__sub(%[[raw_data]], %[[qparam]]#1) : (!tfr.tensor, !tfr.tensor) -> (!tfr.tensor)
CHECK: %[[act_range:.*]]:2 = tfr.quant_act_range(%{{.*}}, %{{.*}}, %{{.*}}) : (!tfr.attr, f32, i64) -> (!tfr.tensor, !tfr.tensor)
CHECK: %[[rescale:.*]] = tfr.quant_rescale(%[[Sub]], %[[factor1]], %{{.*}}) : (!tfr.tensor, !tfr.tensor, i64) -> (!tfr.tensor)
CHECK: %[[attr:.*]] = tfr.constant i16 -> !tfr.attr
CHECK: %[[Cast:.*]] = tfr.call @tf__cast(%[[rescale]], %[[attr]], %{{.*}}) : (!tfr.tensor, !tfr.attr, i1) -> (!tfr.tensor)
CHECK: %[[attr_1:.*]] = tfr.constant i8 -> !tfr.attr
CHECK: tfr.call @tf__cast(%[[Cast]], %[[attr_1]], %{{.*}}) : (!tfr.tensor, !tfr.attr, i1) -> (!tfr.tensor)
CHECK: }
CHECK-LABEL: tfr.func @tf__test_identity_n_op(%x: !tfr.tensor_list) -> (!tfr.tensor_list) {
CHECK-NEXT: %[[raw_data:.*]] = tfr.quant_raw_data(%x) : (!tfr.tensor_list) -> (!tfr.tensor_list)
CHECK: tfr.return %[[raw_data:.*]] : !tfr.tensor_list
CHECK: }
"""
self._check_code(mlir_code, mlir_code_exp)
if __name__ == '__main__':
test.main()
tensorflow/compiler/mlir/tfr/python/tfr_wrapper.cc
+62
View File
@@ -0,0 +1,62 @@
/* 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 <string>
#include "llvm/Support/LogicalResult.h"
#include "llvm/Support/MemoryBuffer.h"
#include "llvm/Support/SMLoc.h"
#include "llvm/Support/SourceMgr.h"
#include "mlir/Dialect/Arith/IR/Arith.h" // from @llvm-project
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/Dialect/SCF/IR/SCF.h" // from @llvm-project
#include "mlir/Dialect/Shape/IR/Shape.h" // from @llvm-project
#include "mlir/IR/AsmState.h" // from @llvm-project
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/IR/Diagnostics.h" // from @llvm-project
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "mlir/IR/Verifier.h" // from @llvm-project
#include "mlir/Parser/Parser.h" // from @llvm-project
#include "pybind11/pybind11.h" // from @pybind11
#include "pybind11/stl.h" // from @pybind11
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_dialect.h"
#include "tensorflow/compiler/mlir/tfr/ir/tfr_ops.h"
#include "tensorflow/python/lib/core/pybind11_lib.h"
PYBIND11_MODULE(tfr_wrapper, m) {
m.def("verify", [](std::string input) {
mlir::DialectRegistry registry;
registry.insert<mlir::arith::ArithDialect, mlir::scf::SCFDialect,
mlir::TF::TensorFlowDialect, mlir::func::FuncDialect,
mlir::shape::ShapeDialect, mlir::TFR::TFRDialect>();
mlir::MLIRContext ctx(registry);
ctx.loadAllAvailableDialects();
llvm::SourceMgr source_mgr = llvm::SourceMgr();
source_mgr.AddNewSourceBuffer(llvm::MemoryBuffer::getMemBuffer(input),
llvm::SMLoc());
auto module = mlir::parseSourceFile<mlir::ModuleOp>(source_mgr, &ctx);
if (!module) {
return false;
}
mlir::SourceMgrDiagnosticHandler sourceMgrHandler(source_mgr, &ctx);
if (failed(mlir::verify(*module))) {
module->emitError("Invalid MLIR module: failed verification.");
return false;
}
return true;
});
}
tensorflow/compiler/mlir/tfr/resources/BUILD
+25
View File
@@ -0,0 +1,25 @@
load("//tensorflow/compiler/mlir/tfr:build_defs.bzl", "gen_op_bindings")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
default_visibility = [
":friends",
],
licenses = ["notice"],
)
package_group(
name = "friends",
packages = [
"//tensorflow/compiler/mlir/...",
],
)
filegroup(
name = "decomposition_lib",
srcs = ["decomposition_lib.mlir"],
)
gen_op_bindings(name = "composite")
gen_op_bindings(name = "test")
tensorflow/compiler/mlir/tfr/resources/composite_ops.cc
+36
View File
@@ -0,0 +1,36 @@
/* 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/core/framework/op.h"
namespace tensorflow {
REGISTER_OP("MyAddN")
.Input("inputs: N * T")
.Output("sum: T")
.Attr("N: int >= 1")
.Attr("T: {numbertype, variant}")
.SetIsCommutative()
.SetIsAggregate();
REGISTER_OP("MyBiasedDense")
.Input("input: T")
.Input("weight: T")
.Input("bias: T")
.Output("out: T")
.Attr("T: {float, int8}")
.Attr("act: {'', 'relu', 'relu6'} = ''");
} // namespace tensorflow
tensorflow/compiler/mlir/tfr/resources/decomposition_lib.mlir
+129
View File
@@ -0,0 +1,129 @@
// 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.
// ==============================================================================
// A test resource file which contains some pre-defined internal tfr.functions
// for decomposition and external tfr.functions for raising the decomposition
// result to the ops in the TF dialect.
//
// All the tfr.func functions are supposed to be translated from the Python
// function with tf.composite annotation.
// All the external tfr.func functions modeles the op signature defined by
// OpDefs.
tfr.func @tf__my_add_n(%values: !tfr.tensor_list,
%n: i64 {tfr.name="N"}) -> !tfr.tensor {
%index = arith.constant 0 : index
%cst = arith.constant 1 : i64
%eq = arith.cmpi eq, %n, %cst : i64
%v1 = tfr.get_element %values[%index] : (!tfr.tensor_list, index) -> !tfr.tensor
%res = scf.if %eq -> !tfr.tensor {
scf.yield %v1 : !tfr.tensor
} else {
%step = arith.index_cast %cst : i64 to index
%end = arith.index_cast %n : i64 to index
%reduce = scf.for %i = %step to %end step %step iter_args(%reduce_iter=%v1) -> !tfr.tensor {
%v = tfr.get_element %values[%i] : (!tfr.tensor_list, index) -> !tfr.tensor
%reduce_next = tfr.call @tf__add(%reduce_iter, %v) : (!tfr.tensor, !tfr.tensor) -> !tfr.tensor
scf.yield %reduce_next : !tfr.tensor
}
scf.yield %reduce : !tfr.tensor
}
tfr.return %res : !tfr.tensor
}
tfr.func @tf__my_add_n_(!tfr.tensor_list<N,T>, i64 {tfr.name="N"}) -> !tfr.tensor attributes {N,T}
// Translated from tf.compose Python function.
tfr.func @tf__my_biased_dense(%input: !tfr.tensor, %weight: !tfr.tensor,
%bias: !tfr.tensor,
%act: !tfr.attr{tfr.name="act", tfr.default=""}) -> !tfr.tensor {
%dot = tfr.call @tf__mat_mul(%input, %weight) : (!tfr.tensor, !tfr.tensor) -> !tfr.tensor
%add = tfr.call @tf__add(%dot, %bias) : (!tfr.tensor, !tfr.tensor) -> !tfr.tensor
%relu = tfr.constant "relu" -> !tfr.attr
%relu6 = tfr.constant "relu6" -> !tfr.attr
%is_relu = tfr.equal %act, %relu -> i1
%res = scf.if %is_relu -> !tfr.tensor {
%applied_relu = tfr.call @tf__relu(%add) : (!tfr.tensor) -> !tfr.tensor
scf.yield %applied_relu : !tfr.tensor
} else {
%is_relu6 = tfr.equal %act, %relu6 -> i1
%res1 = scf.if %is_relu6 -> !tfr.tensor {
%applied_relu6 = tfr.call @tf__relu6(%add) : (!tfr.tensor) -> !tfr.tensor
scf.yield %applied_relu6 : !tfr.tensor
} else {
scf.yield %add : !tfr.tensor
}
scf.yield %res1 : !tfr.tensor
}
tfr.return %res : !tfr.tensor
}
tfr.func @tf__my_biased_dense_(!tfr.tensor<T>, !tfr.tensor<T>, !tfr.tensor<T>,
!tfr.attr{tfr.name="act", tfr.default=""}) -> !tfr.tensor attributes {T}
// This is a wong decomposition and used to verify that tf.Elu isn't decomposed
// since its kernel has been registered.
tfr.func @tf__elu_(%input: !tfr.tensor) -> !tfr.tensor {
tfr.return %input : !tfr.tensor
}
// Translated from:
//
// REGISTER_OP("Add")
// .Input("x: T")
// .Input("y: T")
// .Output("z: T")
// .Attr(
// "T: {bfloat16, half, float, double, uint8, int8, int16, int32, int64, "
// "complex64, complex128, string}")
tfr.func @tf__add_(!tfr.tensor<T>, !tfr.tensor<T>)
-> !tfr.tensor<T> attributes{T}
// Translated from:
//
// REGISTER_OP("MatMul")
// .Input("a: T")
// .Input("b: T")
// .Output("product: T")
// .Attr("transpose_a: bool = false")
// .Attr("transpose_b: bool = false")
// .Attr("T: {bfloat16, half, float, double, int32, int64, complex64, complex128}")
// T is a derived attribute.
// transpose_a and transpose_b is materialized attributes.
tfr.func @tf__mat_mul_(!tfr.tensor<T>, !tfr.tensor<T>,
i1 {tfr.name="transpose_a", tfr.default=false},
i1 {tfr.name="transpose_b", tfr.default=false})
-> !tfr.tensor<T> attributes{T}
// Translated from:
//
// REGISTER_OP("Relu")
// .Input("features: T")
// .Output("activations: T")
// .Attr("T: {realnumbertype, qint8}")
// T is a derived attribute.
tfr.func @tf__relu_(!tfr.tensor<T>) -> !tfr.tensor<T> attributes{T}
// Translated from:
//
// REGISTER_OP("Relu6")
// .Input("features: T")
// .Output("activations: T")
// .Attr("T: {realnumbertype}")
// T is a derived attribute.
tfr.func @tf__relu6_(!tfr.tensor<T>) -> !tfr.tensor<T> attributes{T}
tensorflow/compiler/mlir/tfr/resources/test_ops.cc
+85
View File
@@ -0,0 +1,85 @@
/* 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/core/framework/op.h"
namespace tensorflow {
REGISTER_OP("TestNoOp");
REGISTER_OP("TestIdentityOp")
.Input("input: T")
.Output("output: T")
.Attr("T: numbertype");
REGISTER_OP("TestIdentityNOp")
.Input("input: N * T")
.Output("output: N * T")
.Attr("N: int >= 1")
.Attr("T: numbertype");
REGISTER_OP("TestInputNOp")
.Input("input: N * T")
.Output("output: T")
.Attr("N: int >= 1")
.Attr("T: numbertype");
REGISTER_OP("TestOutputNOp")
.Input("input: T")
.Output("output: N * T")
.Attr("N: int >= 1")
.Attr("T: numbertype");
REGISTER_OP("TestTwoInputsOp")
.Input("lhs: T")
.Input("rhs: T")
.Output("output: T")
.Attr("T: numbertype")
.Attr("pred: bool = false");
REGISTER_OP("TestComplexTFOp")
.Input("lhs: T")
.Input("rhs: Tlen")
.Output("output: N * T")
.Attr("N: int >= 1")
.Attr("T: numbertype")
.Attr("Tlen: {int32, int64} = DT_INT64");
REGISTER_OP("TestNumAttrsOp")
.Attr("x1: int = -10")
.Attr("y1: int = 1")
.Attr("x2: float = 0.0")
.Attr("y2: float = -3.0");
REGISTER_OP("TestNonNumAttrsOp")
.Attr("z: shape")
.Attr("x: string = 'hello'")
.Attr("y: type = DT_FLOAT");
REGISTER_OP("TestThreeInputsOp")
.Input("x: T")
.Input("y: T")
.Input("z: T")
.Output("output: T")
.Attr("T: numbertype")
.Attr("act: {'x', 'y', 'z'} = 'z'");
REGISTER_OP("TestTwoOutputsOp")
.Input("input: T")
.Output("output1: T")
.Output("output2: T")
.Attr("T: numbertype");
} // namespace tensorflow
tensorflow/compiler/mlir/tfr/tests/canonicalize.mlir
+244
View File
@@ -0,0 +1,244 @@
// 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.
// ==============================================================================
// RUN: tfr-opt %s -canonicalize="test-convergence" -verify-diagnostics -split-input-file | FileCheck %s
// Tests for ops with canonicalization patterns.
// CHECK-LABEL: get_real_shape
func.func @get_real_shape(%arg0: tensor<1x2xf32>) -> tensor<2xindex> {
%0 = "tfr.cast"(%arg0) : (tensor<1x2xf32>) -> !tfr.tensor
%1 = tfr.get_shape %0 -> !shape.shape
%2 = shape.to_extent_tensor %1 : !shape.shape -> tensor<2xindex>
func.return %2 : tensor<2xindex>
// CHECK-NEXT: %[[e:.*]] = shape.const_shape [1, 2] : tensor<2xindex>
// CHECK-NEXT: return %[[e]] : tensor<2xindex>
}
// CHECK-LABEL: equal
func.func @equal() -> (i1, i1, i1, i1) {
%0 = tfr.constant f32 -> !tfr.attr
%1 = tfr.constant f32 -> !tfr.attr
%2 = tfr.constant i32 -> !tfr.attr
%same_type = tfr.equal %0,%1 -> i1
%diff_type = tfr.equal %0,%2 -> i1
%3 = tfr.constant "hello" -> !tfr.attr
%4 = tfr.constant "hello" -> !tfr.attr
%5 = tfr.constant "how are you" -> !tfr.attr
%same_str = tfr.equal %3,%4 -> i1
%diff_str = tfr.equal %3,%5 -> i1
func.return %same_type, %diff_type, %same_str, %diff_str : i1, i1, i1, i1
// CHECK-DAG: %true = arith.constant true
// CHECK-DAG: %false = arith.constant false
// CHECK-NEXT: return %true, %false, %true, %false : i1, i1, i1, i1
}
// -----
// CHECK-LABEL: constant_tensor_array
func.func @constant_tensor_array() -> !tfr.tensor {
%0 = tfr.constant [1, -1, 3] -> !tfr.attr
%1 = "tfr.constant_tensor"(%0) : (!tfr.attr) -> !tfr.tensor
func.return %1 : !tfr.tensor
// CHECK-NEXT: %[[RES:.*]] = "tf.Const"() <{value = dense<[1, -1, 3]> : tensor<3xi64>}> : () -> tensor<3xi64>
// CHECK-NEXT: "tfr.cast"(%[[RES]]) : (tensor<3xi64>) -> !tfr.tensor
// CHECK-NEXT: return
}
// -----
// CHECK-LABEL: constant_tensor_scalar
func.func @constant_tensor_scalar() -> !tfr.tensor {
%0 = "arith.constant"() {value = 42 : i32} : () -> i32
%1 = "tfr.constant_tensor"(%0) : (i32) -> !tfr.tensor
func.return %1 : !tfr.tensor
// CHECK-NEXT: %[[RES:.*]] = "tf.Const"() <{value = dense<42> : tensor<i32>}> : () -> tensor<i32>
// CHECK-NEXT: "tfr.cast"(%[[RES]]) : (tensor<i32>) -> !tfr.tensor
// CHECK-NEXT: return
}
// -----
// CHECK-LABEL: quant_raw_data
func.func @quant_raw_data(%arg0: tensor<1x10x!quant.uniform<i8:f32, 0.1:1>>) -> tensor<1x10x!quant.uniform<i8:f32, 0.2:2>> {
%0 = "tfr.cast"(%arg0) : (tensor<1x10x!quant.uniform<i8:f32, 0.1:1>>) -> !tfr.tensor
%1 = tfr.quant_raw_data(%0) : (!tfr.tensor) -> !tfr.tensor
%2 = tfr.call @tf__risc(%1) : (!tfr.tensor) -> !tfr.tensor
%3 = "tfr.cast"(%2) : (!tfr.tensor) -> tensor<1x10x!quant.uniform<i8:f32, 0.3:3>>
%4 = "tfr.cast"(%3) : (tensor<1x10x!quant.uniform<i8:f32, 0.3:3>>) -> !tfr.tensor
%5 = tfr.quant_raw_data(%4) : (!tfr.tensor) -> !tfr.tensor
%6 = tfr.call @tf__risc(%5) : (!tfr.tensor) -> !tfr.tensor
%7 = "tfr.cast"(%6) : (!tfr.tensor) -> tensor<1x10x!quant.uniform<i8:f32, 0.2:2>>
func.return %7 : tensor<1x10x!quant.uniform<i8:f32, 0.2:2>>
// CHECK: %[[c0:.*]] = "tfr.cast"(%[[arg0:.*]]) : (tensor<1x10x!quant.uniform<i8:f32, 1.000000e-01:1>>) -> !tfr.tensor
// CHECK: %[[c1:.*]] = tfr.call @tf__risc(%[[c0]]) : (!tfr.tensor) -> !tfr.tensor
// CHECK: %[[c2:.*]] = tfr.call @tf__risc(%[[c1]]) : (!tfr.tensor) -> !tfr.tensor
// CHECK: %[[c3:.*]] = "tfr.cast"(%[[c2]]) : (!tfr.tensor) -> tensor<1x10x!quant.uniform<i8:f32, 2.000000e-01:2>>
// CHECK: return %[[c3]]
}
// -----
// CHECK-LABEL: quant_raw_data_with_list
func.func @quant_raw_data_with_list(%arg0: !tfr.tensor, %arg1: !tfr.tensor) -> !tfr.tensor {
%cst_1 = "tf.Const"() <{value = dense<1> : tensor<i64>}> : () -> tensor<i64>
%1 = "tfr.cast"(%arg0) : (!tfr.tensor) -> tensor<1x4x4x3x!quant.uniform<i8:f32, 0.0078420601785182952:-1>>
%2 = "tfr.cast"(%arg1) : (!tfr.tensor) -> tensor<1x3x4x3x!quant.uniform<i8:f32, 0.0078420601785182952:-1>>
%3 = "tfr.cast"(%2) : (tensor<1x3x4x3x!quant.uniform<i8:f32, 0.0078420601785182952:-1>>) -> !tfr.tensor
%4 = "tfr.cast"(%1) : (tensor<1x4x4x3x!quant.uniform<i8:f32, 0.0078420601785182952:-1>>) -> !tfr.tensor
%5 = "tfr.build_list"(%3, %4) : (!tfr.tensor, !tfr.tensor) -> !tfr.tensor_list
%6 = tfr.quant_raw_data(%5) : (!tfr.tensor_list) -> !tfr.tensor_list
%7 = "tfr.cast"(%cst_1) : (tensor<i64>) -> !tfr.tensor
%8 = tfr.call @tf__concat(%7, %6) : (!tfr.tensor, !tfr.tensor_list) -> !tfr.tensor
func.return %8 : !tfr.tensor
// CHECK: %[[CONST_0:.*]] = "tf.Const"() <{value = dense<1> : tensor<i64>}> : () -> tensor<i64>
// CHECK: %[[BUILD_LIST_0:.*]] = "tfr.build_list"(%arg1, %arg0) : (!tfr.tensor, !tfr.tensor) -> !tfr.tensor_list
// CHECK: %[[CAST_0:.*]] = "tfr.cast"(%[[CONST_0]]) : (tensor<i64>) -> !tfr.tensor
// CHECK: %[[CONCAT_O:.*]] = tfr.call @tf__concat(%[[CAST_0]], %[[BUILD_LIST_0]]) : (!tfr.tensor, !tfr.tensor_list) -> !tfr.tensor
// CHECK: return %[[CONCAT_O]] : !tfr.tensor
}
// -----
// CHECK-LABEL: cast_with_unranked_quant
func.func @cast_with_unranked_quant(%arg0: tensor<*xi8>, %arg1: tensor<*xi8>) -> tensor<*xf32> {
%0 = "tf.MaximumFloat"(%arg0, %arg1) : (tensor<*xi8>, tensor<*xi8>) -> tensor<*xi8>
%1 = "tfr.cast"(%0) : (tensor<*xi8>) -> !tfr.tensor
%2 = "tfr.cast"(%1) : (!tfr.tensor) -> tensor<*x!quant.uniform<i8:f32, 0.0065901698544621468:-19>>
%3 = "tf.DequantizeFloat"(%2) : (tensor<*x!quant.uniform<i8:f32, 0.0065901698544621468:-19>>) -> tensor<*xf32>
func.return %3 : tensor<*xf32>
// The cast ops should not be removed in this case or it will result in an
// invalid DequantizeFloat op as following:
// %0 = "tf.MaximumFloat"(%arg0, %arg1) : (tensor<*xi8>, tensor<*xi8>) -> tensor<*xi8>
// %1 = "tf.DequantizeFloat"(%0) : (tensor<*xi8>) -> tensor<*xf32>
// CHECK: %[[MAXIMUMFLOAT_0:.*]] = "tf.MaximumFloat"(%arg0, %arg1) : (tensor<*xi8>, tensor<*xi8>) -> tensor<*xi8>
// CHECK: %[[CAST_0:.*]] = "tfr.cast"(%[[MAXIMUMFLOAT_0]]) : (tensor<*xi8>) -> !tfr.tensor
// CHECK: %[[CAST_1:.*]] = "tfr.cast"(%[[CAST_0]]) : (!tfr.tensor) -> tensor<*x!quant.uniform<i8:f32, 0.0065901698544621468:-19>>
// CHECK: %[[DEQUANTIZEFLOAT_0:.*]] = "tf.DequantizeFloat"(%[[CAST_1]]) : (tensor<*x!quant.uniform<i8:f32, 0.0065901698544621468:-19>>) -> tensor<*xf32>
// CHECK: return %[[DEQUANTIZEFLOAT_0]] : tensor<*xf32>
}
// -----
// CHECK-LABEL: quant_qparam
func.func @quant_qparam(%arg0: tensor<1x10x!quant.uniform<i8:f32, 0.1:42>>) -> (tensor<f32>, tensor<i32>) {
%0 = "tfr.cast"(%arg0) : (tensor<1x10x!quant.uniform<i8:f32, 0.1:42>>) -> !tfr.tensor
%scale, %zp = tfr.quant_qparam(%0) : (!tfr.tensor) -> (!tfr.tensor, !tfr.tensor)
%1 = "tfr.cast"(%scale) : (!tfr.tensor) -> tensor<f32>
%2 = "tfr.cast"(%zp) : (!tfr.tensor) -> tensor<i32>
func.return %1, %2 : tensor<f32>, tensor<i32>
// CHECK-DAG: %[[scale:.*]] = "tf.Const"() <{value = dense<1.000000e-01> : tensor<f32>}>
// CHECK-DAG: %[[zp:.*]] = "tf.Const"() <{value = dense<42> : tensor<i32>}> : () -> tensor<i32>
// CHECK: return %[[scale]], %[[zp]]
}
// CHECK-LABEL: quant_qparam_per_channel
func.func @quant_qparam_per_channel(%arg0: tensor<1x3x!quant.uniform<i8:f32:1, {0.1:1, 0.2:2, 0.3:3}>>) -> (tensor<3xf32>, tensor<3xi32>) {
%0 = "tfr.cast"(%arg0) : (tensor<1x3x!quant.uniform<i8:f32:1, {0.1:1, 0.2:2, 0.3:3}>>) -> !tfr.tensor
%scale, %zp = tfr.quant_qparam(%0) : (!tfr.tensor) -> (!tfr.tensor, !tfr.tensor)
%1 = "tfr.cast"(%scale) : (!tfr.tensor) -> tensor<3xf32>
%2 = "tfr.cast"(%zp) : (!tfr.tensor) -> tensor<3xi32>
func.return %1, %2 : tensor<3xf32>, tensor<3xi32>
// CHECK-DAG: %[[scale:.*]] = "tf.Const"() <{value = dense<[1.000000e-01, 2.000000e-01, 3.000000e-01]> : tensor<3xf32>}>
// CHECK-DAG: %[[zp:.*]] = "tf.Const"() <{value = dense<[1, 2, 3]> : tensor<3xi32>}> : () -> tensor<3xi32>
// CHECK: return %[[scale]], %[[zp]]
}
// CHECK-LABEL: quant_qparam_invalid
func.func @quant_qparam_invalid(%arg0: tensor<1x3x!quant.calibrated<f32<-1.0:1.0>>>) -> (!tfr.tensor, !tfr.tensor) {
%0 = "tfr.cast"(%arg0) : (tensor<1x3x!quant.calibrated<f32<-1.0:1.0>>>) -> !tfr.tensor
%scale, %zp = tfr.quant_qparam(%0) : (!tfr.tensor) -> (!tfr.tensor, !tfr.tensor)
func.return %scale, %zp: !tfr.tensor, !tfr.tensor
// CHECK: %[[scale:.*]], %[[zp:.*]] = tfr.quant_qparam(%[[input:.*]]) : (!tfr.tensor) -> (!tfr.tensor, !tfr.tensor)
// CHECK: return %[[scale]], %[[zp]]
}
// -----
// CHECK-LABEL: redundant_cast_with_different_element_type
func.func @redundant_cast_with_different_element_type(%arg0: tensor<*xf32>) -> (tensor<*xi32>, tensor<2xi32>) {
%0 = "tfr.cast"(%arg0) : (tensor<*xf32>) -> !tfr.tensor
%1 = "tfr.cast"(%0) : (!tfr.tensor) -> tensor<*xi32>
%2 = "tfr.cast"(%0) : (!tfr.tensor) -> tensor<2xi32>
func.return %1, %2 : tensor<*xi32>, tensor<2xi32>
// CHECK: %[[tf_cast_unranked:.*]] = "tf.Cast"(%arg0) <{Truncate = false}> : (tensor<*xf32>) -> tensor<*xi32>
// CHECK: %[[ensure_shape:.*]] = "tf.EnsureShape"(%arg0) <{shape = #tf_type.shape<2>}> : (tensor<*xf32>) -> tensor<2xf32>
// CHECK: %[[tf_cast_ranked:.*]] = "tf.Cast"(%[[ensure_shape]]) <{Truncate = false}> : (tensor<2xf32>) -> tensor<2xi32>
// CHECK: return %[[tf_cast_unranked]], %[[tf_cast_ranked]] : tensor<*xi32>, tensor<2xi32>
}
// -----
// CHECK-LABEL: redundant_cast_with_quant_type
func.func @redundant_cast_with_quant_type(%arg0: tensor<10x!quant.uniform<i8:f32, 0.0039133410900831223:-128>>) -> (tensor<10xi32>) {
%0 = "tfr.cast"(%arg0) : (tensor<10x!quant.uniform<i8:f32, 0.0039133410900831223:-128>>) -> !tfr.tensor
%1 = tfr.quant_raw_data(%0) : (!tfr.tensor) -> !tfr.tensor
%2 = "tfr.cast"(%1) : (!tfr.tensor) -> tensor<10xi8>
%3 = "tf.Cast"(%2) {Truncate = false} : (tensor<10xi8>) -> tensor<10xi32>
func.return %3 : tensor<10xi32>
// CHECK: %[[CAST_0:.*]] = "tfr.cast"(%arg0) : (tensor<10x!quant.uniform<i8:f32, 0.0039133410900831223:-128>>) -> !tfr.tensor
// CHECK: %[[CAST_1:.*]] = "tfr.cast"(%[[CAST_0]]) : (!tfr.tensor) -> tensor<10xi8>
// CHECK: %[[CAST_2:.*]] = "tf.Cast"(%[[CAST_1]]) <{Truncate = false}> : (tensor<10xi8>) -> tensor<10xi32>
// CHECK: return %[[CAST_2]] : tensor<10xi32>
}
// -----
// CHECK-LABEL: build_const_list
func.func @build_const_list() -> !tfr.attr {
%0 = "arith.constant"() {value = 42 : i32} : () -> i32
%1 = "arith.constant"() {value = 41 : i32} : () -> i32
%2 = "tfr.build_list"(%0, %1) : (i32, i32) -> !tfr.attr
func.return %2 : !tfr.attr
// CHECK-NEXT: %[[c:.*]] = tfr.constant [42 : i32, 41 : i32] -> !tfr.attr
// CHECK-NEXT: return %[[c]] : !tfr.attr
}
// -----
// CHECK-LABEL: build_high_dim_const_list
func.func @build_high_dim_const_list() -> !tfr.attr {
%0 = "arith.constant"() {value = 42 : i32} : () -> i32
%1 = "arith.constant"() {value = 41 : i32} : () -> i32
%2 = "tfr.build_list"(%0, %1) : (i32, i32) -> !tfr.attr
%3 = "tfr.build_list"(%0, %1) : (i32, i32) -> !tfr.attr
%4 = "tfr.build_list"(%2, %3) : (!tfr.attr, !tfr.attr) -> !tfr.attr
func.return %4 : !tfr.attr
// CHECK-NEXT: %[[c:.*]] = tfr.constant {{\[}}[42 : i32, 41 : i32], [42 : i32, 41 : i32]] -> !tfr.attr
// CHECK-NEXT: return %[[c]] : !tfr.attr
}
// -----
// CHECK-LABEL: get_length
func.func @get_length(%arg0: !tfr.tensor<A>, %arg1: !tfr.tensor<B>) -> index {
%0 = "tfr.build_list"(%arg0, %arg1) : (!tfr.tensor<A>, !tfr.tensor<B>) -> !tfr.tensor_list
%1 = "tfr.get_length"(%0) : (!tfr.tensor_list) -> index
func.return %1 : index
// CHECK-NEXT: %[[c:.*]] = arith.constant 2 : index
// CHECK-NEXT: return %[[c]] : index
}
tensorflow/compiler/mlir/tfr/tests/control_flow.mlir
+71
View File
@@ -0,0 +1,71 @@
// 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.
// ==============================================================================
// RUN: tfr-opt %s -tfr-decompose -verify-diagnostics -split-input-file | FileCheck %s
tfr.func @tf__my_pack(%values: !tfr.tensor_list,
%n: i32 {tfr.name="N"},
%axis: i32 {tfr.name="axis"}) -> !tfr.tensor {
%index = arith.constant 0 : index
%cst = arith.constant 1 : i32
%eq = arith.cmpi eq, %n, %cst : i32
%v1 = tfr.get_element %values[%index] : (!tfr.tensor_list, index) -> !tfr.tensor
%temp = tfr.call @tf__expand_dims(%v1, %axis) : (!tfr.tensor, i32) -> !tfr.tensor
%res = scf.if %eq -> !tfr.tensor {
scf.yield %temp : !tfr.tensor
} else {
%step = arith.index_cast %cst : i32 to index
%end = arith.index_cast %n : i32 to index
%reduce = scf.for %i = %step to %end step %step iter_args(%reduce_iter=%temp) -> !tfr.tensor {
%v = tfr.get_element %values[%i] : (!tfr.tensor_list, index) -> !tfr.tensor
%temp1 = tfr.call @tf__expand_dims(%v, %axis) : (!tfr.tensor, i32) -> !tfr.tensor
%reduce_next = tfr.call @tf__risc_concat(%reduce_iter, %temp1, %axis) : (!tfr.tensor, !tfr.tensor, i32) -> !tfr.tensor
scf.yield %reduce_next : !tfr.tensor
}
scf.yield %reduce : !tfr.tensor
}
tfr.return %res : !tfr.tensor
}
// CHECK-LABEL: pack_one
func.func @pack_one(%arg0: tensor<2x3xf32>) -> tensor<1x2x3xf32> {
%0 = "tf.MyPack"(%arg0) {N=1:i32, axis=0:i32} : (tensor<2x3xf32>) -> tensor<1x2x3xf32>
func.return %0 : tensor<1x2x3xf32>
// CHECK-NEXT: %[[AXIS:.*]] = arith.constant 0 : i32
// CHECK-NEXT: %[[CAST:.*]] = "tfr.cast"(%arg0) : (tensor<2x3xf32>) -> !tfr.tensor
// CHECK-NEXT: %[[ED:.*]] = tfr.call @tf__expand_dims(%[[CAST]], %[[AXIS]]) : (!tfr.tensor, i32) -> !tfr.tensor
// CHECK-NEXT: %[[BACK:.*]] = "tfr.cast"(%[[ED]]) : (!tfr.tensor) -> tensor<1x2x3xf32>
// CHECK-NEXT: return %[[BACK]] : tensor<1x2x3xf32>
}
// CHECK-LABEL: pack_multiple
func.func @pack_multiple(%arg0: tensor<2x3xf32>,
%arg1: tensor<2x3xf32>,
%arg2: tensor<2x3xf32>) -> tensor<3x2x3xf32> {
%0 = "tf.MyPack"(%arg0, %arg1, %arg2) {N=3:i32, axis=0:i32} : (tensor<2x3xf32>, tensor<2x3xf32>, tensor<2x3xf32>) -> tensor<3x2x3xf32>
func.return %0 : tensor<3x2x3xf32>
// CHECK-NEXT: %[[AXIS:.*]] = arith.constant 0 : i32
// CHECK-NEXT: %[[CAST0:.*]] = "tfr.cast"(%arg0) : (tensor<2x3xf32>) -> !tfr.tensor
// CHECK-NEXT: %[[CAST1:.*]] = "tfr.cast"(%arg1) : (tensor<2x3xf32>) -> !tfr.tensor
// CHECK-NEXT: %[[CAST2:.*]] = "tfr.cast"(%arg2) : (tensor<2x3xf32>) -> !tfr.tensor
// CHECK-NEXT: %[[EX0:.*]] = tfr.call @tf__expand_dims(%[[CAST0]], %[[AXIS]]) : (!tfr.tensor, i32) -> !tfr.tensor
// CHECK-NEXT: %[[EX1:.*]] = tfr.call @tf__expand_dims(%[[CAST1]], %[[AXIS]]) : (!tfr.tensor, i32) -> !tfr.tensor
// CHECK-NEXT: %[[CONCAT1:.*]] = tfr.call @tf__risc_concat(%[[EX0]], %[[EX1]], %c0_i32) : (!tfr.tensor, !tfr.tensor, i32) -> !tfr.tensor
// CHECK-NEXT: %[[EX2:.*]] = tfr.call @tf__expand_dims(%[[CAST2]], %[[AXIS]]) : (!tfr.tensor, i32) -> !tfr.tensor
// CHECK-NEXT: %[[CONCAT2:.*]] = tfr.call @tf__risc_concat(%[[CONCAT1]], %[[EX2]], %[[AXIS]]) : (!tfr.tensor, !tfr.tensor, i32) -> !tfr.tensor
// CHECK-NEXT: %[[BACK:.*]] = "tfr.cast"(%[[CONCAT2]]) : (!tfr.tensor) -> tensor<3x2x3xf32>
// CHECK-NEXT: return %[[BACK]] : tensor<3x2x3xf32>
}
tensorflow/compiler/mlir/tfr/tests/decompose.mlir
+303
View File
@@ -0,0 +1,303 @@
// 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.
// ==============================================================================
// RUN: tfr-opt %s -tfr-decompose -verify-diagnostics | FileCheck %s
// Definitions for ops that are being used in the tests.
// ex) tf.MyOp refers to tfr.func @tf__my_op
// CHECK-LABEL: @tf__fake_no_op
tfr.func @tf__fake_no_op(%arg0: !tfr.tensor) -> !tfr.tensor {
tfr.return %arg0 : !tfr.tensor
// CHECK-NEXT: tfr.return %arg0 : !tfr.tensor
}
// CHECK-LABEL: @tf__intermediate
tfr.func @tf__intermediate(%arg0: !tfr.tensor) -> !tfr.tensor {
%0 = tfr.call @tf__risc(%arg0) : (!tfr.tensor) -> !tfr.tensor
tfr.return %0 : !tfr.tensor
// CHECK-NEXT: %[[id:.*]] = tfr.call @tf__risc(%arg0) : (!tfr.tensor) -> !tfr.tensor
// CHECK-NEXT: tfr.return %[[id]] : !tfr.tensor
}
// CHECK-LABEL: @tf__fused_n
tfr.func @tf__fused_n(
%arg0: !tfr.tensor,
%arg1: !tfr.tensor_list,
%arg2: index {tfr.name="A",tfr.default=1:index})
-> !tfr.tensor_list {
%0 = tfr.call @tf__intermediate(%arg0) : (!tfr.tensor) -> !tfr.tensor
%1 = tfr.get_element %arg1[%arg2] : (!tfr.tensor_list, index) -> !tfr.tensor
%2 = tfr.call @tf__intermediate(%1) : (!tfr.tensor) -> !tfr.tensor
%3 = "tfr.build_list"(%0, %2) : (!tfr.tensor, !tfr.tensor) -> !tfr.tensor_list
tfr.return %3 : !tfr.tensor_list
// CHECK-NEXT: %[[id1:.*]] = tfr.call @tf__intermediate(%arg0) : (!tfr.tensor) -> !tfr.tensor
// CHECK-NEXT: %[[ge:.*]] = tfr.get_element %arg1[%arg2] : (!tfr.tensor_list, index) -> !tfr.tensor
// CHECK-NEXT: %[[id2:.*]] = tfr.call @tf__intermediate(%[[ge]]) : (!tfr.tensor) -> !tfr.tensor
// CHECK-NEXT: %[[bl:.*]] = "tfr.build_list"(%[[id1]], %[[id2]]) : (!tfr.tensor, !tfr.tensor) -> !tfr.tensor_list
// CHECK-NEXT: tfr.return %[[bl]] : !tfr.tensor_list
}
// CHECK-LABEL: @tf__my_max_pool
tfr.func @tf__my_max_pool(%input_: !tfr.tensor, %stride_w: i64{tfr.name="stride_w"}, %stride_h: i64{tfr.name="stride_h"}) -> (!tfr.tensor) {
%cst_1 = arith.constant 1 : i64
%stride = "tfr.build_list"(%cst_1, %stride_w, %stride_h, %cst_1) : (i64, i64, i64, i64) -> !tfr.attr
%filter = tfr.constant [1, 2, 2, 1] -> !tfr.attr
%padding = tfr.constant "VALID" -> !tfr.attr
%explicit_paddings = tfr.constant [] -> !tfr.attr
%data_format = tfr.constant "NHWC" -> !tfr.attr
%MaxPool = tfr.call @tf__max_pool(%input_, %stride, %filter, %padding, %explicit_paddings, %data_format) : (!tfr.tensor, !tfr.attr, !tfr.attr, !tfr.attr, !tfr.attr, !tfr.attr) -> (!tfr.tensor)
tfr.return %MaxPool : !tfr.tensor
// CHECK: tf__max_pool
}
// CHECK-LABEL: @tf__cast_float
tfr.func @tf__cast_float(%input_: !tfr.tensor, %out_type: !tfr.attr{tfr.name="out_type"}) -> (!tfr.tensor) {
%false = arith.constant false
%cast = tfr.call @tf__cast(%input_, %out_type, %false) : (!tfr.tensor, !tfr.attr, i1) -> (!tfr.tensor)
tfr.return %cast : !tfr.tensor
}
// end op definitions
// CHECK-LABEL: decompose_tf_no_op
func.func @decompose_tf_no_op(%arg0: tensor<1x2x3x4x!tf_type.string>) -> tensor<1x2x3x4x!tf_type.string> {
%0 = "tf.FakeNoOp"(%arg0) : (tensor<1x2x3x4x!tf_type.string>) -> tensor<1x2x3x4x!tf_type.string>
func.return %0 : tensor<1x2x3x4x!tf_type.string>
// CHECK-NEXT: return %arg0
}
// CHECK-LABEL: decompose_tf_intermediate
func.func @decompose_tf_intermediate(%arg0: tensor<1x2x3x4x!tf_type.string>) -> tensor<1x2x3x4x!tf_type.string> {
%0 = "tf.Intermediate"(%arg0) : (tensor<1x2x3x4x!tf_type.string>) -> tensor<1x2x3x4x!tf_type.string>
func.return %0 : tensor<1x2x3x4x!tf_type.string>
// CHECK-NEXT: %[[casted:.*]] = "tfr.cast"(%arg0) : (tensor<1x2x3x4x!tf_type.string>) -> !tfr.tensor
// CHECK-NEXT: %[[id:.*]] = tfr.call @tf__risc(%[[casted]]) : (!tfr.tensor) -> !tfr.tensor
// CHECK-NEXT: %[[back:.*]] = "tfr.cast"(%[[id]]) : (!tfr.tensor) -> tensor<1x2x3x4x!tf_type.string>
// CHECK-NEXT: return %[[back]]
}
// CHECK-LABEL: decompose_fused_n_default
func.func @decompose_fused_n_default(%arg0: tensor<1x2x3x4x!tf_type.string>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
%0:2 = "tf.FusedN"(%arg0, %arg1, %arg2) : (tensor<1x2x3x4x!tf_type.string>, tensor<f32>, tensor<f32>) -> (tensor<1x2x3x4x!tf_type.string>, tensor<f32>)
func.return %0#1 : tensor<f32>
// CHECK-NEXT: %[[in0:.*]] = "tfr.cast"(%arg0) : (tensor<1x2x3x4x!tf_type.string>) -> !tfr.tensor
// CHECK-NEXT: %[[in2:.*]] = "tfr.cast"(%arg2) : (tensor<f32>) -> !tfr.tensor
// CHECK-NEXT: %[[id0:.*]] = tfr.call @tf__risc(%[[in0]]) : (!tfr.tensor) -> !tfr.tensor
// CHECK-NEXT: %[[id2:.*]] = tfr.call @tf__risc(%[[in2]]) : (!tfr.tensor) -> !tfr.tensor
// CHECK-NEXT: %[[back:.*]] = "tfr.cast"(%[[id2]]) : (!tfr.tensor) -> tensor<f32>
// CHECK-NEXT: return %[[back]] : tensor<f32>
}
// CHECK-LABEL: decompose_fused_n
func.func @decompose_fused_n(%arg0: tensor<1x2x3x4x!tf_type.string>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
%0:2 = "tf.FusedN"(%arg0, %arg1, %arg2) {A=0:index} : (tensor<1x2x3x4x!tf_type.string>, tensor<f32>, tensor<f32>) -> (tensor<1x2x3x4x!tf_type.string>, tensor<f32>)
func.return %0#1 : tensor<f32>
// CHECK-NEXT: %[[in0:.*]] = "tfr.cast"(%arg0) : (tensor<1x2x3x4x!tf_type.string>) -> !tfr.tensor
// CHECK-NEXT: %[[in1:.*]] = "tfr.cast"(%arg1) : (tensor<f32>) -> !tfr.tensor
// CHECK-NEXT: %[[id0:.*]] = tfr.call @tf__risc(%[[in0]]) : (!tfr.tensor) -> !tfr.tensor
// CHECK-NEXT: %[[id1:.*]] = tfr.call @tf__risc(%[[in1]]) : (!tfr.tensor) -> !tfr.tensor
// CHECK-NEXT: %[[back:.*]] = "tfr.cast"(%[[id1]]) : (!tfr.tensor) -> tensor<f32>
// CHECK-NEXT: return %[[back]] : tensor<f32>
}
// CHECK-LABEL: attribute_propagate_direct
func.func @attribute_propagate_direct(%arg0: tensor<1x2x3x4x!tf_type.string>) -> tensor<1x2x3x4x!tf_type.string> {
%0 = "tf.Intermediate"(%arg0) {_tpu_replicate, device="hello"} : (tensor<1x2x3x4x!tf_type.string>) -> tensor<1x2x3x4x!tf_type.string>
func.return %0 : tensor<1x2x3x4x!tf_type.string>
// CHECK-NEXT: %[[casted:.*]] = "tfr.cast"(%arg0) : (tensor<1x2x3x4x!tf_type.string>) -> !tfr.tensor
// CHECK-NEXT: %[[id:.*]] = tfr.call @tf__risc(%[[casted]]) {_tpu_replicate, device = "hello"}
// CHECK-NEXT: %[[back:.*]] = "tfr.cast"(%[[id]]) : (!tfr.tensor) -> tensor<1x2x3x4x!tf_type.string>
// CHECK-NEXT: return %[[back]]
}
// CHECK-LABEL: attribute_propagate
func.func @attribute_propagate(%arg0: tensor<1x2x3x4x!tf_type.string>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
%0:2 = "tf.FusedN"(%arg0, %arg1, %arg2) {A=0:index, _tpu_replicate, device="hello"} : (tensor<1x2x3x4x!tf_type.string>, tensor<f32>, tensor<f32>) -> (tensor<1x2x3x4x!tf_type.string>, tensor<f32>)
func.return %0#1 : tensor<f32>
// CHECK-NEXT: %[[in0:.*]] = "tfr.cast"(%arg0) : (tensor<1x2x3x4x!tf_type.string>) -> !tfr.tensor
// CHECK-NEXT: %[[in1:.*]] = "tfr.cast"(%arg1) : (tensor<f32>) -> !tfr.tensor
// CHECK-NEXT: %[[id0:.*]] = tfr.call @tf__risc(%[[in0]]) {_tpu_replicate, device = "hello"}
// CHECK-NEXT: %[[id1:.*]] = tfr.call @tf__risc(%[[in1]]) {_tpu_replicate, device = "hello"}
// CHECK-NEXT: %[[back:.*]] = "tfr.cast"(%[[id1]]) : (!tfr.tensor) -> tensor<f32>
// CHECK-NEXT: return %[[back]] : tensor<f32>
}
// CHECK: attribute_cast
func.func @attribute_cast(%arg0: tensor<1x4x4x1xf32>) -> tensor<1x2x2x1xf32> {
%0 = "tfr.cast"(%arg0) : (tensor<1x4x4x1xf32>) -> !tfr.tensor
%stride_i32 = arith.constant 2 : i32
%1 = tfr.call @tf__my_max_pool(%0, %stride_i32, %stride_i32) : (!tfr.tensor, i32, i32) -> !tfr.tensor
%2 = "tfr.cast"(%1) : (!tfr.tensor) -> tensor<1x2x2x1xf32>
func.return %2 : tensor<1x2x2x1xf32>
// CHECK: tf__max_pool
}
// CHECK-LABEL: no_tf_canonicalization
func.func @no_tf_canonicalization(%arg0: tensor<8xi1>, %arg1: tensor<8x3xf32>, %arg2: tensor<8x3xf32>) -> tensor<8x3xf32> {
%0 = "tf.Select"(%arg0, %arg1, %arg2) : (tensor<8xi1>, tensor<8x3xf32>, tensor<8x3xf32>) -> tensor<8x3xf32>
func.return %0: tensor<8x3xf32>
// CHECK: "tf.Select"
}
// CHECK-LABEL: denied_attribute
func.func @denied_attribute(%arg0: tensor<1x2x3x4x!tf_type.string>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
// expected-error@+1 {{Denied unregistered attribute was found: denied_attr}}
%0:2 = "tf.FusedN"(%arg0, %arg1, %arg2) {A=0:index, denied_attr} : (tensor<1x2x3x4x!tf_type.string>, tensor<f32>, tensor<f32>) -> (tensor<1x2x3x4x!tf_type.string>, tensor<f32>)
func.return %0#1 : tensor<f32>
// CHECK-NEXT: "tf.FusedN"(%arg0, %arg1, %arg2) {A = 0 : index, denied_attr}
}
// CHECK-LABEL: quantized_tensor
func.func @quantized_tensor(%arg0: tensor<1x10x!quant.uniform<i8:f32, 0.0038396485615521669:-128>>) -> tensor<1x10x!quant.uniform<i8:f32, 3.906250e-03:-128>> {
%0 = "tf.Intermediate"(%arg0) : (tensor<1x10x!quant.uniform<i8:f32, 0.0038396485615521669:-128>>) -> tensor<1x10x!quant.uniform<i8:f32, 3.906250e-03:-128>>
func.return %0 : tensor<1x10x!quant.uniform<i8:f32, 3.906250e-03:-128>>
// CHECK: "tfr.cast"(%[[arg0:.*]]) : (tensor<1x10x!quant.uniform<i8:f32, 0.0038396485615521669:-128>>) -> !tfr.tensor
// CHECK: "tfr.cast"(%[[result:.*]]) : (!tfr.tensor) -> tensor<1x10x!quant.uniform<i8:f32, 3.906250e-03:-128>>
}
// CHECK-LABEL: decompose_quant_act_range
func.func @decompose_quant_act_range() -> !tfr.tensor_list {
%scale = arith.constant 0.1 : f32
%zp = arith.constant 42 : i64
%none_attr = tfr.constant "NONE" -> !tfr.attr
%relu_attr = tfr.constant "RELU" -> !tfr.attr
%relu6_attr = tfr.constant "RELU6" -> !tfr.attr
%reluN1_1_attr = tfr.constant "RELU_N1_TO_1" -> !tfr.attr
%none:2 = "tfr.quant_act_range"(%none_attr, %scale, %zp) : (!tfr.attr, f32, i64) -> (!tfr.tensor, !tfr.tensor)
%relu:2 = "tfr.quant_act_range"(%relu_attr, %scale, %zp) : (!tfr.attr, f32, i64) -> (!tfr.tensor, !tfr.tensor)
%relu6:2 = "tfr.quant_act_range"(%relu6_attr, %scale, %zp) : (!tfr.attr, f32, i64) -> (!tfr.tensor, !tfr.tensor)
%reluN1_1:2 = "tfr.quant_act_range"(%reluN1_1_attr, %scale, %zp) : (!tfr.attr, f32, i64) -> (!tfr.tensor, !tfr.tensor)
%result = "tfr.build_list"(
%none#0, %none#1, %relu#0, %relu#1,
%relu6#0, %relu6#1, %reluN1_1#0, %reluN1_1#1) : (
!tfr.tensor, !tfr.tensor, !tfr.tensor, !tfr.tensor,
!tfr.tensor, !tfr.tensor, !tfr.tensor, !tfr.tensor) -> !tfr.tensor_list
func.return %result : !tfr.tensor_list
// CHECK-DAG: %[[N_128:.*]] = arith.constant -128 : i32
// CHECK-DAG: %[[N32:.*]] = arith.constant 32 : i32
// CHECK-DAG: %[[N42:.*]] = arith.constant 42 : i32
// CHECK-DAG: %[[N52:.*]] = arith.constant 52 : i32
// CHECK-DAG: %[[N102:.*]] = arith.constant 102 : i32
// CHECK-DAG: %[[N127:.*]] = arith.constant 127 : i32
// CHECK-NEXT: %[[none_min:.*]] = "tfr.constant_tensor"(%[[N_128]])
// CHECK-NEXT: %[[none_max:.*]] = "tfr.constant_tensor"(%[[N127]])
// CHECK-NEXT: %[[relu_min:.*]] = "tfr.constant_tensor"(%[[N42]])
// CHECK-NEXT: %[[relu_max:.*]] = "tfr.constant_tensor"(%[[N127]])
// CHECK-NEXT: %[[relu6_min:.*]] = "tfr.constant_tensor"(%[[N42]])
// CHECK-NEXT: %[[relu6_max:.*]] = "tfr.constant_tensor"(%[[N102]])
// CHECK-NEXT: %[[reluN1_1_min:.*]] = "tfr.constant_tensor"(%[[N32]])
// CHECK-NEXT: %[[reluN1_1_max:.*]] = "tfr.constant_tensor"(%[[N52]])
// CHECK-NEXT: %[[result:.*]] = "tfr.build_list"(%[[none_min]], %[[none_max]], %[[relu_min]], %[[relu_max]],
// CHECK-SAME: %[[relu6_min]], %[[relu6_max]], %[[reluN1_1_min]], %[[reluN1_1_max]]
// CHECK-NEXT: return %[[result]]
}
// CHECK-LABEL: decompose_quant_act_range_invalid
func.func @decompose_quant_act_range_invalid() -> (!tfr.tensor, !tfr.tensor) {
%scale = arith.constant 0.1 : f32
%zp = arith.constant 42 : i64
%elu_attr = tfr.constant "ELU" -> !tfr.attr
%min, %max = "tfr.quant_act_range"(%elu_attr, %scale, %zp) : (!tfr.attr, f32, i64) -> (!tfr.tensor, !tfr.tensor)
func.return %min, %max : !tfr.tensor, !tfr.tensor
// CHECK: %[[elu_attr:.*]] = tfr.constant "ELU" -> !tfr.attr
// CHECK: %[[min:.*]], %[[max:.*]] = tfr.quant_act_range(%[[elu_attr]]
// CHECK: return %[[min]], %[[max]]
}
// CHECK-LABEL: decompose_quant_scale_factor
func.func @decompose_quant_scale_factor() -> (!tfr.tensor, !tfr.tensor) {
%output_scale = arith.constant 0.1 : f32
%input_scale = arith.constant 0.25 : f32
%filter_scale = arith.constant 0.4 : f32
%input_scale_tensor = "tfr.constant_tensor"(%input_scale) : (f32) -> !tfr.tensor
%filter_scale_tensor = "tfr.constant_tensor"(%filter_scale) : (f32) -> !tfr.tensor
%list = "tfr.build_list"(%input_scale_tensor, %filter_scale_tensor) : (!tfr.tensor, !tfr.tensor) -> !tfr.tensor_list
%out = "tfr.quant_scale_factor"(%output_scale, %list) : (f32, !tfr.tensor_list) -> !tfr.tensor
%perchannel_scale = arith.constant dense<[0.4, 4.0]> : tensor<2xf32>
%perchannel_scale_tensor = "tfr.cast"(%perchannel_scale) : (tensor<2xf32>) -> !tfr.tensor
%list2 = "tfr.build_list"(%input_scale_tensor, %perchannel_scale_tensor) : (!tfr.tensor, !tfr.tensor) -> !tfr.tensor_list
%perchannel = "tfr.quant_scale_factor"(%output_scale, %list2) : (f32, !tfr.tensor_list) -> !tfr.tensor
func.return %out, %perchannel : !tfr.tensor, !tfr.tensor
// CHECK-DAG: %[[scale_factors:.*]] = "tf.Const"() <{value = dense<[1.000000e+00, 1.000000e+01]> : tensor<2xf32>}> : () -> tensor<2xf32>
// CHECK-DAG: %[[scale_factor:.*]] = "tf.Const"() <{value = dense<1.000000e+00> : tensor<f32>}> : () -> tensor<f32>
// CHECK: %[[cast:.*]] = "tfr.cast"(%[[scale_factor]]) : (tensor<f32>) -> !tfr.tensor
// CHECK: %[[cast_perchannel:.*]] = "tfr.cast"(%[[scale_factors]]) : (tensor<2xf32>) -> !tfr.tensor
// CHECK: return %[[cast]], %[[cast_perchannel]] : !tfr.tensor, !tfr.tensor
}
// CHECK-LABEL: decompose_quant_scale_factor_invalid
func.func @decompose_quant_scale_factor_invalid() -> !tfr.tensor {
%output_scale = arith.constant 0.1 : f32
%input_scale = arith.constant 0.25 : f32
%filter_scale = arith.constant 0.4 : f32
%input_scale_tensor = "tfr.constant_tensor"(%input_scale) : (f32) -> !tfr.tensor
%filter_scale_tensor = "tfr.constant_tensor"(%filter_scale) : (f32) -> !tfr.tensor
%list = "tfr.build_list"(%input_scale_tensor, %filter_scale_tensor, %input_scale_tensor) : (!tfr.tensor, !tfr.tensor, !tfr.tensor) -> !tfr.tensor_list
%out = "tfr.quant_scale_factor"(%output_scale, %list) : (f32, !tfr.tensor_list) -> !tfr.tensor
func.return %out : !tfr.tensor
// CHECK-DAG: %[[cst_0:.*]] = arith.constant 1.000000e-01 : f32
// CHECK-DAG: %[[cst_1:.*]] = "tf.Const"() <{value = dense<2.500000e-01> : tensor<f32>}> : () -> tensor<f32>
// CHECK-DAG: %[[cst_2:.*]] = "tf.Const"() <{value = dense<4.000000e-01> : tensor<f32>}> : () -> tensor<f32>
// CHECK: %[[tfrcast0:.*]] = "tfr.cast"(%[[cst_1]]) : (tensor<f32>) -> !tfr.tensor
// CHECK: %[[tfrcast1:.*]] = "tfr.cast"(%[[cst_2]]) : (tensor<f32>) -> !tfr.tensor
// CHECK: %[[list:.*]] = "tfr.build_list"(%[[tfrcast0]], %[[tfrcast1]], %[[tfrcast0]]) : (!tfr.tensor, !tfr.tensor, !tfr.tensor) -> !tfr.tensor_list
// CHECK: %[[qsf:.*]] = tfr.quant_scale_factor(%[[cst_0]], %[[list]]) : (f32, !tfr.tensor_list) -> !tfr.tensor
// CHECK: return %[[qsf]] : !tfr.tensor
}
// CHECK-LABEL: decompose_quant_rescale
func.func @decompose_quant_rescale(%arg0: tensor<2xi32>) -> !tfr.tensor {
%zp = arith.constant 67 : i64
%cst = "tf.Const"() {value = dense<1.0> : tensor<f32>} : () -> tensor<f32>
%scale_factor = "tfr.cast"(%cst) : (tensor<f32>) -> !tfr.tensor
%input = "tfr.cast"(%arg0) : (tensor<2xi32>) -> !tfr.tensor
%rescaled = "tfr.quant_rescale"(%input, %scale_factor, %zp) : (!tfr.tensor, !tfr.tensor, i64) -> !tfr.tensor
func.return %rescaled : !tfr.tensor
// CHECK-DAG: %[[f32:.*]] = tfr.constant f32 -> !tfr.attr
// CHECK-DAG: %[[i32:.*]] = tfr.constant i32 -> !tfr.attr
// CHECK-DAG: %[[scale_cst:.*]] = "tf.Const"() <{value = dense<1.000000e+00> : tensor<f32>}> : () -> tensor<f32>
// CHECK-DAG: %false = arith.constant false
// CHECK-DAG: %[[zp_cst:.*]] = "tf.Const"() <{value = dense<67> : tensor<i64>}> : () -> tensor<i64>
// CHECK: %[[zp:.*]] = "tfr.cast"(%[[zp_cst]]) : (tensor<i64>) -> !tfr.tensor
// CHECK: %[[scale:.*]] = "tfr.cast"(%[[scale_cst]]) : (tensor<f32>) -> !tfr.tensor
// CHECK: %[[input:.*]] = "tfr.cast"(%arg0) : (tensor<2xi32>) -> !tfr.tensor
// CHECK: %[[cast:.*]] = tfr.call @tf__cast(%[[input]], %[[f32]], %false) : (!tfr.tensor, !tfr.attr, i1) -> !tfr.tensor
// CHECK: %[[rescaled:.*]] = tfr.call @tf__mul(%[[cast]], %[[scale]]) : (!tfr.tensor, !tfr.tensor) -> !tfr.tensor
// CHECK: %[[rounded:.*]] = tfr.call @tf__round(%[[rescaled]]) : (!tfr.tensor) -> !tfr.tensor
// CHECK: %[[zp_cast:.*]] = tfr.call @tf__cast(%[[zp]], %[[f32]], %false) : (!tfr.tensor, !tfr.attr, i1) -> !tfr.tensor
// CHECK: %[[recentered:.*]] = tfr.call @tf__add(%[[rounded]], %[[zp_cast]]) : (!tfr.tensor, !tfr.tensor) -> !tfr.tensor
// CHECK: %[[cast_i32:.*]] = tfr.call @tf__cast(%[[recentered]], %[[i32]], %false) : (!tfr.tensor, !tfr.attr, i1) -> !tfr.tensor
// CHECK: return %[[cast_i32]] : !tfr.tensor
}
// CHECK-LABEL: decompose_output_type
func.func @decompose_output_type(%arg0: tensor<2xf32>) -> tensor<2xi32> {
%0 = "tf.CastFloat"(%arg0) : (tensor<2xf32>) -> tensor<2xi32>
func.return %0: tensor<2xi32>
// CHECK: %[[i32:.*]] = tfr.constant i32 -> !tfr.attr
// CHECK: tfr.call @tf__cast(%[[casted_arg:.*]], %[[i32]], %false) : (!tfr.tensor, !tfr.attr, i1) -> !tfr.tensor
}
tensorflow/compiler/mlir/tfr/tests/end2end.mlir
+249
View File
@@ -0,0 +1,249 @@
// 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.
// ==============================================================================
// RUN: tfr-opt %s -tfr-decompose -tfr-raise-to-tf -canonicalize -verify-diagnostics -split-input-file | FileCheck %s
//=================> User models, from GraphDef <====================
// CHECK-LABEL: my_identity
func.func @my_identity(%arg0: tensor<2x3xf32>) -> tensor<2x3xf32> {
%0 = "tf.MyIdentity"(%arg0) : (tensor<2x3xf32>) -> tensor<2x3xf32>
func.return %0 : tensor<2x3xf32>
// CHECK-NEXT: return %arg0 : tensor<2x3xf32>
}
// CHECK-LABEL: my_rsqrt
func.func @my_rsqrt(%arg0: tensor<2x3xf32>) -> tensor<3x2x3xf32> {
%0 = "tf.MyRsqrt"(%arg0) : (tensor<2x3xf32>) -> tensor<3x2x3xf32>
func.return %0 : tensor<3x2x3xf32>
// CHECK-NEXT: %[[RE:.*]] = "tf.RiscReciprocal"(%arg0) : (tensor<2x3xf32>) -> tensor<*xf32>
// CHECK-NEXT: %[[SQRT:.*]] = "tf.RiscSqrt"(%[[RE]]) : (tensor<*xf32>) -> tensor<*xf32>
// CHECK-NEXT: %[[ES:.*]] = "tf.EnsureShape"(%[[SQRT]]) <{shape = #tf_type.shape<3x2x3>}> : (tensor<*xf32>) -> tensor<3x2x3xf32>
// CHECK-NEXT: return %[[ES]] : tensor<3x2x3xf32>
}
// CHECK-LABEL: my_leaky_relu
func.func @my_leaky_relu(%arg0: tensor<2x3xf32>) -> tensor<3x2x3xf32> {
%0 = "tf.MyLeakyRelu"(%arg0) {alpha=3.0 : f32} : (tensor<2x3xf32>) -> tensor<3x2x3xf32>
func.return %0 : tensor<3x2x3xf32>
// CHECK-NEXT: %[[ALPHA:.*]] = "tf.Const"() <{value = dense<3.000000e+00> : tensor<f32>}> : () -> tensor<f32>
// CHECK-NEXT: %[[SHAPE:.*]] = "tf.RiscShape"(%arg0) {T = i32} : (tensor<2x3xf32>) -> tensor<*xi32>
// CHECK-NEXT: %[[ALPHA1:.*]] = "tf.RiscBroadcast"(%[[ALPHA]], %[[SHAPE]]) : (tensor<f32>, tensor<*xi32>) -> tensor<*xf32>
// CHECK-NEXT: %[[MAX:.*]] = "tf.RiscMaximum"(%arg0, %[[ALPHA1]]) : (tensor<2x3xf32>, tensor<*xf32>) -> tensor<*xf32>
// CHECK-NEXT: %[[ES:.*]] = "tf.EnsureShape"(%[[MAX]]) <{shape = #tf_type.shape<3x2x3>}> : (tensor<*xf32>) -> tensor<3x2x3xf32>
// CHECK-NEXT: return %[[ES]] : tensor<3x2x3xf32>
}
// CHECK-LABEL: my_leaky_relu_with_default
func.func @my_leaky_relu_with_default(%arg0: tensor<2x3xf32>) -> tensor<3x2x3xf32> {
%0 = "tf.MyLeakyRelu"(%arg0) : (tensor<2x3xf32>) -> tensor<3x2x3xf32>
func.return %0 : tensor<3x2x3xf32>
// CHECK-NEXT: %[[ALPHA:.*]] = "tf.Const"() <{value = dense<2.000000e-01> : tensor<f32>}> : () -> tensor<f32>
// CHECK-NEXT: %[[SHAPE:.*]] = "tf.RiscShape"(%arg0) {T = i32} : (tensor<2x3xf32>) -> tensor<*xi32>
// CHECK-NEXT: %[[ALPHA1:.*]] = "tf.RiscBroadcast"(%[[ALPHA]], %[[SHAPE]]) : (tensor<f32>, tensor<*xi32>) -> tensor<*xf32>
// CHECK-NEXT: %[[MAX:.*]] = "tf.RiscMaximum"(%arg0, %[[ALPHA1]]) : (tensor<2x3xf32>, tensor<*xf32>) -> tensor<*xf32>
// CHECK-NEXT: %[[ES:.*]] = "tf.EnsureShape"(%[[MAX]]) <{shape = #tf_type.shape<3x2x3>}> : (tensor<*xf32>) -> tensor<3x2x3xf32>
// CHECK-NEXT: return %[[ES]] : tensor<3x2x3xf32>
}
// CHECK-LABEL: my_cast
func.func @my_cast(%arg0: tensor<2x3xf32>) -> tensor<2x3xi32> {
%0 = "tf.MyCast"(%arg0) {Tout=i32} : (tensor<2x3xf32>) -> tensor<2x3xi32>
func.return %0 : tensor<2x3xi32>
// CHECK-NEXT: %[[CAST:.*]] = "tf.RiscCast"(%arg0) {Tout = i32} : (tensor<2x3xf32>) -> tensor<*xi32>
// CHECK-NEXT: %[[ES:.*]] = "tf.EnsureShape"(%[[CAST]]) <{shape = #tf_type.shape<2x3>}> : (tensor<*xi32>) -> tensor<2x3xi32>
// CHECK-NEXT: return %[[ES]] : tensor<2x3xi32>
}
// CHECK-LABEL: my_pack_single_input
func.func @my_pack_single_input(%arg0: tensor<2x3xf32>) -> tensor<3x2x3xf32> {
%0 = "tf.MyPack"(%arg0) {N=1:i32, axis=0:i32} : (tensor<2x3xf32>) -> tensor<3x2x3xf32>
func.return %0 : tensor<3x2x3xf32>
// CHECK-NEXT: %[[AXIS:.*]] = "tf.Const"() <{value = dense<0> : tensor<i32>}> : () -> tensor<i32>
// CHECK-NEXT: %[[ED:.*]] = "tf.ExpandDims"(%arg0, %[[AXIS]]) : (tensor<2x3xf32>, tensor<i32>) -> tensor<*xf32>
// CHECK-NEXT: %[[ES:.*]] = "tf.EnsureShape"(%[[ED]]) <{shape = #tf_type.shape<3x2x3>}> : (tensor<*xf32>) -> tensor<3x2x3xf32>
// CHECK-NEXT: return %[[ES]] : tensor<3x2x3xf32>
}
// CHECK-LABEL: my_pack_multiple_inputs
func.func @my_pack_multiple_inputs(%arg0: tensor<2x3xf32>, %arg1: tensor<2x3xf32>, %arg2: tensor<2x3xf32>) -> tensor<3x2x3xf32> {
%0 = "tf.MyPack"(%arg0, %arg1, %arg2) {N=3:i32, axis=0:i32} : (tensor<2x3xf32>, tensor<2x3xf32>, tensor<2x3xf32>) -> tensor<3x2x3xf32>
func.return %0 : tensor<3x2x3xf32>
// CHECK-NEXT: %[[AXIS:.*]] = "tf.Const"() <{value = dense<0> : tensor<i32>}> : () -> tensor<i32>
// CHECK-NEXT: %[[ED0:.*]] = "tf.ExpandDims"(%arg0, %[[AXIS]]) : (tensor<2x3xf32>, tensor<i32>) -> tensor<*xf32>
// CHECK-NEXT: %[[ED1:.*]] = "tf.ExpandDims"(%arg1, %[[AXIS]]) : (tensor<2x3xf32>, tensor<i32>) -> tensor<*xf32>
// CHECK-NEXT: %[[CC0:.*]] = "tf.RiscConcat"(%[[ED0]], %[[ED1]]) {axis = 0 : i32} : (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32>
// CHECK-NEXT: %[[ED2:.*]] = "tf.ExpandDims"(%arg2, %[[AXIS]]) : (tensor<2x3xf32>, tensor<i32>) -> tensor<*xf32>
// CHECK-NEXT: %[[CC1:.*]] = "tf.RiscConcat"(%[[CC0]], %[[ED2]]) {axis = 0 : i32} : (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32>
// CHECK-NEXT: %[[ES:.*]] = "tf.EnsureShape"(%[[CC1]]) <{shape = #tf_type.shape<3x2x3>}> : (tensor<*xf32>) -> tensor<3x2x3xf32>
// CHECK-NEXT: return %[[ES]] : tensor<3x2x3xf32>
}
// CHECK-LABEL: my_add_n_single_input
func.func @my_add_n_single_input(%arg0: tensor<2x3xf32>) -> tensor<2x3xf32> {
%0 = "tf.MyAddN"(%arg0) {N=1:i32} : (tensor<2x3xf32>) -> tensor<2x3xf32>
func.return %0 : tensor<2x3xf32>
// CHECK-NEXT: return %arg0 : tensor<2x3xf32>
}
// CHECK-LABEL: my_add_n_multiple_inputs
func.func @my_add_n_multiple_inputs(%arg0: tensor<2x3xf32>, %arg1: tensor<2x3xf32>, %arg2: tensor<2x3xf32>) -> tensor<2x3xf32> {
%0 = "tf.MyAddN"(%arg0, %arg1, %arg2) {N=3:i32} : (tensor<2x3xf32>, tensor<2x3xf32>, tensor<2x3xf32>) -> tensor<2x3xf32>
func.return %0 : tensor<2x3xf32>
// CHECK-NEXT: %[[ADD0:.*]] = "tf.RiscAdd"(%arg0, %arg1) : (tensor<2x3xf32>, tensor<2x3xf32>) -> tensor<*xf32>
// CHECK-NEXT: %[[ADD1:.*]] = "tf.RiscAdd"(%[[ADD0]], %arg2) : (tensor<*xf32>, tensor<2x3xf32>) -> tensor<*xf32>
// CHECK-NEXT: %[[ES:.*]] = "tf.EnsureShape"(%[[ADD1]]) <{shape = #tf_type.shape<2x3>}> : (tensor<*xf32>) -> tensor<2x3xf32>
// CHECK-NEXT: return %[[ES]] : tensor<2x3xf32>
}
// CHECK-LABEL: my_map_and_batch_dataset
func.func @my_map_and_batch_dataset(%input: tensor<*x!tf_type.variant>,
%other1: tensor<*xf32>,
%other2: tensor<*xi32>) -> tensor<*x!tf_type.variant> {
%0 = "tf.MyMapAndBatchDataset"(%input, %other1, %other2)
{batch_size=1000 : i64, num_parallel_calls = 8 : i64, drop_remainder = 0 : i1,
func = @"__some_func", output_types = [f32], output_shapes = [#tf_type.shape<>], preserve_cardinality = true}
: (tensor<*x!tf_type.variant>, tensor<*xf32>, tensor<*xi32>) -> tensor<*x!tf_type.variant>
func.return %0 : tensor<*x!tf_type.variant>
// CHECK-DAG: %[[BATCH:.*]] = "tf.Const"() <{value = dense<1000> : tensor<i64>}> : () -> tensor<i64>
// CHECK-DAG: %[[PARAL:.*]] = "tf.Const"() <{value = dense<8> : tensor<i64>}> : () -> tensor<i64>
// CHECK-DAG: %[[KEEP:.*]] = "tf.Const"() <{value = dense<false> : tensor<i1>}> : () -> tensor<i1>
// CHECK: %[[CAST:.*]] = "tf.Cast"(%arg2) <{Truncate = false}> : (tensor<*xi32>) -> tensor<*xf32>
// CHECK: %[[RET:.*]] = "tf.MapAndBatchDatasetV0"(%arg0, %[[BATCH]], %[[PARAL]], %[[KEEP]], %arg1, %[[CAST]])
// CHECK-SAME: {f = @__some_func, output_shapes = [#tf_type.shape<>], output_types = [f32], preserve_cardinality = true} : (tensor<*x!tf_type.variant>, tensor<i64>, tensor<i64>, tensor<i1>, tensor<*xf32>, tensor<*xf32>) -> tensor<*x!tf_type.variant>
// CHECK: return %[[RET]] : tensor<*x!tf_type.variant>
}
//=================> decomposition functions, translated from tf.compose api <====================
tfr.func @tf__my_identity(%value: !tfr.tensor) -> !tfr.tensor {
tfr.return %value : !tfr.tensor
}
tfr.func @tf__my_cast(%value: !tfr.tensor, %tout: !tfr.attr{tfr.name="Tout"}) -> !tfr.tensor {
%0 = tfr.call @tf__risc_cast(%value, %tout) : (!tfr.tensor, !tfr.attr) -> !tfr.tensor
tfr.return %0 : !tfr.tensor
}
tfr.func @tf__my_rsqrt(%value: !tfr.tensor) -> !tfr.tensor {
%1 = tfr.call @tf__risc_reciprocal(%value) : (!tfr.tensor) -> !tfr.tensor
%2 = tfr.call @tf__risc_sqrt(%1) : (!tfr.tensor) -> !tfr.tensor
tfr.return %2 : !tfr.tensor
}
tfr.func @tf__my_leaky_relu(%value: !tfr.tensor, %alpha: f32 {tfr.name="alpha", tfr.default=0.2:f32}) -> !tfr.tensor {
%1 = tfr.call @tf__risc_shape(%value) : (!tfr.tensor) -> !tfr.tensor
%2 = "tfr.constant_tensor"(%alpha) : (f32) -> tensor<f32>
%t = "tfr.cast"(%2) : (tensor<f32>) -> !tfr.tensor
%3 = tfr.call @tf__risc_broadcast(%t, %1) : (!tfr.tensor, !tfr.tensor) -> !tfr.tensor
%4 = tfr.call @tf__risc_maximum(%value, %3) : (!tfr.tensor, !tfr.tensor) -> !tfr.tensor
tfr.return %4 : !tfr.tensor
}
// TODO(fengliuai): use shape dialect to manipulate the shape then this can be decomposed further.
tfr.func @tf__my_expand_dims(%value: !tfr.tensor, %axis: i32 {tfr.name="axis"}) -> !tfr.tensor {
%axis_cst = "tfr.constant_tensor"(%axis) : (i32) -> tensor<i32>
%dim = "tfr.cast"(%axis_cst) : (tensor<i32>) -> !tfr.tensor
%0 = tfr.call @tf__expand_dims(%value, %dim) : (!tfr.tensor, !tfr.tensor) -> !tfr.tensor
tfr.return %0 : !tfr.tensor
}
tfr.func @tf__my_pack(%values: !tfr.tensor_list,
%n: i32 {tfr.name="N"},
%axis: i32 {tfr.name="axis"}) -> !tfr.tensor {
%index = arith.constant 0 : index
%cst = arith.constant 1 : i32
%eq = arith.cmpi eq, %n, %cst : i32
%v1 = tfr.get_element %values[%index] : (!tfr.tensor_list, index) -> !tfr.tensor
%temp = tfr.call @tf__my_expand_dims(%v1, %axis) : (!tfr.tensor, i32) -> !tfr.tensor
%res = scf.if %eq -> !tfr.tensor {
scf.yield %temp : !tfr.tensor
} else {
%step = arith.index_cast %cst : i32 to index
%end = arith.index_cast %n : i32 to index
%reduce = scf.for %i = %step to %end step %step iter_args(%reduce_iter=%temp) -> !tfr.tensor {
%v = tfr.get_element %values[%i] : (!tfr.tensor_list, index) -> !tfr.tensor
%temp1 = tfr.call @tf__my_expand_dims(%v, %axis) : (!tfr.tensor, i32) -> !tfr.tensor
%reduce_next = tfr.call @tf__risc_concat(%reduce_iter, %temp1, %axis) : (!tfr.tensor, !tfr.tensor, i32) -> !tfr.tensor
scf.yield %reduce_next : !tfr.tensor
}
scf.yield %reduce : !tfr.tensor
}
tfr.return %res : !tfr.tensor
}
tfr.func @tf__my_add_n(%values: !tfr.tensor_list,
%n: i32 {tfr.name="N"}) -> !tfr.tensor {
%index = arith.constant 0 : index
%cst = arith.constant 1 : i32
%eq = arith.cmpi eq, %n, %cst : i32
%v1 = tfr.get_element %values[%index] : (!tfr.tensor_list, index) -> !tfr.tensor
%res = scf.if %eq -> !tfr.tensor {
scf.yield %v1 : !tfr.tensor
} else {
%step = arith.index_cast %cst : i32 to index
%end = arith.index_cast %n : i32 to index
%reduce = scf.for %i = %step to %end step %step iter_args(%reduce_iter=%v1) -> !tfr.tensor {
%v = tfr.get_element %values[%i] : (!tfr.tensor_list, index) -> !tfr.tensor
%reduce_next = tfr.call @tf__risc_add(%reduce_iter, %v) : (!tfr.tensor, !tfr.tensor) -> !tfr.tensor
scf.yield %reduce_next : !tfr.tensor
}
scf.yield %reduce : !tfr.tensor
}
tfr.return %res : !tfr.tensor
}
tfr.func @tf__my_map_and_batch_dataset(
%input_dataset: !tfr.tensor,
%other_arguments: !tfr.tensor_list,
%batch_size: i64 {tfr.name="batch_size"},
%num_parallel_calls: i64 {tfr.name="num_parallel_calls"},
%drop_remainder: i1 {tfr.name="drop_remainder"},
%f: !tfr.attr {tfr.name="func"},
%output_types: !tfr.attr {tfr.name="output_types"},
%output_shapes: !tfr.attr {tfr.name="output_shapes"},
%preserve_cardinality: i1 {tfr.name="preserve_cardinality", tfr.default=false}) -> !tfr.tensor {
%batch = "tfr.constant_tensor"(%batch_size) : (i64) -> tensor<i64>
%batch1 = "tfr.cast"(%batch) : (tensor<i64>) -> !tfr.tensor
%calls = "tfr.constant_tensor"(%num_parallel_calls) : (i64) -> tensor<i64>
%calls1 = "tfr.cast"(%calls) : (tensor<i64>) -> !tfr.tensor
%drop = "tfr.constant_tensor"(%drop_remainder) : (i1) -> tensor<i1>
%drop1 = "tfr.cast"(%drop) : (tensor<i1>) -> !tfr.tensor
%ret = tfr.call @tf__map_and_batch_dataset_v0(%input_dataset, %batch1, %calls1, %drop1, %other_arguments, %f, %output_types, %output_shapes, %preserve_cardinality)
: (!tfr.tensor, !tfr.tensor, !tfr.tensor, !tfr.tensor, !tfr.tensor_list, !tfr.attr, !tfr.attr, !tfr.attr, i1) -> !tfr.tensor
tfr.return %ret : !tfr.tensor
}
//=================> signatures of the primitive ops with kernels, modeled as external TFR function <==
tfr.func @tf__risc_cast_(!tfr.tensor, !tfr.attr{tfr.name="Tout"}) -> !tfr.tensor<Tout> attributes{Tout}
tfr.func @tf__risc_add_(!tfr.tensor<T>, !tfr.tensor<T>) -> !tfr.tensor<T> attributes{T}
tfr.func @tf__risc_concat_(!tfr.tensor<T>, !tfr.tensor<T>, i32{tfr.name="axis"}) -> !tfr.tensor<T> attributes{T}
tfr.func @tf__risc_broadcast_(!tfr.tensor<T>, !tfr.tensor<Tidx>) -> !tfr.tensor<T> attributes{T, Tidx}
tfr.func @tf__risc_reciprocal_(!tfr.tensor<T>) -> !tfr.tensor<T> attributes{T}
tfr.func @tf__risc_sqrt_(!tfr.tensor<T>) -> !tfr.tensor<T> attributes{T}
tfr.func @tf__risc_shape_(!tfr.tensor, !tfr.attr{tfr.name="T", tfr.default=i32}) -> !tfr.tensor<T> attributes{T}
tfr.func @tf__risc_maximum_(!tfr.tensor<T>, !tfr.tensor<T>) -> !tfr.tensor<T> attributes{T}
tfr.func @tf__expand_dims_(!tfr.tensor<T>, !tfr.tensor<Tdim>) -> !tfr.tensor<T> attributes{T, Tdim}
tfr.func @tf__map_and_batch_dataset_v0_(!tfr.tensor<T>, !tfr.tensor, !tfr.tensor, !tfr.tensor, !tfr.tensor_list<Targuments>,
!tfr.attr{tfr.name="f"}, !tfr.attr{tfr.name="output_types"}, !tfr.attr{tfr.name="output_shapes"}, i1{tfr.name="preserve_cardinality"})
-> !tfr.tensor<T> attributes{T, Targuments}
tensorflow/compiler/mlir/tfr/tests/ops.mlir
+442
View File
@@ -0,0 +1,442 @@
// 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.
// ==============================================================================
// RUN: tfr-opt %s -verify-diagnostics -split-input-file | tfr-opt | FileCheck %s
// Tests for types, ops with custom constraints, verifiers, printer or parser
// methods.
// CHECK-LABEL: tensor_type_noconstraint
func.func private @tensor_type_noconstraint() -> !tfr.tensor
// -----
// CHECK-LABEL: tensor_type
func.func private @tensor_type() -> !tfr.tensor<T>
// -----
// CHECK-LABEL: tensor_list_type_noconstraint
func.func private @tensor_list_type_noconstraint() -> !tfr.tensor_list
// -----
// CHECK-LABEL: tensor_list_type_array_like
func.func private @tensor_list_type_array_like() -> !tfr.tensor_list<[N, T]>
// -----
// CHECK-LABEL: tensor_list_type_tuple_like
func.func private @tensor_list_type_tuple_like() -> !tfr.tensor_list<input_T>
// -----
// expected-error@+1 {{unbalanced '[' character in pretty dialect name}}
func.func private @tensor_invalid_1() -> !tfr.tensor<[N, T>
// -----
// expected-error@+1 {{unbalanced}}
func.func @tensor_invalid_2() -> !tfr.tensor<[N, T]
// -----
// CHECK-LABEL: call_op
func.func @call_op(%arg0: !tfr.tensor<T>, %arg1: !tfr.tensor_list<TL>, %arg2: i32) -> !tfr.tensor<K> {
%0 = tfr.call @Foo(%arg0, %arg1, %arg2) : (!tfr.tensor<T>, !tfr.tensor_list<TL>, i32) -> !tfr.tensor<K>
func.return %0 : !tfr.tensor<K>
}
// -----
// CHECK-LABEL: call_op_arg_attr(%arg0: i32) -> !tfr.tensor<K>
func.func @call_op_arg_attr(%arg0: i32) -> !tfr.tensor<K> {
%0 = tfr.call @Bar(%arg0) : (i32) -> !tfr.tensor<K>
func.return %0 : !tfr.tensor<K>
}
// -----
func.func @call_op_invalid_1(%arg0: tensor<?xf32>) -> !tfr.tensor<K> {
// expected-error@+1 {{got 'tensor<?xf32>'}}
%0 = tfr.call @Huu(%arg0) : (tensor<?xf32>) -> !tfr.tensor<K>
func.return %0 : !tfr.tensor<K>
}
// -----
// CHECK-LABEL: get_shape
func.func @get_shape(%arg0: !tfr.tensor) -> (!shape.shape, !shape.shape) {
%0 = tfr.get_shape %arg0 -> !shape.shape
%1 = "tfr.get_shape"(%arg0) : (!tfr.tensor) -> !shape.shape
func.return %0, %1 : !shape.shape, !shape.shape
}
// -----
// CHECK-LABEL: get_real_shape
func.func @get_real_shape(%arg0: tensor<1x2xf32>) -> tensor<2xindex> {
%0 = "tfr.cast"(%arg0) : (tensor<1x2xf32>) -> !tfr.tensor
%1 = tfr.get_shape %0 -> !shape.shape
%2 = shape.to_extent_tensor %1 : !shape.shape -> tensor<2xindex>
func.return %2 : tensor<2xindex>
}
// -----
func.func @get_element_type(%arg0: !tfr.tensor) -> (!tfr.attr, !tfr.attr) {
%0 = tfr.get_element_type %arg0 -> !tfr.attr
%1 = "tfr.get_element_type"(%arg0) : (!tfr.tensor) -> !tfr.attr
func.return %0, %1 : !tfr.attr, !tfr.attr
}
// -----
// CHECK-LABEL: from_tf_tensor
func.func @from_tf_tensor(%arg0: tensor<?xf32>) -> !tfr.tensor<K> {
%0 = "tfr.cast"(%arg0) : (tensor<?xf32>) -> !tfr.tensor<K>
func.return %0 : !tfr.tensor<K>
}
// -----
// CHECK-LABEL: to_tf_tensor
func.func @to_tf_tensor(%arg0: !tfr.tensor<T>) -> tensor<?xi32> {
%0 = "tfr.cast"(%arg0) : (!tfr.tensor<T>) -> tensor<?xi32>
func.return %0 : tensor<?xi32>
}
// -----
// CHECK-LABEL: constant
func.func @constant() -> (!tfr.attr, !tfr.attr, !tfr.attr, !tfr.attr) {
%0 = tfr.constant f32 -> !tfr.attr
%1 = tfr.constant [f32, i32] -> !tfr.attr
%2 = "tfr.constant"() {value = f32} : () -> !tfr.attr
%3 = "tfr.constant"() {value = [f32, i32]} : () -> !tfr.attr
func.return %0, %1, %2, %3 : !tfr.attr, !tfr.attr, !tfr.attr, !tfr.attr
}
// -----
// CHECK-LABEL: equal
func.func @equal() -> (i1, i1, i1, i1) {
%0 = tfr.constant f32 -> !tfr.attr
%1 = tfr.constant f32 -> !tfr.attr
%2 = tfr.constant i32 -> !tfr.attr
%same_type = tfr.equal %0,%1 -> i1
%diff_type = tfr.equal %0,%2 -> i1
%3 = tfr.constant "hello" -> !tfr.attr
%4 = tfr.constant "hello" -> !tfr.attr
%5 = tfr.constant "how are you" -> !tfr.attr
%same_str = tfr.equal %3,%4 -> i1
%diff_str = tfr.equal %3,%5 -> i1
func.return %same_type, %diff_type, %same_str, %diff_str : i1, i1, i1, i1
}
// -----
// CHECK-LABEL: constant_tensor_scalar
func.func @constant_tensor_scalar(%arg0: i32) -> tensor<i32> {
%0 = "tfr.constant_tensor"(%arg0) : (i32) -> tensor<i32>
func.return %0 : tensor<i32>
}
// -----
// CHECK-LABEL: constant_tensor_vector
func.func @constant_tensor_vector(%arg0: vector<1x2xi32>) -> tensor<1x2xi32> {
%0 = "tfr.constant_tensor"(%arg0) : (vector<1x2xi32>) -> tensor<1x2xi32>
func.return %0 : tensor<1x2xi32>
}
// -----
// CHECK-LABEL: constant_tensor_array
func.func @constant_tensor_array() -> !tfr.tensor {
%0 = tfr.constant [1, -1, 3] -> !tfr.attr
%1 = "tfr.constant_tensor"(%0) : (!tfr.attr) -> !tfr.tensor
func.return %1 : !tfr.tensor
}
// -----
// CHECK-LABEL: constant_tensor_scalar
func.func @constant_tensor_scalar() -> !tfr.tensor {
%0 = "arith.constant"() {value = 42 : i32} : () -> i32
%1 = "tfr.constant_tensor"(%0) : (i32) -> !tfr.tensor
func.return %1 : !tfr.tensor
}
// -----
func.func @constant_tensor_invalid_0(%arg0: i32) -> tensor<f32> {
// expected-error@+1 {{input and output should have the same scalar types.}}
%0 = "tfr.constant_tensor"(%arg0) : (i32) -> tensor<f32>
func.return %0 : tensor<f32>
}
// -----
func.func @constant_tensor_invalid_1(%arg0: vector<1xi32>) -> tensor<?xi32> {
// expected-error@+1 {{output type should be static and ranked}}
%0 = "tfr.constant_tensor"(%arg0) : (vector<1xi32>) -> tensor<?xi32>
func.return %0 : tensor<?xi32>
}
// -----
func.func @constant_tensor_invalid_2(%arg0: vector<1xi32>) -> tensor<1xf32> {
// expected-error@+1 {{input and output should have same shape and element type}}
%0 = "tfr.constant_tensor"(%arg0) : (vector<1xi32>) -> tensor<1xf32>
func.return %0 : tensor<1xf32>
}
// -----
func.func @constant_tensor_invalid_3(%arg0: vector<1xi32>) -> tensor<1x1xi32> {
// expected-error@+1 {{input and output should have same shape and element type}}
%0 = "tfr.constant_tensor"(%arg0) : (vector<1xi32>) -> tensor<1x1xi32>
func.return %0 : tensor<1x1xi32>
}
// -----
func.func @constant_tensor_invalid_4(%arg0: i32) -> tensor<1x1xi32> {
// expected-error@+1 {{input can not be converted to an output tensor}}
%0 = "tfr.constant_tensor"(%arg0) : (i32) -> tensor<1x1xi32>
func.return %0 : tensor<1x1xi32>
}
// -----
// CHECK-LABEL: get_element
func.func @get_element(%arg0: !tfr.tensor_list<T>) -> !tfr.tensor {
%cst = "arith.constant"() {value = 1 : index} : () -> index
%0 = tfr.get_element %arg0[%cst] : (!tfr.tensor_list<T>, index) -> !tfr.tensor
func.return %0 : !tfr.tensor
}
// -----
// CHECK-LABEL: build_list
func.func @build_list(%arg0: !tfr.tensor<A>, %arg1: !tfr.tensor<B>) -> !tfr.tensor_list {
%0 = "tfr.build_list"(%arg0, %arg1) : (!tfr.tensor<A>, !tfr.tensor<B>) -> !tfr.tensor_list
func.return %0 : !tfr.tensor_list
}
// -----
// CHECK-LABEL: quant_act_range
func.func @quant_act_range(%arg0: !tfr.attr, %arg1: f32, %arg2: i64) -> !tfr.tensor {
%0:2 = "tfr.quant_act_range"(%arg0, %arg1, %arg2) : (!tfr.attr,f32,i64) -> (!tfr.tensor,!tfr.tensor)
func.return %0#0 : !tfr.tensor
}
// -----
// CHECK-LABEL: quant_rescale
func.func @quant_rescale(%arg0: !tfr.tensor, %arg1: !tfr.tensor, %arg2: i64) -> !tfr.tensor {
%0 = "tfr.quant_rescale"(%arg0, %arg1, %arg2) : (!tfr.tensor, !tfr.tensor, i64) -> (!tfr.tensor)
func.return %0 : !tfr.tensor
}
// -----
// CHECK-LABEL: quant_raw_data
func.func @quant_raw_data(%arg0: !tfr.tensor) -> !tfr.tensor {
%0 = "tfr.quant_raw_data"(%arg0) : (!tfr.tensor) -> (!tfr.tensor)
func.return %0 : !tfr.tensor
}
// -----
// CHECK-LABEL: quant_qparam
func.func @quant_qparam(%arg0: !tfr.tensor) -> (!tfr.tensor, !tfr.tensor) {
%scale, %zp = tfr.quant_qparam(%arg0) : (!tfr.tensor) -> (!tfr.tensor, !tfr.tensor)
func.return %scale, %zp : !tfr.tensor, !tfr.tensor
}
// -----
// CHECK-LABEL: quant_scale_factor
func.func @quant_scale_factor(%arg0: f32, %arg1: !tfr.tensor_list) -> (!tfr.tensor) {
%0 = "tfr.quant_scale_factor"(%arg0, %arg1) : (f32, !tfr.tensor_list) -> (!tfr.tensor)
func.return %0 : !tfr.tensor
}
// -----
// CHECK-LABEL: build_const_list
func.func @build_const_list() -> !tfr.attr {
%0 = "arith.constant"() {value = 42 : i32} : () -> i32
%1 = "arith.constant"() {value = 41 : i32} : () -> i32
%2 = "tfr.build_list"(%0, %1) : (i32, i32) -> !tfr.attr
func.return %2 : !tfr.attr
}
// -----
// CHECK-LABEL: build_high_dim_const_list
func.func @build_high_dim_const_list() -> !tfr.attr {
%0 = "arith.constant"() {value = 42 : i32} : () -> i32
%1 = "arith.constant"() {value = 41 : i32} : () -> i32
%2 = "tfr.build_list"(%0, %1) : (i32, i32) -> !tfr.attr
%3 = "tfr.build_list"(%0, %1) : (i32, i32) -> !tfr.attr
%4 = "tfr.build_list"(%2, %3) : (!tfr.attr, !tfr.attr) -> !tfr.attr
func.return %4 : !tfr.attr
}
// -----
// CHECK-LABEL: get_length
func.func @get_length(%arg0: !tfr.tensor<A>, %arg1: !tfr.tensor<B>) -> index {
%0 = "tfr.build_list"(%arg0, %arg1) : (!tfr.tensor<A>, !tfr.tensor<B>) -> !tfr.tensor_list
%1 = "tfr.get_length"(%0) : (!tfr.tensor_list) -> index
func.return %1 : index
}
// -----
// CHECK-LABEL: tfr.func
tfr.func @External(%arg0: !tfr.tensor<A>,
%arg1: !tfr.tensor_list<C>,
%arg2: i32 {tfr.name = "A"},
%arg3: !tfr.attr {tfr.name = "T"})
-> (!tfr.tensor<A>, !tfr.tensor_list<C>)
attributes {A, C}
// -----
// CHECK-LABEL: tfr.func
tfr.func @Foo(%arg0: !tfr.tensor<A>,
%arg1: !tfr.tensor_list<C>,
%arg2: i32 {tfr.name = "A"},
%arg3: vector<1xi32> {tfr.name = "C"})
-> (!tfr.tensor<A>, !tfr.tensor_list<C>)
attributes {A, C} {
tfr.return %arg0, %arg1 : !tfr.tensor<A>, !tfr.tensor_list<C>
}
// -----
// CHECK-LABEL: tfr.func
tfr.func @Bar(%arg0: !tfr.tensor<A>,
%arg2: i32 {tfr.name = "B"},
%arg3: vector<1xi32> {tfr.name = "C"})
-> (!tfr.tensor<A>, !tfr.tensor<A>)
attributes {A} {
tfr.return %arg0, %arg0 : !tfr.tensor<A>, !tfr.tensor<A>
}
// -----
// expected-error@+1 {{Undefined attributes are used: A}}
tfr.func @Foo_undefined_attr(%arg0: !tfr.tensor<A>,
%arg1: !tfr.tensor_list<A>,
%arg2: i32 {tfr.name = "A"},
%arg3: vector<1xi32> {tfr.name = "C"}) ->
(!tfr.tensor<A>, !tfr.tensor_list<A>) {
tfr.return %arg0, %arg1 : !tfr.tensor<A>, !tfr.tensor_list<A>
}
// -----
// expected-error@+1 {{3 attribute argument doesn't have a tfr.name attribute}}
tfr.func @Foo_unnamed_attr(%arg0: !tfr.tensor<A>,
%arg1: !tfr.tensor_list<A>,
%arg2: i32 {tfr.name = "A"},
%arg3: vector<1xi32>) ->
(!tfr.tensor<A>, !tfr.tensor_list<A>) {
tfr.return %arg0, %arg1 : !tfr.tensor<A>, !tfr.tensor_list<A>
}
// -----
// expected-error@+1 {{tfr.tensor/tfr.tensor_list argument should be before non tensor arguments}}
tfr.func @Foo_invalid_arg_order(%arg0: !tfr.tensor<A>,
%arg2: i32 {tfr.name = "A"},
%arg1: !tfr.tensor_list<A>,
%arg3: vector<1xi32> {tfr.name = "C"}) ->
(!tfr.tensor<A>, !tfr.tensor_list<A>) {
tfr.return %arg0, %arg1 : !tfr.tensor<A>, !tfr.tensor_list<A>
}
// -----
tfr.func @Foo_valid_arg_order0(
%arg1: !tfr.tensor_list,
%arg0: !tfr.tensor<T>,
%arg2: i32 {tfr.name = "A"},
%arg3: vector<1xi32> {tfr.name = "C"}) ->
(!tfr.tensor, !tfr.tensor_list) attributes {T}{
tfr.return %arg0, %arg1 : !tfr.tensor<T>, !tfr.tensor_list
}
// -----
// expected-error@+1 {{tfr.tensor argument should be before tfr.tensor_list argument.}}
tfr.func @Foo_invalid_arg_order0(
%arg1: !tfr.tensor_list,
%arg0: !tfr.tensor<T>,
%arg2: i32 {tfr.name = "A"},
%arg3: vector<1xi32> {tfr.name = "C"}) ->
(!tfr.tensor, !tfr.tensor_list) {
tfr.return %arg0, %arg1 : !tfr.tensor<T>, !tfr.tensor_list
}
// -----
// expected-error@+1 {{tfr.tensor result should be before tfr.tensor_list result}}
tfr.func @Foo_invalid_result_order(%arg0: !tfr.tensor<A>,
%arg1: !tfr.tensor_list<A>,
%arg2: i32 {tfr.name = "A"},
%arg3: vector<1xi32> {tfr.name = "C"}) ->
(!tfr.tensor_list<A>, !tfr.tensor<A>) {
tfr.return %arg1, %arg0 : !tfr.tensor_list<A>, !tfr.tensor<A>
}
// -----
// expected-error@+1 {{More than one tfr.tensor_list argument isn't allowed}}
tfr.func @Foo_multiple_tensor_list_args(%arg0: !tfr.tensor<A>,
%arg1: !tfr.tensor_list<A>,
%arg2: !tfr.tensor_list<A>,
%arg3: i32 {tfr.name = "A"},
%arg4: vector<1xi32> {tfr.name = "C"}) ->
(!tfr.tensor<A>, !tfr.tensor_list<A>) {
tfr.return %arg0, %arg1 : !tfr.tensor<A>, !tfr.tensor_list<A>
}
// -----
// expected-error@+1 {{More than one tfr.tensor_list result isn't allowed}}
tfr.func @Foo_multiple_tensor_list_results(%arg0: !tfr.tensor<C>,
%arg1: !tfr.tensor_list<A>,
%arg2: i32 {tfr.name = "A"},
%arg3: vector<1xi32> {tfr.name = "C"}) ->
(!tfr.tensor_list<A>, !tfr.tensor_list<A>) {
tfr.return %arg1, %arg1 : !tfr.tensor_list<A>, !tfr.tensor_list<A>
}
// -----
// expected-error@+1 {{None tfr.tensor/tfr.tensor_list results aren't allowed as a result}}
tfr.func @Foo_return_attr(%arg0: !tfr.tensor<C>,
%arg1: !tfr.tensor_list<A>,
%arg2: i32 {tfr.name = "A"},
%arg3: vector<1xi32> {tfr.name = "C"}) -> i32 {
tfr.return %arg2 : i32
}
tensorflow/compiler/mlir/tfr/tests/raise_to_tf.mlir
+140
View File
@@ -0,0 +1,140 @@
// 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.
// ==============================================================================
// RUN: tfr-opt %s -tfr-raise-to-tf -verify-diagnostics -split-input-file | FileCheck %s
tfr.func @tf__risc_same_(!tfr.tensor<T>) -> !tfr.tensor<T> attributes {T}
tfr.func @tf__risc_concat_(!tfr.tensor_list<N, T>) -> !tfr.tensor<T> attributes {T, N}
tfr.func @tf__risc_split_(!tfr.tensor<T>, i32 {tfr.name="N"}) -> !tfr.tensor_list<N, T> attributes {T, N}
tfr.func @tf__risc_cast_(!tfr.tensor, !tfr.attr {tfr.name="K"}) -> !tfr.tensor<K> attributes {T, K}
tfr.func @tf__const_(!tfr.attr {tfr.name="value", tfr.type="tensor"},
!tfr.attr {tfr.name="K",tfr.type="dtype"}) -> !tfr.tensor<K> attributes {T, K}
tfr.func @tf__positive_(!tfr.tensor<T>) -> !tfr.tensor<i1_> attributes {T, i1_}
tfr.func @tf__invalid_type_op_(!tfr.tensor<T>) -> !tfr.tensor<i8_> attributes {T, i8_}
// CHECK-LABEL: decompose_tf_same
func.func @decompose_tf_same(%arg0: tensor<1x2x3x4x!tf_type.string>) -> tensor<1x2x3x4x!tf_type.string> {
%0 = "tfr.cast"(%arg0) : (tensor<1x2x3x4x!tf_type.string>) -> !tfr.tensor
%1 = tfr.call @tf__risc_same(%0) : (!tfr.tensor) -> !tfr.tensor
%2 = "tfr.cast"(%1) : (!tfr.tensor) -> tensor<1x2x3x4x!tf_type.string>
func.return %2 : tensor<1x2x3x4x!tf_type.string>
// CHECK: %[[id:.*]] = "tf.RiscSame"(%arg0) : (tensor<1x2x3x4x!tf_type.string>) -> tensor<*x!tf_type.string>
// CHECK: %[[es:.*]] = "tf.EnsureShape"(%[[id]]) <{shape = #tf_type.shape<1x2x3x4>}> : (tensor<*x!tf_type.string>) -> tensor<1x2x3x4x!tf_type.string>
// CHECK: return %[[es]] : tensor<1x2x3x4x!tf_type.string>
}
// CHECK-LABEL: decompose_tf_consecutive
func.func @decompose_tf_consecutive(%arg0: tensor<1x2x3x4x!tf_type.string>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
%0 = "tfr.cast"(%arg0) : (tensor<1x2x3x4x!tf_type.string>) -> !tfr.tensor
%1 = "tfr.cast"(%arg2) : (tensor<f32>) -> !tfr.tensor
%2 = tfr.call @tf__risc_same(%0) : (!tfr.tensor) -> !tfr.tensor
%3 = tfr.call @tf__risc_same(%1) : (!tfr.tensor) -> !tfr.tensor
%4 = "tfr.cast"(%3) : (!tfr.tensor) -> tensor<f32>
func.return %4 : tensor<f32>
// CHECK: %[[id0:.*]] = "tf.RiscSame"(%arg0) : (tensor<1x2x3x4x!tf_type.string>) -> tensor<*x!tf_type.string>
// CHECK: %[[id2:.*]] = "tf.RiscSame"(%arg2) : (tensor<f32>) -> tensor<*xf32>
// CHECK: %[[es:.*]] = "tf.EnsureShape"(%[[id2]]) <{shape = #tf_type.shape<>}> : (tensor<*xf32>) -> tensor<f32>
// CHECK: return %[[es]] : tensor<f32>
}
// CHECK-LABEL: decompose_tf_concat_n
func.func @decompose_tf_concat_n(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<3xf32> {
%0 = "tfr.cast"(%arg0) : (tensor<f32>) -> !tfr.tensor
%1 = "tfr.cast"(%arg1) : (tensor<f32>) -> !tfr.tensor
%2 = "tfr.cast"(%arg2) : (tensor<f32>) -> !tfr.tensor
%3 = "tfr.build_list"(%0, %1, %2) : (!tfr.tensor, !tfr.tensor, !tfr.tensor) -> !tfr.tensor_list
%concat = tfr.call @tf__risc_concat(%3) : (!tfr.tensor_list) -> !tfr.tensor
%4 = "tfr.cast"(%concat) : (!tfr.tensor) -> tensor<3xf32>
func.return %4 : tensor<3xf32>
// CHECK: %[[concat:.*]] = "tf.RiscConcat"(%arg0, %arg1, %arg2) : (tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<*xf32>
// CHECK: %[[es:.*]] = "tf.EnsureShape"(%[[concat]]) <{shape = #tf_type.shape<3>}> : (tensor<*xf32>) -> tensor<3xf32>
// CHECK: return %[[es]] : tensor<3xf32>
}
// CHECK-LABEL: decompose_tf_split
func.func @decompose_tf_split(%arg0: tensor<3xf32>) -> (tensor<f32>) {
%0 = "tfr.cast"(%arg0) : (tensor<3xf32>) -> !tfr.tensor
%n = arith.constant 3: i32
%split = tfr.call @tf__risc_split(%0, %n) : (!tfr.tensor, i32) -> !tfr.tensor_list
%i0 = arith.constant 0: index
%s0 = tfr.get_element %split[%i0] : (!tfr.tensor_list, index) -> !tfr.tensor
%4 = "tfr.cast"(%s0) : (!tfr.tensor) -> tensor<f32>
func.return %4 : tensor<f32>
// CHECK: %[[split:.*]]:3 = "tf.RiscSplit"(%arg0) {N = 3 : i32} : (tensor<3xf32>) -> (tensor<*xf32>, tensor<*xf32>, tensor<*xf32>)
// CHECK: %[[es:.*]] = "tf.EnsureShape"(%[[split]]#0) <{shape = #tf_type.shape<>}> : (tensor<*xf32>) -> tensor<f32>
// CHECK: return %[[es]] : tensor<f32>
}
// CHECK-LABEL: decompose_tf_cast
func.func @decompose_tf_cast(%arg0: tensor<f32>) -> tensor<i32> {
%0 = "tfr.cast"(%arg0) : (tensor<f32>) -> !tfr.tensor
%t = tfr.constant i32 -> !tfr.attr
%concat = tfr.call @tf__risc_cast(%0, %t) : (!tfr.tensor, !tfr.attr) -> !tfr.tensor
%4 = "tfr.cast"(%concat) : (!tfr.tensor) -> tensor<i32>
func.return %4 : tensor<i32>
// CHECK: %[[tfcast:.*]] = "tf.RiscCast"(%arg0) {K = i32} : (tensor<f32>) -> tensor<*xi32>
// CHECK: %[[es:.*]] = "tf.EnsureShape"(%[[tfcast]]) <{shape = #tf_type.shape<>}> : (tensor<*xi32>) -> tensor<i32>
// CHECK: return %[[es]] : tensor<i32>
}
// CHECK-LABEL: convert_to_scalar_tensor
func.func @convert_to_scalar_tensor() -> tensor<f32> {
%0 = arith.constant 3.0: f32
%t = tfr.constant f32 -> !tfr.attr
%cst = tfr.call @tf__const(%0, %t) : (f32, !tfr.attr) -> !tfr.tensor
%4 = "tfr.cast"(%cst) : (!tfr.tensor) -> tensor<f32>
func.return %4 : tensor<f32>
// CHECK: %[[cst:.*]] = "tf.Const"() <{value = dense<3.000000e+00> : tensor<f32>}> : () -> tensor<f32>
// CHECK: return %[[cst]] : tensor<f32>
}
// CHECK-LABEL: attribute_propagate
func.func @attribute_propagate(%arg0: tensor<f32>) -> tensor<i32> {
%0 = "tfr.cast"(%arg0) : (tensor<f32>) -> !tfr.tensor
%t = tfr.constant i32 -> !tfr.attr
%concat = tfr.call @tf__risc_cast(%0, %t) {device = "hello", _tpu_replicate} : (!tfr.tensor, !tfr.attr) -> !tfr.tensor
%4 = "tfr.cast"(%concat) : (!tfr.tensor) -> tensor<i32>
func.return %4 : tensor<i32>
// CHECK: %[[tfcast:.*]] = "tf.RiscCast"(%arg0) {K = i32, _tpu_replicate, device = "hello"} : (tensor<f32>) -> tensor<*xi32>
// CHECK: %[[es:.*]] = "tf.EnsureShape"(%[[tfcast]]) <{shape = #tf_type.shape<>}> : (tensor<*xi32>) -> tensor<i32>
// CHECK: return %[[es]] : tensor<i32>
}
// CHECK-LABEL: fixed_element_attribute
func.func @fixed_element_attribute(%arg0: tensor<2xf32>) -> tensor<2xi1> {
%0 = "tfr.cast"(%arg0) : (tensor<2xf32>) -> !tfr.tensor
%1 = tfr.call @tf__positive(%0) : (!tfr.tensor) -> !tfr.tensor
%2 = "tfr.cast"(%1) : (!tfr.tensor) -> tensor<2xi1>
func.return %2 : tensor<2xi1>
// CHECK: %[[positive:.*]] = "tf.Positive"(%arg0) : (tensor<2xf32>) -> tensor<*xi1>
// CHECK: %[[res:.*]] = "tf.EnsureShape"(%[[positive]]) <{shape = #tf_type.shape<2>}> : (tensor<*xi1>) -> tensor<2xi1>
// CHECK: return %[[res]] : tensor<2xi1>
}
// CHECK-LABEL: fixed_element_attribute_invalid
func.func @fixed_element_attribute_invalid(%arg0: tensor<2xf32>) -> tensor<2xi8> {
%0 = "tfr.cast"(%arg0) : (tensor<2xf32>) -> !tfr.tensor
// expected-error@+1 {{type i8_ can't be resolved for the signature of the op}}
%1 = tfr.call @tf__invalid_type_op(%0) : (!tfr.tensor) -> !tfr.tensor
%2 = "tfr.cast"(%1) : (!tfr.tensor) -> tensor<2xi8>
func.return %2 : tensor<2xi8>
// CHECK: tfr.call @tf__invalid_type_op
}
tensorflow/compiler/mlir/tfr/tests/rewrite_quantized_io.mlir
+58
View File
@@ -0,0 +1,58 @@
// 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.
// ==============================================================================
// RUN: tfr-opt %s -tfr-decompose -tfr-rewrite-quantized-io -verify-diagnostics | FileCheck %s
// CHECK-LABEL: @tf__my_requantize
tfr.func @tf__my_requantize(%input: !tfr.tensor) -> !tfr.tensor {
%raw_data = tfr.quant_raw_data(%input) : (!tfr.tensor) -> !tfr.tensor
%scale, %zp = tfr.quant_qparam(%input) : (!tfr.tensor) -> (!tfr.tensor, !tfr.tensor)
%result = tfr.call @tf__requantize(%raw_data, %scale, %zp) : (!tfr.tensor, !tfr.tensor, !tfr.tensor) -> !tfr.tensor
tfr.return %result : !tfr.tensor
}
// CHECK-LABEL: @tf__intermediate
tfr.func @tf__intermediate(%arg0: !tfr.tensor) -> !tfr.tensor {
%0 = tfr.call @tf__risc(%arg0) : (!tfr.tensor) -> !tfr.tensor
tfr.return %0 : !tfr.tensor
}
// CHECK-LABEL: remove_quantized_io
func.func @remove_quantized_io(
%arg0: tensor<1x10x!quant.uniform<i8:f32, 0.1:-128>>,
%arg1: tensor<1x5xf32>) -> (tensor<1x10x!quant.uniform<i8:f32, 0.2:42>>, tensor<1x5xf32>) {
%0 = "tf.MyRequantize"(%arg0) : (tensor<1x10x!quant.uniform<i8:f32, 0.1:-128>>) -> tensor<1x10x!quant.uniform<i8:f32, 0.2:42>>
%1 = "tf.Intermediate"(%arg1) : (tensor<1x5xf32>) -> tensor<1x5xf32>
func.return %0, %1 : tensor<1x10x!quant.uniform<i8:f32, 0.2:42>>, tensor<1x5xf32>
// CHECK-DAG: %[[scale:.*]] = "tf.Const"() <{value = dense<1.000000e-01> : tensor<f32>}> : () -> tensor<f32>
// CHECK-DAG: %[[zp:.*]] = "tf.Const"() <{value = dense<-128> : tensor<i32>}> : () -> tensor<i32>
// CHECK: %[[quant:.*]] = "tfr.cast"(%arg0) : (tensor<1x10xi8>) -> !tfr.tensor
// CHECK: %[[scale_cast:.*]] = "tfr.cast"(%[[scale]])
// CHECK: %[[zp_cast:.*]] = "tfr.cast"(%[[zp]])
// CHECK: %[[requant:.*]] = tfr.call @tf__requantize(%[[quant]], %[[scale_cast]], %[[zp_cast]])
// CHECK: %[[result:.*]] = "tfr.cast"(%[[requant]])
// CHECK-NOT: quant.uniform
// CHECK: return %[[result]], %[[float_resunt:.*]] : tensor<1x10xi8>, tensor<1x5xf32>
}
// CHECK-LABEL: quant_input_multiple_users
// expected-error@+1 {{The argument with type tensor<1x10x!quant.uniform<i8:f32, 1.000000e-01>> should have one user}}
func.func @quant_input_multiple_users(%arg0: tensor<1x10x!quant.uniform<i8:f32, 0.1>>) -> (!tfr.tensor, !tfr.tensor) {
%0 = "tfr.cast"(%arg0) : (tensor<1x10x!quant.uniform<i8:f32, 0.1>>) -> !tfr.tensor
%1 = "tfr.cast"(%arg0) : (tensor<1x10x!quant.uniform<i8:f32, 0.1>>) -> !tfr.tensor
func.return %0, %1 : !tfr.tensor, !tfr.tensor
}
tensorflow/compiler/mlir/tfr/tfr_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 verify(arg0: str) -> bool: ...
tensorflow/compiler/mlir/tfr/utils/utils.cc
+198
View File
@@ -0,0 +1,198 @@
/* 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/tfr/utils/utils.h"
#include <string>
#include "llvm/ADT/StringExtras.h"
#include "llvm/ADT/StringRef.h"
#include "llvm/ADT/StringSet.h"
#include "mlir/IR/Block.h" // from @llvm-project
#include "mlir/IR/Operation.h" // from @llvm-project
#include "mlir/IR/OperationSupport.h" // from @llvm-project
#include "mlir/IR/Region.h" // from @llvm-project
#include "mlir/Support/LLVM.h" // from @llvm-project
#include "mlir/Support/LogicalResult.h" // from @llvm-project
#include "tensorflow/compiler/mlir/tfr/ir/tfr_ops.h"
namespace mlir {
namespace TFR {
namespace {
// TODO(b/174692018): Use the official allowlist of the unregistered attrs.
const llvm::StringSet<>& GetAllowedAttributes() {
static auto* const ops = new llvm::StringSet<>({"device", "_tpu_replicate"});
return *ops;
}
// Some TFL optional attributes may not appear in their corresponding TF op
// attributes.
const llvm::StringSet<>& GetOptionalAttributes() {
static auto* const ops =
new llvm::StringSet<>({"asymmetric_quantize_inputs"});
return *ops;
}
void CollectAllowedAttrs(CallOp src, NamedAttrList* attrs) {
for (auto& attr : src->getAttrs()) {
if (GetAllowedAttributes().contains(attr.getName().strref())) {
attrs->append(attr);
}
}
}
// Adds `attrs` to all the operations between `begin` and `end` in the same
// block. Does not include `end`.
void AddAttributesInSameBlock(Block::iterator begin, Block::iterator end,
const NamedAttrList& attrs) {
for (Block::iterator it = begin; it != end; ++it) {
for (auto& attr : attrs) {
it->setAttr(attr.getName(), attr.getValue());
}
}
}
// Adds `attrs` to all the operations between `begin` and `end`. Does not
// include `end`. The operations might be across multiple blocks.
void AddAttributes(Block::iterator begin, Block::iterator end,
const NamedAttrList& attrs) {
if (begin->getBlock() == end->getBlock()) {
AddAttributesInSameBlock(begin, end, attrs);
} else {
Region::iterator begin_block = Region::iterator(begin->getBlock());
Region::iterator end_block = Region::iterator(end->getBlock());
AddAttributesInSameBlock(begin, begin_block->end(), attrs);
for (Region::iterator it = ++begin_block; it != end_block; ++it) {
AddAttributesInSameBlock(it->begin(), it->end(), attrs);
}
}
}
} // namespace
std::string GetComposeFuncName(StringRef tf_op_name) {
std::string compose_func_name;
for (int i = 0; i < tf_op_name.size(); ++i) {
if (tf_op_name[i] == '_') {
// The field name must not contain "_"s. "_Arg" and "_RetVal" are special
// op names and we can return empty string to skip the decomposition.
return {};
}
if (tf_op_name[i] == '.') {
compose_func_name.push_back('_');
} else if (llvm::isUpper(tf_op_name[i])) {
compose_func_name.push_back('_');
compose_func_name.push_back(llvm::toLower(tf_op_name[i]));
} else {
compose_func_name.push_back(tf_op_name[i]);
}
}
return compose_func_name;
}
std::string GetTFOpName(StringRef compose_func_name) {
std::string tf_op_name;
bool after_underscore = false;
for (int i = 0; i < compose_func_name.size(); ++i) {
if (llvm::isUpper(compose_func_name[i])) {
// The field name must not contain uppercase letters.
return {};
}
if (after_underscore) {
if (llvm::isLower(compose_func_name[i])) {
tf_op_name.push_back(llvm::toUpper(compose_func_name[i]));
after_underscore = false;
} else {
// The character after a "_" must be a lowercase letter.
return {};
}
} else if (compose_func_name[i] == '_') { // first time visit '_'
if (i + 1 < compose_func_name.size() && compose_func_name[i + 1] == '_') {
tf_op_name.push_back('.');
i++;
}
after_underscore = true;
} else {
tf_op_name.push_back(compose_func_name[i]);
}
}
if (after_underscore) {
// Trailing "_".
return {};
}
return tf_op_name;
}
LogicalResult ValidateAttrs(Operation* src, const StringSet<>& registered) {
for (auto& attr : src->getAttrs()) {
StringRef attr_name = attr.getName().strref();
if (!registered.contains(attr_name) &&
!(GetAllowedAttributes().contains(attr_name) ||
GetOptionalAttributes().contains(attr_name))) {
src->emitError("Denied unregistered attribute was found: " + attr_name);
return failure();
}
}
return success();
}
LogicalResult CopyAllowedUnregisteredAttrs(Operation* src, CallOp dst,
const StringSet<>& registered) {
for (auto& attr : src->getAttrs()) {
StringRef attr_name = attr.getName().strref();
// Skip the registered or optional attribute.
if (registered.contains(attr_name) ||
GetOptionalAttributes().contains(attr_name))
continue;
// Unregistered attribute.
if (GetAllowedAttributes().contains(attr_name)) {
dst->setAttr(attr.getName(), attr.getValue());
} else {
src->emitError("Denied unregistered attribute was found: " + attr_name);
return failure();
}
}
return success();
}
LogicalResult CopyNonSymbolRefAttrs(CallOp src, Operation* dst) {
NamedAttrList attrs;
CollectAllowedAttrs(src, &attrs);
for (auto& attr : attrs) {
dst->setAttr(attr.getName(), attr.getValue());
}
return success();
}
void PropagateAttrsToOperations(CallOp src, Block::iterator begin,
Block::iterator end) {
// Find all the attributes in the call op. These attributes are not in the
// op definition, so needs to be propagated to all the target ops.
NamedAttrList attrs;
CollectAllowedAttrs(src, &attrs);
// Add all the attributes to the operations in the range.
if (!attrs.empty()) {
AddAttributes(begin, end, attrs);
}
}
} // namespace TFR
} // namespace mlir
tensorflow/compiler/mlir/tfr/utils/utils.h
+66
View File
@@ -0,0 +1,66 @@
/* 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_TFR_UTILS_UTILS_H_
#define TENSORFLOW_COMPILER_MLIR_TFR_UTILS_UTILS_H_
#include <string>
#include "mlir/IR/Block.h" // from @llvm-project
#include "mlir/IR/Operation.h" // from @llvm-project
#include "mlir/IR/OperationSupport.h" // from @llvm-project
#include "mlir/Support/LLVM.h" // from @llvm-project
#include "mlir/Support/LogicalResult.h" // from @llvm-project
#include "tensorflow/compiler/mlir/tfr/ir/tfr_ops.h"
namespace mlir {
namespace TFR {
// This is a hardcoded rule for mapping a TF op name to the corresponding
// TFR function name. Examples:
// tf.Pack => tf__pack
// tf.ConcatV2 => tf__concat_v2
// TODO(fengliuai): move to an util file.
std::string GetComposeFuncName(StringRef tf_op_name);
// This is a hardcoded rule for mapping a TFR function op name to the
// corresponding TF opname. Examples:
// tf__pack -> tf.Pack
// tf__concat_v2 => tf.ConcatV2
std::string GetTFOpName(StringRef compose_func_name);
// Validate the attributes of 'src' is either contained in the registered
// attribute sets or in the allowed list.
LogicalResult ValidateAttrs(Operation* src, const StringSet<>& registered);
// Copies all the allowed attributes in 'src' to 'dst'. The copy failed if the
// 'dst' has the attribute. Return a failure if there are any attributes are not
// allowed and also unregistered.
LogicalResult CopyAllowedUnregisteredAttrs(Operation* src, CallOp dst,
const StringSet<>& registered);
// Copies all the allowed attributes in 'src' to 'dst'. FlatSymbolRefAttr is
// excluded.
LogicalResult CopyNonSymbolRefAttrs(CallOp src, Operation* dst);
// Propagates all the attributes in 'src' to the operations between 'begin' and
// 'end'. Operation 'end' is excluded.
void PropagateAttrsToOperations(CallOp src, Block::iterator begin,
Block::iterator end);
} // namespace TFR
} // namespace mlir
#endif // TENSORFLOW_COMPILER_MLIR_TFR_UTILS_UTILS_H_