wehub/tensorflow--tensorflow
1
0
Fork 0
Code Issues Pull Requests Actions 6 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/core/ir/BUILD
+258
View File
@@ -0,0 +1,258 @@
load("@llvm-project//mlir:tblgen.bzl", "gentbl_cc_library", "td_library")
load("@rules_cc//cc:cc_library.bzl", "cc_library")
load("//tensorflow:tensorflow.bzl", "tf_cc_test")
load("//tensorflow:tensorflow.default.bzl", "get_compatible_with_portable")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
default_compatible_with = get_compatible_with_portable(),
default_visibility = [
"//tensorflow/compiler/mlir/tensorflow:__subpackages__",
"//tensorflow/core:__subpackages__",
"//tensorflow/tools/tfg_graph_transforms:__subpackages__",
],
licenses = ["notice"], # Apache 2.0
)
td_library(
name = "InterfacesTdFiles",
srcs = ["interfaces.td"],
deps = [
"@llvm-project//mlir:OpBaseTdFiles",
],
)
gentbl_cc_library(
name = "InterfacesIncGen",
tbl_outs = {
"interfaces.h.inc": ["-gen-op-interface-decls"],
"interfaces.cc.inc": ["-gen-op-interface-defs"],
},
tblgen = "@llvm-project//mlir:mlir-tblgen",
td_file = "interfaces.td",
deps = [
":InterfacesTdFiles",
],
)
# ODS (https://mlir.llvm.org/docs/OpDefinitions/) generation for op and dialect files to include.
td_library(
name = "DialectTdFiles",
srcs = [
"dialect.td",
"ops.td",
],
deps = [
"@llvm-project//mlir:CallInterfacesTdFiles",
"@llvm-project//mlir:ControlFlowInterfacesTdFiles",
"@llvm-project//mlir:FunctionInterfacesTdFiles",
"@llvm-project//mlir:InferTypeOpInterfaceTdFiles",
"@llvm-project//mlir:OpBaseTdFiles",
"@llvm-project//mlir:SideEffectInterfacesTdFiles",
],
)
gentbl_cc_library(
name = "DialectIncGen",
tbl_outs = {
"ops.h.inc": [
"-gen-op-decls",
"-dialect",
"tfg",
],
"ops.cc.inc": [
"-gen-op-defs",
"-dialect",
"tfg",
],
"dialect.h.inc": [
"-gen-dialect-decls",
"-dialect",
"tfg",
],
"dialect.cc.inc": [
"-gen-dialect-defs",
"-dialect",
"tfg",
],
"attributes.h.inc": [
"-gen-attrdef-decls",
"-attrdefs-dialect",
"tfg",
],
"attributes.cc.inc": [
"-gen-attrdef-defs",
"-attrdefs-dialect",
"tfg",
],
},
tblgen = "@llvm-project//mlir:mlir-tblgen",
td_file = "ops.td",
deps = [
":DialectTdFiles",
":InterfacesTdFiles",
"//tensorflow/core/ir/types:DialectTdFiles",
],
)
cc_library(
name = "Dialect",
srcs = [
"interfaces.cc",
"ops.cc",
"tf_op_names.cc",
"tf_op_names.inc",
"tf_op_wrapper.cc",
"utility.cc",
],
hdrs = [
"dialect.h",
"interfaces.h",
"ops.h",
"tf_op_wrapper.h",
"utility.h",
],
deps = [
":DialectIncGen",
":InterfacesIncGen",
"//tensorflow/core/ir/types:Dialect",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:ControlFlowInterfaces",
"@llvm-project//mlir:FunctionInterfaces",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:InferTypeOpInterface",
"@llvm-project//mlir:SideEffectInterfaces",
"@llvm-project//mlir:Support",
],
)
cc_library(
name = "tf_op_registry",
srcs = ["tf_op_registry.cc"],
hdrs = ["tf_op_registry.h"],
deps = [
":Dialect",
"//tensorflow/core:framework",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Support",
],
)
cc_library(
name = "shape_inference_utils",
srcs = ["utils/shape_inference_utils.cc"],
hdrs = ["utils/shape_inference_utils.h"],
deps = [
":Dialect",
"//tensorflow/core:framework",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core/ir/importexport:convert_tensor",
"//tensorflow/core/ir/importexport:convert_types",
"//tensorflow/core/ir/importexport:graphdef_export",
"//tensorflow/core/ir/types:Dialect",
"//tensorflow/core/platform:logging",
"//tensorflow/core/platform:status",
"@com_google_absl//absl/log",
"@com_google_absl//absl/status",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:DerivedAttributeOpInterface",
"@llvm-project//mlir:FunctionInterfaces",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:InferTypeOpInterface",
"@llvm-project//mlir:Support",
],
)
tf_cc_test(
name = "utility_test",
size = "small",
srcs = ["utility_test.cc"],
deps = [
":Dialect",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Parser",
"@llvm-project//mlir:Support",
],
)
tf_cc_test(
name = "tf_op_wrapper_test",
size = "small",
srcs = ["tf_op_wrapper_test.cc"],
deps = [
":Dialect",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Parser",
"@llvm-project//mlir:Support",
],
)
tf_cc_test(
name = "shape_inference_utils_test",
size = "small",
srcs = ["utils/shape_inference_utils_test.cc"],
deps = [
":Dialect",
":shape_inference_utils",
"//tensorflow/core:framework",
"//tensorflow/core:ops",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"//tensorflow/core/platform:errors",
"@com_google_absl//absl/status",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:InferTypeOpInterface",
"@llvm-project//mlir:Parser",
"@llvm-project//mlir:Support",
],
)
tf_cc_test(
name = "ops_test",
size = "small",
srcs = ["ops_test.cc"],
deps = [
":Dialect",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"@llvm-project//mlir:ControlFlowInterfaces",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Parser",
],
)
tf_cc_test(
name = "interfaces_test",
size = "small",
srcs = ["interfaces_test.cc"],
deps = [
":Dialect",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:IR",
],
)
tf_cc_test(
name = "tf_op_registry_test",
size = "small",
srcs = ["tf_op_registry_test.cc"],
deps = [
":Dialect",
":tf_op_registry",
"//tensorflow/core:ops",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Parser",
"@llvm-project//mlir:Support",
],
)
tensorflow/core/ir/README.md
+400
View File
@@ -0,0 +1,400 @@
# TensorFlow Graph IR
This directory contains the definition of the Intermediate Representation (IR)
for TensorFlow graphs using MLIR.
## Introduction
This directory defined an MLIR dialect, the “TensorFlow Graph dialect”, that
represents accurately TensorFlow graphs. Contrary to the previous TensorFlow
dialect which made some opinionated choices that diverged from GraphDef and
TensorFlow Graph semantics, this dialect embraces TensorFlow Graph as it is. In
particular the concepts of control dependencies, requested device, assigned
device, and node name are all first-class attributes on the MLIR operations in
this dialect.
The main principle that drove the development of this dialect has been to ensure
perfect round-trip and general compatibility with existing TensorFlow semantics,
so that this solution can be deployed by default in any situation where "Graph
Optimization" and Grappler transformations are involved today, regardless of
TensorFlow V1 or V2. This new approach is also made possible by evolutions in
MLIR that allow representing graphs in a way that wasn’t possible before (more
in the [Graph operation design](#graph_operation_design) section below).
## History of Dialects for TensorFlow
MLIR started with a basic structure reflecting LLVM in that it defined a
`Module` containing a list of `Functions`. Each of these was defining a body
constrained to be a Control-Flow Graph (CFG): a list of `Blocks`, each of them
containing a list of `Operations`. A fundamental aspect of the CFG
representation is the notion of “control”: the abstract semantic model considers
that a single `Operation` executes at a given time, and the next `Operation` to
execute is necessarily the one listed immediately after[^1]. The last
`Operation` in a `Block` is a `Terminator`: it decides what is the next `Block`
where the control will be transferred (think of a branch).
When MLIR started, a first dialect -- that we were referring to as “TF control
dialect” -- was developed to model TensorFlow graphs. This dialect supported
control dependencies, but didn’t allow cycles in the graph, which forced some
tricks to model TensorFlow V1 loops and in particular the `NextIteration`
operation. While this dialect enabled some experimentation, it wasn’t seen as
really practical and another dialect was co-existing: the “tf” dialect that
we’re using currently. This dialect was designed before TF2.0
[was released](https://blog.tensorflow.org/2019/09/tensorflow-20-is-now-available.html),
and made strong assumptions about TensorFlow evolving towards a world where
eager execution and function execution become unified and V1 specific constructs
would be deprecated and disappear. As such control dependencies are not
supported and are instead implicit, control-flow V1 ops (such as Switch & Merge)
and deadness aren’t supported[^2], new device placement modelling solutions were
considered. These choices in the model enabled us to write graph transformations
as stateless DAG-to-DAG patterns that can be applied to a subgraph, without
considering the entire graph.
## Motivation
The combination of the TensorFlow and executor dialects allows for importing
most TensorFlow graphs and the TensorFlow dialect has proven enough to implement
the TF/XLA bridge, TFLite converter, and TFRT . However, the intent was for
TensorFlow 2.0 to trace TensorFlow functions directly in the TensorFlow dialect,
leaving the executor dialect only as a way to provide limited support for
TensorFlow V1 graphs.
However, the implementation of TensorFlow 2.0 didn't break away from TensorFlow
V1 entirely, instead TensorFlow functions are wrapped above TensorFlow V1 and
expose a leaky abstraction over the classical graph. As a result, the TensorFlow
dialect never got in a position to be enabled by default in TensorFlow. In
particular there are many subtle way in which TensorFlow functions diverges from
the sequential eager interpretation. For example the following pattern has been
recommended to users who intended to call a function `bar` knowing that the
first argument wasn’t necessary if they only used the first result.
```
@tf.function
def foo(z):
x = tf.Placeholder(tf.int32)
y, _ = bar(x, z)
return y
```
The use of a placeholder would throw an exception in eager mode, but “works” in
graph mode as long as inlining and pruning ensure the placeholder is removed
before execution.
Other cases involve the need for control dependencies beyond what the
auto-control-dependency tracking offers. For example the
[tf.recompute\_grad](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/ops/custom_gradient.py#L497)
creates control-dependencies on non-side-effecting ops to have a finer grain
control of memory usage.
Finally, the error modelling in TensorFlow can also be surprising. While in
eager op-by-op mode the execution is interrupted as soon as an error occurs,
`tf.function` tracing does not consider error handling as side-effecting
(otherwise it would have to add a control dependency between every node!) and as
such a program like:
```
@tf.function
def foo(x, y, variable):
b = tf.matmul(x, y)
variable.assign(1.0)
return b
```
Does not guarantee that the assignment to the variable won’t occur if an error
occurs while processing the matmul, so calling:
```
foo(1., 2., variable)
```
Throws an exception because `tf.matmul` expects rank-2 tensors, but the variable
may or may not have been assigned. As such a user may want to opt in a safer
behavior for their function:
```
@tf.function
def foo(x, y, variable):
b = tf.matmul(x, y)
with tf.control_dependencies([b]):
variable.assign(1.0)
return b
```
However, this control dependency cannot be modelled in the TensorFlow dialect:
it will be just dropped! There is no solution today to prevent the variable
assignment to be executed ahead of the `matmul` in the TensorFlow Dialect.
While many of these cases could be modeled with different constructs at the
source level, this would be a major overhaul of TensorFlow itself, and more
importantly its ecosystem. Instead, we recognize that the TensorFlow dialect as
it exists today cannot support all of these use-cases, and it prevented MLIR
from providing a general graph transformation solution for TensorFlow,
contributing to more fragmentation instead of reducing it as promised.
The rest of this document describe how this new dialect follows a more pragmatic
approach to enable MLIR deployment in TensorFlow.
## Design
This new dialect intends to allow us to replace Grappler and existing graph
transformations, for TensorFlow V1 and V2 without constraints. As such the main
principle is to support perfect roundtrip between TensorFlow Graph/GraphDef and
MLIR.
### General Operations
An individual TensorFlow `NodeDef` is translated into an individual MLIR
operation using the following form:
```
%AddV2, %ctl = tfg.AddV2(%placeholder, %placeholder_1) [%ctl_1, %ctl_2]
device("GPU") assigned_device("TPU") name("add")
{some_attribute = "some attr!"}
: (tensor<*xi32>, tensor<*xi32>) -> (tensor<*xi32>)
```
* Each operation returns an optional variadic number of tensors as well as a
control token to express potential control dependencies.
* The node type is carried in the operation mnemonic.
* The list of regular inputs is in-between parentheses.
* Optional control dependencies are exposed after the regular inputs and
printed between square brackets.
* The pre-placement “requested device” as well as the post-placement “assigned
device” information are preserved.
* The node name is carried as a first-class attribute.
* Optional “op specific” attributes can be listed between curly brackets.
* Finally, the type signature follows, omitting the control dependencies.
This structure allows for a perfect round-trip to NodeDef, while still being
ergonomic when manipulating it in MLIR (compared to the `tf\_executor` dialect
for example). The tradeoff we are making here is that we preserve all
attributes, including the “derived” ones[^3], which creates some amount of
redundancy with the signature. We may consider pruning these redundant
attributes in the future in the same way as we do in the TensorFlow dialect.
### Graph Operation
A structural operation is introduced as a container: `tfg.graph` acts as a bag
of unordered TensorFlow operations, and carries a “version” attribute that
corresponds to the
[VersionDef](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/framework/versions.proto)
present in GraphDef:
```
tfg.graph #tfg.version<producer = 42, min_consumer = 33> {
%arg0, %ctl_0 = tfg.placeholder() : () -> (tensor<*xi32>)
%add, %ctl_1 = tfg.AddV2(%arg0, %arg1)
: (tensor<*xi32>, tensor<*xi32>) -> (tensor<*xi32>)
%arg1, %ctl_2 = tfg.placeholder() : () -> (tensor<*xi32>)
}
```
Note that the `AddV2` operation is using the result of a `placeholder` operation
that is defined later in the list. This wasn’t possible in MLIR 2 years ago when
the TensorFlow dialect was designed. It was actually
[attempted to allow such unordered semantics](https://groups.google.com/a/tensorflow.org/g/mlir/c/gPQFIy9XpVw/m/hfxmBGF8AQAJ)
and break away from the CFG-centric representation, but we couldn’t reach a
consensus, and some key members of the team believed that a departure from
CFG/SSA would limit the reusability of many algorithms. On the other hand, this
choice prevented us to design a graph dialect that can just replace TensorFlow
Graph structure as-is. Since then MLIR evolved to become more general and this
feature is now available (it was motivated by the
[support for HW synthesis tools](https://llvm.discourse.group/t/rfc-allowing-dialects-to-relax-the-ssa-dominance-condition/833)).
Another recent development that made it also more friendly is the
[removal of the requirement for terminators](https://llvm.discourse.group/t/rfc-making-terminator-optional-for-single-block-graph-regions/2997):
the `tfg.graph` operation above contains a single block listing operations, and
a terminator does not have any role to play. Finally, a Dialect can now
[act as fallback for OpInterfaces](https://llvm.discourse.group/t/rfc-dialect-fallback-for-opinterface/3074),
which allows us to reuse more of the TensorFlow registry to provide information
to MLIR passes about TensorFlow operation without having to register them with
MLIR in the first place.
The `tfg.graph` operation round-trips almost perfectly to
[Graph](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/graph/graph.h#L504),
except for the `Function Library`, which I address below.
### Function Library
Functions in TensorFlow are stored as
[FunctionDef](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/framework/function.proto),
which has a signature, holds attributes, identifies argument and returned
values, and finally contains a list of nodes for its body. While on the surface
this `repeated NodeDef node_def` field looks identical to the body of
[GraphDef](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/framework/graph.proto#L17),
there are fundamental differences in the representation, and in particular the
format the edges are represented is different.
To understand these differences, it is important to realize that a key aspect of
`FunctionsDef` is that they are stored uninstantiated, and can be considered in
a similar way to a C++ template function. The signature is actually an `OpDef`,
and just like any regular TensorFlow operation the types of the arguments and
the results are encoded and constrained with attributes. These attributes are
only provided or inferred based on the function’s use: the call-site is
responsible for instantiating a function before it’s body can be represented as
a Graph. Because of this, the body of an uninstantiated function is modeled
differently than Graph body:
```
tfg.func generic @foo(%arg0 : !tfg.tensor {tfg.name = "input"},
%arg1 : !tfg.tensor {tfg.name = "another_input"})
-> (!tfg.tensor {tfg.name = "result1"},
!tfg.tensor {tfg.name = "result2"})
attributes {description = "function foo"} {
%Greater, %ctl_0 = tfg.Greater(%arg0, %arg1) name("Greater")
%G_z = tfg.get_result(%Greater) "z" : 0
%Switch, %ctl_1 = tfg.Switch(%G_z, %G_z) name("cond/Switch")
%s_true = tfg.get_result %Switch "output_true" : 0
%s_false = tfg.get_result %Switch "output_false" : 0
tfg.return(%s_true, %s_false) [%ctl_0]
}
```
Note how the tensor types `!tfg.tensor` are opaque, and every operation returns
a single tensor output and a control token. The tensor output is then unpacked
by looking up individual results by name. This is particularly visible with the
`Switch` operation where the two results are accessed using `tfg.get_result`
looking them up by name `output_true:0` and `output_false:0`. This is required
because the OpDef can define the number of output based on the attribute present
on the NodeDef, and these attributes can in turn be dependent on the attributes
added on the function during instantiation (you can read more about it in the
[description of the placeholder attribute value](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/framework/attr_value.proto#L48-L55)).
Post-instantiation, a function body is similar to the one of a graph:
```
tfg.func @foo(%arg0 : tensor<*xf32> {tfg.name = "input"},
%arg1 : tensor<*xf32> {tfg.name = "another_input"})
-> (tensor<*xi1> {tfg.name = "result1"},
tensor<*xi1> {tfg.name = "result2"})
attributes {description = "function foo"} {
%Greater, %ctl_0 = tfg.Greater(%arg0, %arg1) [%arg1.ctl] name("Greater")
: (tensor<*xf32>, tensor<*xf32>) -> tensor<*xi1>
%Switch:2, %ctl_1 = tfg.Switch(%Greater, %Greater) name("cond/Switch")
: (tensor<*xi1>, tensor<*xi1>) -> tensor<*xi1>
tfg.return(%Switch#0, %Switch#1) [%ctl_0]
}
```
The operations aren’t ordered, except for the `tfg.return` which is a terminator
and must be the last operation. The only remaining difference with a graph is in
the handling of the function signature (arguments and returned values), and
attributes.
There is one aspect of the modelling worth mentioning from the MLIR point of
view: FunctionDef allows for nodes in a graph to express input control
dependencies from function arguments. However, in MLIR you need an actual
[SSA](https://en.wikipedia.org/wiki/Static_single_assignment_form) value to add
an edge between two operations. These values are typed and this is why
operations define a control token (like `%ctl_0`). We apply the same recipe for
arguments and for each of them we define a control token. We omit these “shadow
arguments” from the textual form, but in-memory the MLIR function has really 4
arguments:
```
tfg.func @foo(%arg0 : tensor<*xf32> {tfg.name = "input"}, %arg0.ctl : !tfg.control
%arg1 : tensor<*xf32> {tfg.name = "another_input"}, %arg1.ctl : !tfg.control)
-> (tensor<*xi1> {tfg.name = "result1"},
tensor<*xi1> {tfg.name = "result2"})
attributes {description = "function foo"} {
...
```
The convention is that callers are only exposed to the non-control input
(`%arg0` and `%arg1`) while the control tokens are only intended to be visible
and used in the body. This makes it very aligned with how TensorFlow works.
Inside the body, values for the control dependencies on the arguments are
available with a `.ctl` suffix (i.e. `%arg0.ctl` and `%arg1.ctl`).
### Saved Model
The basic blocks above are enough to model `GraphDef`, but not the entirety of
SavedModel. However, most of the use cases that we’re targeting right now are in
the scope of the existing GraphOptimization and Grappler APIs, which aren’t
really coupled to SavedModel. The user can load a SavedModel independently of
MLIR and invoke MLIR transformations on a Function or Graph from there. There is
also already a dialect to model the specific aspects of SavedModel, it is
currently wrapping around the TensorFlow executor dialect and the TensorFlow
dialect, and we may look into integrating it with the `tfg` dialect in the
future. For these reasons, we mostly leave out modeling the Saved Model for
future work right now.
### Future Enhancements
Functional control-flow is modeled with nodes in the graph invoking functions in
the library. MLIR supports `region`s, which is a concept that allows attaching
subgraphs directly inside a graph, making it more friendly to optimizations. For
example a conditional operation can represent the two branches subgraph in the
TensorFlow dialect directly as follows:
```
%0, %1, %2 = "tf.IfRegion"(%arg0) ({
%t0 = "tf.Abs"(%arg1) : (tensor<2xf32>) -> tensor<2xf32>
%t1 = "tf.Acos"(%arg1) : (tensor<2xf32>) -> tensor<2xf32>
%t2 = "tf.Acosh"(%arg1) : (tensor<2xf32>) -> tensor<2xf32>
"tf.Yield"(%t0, %t1, %t2) : (tensor<2xf32>, tensor<2xf32>, tensor<2xf32>) -> ()
}, {
%e0 = "tf.Neg"(%arg1) : (tensor<2xf32>) -> tensor<2xf32>
%e1 = "tf.Relu"(%arg1) : (tensor<2xf32>) -> tensor<2xf32>
%e2 = "tf.Sin"(%arg1) : (tensor<2xf32>) -> tensor<2xf32>
"tf.Yield"(%e0, %e1, %e2) : (tensor<2xf32>, tensor<2xf32>, tensor<2xf32>)
}): (tensor<i1>) -> (tensor<2xf32>, tensor<2xf32>, tensor<2xf32>)
%3 = "tf.Add"(%0, %1) : (tensor<2xf32>, tensor<2xf32>) -> tensor<2xf32>
%4 = "tf.Add"(%2, %3) : (tensor<2xf32>, tensor<2xf32>) -> tensor<2xf32>
```
## Integration
MLIR transformations in this dialect will operate on a module that will contain
at most one `graph` operation as well as a list of functions. This interface
will make such transformations suitable for fit within Grappler or as
GraphOptimization interchangeably.
Instead of a flat graph, an entry function will be provided when feeds/fetches
are available for the main graph (PRE_PLACEMENT graph optimizations execute in
Session before feeds/fetches are provided).
## FAQ
### Why not just use the TensorFlow Executor Dialect?
The executor dialect wasn’t designed to write transformation: it is designed as
a wrapper around the TensorFlow dialect: the intent was for it to be a stepping
stone to integrate MLIR and TensorFlow, and then disappear when TensorFlow V1
graphs would be deprecated. This new dialect embraces TensorFlow as it is
instead of as I wish it would be.
In particular the executor dialect represents each TensorFlow node as an
isolated “subgraph” nested under an “island” operation. This requires 3
operations and an extra region for each TensorFlow node, which is quite
inefficient in memory as well as requiring extra indirection when pattern
matching or updating nodes in the graph.
### What happens to the existing TensorFlow Dialects?
The existing TensorFlow dialect is suitable for representing a large subset of
TensorFlow programs (like models that intend to convert to TFLite, or XLA), and
for such cases we will continue to use it.
### What happens to the existing TensorFlow Executor Dialect?
This new TensorFlow Graph Dialect could be used to replace the Executor Dialect
as the standalone staging importing format. Importing from GraphDef/Graph would
always go through the TensorFlow Graph Dialect before using some clustering or
promotion algorithms to raise some subgraphs to the TensorFlow Dialect, just
like we do now to cluster islands operations in TensorFlow Executor Dialect.
The details of such mechanisms are left for future work.
<!-- Footnotes -->
[^1]: While the semantic model is sequential, this does not prevent an
implementation to execute operation in parallel when proven safe. This is
similar to how a superscalar CPU involves implicit parallelism. For
example when mapping the TensorFlow dialect to TFRT, only side-effecting
operations (Variable accesses for example) are sequenced.
[^2]: One of the first tools built with this was the TF->TFlite converter
(replacing TOCO). Since V1 control-flow isn’t supported on TFLite this
wasn’t a limitation.
[^3]: Derived attributes is a concept used in the TensorFlow dialect: since MLIR
models type and shape information on each individual result produced by an
operation, some attributes that are inserted for the sole purpose of
typing are redundant and eliminated in MLIR.
tensorflow/core/ir/dialect.h
+87
View File
@@ -0,0 +1,87 @@
/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_CORE_IR_DIALECT_H_
#define TENSORFLOW_CORE_IR_DIALECT_H_
#include "mlir/IR/BuiltinTypes.h" // from @llvm-project
#include "mlir/IR/Diagnostics.h" // from @llvm-project
#include "mlir/IR/Dialect.h" // from @llvm-project
#include "mlir/IR/OpImplementation.h" // from @llvm-project
#include "mlir/IR/TypeUtilities.h" // from @llvm-project
#include "tensorflow/core/ir/types/dialect.h"
namespace mlir {
namespace tfg {
// Include the relevant TensorFlow attrs/types directly in the TFG namespace.
using mlir::tf_type::Bfloat16RefType; // NOLINT
using mlir::tf_type::BoolRefType; // NOLINT
using mlir::tf_type::Complex128RefType; // NOLINT
using mlir::tf_type::Complex64RefType; // NOLINT
using mlir::tf_type::ControlType; // NOLINT
using mlir::tf_type::DoubleRefType; // NOLINT
using mlir::tf_type::Float8E4M3B11FNUZRefType; // NOLINT
using mlir::tf_type::Float8E4M3FNRefType; // NOLINT
using mlir::tf_type::Float8E4M3FNUZRefType; // NOLINT
using mlir::tf_type::Float8E5M2FNUZRefType; // NOLINT
using mlir::tf_type::Float8E5M2RefType; // NOLINT
using mlir::tf_type::FloatRefType; // NOLINT
using mlir::tf_type::FuncAttr; // NOLINT
using mlir::tf_type::HalfRefType; // NOLINT
using mlir::tf_type::Int16RefType; // NOLINT
using mlir::tf_type::Int2RefType; // NOLINT
using mlir::tf_type::Int32RefType; // NOLINT
using mlir::tf_type::Int4RefType; // NOLINT
using mlir::tf_type::Int64RefType; // NOLINT
using mlir::tf_type::Int8RefType; // NOLINT
using mlir::tf_type::OpaqueTensorType; // NOLINT
using mlir::tf_type::PlaceholderAttr; // NOLINT
using mlir::tf_type::Qint16RefType; // NOLINT
using mlir::tf_type::Qint16Type; // NOLINT
using mlir::tf_type::Qint32RefType; // NOLINT
using mlir::tf_type::Qint32Type; // NOLINT
using mlir::tf_type::Qint8RefType; // NOLINT
using mlir::tf_type::Qint8Type; // NOLINT
using mlir::tf_type::Quint16RefType; // NOLINT
using mlir::tf_type::Quint16Type; // NOLINT
using mlir::tf_type::Quint8RefType; // NOLINT
using mlir::tf_type::Quint8Type; // NOLINT
using mlir::tf_type::ResourceRefType; // NOLINT
using mlir::tf_type::ResourceType; // NOLINT
using mlir::tf_type::ShapeAttr; // NOLINT
using mlir::tf_type::StringRefType; // NOLINT
using mlir::tf_type::StringType; // NOLINT
using mlir::tf_type::Uint16RefType; // NOLINT
using mlir::tf_type::Uint2RefType; // NOLINT
using mlir::tf_type::Uint32RefType; // NOLINT
using mlir::tf_type::Uint4RefType; // NOLINT
using mlir::tf_type::Uint64RefType; // NOLINT
using mlir::tf_type::Uint8RefType; // NOLINT
using mlir::tf_type::VariantRefType; // NOLINT
using mlir::tf_type::VariantType; // NOLINT
using mlir::tf_type::VersionAttr; // NOLINT
class TFGraphOpAsmInterface;
class TFOp;
} // namespace tfg
} // namespace mlir
// Dialect main class is defined in ODS, we include it here.
#include "tensorflow/core/ir/dialect.h.inc" // IWYU pragma: export
// ODS-generated attribute classes.
#define GET_ATTRDEF_CLASSES
#include "tensorflow/core/ir/attributes.h.inc"
#endif // TENSORFLOW_CORE_IR_DIALECT_H_
tensorflow/core/ir/dialect.td
+177
View File
@@ -0,0 +1,177 @@
/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TFG_DIALECT
#define TFG_DIALECT
include "mlir/IR/OpBase.td"
include "mlir/Interfaces/InferTypeOpInterface.td"
// ODS Definition for the dialect, see https://mlir.llvm.org/docs/OpDefinitions/
// for more information.
//===----------------------------------------------------------------------===//
// TFGraph dialect definitions
//===----------------------------------------------------------------------===//
def TFGraphDialect : Dialect {
let name = "tfg";
let summary = "This dialect models TensorFlow Graphs as encoded in GraphDef.";
let description = [{
This dialect is modeling TensorFlow GraphDefs and intended to provide a high
level of fidelity.
The attribute mappings from GraphDef are listed down below,
Graph/Function Attributes:
FunctionDef.attr will prepand with "tf" prefix
FunctionDef.signature.name <-> "sym_name"
FunctionDef.signature.description <-> "description"
FunctionDef.signature.is_stateful <-> "is_stateful"
FunctionDef.signature.gradient <-> "gradient"
FunctionDef.resource_arg_unique_id <-> "resource_arg_unique_ids_keys"
FunctionDef.resource_arg_unique_id <-> "resource_arg_unique_ids_values"
Input Attributes:
FunctionDef.signature.input_arg.name <-> "tfg.name"
FunctionDef.signature.input_arg.description <-> "tfg.description"
FunctionDef.signature.input_arg.handle_data <-> "tfg.handle_data"
FunctionDef.signature.input_arg.is_ref <-> "tfg.is_ref"
FunctionDef.arg_attr will prepand with "tf" prefix
Output Attributes:
FunctionDef.signature.output_arg.name <-> "tfg.name"
FunctionDef.signature.output_arg.description <-> "tfg.description"
FunctionDef.signature.output_arg.handle_data <-> "tfg.handle_data"
FunctionDef.signature.output_arg.type <-> "tfg.dtype"
FunctionDef.signature.control_output <-> "tfg.control_ret_name_"
Node Attributes:
NodeDef.device <-> "_mlir_device"
NodeDef.name <-> "_mlir_name"
NodeDef.attr <-> "_output_shape"
NodeDef.experimental_type <-> "_mlir_fulltype"
}];
let extraClassDeclaration = [{
StringAttr getNameAttrIdentifier() const { return name_key_; }
static constexpr StringLiteral getNameAttrKey() { return {"_mlir_name"}; }
StringAttr getDeviceAttrIdentifier() const { return device_key_; }
static constexpr StringLiteral getDeviceAttrKey() {
return {"_mlir_device"};
}
StringAttr getAssignedDeviceAttrIdentifier() const {
return assigned_device_key_;
}
static constexpr StringLiteral getAssignedDeviceAttrKey() {
return {"_mlir_assigned_device"};
}
StringAttr getFullTypeAttrIdentifier() const { return fulltype_key_; }
static constexpr StringLiteral getFullTypeAttrKey() {
return {"_mlir_fulltype"};
}
StringAttr getTfgNameAttrIdentifier() const { return tfg_name_key_; }
static constexpr StringRef getTfgNameAttrKey() { return "tfg.name"; }
StringAttr getTfgDescriptionAttrIdentifier() const {
return tfg_description_key_;
}
static constexpr StringRef getTfgDescriptionAttrKey() {
return {"tfg.description"};
}
StringAttr getTfgIsRefAttrIdentifier() const { return tfg_is_ref_key_; }
static constexpr StringRef getTfgIsRefAttrKey() { return {"tfg.is_ref"}; }
StringAttr getTfgHandleDataAttrIdentifier() const {
return tfg_handle_data_key_;
}
static constexpr StringRef getTfgHandleDataAttrKey() {
return {"tfg.handle_data"};
}
StringAttr getTfgFullTypeAttrIdentifier() const {
return tfg_full_type_key_;
}
static constexpr StringRef getTfgFullTypeAttrKey() {
return {"tfg.experimental_full_type"};
}
StringAttr getLiftedGraphFuncNameAttrIdentifier() const {
return lifted_graph_func_name_;
}
static constexpr StringRef getLiftedGraphFuncNameKey() {
return {"_mlir_lifted_graph"};
}
// Cached accessor for the control type.
ControlType getControlType() const { return control_ty_; }
// Print an operation that belongs to this dialect if unregistered.
void printCustomTfOp(Operation *op, OpAsmPrinter &printer) const;
// Returns the hook to parse an operation belonging to this dialect, even
// if unregistered.
std::optional<ParseOpHook> getParseOperationHook(StringRef opName) const
override;
// Returns the took to print an operation belonging to this dialect, even
// if unregistered.
llvm::unique_function<void(Operation *, OpAsmPrinter &)>
getOperationPrinter(Operation *op) const override;
// Functions for checking operation categories.
#define GET_OP_CATEGORIES
#include "tensorflow/core/ir/tf_op_names.inc"
private:
// Fallback implementation of OpAsmOpInterface.
TFGraphOpAsmInterface *fallbackOpAsmInterface_ = nullptr;
// Cached TensorFlow operation names.
#define GET_OP_NAME_DECLS
#include "tensorflow/core/ir/tf_op_names.inc"
// Cached identifier for efficiency purpose.
StringAttr assigned_device_key_;
StringAttr device_key_;
StringAttr fulltype_key_;
StringAttr lifted_graph_func_name_;
StringAttr name_key_;
StringAttr tfg_description_key_;
StringAttr tfg_full_type_key_;
StringAttr tfg_handle_data_key_;
StringAttr tfg_is_ref_key_;
StringAttr tfg_name_key_;
// Cached control type.
ControlType control_ty_;
}];
let cppNamespace = "::mlir::tfg";
let useDefaultAttributePrinterParser = 1;
let hasNonDefaultDestructor = 1;
let hasOperationInterfaceFallback = 1;
}
#endif // TFG_DIALECT
tensorflow/core/ir/importexport/BUILD
+255
View File
@@ -0,0 +1,255 @@
load("//tensorflow:tensorflow.bzl", "tf_cc_binary")
load("//tensorflow/compiler/mlir:glob_lit_test.bzl", "glob_lit_tests")
load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
default_visibility = [
":__subpackages__",
":friends",
"//tensorflow/compiler/mlir:__subpackages__",
"//tensorflow/core:__subpackages__",
"//tensorflow/core:dependency_allowlist",
"//tensorflow/tools/tfg_graph_transforms:__subpackages__",
],
licenses = ["notice"], # Apache 2.0
)
package_group(
name = "friends",
packages = [
"//learning/brain/contrib/tpu_modeling/inference_converter_v2/mlir/...",
"//waymo/ml/cn/...",
"//waymo/ml/compiler/mlir/...",
],
)
cc_library(
name = "parse_text_proto",
srcs = [
"parse_text_proto.cc",
],
hdrs = ["parse_text_proto.h"],
deps = [
"//tensorflow/core:framework_lite",
"//tensorflow/core:lib_proto_parsing",
"//tensorflow/core/platform:casts",
"@com_google_absl//absl/strings",
"@com_google_protobuf//:protobuf",
],
)
cc_library(
name = "mangling",
srcs = [
"mangling.cc",
],
hdrs = ["mangling.h"],
deps = [
":parse_text_proto",
"//tensorflow/core:framework",
"//tensorflow/core:framework_lite",
"//tensorflow/core:lib_proto_parsing",
"//tensorflow/core:protos_all_cc",
"@com_google_absl//absl/strings",
],
)
cc_library(
name = "load_proto",
srcs = [
"load_proto.cc",
],
hdrs = ["load_proto.h"],
deps = [
":parse_text_proto",
"//tensorflow/core:framework_lite",
"//tensorflow/core:lib_proto_parsing",
"@com_google_absl//absl/strings",
"@llvm-project//llvm:Support",
],
)
cc_library(
name = "convert_types",
srcs = [
"convert_types.cc",
],
hdrs = ["convert_types.h"],
deps = [
"//tensorflow/core:framework",
"//tensorflow/core:lib_proto_parsing",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core/ir:Dialect",
"//tensorflow/core/ir/types:Dialect",
"//tensorflow/core/platform:statusor",
"@com_google_absl//absl/strings",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Support",
],
)
cc_library(
name = "convert_tensor",
srcs = [
"convert_tensor.cc",
],
hdrs = ["convert_tensor.h"],
deps = [
":convert_types",
":mangling",
"//tensorflow/core:framework",
"//tensorflow/core:framework_lite",
"//tensorflow/core:lib_proto_parsing",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core:tflite_portable_logging",
"//tensorflow/core/ir:Dialect",
"//tensorflow/core/ir/types:Dialect",
"//tensorflow/core/platform:errors",
"//tensorflow/core/platform:statusor",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Support",
"@tsl//tsl/platform:ml_dtypes",
],
)
cc_library(
name = "savedmodel_import",
srcs = ["savedmodel_import.cc"],
hdrs = ["savedmodel_import.h"],
deps = [
":graphdef_import",
"//tensorflow/core:core_cpu_base",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core/platform:errors",
"//tensorflow/core/platform:statusor",
"@llvm-project//mlir:IR",
],
)
cc_library(
name = "savedmodel_export",
srcs = ["savedmodel_export.cc"],
hdrs = ["savedmodel_export.h"],
deps = [
":graphdef_export",
"//tensorflow/core:core_cpu_base",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core/platform:errors",
"//tensorflow/core/platform:status",
"@llvm-project//mlir:IR",
],
)
cc_library(
name = "graphdef_import",
srcs = [
"functiondef_import.cc",
"graphdef_import.cc",
],
hdrs = [
"functiondef_import.h",
"graphdef_import.h",
],
deps = [
":convert_attributes",
":convert_types",
"//tensorflow/core:core_cpu_base",
"//tensorflow/core:framework",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core/ir:Dialect",
"//tensorflow/core/ir/types:Dialect",
"//tensorflow/core/platform:errors",
"//tensorflow/core/platform:status",
"//tensorflow/core/platform:statusor",
"//tensorflow/core/platform:stringpiece",
"@com_google_absl//absl/container:flat_hash_map",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Support",
],
)
cc_library(
name = "graphdef_export",
srcs = [
"functiondef_export.cc",
"graphdef_export.cc",
],
hdrs = [
"functiondef_export.h",
"graphdef_export.h",
],
deps = [
":convert_attributes",
":convert_types",
"//tensorflow/core:core_cpu_base",
"//tensorflow/core:framework",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core/ir:Dialect",
"//tensorflow/core/ir/types:Dialect",
"//tensorflow/core/platform:errors",
"//tensorflow/core/platform:status",
"//tensorflow/core/platform:statusor",
"@com_google_absl//absl/strings",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Support",
],
)
cc_library(
name = "convert_attributes",
srcs = [
"convert_attributes.cc",
],
hdrs = ["convert_attributes.h"],
deps = [
":convert_tensor",
":convert_types",
":mangling",
"//tensorflow/core:framework",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core/ir:Dialect",
"//tensorflow/core/ir/types:Dialect",
"//tensorflow/core/platform:errors",
"//tensorflow/core/platform:statusor",
"@com_google_absl//absl/strings",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Support",
"@xla//xla:status_macros",
],
)
# Command line tool to convert to/from GraphDef to MLIR Graph.
tf_cc_binary(
name = "tfg-translate",
srcs = ["tfg-translate.cc"],
deps = [
":graphdef_export",
":graphdef_import",
":load_proto",
"//tensorflow/compiler/mlir:init_mlir",
"//tensorflow/core:ops", # Ops need to be registered for import.
"//tensorflow/core/framework:graph_debug_info_proto_cc",
"//tensorflow/core/framework:graph_proto_cc",
"//tensorflow/core/ir:Dialect",
"//tensorflow/core/platform:status",
"@com_google_absl//absl/log",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Support",
"@llvm-project//mlir:TranslateLib",
],
)
glob_lit_tests(
name = "all_tests",
data = ["//tensorflow/core/ir:test_utilities"],
driver = "//tensorflow/core/ir:run_lit.sh",
exclude = [],
test_file_exts = ["mlir"],
)
tensorflow/core/ir/importexport/convert_attributes.cc
+509
View File
@@ -0,0 +1,509 @@
/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/core/ir/importexport/convert_attributes.h"
#include <string>
#include <vector>
#include "llvm/ADT/StringSet.h"
#include "llvm/ADT/TypeSwitch.h"
#include "llvm/Support/raw_ostream.h"
#include "mlir/IR/Builders.h" // from @llvm-project
#include "mlir/IR/BuiltinAttributes.h" // from @llvm-project
#include "mlir/IR/Location.h" // from @llvm-project
#include "mlir/Support/DebugStringHelper.h" // from @llvm-project
#include "mlir/Support/LLVM.h" // from @llvm-project
#include "xla/status_macros.h"
#include "tensorflow/core/framework/attr_value.pb.h"
#include "tensorflow/core/framework/full_type.pb.h"
#include "tensorflow/core/framework/op_def.pb.h"
#include "tensorflow/core/ir/dialect.h"
#include "tensorflow/core/ir/importexport/convert_tensor.h"
#include "tensorflow/core/ir/importexport/convert_types.h"
#include "tensorflow/core/ir/importexport/mangling.h"
#include "tensorflow/core/ir/types/dialect.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/platform/statusor.h"
using tensorflow::AttrValue;
using tensorflow::AttrValueMap;
using tensorflow::DataType;
using tensorflow::NodeDef;
using tensorflow::Status;
using tensorflow::StatusOr;
using tensorflow::TensorProto;
using tensorflow::TensorShapeProto;
using tensorflow::errors::InvalidArgument;
using tensorflow::errors::Unimplemented;
namespace mlir {
namespace tfg {
namespace {
// Converts a location to the debug information for the node def.
Status ConvertLocation(Location inst_loc,
NodeDef::ExperimentalDebugInfo* debug_info) {
if (auto call_site = mlir::dyn_cast<CallSiteLoc>(inst_loc)) {
if (auto name_loc = mlir::dyn_cast<NameLoc>(call_site.getCallee())) {
debug_info->add_original_node_names(name_loc.getName().data());
}
} else if (auto fused = mlir::dyn_cast<FusedLoc>(inst_loc)) {
auto locations = fused.getLocations();
if (locations.size() <= 1)
return absl::InvalidArgumentError("Expected experimental debug info.");
// skip the first one, which is the name of the node_def.
for (int i = 0, end = locations.size() - 1; i < end; ++i) {
TF_RETURN_IF_ERROR(ConvertLocation(locations[i], debug_info));
}
}
return absl::OkStatus();
}
Status ConvertAttribute(BoolAttr attr, AttrValue* value) {
value->set_b(attr.getValue());
return absl::OkStatus();
}
Status ConvertAttribute(IntegerAttr attr, AttrValue* value) {
value->set_i(attr.getInt());
return absl::OkStatus();
}
Status ConvertAttribute(FloatAttr attr, AttrValue* value) {
value->set_f(attr.getValueAsDouble());
return absl::OkStatus();
}
Status ConvertAttribute(ElementsAttr attr, AttrValue* value) {
return ConvertToTensorProto(attr, value->mutable_tensor());
}
Status ConvertAttribute(PlaceholderAttr attr, AttrValue* value) {
value->set_placeholder(attr.getValue().str());
return absl::OkStatus();
}
Status ConvertAttribute(ShapeAttr attr, AttrValue* value) {
SetTensorShapeProto(attr, value->mutable_shape());
return absl::OkStatus();
}
Status ConvertAttribute(FlatSymbolRefAttr attr, AttrValue* value) {
value->mutable_func()->set_name(attr.getValue().str());
return absl::OkStatus();
}
Status ConvertAttribute(FuncAttr attr, bool remove_ref_type, AttrValue* value) {
TF_RETURN_IF_ERROR(
ConvertAttribute(mlir::cast<FlatSymbolRefAttr>(attr.getName()), value));
TF_RETURN_IF_ERROR(ConvertAttributes(attr.getAttrs().getValue(),
/*attrs_to_ignore=*/{}, remove_ref_type,
value->mutable_func()->mutable_attr()));
return absl::OkStatus();
}
Status ConvertAttribute(StringAttr attr, AttrValue* value) {
value->set_s(attr.str());
return absl::OkStatus();
}
Status ConvertAttribute(Type type, bool remove_ref_type, AttrValue* value) {
DataType dtype;
TF_RETURN_IF_ERROR(ConvertToDataType(type, &dtype));
if (tensorflow::IsRefType(dtype)) dtype = tensorflow::RemoveRefType(dtype);
value->set_type(dtype);
return absl::OkStatus();
}
Status ConvertAttribute(const TypeAttr& type, bool remove_ref_type,
AttrValue* value) {
return ConvertAttribute(type.getValue(), remove_ref_type, value);
}
Status ConvertAttribute(const UnitAttr& attr, AttrValue* value) {
value->clear_value();
return absl::OkStatus();
}
Status ConvertAttribute(const ArrayAttr& attr, bool remove_ref_type,
AttrValue* value) {
auto* list = value->mutable_list();
for (Attribute a : attr.getValue()) {
if (auto attr = mlir::dyn_cast<BoolAttr>(a)) {
list->add_b(attr.getValue());
} else if (auto attr = mlir::dyn_cast<IntegerAttr>(a)) {
list->add_i(attr.getInt());
} else if (auto attr = mlir::dyn_cast<FloatAttr>(a)) {
list->add_f(attr.getValueAsDouble());
} else if (auto attr = mlir::dyn_cast<StringAttr>(a)) {
AttrValue nested_value;
TF_RETURN_IF_ERROR(ConvertAttribute(attr, &nested_value));
switch (nested_value.value_case()) {
case AttrValue::kS:
list->add_s(nested_value.s());
break;
case AttrValue::kType:
list->add_type(nested_value.type());
break;
case AttrValue::kShape:
*list->add_shape() = nested_value.shape();
break;
default:
return absl::UnimplementedError("Unhandled nested attribute!");
}
} else if (auto attr = mlir::dyn_cast<ElementsAttr>(a)) {
TensorProto tensor;
TF_RETURN_IF_ERROR(ConvertToTensorProto(attr, &tensor));
*list->add_tensor() = tensor;
} else if (auto attr = mlir::dyn_cast<FlatSymbolRefAttr>(a)) {
AttrValue attr_val;
TF_RETURN_IF_ERROR(ConvertAttribute(attr, &attr_val));
*list->add_func() = attr_val.func();
} else if (auto attr = mlir::dyn_cast<FuncAttr>(a)) {
AttrValue attr_val;
TF_RETURN_IF_ERROR(ConvertAttribute(attr, remove_ref_type, &attr_val));
*list->add_func() = attr_val.func();
} else if (auto attr = mlir::dyn_cast<TypeAttr>(a)) {
AttrValue attr_val;
// For type attributes, we only propagate the element type.
Type elt_type = attr.getValue();
if (auto shaped_type = mlir::dyn_cast<ShapedType>(elt_type)) {
elt_type = shaped_type.getElementType();
}
TF_RETURN_IF_ERROR(
ConvertAttribute(elt_type, remove_ref_type, &attr_val));
list->add_type(attr_val.type());
} else if (auto attr = mlir::dyn_cast<ShapeAttr>(a)) {
AttrValue attr_val;
TF_RETURN_IF_ERROR(ConvertAttribute(attr, &attr_val));
*list->add_shape() = attr_val.shape();
} else {
return absl::UnimplementedError(absl::StrCat(
"Unhandled MLIR attribute in export to graph:", debugString(a)));
}
}
return absl::OkStatus();
}
} // namespace
absl::StatusOr<AttrValue> ConvertAttribute(Attribute attr) {
AttrValue value;
if (auto symbol_ref = mlir::dyn_cast<SymbolRefAttr>(attr)) {
TF_RETURN_IF_ERROR(
ConvertAttribute(mlir::cast<FlatSymbolRefAttr>(symbol_ref), &value));
return value;
}
if (auto func_attr = mlir::dyn_cast<FuncAttr>(attr)) {
TF_RETURN_IF_ERROR(
ConvertAttribute(func_attr, /*remove_ref_type=*/false, &value));
return value;
}
if (mlir::isa<AffineMapAttr>(attr))
return absl::UnimplementedError("AffineMap attribute unimplemented");
TF_RETURN_IF_ERROR(
llvm::TypeSwitch<Attribute, Status>(attr)
.Case<BoolAttr, IntegerAttr, FloatAttr, StringAttr, ElementsAttr,
UnitAttr, ShapeAttr, PlaceholderAttr>([&](auto derived_attr) {
return ConvertAttribute(derived_attr, &value);
})
.Case<ArrayAttr, TypeAttr>([&](auto derived_attr) {
return ConvertAttribute(derived_attr,
/*remove_ref_type=*/false, &value);
})
.Default([&](Attribute attr) {
return absl::UnimplementedError(absl::StrCat(
"Unhandled attribute kind for attribute: ", debugString(attr)));
}));
return value;
}
Status ConvertAttributes(ArrayRef<NamedAttribute> attrs,
ArrayRef<StringRef> attrs_to_ignore,
bool remove_ref_type, AttrValueMap* values) {
StringSet<> ignored_attrs;
ignored_attrs.insert(attrs_to_ignore.begin(), attrs_to_ignore.end());
AttrValueMap func_call_attrs;
for (const NamedAttribute& named_attr : attrs) {
std::string name_str = named_attr.getName().str();
auto attr = named_attr.getValue();
absl::string_view name = name_str;
if (ignored_attrs.contains(name_str)) {
// The name, device spec of a TF op or function are not stored as
// AttrValue inside NodeDef, but we model them using attribute inside
// MLIR. So we need to ignore them when going back to AttrValue here.
continue;
}
if (mangling_util::IsMangledAttributeName(name)) {
// In MLIR, attributes for functions requires dialect prefix. We need to
// remove TF dialect prefix before converting to AttrValue.
name = mangling_util::DemangleAttributeName(name);
}
TF_ASSIGN_OR_RETURN(AttrValue value, ConvertAttribute(attr));
if (mlir::isa<SymbolRefAttr>(attr)) {
func_call_attrs[std::string(name)] = value;
continue;
}
if (mlir::isa<FuncAttr>(attr)) {
func_call_attrs[std::string(name)] = value;
continue;
}
// According to the NodeDef proto definition, an attribute name from the
// input TensorFlow GraphDef shouldn't contain '.'. If it does appear in
// the attribute from MLIR, it is treated as an attribute from function
// calls.
std::vector<std::string> name_tokens =
absl::StrSplit(name, '.', absl::SkipEmpty());
TF_RET_CHECK(!name_tokens.empty());
TF_RET_CHECK(name_tokens.size() <= 2);
auto it = func_call_attrs.find(name_tokens[0]);
if (it == func_call_attrs.end())
(*values)[std::string(name)] = value;
else
(*it->second.mutable_func()->mutable_attr())[name_tokens[1]] = value;
}
for (const auto& it : func_call_attrs) {
(*values)[it.first] = it.second;
}
return absl::OkStatus();
}
Status SetShapeAttribute(absl::string_view name, ShapedType shaped_type,
AttrValueMap* values) {
AttrValue value;
SetTensorShapeProto(shaped_type, value.mutable_shape());
auto result = values->insert({std::string(name), value});
if (!result.second) {
// This should be extremely rare as it means we are adding the same
// attribute multiple times/have some redundancy in representing this
// attribute.
TensorShapeProto actual_shape = result.first->second.shape();
// Just check via string output as we shouldn't get here and if we do they
// should be trivially the same, else fail.
std::string new_shape_string = value.shape().ShortDebugString();
if (actual_shape.ShortDebugString() != new_shape_string) {
return absl::InvalidArgumentError(absl::StrCat(
"Expected ", new_shape_string, " '", name, "' attribute but found ",
actual_shape.ShortDebugString()));
}
}
return absl::OkStatus();
}
// Converts non func AttrValue proto into an MLIR attribute. Func attribute is
// exclused in this function because the function might be renamed when the
// function definition is imported.
absl::StatusOr<Attribute> ConvertNonFuncAttributeValue(const AttrValue& value,
Builder& builder) {
switch (value.value_case()) {
case AttrValue::kI:
return builder.getI64IntegerAttr(value.i());
case AttrValue::kS:
return builder.getStringAttr(value.s());
case AttrValue::kF:
return builder.getFloatAttr(builder.getF32Type(), value.f());
case AttrValue::kB:
return builder.getBoolAttr(value.b());
case AttrValue::kType: {
Type type;
TF_RETURN_IF_ERROR(ConvertDataType(value.type(), builder, &type));
return TypeAttr::get(type);
}
case AttrValue::kShape:
return ConvertTensorShapeProto(value.shape(), builder.getContext());
case AttrValue::kTensor:
return ConvertTensorProto(value.tensor(), builder);
case AttrValue::kList: {
absl::InlinedVector<Attribute, 8> attrs;
for (const auto& item : value.list().i())
attrs.push_back(builder.getI64IntegerAttr(item));
for (const auto& item : value.list().s())
attrs.push_back(builder.getStringAttr(item));
for (const auto& item : value.list().f())
attrs.push_back(builder.getFloatAttr(builder.getF32Type(), item));
for (const auto& item : value.list().b())
attrs.push_back(builder.getBoolAttr(item));
for (const auto& item : value.list().type()) {
Type type;
TF_RETURN_IF_ERROR(ConvertDataType(DataType(item), builder, &type));
attrs.push_back(TypeAttr::get(type));
}
for (const auto& item : value.list().shape()) {
TF_ASSIGN_OR_RETURN(
auto attr, ConvertTensorShapeProto(item, builder.getContext()));
attrs.push_back(attr);
}
for (const auto& item : value.list().tensor()) {
TF_ASSIGN_OR_RETURN(auto attr, ConvertTensorProto(item, builder));
attrs.push_back(attr);
}
for (const auto& func_attr : value.list().func()) {
NamedAttrList subattrs;
for (const auto& subattr : func_attr.attr()) {
TF_ASSIGN_OR_RETURN(auto attr,
ConvertAttributeValue(subattr.second, builder));
if (subattr.first.empty())
return absl::InvalidArgumentError("empty func_attr name");
subattrs.push_back(builder.getNamedAttr(subattr.first, attr));
}
attrs.push_back(FuncAttr::get(builder.getContext(), func_attr.name(),
builder.getDictionaryAttr(subattrs)));
}
return builder.getArrayAttr(llvm::ArrayRef(attrs.begin(), attrs.end()));
}
case AttrValue::VALUE_NOT_SET:
return builder.getUnitAttr();
case AttrValue::kPlaceholder:
return PlaceholderAttr::get(builder.getContext(), value.placeholder());
default:
return absl::UnimplementedError(
absl::StrCat("Attribute ", value.DebugString()));
}
}
absl::StatusOr<Attribute> ConvertAttributeValue(const AttrValue& value,
Builder& builder) {
switch (value.value_case()) {
case AttrValue::kFunc: {
NamedAttrList attrs;
for (const auto& func_attr : value.func().attr()) {
if (func_attr.first.empty())
return absl::InvalidArgumentError("empty attr name");
TF_ASSIGN_OR_RETURN(auto attr,
ConvertAttributeValue(func_attr.second, builder));
attrs.push_back(builder.getNamedAttr(func_attr.first, attr));
}
auto func_attrs = builder.getDictionaryAttr(attrs);
return FuncAttr::get(builder.getContext(), value.func().name(),
func_attrs);
}
default:
return ConvertNonFuncAttributeValue(value, builder);
}
}
absl::StatusOr<tf_type::FullTypeAttr> ConvertAttribute(
const tensorflow::FullTypeDef& full_type, Builder& builder) {
using FullTypeAttr = ::mlir::tf_type::FullTypeAttr;
SmallVector<FullTypeAttr> args;
for (const tensorflow::FullTypeDef& it : full_type.args()) {
TF_ASSIGN_OR_RETURN(FullTypeAttr arg, ConvertAttribute(it, builder));
args.push_back(arg);
}
Attribute attr;
switch (full_type.attr_case()) {
case tensorflow::FullTypeDef::AttrCase::kS:
attr = builder.getStringAttr(full_type.s());
break;
case tensorflow::FullTypeDef::AttrCase::kI:
attr = builder.getI64IntegerAttr(full_type.i());
break;
case tensorflow::FullTypeDef::ATTR_NOT_SET:
break;
default:
return absl::InvalidArgumentError("Unsupported attr kind in FullType");
}
IntegerAttr type_id_attr =
mlir::IntegerAttr::get(mlir::IntegerType::get(builder.getContext(), 32),
static_cast<int32_t>(full_type.type_id()));
return FullTypeAttr::get(builder.getContext(), type_id_attr, args, attr);
}
absl::StatusOr<tensorflow::FullTypeDef> ConvertAttribute(
tf_type::FullTypeAttr full_type) {
using FullTypeDef = tensorflow::FullTypeDef;
FullTypeDef ret;
for (tf_type::FullTypeAttr it : full_type.getArgs()) {
TF_ASSIGN_OR_RETURN(*ret.add_args(), ConvertAttribute(it));
}
if (full_type.getAttr()) {
bool converted = llvm::TypeSwitch<Attribute, bool>(full_type.getAttr())
.Case<StringAttr>([&](StringAttr sattr) {
ret.set_s(sattr.str());
return true;
})
.Case<IntegerAttr>([&](IntegerAttr iattr) {
ret.set_i(iattr.getInt());
return true;
})
.Default([&](Attribute attr) { return false; });
if (!converted)
return absl::InvalidArgumentError(
absl::StrCat("Unsupported attr kind in FullType:",
mlir::debugString(full_type.getAttr())));
}
ret.set_type_id(
static_cast<tensorflow::FullTypeId>(full_type.getTypeId().getInt()));
return ret;
}
absl::StatusOr<ArrayAttr> ConvertHandleData(
Builder builder,
const tensorflow::protobuf::RepeatedPtrField<
tensorflow::ResourceHandleProto_DtypeAndShape>& handle_data) {
SmallVector<Attribute> dtype_and_shape;
for (const auto& handle : handle_data) {
if (handle.dtype() == tensorflow::DT_INVALID)
return absl::InvalidArgumentError("Invalid dtype for handle_data");
Type dtype;
TF_RETURN_IF_ERROR(ConvertDataType(handle.dtype(), builder, &dtype));
TF_ASSIGN_OR_RETURN(
ShapeAttr shape,
ConvertTensorShapeProto(handle.shape(), builder.getContext()));
TensorType handle_type;
if (shape.hasRank()) {
handle_type = RankedTensorType::get(shape.getShape(), dtype);
} else {
handle_type = UnrankedTensorType::get(dtype);
}
dtype_and_shape.push_back(TypeAttr::get(handle_type));
}
return builder.getArrayAttr(dtype_and_shape);
}
Status ConvertHandleData(ArrayAttr handle_data_arr,
tensorflow::OpDef::ArgDef* arg) {
if (!handle_data_arr) return {};
for (auto handle_data_attr : handle_data_arr.getAsRange<TypeAttr>()) {
TensorType handle_type =
mlir::dyn_cast<TensorType>(handle_data_attr.getValue());
if (!handle_type) {
return absl::InvalidArgumentError(
absl::StrCat("Expected an array of tensor types, but got ",
debugString(handle_data_arr)));
}
auto* handle_data = arg->add_handle_data();
if (handle_type.hasRank()) {
ConvertToTensorShapeProto(handle_type.getShape(),
handle_data->mutable_shape());
} else {
handle_data->mutable_shape()->set_unknown_rank(true);
}
DataType dtype;
TF_RETURN_IF_ERROR(ConvertToDataType(handle_type.getElementType(), &dtype));
handle_data->set_dtype(dtype);
}
return {};
}
} // namespace tfg
} // namespace mlir
tensorflow/core/ir/importexport/convert_attributes.h
+86
View File
@@ -0,0 +1,86 @@
/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_CORE_IR_IMPORTEXPORT_CONVERT_ATTRIBUTES_H_
#define TENSORFLOW_CORE_IR_IMPORTEXPORT_CONVERT_ATTRIBUTES_H_
#include <string>
#include "absl/strings/string_view.h"
#include "mlir/IR/Attributes.h" // from @llvm-project
#include "mlir/IR/Builders.h" // from @llvm-project
#include "mlir/IR/BuiltinTypes.h" // from @llvm-project
#include "tensorflow/core/framework/node_def_util.h"
#include "tensorflow/core/framework/op_def.pb.h"
#include "tensorflow/core/framework/resource_handle.pb.h"
#include "tensorflow/core/ir/dialect.h"
#include "tensorflow/core/platform/statusor.h"
namespace mlir {
namespace tfg {
// Convert the list of MLIR Attributes `attrs` to the `tensorflow::AttrValueMap`
// `values`.
absl::Status ConvertAttributes(ArrayRef<NamedAttribute> attrs,
ArrayRef<StringRef> attrs_to_ignore,
bool remove_ref_type,
tensorflow::AttrValueMap* values);
// Convert the MLIR attribute `attr` and return a `tensorflow::AttrValue`.
absl::StatusOr<tensorflow::AttrValue> ConvertAttribute(Attribute attr);
absl::Status SetShapeAttribute(absl::string_view name, ShapedType shaped_type,
tensorflow::AttrValueMap* values);
// Converts an MLIR shaped type to a TensorFlow shape attribute.
ShapeAttr ConvertTypeToTensorShapeAttr(const Type& type);
/// Import from TensorFlow to MLIR
// Converts non func AttrValue proto into an MLIR attribute. Func attribute is
// exclused in this function because the function might be renamed when the
// function definition is imported.
absl::StatusOr<Attribute> ConvertNonFuncAttributeValue(
const tensorflow::AttrValue& value, Builder& builder);
// Converts all kinds of AttrValue proto into an MLIR attribute.
absl::StatusOr<Attribute> ConvertAttributeValue(
const tensorflow::AttrValue& value, Builder& builder);
// Convert the MLIR FullTyoe attribute `attr` and return a
// `tensorflow::FullTypeDef`.
absl::StatusOr<tensorflow::FullTypeDef> ConvertAttribute(
tf_type::FullTypeAttr full_type);
// Converts fulltype proto to attribute.
absl::StatusOr< ::mlir::tf_type::FullTypeAttr> ConvertAttribute(
const tensorflow::FullTypeDef& full_type, Builder& builder);
// Convert an array of handle data (pairs of data types and shapes) to an array
// attribute of tensor types.
absl::StatusOr<ArrayAttr> ConvertHandleData(
Builder builder,
const tensorflow::protobuf::RepeatedPtrField<
tensorflow::ResourceHandleProto_DtypeAndShape>& handle_data);
// Convert an array of handle data into the `handle_data` field of the provided
// ArgDef. Each entry of the array is expected to be a TensorType.
absl::Status ConvertHandleData(ArrayAttr handle_data_arr,
tensorflow::OpDef::ArgDef* arg);
} // namespace tfg
} // namespace mlir
#endif // TENSORFLOW_CORE_IR_IMPORTEXPORT_CONVERT_ATTRIBUTES_H_
tensorflow/core/ir/importexport/convert_tensor.cc
+571
View File
@@ -0,0 +1,571 @@
/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/core/ir/importexport/convert_tensor.h"
#include <optional>
#include <string>
#include <vector>
#include "llvm/ADT/APFloat.h"
#include "llvm/ADT/ArrayRef.h"
#include "llvm/ADT/SmallVector.h"
#include "mlir/IR/Attributes.h" // from @llvm-project
#include "mlir/IR/Builders.h" // from @llvm-project
#include "mlir/IR/BuiltinTypes.h" // from @llvm-project
#include "mlir/IR/Types.h" // from @llvm-project
#include "mlir/Support/LLVM.h" // from @llvm-project
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/framework/tensor.pb.h"
#include "tensorflow/core/framework/tensor_shape.pb.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/ir/dialect.h"
#include "tensorflow/core/ir/importexport/convert_types.h"
#include "tensorflow/core/ir/importexport/mangling.h"
#include "tensorflow/core/ir/types/dialect.h"
#include "tensorflow/core/lib/core/errors.h"
#include "tensorflow/core/platform/bfloat16.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/platform/protobuf.h"
#include "tensorflow/core/platform/statusor.h"
#include "tensorflow/core/platform/tstring.h"
#include "tsl/platform/ml_dtypes.h"
namespace mlir {
namespace tfg {
using tensorflow::bfloat16;
using tensorflow::PartialTensorShape;
using tensorflow::Status;
using tensorflow::Tensor;
using tensorflow::TensorProto;
using tensorflow::TensorShape;
using tensorflow::TensorShapeProto;
using tensorflow::tstring;
using tensorflow::errors::InvalidArgument;
using tensorflow::errors::Unimplemented;
using tensorflow::port::CopyFromArray;
using tensorflow::protobuf::RepeatedField;
using tensorflow::protobuf::RepeatedPtrField;
static TensorProto ConvertToProto(const Tensor& input_tensor,
bool use_tensor_content = true) {
TensorProto tensor_proto;
// Using tensor content (mostly*) reduces serialization overhead during RPC
// calls, but is less human reader friendly. People reading protobufs are less
// frequent than serialization, so default to using tensor content
// representation.
// * For scalars and short strings it may be marginally worse and a more
// intelligent decision could be made by caller.
if (use_tensor_content)
input_tensor.AsProtoTensorContent(&tensor_proto);
else
input_tensor.AsProtoField(&tensor_proto);
return tensor_proto;
}
static std::string MangleTensor(const Tensor& tensor) {
return mangling_util::MangleTensor(ConvertToProto(tensor));
}
// Converts a TensorFlow tensor into an MLIR elements attribute.
template <typename T>
absl::StatusOr<ElementsAttr> ConvertFlatTensor(const Tensor& input_tensor,
ShapedType type) {
auto arr = input_tensor.flat<T>();
return ElementsAttr(
DenseElementsAttr::get(type, llvm::ArrayRef(arr.data(), arr.size())));
}
ElementsAttr ConvertBf16Tensor(const Tensor& input_tensor,
RankedTensorType type) {
auto buffer = llvm::ArrayRef(static_cast<char*>(input_tensor.data()),
input_tensor.TotalBytes());
return DenseElementsAttr::getFromRawBuffer(type, buffer);
}
ElementsAttr ConvertHalfTensor(const Tensor& tensor, RankedTensorType type) {
auto buffer =
llvm::ArrayRef(static_cast<char*>(tensor.data()), tensor.TotalBytes());
return DenseElementsAttr::getFromRawBuffer(type, buffer);
}
absl::StatusOr<ElementsAttr> ConvertStringTensor(const Tensor& input_tensor,
ShapedType type) {
// Extract to a vector of StringRefs for converting.
auto arr = input_tensor.flat<tstring>();
std::vector<StringRef> string_refs;
string_refs.reserve(arr.size());
for (int i = 0; i < arr.size(); i++) {
const auto& val = arr(i);
string_refs.push_back({val.data(), val.size()});
}
return ElementsAttr(DenseStringElementsAttr::get(type, string_refs));
}
absl::StatusOr<ElementsAttr> ConvertTensor(const Tensor& input_tensor,
Builder builder) {
const auto& input_dtype = input_tensor.dtype();
const auto& input_shape = input_tensor.shape();
Type elt_type;
TF_RETURN_IF_ERROR(ConvertDataType(input_dtype, builder, &elt_type));
SmallVector<int64_t, 4> shape;
ConvertToMlirShape(input_shape, &shape);
auto type = RankedTensorType::get(shape, elt_type);
#define CONVERT_FLAT(DTYPE, CTYPE) \
case tensorflow::DTYPE: \
return ConvertFlatTensor<CTYPE>(input_tensor, type);
// TODO(fengliuai): customize the conversions for quantized types.
switch (input_dtype) {
CONVERT_FLAT(DT_BOOL, bool)
CONVERT_FLAT(DT_FLOAT, float)
CONVERT_FLAT(DT_DOUBLE, double)
CONVERT_FLAT(DT_INT8, int8_t)
CONVERT_FLAT(DT_INT16, int16_t)
CONVERT_FLAT(DT_INT32, int32_t)
CONVERT_FLAT(DT_INT64, int64_t)
CONVERT_FLAT(DT_UINT8, uint8_t)
CONVERT_FLAT(DT_UINT16, uint16_t)
CONVERT_FLAT(DT_UINT32, uint32_t)
CONVERT_FLAT(DT_UINT64, uint64_t)
CONVERT_FLAT(DT_COMPLEX64, std::complex<float>)
CONVERT_FLAT(DT_COMPLEX128, std::complex<double>)
// BFLOAT16 is a special case that it needs to be cast to double type to
// match its storage type.
case tensorflow::DT_BFLOAT16:
return ConvertBf16Tensor(input_tensor, type);
case tensorflow::DT_HALF:
return ConvertHalfTensor(input_tensor, type);
case tensorflow::DT_STRING:
return ConvertStringTensor(input_tensor, type);
default:
// TODO(shpeisman): restructure code to reuse dialect pointer across
// calls.
return ElementsAttr(
tf_type::TensorProtoAttr::get(type, MangleTensor(input_tensor)));
}
#undef CONVERT_FLAT
}
// Returns the number of elements present in this TensorProto, or -1 if that
// could not be determined. This might be less than the shape of the proto might
// indicate, if we're storing a splat tensor.
static int NumberOfMaterializedElements(const TensorProto& tensor) {
if (!tensor.tensor_content().empty()) return -1;
// We don't know which element type this protocol buffer is storing, and the
// metaprogramming facilities for TensorProto are too limited to check their
// number without knowing this, so we need to manually dispatch to each
// possible member of TensorProto, depening on its dtype.
#define MATCH(DTYPE, FIELD) \
case tensorflow::DTYPE: \
return tensor.FIELD##_val().size()
switch (tensor.dtype()) {
MATCH(DT_FLOAT, float);
MATCH(DT_DOUBLE, double);
MATCH(DT_INT8, int);
MATCH(DT_UINT8, int);
MATCH(DT_INT16, int);
MATCH(DT_UINT16, int);
MATCH(DT_INT32, int);
MATCH(DT_UINT32, uint32);
MATCH(DT_INT64, int64);
MATCH(DT_UINT64, uint64);
MATCH(DT_BOOL, bool);
MATCH(DT_HALF, half);
MATCH(DT_BFLOAT16, half);
MATCH(DT_STRING, string);
// TODO(b/188995810): DenseElementsAttr::get doesn't support complex
// Attributes being passed, so we bail out for now. This should just be
// MATCH(DT_COMPLEX64, scomplex) / 2;
// MATCH(DT_COMPLEX128, dcomplex) / 2;
// when DenseElementsAttr is updated.
case tensorflow::DT_COMPLEX64:
case tensorflow::DT_COMPLEX128:
default:
return -1;
}
}
absl::StatusOr<ElementsAttr> ConvertTensorProto(const TensorProto& input_tensor,
Builder builder) {
// If there is only one actual element in the proto, but its shape would
// indicate there are more values, then this is representing a splat tensor.
// We can create an MLIR Attribute more efficiently in this case.
TensorShape input_tensor_shape(input_tensor.tensor_shape());
int num_elt = NumberOfMaterializedElements(input_tensor);
if ((num_elt == 1 ||
(num_elt == 0 && input_tensor.tensor_content().empty())) &&
input_tensor_shape.num_elements() > 1) {
// We first convert this TensorProto to one of shape [1]. We then create an
// Attribute for that proto, and finally splat the Attribute.
TensorProto tensor_copy = input_tensor;
auto* shape = tensor_copy.mutable_tensor_shape();
shape->clear_dim();
shape->add_dim()->set_size(1);
TF_ASSIGN_OR_RETURN(ElementsAttr single_attr,
ConvertTensorProto(tensor_copy, builder));
std::vector<int64_t> original_dimensions;
for (auto dim : input_tensor_shape) original_dimensions.push_back(dim.size);
return ElementsAttr(SplatElementsAttr::get(
single_attr.getShapedType().clone(original_dimensions),
single_attr.getValues<Attribute>()[0]));
}
Tensor t;
if (!t.FromProto(input_tensor)) {
return absl::InvalidArgumentError(absl::StrCat(
"Failed to parse input_tensor: ", input_tensor.DebugString()));
}
return ConvertTensor(t, builder);
}
void ConvertToTensorShapeProto(ArrayRef<int64_t> shape,
TensorShapeProto* output_shape) {
for (auto d : shape) {
output_shape->add_dim()->set_size(d);
}
}
PartialTensorShape ConvertTypeToTensorShape(const Type& type) {
if (mlir::isa<UnrankedTensorType>(type)) {
// An empty PartialTensorShape indicates an unranked tensor.
return PartialTensorShape();
}
if (auto tensor_type = mlir::dyn_cast<RankedTensorType>(type)) {
TensorShapeProto tensor_shape_proto;
ConvertToTensorShapeProto(ConvertMlirShapeToTF(tensor_type.getShape()),
&tensor_shape_proto);
return PartialTensorShape(tensor_shape_proto);
}
// If type is not a RankedTensor or UnrankedTensor, it must be a scalar.
// Empty TensorShape indicates a scalar.
return TensorShape();
}
ShapeAttr ConvertTypeToTensorShapeAttr(const Type& type) {
if (mlir::isa<UnrankedTensorType>(type)) {
return ShapeAttr::get(type.getContext(), std::nullopt);
}
if (auto tensor_type = mlir::dyn_cast<RankedTensorType>(type)) {
return ShapeAttr::get(
type.getContext(),
llvm::ArrayRef(ConvertMlirShapeToTF(tensor_type.getShape())));
}
// If type is not a RankedTensor or UnrankedTensor, it must be a scalar.
// Empty TensorShape indicates a scalar.
return ShapeAttr::get(type.getContext(), ArrayRef<int64_t>());
}
// Converts the tensor shape proto into an MLIR shape attribute.
absl::StatusOr<ShapeAttr> ConvertTensorShapeProto(const TensorShapeProto& shape,
MLIRContext* context) {
if (shape.unknown_rank()) return ShapeAttr::get(context, std::nullopt);
SmallVector<int64_t, 4> dims;
dims.reserve(shape.dim_size());
for (const auto& dim : shape.dim()) {
dims.push_back(dim.size());
}
return ShapeAttr::get(context, llvm::ArrayRef(dims));
}
// Converts an MLIR dense string elements attribute to a TensorFlow tensor
// proto.
void ConvertStringElementsAttr(const DenseStringElementsAttr attr,
RepeatedPtrField<std::string>* output) {
for (const auto& val : attr.getRawStringData())
output->Add({val.data(), val.size()});
}
template <typename T>
void ConvertComplexElementsAttr(const DenseElementsAttr attr,
RepeatedField<T>* output) {
for (const auto& val : attr.getValues<std::complex<T>>()) {
output->Add(val.real());
output->Add(val.imag());
}
}
// Converts an Tensor proto attribute to a TensorFlow tensor proto.
Status ConvertTensorProtoAttr(const mlir::tf_type::TensorProtoAttr attr,
TensorProto* output_tensor) {
auto mangled_tensor = attr.getValue();
absl::string_view tensor_view(mangled_tensor.data(), mangled_tensor.size());
return mangling_util::DemangleTensor(tensor_view, output_tensor);
}
template <typename T>
void ConvertElementsAttr(const DenseElementsAttr attr,
RepeatedField<T>* output) {
if (attr.isSplat()) {
if (attr.getSplatValue<T>() != T(0)) output->Add(attr.getSplatValue<T>());
} else {
output->Reserve(attr.getNumElements());
for (auto value : attr.getValues<T>()) output->AddAlreadyReserved(value);
}
}
// Converts an MLIR elements attribute and adds it to specified repeated field.
template <typename T, typename Cord>
void ConvertFloatElementsAttr(const DenseElementsAttr attr,
RepeatedField<T>* output, Cord* tensor_content) {
if (attr.isSplat()) {
auto value = attr.getSplatValue<T>();
// Emit the value if it isn't 0 (default), but be careful about -0.0.
if (value != T(0) || std::signbit(value))
output->Add(attr.getSplatValue<T>());
} else {
CopyFromArray(tensor_content, attr.getRawData().data(),
attr.getRawData().size());
}
}
// Converts an MLIR elements attribute containing half values and adds it to
// specified repeated field.
void ConvertHalfElementsAttr(const DenseElementsAttr attr,
RepeatedField<int>* output) {
// Half values are stored as bit representations in int, requiring a bit_cast.
if (attr.isSplat()) {
uint16_t bits =
Eigen::numext::bit_cast<uint16_t>(attr.getSplatValue<Eigen::half>());
// Only +0 has a 0 bit representation.
if (bits != 0) {
output->Add(bits);
}
} else {
output->Reserve(attr.getNumElements());
for (const Eigen::half value : attr.getValues<Eigen::half>()) {
output->AddAlreadyReserved(Eigen::numext::bit_cast<uint16_t>(value));
}
}
}
// Converts an MLIR elements attribute containing signed int values and adds it
// to specified repeated field.
template <typename T, typename U = T, typename Cord>
void ConvertIntElementsAttr(const DenseElementsAttr attr,
RepeatedField<T>* output, Cord* tensor_content) {
if (attr.isSplat()) {
if (attr.getSplatValue<U>() != U(0))
output->Add(static_cast<T>(attr.getSplatValue<U>()));
} else {
CopyFromArray(tensor_content, attr.getRawData().data(),
attr.getRawData().size());
}
}
// Converts an MLIR elements attribute containing unsigned int values and adds
// it to specified repeated field.
template <typename T, typename U = T, typename Cord>
void ConvertUIntElementsAttr(const DenseElementsAttr attr,
RepeatedField<T>* output, Cord* tensor_content) {
if (attr.isSplat()) {
if (attr.getSplatValue<U>() != U(0))
output->Add(static_cast<T>(attr.getSplatValue<U>()));
} else {
CopyFromArray(tensor_content, attr.getRawData().data(),
attr.getRawData().size());
}
}
void ConvertBfloat16ElementsAttr(const DenseElementsAttr attr,
RepeatedField<int>* output) {
// Bfloat16 values are stored as bit representations in int, requiring a
// bit_cast.
if (attr.isSplat()) {
uint16_t bits =
Eigen::numext::bit_cast<uint16_t>(attr.getSplatValue<bfloat16>());
// Only +0 has a 0 bit representation.
if (bits != 0) {
output->Add(bits);
}
} else {
output->Reserve(attr.getNumElements());
for (const bfloat16 value : attr.getValues<bfloat16>()) {
output->AddAlreadyReserved(Eigen::numext::bit_cast<uint16_t>(value));
}
}
}
template <typename T>
void ConvertFloat8ElementsAttr(const DenseElementsAttr attr,
std::string* output) {
// Float8 values are stored as bit representations in int, requiring a
// bit_cast.
if (attr.isSplat()) {
uint8_t bits = Eigen::numext::bit_cast<uint8_t>(attr.getSplatValue<T>());
// Only +0 has a 0 bit representation.
if (bits != 0) {
output->push_back(bits);
}
} else {
output->reserve(attr.getNumElements());
for (const T value : attr.getValues<T>()) {
output->push_back(Eigen::numext::bit_cast<uint8_t>(value));
}
}
}
Status ConvertToTensorProto(const ElementsAttr attr, TensorProto* output) {
auto type = attr.getShapedType();
auto shape = type.getShape();
tensorflow::DataType output_dtype;
TF_RETURN_IF_ERROR(ConvertToDataType(type, &output_dtype));
output->set_dtype(output_dtype);
ConvertToTensorShapeProto(shape, output->mutable_tensor_shape());
if (auto tensor_attr = mlir::dyn_cast<mlir::tf_type::TensorProtoAttr>(attr))
return ConvertTensorProtoAttr(tensor_attr, output);
auto dense_attr = mlir::dyn_cast<DenseElementsAttr>(attr);
if (!dense_attr)
return absl::InvalidArgumentError("Unsupported elements attr");
switch (output_dtype) {
case tensorflow::DT_BOOL:
ConvertElementsAttr(dense_attr, output->mutable_bool_val());
break;
case tensorflow::DT_BFLOAT16:
ConvertBfloat16ElementsAttr(dense_attr, output->mutable_half_val());
break;
case tensorflow::DT_COMPLEX64:
ConvertComplexElementsAttr(dense_attr, output->mutable_scomplex_val());
break;
case tensorflow::DT_COMPLEX128:
ConvertComplexElementsAttr(dense_attr, output->mutable_dcomplex_val());
break;
case tensorflow::DT_DOUBLE:
ConvertFloatElementsAttr(dense_attr, output->mutable_double_val(),
output->mutable_tensor_content());
break;
case tensorflow::DT_HALF:
ConvertHalfElementsAttr(dense_attr, output->mutable_half_val());
break;
case tensorflow::DT_FLOAT:
ConvertFloatElementsAttr(dense_attr, output->mutable_float_val(),
output->mutable_tensor_content());
break;
case tensorflow::DT_FLOAT8_E5M2:
ConvertFloat8ElementsAttr<tsl::float8_e5m2>(dense_attr,
output->mutable_float8_val());
break;
case tensorflow::DT_FLOAT8_E4M3FN:
ConvertFloat8ElementsAttr<tsl::float8_e4m3fn>(
dense_attr, output->mutable_float8_val());
break;
case tensorflow::DT_INT4:
ConvertIntElementsAttr<int, tsl::int4>(dense_attr,
output->mutable_int_val(),
output->mutable_tensor_content());
break;
case tensorflow::DT_UINT4:
ConvertUIntElementsAttr<int, tsl::uint4>(
dense_attr, output->mutable_int_val(),
output->mutable_tensor_content());
break;
case tensorflow::DT_QUINT8:
case tensorflow::DT_INT8:
ConvertUIntElementsAttr<int, int8_t>(dense_attr,
output->mutable_int_val(),
output->mutable_tensor_content());
break;
case tensorflow::DT_QUINT16:
case tensorflow::DT_INT16:
ConvertIntElementsAttr<int, int16_t>(dense_attr,
output->mutable_int_val(),
output->mutable_tensor_content());
break;
case tensorflow::DT_INT32:
ConvertIntElementsAttr(dense_attr, output->mutable_int_val(),
output->mutable_tensor_content());
break;
case tensorflow::DT_INT64:
ConvertIntElementsAttr(dense_attr, output->mutable_int64_val(),
output->mutable_tensor_content());
break;
case tensorflow::DT_STRING:
ConvertStringElementsAttr(mlir::cast<DenseStringElementsAttr>(dense_attr),
output->mutable_string_val());
break;
case tensorflow::DT_UINT8:
ConvertUIntElementsAttr<int, uint8_t>(dense_attr,
output->mutable_int_val(),
output->mutable_tensor_content());
break;
case tensorflow::DT_UINT16:
ConvertUIntElementsAttr<int, uint16_t>(dense_attr,
output->mutable_int_val(),
output->mutable_tensor_content());
break;
case tensorflow::DT_UINT32:
ConvertUIntElementsAttr(dense_attr, output->mutable_uint32_val(),
output->mutable_tensor_content());
break;
case tensorflow::DT_UINT64:
ConvertUIntElementsAttr(dense_attr, output->mutable_uint64_val(),
output->mutable_tensor_content());
break;
default:
return absl::UnimplementedError(absl::StrCat(
"Unimplemented data type ", DataTypeString(output_dtype)));
}
return absl::OkStatus();
}
Status ConvertToTensor(const ElementsAttr attr, Tensor* output_tensor) {
TensorProto tensor_proto;
TF_RETURN_IF_ERROR(ConvertToTensorProto(attr, &tensor_proto));
if (!output_tensor->FromProto(tensor_proto)) {
return absl::InvalidArgumentError(
"Couldn't convert tensor proto to tensor.");
}
return absl::OkStatus();
}
llvm::SmallVector<int64_t> ConvertMlirShapeToTF(llvm::ArrayRef<int64_t> shape) {
return llvm::to_vector(llvm::map_range(shape, [](int64_t dim) {
return mlir::ShapedType::isDynamic(dim) ? -1 : dim;
}));
}
llvm::SmallVector<int64_t> ConvertTFShapeToMlir(llvm::ArrayRef<int64_t> shape) {
return llvm::to_vector(llvm::map_range(shape, [](int64_t dim) {
return dim == -1 ? mlir::ShapedType::kDynamic : dim;
}));
}
mlir::RankedTensorType GetTypeFromTFTensorShape(llvm::ArrayRef<int64_t> shape,
mlir::Type elementType,
mlir::Attribute encoding) {
return mlir::RankedTensorType::get(ConvertTFShapeToMlir(shape), elementType,
encoding);
}
} // namespace tfg
} // namespace mlir
tensorflow/core/ir/importexport/convert_tensor.h
+93
View File
@@ -0,0 +1,93 @@
/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_CORE_IR_IMPORTEXPORT_CONVERT_TENSOR_H_
#define TENSORFLOW_CORE_IR_IMPORTEXPORT_CONVERT_TENSOR_H_
#include "llvm/ADT/ArrayRef.h"
#include "llvm/ADT/SmallVector.h"
#include "mlir/IR/Attributes.h" // from @llvm-project
#include "mlir/IR/Builders.h" // from @llvm-project
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/framework/tensor.pb.h"
#include "tensorflow/core/framework/tensor_shape.pb.h"
#include "tensorflow/core/ir/dialect.h"
#include "tensorflow/core/ir/types/dialect.h"
#include "tensorflow/core/platform/statusor.h"
namespace mlir {
namespace tfg {
// Converts an TensorFlow tensor proto into an MLIR elements attribute.
absl::StatusOr<ElementsAttr> ConvertTensorProto(
const tensorflow::TensorProto& input_tensor, Builder builder);
// Converts an TensorFlow tensor into an MLIR elements attribute.
absl::StatusOr<ElementsAttr> ConvertTensor(
const tensorflow::Tensor& input_tensor, Builder builder);
// Converts a shape from MLIR to a TensorFlow tensor shape proto.
void ConvertToTensorShapeProto(ArrayRef<int64_t> shape,
tensorflow::TensorShapeProto* output_shape);
// Converts an MLIR type to a TensorFlow tensor shape.
tensorflow::PartialTensorShape ConvertTypeToTensorShape(const Type& type);
// Converts a TensorFlow shape attribute to an MLIR shape attribute.
absl::StatusOr<ShapeAttr> ConvertTensorShapeProto(
const tensorflow::TensorShapeProto& shape, MLIRContext* context);
// Fill in the contents of TensorShapeProto for the given shape.
// ShapeContainerT is any type with the following methods:
// bool hasRank()
// ArrayRef<int64_t> getShape()
// This includes TF::ShapeAttr and ShapedType.
template <typename ShapeContainerT>
void SetTensorShapeProto(ShapeContainerT shape,
tensorflow::TensorShapeProto* proto) {
if (shape.hasRank()) {
for (int64_t dim : shape.getShape()) {
// TODO(hinsu): Use tensorflow::kTFDynamicSize instead of -1 without
// depending on tensorflow/compiler
proto->add_dim()->set_size(mlir::ShapedType::isDynamic(dim) ? -1 : dim);
}
} else {
proto->set_unknown_rank(true);
}
}
// Converts an MLIR elements attribute to a TensorFlow tensor proto.
absl::Status ConvertToTensorProto(ElementsAttr attr,
tensorflow::TensorProto* output_tensor);
// Converts an MLIR elements attribute to a TensorFlow tensor.
absl::Status ConvertToTensor(ElementsAttr attr,
tensorflow::Tensor* output_tensor);
// Converts a TF shape to MLIR shape, i.e. -1 becomes kDynamicSize.
llvm::SmallVector<int64_t> ConvertTFShapeToMlir(llvm::ArrayRef<int64_t> shape);
// Converts an MLIR shape to TF shape, i.e. kDynamicSize becomes -1.
llvm::SmallVector<int64_t> ConvertMlirShapeToTF(llvm::ArrayRef<int64_t> shape);
// Creates a TF TensorShape using MLIR shape, element type and encoding.
mlir::RankedTensorType GetTypeFromTFTensorShape(llvm::ArrayRef<int64_t> shape,
mlir::Type elementType,
mlir::Attribute encoding = {});
} // namespace tfg
} // namespace mlir
#endif // TENSORFLOW_CORE_IR_IMPORTEXPORT_CONVERT_TENSOR_H_
tensorflow/core/ir/importexport/convert_types.cc
+244
View File
@@ -0,0 +1,244 @@
/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/core/ir/importexport/convert_types.h"
#include <limits>
#include "absl/strings/str_cat.h"
#include "llvm/Support/Casting.h"
#include "mlir/IR/BuiltinTypes.h" // from @llvm-project
#include "mlir/IR/Types.h" // from @llvm-project
#include "mlir/Support/DebugStringHelper.h" // from @llvm-project
#include "mlir/Support/LLVM.h" // from @llvm-project
#include "tensorflow/core/framework/types.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/ir/dialect.h"
#include "tensorflow/core/lib/core/errors.h"
namespace mlir {
namespace tfg {
using tensorflow::DataType;
using tensorflow::Status;
using tensorflow::TensorShape;
using tensorflow::TensorShapeProto;
using tensorflow::errors::InvalidArgument;
using tensorflow::errors::Unimplemented;
Status ConvertDataType(DataType dtype, Builder& builder, Type* type) {
switch (dtype) {
case tensorflow::DT_HALF:
*type = builder.getF16Type();
return absl::OkStatus();
case tensorflow::DT_FLOAT:
*type = builder.getF32Type();
return absl::OkStatus();
case tensorflow::DT_DOUBLE:
*type = builder.getF64Type();
return absl::OkStatus();
case tensorflow::DT_BOOL:
*type = builder.getIntegerType(1);
return absl::OkStatus();
case tensorflow::DT_INT8:
*type = builder.getIntegerType(8);
return absl::OkStatus();
case tensorflow::DT_INT16:
*type = builder.getIntegerType(16);
return absl::OkStatus();
case tensorflow::DT_INT32:
*type = builder.getIntegerType(32);
return absl::OkStatus();
case tensorflow::DT_INT64:
*type = builder.getIntegerType(64);
return absl::OkStatus();
case tensorflow::DT_UINT8:
*type = builder.getIntegerType(8, /*isSigned=*/false);
return absl::OkStatus();
case tensorflow::DT_UINT16:
*type = builder.getIntegerType(16, /*isSigned=*/false);
return absl::OkStatus();
case tensorflow::DT_UINT32:
*type = builder.getIntegerType(32, /*isSigned=*/false);
return absl::OkStatus();
case tensorflow::DT_UINT64:
*type = builder.getIntegerType(64, /*isSigned=*/false);
return absl::OkStatus();
case tensorflow::DT_BFLOAT16:
*type = builder.getBF16Type();
return absl::OkStatus();
case tensorflow::DT_COMPLEX64:
*type = ComplexType::get(builder.getF32Type());
return absl::OkStatus();
case tensorflow::DT_COMPLEX128:
*type = ComplexType::get(builder.getF64Type());
return absl::OkStatus();
case tensorflow::DT_FLOAT8_E4M3FN:
*type = builder.getType<Float8E4M3FNType>();
return absl::OkStatus();
case tensorflow::DT_FLOAT8_E5M2:
*type = builder.getType<Float8E5M2Type>();
return absl::OkStatus();
case tensorflow::DT_INT4:
*type = builder.getIntegerType(4, /*isSigned=*/true);
return absl::OkStatus();
case tensorflow::DT_UINT4:
*type = builder.getIntegerType(4, /*isSigned=*/false);
return absl::OkStatus();
case tensorflow::DT_INT2:
*type = builder.getIntegerType(2, /*isSigned=*/true);
return absl::OkStatus();
case tensorflow::DT_UINT2:
*type = builder.getIntegerType(2, /*isSigned=*/false);
return absl::OkStatus();
#define HANDLE_TF_TYPE(tftype, enumerant, name) \
case tensorflow::DT_##enumerant: \
*type = builder.getType<tftype##Type>(); \
return ::tensorflow::OkStatus();
#include "tensorflow/core/ir/types/types.def"
default:
return absl::UnimplementedError(absl::StrCat(
"Converting DataType '", DataTypeString(dtype), "' to MLIR Type"));
}
}
Status ConvertScalarTypeToDataType(Type type, DataType* dtype) {
if (type.isF16()) {
*dtype = tensorflow::DT_HALF;
return absl::OkStatus();
} else if (type.isF32()) {
*dtype = tensorflow::DT_FLOAT;
return absl::OkStatus();
} else if (type.isF64()) {
*dtype = tensorflow::DT_DOUBLE;
return absl::OkStatus();
} else if (type.isBF16()) {
*dtype = tensorflow::DT_BFLOAT16;
return absl::OkStatus();
} else if (llvm::isa<Float8E4M3FNType>(type)) {
*dtype = ::tensorflow::DT_FLOAT8_E4M3FN;
return absl::OkStatus();
} else if (llvm::isa<Float8E5M2FNUZType>(type)) {
*dtype = ::tensorflow::DT_FLOAT8_E5M2;
return absl::OkStatus();
} else if (auto itype = mlir::dyn_cast<IntegerType>(type)) {
switch (itype.getWidth()) {
case 1:
*dtype = tensorflow::DT_BOOL;
return absl::OkStatus();
case 2:
*dtype =
itype.isUnsigned() ? tensorflow::DT_UINT2 : tensorflow::DT_INT2;
return absl::OkStatus();
case 4:
*dtype =
itype.isUnsigned() ? tensorflow::DT_UINT4 : tensorflow::DT_INT4;
return absl::OkStatus();
case 8:
*dtype =
itype.isUnsigned() ? tensorflow::DT_UINT8 : tensorflow::DT_INT8;
return absl::OkStatus();
case 16:
*dtype =
itype.isUnsigned() ? tensorflow::DT_UINT16 : tensorflow::DT_INT16;
return absl::OkStatus();
case 32:
*dtype =
itype.isUnsigned() ? tensorflow::DT_UINT32 : tensorflow::DT_INT32;
return absl::OkStatus();
case 64:
*dtype =
itype.isUnsigned() ? tensorflow::DT_UINT64 : tensorflow::DT_INT64;
return absl::OkStatus();
default:
return absl::UnimplementedError(
absl::StrCat("Converting ", debugString(type), " to DataType"));
}
} else if (auto complex_type = mlir::dyn_cast<ComplexType>(type)) {
auto etype = complex_type.getElementType();
if (etype.isF32()) {
*dtype = tensorflow::DT_COMPLEX64;
return absl::OkStatus();
} else if (etype.isF64()) {
*dtype = tensorflow::DT_COMPLEX128;
return absl::OkStatus();
}
return absl::UnimplementedError(
absl::StrCat("Converting ", debugString(type), " to DataType"));
}
#define HANDLE_TF_TYPE(tftype, enumerant, name) \
if (llvm::isa<tftype##Type>(type)) { \
*dtype = tensorflow::DT_##enumerant; \
return ::tensorflow::OkStatus(); \
}
// NOLINTNEXTLINE
#include "tensorflow/core/ir/types/types.def"
return absl::UnimplementedError(
absl::StrCat("Converting ", debugString(type), " to DataType"));
}
Status ConvertToDataType(Type type, DataType* dtype) {
if (auto stype = mlir::dyn_cast<ShapedType>(type)) {
TF_RETURN_IF_ERROR(
ConvertScalarTypeToDataType(stype.getElementType(), dtype));
} else {
TF_RETURN_IF_ERROR(ConvertScalarTypeToDataType(type, dtype));
}
return absl::OkStatus();
}
void ConvertToMlirShape(const TensorShape& input_shape,
SmallVectorImpl<int64_t>* shape) {
shape->reserve(input_shape.dims());
for (const auto& d : input_shape) {
shape->push_back(d.size);
}
}
Status ConvertToMlirShape(const TensorShapeProto& input_shape,
SmallVectorImpl<int64_t>* shape) {
shape->reserve(input_shape.dim_size());
auto& dims = input_shape.dim();
for (auto& d : dims) {
if (d.size() > std::numeric_limits<int64_t>::max()) {
return absl::InvalidArgumentError("Shape element overflows");
}
// This isn't really expected, but Grappler is using such shapes for its
// symbolic shape analysis and it may spill into here.
if (d.size() < ShapedType::kDynamic)
shape->push_back(ShapedType::kDynamic);
else
shape->push_back(d.size());
}
return absl::OkStatus();
}
absl::StatusOr<Type> ConvertToMlirTensorType(const TensorShapeProto& shape,
DataType dtype, Builder* builder) {
Type element_type;
TF_RETURN_IF_ERROR(ConvertDataType(dtype, *builder, &element_type));
if (shape.unknown_rank()) {
return UnrankedTensorType::get(element_type);
}
SmallVector<int64_t, 4> shape_dims;
TF_RETURN_IF_ERROR(ConvertToMlirShape(shape, &shape_dims));
return RankedTensorType::get(shape_dims, element_type);
}
} // namespace tfg
} // namespace mlir
tensorflow/core/ir/importexport/convert_types.h
+56
View File
@@ -0,0 +1,56 @@
/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_CORE_IR_IMPORTEXPORT_CONVERT_TYPES_H_
#define TENSORFLOW_CORE_IR_IMPORTEXPORT_CONVERT_TYPES_H_
#include "mlir/IR/Builders.h" // from @llvm-project
#include "mlir/IR/Types.h" // from @llvm-project
#include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/framework/tensor_shape.pb.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/platform/statusor.h"
namespace mlir {
namespace tfg {
// Converts the TensorFlow DataType 'dtype' into an MLIR (scalar) type.
absl::Status ConvertDataType(tensorflow::DataType dtype, Builder& builder,
Type* type);
// Converts a scalar MLIR type to a TensorFlow Datatype.
absl::Status ConvertScalarTypeToDataType(Type type,
tensorflow::DataType* dtype);
// Converts an MLIR type to TensorFlow DataType. If 'type' is a scalar type, it
// is converted directly. If it is a shaped type, the element type is converted.
absl::Status ConvertToDataType(Type type, tensorflow::DataType* dtype);
// Converts an TensorFlow shape to the one used in MLIR.
void ConvertToMlirShape(const tensorflow::TensorShape& input_shape,
SmallVectorImpl<int64_t>* shape);
// Converts an TensorFlow shape proto to the one used in MLIR.
absl::Status ConvertToMlirShape(const tensorflow::TensorShapeProto& input_shape,
SmallVectorImpl<int64_t>* shape);
// Given a tensor shape and dtype, get the corresponding MLIR tensor type.
absl::StatusOr<Type> ConvertToMlirTensorType(
const tensorflow::TensorShapeProto& shape, tensorflow::DataType dtype,
Builder* builder);
} // namespace tfg
} // namespace mlir
#endif // TENSORFLOW_CORE_IR_IMPORTEXPORT_CONVERT_TYPES_H_
tensorflow/core/ir/importexport/functiondef_export.cc
+321
View File
@@ -0,0 +1,321 @@
/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/core/ir/importexport/functiondef_export.h"
#include <string>
#include <utility>
#include "mlir/IR/BuiltinAttributes.h" // from @llvm-project
#include "mlir/IR/BuiltinTypes.h" // from @llvm-project
#include "mlir/IR/OperationSupport.h" // from @llvm-project
#include "mlir/Support/LLVM.h" // from @llvm-project
#include "tensorflow/core/framework/attr_value.pb.h"
#include "tensorflow/core/framework/node_def.pb.h"
#include "tensorflow/core/framework/op_def.pb.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/ir/dialect.h"
#include "tensorflow/core/ir/importexport/convert_attributes.h"
#include "tensorflow/core/ir/importexport/convert_types.h"
#include "tensorflow/core/ir/importexport/graphdef_export.h"
#include "tensorflow/core/ir/ops.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/platform/statusor.h"
using tensorflow::FunctionDef;
using tensorflow::OpDef;
using tensorflow::OpDef_AttrDef;
using tensorflow::Status;
using tensorflow::errors::InvalidArgument;
#define DEBUG_TYPE "mlir-to-graphdef"
namespace mlir {
namespace tfg {
// Compute the name to use in FunctionDef for a given Value (either the result
// of an operation or a block operand if a function argument) and store the
// result in the provided name string. The `control_ty` is the instance of the
// `ControlType` to compare against and detect a control dependency case.
static absl::StatusOr<std::string> GetValueName(Value operand,
Type control_ty) {
bool is_control = (operand.getType() == control_ty);
OpResult op_result = mlir::dyn_cast<OpResult>(operand);
if (!op_result) {
BlockArgument block_operand = mlir::dyn_cast<BlockArgument>(operand);
int arg_num = block_operand.getArgNumber();
// Function arguments are coming as pair: the even are the actual tensors
// while the odd position are the associated control input.
std::string name;
if (is_control) name = "^";
DictionaryAttr arg_attrs = function_interface_impl::getArgAttrDict(
FunctionOpInterface(block_operand.getParentBlock()->getParentOp()),
arg_num - is_control);
if (!arg_attrs)
return absl::InvalidArgumentError(
absl::StrCat("Missing attribute for argument #", arg_num));
StringAttr arg_name = arg_attrs.getAs<StringAttr>("tfg.name");
if (!arg_name)
return absl::InvalidArgumentError(absl::StrCat(
"Can't export graph with missing op-name for function parameter #",
arg_num));
absl::StrAppend(&name, arg_name.getValue().str());
return name;
}
GetResultOp get_result = op_result.getDefiningOp<GetResultOp>();
Operation *producer;
if (is_control) {
producer = op_result.getDefiningOp();
} else {
if (!get_result)
return absl::InvalidArgumentError(
"Missing get_result operation as input");
producer = get_result.getValue().getDefiningOp();
if (!producer)
return absl::InvalidArgumentError(
"Expect a tfg operation as input to GetResultOp");
}
auto name_attr =
producer->getAttrOfType<StringAttr>(TFGraphDialect::getNameAttrKey());
if (!name_attr)
return absl::InvalidArgumentError(
"Can't export graph with missing op-name");
std::string name;
if (is_control) name = "^";
absl::StrAppend(&name, name_attr.getValue().str());
if (get_result)
absl::StrAppend(&name, ":", get_result.getName().str(), ":",
get_result.getNumber());
return name;
}
// Export a function argument or returned value as an ArgDef entry.
// If arg_def_attrs is provided, it is populated with the extra attributes
// converted from MLIR to AttrValue proto representation. This is useful only
// for Function arguments to populate the `arg_attr` field.
//
static Status ExportArgDef(OpDef::ArgDef *arg, DictionaryAttr arg_attrs,
FunctionDef::ArgAttrs *arg_def_attrs = nullptr) {
StringAttr arg_name = arg_attrs.getAs<StringAttr>("tfg.name");
if (!arg_name)
return absl::InvalidArgumentError("Missing \"tfg.name\" attribute");
arg->set_name(arg_name.getValue().str());
StringAttr description = arg_attrs.getAs<StringAttr>("tfg.description");
if (description) arg->set_description(description.getValue().str());
TypeAttr input_type = arg_attrs.getAs<TypeAttr>("tfg.type");
if (input_type) {
tensorflow::DataType dtype;
TF_RETURN_IF_ERROR(ConvertToDataType(input_type.getValue(), &dtype));
arg->set_type(dtype);
}
if (StringAttr type_attr = arg_attrs.getAs<StringAttr>("tfg.type_attr"))
arg->set_type_attr(type_attr.getValue().str());
if (StringAttr number_attr = arg_attrs.getAs<StringAttr>("tfg.number_attr"))
arg->set_number_attr(number_attr.getValue().str());
if (StringAttr type_list_attr =
arg_attrs.getAs<StringAttr>("tfg.type_list_attr"))
arg->set_type_attr(type_list_attr.getValue().str());
if (auto full_type = arg_attrs.getAs<tf_type::FullTypeAttr>(
"tfg.experimental_full_type")) {
TF_ASSIGN_OR_RETURN(*arg->mutable_experimental_full_type(),
ConvertAttribute(full_type));
}
TF_RETURN_IF_ERROR(
ConvertHandleData(arg_attrs.getAs<ArrayAttr>("tfg.handle_data"), arg));
if (UnitAttr number_attr = arg_attrs.getAs<UnitAttr>("tfg.is_ref"))
arg->set_is_ref(true);
auto sig_arg_attrs = arg_attrs.getAs<DictionaryAttr>("tfg.arg_attrs");
if (arg_def_attrs && sig_arg_attrs) {
TF_RETURN_IF_ERROR(ConvertAttributes(
sig_arg_attrs.getValue(), /*attrs_to_ignore=*/{},
/*remove_ref_type=*/false, arg_def_attrs->mutable_attr()));
}
return absl::OkStatus();
}
absl::StatusOr<FunctionDef> ConvertGenericFunctionToFunctionDef(
GraphFuncOp func_op) {
if (!func_op.getGeneric())
return absl::InvalidArgumentError(
"Expected a generic function in ConvertGenericFunctionToFunctionDef");
auto control_ty = tfg::ControlType::get(func_op.getContext());
auto *tfg_dialect = cast<TFGraphDialect>(func_op->getDialect());
FunctionDef fdef;
for (Operation &op : func_op.SingleBlock::getBody()->without_terminator()) {
if (op.getDialect() != tfg_dialect)
return absl::InvalidArgumentError(
"Non tfg op encountered when exporting function");
if (isa<GetResultOp>(&op)) continue;
TF_RETURN_IF_ERROR(ConvertToNodeDef(
&op, fdef.add_node_def(), tfg_dialect,
[&](Value value) { return GetValueName(value, control_ty); }));
}
const std::string func_name = func_op.getName().str();
OpDef *signature = fdef.mutable_signature();
signature->set_name(func_name);
if (func_op->getAttr("is_stateful")) signature->set_is_stateful(true);
if (auto description = func_op->getAttrOfType<StringAttr>("description"))
signature->set_description(description.getValue().str());
if (auto attrs = func_op->getAttrOfType<DictionaryAttr>("tfg.func_attrs")) {
for (NamedAttribute attr : attrs) {
OpDef_AttrDef *func_attr = signature->add_attr();
func_attr->set_name(attr.getName().str());
DictionaryAttr dict_attr =
mlir::dyn_cast<DictionaryAttr>(attr.getValue());
if (!dict_attr)
return absl::InvalidArgumentError("Expects dict attribute");
if (StringAttr type = dict_attr.getAs<StringAttr>("function_type"))
func_attr->set_type(type.getValue().str());
if (Attribute default_value = dict_attr.get("default_value")) {
TF_ASSIGN_OR_RETURN((*func_attr->mutable_default_value()),
ConvertAttribute(default_value));
}
if (StringAttr description = dict_attr.getAs<StringAttr>("description"))
func_attr->set_description(description.getValue().str());
if (IntegerAttr minimum = dict_attr.getAs<IntegerAttr>("minimum")) {
func_attr->set_minimum(minimum.getInt());
func_attr->set_has_minimum(true);
}
if (Attribute allowed_values = dict_attr.get("allowed_values")) {
TF_ASSIGN_OR_RETURN((*func_attr->mutable_allowed_values()),
ConvertAttribute(allowed_values));
}
}
}
if (auto control_outputs =
func_op->getAttrOfType<ArrayAttr>("control_output")) {
for (Attribute attr : control_outputs) {
StringAttr output = mlir::dyn_cast<StringAttr>(attr);
if (!output)
return absl::InvalidArgumentError(
"Can't export function with non-string \"control_output\" "
"attribute entry");
signature->add_control_output(output.getValue().str());
}
}
// Convert the function argument into an OpDef::ArgDef in the signature.
ArrayAttr args_attr = func_op.getAllArgAttrs();
for (int arg_num : llvm::seq<int>(0, func_op.getNumArguments())) {
// Odd position are just for control dependencies.
if (arg_num % 2) continue;
OpDef::ArgDef *arg = signature->add_input_arg();
if (arg_num >= args_attr.size())
return absl::InvalidArgumentError(
absl::StrCat("Can't export function ", func_op.getName().str(),
" because missing attributes for arg #", arg_num));
DictionaryAttr arg_attrs = mlir::cast<DictionaryAttr>(args_attr[arg_num]);
FunctionDef::ArgAttrs func_def_arg_attrs;
TF_RETURN_WITH_CONTEXT_IF_ERROR(
ExportArgDef(arg, arg_attrs, &func_def_arg_attrs),
" when exporting argument ", arg_num, " for function ",
func_op.getName().str());
// On top of the signature, function arguments can have attribute directul
// on the FunctionDef.
if (!func_def_arg_attrs.attr().empty())
(*fdef.mutable_arg_attr())[arg_num / 2] = std::move(func_def_arg_attrs);
}
// Handle the results now.
// An ArgDef entry needs to be constructed for all non-control returned value,
// and a mapping from the output name to the signature is also recorded in the
// FunctionDef.
auto return_op = llvm::cast<tfg::ReturnOp>(
func_op.SingleBlock::getBody()->getTerminator());
ArrayAttr results_attr = func_op.getAllResultAttrs();
for (const auto &indexed_result : llvm::enumerate(return_op->getOperands())) {
int res_num = indexed_result.index();
if (res_num >= results_attr.size())
return absl::InvalidArgumentError(
absl::StrCat("Can't export function ", func_op.getName().str(),
" because missing attributes for result #", res_num));
auto res_attrs = mlir::cast<DictionaryAttr>(results_attr[res_num]);
auto name = res_attrs.getAs<StringAttr>("tfg.name");
if (!name)
return absl::InvalidArgumentError(absl::StrCat(
"Can't export function ", func_op.getName().str(),
" because missing \"tfg.name\" attribute for result #", res_num));
Value ret_val = indexed_result.value();
if (ret_val.getType() == control_ty) {
// When we return a control dependency, it is not really a returned value
// but it is added to the `control_ret` field of the FunctionDef.
TF_ASSIGN_OR_RETURN(std::string ret_name,
GetValueName(ret_val, control_ty));
fdef.mutable_control_ret()->insert(
{name.getValue().str(), StringRef(ret_name).drop_front().str()});
continue;
}
// Tensor results are turned into an ArgDef in the `output_arg` field.
OpDef::ArgDef *output = signature->add_output_arg();
TF_RETURN_WITH_CONTEXT_IF_ERROR(ExportArgDef(output, res_attrs),
" when exporting result ", res_num,
" for function ", func_op.getName().str());
// The `ret` field of the FunctionDef keeps a mapping of the returned value
// name to the entried in the FunctionDef signature.
TF_ASSIGN_OR_RETURN(std::string ret_name,
GetValueName(ret_val, control_ty));
fdef.mutable_ret()->insert({name.getValue().str(), ret_name});
}
// Handled the `resource_arg_unique_id` entries. At the moment it is
// represented as two vectors of integers which are expected of the same
// length.
auto unique_ids_keys = func_op->getAttrOfType<DenseIntElementsAttr>(
"resource_arg_unique_ids_keys");
if (unique_ids_keys) {
auto unique_ids_values = func_op->getAttrOfType<DenseIntElementsAttr>(
"resource_arg_unique_ids_values");
if (!unique_ids_values)
return absl::InvalidArgumentError(absl::StrCat(
"Can't export function ", func_name,
" because \"resource_arg_unique_ids_keys\" attribute is present "
"but "
"\"resource_arg_unique_ids_values\" is missing"));
if (unique_ids_keys.size() != unique_ids_values.size())
return absl::InvalidArgumentError(absl::StrCat(
"Can't export function ", func_name,
" because \"resource_arg_unique_ids_keys\" array does not have the "
"same size as \"resource_arg_unique_ids_values\""));
auto *unique_ids_map = fdef.mutable_resource_arg_unique_id();
for (auto key_value : llvm::zip(unique_ids_keys.getValues<int32_t>(),
unique_ids_values.getValues<int32_t>()))
(*unique_ids_map)[std::get<0>(key_value)] = std::get<1>(key_value);
}
// Finally the dialect attributes (prefixed by `tf.` in general) are converted
// as-is and stored on the `attr` field of the FunctionDef.
SmallVector<NamedAttribute, 8> funcAttrs(func_op->getDialectAttrs());
TF_RETURN_IF_ERROR(ConvertAttributes(funcAttrs, {"tfg.func_attrs"},
/*remove_ref_type=*/false,
fdef.mutable_attr()));
return fdef;
}
} // namespace tfg
} // namespace mlir
tensorflow/core/ir/importexport/functiondef_export.h
+35
View File
@@ -0,0 +1,35 @@
/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_CORE_IR_IMPORTEXPORT_FUNCTIONDEF_EXPORT_H_
#define TENSORFLOW_CORE_IR_IMPORTEXPORT_FUNCTIONDEF_EXPORT_H_
#include "mlir/IR/Builders.h" // from @llvm-project
#include "tensorflow/core/framework/function.pb.h"
#include "tensorflow/core/ir/ops.h"
#include "tensorflow/core/platform/statusor.h"
namespace mlir {
namespace tfg {
// Export a generic GraphFuncOp into a FunctionDef. This is intended to be a
// straight serialization, an error is returned in case of failure.
absl::StatusOr<tensorflow::FunctionDef> ConvertGenericFunctionToFunctionDef(
GraphFuncOp func);
} // namespace tfg
} // namespace mlir
#endif // TENSORFLOW_CORE_IR_IMPORTEXPORT_FUNCTIONDEF_EXPORT_H_
tensorflow/core/ir/importexport/functiondef_import.cc
+556
View File
@@ -0,0 +1,556 @@
/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/core/ir/importexport/functiondef_import.h"
#include <string>
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/ADT/StringMap.h"
#include "mlir/IR/Attributes.h" // from @llvm-project
#include "mlir/IR/BuiltinAttributes.h" // from @llvm-project
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "mlir/IR/OperationSupport.h" // from @llvm-project
#include "mlir/Support/LLVM.h" // from @llvm-project
#include "tensorflow/core/framework/attr_value.pb.h"
#include "tensorflow/core/framework/node_def.pb.h"
#include "tensorflow/core/framework/op_def.pb.h"
#include "tensorflow/core/ir/dialect.h"
#include "tensorflow/core/ir/importexport/convert_attributes.h"
#include "tensorflow/core/ir/importexport/convert_types.h"
#include "tensorflow/core/ir/ops.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/platform/statusor.h"
using tensorflow::AttrValue;
using tensorflow::FunctionDef;
using tensorflow::NodeDef;
using tensorflow::OpDef;
using tensorflow::OpDef_AttrDef;
using tensorflow::Status;
using tensorflow::StatusOr;
using tensorflow::errors::InvalidArgument;
using tensorflow::protobuf::RepeatedPtrField;
#define DEBUG_TYPE "graphdef-to-mlir"
namespace mlir {
namespace tfg {
namespace {
class ValueMapManager {
public:
ValueMapManager(
llvm::StringMap<llvm::StringMap<SmallVector<Value, 1>>>& values_map,
OpBuilder& builder, OperationName mlir_placeholder, Type placeholder_ty,
Type control_ty, Location loc)
: values_map_(values_map),
builder_(builder),
loc_(loc),
mlir_placeholder_(mlir_placeholder),
placeholder_ty_(placeholder_ty),
control_ty_(control_ty) {}
Status DefineOperation(Operation* op, StringRef node_name) {
llvm::StringMap<SmallVector<Value, 1>>& op_info = values_map_[node_name];
SmallVector<Value, 1>& base_operation = op_info["^"];
// Replace placeholders.
if (!base_operation.empty()) {
Operation* placeholder = base_operation[0].getDefiningOp();
if (!placeholder ||
placeholder->getName().getStringRef() != "tfg.__mlir_placeholder")
return absl::InvalidArgumentError(absl::StrCat(
"Duplicated node (or function argument) with the same name: `",
node_name.str(), "`"));
op->moveBefore(placeholder);
placeholder->replaceAllUsesWith(op);
placeholder->erase();
base_operation.clear();
}
base_operation.push_back(op->getResult(1));
base_operation.push_back(op->getResult(0));
return absl::OkStatus();
}
absl::StatusOr<Value> GetValueOrCreatePlaceholder(StringRef full_name) {
StringRef node_name;
StringRef output_name = "";
bool is_control_dep = full_name[0] == '^';
size_t output_num = 0;
if (is_control_dep) full_name = full_name.drop_front();
{
size_t colon_sep = full_name.find_first_of(':');
if (colon_sep == StringRef::npos) {
node_name = full_name;
} else {
node_name = full_name.take_front(colon_sep);
output_name = full_name.drop_front(colon_sep + 1);
}
colon_sep = output_name.find_last_of(':');
if (colon_sep != StringRef::npos) {
// NOLINTNEXTLINE: type matching the API taking a reference.
unsigned long long value;
if (!llvm::getAsUnsignedInteger(output_name.drop_front(colon_sep + 1),
10, value)) {
if (LLVM_LIKELY(
value <=
std::numeric_limits<llvm::SmallVectorSizeType<Value>>::max() -
1))
output_num = value;
else
return absl::InvalidArgumentError(absl::StrCat(
"Output index ", value, " is invalid (too large)"));
}
output_name = output_name.take_front(colon_sep);
}
}
llvm::StringMap<SmallVector<Value, 1>>& op_info = values_map_[node_name];
SmallVector<Value, 1>& base_operation = op_info["^"];
if (base_operation.empty()) {
OperationState state(loc_, mlir_placeholder_);
state.addAttribute(TFGraphDialect::getNameAttrKey(),
builder_.getStringAttr(node_name));
state.types.push_back(placeholder_ty_);
state.types.push_back(control_ty_);
Operation* placeholder = builder_.create(state);
base_operation.push_back(placeholder->getResult(1));
base_operation.push_back(placeholder->getResult(0));
}
if (is_control_dep) return base_operation[0];
SmallVector<Value, 1>& value_info = op_info[output_name];
if (value_info.size() <= output_num)
value_info.resize(output_num + 1, Value{});
if (!value_info[output_num]) {
// Guard against accessing OOB. This probably should have been caught
// earlier.
if (base_operation.size() == 1)
return absl::InvalidArgumentError(
"Requested result from op that produces no values, but not "
"considered control dep");
// Create a tfg.get_result for this output.
value_info[output_num] = GetResultOp::create(
builder_, loc_, base_operation[1], output_name, output_num);
}
return value_info[output_num];
}
private:
llvm::StringMap<llvm::StringMap<SmallVector<Value, 1>>>& values_map_;
OpBuilder& builder_;
Location loc_;
OperationName mlir_placeholder_;
Type placeholder_ty_;
Type control_ty_;
};
// Convert the list of `nodes` one by one into MLIR Operations using the
// provided OpBuilder.
// The provided `nodes_map` will be populated with a mapping from the node name
// to the result count and the Operation.
// The supplied `args_map` is looked up for Function arguments when an entry
// cannot be found in the nodes_map.
Status ImportNodes(ValueMapManager value_manager,
const RepeatedPtrField<NodeDef>& nodes, OpBuilder& builder) {
Location unknown_loc = builder.getUnknownLoc();
MLIRContext* context = builder.getContext();
Type placeholder_ty = OpaqueTensorType::get(context);
Type control_ty = ControlType::get(context);
TFGraphDialect* tfgDialect =
cast<TFGraphDialect>(context->getLoadedDialect("tfg"));
StringAttr device_attr = tfgDialect->getDeviceAttrIdentifier();
StringAttr name_attr = tfgDialect->getNameAttrIdentifier();
StringAttr fulltype_attr = tfgDialect->getFullTypeAttrIdentifier();
// Process every node and create a matching MLIR operation
for (const NodeDef& node : nodes) {
DVLOG(1) << "Processing node " << node.name() << "\n";
if (node.op().empty()) return absl::InvalidArgumentError("empty op type");
OperationState state(unknown_loc, absl::StrCat("tfg.", node.op()));
// Fetch the inputs, creating placeholder if an input hasn't been visited.
for (const std::string& input : node.input()) {
if (input.empty())
return absl::InvalidArgumentError(
absl::StrCat("Node '", node.name(), "' has an empty input"));
TF_ASSIGN_OR_RETURN(Value value,
value_manager.GetValueOrCreatePlaceholder(input));
state.operands.push_back(value);
}
// Retrieve the entry in the nodes_map for this node and infer the result
// count from what was inferred during the first traversal above.
state.types.push_back(placeholder_ty);
state.types.push_back(control_ty);
// Handle attributes.
for (const auto& namedAttr : node.attr()) {
const std::string& name = namedAttr.first;
const AttrValue& tf_attr = namedAttr.second;
TF_ASSIGN_OR_RETURN(Attribute attr,
ConvertAttributeValue(tf_attr, builder));
state.addAttribute(name, attr);
}
if (!node.device().empty())
state.addAttribute(device_attr, StringAttr::get(context, node.device()));
if (!node.name().empty())
state.addAttribute(name_attr, StringAttr::get(context, node.name()));
if (node.has_experimental_type()) {
TF_ASSIGN_OR_RETURN(tf_type::FullTypeAttr type,
ConvertAttribute(node.experimental_type(), builder));
state.addAttribute(fulltype_attr, type);
}
Operation* op = builder.create(state);
StringRef node_name = node.name();
{
size_t colon_sep = node_name.find_first_of(':');
if (colon_sep != StringRef::npos)
node_name = node_name.take_front(colon_sep);
}
TF_RETURN_IF_ERROR(value_manager.DefineOperation(op, node_name));
}
// We don't expect any placeholder left at this point, fail if any.
for (Operation& op : *builder.getInsertionBlock()) {
if (op.getName().getStringRef() == "tfg.__mlir_placeholder") {
return absl::InvalidArgumentError(absl::StrCat(
"Couldn't import graph: placeholder left ",
op.getAttrOfType<StringAttr>(name_attr).getValue().str()));
}
}
return absl::OkStatus();
}
absl::StatusOr<NamedAttrList> ConvertArgDefAttributes(const OpDef::ArgDef& arg,
Builder builder) {
NamedAttrList input_attrs;
StringAttr arg_name = builder.getStringAttr(arg.name());
input_attrs.set("tfg.name", arg_name);
if (!arg.description().empty())
input_attrs.append("tfg.description",
builder.getStringAttr(arg.description()));
Type input_type;
if (arg.type() != tensorflow::DT_INVALID) {
TF_RETURN_IF_ERROR(ConvertDataType(arg.type(), builder, &input_type));
input_attrs.append("tfg.type", TypeAttr::get(input_type));
}
if (!arg.type_attr().empty())
input_attrs.append("tfg.type_attr", builder.getStringAttr(arg.type_attr()));
if (!arg.number_attr().empty())
input_attrs.append("tfg.number_attr",
builder.getStringAttr(arg.number_attr()));
if (!arg.type_list_attr().empty())
input_attrs.append("tfg.type_list_attr",
builder.getStringAttr(arg.type_list_attr()));
if (arg.handle_data_size()) {
TF_ASSIGN_OR_RETURN(Attribute handle_data,
ConvertHandleData(builder, arg.handle_data()));
input_attrs.append("tfg.handle_data", handle_data);
}
if (arg.is_ref()) input_attrs.append("tfg.is_ref", builder.getUnitAttr());
if (arg.has_experimental_full_type()) {
TF_ASSIGN_OR_RETURN(
tf_type::FullTypeAttr type,
ConvertAttribute(arg.experimental_full_type(), builder));
input_attrs.append("tfg.experimental_full_type", type);
}
return input_attrs;
}
// Import the given `func` and inser the resulting `GraphFunc`
// operation using the provided `builder`. The `nodes_map` and `args_map` are
// used as scratchpad for the import inside this function. The `gradients` maps
// is provided to
Status ImportGenericFunction(
GraphFuncOp func_op, const FunctionDef& func,
llvm::StringMap<llvm::StringMap<SmallVector<Value, 1>>>& values_map,
OpBuilder& builder) {
const OpDef& signature = func.signature();
Location unknown_loc = builder.getUnknownLoc();
MLIRContext* context = builder.getContext();
NamedAttrList attrs;
DictionaryAttr func_attrs = builder.getDictionaryAttr({});
if (signature.name().empty())
return absl::InvalidArgumentError("generic function without a name");
attrs.append("sym_name", builder.getStringAttr(signature.name()));
attrs.append("generic", builder.getUnitAttr());
if (!signature.description().empty())
attrs.append("description", builder.getStringAttr(signature.description()));
if (signature.is_stateful())
attrs.append("is_stateful", builder.getUnitAttr());
if (signature.control_output_size()) {
SmallVector<Attribute> control_outputs;
for (const std::string& output : signature.control_output())
control_outputs.push_back(builder.getStringAttr(output));
attrs.append("control_output", builder.getArrayAttr(control_outputs));
}
{
NamedAttrList attr_defs;
for (const OpDef_AttrDef& attr : signature.attr()) {
NamedAttrList attr_def;
if (attr.name().empty())
return absl::InvalidArgumentError(
"Missing name for function attribute");
if (!attr.type().empty())
attr_def.append(builder.getNamedAttr(
"function_type", builder.getStringAttr(attr.type())));
if (attr.has_default_value()) {
TF_ASSIGN_OR_RETURN(Attribute attr, ConvertAttributeValue(
attr.default_value(), builder));
attr_def.append(builder.getNamedAttr("default_value", attr));
}
if (!attr.description().empty())
attr_def.append(builder.getNamedAttr(
"description", builder.getStringAttr(attr.description())));
if (attr.has_minimum() || attr.minimum())
attr_def.append(builder.getNamedAttr(
"minimum", builder.getI32IntegerAttr(attr.minimum())));
if (attr.has_allowed_values()) {
TF_ASSIGN_OR_RETURN(
Attribute attr,
ConvertAttributeValue(attr.allowed_values(), builder));
attr_def.append(builder.getNamedAttr("allowed_values", attr));
}
attr_defs.append(builder.getNamedAttr(
attr.name(), attr_def.getDictionary(builder.getContext())));
}
if (!attr_defs.empty()) {
func_attrs = attr_defs.getDictionary(builder.getContext());
attrs.append("tfg.func_attrs", func_attrs);
}
}
// The resource_arg_unique_id is a list of `pair<int, int>`, we import it
// as two arrays of integer right now.
if (func.resource_arg_unique_id_size()) {
SmallVector<int32_t> resource_arg_unique_ids_keys;
SmallVector<int32_t> resource_arg_unique_ids_values;
for (const auto& unique_id : func.resource_arg_unique_id()) {
resource_arg_unique_ids_keys.push_back(unique_id.first);
resource_arg_unique_ids_values.push_back(unique_id.second);
}
attrs.append("resource_arg_unique_ids_keys",
builder.getI32TensorAttr(resource_arg_unique_ids_keys));
attrs.append("resource_arg_unique_ids_values",
builder.getI32TensorAttr(resource_arg_unique_ids_values));
}
// Import the function attributes with a `tf.` prefix to match the current
// infrastructure expectations.
for (const auto& namedAttr : func.attr()) {
if (namedAttr.first.empty())
return absl::InvalidArgumentError("Invalid function attribute name");
const std::string& name = "tf." + namedAttr.first;
const AttrValue& tf_attr = namedAttr.second;
TF_ASSIGN_OR_RETURN(Attribute attr,
ConvertAttributeValue(tf_attr, builder));
attrs.append(name, attr);
}
SmallString<8> arg_or_res_attr_name;
SmallString<8> sub_arg_attr_name;
// Iterate of the input in the signature. Each input will correspond to
// potentially multiple arguments because of how the OpDef allows repeated
// arguments controlled by `number_attr` for example.
// We populate the `arg_names` vector with the name of each input at each
// position, and `arg_types` with the matching type.
int arg_num = 0;
SmallVector<StringRef> arg_names;
SmallVector<Type> arg_types;
SmallVector<Attribute> args_attrs;
SmallVector<Attribute> res_attrs;
for (const auto& enumerated_input : llvm::enumerate(signature.input_arg())) {
const OpDef::ArgDef& input = enumerated_input.value();
TF_ASSIGN_OR_RETURN(NamedAttrList input_attrs,
ConvertArgDefAttributes(input, builder));
auto it = func.arg_attr().find(enumerated_input.index());
if (it != func.arg_attr().end()) {
NamedAttrList arg_attr;
for (const auto& named_attr : it->second.attr()) {
TF_ASSIGN_OR_RETURN(Attribute attr,
ConvertAttributeValue(named_attr.second, builder));
arg_attr.append(named_attr.first, attr);
}
input_attrs.append("tfg.arg_attrs",
arg_attr.getDictionary(builder.getContext()));
}
arg_names.push_back(builder.getStringAttr(input.name()).getValue());
arg_types.push_back(OpaqueTensorType::get(context));
args_attrs.push_back(input_attrs.getDictionary(context));
args_attrs.push_back(NamedAttrList{}.getDictionary(context));
arg_num++;
}
attrs.push_back(builder.getNamedAttr(func_op.getArgAttrsAttrName(),
builder.getArrayAttr(args_attrs)));
// Process the results attributes now.
int res_num = 0;
for (const OpDef::ArgDef& output : signature.output_arg()) {
TF_ASSIGN_OR_RETURN(NamedAttrList output_attrs,
ConvertArgDefAttributes(output, builder));
res_attrs.push_back(output_attrs.getDictionary(context));
++res_num;
}
// Process the control output metadata and store them as attributes.
for (const std::string& output : signature.control_output()) {
NamedAttrList output_attrs;
output_attrs.append("tfg.name", builder.getStringAttr(output));
res_attrs.push_back(output_attrs.getDictionary(context));
++res_num;
}
attrs.push_back(builder.getNamedAttr(func_op.getResAttrsAttrName(),
builder.getArrayAttr(res_attrs)));
values_map.clear();
Block* body = new Block();
func_op.getBody().push_back(body);
Type control_ty = ControlType::get(context);
// Create the block arguments and populate the `values_map` with the matching
// input names.
for (auto type_and_name : llvm::zip(arg_types, arg_names)) {
Value arg = body->addArgument(std::get<0>(type_and_name), unknown_loc);
llvm::StringMap<SmallVector<Value, 1>>& values =
values_map[std::get<1>(type_and_name)];
Value ctl = body->addArgument(control_ty, unknown_loc);
values[""].push_back(arg);
values["^"].push_back(ctl);
}
// Pre-populate the nodes_map with the needed slots for the return.
OpBuilder body_builder = OpBuilder::atBlockEnd(body);
// We use placeholders during the import to create "fake" operations to break
// cycles: we need operands to feed to the users.
OperationName mlir_placeholder("tfg.__mlir_placeholder", context);
Type placeholder_ty = OpaqueTensorType::get(context);
ValueMapManager value_manager(values_map, body_builder, mlir_placeholder,
placeholder_ty, control_ty, unknown_loc);
// Import the function body here, after this we have a function with all
// the nodes, and the nodes_map contains the mapping from node_name to actual
// MLIR Operations.
TF_RETURN_WITH_CONTEXT_IF_ERROR(
ImportNodes(value_manager, func.node_def(), body_builder),
" when importing function ", func.signature().name());
// After the body, the final part is to setup the return. It comes in two
// parts: the `ret` field from the FunctionDef for the regular output and the
// `control_ret` field for the control output.
//
// Because `ret` and `control_ret` aren't ordered, there is an indirection to
// the FunctionDef signature to retrieve the position of each `ret` and
// `control_ret` entry by name. We compute this mapping from the name of an
// output to the position in the result array first.
res_num = 0;
llvm::StringMap<int> output_name_to_position;
for (const OpDef::ArgDef& output : signature.output_arg()) {
if (output_name_to_position.count(output.name()))
return absl::InvalidArgumentError(
absl::StrCat("Duplicated output_arg entry", output.name()));
output_name_to_position[output.name()] = res_num;
++res_num;
}
res_num = 0;
llvm::StringMap<int> control_output_to_position;
for (const std::string& output : signature.control_output()) {
if (control_output_to_position.count(output))
return absl::InvalidArgumentError(
absl::StrCat("Duplicated control_output entry", output));
control_output_to_position[output] = res_num;
++res_num;
}
// We pre-allocate the array of operands and populate it using the
// `output_name_to_position` and `control_output_to_position` populated
// previously.
SmallVector<Value> ret_vals(func.ret_size() + func.control_ret_size(),
Value());
for (const auto& ret_val : func.ret()) {
auto position = output_name_to_position.find(ret_val.first);
if (position == output_name_to_position.end()) {
return absl::InvalidArgumentError(absl::StrCat(
"Can't import function, returned value references unknown output "
"argument ",
ret_val.first));
}
if (ret_val.second.empty()) {
return absl::InvalidArgumentError(absl::StrCat(
"Function '", func.signature().name(), "' has empty result name"));
}
TF_ASSIGN_OR_RETURN(
ret_vals[position->second],
value_manager.GetValueOrCreatePlaceholder(ret_val.second));
}
for (const auto& ret_val : func.control_ret()) {
auto position = control_output_to_position.find(ret_val.first);
if (position == control_output_to_position.end()) {
return absl::InvalidArgumentError(absl::StrCat(
"Can't import function, returned value references unknown output "
"argument ",
ret_val.first));
}
if (ret_val.second.empty()) {
return absl::InvalidArgumentError(
absl::StrCat("Function '", func.signature().name(),
"' has empty control result name"));
}
TF_ASSIGN_OR_RETURN(Value result, value_manager.GetValueOrCreatePlaceholder(
(Twine("^") + ret_val.second).str()));
if (!mlir::isa<ControlType>(result.getType()))
return absl::InvalidArgumentError(
absl::StrCat("failed to map returned value ", ret_val.second,
", isn't a control output"));
ret_vals[func.ret_size() + position->second] = result;
}
// Check that all the of the return operands have been populated.
for (const auto& indexed_val : llvm::enumerate(ret_vals)) {
if (indexed_val.value()) continue;
return absl::InvalidArgumentError(
absl::StrCat("Failed to import function, missing output for position ",
indexed_val.index()));
}
MutableArrayRef<Value> operands = ret_vals;
ReturnOp ret_op = ReturnOp::create(body_builder, unknown_loc,
operands.slice(0, func.ret_size()),
operands.slice(func.ret_size()));
// Now that we have all the types, set the function signature as the
// "function_type" attribute.
{
SmallVector<Type> arg_types_with_ctl;
for (Type type : arg_types) {
arg_types_with_ctl.push_back(type);
arg_types_with_ctl.push_back(control_ty);
}
attrs.append("function_type",
TypeAttr::get(builder.getFunctionType(
arg_types_with_ctl, ret_op.getOperandTypes())));
}
func_op->setAttrs(attrs);
return absl::OkStatus();
}
} // namespace
Status ConvertGenericFunction(GraphFuncOp func_op, const FunctionDef& func,
OpBuilder& builder) {
llvm::StringMap<llvm::StringMap<SmallVector<Value, 1>>> values_map;
return ImportGenericFunction(func_op, func, values_map, builder);
}
} // namespace tfg
} // namespace mlir
tensorflow/core/ir/importexport/functiondef_import.h
+36
View File
@@ -0,0 +1,36 @@
/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_CORE_IR_IMPORTEXPORT_FUNCTIONDEF_IMPORT_H_
#define TENSORFLOW_CORE_IR_IMPORTEXPORT_FUNCTIONDEF_IMPORT_H_
#include "mlir/IR/Builders.h" // from @llvm-project
#include "tensorflow/core/framework/function.pb.h"
#include "tensorflow/core/ir/ops.h"
#include "tensorflow/core/platform/status.h"
namespace mlir {
namespace tfg {
// Import the FunctionDef `func` as a TFG generic function (see GraphFuncOp
// documentation). The function will be inserted using the provided `builder`.
absl::Status ConvertGenericFunction(GraphFuncOp func_op,
const tensorflow::FunctionDef& func,
OpBuilder& builder);
} // namespace tfg
} // namespace mlir
#endif // TENSORFLOW_CORE_IR_IMPORTEXPORT_FUNCTIONDEF_IMPORT_H_
tensorflow/core/ir/importexport/graphdef_export.cc
+661
View File
@@ -0,0 +1,661 @@
/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/core/ir/importexport/graphdef_export.h"
#include <optional>
#include <string>
#include <tuple>
#include <utility>
#include <vector>
#include "absl/strings/str_cat.h"
#include "llvm/ADT/PointerUnion.h"
#include "llvm/ADT/ScopeExit.h"
#include "llvm/ADT/Sequence.h"
#include "mlir/IR/Attributes.h" // from @llvm-project
#include "mlir/IR/BuiltinAttributes.h" // from @llvm-project
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/IR/Location.h" // from @llvm-project
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "mlir/IR/Operation.h" // from @llvm-project
#include "mlir/IR/SymbolTable.h" // from @llvm-project
#include "mlir/IR/Threading.h" // from @llvm-project
#include "mlir/IR/Value.h" // from @llvm-project
#include "mlir/Support/LLVM.h" // from @llvm-project
#include "tensorflow/core/framework/attr_value.pb.h"
#include "tensorflow/core/framework/function.h"
#include "tensorflow/core/framework/function.pb.h"
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/framework/node_def.pb.h"
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/op_def.pb.h"
#include "tensorflow/core/framework/op_def_builder.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/framework/versions.pb.h"
#include "tensorflow/core/ir/dialect.h"
#include "tensorflow/core/ir/importexport/convert_attributes.h"
#include "tensorflow/core/ir/importexport/convert_types.h"
#include "tensorflow/core/ir/importexport/functiondef_export.h"
#include "tensorflow/core/ir/ops.h"
#include "tensorflow/core/ir/types/dialect.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/platform/statusor.h"
using tensorflow::AttrValue;
using tensorflow::DataType;
using tensorflow::FunctionDef;
using tensorflow::FunctionLibraryDefinition;
using tensorflow::GradientDef;
using tensorflow::GraphDef;
using tensorflow::NodeDef;
using tensorflow::OpDef;
using tensorflow::OpRegistrationData;
using tensorflow::OpRegistry;
using tensorflow::Status;
using tensorflow::StatusOr;
using tensorflow::VersionDef;
using tensorflow::errors::InvalidArgument;
namespace mlir {
namespace tfg {
namespace {
// This class implements an exporter for TFG directly to GraphDef.
class GraphDefExporter {
public:
GraphDefExporter(
TFGraphDialect *dialect, const OpRegistry &registry,
llvm::PointerUnion<SymbolTable *, const FunctionLibraryDefinition *>
function_table)
: ctx_(dialect->getContext()),
dialect_(dialect),
registry_(registry),
function_table_(function_table) {}
// Export a TFG module to GraphDef. The module may contain at most one GraphOp
// and only GraphFuncOp otherwise.
Status ExportToGraphDef(ModuleOp module, GraphDef *graph);
// Export a TFG graph function to a FunctionDef. If the function has a
// gradient, add it to the graph afterwards to preserve thread-safety.
absl::StatusOr<std::optional<GradientDef>> ExportFunction(GraphFuncOp func,
FunctionDef *def);
private:
// Export just the input and outputs of a function signature. When
// fully-qualifying result names, this must be done before any nodes are
// Convert argument attributes to an ArgDef.
absl::StatusOr<OpDef::ArgDef> ConvertArgumentAttributes(DictionaryAttr attrs);
// Convert a TFG op to a node. When converting a function, fully-qualified
// result names must be used.
Status ConvertOperation(Operation *op, NodeDef *node, bool is_func);
// Get the name associated with a value.
absl::StatusOr<std::string> GetEdgeName(Value value, bool is_func);
// Get the name and index of an output segment to fully qualify result names.
// This requires querying the op registry.
absl::StatusOr<std::pair<StringRef, unsigned int>> GetOutputSegment(
OpResult result);
// Get the name of a function argument from a function in the symbol table.
absl::StatusOr<StringRef> GetFunctionOutputName(unsigned result_idx,
const std::string &op_name,
SymbolTable &table);
// Get the name of a function argument from a function in the library.
static absl::StatusOr<StringRef> GetFunctionOutputName(
unsigned result_idx, const std::string &op_name,
const FunctionLibraryDefinition &library);
// The current MLIR context.
MLIRContext *ctx_;
// The TFG dialect instance.
TFGraphDialect *dialect_;
// The TF op registry to use.
const OpRegistry &registry_;
// A lookup table for functions.
llvm::PointerUnion<SymbolTable *, const FunctionLibraryDefinition *>
function_table_;
};
} // namespace
// Returns a validated graph to export. A TFG module is valid for export if it
// contains at most one graph operation and any number of graph functions.
// Otherwise, returns an error.
static absl::StatusOr<GraphOp> ValidateModuleForExport(ModuleOp module) {
GraphOp graph_op;
for (Operation &op : *module.getBody()) {
if (isa<GraphFuncOp>(op)) continue;
if (auto new_graph_op = dyn_cast<GraphOp>(op)) {
if (graph_op) {
return absl::InvalidArgumentError(
"Can't export module with two different tfg.graph");
}
graph_op = new_graph_op;
continue;
}
return absl::InvalidArgumentError(absl::StrCat(
"Can't export module with other ops than tfg.graph or tfg.func, has: ",
op.getName().getStringRef().str()));
}
return graph_op;
}
// Converts a version attribute to VersionDef.
static void ExportVersionAttr(VersionAttr attr, VersionDef *version) {
version->set_producer(attr.getProducer());
version->set_min_consumer(attr.getMinConsumer());
for (int32_t bad_consumer : attr.getBadConsumers())
version->add_bad_consumers(bad_consumer);
}
Status GraphDefExporter::ExportToGraphDef(ModuleOp module, GraphDef *graph) {
TF_ASSIGN_OR_RETURN(GraphOp graph_op, ValidateModuleForExport(module));
if (graph_op) {
ExportVersionAttr(graph_op.getVersion(), graph->mutable_versions());
for (Operation &op : *graph_op.getBody()) {
TF_RETURN_IF_ERROR(ConvertOperation(&op, graph->mutable_node()->Add(),
/*is_func=*/false));
}
}
const auto convert_func = [this](GraphFuncOp func, FunctionDef *def,
std::optional<GradientDef> &gradient) {
// Generic functions are not on the hot path and skip the conversion to
// Graph so just call the existing exporter.
if (func.getGeneric()) {
TF_ASSIGN_OR_RETURN(*def, ConvertGenericFunctionToFunctionDef(func));
} else {
TF_ASSIGN_OR_RETURN(gradient, ExportFunction(func, def));
}
return absl::OkStatus();
};
// TODO(jeffniu): Don't export functions in parallel if there are too few or
// they are too small.
if (ctx_->isMultithreadingEnabled()) {
ctx_->enterMultiThreadedExecution();
auto exit =
llvm::make_scope_exit([this] { ctx_->exitMultiThreadedExecution(); });
// Prepare the arguments to parallel for each.
struct Argument {
GraphFuncOp func;
FunctionDef *def;
Status status;
std::optional<GradientDef> gradient;
};
std::vector<Argument> args;
for (auto func : module.getOps<GraphFuncOp>())
args.push_back(Argument{func, graph->mutable_library()->add_function()});
const auto process_func = [&convert_func](Argument &arg) {
arg.status = convert_func(arg.func, arg.def, arg.gradient);
return success(arg.status.ok());
};
// Execute the exports in parallel.
if (failed(failableParallelForEach(ctx_, args, process_func))) {
Status result;
for (const Argument &arg : args) {
result.Update(arg.status);
}
return result;
}
} else {
for (auto func : module.getOps<GraphFuncOp>()) {
std::optional<GradientDef> gradient;
TF_RETURN_IF_ERROR(convert_func(
func, graph->mutable_library()->add_function(), gradient));
if (gradient)
*graph->mutable_library()->add_gradient() = std::move(*gradient);
}
}
return absl::OkStatus();
}
// The only dialect attributes allowed have the "tf." prefix. This is a slightly
// faster check that an attribute is a dialect attribute.
static bool IsDialectAttr(const NamedAttribute &attr) {
return attr.getName().getValue().starts_with("tf.");
}
// Export the given attribute list.
static Status ConvertAttributes(
tensorflow::protobuf::Map<std::string, AttrValue> *map,
ArrayRef<NamedAttribute> attrs) {
for (const NamedAttribute &attr : attrs) {
if (!IsDialectAttr(attr)) continue;
StringRef name = attr.getName().strref().drop_front(/*strlen("tf.")=*/3);
TF_ASSIGN_OR_RETURN((*map)[name.str()], ConvertAttribute(attr.getValue()));
}
return absl::OkStatus();
}
absl::StatusOr<std::optional<GradientDef>> GraphDefExporter::ExportFunction(
GraphFuncOp func, FunctionDef *def) {
std::string func_name = func.getSymName().str();
// TODO(jeffniu): Exploit the sorted order of the function attributes.
// Get a gradient, if there is one.
std::optional<GradientDef> gradient;
if (std::optional<StringRef> gradient_name = func.getGradient()) {
gradient.emplace();
gradient->set_gradient_func(gradient_name->str());
gradient->set_function_name(func_name);
}
// Convert the first-class attributes.
OpDef *signature = def->mutable_signature();
signature->set_name(func_name);
if (std::optional<StringRef> description = func.getDescription())
signature->set_description(description->str());
signature->set_is_stateful(func.getIsStateful());
if (DenseIntElementsAttr keys = func.getResourceArgUniqueIdsKeysAttr()) {
DenseIntElementsAttr values = func.getResourceArgUniqueIdsValuesAttr();
if (!values) {
return absl::InvalidArgumentError(
"'resource_arg_unique_ids_keys' is present but "
"'resource_arg_unique_ids_values' is missing");
}
if (keys.size() != values.size()) {
return absl::InvalidArgumentError(
"'resource_arg_unique_ids_keys' is not the same size as "
"'resource_arg_unique_ids_values'");
}
auto *id_map = def->mutable_resource_arg_unique_id();
for (auto kv :
llvm::zip(keys.getValues<int32_t>(), values.getValues<int32_t>()))
(*id_map)[std::get<0>(kv)] = std::get<1>(kv);
}
// Convert other attributes with the "tf." prefix.
TF_RETURN_IF_ERROR(ConvertAttributes(def->mutable_attr(), func->getAttrs()));
// Convert the arguments.
for (int i = 0, e = func.getNumArguments(); i < e; i += 2) {
auto attrs = mlir::cast<DictionaryAttr>(func.getArgAttrs().value()[i]);
TF_ASSIGN_OR_RETURN(OpDef::ArgDef &arg = *signature->add_input_arg(),
ConvertArgumentAttributes(attrs));
DataType dtype;
TF_RETURN_IF_ERROR(ConvertToDataType(
mlir::cast<TensorType>(func.getArgument(i).getType()).getElementType(),
&dtype));
arg.set_type(dtype);
// Convert the attributes.
if (llvm::any_of(attrs, [](const NamedAttribute &attr) {
return IsDialectAttr(attr);
})) {
auto *map = (*def->mutable_arg_attr())[i / 2].mutable_attr();
TF_RETURN_IF_ERROR(ConvertAttributes(map, attrs.getValue()));
}
}
// Convert the results.
auto return_op = cast<ReturnOp>(func.SingleBlock::getBody()->getTerminator());
if (func.getAllResultAttrs()) {
for (auto it :
llvm::zip(func.getResultTypes(),
func.getAllResultAttrs().getAsRange<DictionaryAttr>(),
TFOp(return_op).getNonControlOperands())) {
TF_ASSIGN_OR_RETURN(OpDef::ArgDef& arg = *signature->add_output_arg(),
ConvertArgumentAttributes(std::get<1>(it)));
DataType dtype;
TF_RETURN_IF_ERROR(ConvertToDataType(
mlir::cast<TensorType>(std::get<0>(it)).getElementType(), &dtype));
arg.set_type(dtype);
// Map the result.
TF_ASSIGN_OR_RETURN((*def->mutable_ret())[arg.name()],
GetEdgeName(std::get<2>(it), /*is_func=*/true));
}
}
// Convert the control results.
for (auto it :
llvm::zip(return_op.getControlRetAttrs().getAsRange<DictionaryAttr>(),
TFOp(return_op).getControlOperands())) {
// The control result attributes contain only the name.
DictionaryAttr attrs = std::get<0>(it);
if (attrs.empty())
return absl::InvalidArgumentError("Control result is missing 'tfg.name'");
assert(attrs.begin()->getName() == dialect_->getTfgNameAttrIdentifier());
std::string name = mlir::cast<StringAttr>(attrs.begin()->getValue()).str();
signature->add_control_output(name);
// Map the control result.
TF_ASSIGN_OR_RETURN(std::string value_name,
GetEdgeName(std::get<1>(it), /*is_func=*/true));
// Add the control result name without '^'.
def->mutable_control_ret()->insert({std::move(name), value_name.substr(1)});
}
// Convert the body.
for (Operation &op : func.SingleBlock::getBody()->without_terminator())
TF_RETURN_IF_ERROR(
ConvertOperation(&op, def->add_node_def(), /*is_func=*/true));
return gradient;
}
absl::StatusOr<OpDef::ArgDef> GraphDefExporter::ConvertArgumentAttributes(
DictionaryAttr attrs) {
OpDef::ArgDef arg;
auto name = attrs.getAs<StringAttr>(dialect_->getTfgNameAttrIdentifier());
if (!name)
return absl::InvalidArgumentError("argument is missing 'tfg.name'");
arg.set_name(name.str());
if (auto description =
attrs.getAs<StringAttr>(dialect_->getTfgDescriptionAttrIdentifier()))
arg.set_description(description.str());
arg.set_is_ref(!!attrs.get(dialect_->getTfgIsRefAttrIdentifier()));
TF_RETURN_IF_ERROR(ConvertHandleData(
attrs.getAs<ArrayAttr>(dialect_->getTfgHandleDataAttrIdentifier()),
&arg));
if (auto full_type = attrs.getAs<tf_type::FullTypeAttr>(
dialect_->getTfgFullTypeAttrIdentifier())) {
TF_ASSIGN_OR_RETURN(*arg.mutable_experimental_full_type(),
ConvertAttribute(full_type));
}
return arg;
}
// Converts a location to the debug information for the node def, if we find
// supported location, that is a top-level NameLoc or any NameLoc nested inside
// a FusedLoc. Other kind of location are ignored. If a NameLoc is of the form
// "name@func" we parse it and import the two appropriately.
static void ExtractExperimentalDebugInfoFromLocation(
Location inst_loc, NodeDef::ExperimentalDebugInfo *debug_info) {
auto add_name_loc = [&](mlir::NameLoc name_loc) {
StringRef node, func;
std::tie(node, func) = name_loc.getName().strref().split('@');
debug_info->add_original_node_names(node.str());
if (!func.empty()) debug_info->add_original_func_names(func.str());
};
if (auto fused = mlir::dyn_cast<mlir::FusedLoc>(inst_loc)) {
for (Location loc : fused.getLocations())
if (auto name_loc = mlir::dyn_cast<mlir::NameLoc>(loc))
add_name_loc(name_loc);
return;
}
if (auto name_loc = mlir::dyn_cast<mlir::NameLoc>(inst_loc))
add_name_loc(name_loc);
}
Status ConvertToNodeDef(
Operation *op, NodeDef *node, TFGraphDialect *dialect,
function_ref<absl::StatusOr<std::string>(Value)> get_value_name) {
// Convert first-class attributes.
if (auto name =
op->getAttrOfType<StringAttr>(dialect->getNameAttrIdentifier()))
node->set_name(name.str());
if (auto device =
op->getAttrOfType<StringAttr>(dialect->getDeviceAttrIdentifier()))
node->set_device(device.str());
if (auto full_type = op->getAttrOfType<tf_type::FullTypeAttr>(
dialect->getFullTypeAttrIdentifier())) {
TF_ASSIGN_OR_RETURN(*node->mutable_experimental_type(),
ConvertAttribute(full_type));
}
{
if (auto assigned_device = op->getAttrOfType<StringAttr>(
dialect->getAssignedDeviceAttrIdentifier())) {
if (!assigned_device.getValue().empty()) {
(*node->mutable_attr())[dialect->getAssignedDeviceAttrKey().str()]
.set_s(assigned_device.str());
}
}
}
// Convert other attributes.
for (const NamedAttribute &attr : op->getAttrs()) {
if (attr.getName() == dialect->getAssignedDeviceAttrIdentifier() ||
attr.getName() == dialect->getDeviceAttrIdentifier() ||
attr.getName() == dialect->getFullTypeAttrIdentifier() ||
attr.getName() == dialect->getNameAttrIdentifier())
continue;
TF_ASSIGN_OR_RETURN((*node->mutable_attr())[attr.getName().str()],
ConvertAttribute(attr.getValue()));
}
// Set the op name.
node->set_op(op->getName().stripDialect().str());
// Set the input names.
for (Value operand : op->getOperands()) {
TF_ASSIGN_OR_RETURN(std::string input_name, get_value_name(operand));
node->add_input(std::move(input_name));
}
// Export the location as debug info.
if (!mlir::isa<UnknownLoc>(op->getLoc())) {
ExtractExperimentalDebugInfoFromLocation(
op->getLoc(), node->mutable_experimental_debug_info());
if (node->experimental_debug_info().original_node_names().empty())
node->clear_experimental_debug_info();
}
return absl::OkStatus();
}
Status GraphDefExporter::ConvertOperation(Operation *op, NodeDef *node,
bool is_func) {
return ConvertToNodeDef(op, node, dialect_, [&](Value value) {
return GetEdgeName(value, is_func);
});
}
// Get the edge name of a value. If `get_output_segment` is specified, it means
// the name should be fully qualified if it is an operation result for exporting
// a function.
static absl::StatusOr<std::string> GetValueName(
Value value, TFGraphDialect *dialect,
function_ref<absl::StatusOr<std::pair<StringRef, unsigned int>>(OpResult)>
get_output_segment) {
std::string name;
bool is_control = value.getType() == dialect->getControlType();
if (auto arg = mlir::dyn_cast<BlockArgument>(value)) {
auto func = dyn_cast<GraphFuncOp>(arg.getOwner()->getParentOp());
if (!func)
return absl::InvalidArgumentError(
"Expected block argument owner to be tfg.func");
// If the block argument is a control token, use the attributes of the
// associated data argument (which preceeds it).
auto attrs = mlir::cast<DictionaryAttr>(
func.getArgAttrs().value()[arg.getArgNumber() - is_control]);
auto name_attr =
attrs.getAs<StringAttr>(dialect->getTfgNameAttrIdentifier());
if (!name_attr) {
return absl::InvalidArgumentError(absl::StrCat(
"Can't export graph with missing op-name for function parameter #",
arg.getArgNumber()));
}
name.reserve(name_attr.size() + 1);
if (is_control) name.push_back('^');
name.append(name_attr.data(), name_attr.size());
return name;
}
auto result = mlir::cast<OpResult>(value);
auto name_attr = result.getOwner()->getAttrOfType<StringAttr>(
dialect->getNameAttrIdentifier());
if (!name_attr)
return absl::InvalidArgumentError(
"Can't export graph with missing op-name");
if (is_control) {
name.reserve(1 + name_attr.size());
name.push_back('^');
name.append(name_attr.data(), name_attr.size());
return name;
}
if (!get_output_segment) {
name.reserve(name_attr.size() + 3);
name.append(name_attr.data(), name_attr.size());
if (result.getResultNumber()) {
name.push_back(':');
absl::StrAppend(&name, result.getResultNumber());
}
return name;
}
TF_ASSIGN_OR_RETURN(auto segment, get_output_segment(result));
name.reserve(name_attr.size() + segment.first.size() + 4);
name.append(name_attr.data(), name_attr.size());
name.push_back(':');
name.append(segment.first.data(), segment.first.size());
name.push_back(':');
absl::StrAppend(&name, segment.second);
return name;
}
absl::StatusOr<std::string> GetValueName(Value value, TFGraphDialect *dialect) {
return GetValueName(value, dialect, /*get_output_segment=*/nullptr);
}
absl::StatusOr<std::string> GraphDefExporter::GetEdgeName(Value value,
bool is_func) {
if (!is_func) return GetValueName(value, dialect_);
return GetValueName(value, dialect_, [&](OpResult result) {
return GetOutputSegment(result);
});
}
// Get the segment size of an op's output.
static absl::StatusOr<unsigned int> GetOutputSegmentSize(
Operation *op, const OpDef::ArgDef &arg) {
if (!arg.type_list_attr().empty()) {
if (auto v = mlir::dyn_cast<ArrayAttr>(op->getAttr(arg.type_list_attr())))
return v.size();
return absl::InvalidArgumentError(
absl::StrCat("Type attr not found: ", arg.type_list_attr()));
}
if (arg.number_attr().empty()) return 1;
if (auto v = mlir::dyn_cast<IntegerAttr>(op->getAttr(arg.number_attr())))
return v.getValue().getZExtValue();
return absl::InvalidArgumentError(
absl::StrCat("Type attr not found: ", arg.number_attr()));
}
absl::StatusOr<StringRef> GraphDefExporter::GetFunctionOutputName(
unsigned result_idx, const std::string &op_name, SymbolTable &table) {
if (auto func = table.lookup<GraphFuncOp>(op_name)) {
if (result_idx >= func.getNumResults()) {
return absl::InvalidArgumentError(absl::StrCat("Result #", result_idx,
" of function '", op_name,
"' is out of range"));
}
if (auto name = func.getResultAttrOfType<StringAttr>(
result_idx, dialect_->getTfgNameAttrIdentifier())) {
return name.getValue();
}
return absl::InvalidArgumentError(absl::StrCat("Function '", op_name,
"' result #", result_idx,
"' is missing 'tfg.name'"));
}
return absl::InvalidArgumentError(
absl::StrCat("Op '", op_name, "' is neither registered nor a function"));
}
// Get the name of a function argument from a function in the library.
absl::StatusOr<StringRef> GraphDefExporter::GetFunctionOutputName(
unsigned result_idx, const std::string &op_name,
const FunctionLibraryDefinition &library) {
if (const FunctionDef *function = library.Find(op_name)) {
if (result_idx >= function->signature().output_arg_size()) {
return absl::InvalidArgumentError(absl::StrCat("Result #", result_idx,
" of function '", op_name,
"' is out of range"));
}
return {function->signature().output_arg(result_idx).name()};
}
return absl::InvalidArgumentError(
absl::StrCat("Op '", op_name, "' is neither registered nor a function"));
}
absl::StatusOr<std::pair<StringRef, unsigned int>>
GraphDefExporter::GetOutputSegment(OpResult result) {
// TODO(jeffniu): OpRegistry::LookUp should accept `string_view`.
Operation *op = result.getOwner();
std::string op_name = op->getName().stripDialect().str();
unsigned result_idx = result.getResultNumber();
// Only edges in functions need to have fully-qualified names. Get the segment
// name using the op definition.
if (const OpRegistrationData *op_reg_data = registry_.LookUp(op_name)) {
const OpDef &op_def = op_reg_data->op_def;
for (const OpDef::ArgDef &arg : op_def.output_arg()) {
TF_ASSIGN_OR_RETURN(unsigned size, GetOutputSegmentSize(op, arg));
if (size > result_idx)
return std::pair<StringRef, unsigned>(arg.name(), result_idx);
result_idx -= size;
}
return absl::InvalidArgumentError(absl::StrCat(
"Result #", result_idx, " of op '", op_name, "' is out of range"));
}
// Try to find a function for a legacy call. Function output segments have
// exactly one element each.
StringRef arg_name;
if (auto *table = function_table_.dyn_cast<SymbolTable *>()) {
TF_ASSIGN_OR_RETURN(arg_name,
GetFunctionOutputName(result_idx, op_name, *table));
} else {
TF_ASSIGN_OR_RETURN(
arg_name,
GetFunctionOutputName(
result_idx, op_name,
*function_table_.get<const FunctionLibraryDefinition *>()));
}
return std::pair<StringRef, unsigned>(arg_name, 0);
}
// Convert a TFG graph directly to GraphDef.
Status ConvertToGraphDef(ModuleOp module, tensorflow::GraphDef *graph) {
SymbolTable table(module);
GraphDefExporter exporter(
module.getContext()->getOrLoadDialect<TFGraphDialect>(),
*OpRegistry::Global(), &table);
return exporter.ExportToGraphDef(module, graph);
}
// Convert a single TFG function to a FunctionDef and add it to the function
// library. If a function with the same name already exists, replace it.
Status ConvertToFunctionDef(GraphFuncOp func,
FunctionLibraryDefinition &library) {
GraphDefExporter exporter(func.getDialect(), *OpRegistry::Global(), &library);
FunctionDef def;
TF_ASSIGN_OR_RETURN(std::optional<GradientDef> gradient,
exporter.ExportFunction(func, &def));
const std::string &name = def.signature().name();
if (library.Contains(name)) {
TF_RETURN_IF_ERROR(library.ReplaceFunction(name, def));
} else {
TF_RETURN_IF_ERROR(library.AddFunctionDef(def));
}
if (gradient) {
if (library.FindGradient(name).empty()) {
TF_RETURN_IF_ERROR(library.AddGradientDef(*gradient));
} else {
TF_RETURN_IF_ERROR(library.ReplaceGradient(*gradient));
}
}
return absl::OkStatus();
}
} // namespace tfg
} // namespace mlir
tensorflow/core/ir/importexport/graphdef_export.h
+56
View File
@@ -0,0 +1,56 @@
/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_CORE_IR_IMPORTEXPORT_GRAPHDEF_EXPORT_H_
#define TENSORFLOW_CORE_IR_IMPORTEXPORT_GRAPHDEF_EXPORT_H_
#include <string>
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/IR/Value.h" // from @llvm-project
#include "mlir/Support/LLVM.h" // from @llvm-project
#include "tensorflow/core/framework/function.h"
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/framework/node_def.pb.h"
#include "tensorflow/core/ir/dialect.h"
#include "tensorflow/core/ir/ops.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/platform/statusor.h"
namespace mlir {
namespace tfg {
// Get the name of a value as if it were an edge in a graph.
absl::StatusOr<std::string> GetValueName(Value value, TFGraphDialect *dialect);
// Convert a TFG graph directly to GraphDef. Graph functions in the module are
// added to the GraphDef's function library.
absl::Status ConvertToGraphDef(ModuleOp module, tensorflow::GraphDef *graph);
// Convert a single TFG op to NodeDef. This utliity function requires a callback
// `get_value_name` that returns the edge name of the given operand.
absl::Status ConvertToNodeDef(
Operation *op, tensorflow::NodeDef *node, TFGraphDialect *dialect,
function_ref<absl::StatusOr<std::string>(Value)> get_value_name);
// Convert a single TFG function to a FunctionDef and add it to the function
// library. If a function with the same name already exists, replace it.
absl::Status ConvertToFunctionDef(
GraphFuncOp func, tensorflow::FunctionLibraryDefinition &library);
} // namespace tfg
} // namespace mlir
#endif // TENSORFLOW_CORE_IR_IMPORTEXPORT_GRAPHDEF_EXPORT_H_
tensorflow/core/ir/importexport/graphdef_import.cc
+952
View File
@@ -0,0 +1,952 @@
/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/core/ir/importexport/graphdef_import.h"
#include <iterator>
#include <memory>
#include <string>
#include <tuple>
#include <utility>
#include <vector>
#include "absl/container/flat_hash_map.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/ScopeExit.h"
#include "llvm/Support/raw_ostream.h"
#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/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/Threading.h" // from @llvm-project
#include "mlir/Support/LLVM.h" // from @llvm-project
#include "mlir/Support/LogicalResult.h" // from @llvm-project
#include "tensorflow/core/framework/full_type.pb.h"
#include "tensorflow/core/framework/function.h"
#include "tensorflow/core/framework/function.pb.h"
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/framework/graph_debug_info.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/op_def_builder.h"
#include "tensorflow/core/framework/versions.pb.h"
#include "tensorflow/core/graph/graph.h"
#include "tensorflow/core/graph/graph_debug_info_builder.h"
#include "tensorflow/core/graph/tensor_id.h"
#include "tensorflow/core/ir/dialect.h"
#include "tensorflow/core/ir/importexport/convert_attributes.h"
#include "tensorflow/core/ir/importexport/convert_types.h"
#include "tensorflow/core/ir/importexport/functiondef_import.h"
#include "tensorflow/core/ir/ops.h"
#include "tensorflow/core/ir/types/dialect.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/platform/statusor.h"
#include "tensorflow/core/platform/stringpiece.h"
using tensorflow::DataType;
using tensorflow::DataTypeVector;
using tensorflow::FullTypeDef;
using tensorflow::FunctionDef;
using tensorflow::FunctionLibraryDefinition;
using tensorflow::Graph;
using tensorflow::GraphDebugInfo;
using tensorflow::GraphDef;
using tensorflow::NodeDef;
using tensorflow::OpDef;
using tensorflow::OpRegistrationData;
using tensorflow::OpRegistry;
using tensorflow::StackTracesMap;
using tensorflow::Status;
using tensorflow::StatusOr;
using tensorflow::StringPiece;
using tensorflow::TensorId;
using tensorflow::VersionDef;
using tensorflow::errors::InvalidArgument;
using tensorflow::errors::NotFound;
namespace mlir {
namespace tfg {
namespace {
// This class implements an importer for GraphDef directly to TFG.
class GraphDefImporter {
public:
// Initialize the importer.
GraphDefImporter(TFGraphDialect *dialect, const OpRegistry &registry,
const GraphDebugInfo &debug_info)
: ctx_(dialect->getContext()),
dialect_(dialect),
b_(ctx_),
registry_(registry),
debug_info_(debug_info),
unknown_loc_(UnknownLoc::get(ctx_)),
placeholder_state_(unknown_loc_, "tfg._mlir_placeholder") {
placeholder_state_.addTypes(dialect_->getControlType());
stack_traces_ = LoadTracesFromDebugInfo(debug_info);
}
// Convert a GraphDef to MLIR module.
absl::StatusOr<OwningOpRef<ModuleOp>> ConvertGraphDef(const GraphDef &graph);
private:
// Convert a function. This function must be thread-safe.
Status ConvertFunctionDef(
GraphFuncOp func_op,
const absl::flat_hash_map<StringPiece, StringPiece> &gradient_map,
const FunctionDef &function);
// A result ID representing an output of `node`. E.g.
// "foo" -> {0, "foo", ""}
// "foo:2" -> {2, "foo", ""}
// "foo:output:0" -> {0, "foo", "output"}
struct ResultId {
// The result or result segment index.
int index;
// The name of the parent node.
StringRef node;
// An optional result segment name.
StringRef output;
// Returns true if the result ID references the control token.
bool IsControl() const { return index == tensorflow::Graph::kControlSlot; }
};
// An unresolved backedge.
struct Backedge {
// The edge name and index.
ResultId id;
// The OpOperand to resolve;
OpOperand *operand;
};
// Cached info about the result of an operation.
struct ResultInfo {
// This flag is true if the results of the operation have been resolved; the
// operation has been created and its `data` and `control` results have been
// populated. If false, the placeholder should be used.
bool resolved = false;
// The control result.
Value control;
// All data results.
ValueRange data;
// Data results organized by output name.
absl::flat_hash_map<StringPiece, ValueRange> outputs;
// A list of unresolved backedges.
std::vector<Backedge> backedges;
};
// State when converting a list of nodes.
class ConversionState
: public absl::flat_hash_map<StringPiece, std::unique_ptr<ResultInfo>> {
public:
// Create a conversion state with a placeholder value. Put the plaecholder
// in the block so that it is owned.
explicit ConversionState(Block *block,
const OperationState &placeholder_state)
: placeholder_op_(
OpBuilder::atBlockBegin(block).create(placeholder_state)),
placeholder_(placeholder_op_->getResult(0)) {}
// Get the placeholder value.
Value GetPlaceholder() { return placeholder_; }
// Finalize the conversion. The placeholder is destroyed.
void Finalize() { placeholder_op_->erase(); }
private:
// The placeholder operation.
Operation *placeholder_op_;
// The placeholder value.
Value placeholder_;
};
// Convert a list a nodes to operations.
Status ConvertNodes(
OpBuilder &builder, ConversionState &s,
const tensorflow::protobuf::RepeatedPtrField<NodeDef> &nodes,
Block *block);
// Convert a node to an operation.
Status ConvertNodeDef(OpBuilder &builder, ConversionState &s,
const NodeDef &node);
// Resolve a data result reference.
static absl::StatusOr<Value> ResolveDataResult(const ResultId &id,
ResultInfo *info);
// Get a named result.
struct Result {
Value control;
Value data;
ResultId id;
ResultInfo *info = nullptr;
};
absl::StatusOr<Result> GetResult(ConversionState &s, StringPiece name);
// Convert TF datatypes to unranked MLIR tensor types.
Status ConvertDataTypesToUnrankedTensorTypes(const DataTypeVector &dtypes,
SmallVectorImpl<Type> &results);
// Extracts the actual data types from `attrs` based on its definition in
// `arg_def` and converts them to unranked tensors. Returns the number of
// added types.
//
// TODO(jeffniu): This is a re-implementation of `ArgNumType` in
// `core/framework/function.cc` on `NamedAttrList` because the default
// attributes need to be added. Find a way to do this in one pass.
absl::StatusOr<unsigned int> ArgNumType(const NamedAttrList &attrs,
const OpDef::ArgDef &arg_def,
SmallVectorImpl<Type> &types);
// Convert function attributes to MLIR attributes.
Status ConvertFunctionAttributes(
const absl::flat_hash_map<StringPiece, StringPiece> &gradient_map,
const FunctionDef &function, GraphFuncOp op, NamedAttrList &attrs);
// Convert function argument attributes to MLIR attributes.
Status ConvertArgumentAttributes(const OpDef::ArgDef &def,
NamedAttrList &attrs);
// Create a location for a node.
Location ConvertLocation(const NodeDef &node);
// Convert the location of a node from the debug info. If it has no debug
// info, return a NameLoc.
Location ConvertLocation(StringRef node_name, StringRef func_name);
// The MLIR context.
MLIRContext *ctx_;
// Reference to the TFG dialect.
TFGraphDialect *dialect_;
// The builder instance.
Builder b_;
// The TF op registry to use.
const OpRegistry &registry_;
// The debug info about the graph.
const GraphDebugInfo &debug_info_;
// Cached unknown location.
Location unknown_loc_;
// Operation state for creating placeholder ops.
OperationState placeholder_state_;
StackTracesMap stack_traces_;
// Map of function OpDefs.
absl::flat_hash_map<StringPiece, const OpDef *> function_op_defs_;
};
} // namespace
// Convert a VersionDef to an MLIR version attribute.
static VersionAttr ConvertVersionAttr(MLIRContext *context,
const VersionDef &version) {
ArrayRef<int32_t> bad_consumers(version.bad_consumers().data(),
version.bad_consumers().size());
return VersionAttr::get(context, version.producer(), version.min_consumer(),
bad_consumers);
}
// Returns true if the function is a generic function, i.e. it contains
// placeholder attributes.
//
// TODO(jeffniu): Having to iterate over every function just to check for
// placeholder attributes is slow. Since most functions are not generic, we can
// speculate by converting all functions as non-generic until we see a
// placeholder attribute, bail out, and fall back to the generic function
// converter.
static bool IsGenericFunction(const FunctionDef &fdef) {
for (const NodeDef &node : fdef.node_def())
for (const auto &named_attr : node.attr())
if (!named_attr.second.placeholder().empty()) return true;
return false;
}
absl::StatusOr<OwningOpRef<ModuleOp>> GraphDefImporter::ConvertGraphDef(
const GraphDef &graph) {
// Create the module.
OwningOpRef<ModuleOp> module = ModuleOp::create(unknown_loc_);
// Create the graph op.
auto builder = OpBuilder::atBlockBegin(module->getBody());
auto graph_op = GraphOp::create(builder, module->getLoc(),
ConvertVersionAttr(ctx_, graph.versions()));
graph_op.getNodes().push_back(new Block);
// Populate the function op defs.
function_op_defs_.reserve(graph.library().function_size());
for (const FunctionDef &function : graph.library().function()) {
function_op_defs_.emplace(function.signature().name(),
&function.signature());
}
// Build a map from function name to gradient function name.
absl::flat_hash_map<StringPiece, StringPiece> gradient_map;
gradient_map.reserve(graph.library().gradient_size());
for (const tensorflow::GradientDef &gradient : graph.library().gradient())
gradient_map.emplace(gradient.function_name(), gradient.gradient_func());
// Convert the graph.
ConversionState s(&graph_op.getNodes().front(), placeholder_state_);
TF_RETURN_IF_ERROR(
ConvertNodes(builder, s, graph.node(), &graph_op.getNodes().front()));
// A function to convert a generic or non-generic function.
const auto convert_func = [this, &gradient_map](GraphFuncOp func_op,
const FunctionDef &function) {
if (IsGenericFunction(function)) {
// Generic functions aren't on the hot path so just call the old
// importer.
OpBuilder builder(ctx_);
TF_RETURN_WITH_CONTEXT_IF_ERROR(
ConvertGenericFunction(func_op, function, builder),
"While importing generic function: ", function.signature().name());
} else {
TF_RETURN_WITH_CONTEXT_IF_ERROR(
ConvertFunctionDef(func_op, gradient_map, function),
"While importing function: ", function.signature().name());
}
return absl::OkStatus();
};
// TODO(jeffniu): Don't import functions in parallel if there are too few (how
// few?) or if the functions are too small (how small?).
if (ctx_->isMultithreadingEnabled()) {
ctx_->enterMultiThreadedExecution();
auto exit =
llvm::make_scope_exit([this] { ctx_->exitMultiThreadedExecution(); });
// Prepare the arguments to parallel for each.
struct Argument {
GraphFuncOp func;
const FunctionDef &def;
Status status;
};
std::vector<Argument> args;
args.reserve(graph.library().function_size());
for (const FunctionDef &function : graph.library().function()) {
args.push_back(
Argument{GraphFuncOp::create(builder, unknown_loc_), function});
}
const auto process_func = [&convert_func](Argument &arg) {
arg.status = convert_func(arg.func, arg.def);
return success(arg.status.ok());
};
// Execute the imports in parallel.
if (failed(failableParallelForEach(ctx_, args, process_func))) {
Status result;
for (const Argument &arg : args) {
result.Update(arg.status);
}
return result;
}
} else {
// Convert the functions.
for (const FunctionDef &function : graph.library().function()) {
auto func_op = GraphFuncOp::create(builder, unknown_loc_);
TF_RETURN_IF_ERROR(convert_func(func_op, function));
}
}
return module;
}
Status GraphDefImporter::ConvertFunctionAttributes(
const absl::flat_hash_map<StringPiece, StringPiece> &gradient_map,
const FunctionDef &function, GraphFuncOp op, NamedAttrList &attrs) {
// Import the function attributes with a `tf.` prefix to match the current
// infratructure expectations.
for (const auto &name_attr : function.attr()) {
if (name_attr.first.empty()) {
return absl::InvalidArgumentError(absl::StrCat(
"Function ", function.signature().name(), " has an empty attr name"));
}
// TODO(b/230143351): `ConvertAttributeValue` is a little slow due to
// `ConvertTensorProto` and `ConvertTensorShapeProto`.
TF_ASSIGN_OR_RETURN(Attribute attr,
ConvertAttributeValue(name_attr.second, b_));
attrs.append(absl::StrCat("tf.", name_attr.first), attr);
}
// Convert the first-class attributes.
const tensorflow::OpDef &signature = function.signature();
if (signature.name().empty())
return absl::InvalidArgumentError("Function without a name");
attrs.append(op.getSymNameAttrName(), b_.getStringAttr(signature.name()));
if (!signature.description().empty()) {
attrs.append(op.getDescriptionAttrName(),
b_.getStringAttr(signature.description()));
}
if (signature.is_stateful())
attrs.append(op.getIsStatefulAttrName(), b_.getUnitAttr());
auto grad_it = gradient_map.find(signature.name());
if (grad_it != gradient_map.end()) {
StringPiece name = grad_it->second;
attrs.append(op.getGradientAttrName(),
FlatSymbolRefAttr::get(ctx_, {name.data(), name.size()}));
}
// The resource_arg_unique_id is a list of `pair<int, int>`, we import it
// as two arrays of integer right now.
if (function.resource_arg_unique_id_size()) {
SmallVector<int32_t> resource_arg_unique_ids_keys;
SmallVector<int32_t> resource_arg_unique_ids_values;
resource_arg_unique_ids_keys.reserve(
function.resource_arg_unique_id_size());
resource_arg_unique_ids_values.reserve(
function.resource_arg_unique_id_size());
for (const auto &unique_id : function.resource_arg_unique_id()) {
resource_arg_unique_ids_keys.push_back(unique_id.first);
resource_arg_unique_ids_values.push_back(unique_id.second);
}
attrs.append(op.getResourceArgUniqueIdsKeysAttrName(),
b_.getI32TensorAttr(resource_arg_unique_ids_keys));
attrs.append(op.getResourceArgUniqueIdsValuesAttrName(),
b_.getI32TensorAttr(resource_arg_unique_ids_values));
}
return absl::OkStatus();
}
Status GraphDefImporter::ConvertArgumentAttributes(const OpDef::ArgDef &def,
NamedAttrList &attrs) {
attrs.append(dialect_->getTfgNameAttrIdentifier(),
b_.getStringAttr(def.name()));
if (!def.description().empty()) {
attrs.append(dialect_->getTfgDescriptionAttrIdentifier(),
b_.getStringAttr(def.description()));
}
if (def.is_ref())
attrs.append(dialect_->getTfgIsRefAttrIdentifier(), b_.getUnitAttr());
if (def.handle_data_size()) {
TF_ASSIGN_OR_RETURN(Attribute handle_data,
ConvertHandleData(b_, def.handle_data()));
attrs.append(dialect_->getTfgHandleDataAttrIdentifier(), handle_data);
}
if (def.has_experimental_full_type()) {
TF_ASSIGN_OR_RETURN(tf_type::FullTypeAttr full_type,
ConvertAttribute(def.experimental_full_type(), b_));
attrs.append(dialect_->getTfgFullTypeAttrIdentifier(), full_type);
}
return absl::OkStatus();
}
Location GraphDefImporter::ConvertLocation(const NodeDef &node) {
if (!node.has_experimental_debug_info()) return unknown_loc_;
const auto &debug_info = node.experimental_debug_info();
const auto &original_nodes = debug_info.original_node_names();
const auto &original_funcs = debug_info.original_func_names();
if (original_nodes.empty()) return unknown_loc_;
SmallVector<Location> node_locs;
node_locs.reserve(original_nodes.size());
for (const auto &it : llvm::enumerate(original_nodes)) {
std::string func_name =
it.index() < original_funcs.size() ? original_funcs[it.index()] : "";
node_locs.push_back(ConvertLocation(it.value(), func_name));
}
return b_.getFusedLoc(node_locs);
}
// This is a re-implementation of GetLocation in `import.cc`.
Location GraphDefImporter::ConvertLocation(StringRef node_name,
StringRef func_name) {
// Concatenate the node name with the function name to match how the key is
// formed in Python.
std::string debug_info_key = (node_name + "@" + func_name).str();
std::string name_loc = func_name.empty() ? node_name.str() : debug_info_key;
auto name_loc_id = b_.getStringAttr(name_loc);
SmallVector<Location> locs;
// Try to find a stack trace to convert to locations.
auto it = stack_traces_.find(name_loc);
if (it == stack_traces_.end()) {
it = stack_traces_.find(debug_info_key);
}
if (it != stack_traces_.end()) {
std::shared_ptr<tensorflow::AbstractStackTrace> trace = it->second;
auto frames = trace->ToFrames();
locs.reserve(frames.size());
for (const auto &frame : frames) {
auto file_attr = b_.getStringAttr(frame.file_name);
locs.push_back(FileLineColLoc::get(file_attr, frame.line_number, 1));
}
}
if (locs.empty()) return NameLoc::get(name_loc_id);
// Use the first location to generate a name location.
Location node_name_loc = NameLoc::get(name_loc_id, locs.front());
// Generate a stack trace using the remaining locations.
ArrayRef<Location> callsite_locs = llvm::ArrayRef(locs).drop_front();
return callsite_locs.empty() ? node_name_loc
: CallSiteLoc::get(node_name_loc, callsite_locs);
}
absl::StatusOr<Value> GraphDefImporter::ResolveDataResult(const ResultId &id,
ResultInfo *info) {
if (id.output.empty()) {
if (id.index >= info->data.size()) {
return absl::InvalidArgumentError(
absl::StrCat("Result #", id.index, " of node '", id.node.str(),
"' is out of bounds"));
}
return info->data[id.index];
}
auto it = info->outputs.find({id.output.data(), id.output.size()});
if (it == info->outputs.end()) {
return absl::InvalidArgumentError(absl::StrCat("Node '", id.node.str(),
"' has no output called '",
id.output.str(), "'"));
}
if (id.index >= it->second.size()) {
return absl::InvalidArgumentError(
absl::StrCat("Result #", id.index, " of segment '", id.node.str(), ":",
id.output.str(), "' is out of bounds"));
}
return it->second[id.index];
}
absl::StatusOr<GraphDefImporter::Result> GraphDefImporter::GetResult(
ConversionState &s, StringPiece name) {
TensorId tensor_id = tensorflow::ParseTensorName(name);
ResultId id{tensor_id.index()};
std::tie(id.node, id.output) =
StringRef(tensor_id.node().data(), tensor_id.node().size()).split(':');
std::unique_ptr<ResultInfo> &info = s[{id.node.data(), id.node.size()}];
if (!info) {
info = std::make_unique<ResultInfo>();
}
// If the result is unresolved, return the placeholder;
if (!info->resolved) {
if (id.IsControl()) {
return Result{s.GetPlaceholder(), Value(), id, info.get()};
}
return Result{Value(), s.GetPlaceholder(), id, info.get()};
}
// If the result is the control token, return it.
if (id.IsControl()) {
return Result{info->control, Value()};
}
TF_ASSIGN_OR_RETURN(Value value, ResolveDataResult(id, info.get()));
return Result{Value(), value};
}
Status GraphDefImporter::ConvertFunctionDef(
GraphFuncOp func_op,
const absl::flat_hash_map<StringPiece, StringPiece> &gradient_map,
const FunctionDef &function) {
const OpDef &signature = function.signature();
// TODO(jeffniu): Does the name need to be mangled?
func_op.getBody().push_back(new Block);
Block *body = &func_op.getBody().front();
auto builder = OpBuilder::atBlockBegin(func_op.SingleBlock::getBody());
// Convert the attributes.
NamedAttrList func_attrs;
TF_RETURN_IF_ERROR(
ConvertFunctionAttributes(gradient_map, function, func_op, func_attrs));
SmallVector<Attribute> arg_attrs, res_attrs, control_ret_attrs;
SmallVector<Type> arg_types, res_types;
// Convert the arguments and argument attributes.
for (const auto &it : llvm::enumerate(signature.input_arg())) {
Type dtype;
TF_RETURN_IF_ERROR(ConvertDataType(it.value().type(), b_, &dtype));
BlockArgument data =
body->addArgument(UnrankedTensorType::get(dtype), unknown_loc_);
BlockArgument ctl =
body->addArgument(dialect_->getControlType(), data.getLoc());
NamedAttrList attrs;
TF_RETURN_IF_ERROR(ConvertArgumentAttributes(it.value(), attrs));
auto attr_it = function.arg_attr().find(it.index());
if (attr_it != function.arg_attr().end()) {
for (const auto &name_attr : attr_it->second.attr()) {
TF_ASSIGN_OR_RETURN(Attribute attr,
ConvertAttributeValue(name_attr.second, b_));
attrs.append("tf." + name_attr.first, attr);
}
}
arg_attrs.append({attrs.getDictionary(ctx_), b_.getDictionaryAttr({})});
arg_types.append({data.getType(), ctl.getType()});
}
// Iterate over the arguments again and map them. We have to add them first
// otherwise the ranges will be invalidated.
ConversionState s(body, placeholder_state_);
for (const auto &it : llvm::enumerate(signature.input_arg())) {
s.emplace(it.value().name(),
absl::WrapUnique(new ResultInfo{
/*resolved=*/true, body->getArgument(it.index() * 2 + 1),
body->getArguments().slice(it.index() * 2, 1)}));
}
TF_RETURN_IF_ERROR(ConvertNodes(builder, s, function.node_def(), body));
// Convert the results and the result attributes.
SmallVector<Value> return_operands;
return_operands.reserve(signature.output_arg_size() +
signature.control_output_size());
for (const OpDef::ArgDef &def : function.signature().output_arg()) {
Type dtype;
TF_RETURN_IF_ERROR(ConvertDataType(def.type(), b_, &dtype));
NamedAttrList attrs;
TF_RETURN_IF_ERROR(ConvertArgumentAttributes(def, attrs));
res_attrs.push_back(attrs.getDictionary(ctx_));
res_types.push_back(UnrankedTensorType::get(dtype));
auto ret_it = function.ret().find(def.name());
if (ret_it == function.ret().end()) {
return absl::InvalidArgumentError(
absl::StrCat("Output '", def.name(), "' was not found in 'ret'"));
}
TF_ASSIGN_OR_RETURN(Result result, GetResult(s, ret_it->second));
if (result.info)
return absl::InvalidArgumentError(
absl::StrCat("Return '", ret_it->second, "' was not found"));
if (result.control)
return absl::InvalidArgumentError(
absl::StrCat("Unexpected control result: ", ret_it->second));
return_operands.push_back(result.data);
}
// Convert the control results.
for (const std::string &control_ret : signature.control_output()) {
auto ret_it = function.control_ret().find(control_ret);
if (ret_it == function.control_ret().end()) {
return absl::InvalidArgumentError(absl::StrCat(
"Control output '", control_ret, "' was not found in 'control_ret'"));
}
std::unique_ptr<ResultInfo> &result = s[ret_it->second];
if (!result || !result->resolved) {
return absl::InvalidArgumentError(
absl::StrCat("Control return ", ret_it->second, " was not found"));
}
return_operands.push_back(result->control);
control_ret_attrs.push_back(b_.getDictionaryAttr(NamedAttribute(
dialect_->getTfgNameAttrIdentifier(), b_.getStringAttr(control_ret))));
}
ReturnOp::create(builder, unknown_loc_, return_operands,
b_.getArrayAttr(control_ret_attrs));
// Finalize the function attributes.
func_attrs.append(func_op.getArgAttrsAttrName(), b_.getArrayAttr(arg_attrs));
func_attrs.append(func_op.getResAttrsAttrName(), b_.getArrayAttr(res_attrs));
func_attrs.append(func_op.getFunctionTypeAttrName(),
TypeAttr::get(b_.getFunctionType(arg_types, res_types)));
func_op->setAttrs(func_attrs.getDictionary(ctx_));
return absl::OkStatus();
}
Status GraphDefImporter::ConvertNodes(
OpBuilder &builder, ConversionState &s,
const tensorflow::protobuf::RepeatedPtrField<NodeDef> &nodes,
Block *block) {
OpBuilder::InsertionGuard ig(builder);
builder.setInsertionPointToStart(block);
for (const NodeDef &node : nodes) {
TF_RETURN_IF_ERROR(ConvertNodeDef(builder, s, node));
}
// If the placeholder has remaining uses, then an input is missing.
if (TF_PREDICT_FALSE(!s.GetPlaceholder().use_empty())) {
// Stringify a result ID.
const auto id_to_str = [](const ResultId &id) {
std::string name = id.node.str();
if (id.IsControl()) return absl::StrCat("^", name);
if (id.output.empty())
return id.index ? absl::StrCat(id.node.str(), ":", id.index) : name;
return absl::StrCat(name, ":", id.output.str(), ":", id.index);
};
// Gather all missing input edges.
std::vector<std::pair<std::string, std::string>> missing_edges;
for (const ResultInfo &info :
llvm::make_pointee_range(llvm::make_second_range(s))) {
if (info.backedges.empty()) continue;
const Backedge &edge = info.backedges.front();
missing_edges.emplace_back(id_to_str(edge.id),
TFOp(edge.operand->getOwner()).name().str());
}
assert(!missing_edges.empty() &&
"placeholder had remaining uses but found no unresolved backedges");
// Destroy the invalid IR.
block->erase();
// Report the missing edges in alphabetical order.
llvm::sort(missing_edges);
std::string error_message;
llvm::raw_string_ostream os(error_message);
llvm::interleave(
missing_edges, os,
[&](const auto &edge) {
os << "Non-existent input " << edge.first << " in node "
<< edge.second;
},
"\n");
return absl::InvalidArgumentError(os.str());
}
// The placeholder has no uses and should not acquire any more uses. Safely
// delete it from the IR.
s.Finalize();
return absl::OkStatus();
}
absl::StatusOr<unsigned int> GraphDefImporter::ArgNumType(
const NamedAttrList &attrs, const OpDef::ArgDef &arg_def,
SmallVectorImpl<Type> &types) {
// Check whether a type list attribute is specified.
if (!arg_def.type_list_attr().empty()) {
if (auto v = mlir::dyn_cast_or_null<ArrayAttr>(
attrs.get(arg_def.type_list_attr()))) {
for (Attribute attr : v) {
if (auto dtype = mlir::dyn_cast<TypeAttr>(attr)) {
types.push_back(UnrankedTensorType::get(dtype.getValue()));
} else {
return absl::InvalidArgumentError(
absl::StrCat("Expected '", arg_def.type_list_attr(),
"' to be a list of types"));
}
}
return v.size();
}
return absl::NotFoundError(
absl::StrCat("Type attr not found: ", arg_def.type_list_attr()));
}
unsigned num = 1;
// Check whether a number attribute is specified.
if (!arg_def.number_attr().empty()) {
if (auto v = mlir::dyn_cast_or_null<IntegerAttr>(
attrs.get(arg_def.number_attr()))) {
num = v.getValue().getZExtValue();
} else {
return absl::NotFoundError(
absl::StrCat("Type attr not found: ", arg_def.number_attr()));
}
}
// Check for a type or type attribute.
Type dtype;
if (arg_def.type() != DataType::DT_INVALID) {
TF_RETURN_IF_ERROR(ConvertDataType(arg_def.type(), b_, &dtype));
} else if (arg_def.type_attr().empty()) {
return absl::InvalidArgumentError(absl::StrCat(
"Arg '", arg_def.name(), "' has invalid type and no type attribute"));
} else {
if (auto v =
mlir::dyn_cast_or_null<TypeAttr>(attrs.get(arg_def.type_attr()))) {
dtype = v.getValue();
} else {
return absl::NotFoundError(
absl::StrCat("Type attr not found: ", arg_def.type_attr()));
}
}
types.append(num, UnrankedTensorType::get(dtype));
return num;
}
Status GraphDefImporter::ConvertNodeDef(OpBuilder &builder, ConversionState &s,
const NodeDef &node) {
VLOG(4) << "Importing: " << node.name();
if (node.op().empty())
return absl::InvalidArgumentError(
absl::StrCat("Node ", node.name(), " has an empty op name"));
OperationState state(ConvertLocation(node), absl::StrCat("tfg.", node.op()));
// The GraphImporter does light shape inference, but here we will defer all of
// that to the shape inference pass.
const OpDef *op_def;
const OpRegistrationData *op_reg_data = nullptr;
if ((op_reg_data = registry_.LookUp(node.op()))) {
op_def = &op_reg_data->op_def;
} else {
auto it = function_op_defs_.find(node.op());
if (it == function_op_defs_.end())
return absl::InvalidArgumentError(
absl::StrCat("Unable to find OpDef for ", node.op()));
op_def = it->second;
}
// Import the attributes. Reserve `+3` for `device`,`name`, and `fulltype`.
state.attributes.reserve(node.attr_size() + 3);
if (!node.device().empty()) {
state.addAttribute(dialect_->getDeviceAttrIdentifier(),
b_.getStringAttr(node.device()));
}
if (!node.name().empty()) {
state.addAttribute(dialect_->getNameAttrIdentifier(),
b_.getStringAttr(node.name()));
}
// If the op doesn't have a FullType, try to infer one.
const auto add_full_type = [&](const FullTypeDef &full_type_def) {
TF_ASSIGN_OR_RETURN(tf_type::FullTypeAttr full_type,
ConvertAttribute(full_type_def, b_));
state.addAttribute(dialect_->getFullTypeAttrIdentifier(), full_type);
return absl::OkStatus();
};
if (node.has_experimental_type()) {
TF_RETURN_IF_ERROR(add_full_type(node.experimental_type()));
} else if (op_reg_data && op_reg_data->type_ctor) {
FullTypeDef full_type_def;
TF_RETURN_IF_ERROR(
tensorflow::full_type::SpecializeType(node, *op_def, full_type_def));
TF_RETURN_IF_ERROR(add_full_type(full_type_def));
}
for (auto &name_attr : node.attr()) {
if (name_attr.first.empty())
return absl::InvalidArgumentError(
absl::StrCat("Node ", node.name(), " has an empty attr name"));
TF_ASSIGN_OR_RETURN(Attribute attr,
ConvertAttributeValue(name_attr.second, b_));
state.addAttribute(name_attr.first, attr);
}
// Add missing default attributes.
for (const auto &attr_def : op_def->attr()) {
if (attr_def.has_default_value() &&
!state.attributes.get(attr_def.name())) {
TF_ASSIGN_OR_RETURN(Attribute attr,
ConvertAttributeValue(attr_def.default_value(), b_));
state.addAttribute(attr_def.name(), attr);
}
}
// Get the result types. Ops can have multiple named results. Track the
// segment sizes.
SmallVector<std::pair<unsigned, unsigned>> result_segments;
if (op_def->output_arg_size() < 0)
return absl::InvalidArgumentError(
absl::StrCat("Node ", node.name(), " output arg size < 0"));
result_segments.reserve(op_def->output_arg_size());
state.types.reserve(op_def->output_arg_size() + 1);
for (const OpDef::ArgDef &def : op_def->output_arg()) {
unsigned index = state.types.size();
TF_ASSIGN_OR_RETURN(unsigned size,
ArgNumType(state.attributes, def, state.types));
result_segments.emplace_back(index, size);
}
state.types.push_back(dialect_->getControlType());
// Collect the operands. Set backedges to a placeholder and resolve them
// later.
state.operands.reserve(node.input_size());
SmallVector<Value> control_operands;
struct BackedgeResolution {
ResultInfo *info;
size_t operand_index;
ResultId id;
};
SmallVector<BackedgeResolution> unresolved_data_operands,
unresolved_control_operands;
for (const std::string &input : node.input()) {
TF_ASSIGN_OR_RETURN(Result result, GetResult(s, input));
if (result.control) {
if (result.info) {
unresolved_control_operands.push_back(BackedgeResolution{
result.info, control_operands.size(), result.id});
}
control_operands.push_back(result.control);
} else {
if (result.info) {
unresolved_data_operands.push_back(
BackedgeResolution{result.info, state.operands.size(), result.id});
}
state.operands.push_back(result.data);
}
}
unsigned num_data_operands = state.operands.size();
state.addOperands(control_operands);
// Create the op and record any unresolved operands.
Operation *op = builder.create(state);
for (const BackedgeResolution &r : unresolved_data_operands) {
r.info->backedges.push_back(
Backedge{r.id, &op->getOpOperand(r.operand_index)});
}
for (const BackedgeResolution &r : unresolved_control_operands) {
r.info->backedges.push_back(
Backedge{r.id, &op->getOpOperand(num_data_operands + r.operand_index)});
}
std::unique_ptr<ResultInfo> &info = s[node.name()];
if (!info) {
info = std::make_unique<ResultInfo>();
}
info->resolved = true;
info->control = *std::prev(op->result_end());
info->data = op->getResults().drop_back();
for (auto it : llvm::zip(result_segments, op_def->output_arg())) {
const std::pair<unsigned, unsigned> &segment = std::get<0>(it);
info->outputs.emplace(std::get<1>(it).name(),
info->data.slice(segment.first, segment.second));
}
// Resolve any associated backedges.
for (const Backedge &backedge : info->backedges) {
Value value;
if (backedge.id.IsControl()) {
value = info->control;
} else {
TF_ASSIGN_OR_RETURN(value, ResolveDataResult(backedge.id, info.get()));
}
backedge.operand->set(value);
}
info->backedges.clear();
return absl::OkStatus();
}
Status GraphDefImporter::ConvertDataTypesToUnrankedTensorTypes(
const DataTypeVector &dtypes, SmallVectorImpl<Type> &results) {
Type dtype;
for (DataType tf_dtype : dtypes) {
TF_RETURN_IF_ERROR(ConvertDataType(tf_dtype, b_, &dtype));
results.push_back(UnrankedTensorType::get(dtype));
}
return absl::OkStatus();
}
absl::StatusOr<OwningOpRef<ModuleOp>> ImportGraphDef(
MLIRContext *context, const GraphDebugInfo &debug_info,
const GraphDef &graph_def) {
GraphDefImporter importer(context->getOrLoadDialect<TFGraphDialect>(),
*OpRegistry::Global(), debug_info);
return importer.ConvertGraphDef(graph_def);
}
absl::StatusOr<OwningOpRef<ModuleOp>> ImportGraphAndFunctionsToMlir(
MLIRContext *context, const GraphDebugInfo &debug_info, const Graph &graph,
const FunctionLibraryDefinition &flib_def) {
// TODO(b/231723721): This conversion path is slow because both the graph and
// the function library are converted to GraphDef.
GraphDef graph_def;
graph.ToGraphDef(&graph_def);
*graph_def.mutable_library() = flib_def.ToProto();
return ImportGraphDef(context, debug_info, graph_def);
}
} // namespace tfg
} // namespace mlir
tensorflow/core/ir/importexport/graphdef_import.h
+45
View File
@@ -0,0 +1,45 @@
/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_CORE_IR_IMPORTEXPORT_GRAPHDEF_IMPORT_H_
#define TENSORFLOW_CORE_IR_IMPORTEXPORT_GRAPHDEF_IMPORT_H_
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "mlir/IR/OwningOpRef.h" // from @llvm-project
#include "tensorflow/core/framework/function.h"
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/framework/graph_debug_info.pb.h"
#include "tensorflow/core/graph/graph.h"
#include "tensorflow/core/platform/statusor.h"
namespace mlir {
namespace tfg {
// Convert a GraphDef directly to TFG.
absl::StatusOr<OwningOpRef<ModuleOp>> ImportGraphDef(
MLIRContext *context, const tensorflow::GraphDebugInfo &debug_info,
const tensorflow::GraphDef &graph_def);
// Converts a graph and function library to a TFG module.
absl::StatusOr<OwningOpRef<ModuleOp>> ImportGraphAndFunctionsToMlir(
MLIRContext *context, const tensorflow::GraphDebugInfo &debug_info,
const tensorflow::Graph &graph,
const tensorflow::FunctionLibraryDefinition &flib_def);
} // namespace tfg
} // namespace mlir
#endif // TENSORFLOW_CORE_IR_IMPORTEXPORT_GRAPHDEF_IMPORT_H_
tensorflow/core/ir/importexport/load_proto.cc
+77
View File
@@ -0,0 +1,77 @@
/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/core/ir/importexport/load_proto.h"
#include "llvm/Support/FileUtilities.h"
#include "llvm/Support/MemoryBuffer.h"
#include "llvm/Support/ToolOutputFile.h"
#include "llvm/Support/raw_ostream.h"
#include "tensorflow/core/ir/importexport/parse_text_proto.h"
#include "tensorflow/core/lib/core/errors.h"
#include "tensorflow/core/lib/core/status.h"
#include "tensorflow/core/platform/protobuf.h"
namespace tensorflow {
namespace {
inline llvm::StringRef StringViewToRef(absl::string_view view) {
return {view.data(), view.size()};
}
} // namespace
absl::Status LoadProtoFromBuffer(absl::string_view input,
protobuf::Message* proto) {
// Attempt to parse as text.
if (mlir::tfg::ParseTextProto(input, "", proto).ok()) return absl::OkStatus();
// Else attempt to parse as binary.
return LoadProtoFromBuffer(input, static_cast<protobuf::MessageLite*>(proto));
}
absl::Status LoadProtoFromBuffer(absl::string_view input,
protobuf::MessageLite* proto) {
// Attempt to parse as binary.
protobuf::io::ArrayInputStream binary_stream(input.data(), input.size());
if (proto->ParseFromZeroCopyStream(&binary_stream)) return absl::OkStatus();
LOG(ERROR) << "Error parsing Protobuf";
return absl::InvalidArgumentError("Could not parse input proto");
}
template <class T>
absl::Status LoadProtoFromFileImpl(absl::string_view input_filename, T* proto) {
const auto file_or_err =
llvm::MemoryBuffer::getFileOrSTDIN(StringViewToRef(input_filename));
if (std::error_code error = file_or_err.getError()) {
return absl::InvalidArgumentError(
absl::StrCat("Could not open input file ", input_filename));
}
const auto& input_file = *file_or_err;
absl::string_view content(input_file->getBufferStart(),
input_file->getBufferSize());
return LoadProtoFromBuffer(content, proto);
}
absl::Status LoadProtoFromFile(absl::string_view input_filename,
protobuf::Message* proto) {
return LoadProtoFromFileImpl(input_filename, proto);
}
absl::Status LoadProtoFromFile(absl::string_view input_filename,
protobuf::MessageLite* proto) {
return LoadProtoFromFileImpl(input_filename, proto);
}
} // namespace tensorflow
tensorflow/core/ir/importexport/load_proto.h
+45
View File
@@ -0,0 +1,45 @@
/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_CORE_IR_IMPORTEXPORT_LOAD_PROTO_H_
#define TENSORFLOW_CORE_IR_IMPORTEXPORT_LOAD_PROTO_H_
#include "absl/strings/string_view.h"
#include "tensorflow/core/lib/core/status.h"
#include "tensorflow/core/platform/protobuf.h"
namespace tensorflow {
// Reads text (.pbtext) or binary (.pb) format of a proto message from the given
// buffer. Returns error status of the file is not found or malformed proto.
// Note that text protos can only be parsed when full protobuf::Message protos
// are used, and will fail for protobuf::MessageLite protos.
absl::Status LoadProtoFromBuffer(absl::string_view input,
protobuf::Message* proto);
absl::Status LoadProtoFromBuffer(absl::string_view input,
protobuf::MessageLite* proto);
// Reads text (.pbtext) or binary (.pb) format of a proto message from the given
// file path. Returns error status of the file is not found or malformed proto.
// Note that text protos can only be parsed when full protobuf::Message protos
// are used, and will fail for protobuf::MessageLite protos.
absl::Status LoadProtoFromFile(absl::string_view input_filename,
protobuf::Message* proto);
absl::Status LoadProtoFromFile(absl::string_view input_filename,
protobuf::MessageLite* proto);
} // namespace tensorflow
#endif // TENSORFLOW_CORE_IR_IMPORTEXPORT_LOAD_PROTO_H_
tensorflow/core/ir/importexport/mangling.cc
+131
View File
@@ -0,0 +1,131 @@
/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/core/ir/importexport/mangling.h"
#include <cstring>
#include <string>
#include "absl/strings/match.h"
#include "absl/strings/str_cat.h"
#include "tensorflow/core/framework/tensor.pb.h"
#include "tensorflow/core/framework/tensor_shape.pb.h"
#include "tensorflow/core/framework/types.h"
#include "tensorflow/core/ir/importexport/parse_text_proto.h"
#include "tensorflow/core/lib/core/errors.h"
#include "tensorflow/core/platform/protobuf.h"
using tensorflow::DataType;
using tensorflow::Status;
using tensorflow::TensorProto;
using tensorflow::TensorShapeProto;
using tensorflow::errors::FailedPrecondition;
namespace mlir {
namespace tfg {
namespace mangling_util {
namespace {
const char kAttributePrefix[] = "tf.";
const char kDataTypePrefix[] = "tfdtype$";
const char kTensorShapePrefix[] = "tfshape$";
const char kTensorPrefix[] = "tftensor$";
} // namespace
std::string PrintShortTextProto(
const ::tensorflow::protobuf::MessageLite& message) {
// google::protobuf::TextFormat::Printer::PrintToString does not have
// a overload for MessageLite so here to be consistent with the existing
// behavior we use MessageLite::ShortDebugString().
return message.ShortDebugString();
}
std::string PrintShortTextProto(
const ::tensorflow::protobuf::Message& message) {
std::string message_short_text;
::tensorflow::protobuf::TextFormat::Printer printer;
printer.SetSingleLineMode(true);
printer.SetExpandAny(true);
printer.PrintToString(message, &message_short_text);
// Single line mode currently might have an extra space at the end.
if (!message_short_text.empty() && message_short_text.back() == ' ') {
message_short_text.pop_back();
}
return message_short_text;
}
std::string MangleAttributeName(absl::string_view str) {
return absl::StrCat(kAttributePrefix, str);
}
bool IsMangledAttributeName(absl::string_view str) {
return absl::StartsWith(str, kAttributePrefix);
}
absl::string_view DemangleAttributeName(absl::string_view str) {
DCHECK(IsMangledAttributeName(str));
return str.substr(std::strlen(kAttributePrefix));
}
MangledKind GetMangledKind(absl::string_view str) {
if (absl::StartsWith(str, kDataTypePrefix)) {
return MangledKind::kDataType;
} else if (absl::StartsWith(str, kTensorShapePrefix)) {
return MangledKind::kTensorShape;
} else if (absl::StartsWith(str, kTensorPrefix)) {
return MangledKind::kTensor;
} else {
return MangledKind::kUnknown;
}
}
std::string MangleShape(const TensorShapeProto& shape) {
return absl::StrCat(kTensorShapePrefix, shape.ShortDebugString());
}
Status DemangleShape(absl::string_view str, TensorShapeProto* proto) {
return ParseTextProto(str, kTensorShapePrefix, proto);
}
std::string MangleTensor(const TensorProto& tensor) {
return absl::StrCat(kTensorPrefix, PrintShortTextProto(tensor));
}
Status DemangleTensor(absl::string_view str, TensorProto* proto) {
return ParseTextProto(str, kTensorPrefix, proto);
}
std::string MangleDataType(const DataType& dtype) {
return absl::StrCat(kDataTypePrefix, DataType_Name(dtype));
}
Status DemangleDataType(absl::string_view str, DataType* proto) {
absl::string_view pbtxt;
TF_RETURN_IF_ERROR(ConsumePrefix(str, kDataTypePrefix, &pbtxt));
// NOLINTNEXTLINE: redundant string conversion for divergence in OSS API.
if (!DataType_Parse(std::string(pbtxt), proto)) {
return absl::FailedPreconditionError(
"Could not parse TFDataType mangled proto");
}
return absl::OkStatus();
}
} // namespace mangling_util
} // namespace tfg
} // namespace mlir
tensorflow/core/ir/importexport/mangling.h
+76
View File
@@ -0,0 +1,76 @@
/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_CORE_IR_IMPORTEXPORT_MANGLING_H_
#define TENSORFLOW_CORE_IR_IMPORTEXPORT_MANGLING_H_
#include <string>
#include "absl/strings/string_view.h"
#include "tensorflow/core/framework/tensor.pb.h"
#include "tensorflow/core/framework/tensor_shape.pb.h"
#include "tensorflow/core/framework/types.h"
#include "tensorflow/core/lib/core/status.h"
namespace mlir {
namespace tfg {
namespace mangling_util {
// The type of a mangled string.
enum class MangledKind { kUnknown, kDataType, kTensorShape, kTensor };
// Print proto in TextFormat in the single-line mode.
std::string PrintShortTextProto(const ::tensorflow::protobuf::Message& message);
// The MessageLite interface does not support reflection so this API
// will only print a summary of the proto. This API is needed for code
// that may work with both Message and MessageLite.
std::string PrintShortTextProto(
const ::tensorflow::protobuf::MessageLite& message);
// Mangles an attribute name, marking the attribute as a TensorFlow attribute.
std::string MangleAttributeName(absl::string_view str);
// Returns true if 'str' was mangled with MangleAttributeName.
bool IsMangledAttributeName(absl::string_view str);
// Demangles an attribute name that was manged with MangleAttributeName.
// REQUIRES: IsMangledAttributeName returns true.
absl::string_view DemangleAttributeName(absl::string_view str);
// Returns the type of a mangled string, or kUnknown.
MangledKind GetMangledKind(absl::string_view str);
// Return a TensorShapeProto mangled as a string.
std::string MangleShape(const tensorflow::TensorShapeProto& shape);
// Demangle a string mangled with MangleShape.
absl::Status DemangleShape(absl::string_view str,
tensorflow::TensorShapeProto* proto);
// Return a TensorProto mangled as a string.
std::string MangleTensor(const tensorflow::TensorProto& tensor);
// Demangle a string mangled with MangleTensor.
absl::Status DemangleTensor(absl::string_view str,
tensorflow::TensorProto* proto);
// Return a DataType mangled as a string.
std::string MangleDataType(const tensorflow::DataType& dtype);
// Demangle a string mangled with MangleDataType.
absl::Status DemangleDataType(absl::string_view str,
tensorflow::DataType* proto);
} // namespace mangling_util
} // namespace tfg
} // namespace mlir
#endif // TENSORFLOW_CORE_IR_IMPORTEXPORT_MANGLING_H_
tensorflow/core/ir/importexport/parse_text_proto.cc
+79
View File
@@ -0,0 +1,79 @@
/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/core/ir/importexport/parse_text_proto.h"
#include <string>
#include "absl/strings/match.h"
#include "absl/strings/string_view.h"
#include "google/protobuf/io/tokenizer.h"
#include "tensorflow/core/lib/core/errors.h"
#include "tensorflow/core/lib/core/status.h"
#include "tensorflow/core/platform/casts.h"
#include "tensorflow/core/platform/protobuf.h"
using tensorflow::Status;
using tensorflow::errors::InvalidArgument;
using tensorflow::errors::NotFound;
namespace mlir {
namespace tfg {
namespace {
// Error collector that simply ignores errors reported.
class NoOpErrorCollector : public tensorflow::protobuf::io::ErrorCollector {
public:
void RecordError(int line, tensorflow::protobuf::io::ColumnNumber column,
absl::string_view message) override {}
};
} // namespace
Status ConsumePrefix(absl::string_view str, absl::string_view prefix,
absl::string_view* output) {
if (absl::StartsWith(str, prefix)) {
*output = str.substr(prefix.size());
return absl::OkStatus();
}
return absl::NotFoundError(
absl::StrCat("No prefix \"", prefix, "\" in \"", str, "\""));
}
Status ParseTextProto(absl::string_view text_proto,
absl::string_view prefix_to_strip,
tensorflow::protobuf::Message* parsed_proto) {
tensorflow::protobuf::TextFormat::Parser parser;
// Don't produce errors when attempting to parse text format as it would fail
// when the input is actually a binary file.
NoOpErrorCollector collector;
parser.RecordErrorsTo(&collector);
// Attempt to parse as text.
absl::string_view text_proto_without_prefix = text_proto;
if (!prefix_to_strip.empty()) {
TF_RETURN_IF_ERROR(
ConsumePrefix(text_proto, prefix_to_strip, &text_proto_without_prefix));
}
tensorflow::protobuf::io::ArrayInputStream input_stream(
text_proto_without_prefix.data(), text_proto_without_prefix.size());
if (parser.Parse(&input_stream, parsed_proto)) {
return absl::OkStatus();
}
parsed_proto->Clear();
return absl::InvalidArgumentError(
absl::StrCat("Could not parse text proto: ", text_proto));
}
} // namespace tfg
} // namespace mlir
tensorflow/core/ir/importexport/parse_text_proto.h
+46
View File
@@ -0,0 +1,46 @@
/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_CORE_IR_IMPORTEXPORT_PARSE_TEXT_PROTO_H_
#define TENSORFLOW_CORE_IR_IMPORTEXPORT_PARSE_TEXT_PROTO_H_
#include "absl/strings/string_view.h"
#include "tensorflow/core/lib/core/errors.h"
#include "tensorflow/core/lib/core/status.h"
#include "tensorflow/core/platform/protobuf.h"
namespace mlir {
namespace tfg {
// Sets output to the given input with `prefix` stripped, or returns an error if
// the prefix doesn't exist.
absl::Status ConsumePrefix(absl::string_view str, absl::string_view prefix,
absl::string_view* output);
// Strips `prefix_to_strip` from `text_proto`, parses, and returns the parsed
// proto.
absl::Status ParseTextProto(absl::string_view text_proto,
absl::string_view prefix_to_strip,
tensorflow::protobuf::Message* parsed_proto);
inline absl::Status ParseTextProto(
absl::string_view /* text_proto */, absl::string_view /* prefix_to_strip */,
tensorflow::protobuf::MessageLite* /* parsed_proto */) {
return absl::UnavailableError("Cannot parse text protos on mobile.");
}
} // namespace tfg
} // namespace mlir
#endif // TENSORFLOW_CORE_IR_IMPORTEXPORT_PARSE_TEXT_PROTO_H_
tensorflow/core/ir/importexport/savedmodel_export.cc
+50
View File
@@ -0,0 +1,50 @@
/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/core/ir/importexport/savedmodel_export.h"
#include <utility>
#include "tensorflow/core/ir/importexport/graphdef_export.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/protobuf/meta_graph.pb.h"
namespace mlir {
namespace tfg {
absl::Status ExportMlirToSavedModel(
mlir::ModuleOp module, const tensorflow::SavedModel &original_saved_model,
tensorflow::SavedModel *output_saved_model) {
if (original_saved_model.meta_graphs_size() == 0) {
return absl::InvalidArgumentError(
"Original saved model has no meta graphs");
}
tensorflow::GraphDef new_graphdef;
TF_RETURN_WITH_CONTEXT_IF_ERROR(ConvertToGraphDef(module, &new_graphdef),
"while converting TFG to GraphDef");
// Overwrite the graph def portion of the saved model with the new one.
tensorflow::MetaGraphDef meta_graph_def = original_saved_model.meta_graphs(0);
*(meta_graph_def.mutable_graph_def()) = std::move(new_graphdef);
*output_saved_model = original_saved_model;
*(output_saved_model->mutable_meta_graphs(0)) = std::move(meta_graph_def);
return absl::OkStatus();
}
} // namespace tfg
} // namespace mlir
tensorflow/core/ir/importexport/savedmodel_export.h
+39
View File
@@ -0,0 +1,39 @@
/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_CORE_IR_IMPORTEXPORT_SAVEDMODEL_EXPORT_H_
#define TENSORFLOW_CORE_IR_IMPORTEXPORT_SAVEDMODEL_EXPORT_H_
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "tensorflow/core/framework/graph_debug_info.pb.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/protobuf/saved_model.pb.h"
namespace mlir {
namespace tfg {
// Given an MLIR module, returns a `output_saved_model` SavedModel.
// The module must contain at most a single Graph operation and zero or more
// TFFunc operations. `original_saved_model` is used as only a GraphDef portion
// of a saved model represented in the MLIR module.
absl::Status ExportMlirToSavedModel(
mlir::ModuleOp module, const tensorflow::SavedModel &original_saved_model,
tensorflow::SavedModel *output_saved_model);
} // namespace tfg
} // namespace mlir
#endif // TENSORFLOW_CORE_IR_IMPORTEXPORT_SAVEDMODEL_EXPORT_H_
tensorflow/core/ir/importexport/savedmodel_import.cc
+43
View File
@@ -0,0 +1,43 @@
/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/core/ir/importexport/savedmodel_import.h"
#include "tensorflow/core/ir/importexport/graphdef_import.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/protobuf/meta_graph.pb.h"
namespace mlir {
namespace tfg {
absl::StatusOr<OwningOpRef<mlir::ModuleOp>> ImportSavedModelToMlir(
mlir::MLIRContext *context, const tensorflow::GraphDebugInfo &debug_info,
const tensorflow::SavedModel &saved_model) {
if (saved_model.meta_graphs_size() == 0) {
return absl::InvalidArgumentError("Input saved model has no meta graphs");
}
if (saved_model.meta_graphs_size() > 1) {
return absl::InvalidArgumentError(
"Input saved model has more than one meta graph, currently not "
"supported");
}
const auto &graphdef = saved_model.meta_graphs(0).graph_def();
return ImportGraphDef(context, debug_info, graphdef);
}
} // namespace tfg
} // namespace mlir
tensorflow/core/ir/importexport/savedmodel_import.h
+40
View File
@@ -0,0 +1,40 @@
/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_CORE_IR_IMPORTEXPORT_SAVEDMODEL_IMPORT_H_
#define TENSORFLOW_CORE_IR_IMPORTEXPORT_SAVEDMODEL_IMPORT_H_
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "mlir/IR/OwningOpRef.h" // from @llvm-project
#include "tensorflow/core/framework/graph_debug_info.pb.h"
#include "tensorflow/core/platform/statusor.h"
#include "tensorflow/core/protobuf/saved_model.pb.h"
namespace mlir {
namespace tfg {
// Converts a saved model to a MLIR module expressed in TFG dialect.
// Only the root graph and function library of the saved model gets imported
// into MLIR TFG dialect.
// TODO(b/218882780): Consider importing SignatureDefs from the SavedModel.
absl::StatusOr<OwningOpRef<mlir::ModuleOp>> ImportSavedModelToMlir(
mlir::MLIRContext* context, const tensorflow::GraphDebugInfo& debug_info,
const tensorflow::SavedModel& saved_model);
} // namespace tfg
} // namespace mlir
#endif // TENSORFLOW_CORE_IR_IMPORTEXPORT_SAVEDMODEL_IMPORT_H_
tensorflow/core/ir/importexport/tests/BUILD
+33
View File
@@ -0,0 +1,33 @@
load("//tensorflow:tensorflow.default.bzl", "get_compatible_with_portable")
load("//tensorflow/compiler/mlir:glob_lit_test.bzl", "glob_lit_tests")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
default_compatible_with = get_compatible_with_portable(),
default_visibility = [":__subpackages__"],
licenses = ["notice"], # Apache 2.0
)
# Bundle together all of the test utilities that are used by tests.
exports_files(["run_lit.sh"])
filegroup(
name = "test_utilities",
testonly = True,
data = [
"//tensorflow/core/ir/importexport:tfg-translate",
"@llvm-project//llvm:FileCheck",
"@llvm-project//llvm:not",
],
)
glob_lit_tests(
name = "all_tests",
data = [":test_utilities"],
driver = "@llvm-project//mlir:run_lit.sh",
exclude = [],
test_file_exts = [
"mlir",
"pbtxt",
],
)
tensorflow/core/ir/importexport/tests/graphdef_to_mlir/BUILD
+33
View File
@@ -0,0 +1,33 @@
load("//tensorflow:tensorflow.default.bzl", "get_compatible_with_portable")
load("//tensorflow/compiler/mlir:glob_lit_test.bzl", "glob_lit_tests")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
default_compatible_with = get_compatible_with_portable(),
default_visibility = [":__subpackages__"],
licenses = ["notice"], # Apache 2.0
)
# Bundle together all of the test utilities that are used by tests.
exports_files(["run_lit.sh"])
filegroup(
name = "test_utilities",
testonly = True,
data = [
"//tensorflow/core/ir/importexport:tfg-translate",
"@llvm-project//llvm:FileCheck",
"@llvm-project//llvm:not",
],
)
glob_lit_tests(
name = "all_tests",
data = [":test_utilities"],
driver = "@llvm-project//mlir:run_lit.sh",
exclude = [],
test_file_exts = [
"mlir",
"pbtxt",
],
)
tensorflow/core/ir/importexport/tests/graphdef_to_mlir/arg_as_control_ret.pbtxt
+66
View File
@@ -0,0 +1,66 @@
# RUN: tfg-translate --graphdef-to-mlir %s | FileCheck %s
# CHECK: return
# CHECK-SAME: %a.ctl
library {
function {
signature {
name: "test"
input_arg {
name: "a"
type: DT_INT32
}
output_arg {
name: "b"
type: DT_INT32
}
control_output: "const"
control_output: "a"
}
node_def {
name: "const"
op: "Const"
attr {
key: "dtype"
value {
type: DT_INT32
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_INT32
tensor_shape {
}
}
}
}
}
ret {
key: "b"
value: "const:output:0"
}
control_ret {
key: "a"
value: "a"
}
control_ret {
key: "const"
value: "const"
}
arg_attr {
key: 0
value {
attr {
key: "_output_shapes"
value {
shape {
}
}
}
}
}
}
}
tensorflow/core/ir/importexport/tests/graphdef_to_mlir/backedge_segment.pbtxt
+234
View File
@@ -0,0 +1,234 @@
# RUN: tfg-translate --graphdef-to-mlir %s | FileCheck %s
# Test that result segments are correctly resolved even in backedges.
# CHECK: %[[SWITCH:.*]]:2, %{{.*}} = Switch{{.*}}name("cond/Switch")
# CHECK: Identity(%[[SWITCH]]#1) name("cond/switch_t")
# CHECK: %[[IDENTITY:.*]], %{{.*}} = Identity{{.*}}name("cond/pred_id")
# CHECK: Switch(%[[SWITCH7:.*]]#1, %{{.*}}) name("cond/Switch_1")
# CHECK: %[[SWITCH7]]:2, %{{.*}} = Switch{{.*}}name("cond/Switch_1/Switch")
library {
function {
signature {
name: "foo"
input_arg {
name: "pred"
type: DT_BOOL
}
input_arg {
name: "cond_switch_1_placeholder"
type: DT_INT32
}
output_arg {
name: "cond_merge"
type: DT_INT32
}
}
node_def {
name: "cond/Switch"
op: "Switch"
input: "pred"
input: "pred"
attr {
key: "T"
value {
type: DT_BOOL
}
}
}
node_def {
name: "cond/switch_t"
op: "Identity"
input: "cond/Switch:output_true:0"
attr {
key: "T"
value {
type: DT_BOOL
}
}
}
node_def {
name: "cond/switch_f"
op: "Identity"
input: "cond/Switch:output_false:0"
attr {
key: "T"
value {
type: DT_BOOL
}
}
}
node_def {
name: "cond/pred_id"
op: "Identity"
input: "pred"
attr {
key: "T"
value {
type: DT_BOOL
}
}
}
node_def {
name: "cond/Switch_1"
op: "Switch"
input: "cond/Switch_1/Switch:output_true:0"
input: "cond/pred_id:output:0"
attr {
key: "T"
value {
type: DT_INT32
}
}
attr {
key: "_class"
value {
list {
s: "loc:@Const"
}
}
}
}
node_def {
name: "cond/Switch_1/Switch"
op: "Switch"
input: "cond_switch_1_placeholder"
input: "cond/pred_id:output:0"
attr {
key: "T"
value {
type: DT_INT32
}
}
attr {
key: "_class"
value {
list {
s: "loc:@Const"
}
}
}
}
node_def {
name: "cond/add/y"
op: "Const"
input: "^cond/switch_f"
attr {
key: "dtype"
value {
type: DT_INT32
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_INT32
tensor_shape {
}
int_val: 1
}
}
}
}
node_def {
name: "cond/add"
op: "AddV2"
input: "cond/add/Switch:output_false:0"
input: "cond/add/y:output:0"
attr {
key: "T"
value {
type: DT_INT32
}
}
}
node_def {
name: "cond/add/Switch"
op: "Switch"
input: "cond_switch_1_placeholder"
input: "cond/pred_id:output:0"
attr {
key: "T"
value {
type: DT_INT32
}
}
attr {
key: "_class"
value {
list {
s: "loc:@Const"
}
}
}
}
node_def {
name: "cond/Merge"
op: "Merge"
input: "cond/add:z:0"
input: "cond/Switch_1:output_true:0"
attr {
key: "N"
value {
i: 2
}
}
attr {
key: "T"
value {
type: DT_INT32
}
}
}
ret {
key: "cond_merge"
value: "cond/Merge:output:0"
}
attr {
key: "_disable_call_shape_inference"
value {
b: true
}
}
arg_attr {
key: 0
value {
attr {
key: "_output_shapes"
value {
list {
shape {
unknown_rank: true
}
}
}
}
}
}
arg_attr {
key: 1
value {
attr {
key: "_class"
value {
list {
s: "loc:@Const"
}
}
}
attr {
key: "_output_shapes"
value {
list {
shape {
}
}
}
}
}
}
}
}
tensorflow/core/ir/importexport/tests/graphdef_to_mlir/empty.pbtxt
+4
View File
@@ -0,0 +1,4 @@
# RUN: tfg-translate -graphdef-to-mlir %s | FileCheck %s
# CHECK: tfg.graph #tf_type.version<producer = 0, min_consumer = 0> {
# CHECK-NEXT: }
tensorflow/core/ir/importexport/tests/graphdef_to_mlir/error_during_backedge.pbtxt
+67
View File
@@ -0,0 +1,67 @@
# RUN: not tfg-translate -graphdef-to-mlir %s 2>&1 | FileCheck %s
# Test that an error while a backedge is unresolved (and thus the placeholder
# has uses) that causes the importer to exit does not trigger an assert.
# CHECK: Unable to find OpDef for __UnknownOp
node {
name: "merge"
op: "Merge"
# Both inputs are created as backedges.
input: "const1:0"
input: "const2:1"
attr {
key: "N"
value {
i: 2
}
}
attr {
key: "T"
value {
type: DT_DOUBLE
}
}
}
node {
name: "const1"
# Trigger a failure
op: "__UnknownOp"
attr {
key: "dtype"
value {
type: DT_INT32
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_INT32
tensor_shape {
}
}
}
}
}
node {
name: "const2"
op: "Const"
attr {
key: "dtype"
value {
type: DT_INT32
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_INT32
tensor_shape {
}
}
}
}
}
tensorflow/core/ir/importexport/tests/graphdef_to_mlir/import_case_with_attr_inference.pbtxt
+107
View File
@@ -0,0 +1,107 @@
# RUN: tfg-translate -graphdef-to-mlir %s | FileCheck %s
# CHECK: Case
# CHECK-SAME: -> (tensor<*xi32>)
node {
name: "Const"
op: "Const"
attr {
key: "dtype"
value {
type: DT_INT32
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_INT32
tensor_shape {
}
int_val: 1
}
}
}
}
node {
name: "Const_1"
op: "Const"
attr {
key: "dtype"
value {
type: DT_INT32
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_INT32
tensor_shape {
}
int_val: 0
}
}
}
}
node {
name: "indexed_case"
op: "StatelessCase"
input: "Const_1"
input: "Const"
attr {
key: "Tin"
value {
list {
type: DT_INT32
}
}
}
attr {
key: "Tout"
value {
list {
type: DT_INT32
}
}
}
attr {
key: "_lower_using_switch_merge"
value {
b: true
}
}
attr {
key: "_read_only_resource_inputs"
value {
list {
}
}
}
attr {
key: "branches"
value {
list {
func {
name: "indexed_case_branch0_4"
}
func {
name: "indexed_case_branch1_5"
}
}
}
}
attr {
key: "output_shapes"
value {
list {
shape {
dim {
size: 4
}
}
}
}
}
}
tensorflow/core/ir/importexport/tests/graphdef_to_mlir/import_if_with_attr_inference.pbtxt
+336
View File
@@ -0,0 +1,336 @@
# RUN: tfg-translate -graphdef-to-mlir %s | FileCheck %s
# CHECK: If
# CHECK-SAME: else_branch = #tf_type.func<@XTimesTwo, {}>,
# CHECK-SAME: then_branch = #tf_type.func<@IfWithWhileThen, {}>
# CHECK-SAME: -> (tensor<*xi32>)
# CHECK: func generic @XTimesTwo
node {
name: "pred"
op: "Placeholder"
attr {
key: "dtype"
value {
type: DT_BOOL
}
}
attr {
key: "shape"
value {
shape {
unknown_rank: true
}
}
}
}
node {
name: "A"
op: "Placeholder"
attr {
key: "dtype"
value {
type: DT_INT32
}
}
attr {
key: "shape"
value {
shape {
unknown_rank: true
}
}
}
}
node {
name: "if"
op: "If"
input: "pred"
input: "A"
attr {
key: "Tcond"
value {
type: DT_BOOL
}
}
attr {
key: "Tin"
value {
list {
type: DT_INT32
}
}
}
attr {
key: "Tout"
value {
list {
type: DT_INT32
}
}
}
attr {
key: "_lower_using_switch_merge"
value {
b: true
}
}
attr {
key: "else_branch"
value {
func {
name: "XTimesTwo"
}
}
}
attr {
key: "output_shapes"
value {
list {
shape {
dim {
size: 4
}
}
}
}
}
attr {
key: "then_branch"
value {
func {
name: "IfWithWhileThen"
}
}
}
}
library {
function {
signature {
name: "LessThanOrEqualToN"
input_arg {
name: "x"
type_attr: "T"
}
output_arg {
name: "z"
type: DT_BOOL
}
attr {
name: "T"
type: "type"
allowed_values {
list {
type: DT_FLOAT
type: DT_DOUBLE
type: DT_INT32
type: DT_INT64
}
}
}
}
node_def {
name: "N"
op: "Const"
attr {
key: "dtype"
value {
type: DT_INT64
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_INT64
tensor_shape {
}
int64_val: 8
}
}
}
}
node_def {
name: "y"
op: "Cast"
input: "N:output:0"
attr {
key: "DstT"
value {
placeholder: "T"
}
}
attr {
key: "SrcT"
value {
type: DT_INT64
}
}
}
node_def {
name: "z"
op: "LessEqual"
input: "x"
input: "y:y:0"
attr {
key: "T"
value {
placeholder: "T"
}
}
}
ret {
key: "z"
value: "z:z:0"
}
}
function {
signature {
name: "XTimesTwo"
input_arg {
name: "x"
type_attr: "T"
}
output_arg {
name: "y"
type_attr: "T"
}
attr {
name: "T"
type: "type"
allowed_values {
list {
type: DT_FLOAT
type: DT_DOUBLE
type: DT_INT32
type: DT_INT64
}
}
}
}
node_def {
name: "two"
op: "Const"
attr {
key: "dtype"
value {
type: DT_INT64
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_INT64
tensor_shape {
}
int64_val: 2
}
}
}
}
node_def {
name: "scale"
op: "Cast"
input: "two:output:0"
attr {
key: "DstT"
value {
placeholder: "T"
}
}
attr {
key: "SrcT"
value {
type: DT_INT64
}
}
}
node_def {
name: "y"
op: "Mul"
input: "x"
input: "scale:y:0"
attr {
key: "T"
value {
placeholder: "T"
}
}
}
ret {
key: "y"
value: "y:z:0"
}
}
function {
signature {
name: "IfWithWhileThen"
input_arg {
name: "x"
type: DT_INT32
}
output_arg {
name: "while"
type: DT_INT32
}
is_stateful: true
}
node_def {
name: "while"
op: "While"
input: "x"
attr {
key: "T"
value {
list {
type: DT_INT32
}
}
}
attr {
key: "_lower_using_switch_merge"
value {
b: true
}
}
attr {
key: "body"
value {
func {
name: "XTimesTwo"
}
}
}
attr {
key: "cond"
value {
func {
name: "LessThanOrEqualToN"
}
}
}
attr {
key: "output_shapes"
value {
list {
shape {
dim {
size: 4
}
}
}
}
}
}
ret {
key: "while"
value: "while:output:0"
}
}
}
versions {
producer: 764
min_consumer: 12
}
tensorflow/core/ir/importexport/tests/graphdef_to_mlir/import_iterator_get_next_attr_inference.pbtxt
+207
View File
@@ -0,0 +1,207 @@
# RUN: tfg-translate -graphdef-to-mlir %s | FileCheck %s
# CHECK: IteratorGetNext
# CHECK-SAME: -> (tensor<*xbf16>, tensor<*xf32>, tensor<*xf64>)
# CHECK: IteratorGetNextSync
# CHECK-SAME: -> (tensor<*xi16>, tensor<*xi32>, tensor<*xi64>)
# CHECK: MultiDeviceIteratorGetNextFromShard
# CHECK-SAME: -> (tensor<*xf16>, tensor<*xcomplex<f32>>, tensor<*xcomplex<f64>>)
node {
name: "args_0"
op: "_Arg"
attr {
key: "T"
value {
type: DT_RESOURCE
}
}
attr {
key: "index"
value {
i: 0
}
}
}
node {
name: "IteratorGetNext"
op: "IteratorGetNext"
input: "args_0"
attr {
key: "output_shapes"
value {
list {
shape {
dim {
size: 1
}
dim {
size: 8
}
}
shape {
dim {
size: 2
}
dim {
size: -1
}
dim {
size: 16
}
}
shape {
unknown_rank: true
}
}
}
}
attr {
key: "output_types"
value {
list {
type: DT_BFLOAT16
type: DT_FLOAT
type: DT_DOUBLE
}
}
}
}
node {
name: "IteratorGetNextSync"
op: "IteratorGetNextSync"
input: "args_0"
attr {
key: "output_shapes"
value {
list {
shape {
unknown_rank: true
}
shape {
dim {
size: 3
}
dim {
size: 24
}
}
shape {
dim {
size: -1
}
dim {
size: 4
}
dim {
size: 32
}
}
}
}
}
attr {
key: "output_types"
value {
list {
type: DT_INT16
type: DT_INT32
type: DT_INT64
}
}
}
}
node {
name: "args_2"
op: "_Arg"
attr {
key: "T"
value {
type: DT_RESOURCE
}
}
attr {
key: "index"
value {
i: 2
}
}
}
node {
name: "args_3"
op: "_Arg"
attr {
key: "T"
value {
type: DT_INT32
}
}
attr {
key: "index"
value {
i: 3
}
}
}
node {
name: "args_4"
op: "_Arg"
attr {
key: "T"
value {
type: DT_INT64
}
}
attr {
key: "index"
value {
i: 4
}
}
}
node {
name: "MultiDeviceIteratorGetNextFromShard"
op: "MultiDeviceIteratorGetNextFromShard"
input: "args_2"
input: "args_3"
input: "args_4"
attr {
key: "output_shapes"
value {
list {
shape {
dim {
size: 5
}
dim {
size: 40
}
}
shape {
unknown_rank: true
}
shape {
dim {
size: 6
}
dim {
size: 48
}
dim {
size: -1
}
}
}
}
}
attr {
key: "output_types"
value {
list {
type: DT_HALF
type: DT_COMPLEX64
type: DT_COMPLEX128
}
}
}
}
tensorflow/core/ir/importexport/tests/graphdef_to_mlir/import_underscore_output_shapes.pbtxt
+49
View File
@@ -0,0 +1,49 @@
# RUN: tfg-translate -graphdef-to-mlir %s | FileCheck %s
# CHECK: Add
# CHECK-SAME: -> (tensor<*xi32>)
node {
name: "Add"
op: "Add"
input: "input0"
input: "input1"
attr {
key: "T"
value {
type: DT_INT32
}
}
attr {
key: "_output_shapes"
value {
list {
shape {
dim {
size: 4
}
}
}
}
}
}
node {
name: "input0"
op: "Placeholder"
attr {
key: "dtype"
value {
type: DT_INT32
}
}
}
node {
name: "input1"
op: "Placeholder"
attr {
key: "dtype"
value {
type: DT_INT32
}
}
}
tensorflow/core/ir/importexport/tests/graphdef_to_mlir/import_while_with_attr_inference.pbtxt
+73
View File
@@ -0,0 +1,73 @@
# RUN: tfg-translate -graphdef-to-mlir %s | FileCheck %s
# CHECK: While
# CHECK-SAME: -> (tensor<*xi32>, tensor<*xf32>)
node {
name: "StatefulWhile"
op: "While"
input: "iter"
input: "val"
attr {
key: "T"
value {
list {
type: DT_INT32
type: DT_FLOAT
}
}
}
attr {
key: "body"
value {
func {
name: "body"
}
}
}
attr {
key: "cond"
value {
func {
name: "cond"
}
}
}
attr {
key: "output_shapes"
value {
list {
shape {
dim {
size: 4
}
}
shape {
dim {
size: 4
}
}
}
}
}
}
node {
name: "iter"
op: "Placeholder"
attr {
key: "dtype"
value {
type: DT_INT32
}
}
}
node {
name: "val"
op: "Placeholder"
attr {
key: "dtype"
value {
type: DT_FLOAT
}
}
}
tensorflow/core/ir/importexport/tests/graphdef_to_mlir/infeed_dequeue.pbtxt
+76
View File
@@ -0,0 +1,76 @@
# RUN: tfg-translate -graphdef-to-mlir %s | FileCheck %s
# CHECK: InfeedDequeue
# CHECK-SAME: tensor<*xi8>
# CHECK: InfeedDequeueTuple
# CHECK-SAME: -> (tensor<*xui16>, tensor<*xui32>, tensor<*xui64>)
node {
name: "InfeedDequeue_0"
op: "InfeedDequeue"
attr {
key: "shape"
value {
shape {
dim {
size: -1
}
dim {
size: 8
}
dim {
size: -1
}
}
}
}
attr {
key: "dtype"
value {
type: DT_INT8
}
}
}
node {
name: "InfeedDequeueTuple"
op: "InfeedDequeueTuple"
attr {
key: "shapes"
value {
list {
shape {
dim {
size: -1
}
dim {
size: -1
}
dim {
size: -1
}
}
shape {
unknown_rank: true
}
shape {
dim {
size: 7
}
dim {
size: 56
}
}
}
}
}
attr {
key: "dtypes"
value {
list {
type: DT_UINT16
type: DT_UINT32
type: DT_UINT64
}
}
}
}
tensorflow/core/ir/importexport/tests/graphdef_to_mlir/infer_arg_handle_type.pbtxt
+140
View File
@@ -0,0 +1,140 @@
# RUN: tfg-translate -graphdef-to-mlir %s | FileCheck %s
# CHECK: tensor<*xf32>
# CHECK: tensor<*x!tf_type.resource>
# CHECK: tensor<*x!tf_type.resource>
# CHECK: tensor<*xf32>
# CHECK: tensor<*x!tf_type.resource>
node {
name: "arg0"
op: "_Arg"
attr {
key: "T"
value {
type: DT_FLOAT
}
}
attr {
key: "index"
value {
i: 0
}
}
}
node {
name: "arg1"
op: "_Arg"
attr {
key: "T"
value {
type: DT_RESOURCE
}
}
attr {
key: "_handle_dtypes"
value {
list {
type: DT_FLOAT
}
}
}
attr {
key: "index"
value {
i: 1
}
}
}
node {
name: "arg2"
op: "_Arg"
attr {
key: "T"
value {
type: DT_RESOURCE
}
}
attr {
key: "_handle_dtypes"
value {
list {
type: DT_FLOAT
}
}
}
attr {
key: "_handle_shapes"
value {
list {
shape {
dim {
size: 4
}
}
}
}
}
attr {
key: "index"
value {
i: 2
}
}
}
node {
name: "arg3"
op: "_Arg"
attr {
key: "T"
value {
type: DT_FLOAT
}
}
attr {
key: "index"
value {
i: 3
}
}
attr {
key: "_output_shapes"
value {
list {
shape {
dim {
size: 8
}
}
}
}
}
}
node {
name: "arg4"
op: "_Arg"
attr {
key: "T"
value {
type: DT_RESOURCE
}
}
attr {
key: "_handle_shapes"
value {
list {
shape {
dim {
size: 4
}
}
}
}
}
attr {
key: "index"
value {
i: 4
}
}
}
tensorflow/core/ir/importexport/tests/graphdef_to_mlir/infer_with_output_shapes.pbtxt
+102
View File
@@ -0,0 +1,102 @@
# RUN: tfg-translate -graphdef-to-mlir %s | FileCheck %s
# CHECK: Add
# CHECK-SAME: -> (tensor<*xi32>)
# CHECK: Add
# CHECK-SAME: -> (tensor<*xi32>)
node {
name: "Const"
op: "Const"
attr {
key: "dtype"
value {
type: DT_INT32
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_INT32
tensor_shape {
}
int_val: 1
}
}
}
}
node {
name: "Const_1"
op: "Const"
attr {
key: "dtype"
value {
type: DT_INT32
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_INT32
tensor_shape {
}
int_val: 0
}
}
}
}
node {
name: "Add"
op: "Add"
input: "Const"
input: "Const_1"
attr {
key: "T"
value {
type: DT_INT32
}
}
attr {
key: "_output_shapes"
value {
list {
shape {
dim {
size: 4
}
}
}
}
}
}
node {
name: "Add_1"
op: "Add"
input: "Const"
input: "Const_1"
attr {
key: "T"
value {
type: DT_INT32
}
}
attr {
key: "_output_shapes"
value {
list {
shape {
dim {
size: 4
}
}
shape {
dim {
size: 4
}
}
}
}
}
}
tensorflow/core/ir/importexport/tests/graphdef_to_mlir/invalid_arg_name.pbtxt
+50
View File
@@ -0,0 +1,50 @@
# RUN: not tfg-translate -graphdef-to-mlir %s 2>&1 | FileCheck %s
# CHECK: INVALID_ARGUMENT: Duplicated node (or function argument) with the same name: `x`
# CHECK-NEXT: when importing function DuplicatedName
library {
function {
signature {
name: "DuplicatedName"
input_arg {
name: "x"
type_attr: "T"
}
output_arg {
name: "y"
type_attr: "T"
}
attr {
name: "T"
type: "type"
allowed_values {
list {
type: DT_FLOAT
type: DT_DOUBLE
type: DT_INT32
type: DT_INT64
}
}
}
}
node_def {
name: "x"
op: "Identity"
input: "x"
attr {
key: "T"
value {
placeholder: "T"
}
}
}
ret {
key: "y"
value: "x"
}
}
}
versions {
producer: 762
}
tensorflow/core/ir/importexport/tests/graphdef_to_mlir/invalid_backedge_input_size.pbtxt
+21
View File
@@ -0,0 +1,21 @@
# RUN: not tfg-translate -graphdef-to-mlir %s 2>&1 | FileCheck %s
# CHECK: Type attr not found
library {
function {
signature {
name: "\344\264\264"
description: "value"
is_distributed_communication: true
}
node_def {
op: "Const"
input: "|"
}
control_ret {
key: ""
value: ""
}
}
}
tensorflow/core/ir/importexport/tests/graphdef_to_mlir/invalid_duplicated_node_name.pbtxt
+61
View File
@@ -0,0 +1,61 @@
# RUN: not tfg-translate -graphdef-to-mlir %s 2>&1 | FileCheck %s
# CHECK: INVALID_ARGUMENT: Duplicated node (or function argument) with the same name: `y`
# CHECK-NEXT: when importing function DuplicatedName
library {
function {
signature {
name: "DuplicatedName"
input_arg {
name: "x"
type_attr: "T"
}
output_arg {
name: "y"
type_attr: "T"
}
attr {
name: "T"
type: "type"
allowed_values {
list {
type: DT_FLOAT
type: DT_DOUBLE
type: DT_INT32
type: DT_INT64
}
}
}
}
node_def {
name: "y"
op: "Identity"
input: "x"
attr {
key: "T"
value {
placeholder: "T"
}
}
}
node_def {
name: "y"
op: "Identity"
input: "x"
attr {
key: "T"
value {
placeholder: "T"
}
}
}
ret {
key: "y"
value: "y"
}
}
}
versions {
producer: 762
}
tensorflow/core/ir/importexport/tests/graphdef_to_mlir/invalid_edge_index.pbtxt
+62
View File
@@ -0,0 +1,62 @@
# RUN: not tfg-translate -graphdef-to-mlir %s 2>&1 | FileCheck %s
# CHECK: Result #1 of node 'const2' is out of bounds
node {
name: "const1"
op: "Const"
attr {
key: "dtype"
value {
type: DT_INT32
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_INT32
tensor_shape {
}
}
}
}
}
node {
name: "const2"
op: "Const"
attr {
key: "dtype"
value {
type: DT_INT32
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_INT32
tensor_shape {
}
}
}
}
}
node {
name: "merge"
op: "Merge"
input: "const1:0"
input: "const2:1"
attr {
key: "N"
value {
i: 2
}
}
attr {
key: "T"
value {
type: DT_DOUBLE
}
}
}
tensorflow/core/ir/importexport/tests/graphdef_to_mlir/invalid_edge_name.pbtxt
+62
View File
@@ -0,0 +1,62 @@
# RUN: not tfg-translate -graphdef-to-mlir %s 2>&1 | FileCheck %s
# CHECK: Node 'const2' has no output called 'invalid'
node {
name: "const1"
op: "Const"
attr {
key: "dtype"
value {
type: DT_INT32
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_INT32
tensor_shape {
}
}
}
}
}
node {
name: "const2"
op: "Const"
attr {
key: "dtype"
value {
type: DT_INT32
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_INT32
tensor_shape {
}
}
}
}
}
node {
name: "merge"
op: "Merge"
input: "const1:output:0"
input: "const2:invalid:0"
attr {
key: "N"
value {
i: 2
}
}
attr {
key: "T"
value {
type: DT_DOUBLE
}
}
}
tensorflow/core/ir/importexport/tests/graphdef_to_mlir/invalid_empty_attr_key.pbtxt
+55
View File
@@ -0,0 +1,55 @@
# RUN: not tfg-translate -graphdef-to-mlir %s 2>&1 | FileCheck %s
# CHECK: empty attr name
node {
name: "SaveV/"
op: "PartitionedCall"
device: "/job:localhost/replica:0/task:0/device:CPU:0"
attr {
key: ""
value {
}
}
attr {
key: "Tin"
value {
}
}
attr {
key: "Tout"
value {
}
}
attr {
key: "config"
value {
s: ""
}
}
attr {
key: "config_proto"
value {
s: ""
}
}
attr {
key: "executor_type"
value {
s: ""
}
}
attr {
key: "f"
value {
func {
}
}
}
}
library {
function {
}
}
versions {
}
tensorflow/core/ir/importexport/tests/graphdef_to_mlir/invalid_empty_func_attr_key.pbtxt
+165
View File
@@ -0,0 +1,165 @@
# RUN: not tfg-translate -graphdef-to-mlir %s 2>&1 | FileCheck %s
# CHECK: empty attr name
node {
name: "SaveV/"
op: "PartitionedCall"
device: "/job:localhost/replica:0/task:0/device:CPU:0"
attr {
key: "T"
value {
func {
attr {
key: ""
value {
}
}
attr {
key: "\036"
value {
}
}
attr {
key: " "
value {
}
}
attr {
key: "1"
value {
}
}
attr {
key: "2"
value {
}
}
attr {
key: "loc:@c"
value {
}
}
attr {
key: "\177"
value {
func {
attr {
key: ""
value {
}
}
attr {
key: "\036"
value {
}
}
attr {
key: "1"
value {
}
}
attr {
key: "2"
value {
}
}
attr {
key: "_class"
value {
}
}
attr {
key: "\177"
value {
}
}
attr {
key: "\177\177"
value {
}
}
}
}
}
}
}
}
attr {
key: "Tin"
value {
}
}
attr {
key: "Tout"
value {
}
}
attr {
key: "config"
value {
s: ""
}
}
attr {
key: "config_proto"
value {
s: ""
}
}
attr {
key: "executor_type"
value {
s: ""
}
}
attr {
key: "f"
value {
func {
attr {
key: "E"
value {
}
}
}
}
}
attr {
key: "shape"
value {
}
}
}
node {
name: "serving_default_x"
op: "Placeholder"
device: "/job:localhost/replica:0/task:0/device:CPU:0"
attr {
key: "dtype"
value {
type: DT_STRING
}
}
attr {
key: "shape"
value {
shape {
unknown_rank: true
}
}
}
}
library {
function {
signature {
is_distributed_communication: true
}
}
gradient {
gradient_func: "\001\000\000\000\000\000\000\000"
}
}
versions {
}
tensorflow/core/ir/importexport/tests/graphdef_to_mlir/invalid_empty_func_attr_name.pbtxt
+38
View File
@@ -0,0 +1,38 @@
# RUN: not tfg-translate -graphdef-to-mlir %s 2>&1 | FileCheck %s
# CHECK: empty func_attr name
node {
name: "NoOp"
op: "NoOp"
attr {
key: "dense_shapes"
value {
list {
shape {
dim {
size: 1
}
}
shape {
dim {
size: 1
}
}
func {
attr {
key: ""
value {
type: DT_QINT16
}
}
}
}
}
}
}
library {
}
versions {
producer: 27
}
tensorflow/core/ir/importexport/tests/graphdef_to_mlir/invalid_empty_op_type.pbtxt
+22
View File
@@ -0,0 +1,22 @@
# RUN: not tfg-translate -graphdef-to-mlir %s 2>&1 | FileCheck %s
# CHECK: empty op type
library {
function {
signature {
name: "XTimesTwo"
}
node_def {
name: "value"
attr {
key: "Tin"
value {
placeholder: "\t"
}
}
}
}
}
versions {
}
tensorflow/core/ir/importexport/tests/graphdef_to_mlir/invalid_func_with_empty_name.pbtxt
+10
View File
@@ -0,0 +1,10 @@
# RUN: not tfg-translate -graphdef-to-mlir %s 2>&1 | FileCheck %s
# CHECK: Function without a name
library {
function {
signature {
}
}
}
tensorflow/core/ir/importexport/tests/graphdef_to_mlir/invalid_function_import.pbtxt
+49
View File
@@ -0,0 +1,49 @@
# RUN: not tfg-translate -graphdef-to-mlir %s 2>&1 | FileCheck %s
# Check that the error message includes the function name. The function is
# invalid because the `control_output` listed in its signature does not appear
# in the function's `control_ret` map.
# CHECK: INVALID_ARGUMENT: Control output 'const' was not found in 'control_ret'
# CHECK: While importing function: test
library {
function {
signature {
name: "test"
input_arg {
name: "a"
type: DT_INT32
}
output_arg {
name: "b"
type: DT_INT32
}
control_output: "const"
}
node_def {
name: "const"
op: "Const"
attr {
key: "dtype"
value {
type: DT_INT32
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_INT32
tensor_shape {
}
}
}
}
}
ret {
key: "b"
value: "const:output:0"
}
}
}
tensorflow/core/ir/importexport/tests/graphdef_to_mlir/invalid_generic_func_with_empty_control_result.pbtxt
+26
View File
@@ -0,0 +1,26 @@
# RUN: not tfg-translate -graphdef-to-mlir %s 2>&1 | FileCheck %s
# CHECK: Function 'foo' has empty control result name
library {
function {
signature {
name: "foo"
control_output: "output"
}
node_def {
name: "y"
op: "NoOp"
attr {
key: "T"
value {
placeholder: "T"
}
}
}
control_ret {
key: "output"
value: ""
}
}
}
tensorflow/core/ir/importexport/tests/graphdef_to_mlir/invalid_generic_func_with_empty_input.pbtxt
+22
View File
@@ -0,0 +1,22 @@
# RUN: not tfg-translate -graphdef-to-mlir %s 2>&1 | FileCheck %s
# CHECK: Node 'y' has an empty input
library {
function {
signature {
name: "foo"
}
node_def {
name: "y"
input: ""
op: "Identity"
attr {
key: "T"
value {
placeholder: "T"
}
}
}
}
}
tensorflow/core/ir/importexport/tests/graphdef_to_mlir/invalid_generic_func_with_empty_name.pbtxt
+20
View File
@@ -0,0 +1,20 @@
# RUN: not tfg-translate -graphdef-to-mlir %s 2>&1 | FileCheck %s
# CHECK: generic function without a name
library {
function {
signature {
}
node_def {
name: "y"
op: "NoOp"
attr {
key: "T"
value {
placeholder: "T"
}
}
}
}
}
tensorflow/core/ir/importexport/tests/graphdef_to_mlir/invalid_generic_func_with_empty_result.pbtxt
+29
View File
@@ -0,0 +1,29 @@
# RUN: not tfg-translate -graphdef-to-mlir %s 2>&1 | FileCheck %s
# CHECK: Function 'foo' has empty result name
library {
function {
signature {
name: "foo"
output_arg {
name: "output"
type: DT_INT32
}
}
node_def {
name: "y"
op: "NoOp"
attr {
key: "T"
value {
placeholder: "T"
}
}
}
ret {
key: "output"
value: ""
}
}
}
tensorflow/core/ir/importexport/tests/graphdef_to_mlir/invalid_generic_function_attr_name.pbtxt
+52
View File
@@ -0,0 +1,52 @@
# RUN: not tfg-translate -graphdef-to-mlir %s 2>&1 | FileCheck %s
# CHECK: Invalid function attribute name
library {
function {
signature {
name: "foo"
input_arg {
name: "a"
}
output_arg {
name: "d"
}
}
node_def {
op: "Const"
attr {
key: "_b"
value {
placeholder: "T"
}
}
attr {
key: "dtype"
value {
type: DT_INT32
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_INT32
tensor_shape {
}
}
}
}
}
ret {
key: "d"
value: "a"
}
attr {
key: ""
value {
s: "a"
}
}
}
}
tensorflow/core/ir/importexport/tests/graphdef_to_mlir/invalid_generic_function_named_edge_index.pbtxt
+52
View File
@@ -0,0 +1,52 @@
# RUN: not tfg-translate -graphdef-to-mlir %s 2>&1 | FileCheck %s
# CHECK: Output index 18446744073709551615 is invalid (too large)
library {
function {
signature {
name: "foo"
attr {
name: "T"
type: "type"
}
}
node_def {
name: "two"
op: "Const"
attr {
key: "dtype"
value {
type: DT_INT64
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_INT64
tensor_shape {}
int64_val: 2
}
}
}
}
node_def {
name: "a"
op: "Cast"
input: "two:output:18446744073709551615"
attr {
key: "DstT"
value {
placeholder: "T"
}
}
attr {
key: "SrcT"
value {
type: DT_INT64
}
}
}
}
}
tensorflow/core/ir/importexport/tests/graphdef_to_mlir/invalid_handle_data.pbtxt
+92
View File
@@ -0,0 +1,92 @@
# RUN: not tfg-translate -graphdef-to-mlir %s 2>&1 | FileCheck %s
# CHECK: INVALID_ARGUMENT: Invalid dtype for handle_data
library {
function {
signature {
name: "XTimesTwo"
input_arg {
name: "x"
type_attr: "T"
}
output_arg {
name: "y"
type_attr: "T"
# This empty handle_data is invalid.
handle_data {
}
}
attr {
name: "T"
type: "type"
allowed_values {
list {
type: DT_FLOAT
type: DT_DOUBLE
type: DT_INT32
type: DT_INT64
}
}
}
}
node_def {
name: "two"
op: "Const"
attr {
key: "dtype"
value {
type: DT_INT64
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_INT64
tensor_shape {
}
int64_val: 2
}
}
}
}
node_def {
name: "scale"
op: "Cast"
input: "two:output:0"
attr {
key: "DstT"
value {
placeholder: "T"
}
}
attr {
key: "SrcT"
value {
type: DT_INT64
}
}
}
node_def {
name: "y"
op: "Mul"
input: "x"
input: "scale:y:0"
attr {
key: "T"
value {
placeholder: "T"
}
}
}
ret {
key: "y"
value: "y:z:0"
}
}
}
versions {
producer: 762
}
tensorflow/core/ir/importexport/tests/graphdef_to_mlir/invalid_missing_control_input.pbtxt
+27
View File
@@ -0,0 +1,27 @@
# RUN: not tfg-translate -graphdef-to-mlir %s 2>&1 | FileCheck %s
# CHECK: Non-existent input ^case/cond/cond/Merge in node ConstantFolding/case/cond/cond/Merge_const
node {
name: "ConstantFolding/case/cond/cond/Merge_const"
op: "Const"
input: "^case/cond/cond/Merge"
device: "/job:localhost/replica:0/task:0/device:CPU:0"
attr {
key: "dtype"
value {
type: DT_DOUBLE
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_DOUBLE
tensor_shape {
}
tensor_content: "\342\351\225\262\014q\274?"
}
}
}
}
tensorflow/core/ir/importexport/tests/graphdef_to_mlir/invalid_missing_control_result.pbtxt
+44
View File
@@ -0,0 +1,44 @@
# RUN: not tfg-translate -graphdef-to-mlir %s 2>&1 | FileCheck %s
# CHECK: Control output 'missing' was not found in 'control_ret'
library {
function {
signature {
name: "test"
output_arg {
name: "output"
type: DT_INT32
}
control_output: "missing"
}
node_def {
name: "const"
op: "Const"
attr {
key: "dtype"
value {
type: DT_INT32
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_INT32
tensor_shape {
}
}
}
}
}
ret {
key: "output"
value: "const"
}
control_ret {
key: "control_output"
value: "const"
}
}
}
tensorflow/core/ir/importexport/tests/graphdef_to_mlir/invalid_missing_control_result_value.pbtxt
+44
View File
@@ -0,0 +1,44 @@
# RUN: not tfg-translate -graphdef-to-mlir %s 2>&1 | FileCheck %s
# CHECK: Control return missing was not found
library {
function {
signature {
name: "test"
output_arg {
name: "output"
type: DT_INT32
}
control_output: "control_output"
}
node_def {
name: "const"
op: "Const"
attr {
key: "dtype"
value {
type: DT_INT32
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_INT32
tensor_shape {
}
}
}
}
}
ret {
key: "output"
value: "const"
}
control_ret {
key: "control_output"
value: "missing"
}
}
}
tensorflow/core/ir/importexport/tests/graphdef_to_mlir/invalid_missing_data_result.pbtxt
+44
View File
@@ -0,0 +1,44 @@
# RUN: not tfg-translate -graphdef-to-mlir %s 2>&1 | FileCheck %s
# CHECK: Output 'missing' was not found in 'ret'
library {
function {
signature {
name: "test"
output_arg {
name: "missing"
type: DT_INT32
}
control_output: "control_output"
}
node_def {
name: "const"
op: "Const"
attr {
key: "dtype"
value {
type: DT_INT32
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_INT32
tensor_shape {
}
}
}
}
}
ret {
key: "output"
value: "const:1"
}
control_ret {
key: "control_output"
value: "const"
}
}
}
tensorflow/core/ir/importexport/tests/graphdef_to_mlir/invalid_missing_data_result_value.pbtxt
+44
View File
@@ -0,0 +1,44 @@
# RUN: not tfg-translate -graphdef-to-mlir %s 2>&1 | FileCheck %s
# CHECK: Return 'missing' was not found
library {
function {
signature {
name: "test"
output_arg {
name: "output"
type: DT_INT32
}
control_output: "control_output"
}
node_def {
name: "const"
op: "Const"
attr {
key: "dtype"
value {
type: DT_INT32
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_INT32
tensor_shape {
}
}
}
}
}
ret {
key: "output"
value: "missing"
}
control_ret {
key: "control_output"
value: "const"
}
}
}
tensorflow/core/ir/importexport/tests/graphdef_to_mlir/invalid_missing_input.pbtxt
+152
View File
@@ -0,0 +1,152 @@
# RUN: not tfg-translate -graphdef-to-mlir %s 2>&1 | FileCheck %s
# CHECK: Non-existent input case/cond/cond/Switch_3 in node case/cond/cond/Merge
node {
name: "ConstantFolding/case/cond/cond/Merge_const"
op: "Const"
input: "^case/cond/cond/Merge"
device: "/job:localhost/replica:0/task:0/device:CPU:0"
attr {
key: "dtype"
value {
type: DT_DOUBLE
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_DOUBLE
tensor_shape {
}
tensor_content: "\342\351\225\262\014q\274?"
}
}
}
}
node {
name: "ConstantFolding/case/cond/cond/Switch_3-folded-1"
op: "Const"
device: "/job:localhost/replica:0/task:0/device:CPU:0"
attr {
key: "dtype"
value {
type: DT_DOUBLE
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_DOUBLE
tensor_shape {
}
tensor_content: "\342\351\225\262\014q\274?"
}
}
}
}
node {
name: "Const_3"
op: "Const"
device: "/job:localhost/replica:0/task:0/device:CPU:0"
attr {
key: "dtype"
value {
type: DT_DOUBLE
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_DOUBLE
tensor_shape {
}
double_val: 0.1234
}
}
}
}
node {
name: "case/cond/cond/Merge"
op: "Merge"
input: "case/cond/cond/Switch_3"
input: "ConstantFolding/case/cond/cond/Switch_3-folded-1"
device: "/job:localhost/replica:0/task:0/device:CPU:0"
attr {
key: "N"
value {
i: 2
}
}
attr {
key: "T"
value {
type: DT_DOUBLE
}
}
}
node {
name: "Less"
op: "Const"
device: "/job:localhost/replica:0/task:0/device:CPU:0"
attr {
key: "dtype"
value {
type: DT_BOOL
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_BOOL
tensor_shape {
}
tensor_content: "\000"
}
}
}
}
node {
name: "case/cond/Switch_1"
op: "Switch"
input: "Const_3"
input: "Less"
device: "/job:localhost/replica:0/task:0/device:CPU:0"
attr {
key: "T"
value {
type: DT_DOUBLE
}
}
attr {
key: "_class"
value {
list {
s: "loc:@Const_3"
}
}
}
}
node {
name: "case/cond/Merge"
op: "Merge"
input: "ConstantFolding/case/cond/cond/Merge_const"
input: "case/cond/Switch_1:1"
device: "/job:localhost/replica:0/task:0/device:CPU:0"
attr {
key: "N"
value {
i: 2
}
}
attr {
key: "T"
value {
type: DT_DOUBLE
}
}
}
tensorflow/core/ir/importexport/tests/graphdef_to_mlir/invalid_missing_two_inputs.pbtxt
+62
View File
@@ -0,0 +1,62 @@
# RUN: not tfg-translate -graphdef-to-mlir %s 2>&1 | FileCheck %s
# CHECK: Non-existent input invalid_1 in node Merge
# CHECK: Non-existent input invalid_2 in node Merge_2
node {
name: "Merge"
op: "Merge"
input: "invalid_1"
input: "Const_3"
attr {
key: "N"
value {
i: 2
}
}
attr {
key: "T"
value {
type: DT_DOUBLE
}
}
}
node {
name: "Merge_2"
op: "Merge"
input: "invalid_2"
input: "Const_3"
attr {
key: "N"
value {
i: 2
}
}
attr {
key: "T"
value {
type: DT_DOUBLE
}
}
}
node {
name: "Const_3"
op: "Const"
attr {
key: "dtype"
value {
type: DT_DOUBLE
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_DOUBLE
tensor_shape {
}
double_val: 0.1234
}
}
}
}
tensorflow/core/ir/importexport/tests/graphdef_to_mlir/invalid_named_edge_index.pbtxt
+62
View File
@@ -0,0 +1,62 @@
# RUN: not tfg-translate -graphdef-to-mlir %s 2>&1 | FileCheck %s
# CHECK: Result #1 of segment 'const2:output' is out of bounds
node {
name: "const1"
op: "Const"
attr {
key: "dtype"
value {
type: DT_INT32
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_INT32
tensor_shape {
}
}
}
}
}
node {
name: "const2"
op: "Const"
attr {
key: "dtype"
value {
type: DT_INT32
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_INT32
tensor_shape {
}
}
}
}
}
node {
name: "merge"
op: "Merge"
input: "const1:output:0"
input: "const2:output:1"
attr {
key: "N"
value {
i: 2
}
}
attr {
key: "T"
value {
type: DT_DOUBLE
}
}
}
tensorflow/core/ir/importexport/tests/graphdef_to_mlir/invalid_op_name.pbtxt
+20
View File
@@ -0,0 +1,20 @@
# RUN: not tfg-translate -graphdef-to-mlir %s 2>&1 | FileCheck %s
# CHECK: Node has an empty op name
library {
function {
signature {
name: "\\344\\264\\264"
description: "value"
is_distributed_communication: true
}
node_def {
input: "|"
}
control_ret {
key: ""
value: ""
}
}
}
tensorflow/core/ir/importexport/tests/graphdef_to_mlir/invalid_type_list.pbtxt
+78
View File
@@ -0,0 +1,78 @@
# RUN: not tfg-translate -graphdef-to-mlir %s 2>&1 | FileCheck %s
# CHECK: Expected 'Tout' to be a list of types
node {
name: "SaveV/"
op: "PartitionedCall"
device: "/job:localhost/replica:0/task:0/device:CPU:0"
attr {
key: "N"
value {
placeholder: ""
}
}
attr {
key: "Tin"
value {
placeholder: ""
}
}
attr {
key: "Tout"
value {
list {
shape {
dim {
size: -1
}
dim {
size: 1
}
}
shape {
dim {
size: -1
name: "T"
}
dim {
size: 1
}
}
}
}
}
attr {
key: "config"
value {
s: ""
}
}
attr {
key: "config_proto"
value {
s: ""
}
}
attr {
key: "executor_type"
value {
s: ""
}
}
attr {
key: "f"
value {
func {
}
}
}
attr {
key: "shape"
value {
type: DT_STRING
}
}
experimental_debug_info {
}
}
tensorflow/core/ir/importexport/tests/graphdef_to_mlir/legacy_call.pbtxt
+134
View File
@@ -0,0 +1,134 @@
# RUN: tfg-translate -graphdef-to-mlir %s | FileCheck %s
# Importing a "direct call" is converted into `tfg.LegacyCall` in MLIR.
# DISABLED-CHECK: %LegacyCall, %ctl_0 = tfg.LegacyCall(%Placeholder) device = "/job:localhost/replica:0/task:0/device:CPU:0" name = "y" {callee = #tfg.func<@XTimesTwo, {T = f32}>} : (tensor<*xf32>) -> (tensor<*xf32>)
# CHECK: tfg.func generic @XTimesTwo
node {
name: "x"
op: "Placeholder"
device: "/job:localhost/replica:0/task:0/device:CPU:0"
attr {
key: "dtype"
value {
type: DT_FLOAT
}
}
attr {
key: "shape"
value {
shape {
unknown_rank: true
}
}
}
}
node {
name: "y"
op: "XTimesTwo"
input: "x"
device: "/job:localhost/replica:0/task:0/device:CPU:0"
attr {
key: "T"
value {
type: DT_FLOAT
}
}
}
node {
name: "z"
op: "Identity"
input: "y"
device: "/job:localhost/replica:0/task:0/device:CPU:0"
attr {
key: "T"
value {
type: DT_FLOAT
}
}
}
library {
function {
signature {
name: "XTimesTwo"
input_arg {
name: "x"
type_attr: "T"
}
output_arg {
name: "y"
type_attr: "T"
}
attr {
name: "T"
type: "type"
allowed_values {
list {
type: DT_FLOAT
type: DT_DOUBLE
type: DT_INT32
type: DT_INT64
}
}
}
}
node_def {
name: "two"
op: "Const"
attr {
key: "dtype"
value {
type: DT_INT64
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_INT64
tensor_shape {
}
int64_val: 2
}
}
}
}
node_def {
name: "scale"
op: "Cast"
input: "two:output:0"
attr {
key: "DstT"
value {
placeholder: "T"
}
}
attr {
key: "SrcT"
value {
type: DT_INT64
}
}
}
node_def {
name: "y"
op: "Mul"
input: "x"
input: "scale:y:0"
attr {
key: "T"
value {
placeholder: "T"
}
}
}
ret {
key: "y"
value: "y:z:0"
}
}
}
versions {
producer: 762
}
tensorflow/core/ir/importexport/tests/graphdef_to_mlir/negative_shape.pbtxt
+42
View File
@@ -0,0 +1,42 @@
# RUN: tfg-translate -graphdef-to-mlir %s | FileCheck %s
# This node has an attribute with a "shape" that has a -5 dimension. This use to
# crash the importer.
# Such shapes don't really make sense in TensorFlow, but grappler is using this
# during some "symbolic shape analysis".
# CHECK: _Arg name("test_model")
# CHECK-SAME: #tf_type.shape<-5>
node {
name: "test_model"
op: "_Arg"
attr {
key: "T"
value {
type: DT_FLOAT
}
}
attr {
key: "_output_shapes"
value {
list {
shape {
dim {
size: -5
}
}
}
}
}
attr {
key: "index"
value {
i: 8
}
}
}
library {
}
versions {
producer: 886
}
tensorflow/core/ir/importexport/tests/graphdef_to_mlir/negative_zero_constant.pbtxt
+72
View File
@@ -0,0 +1,72 @@
# RUN: tfg-translate -graphdef-to-mlir %s | FileCheck %s
# Test that we import correctly a negative zero.
# CHECK: Const name("float") {{.*}} value = dense<-0.000000e+00> : tensor<f32>
# CHECK: Const name("half") {{.*}} value = dense<-0.000000e+00> : tensor<f16>
# CHECK: Const name("complex") {{.*}} value = dense<(-0.000000e+00,-0.000000e+00)> : tensor<complex<f32>>
node {
name: "float"
op: "Const"
attr {
key: "dtype"
value {
type: DT_FLOAT
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_FLOAT
tensor_shape {
}
float_val: -0
}
}
}
}
node {
name: "half"
op: "Const"
attr {
key: "dtype"
value {
type: DT_HALF
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_HALF
tensor_shape {
}
half_val: 0x8000 # -0.0 in half precision
}
}
}
}
node {
name: "complex"
op: "Const"
attr {
key: "dtype"
value {
type: DT_COMPLEX64
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_COMPLEX64
tensor_shape {
}
scomplex_val: -0
scomplex_val: -0
}
}
}
}
tensorflow/core/ir/importexport/tests/graphdef_to_mlir/three_nodes_with_attrs.pbtxt
+49
View File
@@ -0,0 +1,49 @@
# RUN: tfg-translate -graphdef-to-mlir %s | FileCheck %s
# Since there are no uses of the output for the Add node, it is imported without any output.
# CHECK: tfg.graph
# CHECK-DAG: %[[ADD:.*]], %{{.*}} = Add(%[[PLACEHOLDER:.*]], %[[PLACEHOLDER0:.*]]) name("Add") {T = i32, _mlir_fulltype = #tf_type.full_type<tensor>, _tpu_replicate = "cluster"} : (tensor<*xi32>, tensor<*xi32>) -> (tensor<*xi32>)
# CHECK-DAG: %[[PLACEHOLDER]], %{{.*}} = Placeholder name("input0") {dtype = i32, shape = #tf_type.shape<*>} : () -> (tensor<*xi32>)
# CHECK-DAG: %[[PLACEHOLDER0]], %{{.*}} = Placeholder name("input1") {dtype = i32, shape = #tf_type.shape<*>} : () -> (tensor<*xi32>)
node {
name: "Add"
op: "Add"
input: "input0"
input: "input1"
attr {
key: "T"
value {
type: DT_INT32
}
}
attr {
key: "_tpu_replicate"
value {
s: "cluster"
}
}
experimental_type {
type_id: TFT_TENSOR
}
}
node {
name: "input0"
op: "Placeholder"
attr {
key: "dtype"
value {
type: DT_INT32
}
}
}
node {
name: "input1"
op: "Placeholder"
attr {
key: "dtype"
value {
type: DT_INT32
}
}
}
tensorflow/core/ir/importexport/tests/graphdef_to_mlir/version.pbtxt
+9
View File
@@ -0,0 +1,9 @@
# RUN: tfg-translate -graphdef-to-mlir %s | FileCheck %s
#CHECK: tfg.graph #tf_type.version<producer = 27, min_consumer = 20, bad_consumers = [0, 1, 1, 2, 3, 5, 8, 13]> {
versions {
producer: 27
min_consumer: 20
bad_consumers: [0, 1, 1, 2, 3, 5, 8, 13]
}
tensorflow/core/ir/importexport/tests/mlir_to_graphdef/BUILD
+33
View File
@@ -0,0 +1,33 @@
load("//tensorflow:tensorflow.default.bzl", "get_compatible_with_portable")
load("//tensorflow/compiler/mlir:glob_lit_test.bzl", "glob_lit_tests")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
default_compatible_with = get_compatible_with_portable(),
default_visibility = [":__subpackages__"],
licenses = ["notice"], # Apache 2.0
)
# Bundle together all of the test utilities that are used by tests.
exports_files(["run_lit.sh"])
filegroup(
name = "test_utilities",
testonly = True,
data = [
"//tensorflow/core/ir/importexport:tfg-translate",
"@llvm-project//llvm:FileCheck",
"@llvm-project//llvm:not",
],
)
glob_lit_tests(
name = "all_tests",
data = [":test_utilities"],
driver = "@llvm-project//mlir:run_lit.sh",
exclude = [],
test_file_exts = [
"mlir",
"pbtxt",
],
)
tensorflow/core/ir/importexport/tests/mlir_to_graphdef/empty.mlir
+16
View File
@@ -0,0 +1,16 @@
// 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: tfg-translate -mlir-to-graphdef %s
tensorflow/core/ir/importexport/tests/mlir_to_graphdef/fulltype.mlir
+29
View File
@@ -0,0 +1,29 @@
// 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: tfg-translate -mlir-to-graphdef %s | FileCheck %s
// Verifies integer fulltype correctly exported.
// CHECK: experimental_type {
// CHECK-NEXT: i: 0
// CHECK: }
tfg.graph #tf_type.version<producer = 527, min_consumer = 12> {
%VarHandleOp, %ctl = VarHandleOp device("/job:localhost/replica:0/task:0/device:CPU:0") name("Variable") {_class = ["loc:@Variable"], _output_shapes = [#tf_type.shape<*>], allowed_devices = [], container = "", dtype = !tf_type.string, shape = #tf_type.shape<>, shared_name = "VOriabAle"} : () -> (tensor<*x!tf_type.resource>)
%Placeholder, %ctl_0 = Placeholder device("/job:localhost/replica:0/task:0/device:CPU:0") name("asset_path_initializer") {_output_shapes = [#tf_type.func<@__inference__traced_restore_52, {}>], dtype = !tf_type.string, shape = #tf_type.shape<>} : () -> (tensor<!tf_type.string>)
%ctl_1 = AssignVariableOp(%VarHandleOp, %Placeholder) device("/job:localhost/replica:0/task:0/device:CPU:0") name("Variable/Assign") {dtype = !tf_type.string, validate_shape = false} : tensor<*x!tf_type.resource>, tensor<!tf_type.string>
%ctl_2 = NoOp [%ctl_1] device("/device:CPU:0") name("NoOp") {_mlir_fulltype = #tf_type.full_type<unset 0 : i64>}
}
tensorflow/core/ir/importexport/tests/mlir_to_graphdef/func_with_no_args_or_results.mlir
+44
View File
@@ -0,0 +1,44 @@
// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
// RUN: tfg-translate -mlir-to-graphdef --split-input-file %s | FileCheck %s
// CHECK: signature {
// CHECK-NEXT: name: "func_no_args_no_results"
// CHECK-NOT: input_arg
// CHECK-NOT: output_arg
tfg.func @func_no_args_no_results() -> () {
return
}
// -----
// CHECK: signature {
// CHECK-NEXT: name: "func_no_args"
// CHECK-NEXT: output_arg
// CHECK-NOT: input_arg
tfg.func @func_no_args() -> (tensor<1xi32> {tfg.name = "ret1"}) {
%Const, %ctl_2 = Const name("c") {dtype = i32, value = dense<0> : tensor<1xi32>} : () -> (tensor<1xi32>)
return (%Const) : tensor<1xi32>
}
// -----
// CHECK: signature {
// CHECK-NEXT: name: "func_no_results"
// CHECK-NEXT: input_arg
// CHECK-NOT: output_arg
tfg.func @func_no_results(%arg : tensor<i32> {tfg.name = "arg"}) -> () {
return
}
tensorflow/core/ir/importexport/tests/mlir_to_graphdef/negative_zero_constant.mlir
+28
View File
@@ -0,0 +1,28 @@
// 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: tfg-translate -mlir-to-graphdef %s | FileCheck %s
tfg.graph #tf_type.version<producer = 0, min_consumer = 0> {
// CHECK: name: "float"
// CHECK: float_val: -0
%Const, %ctl = Const name("float") {dtype = f32, value = dense<-0.000000e+00> : tensor<f32>} : () -> (tensor<f32>)
// CHECK: name: "half"
// CHECK: half_val: 32768
%Const_0, %ctl_1 = Const name("half") {dtype = f16, value = dense<-0.000000e+00> : tensor<f16>} : () -> (tensor<f16>)
// CHECK: name: "complex"
// CHECK: scomplex_val: -0
// CHECK: scomplex_val: -0
%Const_2, %ctl_3 = Const name("complex") {dtype = complex<f32>, value = dense<(-0.000000e+00,-0.000000e+00)> : tensor<complex<f32>>} : () -> (tensor<complex<f32>>)
}
tensorflow/core/ir/importexport/tests/mlir_to_graphdef/nested_legacy_call.mlir
+50
View File
@@ -0,0 +1,50 @@
// 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: tfg-translate -mlir-to-graphdef %s | FileCheck %s
tfg.graph #tf_type.version<producer = 1, min_consumer = 1> {
// CHECK: name: "foo1"
%Foo:2, %ctl = Foo name("foo1") : () -> (tensor<*xi32>, tensor<*xi32>)
// CHECK: name: "id1"
// CHECK-NEXT: op: "Identity"
// CHECK-NEXT: input: "foo1"
%Identity, %ctl_0 = Identity(%Foo#0) name("id1") {T = i32} : (tensor<*xi32>) -> (tensor<*xi32>)
// CHECK: name: "id2"
// CHECK-NEXT: op: "Identity"
// CHECK-NEXT: input: "foo1:1"
%Identity_1, %ctl_2 = Identity(%Foo#1) name("id2") {T = i32} : (tensor<*xi32>) -> (tensor<*xi32>)
}
// CHECK: library
// CHECK: name: "Foo"
tfg.func @Foo() -> (tensor<*xi32> {tfg.name = "ret1"},
tensor<*xi32> {tfg.name = "ret2"}) {
// CHECK: name: "bar"
%Bar:2, %ctl = Bar name("bar") : () -> (tensor<*xi32>, tensor<*xi32>)
// CHECK: value: "bar:ret1:0"
// CHECK: value: "bar:ret2:0"
return(%Bar#0, %Bar#1) : tensor<*xi32>, tensor<*xi32>
}
// CHECK: name: "Bar"
tfg.func @Bar() -> (tensor<*xi32> {tfg.name = "ret1"},
tensor<*xi32> {tfg.name = "ret2"}) {
// CHECK: name: "const"
%Const, %ctl = Const name("const") {value = dense<0> : tensor<i32>, dtype = i32} : () -> (tensor<*xi32>)
// CHECK: value: "const:output:0"
// CHECK: value: "const:output:0"
return(%Const, %Const) : tensor<*xi32>, tensor<*xi32>
}
tensorflow/core/ir/importexport/tests/mlir_to_graphdef/three_nodes_with_attrs.mlir
+64
View File
@@ -0,0 +1,64 @@
// 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: tfg-translate -mlir-to-graphdef %s | FileCheck %s
tfg.graph #tf_type.version<producer = 34, min_consumer = 5> {
%ctl = tfg.Add(%Placeholder, %Placeholder_1) name("SomeAdd") {T = i32, _mlir_fulltype = #tf_type.full_type<tensor>, _tpu_replicate = "cluster"} : (!tf_type.tensor, !tf_type.tensor) -> ()
%Placeholder, %ctl_0 = tfg.Placeholder name("Placeholder1") {dtype = i32} : () -> (!tf_type.tensor)
%Placeholder_1, %ctl_2 = tfg.Placeholder name("Placeholder2") {dtype = i32} : () -> (!tf_type.tensor)
}
// CHECK: node {
// CHECK-NEXT: name: "SomeAdd"
// CHECK-NEXT: op: "Add"
// CHECK-NEXT: input: "Placeholder1"
// CHECK-NEXT: input: "Placeholder2"
// CHECK-NEXT: attr {
// CHECK-NEXT: key: "T"
// CHECK-NEXT: value {
// CHECK-NEXT: type: DT_INT32
// CHECK-NEXT: }
// CHECK-NEXT: }
// CHECK-NEXT: attr {
// CHECK-NEXT: key: "_tpu_replicate"
// CHECK-NEXT: value {
// CHECK-NEXT: s: "cluster"
// CHECK-NEXT: }
// CHECK-NEXT: }
// CHECK-NEXT: experimental_type {
// CHECK-NEXT: type_id: TFT_TENSOR
// CHECK-NEXT: }
// CHECK: node {
// CHECK-NEXT: name: "Placeholder1"
// CHECK-NEXT: op: "Placeholder"
// CHECK-NEXT: attr {
// CHECK-NEXT: key: "dtype"
// CHECK-NEXT: value {
// CHECK-NEXT: type: DT_INT32
// CHECK-NEXT: }
// CHECK-NEXT: }
// CHECK: node {
// CHECK-NEXT: name: "Placeholder2"
// CHECK-NEXT: op: "Placeholder"
// CHECK-NEXT: attr {
// CHECK-NEXT: key: "dtype"
// CHECK-NEXT: value {
// CHECK-NEXT: type: DT_INT32
// CHECK-NEXT: }
// CHECK-NEXT: }
// CHECK: versions {
// CHECK-NEXT: producer: 34
// CHECK-NEXT: min_consumer: 5
// CHECK-NEXT: }
tensorflow/core/ir/importexport/tests/mlir_to_graphdef/version.mlir
+25
View File
@@ -0,0 +1,25 @@
// 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: tfg-translate -mlir-to-graphdef %s | FileCheck %s
tfg.graph #tf_type.version<producer = 42, min_consumer = 21, bad_consumers = [1, 2, 5, 12]> {
}
// CHECK: producer: 42
// CHECK: min_consumer: 21
// CHECK: bad_consumers: 1
// CHECK: bad_consumers: 2
// CHECK: bad_consumers: 5
// CHECK: bad_consumers: 12
tensorflow/core/ir/importexport/tests/roundtrip/BUILD
+55
View File
@@ -0,0 +1,55 @@
load("@rules_cc//cc:cc_library.bzl", "cc_library")
load("//tensorflow:tensorflow.bzl", "tf_cc_binary")
# Regression tests for bridge.
load(":roundtrip.bzl", "glob_roundtrip_tests")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
default_visibility = ["//tensorflow/core/ir/importexport:__subpackages__"],
licenses = ["notice"], # Apache 2.0
)
cc_library(
name = "roundtrip",
srcs = [
"roundtrip.cc",
],
hdrs = ["roundtrip.h"],
deps = [
"//tensorflow/core:core_cpu_base",
"//tensorflow/core:framework",
"//tensorflow/core:protos_all_cc",
"@llvm-project//llvm:Support",
],
)
# This will run the `verify-roundtrip` binary below on every single pbtxt in this directory.
glob_roundtrip_tests(
exclude = [],
test_file_exts = ["pbtxt"],
)
tf_cc_binary(
name = "verify-roundtrip",
testonly = 1,
srcs = ["verify_roundtrip.cc"],
deps = [
":roundtrip",
"//tensorflow/compiler/mlir:init_mlir",
"//tensorflow/core:core_cpu_base",
"//tensorflow/core:framework",
"//tensorflow/core:ops", # Ops need to be registered for import.
"//tensorflow/core/ir/importexport:graphdef_export",
"//tensorflow/core/ir/importexport:graphdef_import",
"//tensorflow/core/ir/importexport:load_proto",
"//tensorflow/core/transforms/consolidate_attrs:Pass",
"@com_google_googletest//:gtest",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Parser",
"@llvm-project//mlir:Pass",
"@llvm-project//mlir:Support",
"@llvm-project//mlir:TranslateLib",
],
)
tensorflow/core/ir/importexport/tests/roundtrip/add.pbtxt
+35
View File
@@ -0,0 +1,35 @@
node {
name: "Add"
op: "Add"
input: "input0"
input: "input1"
attr {
key: "T"
value {
type: DT_INT32
}
}
}
node {
name: "input0"
op: "Placeholder"
attr {
key: "dtype"
value {
type: DT_INT32
}
}
}
node {
name: "input1"
op: "Placeholder"
attr {
key: "dtype"
value {
type: DT_INT32
}
}
}
versions {
producer: 27
}
tensorflow/core/ir/importexport/tests/roundtrip/arg-as-fetch.pbtxt
+19
View File
@@ -0,0 +1,19 @@
node {
name: "arg"
op: "_Arg"
attr {
key: "T"
value {
type: DT_INT32
}
}
attr {
key: "index"
value {
i: 0
}
}
}
versions {
producer: 27
}
tensorflow/core/ir/importexport/tests/roundtrip/arg-control-dep.pbtxt
+71
View File
@@ -0,0 +1,71 @@
node {
name: "Constant"
op: "Const"
attr {
key: "dtype"
value {
type: DT_INT32
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_INT32
tensor_shape {
}
int_val: 0
}
}
}
}
node {
name: "_tf.foo"
op: "foo"
input: "Constant"
}
library {
function {
signature {
name: "foo"
input_arg {
name: "arg"
type: DT_INT32
}
output_arg {
name: "return_value"
type: DT_INT32
}
}
node_def {
name: "test"
op: "Const"
input: "^arg"
attr {
key: "dtype"
value {
type: DT_INT32
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_INT32
tensor_shape {
}
int_val: 0
}
}
}
}
ret {
key: "return_value"
value: "arg"
}
}
}
versions {
producer: 62
min_consumer: 12
}
tensorflow/core/ir/importexport/tests/roundtrip/arg-data-type.pbtxt
+39
View File
@@ -0,0 +1,39 @@
node {
name: "p"
op: "Placeholder"
attr {
key: "dtype"
value {
type: DT_BOOL
}
}
attr {
key: "shape"
value {
shape {
unknown_rank: true
}
}
}
}
node {
name: "x"
op: "Placeholder"
attr {
key: "dtype"
value {
type: DT_INT32
}
}
attr {
key: "shape"
value {
shape {
unknown_rank: true
}
}
}
}
versions {
producer: 216
}
tensorflow/core/ir/importexport/tests/roundtrip/arg-retval-attrs.pbtxt
+151
View File
@@ -0,0 +1,151 @@
node {
name: "arg0"
op: "_Arg"
attr {
key: "T"
value {
type: DT_FLOAT
}
}
attr {
key: "index"
value {
i: 0
}
}
experimental_type {
type_id: TFT_TENSOR
}
}
node {
name: "arg1"
op: "_Arg"
attr {
key: "T"
value {
type: DT_INT32
}
}
attr {
key: "_arg1_attr0"
value {
s: "_arg1_attr0_value"
}
}
attr {
key: "_arg1_attr1"
value {
f: 8.0
}
}
attr {
key: "index"
value {
i: 1
}
}
}
node {
name: "arg2"
op: "_Arg"
attr {
key: "T"
value {
type: DT_BOOL
}
}
attr {
key: "index"
value {
i: 2
}
}
}
node {
name: "ret0"
op: "_Retval"
input: "arg0"
attr {
key: "T"
value {
type: DT_FLOAT
}
}
attr {
key: "_ret0_attr0"
value {
i: 8
}
}
attr {
key: "_ret0_attr1"
value {
b: false
}
}
attr {
key: "index"
value {
i: 0
}
}
}
node {
name: "ret1"
op: "_Retval"
input: "arg1"
attr {
key: "T"
value {
type: DT_INT32
}
}
attr {
key: "index"
value {
i: 1
}
}
}
node {
name: "ret2"
op: "_Retval"
input: "arg2"
attr {
key: "T"
value {
type: DT_BOOL
}
}
attr {
key: "_ret2_attr0"
value {
type: DT_VARIANT
}
}
attr {
key: "_ret2_attr1"
value {
shape {
dim {
size: 128
}
dim {
size: 1024
}
}
}
}
attr {
key: "index"
value {
i: 2
}
}
experimental_type {
type_id: TFT_TENSOR
}
}
versions {
producer: 121
}
tensorflow/core/ir/importexport/tests/roundtrip/case_op.pbtxt
+256
View File
@@ -0,0 +1,256 @@
node {
name: "Const"
op: "Const"
attr {
key: "dtype"
value {
type: DT_INT32
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_INT32
tensor_shape {
}
int_val: 1
}
}
}
}
node {
name: "Const_1"
op: "Const"
attr {
key: "dtype"
value {
type: DT_INT32
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_INT32
tensor_shape {
}
int_val: 0
}
}
}
}
node {
name: "indexed_case"
op: "StatelessCase"
input: "Const_1"
input: "Const"
attr {
key: "Tin"
value {
list {
type: DT_INT32
}
}
}
attr {
key: "Tout"
value {
list {
type: DT_INT32
}
}
}
attr {
key: "_lower_using_switch_merge"
value {
b: true
}
}
attr {
key: "_read_only_resource_inputs"
value {
list {
}
}
}
attr {
key: "branches"
value {
list {
func {
name: "indexed_case_branch0_4"
}
func {
name: "indexed_case_branch1_5"
}
}
}
}
attr {
key: "output_shapes"
value {
list {
shape {
}
}
}
}
}
node {
name: "indexed_case/Identity"
op: "Identity"
input: "indexed_case"
attr {
key: "T"
value {
type: DT_INT32
}
}
}
library {
function {
signature {
name: "indexed_case_branch0_4"
input_arg {
name: "add_const"
type: DT_INT32
}
output_arg {
name: "add"
type: DT_INT32
}
}
node_def {
name: "add/y"
op: "Const"
attr {
key: "dtype"
value {
type: DT_INT32
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_INT32
tensor_shape {
}
int_val: 1
}
}
}
experimental_debug_info {
original_node_names: "add/y"
}
}
node_def {
name: "add_0"
op: "AddV2"
input: "add_const"
input: "add/y:output:0"
attr {
key: "T"
value {
type: DT_INT32
}
}
experimental_debug_info {
original_node_names: "add"
}
}
ret {
key: "add"
value: "add_0:z:0"
}
arg_attr {
key: 0
value {
attr {
key: "_output_shapes"
value {
list {
shape {
}
}
}
}
}
}
}
function {
signature {
name: "indexed_case_branch1_5"
input_arg {
name: "add_const"
type: DT_INT32
}
output_arg {
name: "add"
type: DT_INT32
}
}
node_def {
name: "add/y"
op: "Const"
attr {
key: "dtype"
value {
type: DT_INT32
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_INT32
tensor_shape {
}
int_val: 2
}
}
}
experimental_debug_info {
original_node_names: "add/y"
}
}
node_def {
name: "add_0"
op: "AddV2"
input: "add_const"
input: "add/y:output:0"
attr {
key: "T"
value {
type: DT_INT32
}
}
experimental_debug_info {
original_node_names: "add"
}
}
ret {
key: "add"
value: "add_0:z:0"
}
arg_attr {
key: 0
value {
attr {
key: "_output_shapes"
value {
list {
shape {
}
}
}
}
}
}
}
}
versions {
producer: 486
min_consumer: 12
}
tensorflow/core/ir/importexport/tests/roundtrip/const-values.pbtxt
+158
View File
@@ -0,0 +1,158 @@
node {
name: "bf16_scalar"
op: "Const"
attr {
key: "dtype"
value {
type: DT_BFLOAT16
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_BFLOAT16
tensor_shape {
}
half_val: 0
}
}
}
}
node {
name: "bf16_vector"
op: "Const"
attr {
key: "dtype"
value {
type: DT_BFLOAT16
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_BFLOAT16
tensor_shape {
dim {
size: 2
}
}
half_val: 16964
half_val: 17485
}
}
}
}
node {
name: "double"
op: "Const"
device: "/device:CPU:0"
attr {
key: "dtype"
value {
type: DT_DOUBLE
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_DOUBLE
tensor_shape {
dim {
size: 2
}
dim {
size: 3
}
}
tensor_content: "\000\000\000\000\000\000\360?\000\000\000\000\000\000\000@\000\000\000\000\000\000\010@\000\000\000\000\000\000\020@\000\000\000\000\000\000\024@\000\000\000\000\000\000\030@"
}
}
}
}
node {
name: "x"
op: "Const"
device: "/device:CPU:0"
attr {
key: "dtype"
value {
type: DT_FLOAT
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_FLOAT
tensor_shape {
dim {
size: 2
}
dim {
size: 3
}
}
tensor_content: "\x00\x00\x80\x3F\x00\x00\x00\x40\x00\x00\x40\x40\x00\x00\x80\x40\x00\x00\xA0\x40\x00\x00\xC0\x40"
}
}
}
}
node {
name: "y"
op: "Const"
device: "/device:CPU:0"
attr {
key: "dtype"
value {
type: DT_INT64
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_INT64
tensor_shape {
dim {
size: 2
}
dim {
size: 3
}
}
tensor_content: "\x01\x00\x00\x00\x00\x00\x00\x00\x03\x00\x00\x00\x00\x00\x00\x00\x02\x00\x00\x00\x00\x00\x00\x00\x05\x00\x00\x00\x00\x00\x00\x00\x04\x00\x00\x00\x00\x00\x00\x00\x07\x00\x00\x00\x00\x00\x00\x00"
}
}
}
}
node {
name: "z"
op: "Const"
device: "/device:CPU:0"
attr {
key: "dtype"
value {
type: DT_INT32
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_INT32
tensor_shape {
dim {
size: 2
}
dim {
size: 3
}
}
tensor_content: "\x01\x00\x00\x00\x03\x00\x00\x00\x02\x00\x00\x00\x05\x00\x00\x00\x04\x00\x00\x00\x07\x00\x00\x00"
}
}
}
}
tensorflow/core/ir/importexport/tests/roundtrip/device-arg-retval-attr.pbtxt
+107
View File
@@ -0,0 +1,107 @@
node {
name: "args_0"
op: "_Arg"
device: "/CPU:0"
attr {
key: "T"
value {
type: DT_FLOAT
}
}
attr {
key: "index"
value {
i: 0
}
}
}
node {
name: "args_1"
op: "_Arg"
attr {
key: "T"
value {
type: DT_INT32
}
}
attr {
key: "_output_shapes"
value {
list {
shape {
dim {
size: 2
}
dim {
size: 4
}
dim {
size: 6
}
dim {
size: 8
}
}
}
}
}
attr {
key: "index"
value {
i: 1
}
}
}
node {
name: "identity"
op: "IdentityN"
input: "args_0"
input: "args_1"
attr: {
key: "T"
value: {
list: {
type: DT_FLOAT
type: DT_INT32
}
}
}
}
node {
name: "rets_0"
op: "_Retval"
input: "identity:0"
attr {
key: "T"
value {
type: DT_FLOAT
}
}
attr {
key: "index"
value {
i: 0
}
}
}
node {
name: "rets_1"
op: "_Retval"
input: "identity:1"
device: "/CPU:1"
attr {
key: "T"
value {
type: DT_INT32
}
}
attr {
key: "index"
value {
i: 1
}
}
}
versions {
producer: 121
}
tensorflow/core/ir/importexport/tests/roundtrip/empty-input-shapes.pbtxt
+39
View File
@@ -0,0 +1,39 @@
node {
name: "input"
op: "Placeholder"
attr {
key: "dtype"
value {
type: DT_BOOL
}
}
}
node {
name: "func0"
op: "func_name"
input: "input"
}
library {
function {
signature {
name: "func_name"
input_arg {
name: "arg0"
type: DT_BOOL
}
output_arg {
name: "retval0"
type: DT_BOOL
}
}
ret {
key: "retval0"
value: "arg0"
}
attr: {
key: "_input_shapes"
value: {
}
}
}
}
tensorflow/core/ir/importexport/tests/roundtrip/empty-value-attr.pbtxt
+72
View File
@@ -0,0 +1,72 @@
node {
name: "_tf.PartitionedCall"
op: "PartitionedCall"
attr {
key: "Tin"
value {
list {
}
}
}
attr {
key: "Tout"
value {
list {
type: DT_FLOAT
}
}
}
attr {
key: "f"
value {
func {
name: "foo"
}
}
}
}
library {
function {
signature {
name: "foo"
output_arg {
name: "constant"
type: DT_FLOAT
}
}
node_def {
name: "X"
op: "Const"
attr {
key: "dtype"
value {
type: DT_FLOAT
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_FLOAT
tensor_shape {
}
float_val: 7
}
}
}
}
ret {
key: "constant"
value: "X:output:0"
}
attr {
key: "_input_shapes"
value {
}
}
}
}
versions {
producer: 85
min_consumer: 12
}
tensorflow/core/ir/importexport/tests/roundtrip/error-message-with-source-info.pbtxt
+55
View File
@@ -0,0 +1,55 @@
node: {
name: "x"
op: "Placeholder"
attr: {
key: "shape"
value: {
shape: {
dim: {
size: -1
}
}
}
}
attr: {
key: "dtype"
value: {
type: DT_INT32
}
}
}
node: {
name: "y"
op: "Placeholder"
attr: {
key: "shape"
value: {
shape: {
dim: {
size: -1
}
}
}
}
attr: {
key: "dtype"
value: {
type: DT_INT32
}
}
}
node: {
name: "x_y_sum"
op: "Add"
input: "x"
input: "y"
attr: {
key: "T"
value: {
type: DT_INT32
}
}
}
versions: {
producer: 321
}
tensorflow/core/ir/importexport/tests/roundtrip/error-message-with-source-info.pbtxt.debug
+32
View File
@@ -0,0 +1,32 @@
files : [ "third_party/tensorflow/compiler/mlir/tensorflow/tests/graphdef2mlir/error-message-with-source-info.pbtxt.fake_py.debug"]
traces: {
key : "x@"
value: {
file_line_cols: {
line : 1
col : 1
}
}
}
traces: {
key : "x_y_sum@"
value: {
file_line_cols: {
line : 3
col : 1
}
file_line_cols: {
line : 4
col : 1
}
}
}
traces: {
key : "y@"
value: {
file_line_cols: {
line : 2
col : 1
}
}
}
tensorflow/core/ir/importexport/tests/roundtrip/error-message-with-source-info.pbtxt.fake_py.debug
+4
View File
@@ -0,0 +1,4 @@
x = value
y = value
math_ops.add(x, y, name='x_y_sum')
build_graph(out_dir)
tensorflow/core/ir/importexport/tests/roundtrip/feed-as-fetch.pbtxt
+13
View File
@@ -0,0 +1,13 @@
node {
name: "input"
op: "Placeholder"
attr {
key: "dtype"
value {
type: DT_INT32
}
}
}
versions {
producer: 27
}
tensorflow/core/ir/importexport/tests/roundtrip/feed-control-dep.pbtxt
+58
View File
@@ -0,0 +1,58 @@
node {
name: "input"
op: "Placeholder"
attr {
key: "dtype"
value {
type: DT_FLOAT
}
}
attr {
key: "shape"
value {
shape {
dim {
size: 1
}
}
}
}
}
node {
name: "variable_node"
op: "Const"
input: "^input"
attr {
key: "dtype"
value {
type: DT_FLOAT
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_FLOAT
tensor_shape {
}
float_val: 1
}
}
}
}
node {
name: "output_node"
op: "Identity"
input: "variable_node"
input: "^input"
attr {
key: "T"
value {
type: DT_FLOAT
}
}
}
library {
}
versions {
}

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