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

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

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 12:14:16 +08:00
commit 8a852e4b4e
36502 changed files with 9277225 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tensorflow/compiler/mlir/lite/experimental/common/BUILD
+22
View File
@@ -0,0 +1,22 @@
load("@rules_cc//cc:cc_library.bzl", "cc_library")
load("//tensorflow:tensorflow.default.bzl", "get_compatible_with_portable")
# copybara:uncomment package(default_applicable_licenses = ["//tensorflow:LICENSE"])
cc_library(
name = "outline_operations",
srcs = ["outline_operations.cc"],
hdrs = ["outline_operations.h"],
compatible_with = get_compatible_with_portable(),
visibility = ["//visibility:public"],
deps = [
"//tensorflow/compiler/mlir/lite:tensorflow_lite",
"//tensorflow/compiler/mlir/tensorflow:cluster_util",
"@com_google_absl//absl/strings",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Pass",
"@llvm-project//mlir:Support",
],
)
tensorflow/compiler/mlir/lite/experimental/common/outline_operations.cc
+212
View File
@@ -0,0 +1,212 @@
/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/compiler/mlir/lite/experimental/common/outline_operations.h"
#include <cassert>
#include <string>
#include "absl/strings/str_cat.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SetVector.h"
#include "llvm/ADT/SmallVector.h"
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/Block.h" // from @llvm-project
#include "mlir/IR/Builders.h" // from @llvm-project
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/IR/Matchers.h" // from @llvm-project
#include "mlir/IR/OpDefinition.h" // from @llvm-project
#include "mlir/IR/Operation.h" // from @llvm-project
#include "mlir/IR/Value.h" // from @llvm-project
#include "mlir/IR/Visitors.h" // from @llvm-project
#include "mlir/Support/LLVM.h" // from @llvm-project
#include "tensorflow/compiler/mlir/lite/ir/tfl_ops.h"
#include "tensorflow/compiler/mlir/lite/utils/utils.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/cluster_util.h"
namespace mlir {
namespace TFL {
namespace common {
bool IsConstantOrNone(Operation* op) {
return (op->getNumResults() == 1 &&
mlir::isa<NoneType>(op->getResult(0).getType())) ||
matchPattern(op, m_Constant()) || isa<QConstOp>(op);
}
// Pre-order traverse, adding results and BlockArgs to `been_defined` and
// collecting operands not contained within `been_defined`. If we encounter an
// operand that references a Value that has been defined (and added to
// `been_defined`) it is garuanteed that the Value definition is not contained
// in descedant node of reference, and given that the input DAG is valid, the
// definition is self-contained within `op` so it is not depended upon.
// Otherwise, the operand must have been defined somewhere above the Subgraph,
// so union with other operand dependencies.
llvm::SmallVector<Value> AccumulateOperandsDefinedAbove(
const llvm::SetVector<Operation*>& partition_ops) {
// Assuming that all are topologically sorted.
llvm::SetVector<Value> been_defined;
llvm::SetVector<Value> results;
auto update_from_op = [&](Operation* op) {
been_defined.insert(op->getResults().begin(), op->getResults().end());
for (Value input : op->getOperands()) {
if (been_defined.contains(input)) {
continue;
}
results.insert(input);
}
};
for (Operation* op : partition_ops) {
update_from_op(op);
op->walk<WalkOrder::PreOrder>([&](Block* nested_block) {
been_defined.insert(nested_block->getArguments().begin(),
nested_block->getArguments().end());
for (Operation& op : nested_block->getOperations()) update_from_op(&op);
});
}
return SmallVector<Value>(results.getArrayRef());
}
llvm::SmallVector<Value> AccumulateResultsDefinedWithin(
const llvm::SetVector<Operation*>& partition_ops) {
llvm::SmallVector<Value> values_for_results;
for (Operation* op : partition_ops) {
if (IsConstantOrNone(op)) {
continue;
}
for (Value output : op->getResults()) {
bool output_consumed_outside_subgraph = false;
for (Operation* consumer : output.getUsers()) {
if (llvm::all_of(partition_ops, [&](Operation* op) {
return !op->isAncestor(consumer);
})) {
output_consumed_outside_subgraph = true;
}
}
if (output_consumed_outside_subgraph) {
values_for_results.push_back(output);
}
}
}
return values_for_results;
}
// Compute signature for raised func from arugments and outputs of
// Operation partition.
llvm::SmallVector<Type> TypesFromValues(
const llvm::SmallVector<Value>& values) {
llvm::SmallVector<Type> types;
for (auto value : values) {
types.push_back(value.getType());
}
return types;
}
func::FuncOp BuildFuncOp(const Subgraph& subgraph, OpBuilder& builder,
ModuleOp& module, OpsAdded& ops_added) {
// The parameters of the new MLIR function are taken to be the union
// of all operands referenced by Operations within the subraph.
// Likewise the results of the function are any Value(s) that are defined
// within the subgraph and are referenced outside the subgraph.
llvm::SmallVector<Type> input_types =
TypesFromValues(subgraph.FuncArguments());
llvm::SmallVector<Type> return_types =
TypesFromValues(subgraph.FuncOutputs());
FunctionType function_type =
builder.getFunctionType(input_types, return_types);
std::string function_name = absl::StrCat("func_", subgraph.subgraph_id_);
func::FuncOp new_func = func::FuncOp::create(builder.getUnknownLoc(),
function_name, function_type);
new_func.setVisibility(func::FuncOp::Visibility::Private);
new_func.addEntryBlock();
// To form the body of the new function we need to clone each
// Operation along with its respective operands and result Values(s).
// The semantic of `Operation::clone` is copying given entity *into* this
// entity. The new FuncOp body is populated by cloning partitioned ops into
// it. Cloning Operation(s) will create cloned Value(s) for the results of a
// cloned op, but it needs a reference to the new operand Value(s) which are
// the result of the cloned ops. The approach is to traverse the subgraph in
// order, accumulating clones of defined Values into a `IRMapping`
// and pass that map to calls to clone ops.
OpBuilder function_builder(new_func.getBody());
// Preferred data structure for mapping MLIR values.
IRMapping values_in_scope;
// Function arguments can appear as operands, so they clone should
// be aware of them.
assert(subgraph.FuncArguments().size() == new_func.getNumArguments());
for (int i = 0; i < subgraph.FuncArguments().size(); ++i) {
Value original_value = subgraph.FuncArguments()[i];
Value new_func_arg = new_func.getArgument(i);
values_in_scope.map(original_value, new_func_arg);
}
for (Operation* op : subgraph.partition_ops_) {
function_builder.clone(*op, values_in_scope);
}
SmallVector<Value> return_operands;
for (Value result : subgraph.FuncOutputs()) {
Value cloned_output = values_in_scope.lookup(result);
return_operands.push_back(cloned_output);
}
mlir::func::ReturnOp::create(function_builder, new_func.getLoc(),
return_operands);
ops_added.func_op = new_func;
module.push_back(new_func);
return new_func;
}
void ExtractSubgraphToFunc(const Subgraph& subgraph, OpBuilder& builder,
ModuleOp& module, OpsAdded& ops_added) {
func::FuncOp func = BuildFuncOp(subgraph, builder, module, ops_added);
// We just use the location of the last ops in the subgraph as the location
// for the call_op.
Operation* last_output = subgraph.partition_ops_.back();
builder.setInsertionPoint(last_output);
auto call_op = func::CallOp::create(builder, last_output->getLoc(), func,
subgraph.FuncArguments());
ops_added.call_op = call_op;
// FuncOutputs refer to the original `Values` in input module which are now
// invalid after pulling out the defining ops. The values in
// `call_ops.getResult` refer to the clones of original `Values` which are now
// returned by the new `FuncOp`. We can replace each in `FuncOutputs` with
// clone in `call_op` to fix up.
for (int i = 0; i < subgraph.FuncOutputs().size(); ++i) {
Value output = subgraph.FuncOutputs()[i];
output.replaceAllUsesWith(call_op.getResult(i));
}
// Clear the subgraph.
// Those ops should be removed.
for (auto* op : subgraph.partition_ops_) {
if (IsConstantOrNone(op)) {
continue;
}
op->dropAllDefinedValueUses();
op->dropAllReferences();
op->erase();
}
// Ensure that users of the call op's results appear after the launch op in
// order to preserve the dominance property.
TF::ReorderOpResultUses(call_op);
}
} // namespace common
} // namespace TFL
} // namespace mlir
tensorflow/compiler/mlir/lite/experimental/common/outline_operations.h
+133
View File
@@ -0,0 +1,133 @@
/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_COMMON_OUTLINE_OPERATIONS_H_
#define TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_COMMON_OUTLINE_OPERATIONS_H_
#include <memory>
#include <string>
#include <utility>
#include "llvm/ADT/SetVector.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/ADT/StringRef.h"
#include "llvm/Support/Casting.h"
#include "llvm/Support/raw_os_ostream.h"
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/Attributes.h" // from @llvm-project
#include "mlir/IR/Block.h" // from @llvm-project
#include "mlir/IR/Builders.h" // from @llvm-project
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "mlir/IR/OpDefinition.h" // from @llvm-project
#include "mlir/IR/Operation.h" // from @llvm-project
#include "mlir/IR/UseDefLists.h" // from @llvm-project
#include "mlir/IR/Value.h" // from @llvm-project
#include "mlir/IR/ValueRange.h" // from @llvm-project
#include "mlir/IR/Visitors.h" // from @llvm-project
#include "mlir/Pass/Pass.h" // from @llvm-project
#include "mlir/Pass/PassRegistry.h" // from @llvm-project
#include "mlir/Support/LLVM.h" // from @llvm-project
#include "tensorflow/compiler/mlir/lite/ir/tfl_ops.h"
#include "tensorflow/compiler/mlir/lite/utils/utils.h"
namespace mlir {
namespace TFL {
namespace common {
// Returns true if the `op` is a constant-like op or produces none type.
bool IsConstantOrNone(Operation* op);
// Computes the list of Value(s) referenced by Subgraph Operations that are
// not defined within the Subgraph. Any such Value(s)
// are validly in-scope for the initial Operation. They must be either
// defined above the subgraph or appear as an argument to the containing func.
// These Value(s) are taken to be the arguments of the new raised func.
// An operand dependency is a Value referenced anywhere in an Op
// that is defined above the Op. All SSA Values are assigned/defined in a
// BlockArg or as a result of an Operation.
llvm::SmallVector<Value> AccumulateOperandsDefinedAbove(
const llvm::SetVector<Operation*>& partition_ops);
// Similar to `AccumulateOperandsDefinedAbove()`, computes the Value(s) that are
// defined within a Subgraph and referenced in a descendant Operation. These
// Values(s) are to be returned by the new raised function.
llvm::SmallVector<Value> AccumulateResultsDefinedWithin(
const llvm::SetVector<Operation*>& partition_ops);
// Represents a view of a set of mlir Operations that form a subgraph of the
// entire Module's DAG. `Subgraph` can be thought of as segment of sequential
// Operations within a func definition. Additional facts:
// 1. Subgraphs are restricted to a single Block. They do not span
// branching instructions. Thus the subgraph is a simple 1-degree path.
// 2. All Operations in a subgraph belong to the same block in a
// funtion body.
// 3. Function bodies are assumed to have only one block in some places.
class Subgraph {
// Set vector preserves insertion order, must insert Ops in topological order.
public:
const llvm::SetVector<Operation*> partition_ops_;
// Subgraphs are given a unique incremented integer id based on when
// they were encountered in this pass.
const int subgraph_id_;
const llvm::StringRef dialect_namespace_;
Subgraph(const llvm::SetVector<Operation*> partition_ops, int num_subgraphs)
: partition_ops_(partition_ops),
subgraph_id_(num_subgraphs),
func_arguments_(AccumulateOperandsDefinedAbove(partition_ops)),
func_outputs_(AccumulateResultsDefinedWithin(partition_ops)) {}
const llvm::SmallVector<Value>& FuncArguments() const {
// `Value`s in MLIR library are implemented as having "value semantics"
// see "llvm/llvm-project/mlir/include/mlir/IR/Value.h" so copying is fine.
return func_arguments_;
}
const llvm::SmallVector<Value>& FuncOutputs() const { return func_outputs_; }
private:
// Compute once at construction and save as field.
const llvm::SmallVector<Value> func_arguments_;
const llvm::SmallVector<Value> func_outputs_;
};
// Helper data structure for output parameters to `ExtractSubgraphToFunc`.
// `ExtractSubgraphToFunc` adds exactly two "new" `Operations`, a FuncOp and
// a CallOp. Pass these back to the caller for setting more specific attributes
// after graph mutation has taken place.
struct OpsAdded {
mlir::func::FuncOp func_op;
mlir::func::CallOp call_op;
};
// Given a `Subgraph` containing a sequence of adjacent `Operations` from
// the `module`, raise these `Operations` (and any ops contained nested within)
// to the body of a new seperate root level function. Replace in their current
// location with a `CallOp` which invokes said `FuncOp`. The inputs to
// this new functions are taken to be the `Values` that appear as operands
// to ops in the subgraph, which are not self-contained within the subgraph.
// The outputs of this function are taken to be the results of ops in the
// subgraph which are referenced as operands outside of the subgraph.
// Also refer to documention of `AccumulateOperandsDefinedAbove` &
// `AccumulateResultsDefinedWithin`.
void ExtractSubgraphToFunc(const Subgraph& subgraph, OpBuilder& builder,
ModuleOp& module, OpsAdded& ops_added);
} // namespace common
} // namespace TFL
} // namespace mlir
#endif // TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_COMMON_OUTLINE_OPERATIONS_H_
tensorflow/compiler/mlir/lite/experimental/remat/BUILD
+44
View File
@@ -0,0 +1,44 @@
load("@rules_cc//cc:cc_library.bzl", "cc_library")
load("//tensorflow:tensorflow.bzl", "tf_cc_test")
load("//tensorflow:tensorflow.default.bzl", "get_compatible_with_portable")
# copybara:uncomment package(default_applicable_licenses = ["//tensorflow:LICENSE"])
cc_library(
name = "rematerializer",
srcs = ["rematerializer.cc"],
hdrs = ["rematerializer.h"],
deps = [
],
)
tf_cc_test(
name = "rematerializer_test",
size = "small",
srcs = ["rematerializer_test.cc"],
deps = [
":rematerializer",
"@com_google_googletest//:gtest_main",
],
)
cc_library(
name = "metadata_util",
srcs = ["metadata_util.cc"],
hdrs = ["metadata_util.h"],
compatible_with = get_compatible_with_portable(),
visibility = ["//visibility:public"],
deps = [
"//tensorflow/compiler/mlir/lite:control_edges",
],
)
tf_cc_test(
name = "metadata_util_test",
size = "small",
srcs = ["metadata_util_test.cc"],
deps = [
":metadata_util",
"@com_google_googletest//:gtest_main",
],
)
tensorflow/compiler/mlir/lite/experimental/remat/metadata_util.cc
+124
View File
@@ -0,0 +1,124 @@
/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/compiler/mlir/lite/experimental/remat/metadata_util.h"
#include <cstddef>
#include <cstdint>
#include <string>
#include <utility>
#include <vector>
namespace {
// We serialize unsigneds as protobuf varints, i.e., in chunks of 7 bits each.
constexpr int kMod = (1 << 7);
void Serialize(std::string* out, uint32_t value) {
for (; value >= kMod; value /= kMod) {
out->push_back(value % kMod + kMod);
}
out->push_back(value);
}
bool Parse(const char** data, size_t* size, uint32_t* out) {
*out = 0;
uint32_t mul = 1;
for (bool done = false; !done;
mul *= kMod, done = !(**data & kMod), ++*data, --*size) {
if (*size == 0) {
return false;
}
*out += static_cast<unsigned char>(**data) % kMod * mul;
}
return true;
}
// Signed ints are zigzag-encoded as unsigned varints, [..., -2, -1, 0, 1, 2,
// ...] ->
// [..., 3, 1, 0, 2, 4, ...].
void Serialize(std::string* out, int32_t value) {
Serialize(out, static_cast<uint32_t>(
value < 0 ? static_cast<uint32_t>(-(value + 1)) * 2 + 1
: static_cast<uint32_t>(value) * 2));
}
bool Parse(const char** data, size_t* size, int32_t* out) {
uint32_t value = 0;
if (!Parse(data, size, &value)) {
return false;
}
const int32_t magnitude = value / 2;
*out = (value % 2) ? (-magnitude - 1) : magnitude;
return true;
}
// Pairs are serialized as the concatenation of their elements' serialization.
template <class First, class Second>
void Serialize(std::string* out, const std::pair<First, Second>& in) {
Serialize(out, in.first);
Serialize(out, in.second);
}
template <class First, class Second>
bool Parse(const char** data, size_t* size, std::pair<First, Second>* out) {
return Parse(data, size, &(out->first)) && Parse(data, size, &(out->second));
}
// Vectors are serialized as the concetation of the serialization of their size
// and the the serializations of their elements.
template <class Value>
void Serialize(std::string* out, const std::vector<Value>& in) {
Serialize(out, static_cast<uint32_t>(in.size()));
for (const auto& val : in) {
Serialize(out, val);
}
}
template <class T>
bool Parse(const char** data, size_t* size, std::vector<T>* out) {
uint32_t num_elems = 0;
if (!Parse(data, size, &num_elems)) {
return false;
}
out->assign(num_elems, T{});
for (auto& elem : *out) {
if (!Parse(data, size, &elem)) {
return false;
}
}
return true;
}
} // namespace
namespace tflite {
std::string SerializeModelControlDependencies(
const ModelControlDependencies& in) {
std::string out;
Serialize(&out, kModelControlDependenciesMetadataVersion);
Serialize(&out, in);
return out;
}
bool ParseModelControlDependencies(const char* data, size_t size,
ModelControlDependencies* out) {
out->clear();
uint32_t version = 0;
return Parse(&data, &size, &version) &&
(version == kModelControlDependenciesMetadataVersion) &&
Parse(&data, &size, out) && (size == 0);
}
} // namespace tflite
tensorflow/compiler/mlir/lite/experimental/remat/metadata_util.h
+62
View File
@@ -0,0 +1,62 @@
/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
/// \file
///
/// Functions for serializiation/deserialization of control dependency
/// information to/from model metadata.
///
#ifndef TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_REMAT_METADATA_UTIL_H_
#define TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_REMAT_METADATA_UTIL_H_
#include <cstddef>
#include <cstdint>
#include <string>
#include <vector>
#include "tensorflow/compiler/mlir/lite/utils/control_edges.h"
namespace tflite {
/// Control dependencies for the model is the collection of control dependencies
/// for its subgraphs.
using ModelControlDependencies = std::vector<ControlEdges>;
/// Serializes `in` into the returned string. The result is parseable with
/// ParseModelControlDependencies.
std::string SerializeModelControlDependencies(
const ModelControlDependencies& in);
/// Deserializes `*out` from a character buffer of size `size` at `data`.
/// Returns true iff successful. `*out` needn't be empty before invocation.
/// When returning false, `*out`'s state is undefined.
bool ParseModelControlDependencies(const char* data, size_t size,
ModelControlDependencies* out);
/// The key under which to store the serialized control dependencies in the
/// model's metadata.
constexpr char kModelControlDependenciesMetadataKey[] =
"model_control_dependencies";
/// To allow future changes to the format, serialized control dependency data
/// will contain a version; this constant is the version that will be used for
/// serialization. For deserialization, past versions should remain parseable.
constexpr uint32_t kModelControlDependenciesMetadataVersion = 1;
inline constexpr char kModelUseStablehloTensorKey[] = "keep_stablehlo_constant";
} // namespace tflite
#endif // TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_REMAT_METADATA_UTIL_H_
tensorflow/compiler/mlir/lite/experimental/remat/metadata_util_test.cc
+55
View File
@@ -0,0 +1,55 @@
/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/compiler/mlir/lite/experimental/remat/metadata_util.h"
#include <cstdint>
#include <limits>
#include <string>
#include <gmock/gmock.h>
#include <gtest/gtest.h>
namespace tflite {
namespace {
class MetadataSerializerTest : public ::testing::Test {
protected:
static constexpr auto kHuge = std::numeric_limits<int32_t>::max();
static constexpr auto kTiny = std::numeric_limits<int32_t>::min();
std::string RoundTrip(const ModelControlDependencies &in) const {
ModelControlDependencies out = {{{-1, -1}}};
const std::string serialized =
tflite::SerializeModelControlDependencies(in);
return tflite::ParseModelControlDependencies(serialized.data(),
serialized.size(), &out)
? (out == in) ? "ok" : "mismatch"
: "malformed";
}
};
TEST_F(MetadataSerializerTest, nothing) { EXPECT_THAT(RoundTrip({}), "ok"); }
TEST_F(MetadataSerializerTest, something) {
EXPECT_THAT(
RoundTrip({{{1, 2}, {2, 3}, {4, 5}},
{},
{{kHuge, kTiny}, {kTiny, kHuge}, {kHuge - 1, kTiny + 1}},
{{1, 0}}}),
"ok");
}
} // namespace
} // namespace tflite
tensorflow/compiler/mlir/lite/experimental/remat/rematerializer.cc
+381
View File
@@ -0,0 +1,381 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/compiler/mlir/lite/experimental/remat/rematerializer.h"
#include <algorithm>
#include <map>
#include <tuple>
#include <utility>
#include <vector>
namespace mlir {
namespace TFL {
namespace {
// Helper functions for sorted + deduped int vectors.
std::tuple<std::vector<int>::iterator, bool> Find(const int item,
std::vector<int>& items) {
const auto iter = std::lower_bound(items.begin(), items.end(), item);
return std::make_tuple(iter, iter != items.end() && *iter == item);
}
void Insert(const int item, std::vector<int>& items) {
const auto [iter, found] = Find(item, items);
if (!found) items.insert(iter, item);
}
void Erase(const int item, std::vector<int>& items) {
const auto [iter, found] = Find(item, items);
if (found) items.erase(iter);
}
} // namespace
int Rematerializer::AddOperation(const bool is_stateful) {
operations_.emplace_back();
operations_.back().is_stateful = is_stateful;
return operations_.size() - 1;
}
int Rematerializer::AddTensor(const SizeT size) {
tensors_.emplace_back();
tensors_.back().size = size;
return tensors_.size() - 1;
}
void Rematerializer::DelUse(const int ioperation, const int itensor) {
auto& tensor = tensors_[itensor];
auto& operation = operations_[ioperation];
const auto& size = tensor.size;
// Was the dependence to be deleted the first/last (or both) use of this
// tensor?
const bool was_first_use =
(!tensor.operations.empty() && ioperation == tensor.first_use());
const bool was_last_use =
(!tensor.operations.empty() && ioperation == tensor.last_use());
Erase(ioperation, tensor.operations);
Erase(itensor, operation.tensors);
if (was_first_use) {
operation.alloc -= size;
if (!was_last_use) {
operations_[tensor.first_use()].alloc += size;
}
}
if (was_last_use) {
operation.dealloc -= size;
if (!was_first_use) {
operations_[tensor.last_use()].dealloc += size;
}
}
}
void Rematerializer::AddUse(const int ioperation, const int itensor) {
auto& tensor = tensors_[itensor];
auto& operation = operations_[ioperation];
const auto& size = tensor.size;
const bool will_be_first_use =
tensor.operations.empty() || ioperation < tensor.first_use();
const bool will_be_last_use =
tensor.operations.empty() || ioperation > tensor.last_use();
if (will_be_first_use) {
operation.alloc += size;
if (!will_be_last_use) {
operations_[tensor.first_use()].alloc -= size;
}
}
if (will_be_last_use) {
operation.dealloc += size;
if (!will_be_first_use) {
operations_[tensor.last_use()].dealloc -= size;
}
}
Insert(ioperation, tensor.operations);
Insert(itensor, operation.tensors);
}
Rematerializer::SizeT Rematerializer::MaxSavings(const int begin, const int end,
const int peak_loc) const {
SizeT max_savings = 0;
// We're looking at the outputs of all operators in [`begin`, `end`). If an
// output is alive after `peak_loc`, rematerializing the operator anywhere
// after `peak_loc` would lower the memory profile at `peak_loc`.
for (int ioperation = begin; ioperation != end; ++ioperation) {
for (const int itensor : operations_[ioperation].tensors) {
if (const Tensor& tensor = tensors_[itensor];
tensor.first_use() == ioperation /* output */ &&
tensor.last_use() > peak_loc /* used later */) {
max_savings += tensor.size;
}
}
}
return max_savings;
}
std::tuple<Rematerializer::SizeT, Rematerializer::RematSpec>
Rematerializer::FindBestRemat(const SizeT min_savings, const int begin_len,
const int end_len) const {
const auto peak = GetPeakMemory();
SizeT best_peak_mem = peak.size;
RematSpec best_remat = {};
for (int len = begin_len; len < end_len; ++len) {
std::vector<std::tuple<SizeT, int, int>> pre_screen;
for (int begin = 0, end = begin + len; end <= peak.op_index;
++begin, ++end) {
if (!std::any_of(operations_.begin() + begin, operations_.begin() + end,
[](const Operation& s) { return s.is_stateful; })) {
if (const auto max_savings = MaxSavings(begin, end, peak.op_index);
max_savings >= min_savings) {
pre_screen.emplace_back(max_savings, begin, end);
}
}
}
std::sort(pre_screen.begin(), pre_screen.end());
for (; !pre_screen.empty(); pre_screen.pop_back()) {
const auto& [max_savings, begin, end] = pre_screen.back();
const auto insert_before = FindBestRematPoint(begin, end, peak.op_index);
if (insert_before == operations_.size()) {
continue;
}
const RematSpec this_remat = {begin, end, insert_before};
if (const auto new_peak = GetPeakMemory(this_remat);
new_peak.size < best_peak_mem &&
peak.size >= new_peak.size + min_savings) {
best_peak_mem = new_peak.size;
best_remat = this_remat;
}
// If the actual savings achieved is bigger than the maximal savings that
// can be possibly achieved, leave early.
if (peak.size >= max_savings + best_peak_mem) {
break;
}
}
// We already found one savings for this length.
if (peak.size >= min_savings + best_peak_mem) {
break;
}
}
return std::make_tuple(best_peak_mem, best_remat);
}
std::vector<Rematerializer::MemSpec> Rematerializer::GetDeltas(
const RematSpec& remat) const {
// We're cloning operations [remat.begin, remat.end) at position
// remat.insert. We store changes to the alloc/dealloc sizes due to the
// insertion in a vector `delta`: A change `c_alloc` of `operations_[i].alloc`
// as `delta[i] += c_alloc`, and a change `c_dealloc` of
// `operations_[i].dealloc` as `delta[i+1] -= c_dealloc`.
std::vector<MemSpec> deltas;
if (remat.begin == remat.end) {
return deltas;
}
// Helper lambda: converts an operation index in the original range to its
// equivalent in the cloned range.
const auto source_to_target = [&](int i) {
return i + (remat.insert - remat.begin);
};
// Helper struct: bundles first and last use of a tensor within
// a contiguous range of operations.
struct TensorUse {
int first_use;
int last_use;
};
// For all tensors in the operation range to be cloned, store their first and
// last use within that range. This will be needed below to decide whether a
// tensor's life will be extended, or if it is to be rematerialized.
std::map<int, TensorUse> source_uses;
for (int ioperation = remat.begin; ioperation < remat.end; ++ioperation) {
const auto& operation = operations_[ioperation];
for (const int itensor : operation.tensors) {
// insertion will only happen for the first use.
const auto [iter, inserted] = source_uses.emplace(
itensor,
TensorUse{/*first_use=*/ioperation, /*last_use=*/ioperation});
if (!inserted) {
// Otherwise, update last_use.
iter->second.last_use = ioperation;
}
}
}
deltas.reserve(2 * source_uses.size());
for (const auto& [itensor, source] : source_uses) {
auto& tensor = tensors_[itensor];
const TensorUse global = {tensor.first_use(), tensor.last_use()};
auto add_alloc = [&](int pos) { deltas.emplace_back(pos, tensor.size); };
auto add_dealloc = [&](int pos) {
deltas.emplace_back(pos + 1, -tensor.size);
};
auto del_dealloc = [&](int pos) {
deltas.emplace_back(pos + 1, tensor.size);
};
if (global.first_use < remat.begin) {
// The tensor is created before the source range.
if (global.last_use < remat.insert) {
// It currently gets deallocated before the newly inserted range, so we
// need to extend its lifetime: It will now be deallocated at its last
// use in the inserted range.
del_dealloc(global.last_use);
add_dealloc(source_to_target(source.last_use));
}
} else {
// Tensor is created in the source range. It will be rematerialized in the
// cloned range.
add_alloc(source_to_target(source.first_use));
if (global.last_use < remat.insert) {
// The last use of the original tensor is before the target range, so
// the lifetime of the rematerialized tensor ends with the target range.
add_dealloc(source_to_target(source.last_use));
} else {
// There are uses of the original tensor after the cloned range. They
// can be replaced with uses of the rematerialized tensor, and the
// original tensor can be deallocated after its last use before the
// rematerialization.
add_dealloc(*std::partition_point(
tensor.operations.rbegin(), tensor.operations.rend(),
[&](int i) { return i >= remat.insert; }));
}
}
}
std::sort(deltas.begin(), deltas.end(), ByOpIndex);
return deltas;
}
Rematerializer::MemProfile Rematerializer::GetMemProfile(
const RematSpec& remat) const {
const auto num_inserted = remat.end - remat.begin;
std::vector<SizeT> profile(operations_.size() + num_inserted);
MapMem([&](const MemSpec& m) { profile[m.op_index] = m.size; }, remat);
return profile;
}
Rematerializer::MemSpec Rematerializer::GetPeakMemory(
const RematSpec& remat) const {
MemSpec peak;
MapMem([&](const MemSpec& m) { peak = std::max(m, peak, BySize); }, remat);
return peak;
}
int Rematerializer::FindBestRematPoint(const int begin, const int end,
const int peak_loc) const {
// All tensors necessarily have their last use before or at the last operator
// (which is the yield operation of a function), so an unchanged best ==
// operations_.size() value means that there is no valid configuration.
int best = operations_.size();
// All tensors whose first use is in [begin, end) and that are not destroyed
// until after peak_loc.
for (int ioperation = begin; ioperation < end; ++ioperation) {
for (const int itensor : operations_[ioperation].tensors) {
if (const auto& tensor = tensors_[itensor];
tensor.first_use() >= begin && tensor.first_use() < end &&
tensor.last_use() > peak_loc) {
for (const int ioperation : tensor.operations) {
// We return the earliest dependence on any output tensor.
if (ioperation > peak_loc && ioperation < best) {
best = ioperation;
break;
}
}
}
}
}
return best;
}
void Rematerializer::Remat(const RematSpec& remat) {
const int num_inserted = remat.end - remat.begin;
// Rewrite all operation indices.
for (auto& tensor : tensors_) {
std::for_each(std::lower_bound(tensor.operations.begin(),
tensor.operations.end(), remat.insert),
tensor.operations.end(),
[&](int& iop) { iop += num_inserted; });
}
operations_.insert(operations_.begin() + remat.insert, num_inserted, {});
// Copy all tensors dependencies of the old region to the new region.
// For any output tensor, a new tensor will be created.
std::vector<std::pair<int, int>> new_tensors;
for (int iop_old = remat.begin, iop_new = remat.insert; iop_old < remat.end;
++iop_old, ++iop_new) {
for (const auto itensor : operations_[iop_old].tensors) {
if (tensors_[itensor].first_use() == iop_old) {
new_tensors.emplace_back(itensor, AddTensor(tensors_[itensor].size));
}
AddUse(iop_new, itensor);
}
}
std::sort(new_tensors.begin(), new_tensors.end());
// Let all inserted + downstream operations refer to the new tensors.
for (int iop = remat.insert; iop < operations_.size(); ++iop) {
// We have to copy, otherwise we're modifying the set during the iteration.
for (const int old_tensor : std::vector<int>(operations_[iop].tensors)) {
const auto new_tensor =
std::lower_bound(new_tensors.begin(), new_tensors.end(),
std::make_pair(old_tensor, 0));
if (new_tensor != new_tensors.end() && new_tensor->first == old_tensor) {
DelUse(iop, old_tensor);
AddUse(iop, new_tensor->second);
}
}
}
}
void Rematerializer::RunGreedyAlgorithm(const int max_cost,
const int max_block_length,
const SizeT min_savings) {
const bool unlimited_cost = (max_cost < 0);
for (int min_block_length = 1, cost = 0;
min_block_length <= max_block_length &&
(unlimited_cost || cost <= max_cost);
min_block_length *= 2) {
while (unlimited_cost || cost <= max_cost) {
const auto [peak, remat] = FindBestRemat(
/*min_savings*/ min_savings,
/*begin_len=*/min_block_length,
/*end_len=*/
std::min(1 + (unlimited_cost
? max_block_length
: std::min(max_block_length, max_cost - cost)),
2 * min_block_length));
if (remat.begin == remat.end) break;
Remat(remat);
ApplyRemat(remat);
cost += (remat.end - remat.begin);
}
}
}
} // namespace TFL
} // namespace mlir
tensorflow/compiler/mlir/lite/experimental/remat/rematerializer.h
+262
View File
@@ -0,0 +1,262 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_REMAT_REMATERIALIZER_H_
#define TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_REMAT_REMATERIALIZER_H_
// This file declares the Rematerializer class, which is used by an MLIR-based
// set of transformations for TFLite IR that lower memory usage by redoing
// operations with small inputs and large outputs instead of keeping the result
// in memory. This class allows us to compactly and efficiently represent the
// (idealized) memory profile of a TFLite graph and simulate the effect of
// re-inserting operations on that memory profile.
#include <algorithm>
#include <cinttypes>
#include <cstdint>
#include <tuple>
#include <vector>
namespace mlir {
namespace TFL {
// A class that
// (1) Encodes in concise form the memory requirements of a computational graph
// (2) Allows for the fast simulation of changes to the peak memory requirement
// under rematerialization of intermediate results in the graph
// (3) Implements a greedy algorithm for finding rematerializations of
// intermediate results in that graph to lower peak memory requirements.
class Rematerializer {
public:
Rematerializer() = default;
virtual ~Rematerializer() = default;
// The type used for memory sizes (in bytes) and differences thereof.
using SizeT = int64_t;
// The memory profile: The i-th element gives the amount of memory
// that is needed when performing the i-th operation. This is the
// sum of the sizes of
//
// (1) input tensors of that operation,
// (2) output tensors of that operation,
// (3) output tensors of preceding operations that are input tensors
// of subsequent operations.
using MemProfile = std::vector<SizeT>;
// Used for specifying memory consumption at a certain operation in the
// computational graph.
struct MemSpec {
int op_index; // The index of the operation
SizeT size; // The amount of memory needed in order to execute this
// operation, i.e., the sum of input and output sizes and the
// sizes of outputs of previous operations that are needed as
// inputs of subsequent operations.
explicit MemSpec(int op_index = 0, SizeT size = 0)
: op_index(op_index), size(size) {}
};
static bool BySize(const MemSpec& a, const MemSpec& b) {
return std::tie(a.size, a.op_index) < std::tie(b.size, b.op_index);
}
static bool ByOpIndex(const MemSpec& a, const MemSpec& b) {
return std::tie(a.op_index, a.size) < std::tie(b.op_index, b.size);
}
// Specifies an elementary rematerialization operation: The operations in
// operations [`begin`, `end`) will be rescheduled before operation `insert`.
// A valid `RematSpec` requires begin <= end <= insert <= number of
// operations. Note that (1) `end` is exclusive -- begin == end signifies a
// trivial RematSpec (no operation will be rescheduled), (2) the
// zero-initialized RematSpec {} is trivial and always valid.
struct RematSpec {
int begin;
int end;
int insert;
};
// Gives the peak memory location and size after inserting operations
// according to `remat` (but doesn't actually insert them.) Ties are broken
// towards later locations. `remat` must be valid (see above).
MemSpec GetPeakMemory(const RematSpec& remat = {}) const;
// Gives memory profile after inserting operations according to `remat` (but
// doesn't actually insert them). `remat` must be valid (see above).
MemProfile GetMemProfile(const RematSpec& remat = {}) const;
// Runs the greedy incremental block algorithm: Finds a sequence of
// rematerializations of block size up to max_block_length, each reducing peak
// memory by at least min_savings. If max_cost >= 0, at most max_cost
// operations will be re-inserted. For each rematerialization found,
// ApplyRemat is invoked (which can be used to apply the rematerialization to
// the higher- level representation, e.g., MLIR, flatbuffer, ...)
void RunGreedyAlgorithm(int max_cost, int max_block_length,
SizeT min_savings);
virtual void ApplyRemat(const RematSpec& remat) {}
protected:
// Rematerializes the outputs of the operations [`remat.begin`, `remat.end`)
// before operation remat.insert by copying that operation range before
// remat.insert and updating tensor references so that any operation that can
// will make use of the rematerialized outputs rather than the original ones.
// `remat` must be valid (see above).
void Remat(const RematSpec& remat);
// The protected methods below are to be used by derived classes to create the
// low-level (this class) representation from a high-level one.
// Creates a new tensor-like object that takes `size` bytes. Returns a
// contiguous increasing index for each new object, starting at 0.
int AddTensor(SizeT size);
// Creates an operation. If `is_stateful`, the operation (and any block of
// operations containing it) will never be considered for rematerialization.
// Returns a contiguous increasing index for each new object, starting at 0.
int AddOperation(bool is_stateful);
// The operator with index `ioperation` will be assumed to produce and/or
// consume the tensor with index `itensor`. NoOp if that's already the case.
// The arguments must be valid indices (i.e., obtained with
// `AddOperation`/`AddTensor`).
void AddUse(int ioperation, int itensor);
// Undoes an AddUse(ioperation, itensor). NoOp if there was no prior `AddUse`.
// The arguments must be valid indices (i.e., obtained with
// `AddOperation`/`AddTensor`).
void DelUse(int ioperation, int itensor);
private:
// Find the best remat operation that saves at least `min_savings` bytes for a
// block of operators with a length is between [`begin_len`, `end_len`).
// 'Best' means with the highest savings, ties are broken towards shorter
// blocks.
std::tuple<SizeT, RematSpec> FindBestRemat(SizeT min_savings, int begin_len,
int end_len) const;
// Optimization: Estimate (from above) the remat savings of instruction block
// [begin, end) after operation `peak_location`
SizeT MaxSavings(int begin, int end, int peak_loc) const;
// If I want to remat ops [begin, end) after the op at operation `peak_loc`,
// find the latest point at which to reinsert them (the op before which to
// insert.)
int FindBestRematPoint(int begin, int end, int peak_loc) const;
// The memory objects.
struct Tensor {
SizeT size; // The size of the object (in bytes.)
std::vector<int> operations; // The operations it is used in. This vector
// is kept sorted + unique.
// The operation that makes the first use of this tensor.
int first_use() const { return *operations.begin(); }
// The operation that makes the last use of this tensor.
int last_use() const { return *operations.rbegin(); }
};
// The operators.
struct Operation {
bool is_stateful = false; // Results of an Operation can be rematerialized
// only if `!is_stateful`. This probably should
// be replaced with a more-fine grained
// approach--for example, the results of a "read
// resource variable" operation can be
// rematerialized as long as this doesn't happen
// after the corresponding "write resource
// variable" operation.
std::vector<int> tensors; // The tensors that are used (input or output) by
// this operation. They needn't correspond to
// tensors in the TF graph -- we may add fake
// tensors to model memory consumed in addition
// to input and output tensors. This vector is
// kept sorted + unique.
SizeT alloc = 0; // The number of bytes that need to be allocated before
// this operation.
SizeT dealloc = 0; // The number of bytes that can be deallocated after
// this operation.
};
// Given the current state of `operations_` and `tensors_`, return a vector of
// corrections that transform the current memory profile into the one that we
// would get after applying `remat`.
//
// The memory profile of a sequence of operations is the partial sum of the
// sizes of the allocations that are necessary before an operation and the
// negative sizes of the deallocations that are possible after the previous
// operation.
//
// If we modify the operation sequence by cloning an operation range, that
// memory profile will change--cloning makes it necessary to extend the
// lifetime of some tensors, while other tensors can be deallocated early and
// rematerialized later.
//
// This method represents these changes in compact form: It returns a vector
// of (position of operation, delta) pairs in lexicographic order; one
// obtains the memory profile after `remat` by adding the deltas from any
// entries (i, delta) to the i-th entry of the partial sum.
//
// This allows us to efficiently compute the change to the peak of a memory
// profile due to cloning an operation range without having to actually clone
// that range and without having to build a profile vector.
//
// The returned vector has at most 2 entries for each tensor referenced in
// [remat.begin, remat.end). There may be multiple entries for a single
// operation position; operation positions refer to the sequence *after*
// cloning [`remat.begin`, `remat.end`) before `remat.insert`.
std::vector<MemSpec> GetDeltas(const RematSpec& remat) const;
// Helper template: Iterates through all `MemSpec`s (i.e., operation
// index/memory usage pairs) for the current graph in operation order and
// calls `mapper` on them. This is an optimization -- by instantiating with an
// appropriate `Mapper`, it allows us to e.g. compute the peak memory without
// having to instantiate an actual memory profile vector.
template <class Mapper>
void MapMem(const Mapper& mapper, const RematSpec& remat) const {
const auto deltas = GetDeltas(remat);
const auto len = (remat.end - remat.begin);
auto idelta = deltas.begin();
for (MemSpec m; m.op_index < operations_.size() + len; ++m.op_index) {
// Are we in the cloned portion of the new operation sequence?
// Then all alloc/dealloc information must come from deltas.
const bool patch =
(m.op_index >= remat.insert) && (m.op_index < remat.insert + len);
// Are we past the insertion portion of the new operation sequence?
// Then we need to convert indices back to the original sequence.
const int shift = (m.op_index >= remat.insert + len) ? len : 0;
m.size += patch ? 0 : operations_[m.op_index - shift].alloc;
// deltas is sorted by location; apply any corrections to the current
// operator.
for (; idelta != deltas.end() && idelta->op_index == m.op_index;
++idelta) {
m.size += idelta->size;
}
mapper(m);
m.size -= patch ? 0 : operations_[m.op_index - shift].dealloc;
}
}
std::vector<Operation> operations_;
std::vector<Tensor> tensors_;
};
} // namespace TFL
} // namespace mlir
#endif // TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_REMAT_REMATERIALIZER_H_
tensorflow/compiler/mlir/lite/experimental/remat/rematerializer_test.cc
+519
View File
@@ -0,0 +1,519 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/compiler/mlir/lite/experimental/remat/rematerializer.h"
#include <algorithm>
#include <array>
#include <cstdlib>
#include <initializer_list>
#include <random>
#include <vector>
#include <gmock/gmock.h>
#include <gtest/gtest.h>
namespace mlir {
namespace TFL {
namespace {
using ::testing::ElementsAre;
using ::testing::ElementsAreArray;
using ::testing::Eq;
using ::testing::FieldsAre;
using ::testing::StrictMock;
class RematTest : public ::testing::Test {
protected:
class TestableRematerializer : public Rematerializer {
public:
using Rematerializer::AddOperation;
using Rematerializer::AddTensor;
using Rematerializer::AddUse;
using Rematerializer::DelUse;
using Rematerializer::Remat;
};
TestableRematerializer r_;
};
TEST_F(RematTest, TensorUseSimple) {
for (int i = 0; i < 6; ++i) {
r_.AddOperation(/*is_stateful=*/false);
r_.AddTensor(/*size=*/1 << i);
}
r_.AddUse(/*ioperation=*/2, /*itensor=*/2);
EXPECT_THAT(r_.GetMemProfile(), ElementsAre(0, 0, 4, 0, 0, 0));
EXPECT_THAT(r_.GetPeakMemory(), FieldsAre(Eq(2), Eq(4)));
r_.AddUse(/*ioperation=*/2, /*itensor=*/2);
EXPECT_THAT(r_.GetMemProfile(), ElementsAre(0, 0, 4, 0, 0, 0));
EXPECT_THAT(r_.GetPeakMemory(), FieldsAre(Eq(2), Eq(4)));
r_.AddUse(/*ioperation=*/4, /*itensor=*/2);
EXPECT_THAT(r_.GetMemProfile(), ElementsAre(0, 0, 4, 4, 4, 0));
EXPECT_THAT(r_.GetPeakMemory(), FieldsAre(Eq(4), Eq(4)));
r_.DelUse(/*ioperation=*/2, /*itensor=*/2);
EXPECT_THAT(r_.GetMemProfile(), ElementsAre(0, 0, 0, 0, 4, 0));
EXPECT_THAT(r_.GetPeakMemory(), FieldsAre(Eq(4), Eq(4)));
r_.DelUse(/*ioperation=*/2, /*itensor=*/2);
EXPECT_THAT(r_.GetMemProfile(), ElementsAre(0, 0, 0, 0, 4, 0));
EXPECT_THAT(r_.GetPeakMemory(), FieldsAre(Eq(4), Eq(4)));
r_.DelUse(/*ioperation=*/4, /*itensor=*/2);
EXPECT_THAT(r_.GetMemProfile(), ElementsAre(0, 0, 0, 0, 0, 0));
EXPECT_THAT(r_.GetPeakMemory(), FieldsAre(Eq(5), Eq(0)));
}
TEST_F(RematTest, TensorUseMany) {
constexpr int n = 6;
for (int i = 0; i < n; ++i) {
r_.AddUse(/*ioperation=*/r_.AddOperation(/*is_stateful=*/false),
/*itensor=*/r_.AddTensor(1 << (n - i - 1)));
}
for (int i = 0; i < n; ++i) {
r_.AddUse(/*ioperation=*/r_.AddOperation(/*is_stateful=*/false),
/*itensor=*/n - 1 - i);
}
EXPECT_THAT(r_.GetMemProfile(), ElementsAreArray({32, 48, 56, 60, 62, 63, 63,
62, 60, 56, 48, 32}));
EXPECT_THAT(r_.GetPeakMemory(), FieldsAre(Eq(6), Eq(63)));
}
TEST_F(RematTest, PeakTiesAreBrokenInFavorOfLaterOperations) {
r_.AddUse(/*ioperation=*/r_.AddOperation(/*is_stateful=*/false),
/*itensor=*/r_.AddTensor(100));
r_.AddUse(/*ioperation=*/r_.AddOperation(/*is_stateful=*/false),
/*itensor=*/r_.AddTensor(1));
r_.AddUse(/*ioperation=*/r_.AddOperation(/*is_stateful=*/false),
/*itensor=*/r_.AddTensor(100));
ASSERT_THAT(r_.GetMemProfile(), ElementsAreArray({100, 1, 100}));
EXPECT_THAT(r_.GetPeakMemory(), FieldsAre(Eq(2), Eq(100)));
}
TEST_F(RematTest, RematRecreatesOutput) {
r_.AddUse(r_.AddOperation(/*is_stateful=*/false), r_.AddTensor(100));
r_.AddOperation(/*is_stateful=*/false);
// /* before: */
// %0 = f1()
// f2()
ASSERT_THAT(r_.GetMemProfile(), ElementsAre(100, 0));
EXPECT_THAT(r_.GetMemProfile({/*begin=*/0, /*end=*/1, /*insert=*/2}),
ElementsAre(100, 0, 100));
r_.Remat({/*begin=*/0, /*end=*/1, /*insert=*/2});
// /* after: */
// %0 = f1()
// f2()
// %1 = f1()
EXPECT_THAT(r_.GetMemProfile(), ElementsAre(100, 0, 100));
EXPECT_THAT(r_.AddTensor(0), 2);
}
TEST_F(RematTest, RematExtendsInputAndRecreatesOutput) {
r_.AddUse(/*ioperation=*/r_.AddOperation(/*is_stateful=*/false),
/*itensor=*/r_.AddTensor(1));
r_.AddUse(/*ioperation=*/r_.AddOperation(/*is_stateful=*/false),
/*itensor=*/r_.AddTensor(100));
r_.AddUse(/*ioperation=*/1, 0);
r_.AddOperation(/*is_stateful=*/false);
r_.AddOperation(/*is_stateful=*/false);
// /* before: */
// %0 = f1()
// %1 = f2(%0)
// f3()
// f4()
ASSERT_THAT(r_.GetMemProfile(), ElementsAre(1, 101, 0, 0));
EXPECT_THAT(r_.GetMemProfile({/*begin=*/1, /*end=*/2, /*insert=*/3}),
ElementsAre(1, 101, 1, 101, 0));
r_.Remat({/*begin=*/1, /*end=*/2, /*insert=*/3});
// /* after: */
// %0 = f1()
// %1 = f2(%0)
// f3() /* %0 is kept alive */
// %2 = f2(%0)
// f4()
EXPECT_THAT(r_.GetMemProfile(), ElementsAre(1, 101, 1, 101, 0));
EXPECT_THAT(r_.AddTensor(0), 3);
}
TEST_F(RematTest, BlockRematDuplicatesIntraBlockValues) {
r_.AddUse(/*ioperation=*/r_.AddOperation(/*is_stateful=*/false),
/*itensor=*/r_.AddTensor(1));
r_.AddUse(/*ioperation=*/r_.AddOperation(/*is_stateful=*/false),
/*itensor=*/r_.AddTensor(10));
r_.AddUse(/*ioperation=*/r_.AddOperation(/*is_stateful=*/false),
/*itensor=*/r_.AddTensor(100));
r_.AddUse(/*ioperation=*/r_.AddOperation(/*is_stateful=*/false),
/*itensor=*/r_.AddTensor(1000));
r_.AddOperation(/*is_stateful=*/false);
r_.AddUse(/*ioperation=*/1, /*itensor=*/0);
r_.AddUse(/*ioperation=*/2, /*itensor=*/0);
r_.AddUse(/*ioperation=*/2, /*itensor=*/1);
r_.AddUse(/*ioperation=*/3, /*itensor=*/0);
r_.AddUse(/*ioperation=*/3, /*itensor=*/1);
r_.AddUse(/*ioperation=*/3, /*itensor=*/2);
// /* before */
// %0 = f1()
// %1 = f2(%0)
// %2 = f3(%0,%1)
// %3 = f4(%0,%1,%2)
// f5()
ASSERT_THAT(r_.GetMemProfile(), ElementsAre(1, 11, 111, 1111, 0));
EXPECT_THAT(r_.GetMemProfile({/*begin=*/1, /*end=*/4, /*insert=*/5}),
ElementsAre(1, 11, 111, 1111, 1, 11, 111, 1111));
r_.Remat({/*begin=*/1, /*end=*/4, /*insert=*/5});
EXPECT_THAT(r_.GetMemProfile(),
ElementsAre(1, 11, 111, 1111, 1, 11, 111, 1111));
EXPECT_THAT(r_.AddTensor(0), 7);
// /* after */
// %0 = f1()
// %1 = f2(%0)
// %2 = f3(%0,%1)
// %3 = f4(%0,%1,%2)
// f5()
// %4 = f2(%0)
// %5 = f3(%0,%4)
// %6 = f4(%0,%4,%5)
}
class RematSimulationTest : public testing::Test {
protected:
class RandomRemat : public Rematerializer {
public:
using Rematerializer::Remat; // For testing.
RandomRemat(const int num_operations, const int num_tensors,
const int num_uses, std::mt19937& rng) {
std::uniform_int_distribution<int> some_size_log(0, 16);
std::uniform_int_distribution<int> some_tensor(0, num_tensors - 1);
std::uniform_int_distribution<int> some_operation(0, num_operations - 1);
for (int i = 0; i < num_tensors; ++i) {
AddTensor(SizeT{1} << some_size_log(rng));
}
for (int i = 0; i < num_operations; ++i) {
AddOperation(/*is_stateful=*/false);
}
for (int i = 0; i < num_uses; ++i) {
AddUse(some_operation(rng), some_tensor(rng));
}
}
};
};
TEST_F(RematSimulationTest, SimulationAgreesWithReality) {
constexpr int kNumOperations = 128;
constexpr int kNumTensors = 32;
constexpr int kNumUses = kNumOperations * kNumTensors / 4;
std::mt19937 rng;
for (int i = 0; i < 1024; ++i) {
RandomRemat remat(kNumOperations, kNumTensors, kNumUses, rng);
// Worst-case scenario: we might double the length of the computation
// schedule each time we remat, so only a few iterations...
std::array<int, 3> randos;
const auto& [begin, end, insert] = randos;
for (int i = 0, num_operations = kNumOperations; i < 4;
++i, num_operations += end - begin) {
std::uniform_int_distribution<int> some_op(0, num_operations - 1);
for (auto& rando : randos) {
rando = some_op(rng);
}
// We need begin <= end <= insert.
std::sort(randos.begin(), randos.end());
const Rematerializer::RematSpec spec{begin, end, insert};
const auto simulated_profile = remat.GetMemProfile(spec);
remat.Remat(spec);
const auto actual_profile = remat.GetMemProfile();
EXPECT_THAT(simulated_profile, ElementsAreArray(actual_profile));
}
}
}
class GreedyRematTest : public testing::Test {
protected:
// A Rematerializer for a graph whose memory profile is a sequence of
// "rainbows" due to nested ops (similar to what happens in backwards pass
// gradient calculations.)
class RainbowRemat : public Rematerializer {
public:
// `sizes` is a vector of vector of sizes. Each sizes vector of length n
// generates a `rainbow` profile of the SSA form
//
// %1 = f1()
// %2 = f2()
// ...
// %n = fn()
// gn(%n)
// ...
// g2(%2)
// g1(%1)
//
// with the sizes of intermediates %1,...%n given by the sizes entries.
//
// If extra_ops > 0, each assignment above becomes a sequence
// %i.1 = fi.1()
// %i.2 = fi.1(%i.1)
// %i.3 = fi.1(%i.2)
// ...
// %i = f1(%i.m)
// where the 'dotted' intermediates have size `extra_size`. This
// allows for testing the effect of window sizes.
explicit RainbowRemat(const std::vector<std::vector<int>>& sizes,
int extra_ops = 0, SizeT extra_size = 0) {
for (const auto& rainbow : sizes) {
int tensor = 0;
int op = 0;
for (const auto& size : rainbow) {
for (int i = 0; i < extra_ops; ++i) {
op = AddOperation(/*is_stateful=*/false);
if (i != 0) {
AddUse(op, tensor);
}
tensor = AddTensor(extra_size);
AddUse(op, tensor);
}
// using negative sizes to signal forbidden operations.
op = AddOperation(/*is_stateful=*/size < 0);
if (extra_ops > 0) {
AddUse(op, tensor);
}
tensor = AddTensor(std::abs(size));
AddUse(op, tensor);
}
for (int i = 0; i < rainbow.size(); ++i) {
op = AddOperation(/*is_stateful=*/false);
AddUse(op, tensor - i);
}
}
}
};
// Similar to above: A multilayer perceptron. The forward pass is
// f[n](f[n-1](...(f[1]())), with dimensions |f[1]|...|f[n]| handed in the
// constructor. The backward pass calculates the loss gradient for a single
// weight in the first layer.
class MlpRemat : public Rematerializer {
public:
explicit MlpRemat(const std::vector<int>& sizes) {
int forward_tensor = -1;
int backward_tensor = -1;
int op = -1;
// Forward pass:
for (const int size : sizes) {
op = AddOperation(/*is_stateful=*/false);
if (forward_tensor >= 0) AddUse(op, forward_tensor);
forward_tensor = AddTensor(size);
AddUse(op, forward_tensor);
}
// Backward pass: Right-multiply the jacobian from the outside in.
// dLoss/df[n] * df[n]/df[n-1] * ...
// The i-th term g[i] depends on g[i-1] and f[n-i] and has size |f[n-i]|.
for (; forward_tensor >= 0; --forward_tensor) {
op = AddOperation(/*is_stateful=*/false);
AddUse(op, forward_tensor);
if (backward_tensor >= 0) AddUse(op, backward_tensor);
backward_tensor = AddTensor(sizes[forward_tensor]);
AddUse(op, backward_tensor);
}
}
// We will also instrument the actual optimizations performed.
MOCK_METHOD(void, ApplyRemat, (const RematSpec&));
};
};
TEST_F(GreedyRematTest, MlpBasic) {
StrictMock<MlpRemat> remat(std::vector<int>({1, 1, 1}));
// (o)ut, (i)n, (l)ive: 0 1 2 3 4 5 Sum
// %0 = f0() o 1
// %1 = f1(%0) i o 2
// %2 = f2(%1) l i o 3
// %3 = g2( %2) l l i o 4
// %4 = g1(%3,%1) l i i o 4
// %5 = g0(%4,%0) i i o 3
ASSERT_THAT(remat.GetMemProfile(), ElementsAreArray({1, 2, 3, 4, 4, 3}));
// Little we can do here -- remat %0 before %5
EXPECT_CALL(remat, ApplyRemat(FieldsAre(/*begin=*/0,
/*end=*/1,
/*insert=*/5)));
// (o)ut, (i)n, (l)ive: 0 1 2 3 4 5 Sum of live sizes
// %0 = f0() o 1
// %1 = f1(%0) i o 2
// %2 = f2(%1) i o 2
// %3 = g2( %2) l i o 3
// %4 = g1(%3,%1) i i o 3
// %0' = f0() o l 2
// %5 = g0(%4,%0') i i o 3
remat.RunGreedyAlgorithm(/*max_cost=*/-1, /*max_block_length=*/1,
/*min_savings=*/1);
EXPECT_THAT(remat.GetMemProfile(), ElementsAreArray({1, 2, 2, 3, 3, 2, 3}));
}
TEST_F(GreedyRematTest, MlpBinary) {
StrictMock<MlpRemat> remat(std::vector<int>({1, 2, 4, 8}));
// (o)ut, (i)n, (l)ive: 0 1 2 3 4 5 6 7 Sum of live sizes
// %0 = f0() o 1
// %1 = f1(%0) i o 3
// %2 = f2(%1) l i o 7
// %3 = f3(%2) l l i o 15
// %4 = g3( %3) l l l i o 23
// %5 = g2(%4 %2) l l i i o 19
// %6 = g1(%5,%1) l i i o 9
// %7 = g0(%6,%0) i i o 4
ASSERT_THAT(remat.GetMemProfile(),
ElementsAreArray({1, 3, 7, 15, 23, 19, 9, 4}));
EXPECT_CALL(remat, ApplyRemat(FieldsAre(/*begin=*/2,
/*end=*/3,
/*insert=*/5)));
EXPECT_CALL(remat, ApplyRemat(FieldsAre(/*begin=*/0,
/*end=*/1,
/*insert=*/8)));
// (o)ut, (i)n, (l)ive: 0 1 2 3 4 5 6 7 Sum of live sizes
// %0 = f0() o 1
// %1 = f1(%0) i o 3
// %2 = f2(%1) i o 6
// %3 = f3(%2) l i o 14
// %4 = g3( %3) l i o 18
// %2'= f2(%1) l o l 14
// %5 = g2(%4 %2) l i i o 18
// %6 = g1(%5,%1) i i o 8
// %0'= f(0) o l 3
// %7 = g0(%6,%0) i i o 4
remat.RunGreedyAlgorithm(/*max_cost=*/-1, /*max_block_length=*/4,
/*min_savings=*/1);
EXPECT_THAT(remat.GetMemProfile(),
ElementsAreArray({1, 3, 6, 14, 18, 14, 18, 8, 3, 4}));
}
TEST_F(GreedyRematTest, SimpleMax) {
RainbowRemat remat({{1, 2, 4, 8, 16}});
ASSERT_THAT(remat.GetMemProfile(),
ElementsAreArray({1, 3, 7, 15, 31, 31, 15, 7, 3, 1}));
remat.RunGreedyAlgorithm(/*max_cost=*/-1, /*max_block_length=*/1,
/*min_savings=*/1);
// Profile is flattened to its minimum--16.
EXPECT_THAT(remat.GetMemProfile(),
ElementsAreArray({1, 2, 4, 8, 16, 16, 8, 8, 4, 4, 2, 2, 1, 1}));
}
TEST_F(GreedyRematTest, SimpleMaxLongWindow) {
RainbowRemat remat({{1, 2, 4, 8, 16}});
ASSERT_THAT(remat.GetMemProfile(),
ElementsAreArray({1, 3, 7, 15, 31, 31, 15, 7, 3, 1}));
remat.RunGreedyAlgorithm(/*max_cost=*/-1, /*max_block_length=*/4,
/*min_savings=*/1);
// Profile is flattened to its minimum--16.
EXPECT_THAT(remat.GetMemProfile(),
ElementsAreArray({1, 2, 4, 8, 16, 16, 8, 8, 4, 4, 2, 2, 1, 1}));
}
TEST_F(GreedyRematTest, SimpleSizeThreshold) {
RainbowRemat remat({{1, 2, 4, 8, 16}});
ASSERT_THAT(remat.GetMemProfile(),
ElementsAreArray({1, 3, 7, 15, 31, 31, 15, 7, 3, 1}));
// Only do remats of at least 4 bytes of savings -- this will lower the
// profile by 4 + 8 instead of 1 + 2 + 4 + 8 as before.
remat.RunGreedyAlgorithm(/*max_cost=*/-1, /*max_block_length=*/1,
/*min_savings=*/4);
EXPECT_THAT(remat.GetMemProfile(),
ElementsAreArray({1, 3, 7, 11, 19, 19, 11, 11, 7, 7, 3, 1}));
}
TEST_F(GreedyRematTest, SimpleCostThreshold) {
RainbowRemat remat({{1, 2, 4, 8, 16}});
ASSERT_THAT(remat.GetMemProfile(),
ElementsAreArray({1, 3, 7, 15, 31, 31, 15, 7, 3, 1}));
// Do at most 1 remat -- this will lower the profile by 8.
remat.RunGreedyAlgorithm(/*max_cost=*/1, /*max_block_length=*/1,
/*min_savings=*/1);
// Only as single remat is done -- this will be the best one possible,
// lowering the profile by 8 instead of the maximum 1 + 2 + 4 + 8.
EXPECT_THAT(remat.GetMemProfile(),
ElementsAreArray({1, 3, 7, 15, 23, 23, 15, 15, 7, 3, 1}));
}
TEST_F(GreedyRematTest, SimpleForbiddenOps) {
// Operator generating size-4 tensor is stateful, so it won't be materialized.
RainbowRemat remat({{1, 2, -4, 8, 16}});
ASSERT_THAT(remat.GetMemProfile(),
ElementsAreArray({1, 3, 7, 15, 31, 31, 15, 7, 3, 1}));
// Best we can do is lower the profile to 8 + 2 + 1.
remat.RunGreedyAlgorithm(/*max_cost=*/-1, /*max_block_length=*/1,
/*min_savings=*/1);
EXPECT_THAT(remat.GetMemProfile(),
ElementsAreArray({1, 2, 4, 12, 20, 20, 12, 12, 4, 2, 2, 1, 1}));
}
TEST_F(GreedyRematTest, DoubleMax) {
// Generate a profile with two local maxima of size 31 and 28.
RainbowRemat remat({{1, 2, 4, 8, 16}, {4, 8, 16}});
ASSERT_THAT(remat.GetMemProfile(),
ElementsAreArray(
{1, 3, 7, 15, 31, 31, 15, 7, 3, 1, 4, 12, 28, 28, 12, 4}));
remat.RunGreedyAlgorithm(/*max_cost=*/-1, /*max_block_length=*/1,
/*min_savings=*/1);
// Two rainbows; the first is lowered by 1 + 2 + 4 + 8 == 15, the second by 4
// + 8 == 12.
EXPECT_THAT(remat.GetMemProfile(),
ElementsAreArray({1, 2, 4, 8, 16, 16, 8, 8, 4, 4, 2,
2, 1, 1, 4, 8, 16, 16, 8, 8, 4, 4}));
}
TEST_F(GreedyRematTest, DoubleCostThreshold) {
// Generate a profile with two local maxima of size 31 and 28.
RainbowRemat remat({{1, 2, 4, 8, 16}, {4, 8, 16}});
ASSERT_THAT(remat.GetMemProfile(),
ElementsAreArray(
{1, 3, 7, 15, 31, 31, 15, 7, 3, 1, 4, 12, 28, 28, 12, 4}));
remat.RunGreedyAlgorithm(/*max_cost=*/2, /*max_block_length=*/1,
/*min_savings=*/1);
// Profile can be flattened twice--first, the global maximum of 31 is reduced
// by 8; then the newly-global maximum of 28 is reduced by 8.
EXPECT_THAT(remat.GetMemProfile(),
ElementsAreArray({1, 3, 7, 15, 23, 23, 15, 15, 7, 3, 1, 4, 12, 20,
20, 12, 12, 4}));
}
TEST_F(GreedyRematTest, SingleLongerBlocksByWindowSize) {
std::vector<Rematerializer::SizeT> best_for_window_size;
for (int window_size : {0, 1, 2, 3, 4, 5}) {
RainbowRemat remat({{1, 2, 4, 8}}, 2, 16);
remat.RunGreedyAlgorithm(/*max_cost=*/-1, /*max_block_length=*/window_size,
/*min_savings=*/1);
best_for_window_size.push_back(remat.GetPeakMemory().size);
}
EXPECT_THAT(best_for_window_size, ElementsAreArray({44, 36, 36, 32, 32, 32}));
}
} // namespace
} // namespace TFL
} // namespace mlir
tensorflow/compiler/mlir/lite/experimental/tac/BUILD
+399
View File
@@ -0,0 +1,399 @@
load("@flatbuffers//:build_defs.bzl", "flatbuffer_cc_library")
load(
"@llvm-project//mlir:tblgen.bzl",
"gentbl_cc_library",
)
load("@xla//third_party/rules_python/python:py_library.bzl", "py_library")
load("//tensorflow:tensorflow.bzl", "if_google", "tf_cc_binary", "tf_cc_test")
load("//tensorflow:tensorflow.default.bzl", "get_compatible_with_portable")
load(
"//tensorflow/core/platform:build_config.bzl",
"tf_proto_library",
)
load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:LICENSE"],
default_visibility = [
"//visibility:public",
],
licenses = ["notice"],
)
flatbuffer_cc_library(
name = "runtime_metadata_fbs",
srcs = ["runtime_metadata.fbs"],
compatible_with = get_compatible_with_portable(),
)
cc_library(
name = "execution_metadata_exporter",
srcs = [
"execution_metadata_exporter.cc",
],
hdrs = [
"execution_metadata_exporter.h",
],
deps = [
":common",
":runtime_metadata_fbs",
"//tensorflow/compiler/mlir/lite:tensorflow_lite",
"//tensorflow/compiler/mlir/lite/experimental/tac/hardwares:target_hardware",
"//tensorflow/compiler/mlir/lite/quantization/ir:QuantOps",
"//tensorflow/compiler/mlir/tensorflow",
"@flatbuffers",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:ArithDialect",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Support",
],
)
tf_cc_test(
name = "execution_metadata_exporter_test",
srcs = [
"execution_metadata_exporter_test.cc",
],
deps = [
":execution_metadata_exporter",
":runtime_metadata_fbs",
"//tensorflow/compiler/mlir/lite:tensorflow_lite",
"//tensorflow/compiler/mlir/lite/experimental/tac/hardwares:all-target-hardwares",
"@com_google_googletest//:gtest_main",
"@flatbuffers",
"@llvm-project//mlir:ArithDialect",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Parser",
],
)
cc_library(
name = "common",
srcs = [
"common/utils.cc",
],
hdrs = [
"common/cost.h",
"common/subgraph.h",
"common/targets.h",
"common/utils.h",
],
deps = [
"//tensorflow/compiler/mlir/lite:tensorflow_lite",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:ArithDialect",
"@llvm-project//mlir:BytecodeOpInterface",
"@llvm-project//mlir:CallOpInterfaces",
"@llvm-project//mlir:CastInterfaces",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:QuantOps",
"@llvm-project//mlir:SideEffectInterfaces",
"@llvm-project//mlir:Support",
],
)
gentbl_cc_library(
name = "transform_patterns_inc_gen",
compatible_with = get_compatible_with_portable(),
tbl_outs = {"transforms/generated_transform_patterns.inc": ["-gen-rewriters"]},
tblgen = "@llvm-project//mlir:mlir-tblgen",
td_file = "transforms/transform_patterns.td",
deps = [
"//tensorflow/compiler/mlir/lite:tensorflow_lite_ops_td_files",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_ops_td_files",
"@llvm-project//mlir:ArithOpsTdFiles",
"@llvm-project//mlir:FuncTdFiles",
],
)
cc_library(
name = "device_transform_patterns",
srcs = [
"transforms/device_transform_patterns.cc",
],
hdrs = [
"transforms/device_transform_patterns.h",
],
textual_hdrs = [
"transforms/generated_transform_patterns.inc",
],
deps = [
":common",
"//tensorflow/compiler/mlir/lite:tensorflow_lite",
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/compiler/mlir/tensorflow:verification_utils",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Pass",
"@llvm-project//mlir:QuantOps",
"@llvm-project//mlir:Support",
],
)
cc_library(
name = "device_transform_gpu",
srcs = [
"transforms/device_transform_gpu.cc",
],
hdrs = [
"transforms/device_transform_gpu.h",
],
deps = [
":common",
"//tensorflow/compiler/mlir/lite:tensorflow_lite",
"//tensorflow/compiler/mlir/lite/experimental/tac/hardwares:gpu_hardware",
"//tensorflow/compiler/mlir/tensorflow",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Pass",
"@llvm-project//mlir:Support",
"@llvm-project//mlir:TransformUtils",
],
alwayslink = 1,
)
cc_library(
name = "device_transform_nnapi",
srcs = [
"transforms/device_transform_nnapi.cc",
],
deps = [
":common",
"//tensorflow/compiler/mlir/lite:tensorflow_lite",
"//tensorflow/compiler/mlir/lite/experimental/tac/hardwares:nnapi_hardware",
"//tensorflow/compiler/mlir/tensorflow",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Pass",
"@llvm-project//mlir:Support",
"@llvm-project//mlir:TransformUtils",
],
alwayslink = 1,
)
cc_library(
name = "device_transform",
srcs = [
"transforms/device_transform.cc",
],
hdrs = [
"transforms/device_transform.h",
],
deps = [
":common",
":device_transform_gpu",
":device_transform_nnapi",
":device_transform_patterns",
"//tensorflow/compiler/mlir/lite:tensorflow_lite",
"//tensorflow/compiler/mlir/lite/experimental/tac/hardwares:target_hardware",
"//tensorflow/compiler/mlir/tensorflow",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:QuantOps",
"@llvm-project//mlir:Support",
"@llvm-project//mlir:TransformUtils",
"@llvm-project//mlir:Transforms",
],
)
cc_library(
name = "cost_model",
srcs = [
"transforms/cost_model.cc",
],
hdrs = [
"transforms/cost_model.h",
],
deps = [
":common",
"//tensorflow/compiler/mlir/lite:tensorflow_lite",
"//tensorflow/compiler/mlir/lite/experimental/tac/hardwares:target_hardware",
"@com_google_absl//absl/strings",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:CallOpInterfaces",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Pass",
"@llvm-project//mlir:Support",
],
alwayslink = 1,
)
# TODO(b/177376459): split tac_module and passes dependency to separate libraries.
cc_library(
name = "target_aware_conversion",
srcs = [
"tac_module.cc",
"transforms/compute_cost.cc",
"transforms/fold_constants_to_subgraph.cc",
"transforms/get_alternative_subgraph.cc",
"transforms/pick_subgraphs.cc",
"transforms/raise_target_subgraphs.cc",
"transforms/tac_filter.cc",
"transforms/target_annotation.cc",
],
hdrs = [
"tac_module.h",
"transforms/passes.h",
"transforms/tac_pass.h",
],
deps = [
":common",
":cost_model",
":device_transform",
":tac_filter_proto_cc",
":tac_importer_exporter",
"//tensorflow/compiler/mlir/lite:tensorflow_lite",
"//tensorflow/compiler/mlir/lite:tf_tfl_passes",
"//tensorflow/compiler/mlir/lite/experimental/common:outline_operations",
"//tensorflow/compiler/mlir/lite/experimental/tac/hardwares:target_hardware",
"//tensorflow/compiler/mlir/tensorflow:cluster_util",
"//tensorflow/compiler/mlir/tensorflow:error_util",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_analysis",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings",
"@com_google_protobuf//:protobuf_headers",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:ArithDialect",
"@llvm-project//mlir:CallOpInterfaces",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Pass",
"@llvm-project//mlir:SideEffectInterfaces",
"@llvm-project//mlir:Support",
"@llvm-project//mlir:TransformUtils",
"@llvm-project//mlir:Transforms",
"@stablehlo//:stablehlo_ops",
],
alwayslink = 1,
)
cc_library(
name = "tac-opt_lib",
testonly = True,
deps = [
":target_aware_conversion",
"//tensorflow/compiler/mlir/lite:litert_mlir_opt_main",
],
alwayslink = 1,
)
# Binary with no hardwares linked.
tf_cc_binary(
name = "tac-opt",
testonly = True,
deps = [
":tac-opt_lib",
],
)
# Binary with all backends linked.
tf_cc_binary(
name = "tac-opt-all-backends",
testonly = True,
deps = [
":tac-opt_lib",
"//tensorflow/compiler/mlir/lite/experimental/common:outline_operations",
"//tensorflow/compiler/mlir/lite/experimental/tac/hardwares:all-target-hardwares",
],
)
cc_library(
name = "tac-translate-lib",
srcs = [
"tac_translate.cc",
],
deps = [
":common",
":execution_metadata_exporter",
":target_aware_conversion",
":tflite_importer_exporter",
"//tensorflow/compiler/mlir:init_mlir",
"//tensorflow/compiler/mlir/lite:tensorflow_lite_legalize_tf",
"//tensorflow/compiler/mlir/lite:tensorflow_lite_optimize",
"//tensorflow/compiler/mlir/lite/experimental/tac/hardwares:target_hardware",
"//tensorflow/compiler/mlir/lite/experimental/tac/utils",
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/compiler/mlir/tensorflow:error_util",
"//tensorflow/compiler/mlir/tensorflow/transforms:tensorflow_passes",
"//tensorflow/core:lib",
"@com_google_absl//absl/log",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/strings:str_format",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Parser",
"@llvm-project//mlir:Pass",
"@llvm-project//mlir:Support",
"@llvm-project//mlir:Transforms",
"@llvm-project//mlir:TranslateLib",
],
)
tf_cc_binary(
name = "tac-translate",
deps = [
":tac-translate-lib",
"//tensorflow/compiler/mlir/lite/experimental/tac/hardwares:all-target-hardwares",
],
)
cc_library(
name = "tac_importer_exporter",
hdrs = ["tac_importer_exporter.h"],
deps = [
"@com_google_absl//absl/status:statusor",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
],
)
cc_library(
name = "tflite_importer_exporter",
srcs = ["tflite_import_export.cc"],
hdrs = ["tflite_import_export.h"],
deps = [
":common",
":execution_metadata_exporter",
":tac_importer_exporter",
"//tensorflow/compiler/mlir/lite/experimental/tac/hardwares:target_hardware",
"//tensorflow/compiler/mlir/lite/experimental/tac/utils",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:Support",
],
)
exports_files([
"run_lit.sh",
])
py_library(
name = "tac",
srcs = [
"tac.py",
],
strict_deps = True,
visibility = ["//visibility:public"],
deps = [
"//tensorflow/compiler/mlir/lite/experimental/tac/py_wrapper:_pywrap_tac_wrapper",
],
)
tf_proto_library(
name = "tac_filter_proto",
srcs = ["tac_filter.proto"],
protodeps = if_google(["//google/protobuf:any"]),
visibility = ["//visibility:public"],
)
tensorflow/compiler/mlir/lite/experimental/tac/README.md
+271
View File
@@ -0,0 +1,271 @@
# Target Aware Conversion (TAC)
Different hardwares have different capabilities and restrictions.
TAC is designed to leverage hardwares' capabilities to:
* Perform device-specific optimizations (such as unsupported ops lowering,
layout transformations, etc.)
* Graph partitioning based on the hardware costs modeling.
* It supports general import/export where you can hook your own
importer/exporter from any format to MLIR and export MLIR to anything.
For more details, please checkout the
[TAC workflow](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/compiler/mlir/lite/experimental/tac/README.md#tac-workflow)
section
## How to use
Once you have a converted TfLite model ready, you can use the following command
to use TAC to optimize for your model:
```
bazel run -c opt //tensorflow/compiler/mlir/lite/experimental/tac:tac-translate -- <PATH_TO_YOUR_MODEL> -o=<OUTPUT_PATH> -device-specs=<HARDWARE_BACKENDS>
```
The devices_specs is a list of the names of the desired hardware backends,
separated by comma, e.g., "GPU,CPU".
If you're interested in what are the subgraphs being explored for different
backends, you can pass in `-output-mlir -inline-subgraphs=false` and check out
the output mlir file.
## How to add a hardware backend
If you want to add a hardware backend for TAC, you can start with the
`SimpleHardware` interface.
For example:
```
class FooHardware : public SimpleHardware {
public:
static constexpr char kId[] = "FOO";
mlir::RewritePatternSet GetTransformations(
MLIRContext* context) const override {
mlir::RewritePatternSet patterns;
// Pick the transformations that we want to perform,
// We can add other transformations we like here.
patterns.add<LowerPackIntoConcatReshape, UnrollSplit, UnrollSplitV,
PadSlice>(context);
return patterns;
}
mlir::TypeID GetTypeId() const override {
return mlir::TypeID::get<FooHardware>();
}
// We can specify what ops are not supported here.
bool IsNotSupportedOp(mlir::Operation* op) const override { return false; }
// This is basically saying how fast are we comparing to CPU.
// The larger the value the better.
float AdvantageOverCPU() const override { return 5.0; }
};
```
Then we need to register our hardware like below:
```
std::unique_ptr<TargetHardware> CreateFooHardware() {
return std::make_unique<FooHardware>();
}
TargetHardwareRegistration<FooHardware> foo_hardware(
"Target device for FOO", CreateFooHardware);
```
### Advanced user
For advanced users (e.g., you may already have your own hardware dialect
defined), please just use `TargetHardware` directly. See the following code
snippet for reference.
```
class MyCustomHardware : public TargetHardware {
public:
static constexpr char kId[] = "MY_CUSTOM_HARDWARE";
mlir::TypeID GetTypeId() const override {
return mlir::TypeID::get<MyCustomHardware>();
}
bool IsOpSupported(mlir::Operation* op) const override {
// check whether the op is supported, if the user has they own dialect,
// this can be target dialect legalization process.
}
double GetHardwareSwitchingCost(const TargetHardware* from,
size_t buffer_size) const override {
// Get the hardware switching cost from the source hardware.
}
double GetOpCost(mlir::Operation* op) const override {
// call customized cost model.
}
mlir::RewritePatternSet GetTransformations(
MLIRContext* context) const override {
// customized transformations patterns: ops lowering/fusion, layout
// transformation, etc.
}
};
```
## TAC workflow
The workflow of target-aware-conversion is as followed:
1 Try to break down the whole graph into several subgraphs based on hardwares'
capabilities. See the diagram below, let's say our desired target backends are
"GPU" and "CPU", and currently "C" is not supported on "GPU", but the rest are
supported by "GPU". So we will end up with 3 subgraphs as shown in the diagram.
![Target Annotation](g3doc/images/target_annotation.png)
2 Perform ops-lowering & target-specific optimizations for
different hardware backends. As shown in the below diagram, the red & the
yellow subgraph will be duplicated as "alternative subgraph view" for "CPU".
"C" op can be lowered into "G" + "H" op which can be supported by "GPU".
![Target Optimization](g3doc/images/target_optimization.png)
3 Estimate the costs for each subgraph (and their alternative views)
based on the hardware cost model. See the following diagram.
![Estimate costs](g3doc/images/compute_cost.png)
4 Pick the proper subgraphs from the alternative views for execution based on
costs(computation costs, transfer costs, quant/dequant costs). As shown in the
diagram below, since cross-device data transferring cost is high, even "G" + "H"
running on GPU maybe less efficient than "C" running on "CPU", we will still
pick "G" + "H" subgraph.
![Pick subgraphs](g3doc/images/pick_subgraphs.png)
The final graph looks like below:
![Final graph](g3doc/images/final_graph.png)
## TAC components
### Hardwares
Hardwares are used to modeling target device capabilities & also ops cost for
the target devices.
We have already modeled `cpu_hardware` & `gpu_hardware` as well as the
`nnapi_hardware`.
### Passes
#### Target Annotation Pass
In this pass, every op will be targeted with the user specified targets based on
the device capabilites. For example, If the user specified the desired targets
are "GPU", "CPU", `conv2d` can run on both "GPU" and "CPU", we will annotate
the op `conv2d` with "GPU" since it's preferred; `pack` can only run on "CPU",
so we will annotate the op with "CPU" since "GPU" does not support this op.
#### Raise Target Subgraphs Pass
In this pass, ops will be broken down into subgraph. Those ops have the same
target annotation will be raised as subgraphs.
In this pass, subgraph is actually implemented with `FuncOp`.
Take the following code as an example:
```
func @simpleTest(%arg0: tensor<1xf32>, %arg1: tensor<1xf32>, %arg2: tensor<1xf32>, %arg3: tensor<1xf32>) -> tensor<2x1xf32> {
%0 = "tfl.add"(%arg0, %arg1) {tac.device = "GPU", fused_activation_function = "RELU6", tac.inference_type = "FLOAT"} : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32>
%1 = "tfl.mul"(%0, %arg2) {tac.device = "GPU", fused_activation_function = "RELU6", tac.inference_type = "FLOAT"} : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32>
%2 = "tfl.add"(%arg0, %arg3) {tac.device = "GPU", fused_activation_function = "RELU6", tac.inference_type = "FLOAT"} : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32>
%3 = "tfl.pack"(%1, %2) {tac.device = "CPU", tac.inference_type = "FLOAT", axis = 0 : i32, values_count = 2 : i32} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2x1xf32>
return %3 : tensor<2x1xf32>
}
```
In this code, `%3` is annotated with "CPU", while others are annotated with
"GPU", in this case, `%3` will be raised as a separate function like below:
```
func private @func_1_GPU_FLOAT(%arg0: tensor<1xf32>, %arg1: tensor<1xf32>) -> tensor<1xf32> attributes {tac.device = "GPU", tac.inference_type = "FLOAT", tac.interface_name = "func_1"} {
%0 = tfl.add %arg0, %arg1 {fused_activation_function = "RELU6", tac.device = "GPU", tac.inference_type = "FLOAT"} : tensor<1xf32>
return %0 : tensor<1xf32>
}
```
And the rest ops will be raised as below:
```
func private @func_2_CPU_FLOAT(%arg0: tensor<1xf32>, %arg1: tensor<1xf32>) -> tensor<2x1xf32> attributes {tac.device = "CPU", tac.inference_type = "FLOAT", tac.interface_name = "func_2"} {
%0 = "tfl.pack"(%arg0, %arg1) {axis = 0 : i32, tac.device = "CPU", tac.inference_type = "FLOAT", values_count = 2 : i32} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2x1xf32>
return %0 : tensor<2x1xf32>
}
func private @func_0_GPU_FLOAT(%arg0: tensor<1xf32>, %arg1: tensor<1xf32>, %arg2: tensor<1xf32>) -> tensor<1xf32> attributes {tac.device = "GPU", tac.inference_type = "FLOAT", tac.interface_name = "func_0"} {
%0 = tfl.add %arg0, %arg1 {fused_activation_function = "RELU6", tac.device = "GPU", tac.inference_type = "FLOAT"} : tensor<1xf32>
%1 = tfl.mul %0, %arg2 {fused_activation_function = "RELU6", tac.device = "GPU", tac.inference_type = "FLOAT"} : tensor<1xf32>
return %1 : tensor<1xf32>
}
```
And the original function will be replaced by `CallOps` to those `FuncOps`:
```
func @simpleTest(%arg0: tensor<1xf32>, %arg1: tensor<1xf32>, %arg2: tensor<1xf32>, %arg3: tensor<1xf32>) -> tensor<2x1xf32> {
%0 = call @func_0_GPU_FLOAT(%arg0, %arg1, %arg2) {tac.device = "GPU", tac.inference_type = "FLOAT", tac.interface_name = "func_0"} : (tensor<1xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32>
%1 = call @func_1_GPU_FLOAT(%arg0, %arg3) {tac.device = "GPU", tac.inference_type = "FLOAT", tac.interface_name = "func_1"} : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32>
%2 = call @func_2_CPU_FLOAT(%0, %1) {tac.device = "CPU", tac.inference_type = "FLOAT", tac.interface_name = "func_2"} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2x1xf32>
return %2 : tensor<2x1xf32>
}
```
Why we need to raise those ops into `FuncOps`? Please see the following section.
#### Get Alternative Subgraph View Pass
In the Get Alternative Subgraph View Pass, we will essentially duplicate those
`FuncOps` and perform unsupported ops lowering & target-specific optimization.
For example, `Pack` is not supported by "GPU", but it can be lowered into
`Concat` + `Reshape` which can be supported by "GPU".
So the original example:
```
func private @func_1_GPU_FLOAT(%arg0: tensor<1xf32>, %arg1: tensor<1xf32>) -> tensor<1xf32> attributes {tac.device = "GPU", tac.inference_type = "FLOAT", tac.interface_name = "func_1"} {
%0 = tfl.add %arg0, %arg1 {fused_activation_function = "RELU6", tac.device = "GPU", tac.inference_type = "FLOAT"} : tensor<1xf32>
return %0 : tensor<1xf32>
}
```
Will be transformed into:
```
func private @func_2_CPU_FLOAT(%arg0: tensor<1xf32>, %arg1: tensor<1xf32>) -> tensor<2x1xf32> attributes {tac.device = "CPU", tac.inference_type = "FLOAT", tac.interface_name = "func_2"} {
%0 = "tfl.pack"(%arg0, %arg1) {axis = 0 : i32, tac.device = "CPU", tac.inference_type = "FLOAT", values_count = 2 : i32} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2x1xf32>
return %0 : tensor<2x1xf32>
}
func private @func_2_GPU_FLOAT(%arg0: tensor<1xf32>, %arg1: tensor<1xf32>) -> tensor<2x1xf32> attributes {tac.device = "GPU", tac.inference_type = "FLOAT", tac.interface_name = "func_2"} {
%cst = arith.constant dense<1> : tensor<4xi32>
%cst_0 = arith.constant dense<2> : tensor<1xi32>
%cst_1 = arith.constant dense<[2, 1]> : tensor<2xi32>
%0 = "tfl.reshape"(%arg0, %cst) {tac.device = "GPU", tac.inference_type = "FLOAT"} : (tensor<1xf32>, tensor<4xi32>) -> tensor<1x1x1x1xf32>
%1 = "tfl.reshape"(%arg1, %cst) {tac.device = "GPU", tac.inference_type = "FLOAT"} : (tensor<1xf32>, tensor<4xi32>) -> tensor<1x1x1x1xf32>
%2 = "tfl.concatenation"(%0, %1) {axis = 3 : i32, fused_activation_function = "NONE", tac.device = "GPU", tac.inference_type = "FLOAT"} : (tensor<1x1x1x1xf32>, tensor<1x1x1x1xf32>) -> tensor<1x1x1x2xf32>
%3 = "tfl.reshape"(%2, %cst_0) {tac.device = "GPU", tac.inference_type = "FLOAT"} : (tensor<1x1x1x2xf32>, tensor<1xi32>) -> tensor<2xf32>
%4 = "tfl.reshape"(%3, %cst_1) {tac.device = "GPU", tac.inference_type = "FLOAT"} : (tensor<2xf32>, tensor<2xi32>) -> tensor<2x1xf32>
return %4 : tensor<2x1xf32>
}
```
#### Compute Costs Pass
In the compute cost pass, we will essentially compute the cost of each op within
the `FuncOp` based on the target-device cost model and sum them together.
#### Pick Subgraphs Pass
In the pick subgraphs pass, we will pick those subgraphs which can minimize the
global costs (we will take the tensor transferring costs as well).
tensorflow/compiler/mlir/lite/experimental/tac/common/cost.h
+50
View File
@@ -0,0 +1,50 @@
/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_TAC_COMMON_COST_H_
#define TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_TAC_COMMON_COST_H_
#include <string>
#include "mlir/IR/Builders.h" // from @llvm-project
#include "mlir/IR/Operation.h" // from @llvm-project
namespace mlir {
namespace TFL {
namespace tac {
// Cost attribute string on the TFL dialect.
constexpr char kCost[] = "tac.cost";
inline void UpdateCost(Operation* op, float cost, OpBuilder* builder) {
op->setAttr(kCost, builder->getF32FloatAttr(cost));
}
// Get the cost annotated with kCost.
inline bool GetCostOnOp(Operation* op, float* cost) {
auto cost_type = op->getAttrOfType<FloatAttr>(kCost);
if (cost_type == nullptr) {
return false;
}
*cost = cost_type.getValueAsDouble();
return true;
}
} // namespace tac
} // namespace TFL
} // namespace mlir
#endif // TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_TAC_COMMON_COST_H_
tensorflow/compiler/mlir/lite/experimental/tac/common/subgraph.h
+52
View File
@@ -0,0 +1,52 @@
/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_TAC_COMMON_SUBGRAPH_H_
#define TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_TAC_COMMON_SUBGRAPH_H_
#include <optional>
#include <string>
#include "llvm/ADT/StringRef.h"
#include "mlir/IR/Operation.h" // from @llvm-project
namespace mlir {
namespace TFL {
namespace tac {
// Interface name here is the "hook" between the CallOp and FuncOps.
// Take the following example:
//
// call @func_1_CPU {tac.interface_name = "func_1"}
//
// "func_1" is the interface name where "func_1_cpu" is the real implementation
// we can have multiple FuncOps like "func_1_cpu" and "func_1_gpu" and they
// both implement "func_1".
//
// The attribute on the FuncOp means what it actually implements while the
// attribute on the CallOp means what it actually looks for.
constexpr char kInterfaceNameAttr[] = "tac.interface_name";
inline std::optional<std::string> GetInterFaceName(Operation* op) {
auto name_attr = op->getAttrOfType<StringAttr>(kInterfaceNameAttr);
if (!name_attr) return std::nullopt;
return name_attr.getValue().str();
}
} // namespace tac
} // namespace TFL
} // namespace mlir
#endif // TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_TAC_COMMON_SUBGRAPH_H_
tensorflow/compiler/mlir/lite/experimental/tac/common/targets.h
+158
View File
@@ -0,0 +1,158 @@
/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_TAC_COMMON_TARGETS_H_
#define TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_TAC_COMMON_TARGETS_H_
#include <algorithm>
#include <cctype>
#include <cstddef>
#include <functional>
#include <optional>
#include <string>
#include <vector>
#include "llvm/ADT/ArrayRef.h"
#include "llvm/ADT/StringRef.h"
#include "mlir/IR/BuiltinAttributes.h" // from @llvm-project
#include "mlir/IR/Operation.h" // from @llvm-project
namespace mlir {
namespace TFL {
namespace tac {
// Device attribute string on the TFL dialect.
constexpr char kDevice[] = "tac.device";
// Inference type.
constexpr char kInferenceType[] = "tac.inference_type";
// Inference type.
constexpr char kSkipTargetAnnotation[] = "tac.skip_target_annotation";
// Custom options fingerprint to apply different options for different filters.
constexpr char kCustomOptionsFingerprint[] = "tac.custom_options_fingerprint";
// Delegate compiler stack version.
constexpr char kDelegateCompilerVersion[] = "tac.delegate_compiler_version";
// TODO(renjieliu): Add more inference types.
enum InferenceType {
UNKNOWN = 0,
FLOAT = 1,
QUANTIZED_INT8 = 2,
QUANTIZED_UINT8 = 3,
HYBRID = 4
};
inline InferenceType GetInferenceTypeEnum(llvm::StringRef inference_type_str) {
if (inference_type_str == "FLOAT") {
return FLOAT;
} else if (inference_type_str == "QUANTIZED_INT8") {
return QUANTIZED_INT8;
} else if (inference_type_str == "QUANTIZED_UINT8") {
return QUANTIZED_UINT8;
} else if (inference_type_str == "HYBRID") {
return HYBRID;
} else {
return UNKNOWN;
}
}
inline std::string GetInferenceString(InferenceType inference_type) {
if (inference_type == FLOAT) {
return "FLOAT";
} else if (inference_type == QUANTIZED_INT8) {
return "QUANTIZED_INT8";
} else if (inference_type == QUANTIZED_UINT8) {
return "QUANTIZED_UINT8";
} else if (inference_type == HYBRID) {
return "HYBRID";
} else {
return "UNKNOWN";
}
}
// Returns canonical representation for hardware name (All uppercase).
// TODO(b/177376459): Remove this in favor of the string defined by hardwares
// MyHardware::kId.
inline std::string GetCanonicalHardwareName(const std::string& hardware_name) {
std::string name = hardware_name;
std::transform(
name.begin(), name.end(), name.begin(),
[](unsigned char c) -> unsigned char { return std::toupper(c); });
return name;
}
// Get the target annotation form the op.
inline std::optional<std::string> GetTargetAnnotation(Operation* op) {
auto device = op->getAttrOfType<StringAttr>(kDevice);
if (device == nullptr || device.getValue().empty()) return std::nullopt;
return GetCanonicalHardwareName(device.getValue().str());
}
// Get inference type attribute from the operation if available.
inline std::optional<InferenceType> GetInferenceTypeAnnotation(Operation* op) {
auto inference_type = op->getAttrOfType<StringAttr>(kInferenceType);
if (inference_type == nullptr) return std::nullopt;
llvm::StringRef device_name_str = inference_type.getValue();
return GetInferenceTypeEnum(device_name_str);
}
// InferenceDeviceType is a combination of the hardware with inference type.
struct InferenceDeviceType {
std::string hardware;
InferenceType inference_type;
bool operator==(const InferenceDeviceType& other) const {
return (hardware == other.hardware) &&
(inference_type == other.inference_type);
}
bool operator!=(const InferenceDeviceType& other) const {
return !(*this == other);
}
struct inference_device_type_hash {
size_t operator()(const InferenceDeviceType& p) const {
auto hash1 = std::hash<std::string>{}(p.hardware);
auto hash2 = std::hash<InferenceType>{}(p.inference_type);
return hash1 ^ hash2;
}
};
};
// Get InferenceDeviceType attribute from the operation if available.
inline std::optional<InferenceDeviceType> GetInferenceDeviceTypeForOp(
Operation* op) {
auto hardware = GetTargetAnnotation(op);
if (!hardware.has_value()) return std::nullopt;
auto inference_type = GetInferenceTypeAnnotation(op);
if (!inference_type.has_value()) return std::nullopt;
InferenceDeviceType inference_device_type;
inference_device_type.hardware = hardware.value();
inference_device_type.inference_type = inference_type.value();
return inference_device_type;
}
} // namespace tac
} // namespace TFL
} // namespace mlir
#endif // TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_TAC_COMMON_TARGETS_H_
tensorflow/compiler/mlir/lite/experimental/tac/common/utils.cc
+80
View File
@@ -0,0 +1,80 @@
/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/compiler/mlir/lite/experimental/tac/common/utils.h"
#include "llvm/Support/Casting.h"
#include "mlir/IR/OpDefinition.h" // from @llvm-project
#include "tensorflow/compiler/mlir/lite/experimental/tac/common/targets.h"
#include "tensorflow/compiler/mlir/lite/ir/tfl_ops.h"
#include "tensorflow/compiler/mlir/lite/utils/utils.h"
namespace mlir {
namespace TFL {
namespace tac {
bool NotTFLQuantDequantizeOp(Operation* op) {
if (!op) return false;
if (llvm::isa<TFL::QuantizeOp, TFL::DequantizeOp>(op)) return false;
return true;
}
bool IsTerminatorOp(Operation* op) {
if (!op) return false;
return op->hasTrait<OpTrait::IsTerminator>();
}
// Try to guess the inference type of the op.
InferenceType GetInferenceType(Operation* op) {
bool float_type_observed = false;
bool int8_type_observed = false;
bool uint8_type_observed = false;
for (auto& input : op->getOpOperands()) {
auto input_type = input.get().getType();
if (IsF32ShapedType(input_type)) {
float_type_observed = true;
} else if (IsQI8Type(input_type)) {
int8_type_observed = true;
} else if (IsQUI8Type(input_type)) {
uint8_type_observed = true;
}
}
// We should not observe both uint8 & int8.
if (int8_type_observed && uint8_type_observed) return UNKNOWN;
if (float_type_observed) {
if (int8_type_observed || uint8_type_observed) {
return HYBRID;
} else {
return FLOAT;
}
}
if (int8_type_observed) {
return QUANTIZED_INT8;
}
if (uint8_type_observed) {
return QUANTIZED_UINT8;
}
// Default to float inference.
return FLOAT;
}
} // namespace tac
} // namespace TFL
} // namespace mlir
tensorflow/compiler/mlir/lite/experimental/tac/common/utils.h
+94
View File
@@ -0,0 +1,94 @@
/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_TAC_COMMON_UTILS_H_
#define TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_TAC_COMMON_UTILS_H_
#include "llvm/Support/Casting.h"
#include "mlir/Bytecode/BytecodeOpInterface.h" // from @llvm-project
#include "mlir/Dialect/Arith/IR/Arith.h" // from @llvm-project
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/Dialect/Quant/IR/QuantTypes.h" // from @llvm-project
#include "mlir/IR/BuiltinTypes.h" // from @llvm-project
#include "mlir/IR/OpDefinition.h" // from @llvm-project
#include "mlir/IR/Types.h" // from @llvm-project
#include "mlir/Interfaces/CallInterfaces.h" // from @llvm-project
#include "mlir/Interfaces/CastInterfaces.h" // from @llvm-project
#include "mlir/Support/LLVM.h" // from @llvm-project
#include "tensorflow/compiler/mlir/lite/experimental/tac/common/targets.h"
#include "tensorflow/compiler/mlir/lite/ir/tfl_ops.h"
#include "tensorflow/compiler/mlir/lite/utils/utils.h"
namespace mlir {
namespace TFL {
namespace tac {
// Returns true if 'op' is non const op. Returns false otherwise or if
// 'op' is null.
inline bool IsNonConstOp(Operation* op) {
if (!op) return false;
if (llvm::isa<arith::ConstantOp, mlir::func::ConstantOp>(op)) return false;
if (op->hasTrait<OpTrait::ConstantLike>()) return false;
if (llvm::isa<TFL::ConstOp, TFL::QConstOp>(op)) return false;
return true;
}
// Returns true if 'op' is a terminator op, otherwise false.
bool IsTerminatorOp(Operation* op);
// Returns true if 'op' is not TFL Quant / Dequant op. Returns False otherwise
// or if 'op' is null.
bool NotTFLQuantDequantizeOp(Operation* op);
// Returns true if it is a shaped type of f32 elements.
inline bool IsF32ShapedType(Type t) {
if (auto shaped_type = mlir::dyn_cast_or_null<ShapedType>(t)) {
return shaped_type.getElementType().isF32();
}
return false;
}
// Return true when the given element_type is QI8.
inline bool IsQI8Type(Type t) {
auto quantized_type = quant::QuantizedType::getQuantizedElementType(t);
return quantized_type != nullptr &&
quantized_type.getStorageTypeIntegralWidth() == 8 &&
quantized_type.isSigned();
}
// Return true when the given element_type is QUI8.
inline bool IsQUI8Type(Type t) {
auto quantized_type = quant::QuantizedType::getQuantizedElementType(t);
return quantized_type != nullptr &&
quantized_type.getStorageTypeIntegralWidth() == 8 &&
!quantized_type.isSigned();
}
// Return true when the given element_type is QI32.
inline bool IsQI32Type(Type t) {
auto quantized_type = quant::QuantizedType::getQuantizedElementType(t);
return quantized_type != nullptr &&
quantized_type.getStorageTypeIntegralWidth() == 32 &&
quantized_type.isSigned();
}
// Try to guess the inference type of the op.
InferenceType GetInferenceType(Operation* op);
} // namespace tac
} // namespace TFL
} // namespace mlir
#endif // TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_TAC_COMMON_UTILS_H_
tensorflow/compiler/mlir/lite/experimental/tac/examples/BUILD
+31
View File
@@ -0,0 +1,31 @@
load("//tensorflow:tensorflow.bzl", "tf_cc_binary")
load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:LICENSE"],
default_visibility = [
"//visibility:public",
],
licenses = ["notice"], # Apache 2.0
)
cc_library(
name = "example_hardware",
srcs = ["example_hardware.cc"],
hdrs = ["example_hardware.h"],
deps = [
"//tensorflow/compiler/mlir/lite:tensorflow_lite",
"//tensorflow/compiler/mlir/lite/experimental/tac:device_transform_patterns",
"//tensorflow/compiler/mlir/lite/experimental/tac/hardwares:simple_hardware",
],
alwayslink = 1,
)
tf_cc_binary(
name = "example-hardware-translate",
deps = [
":example_hardware",
"//tensorflow/compiler/mlir/lite/experimental/tac:tac-translate-lib",
"//tensorflow/compiler/mlir/lite/experimental/tac/hardwares:all-target-hardwares",
],
)
tensorflow/compiler/mlir/lite/experimental/tac/examples/example_hardware.cc
+47
View File
@@ -0,0 +1,47 @@
/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/compiler/mlir/lite/experimental/tac/examples/example_hardware.h"
#include <memory>
#include "tensorflow/compiler/mlir/lite/experimental/tac/transforms/device_transform_patterns.h"
#include "tensorflow/compiler/mlir/lite/ir/tfl_ops.h"
namespace mlir {
namespace TFL {
namespace tac {
constexpr char ExampleHardware::kId[]; // Define kId.
mlir::RewritePatternSet ExampleHardware::GetTransformations(
MLIRContext* context) const {
mlir::RewritePatternSet patterns(context);
patterns.add<LowerPackIntoConcatReshape, UnrollSplit, UnrollSplitV, PadSlice,
PadConcat>(context);
return patterns;
}
std::unique_ptr<TargetHardware> CreateExampleHardware() {
return std::make_unique<ExampleHardware>();
}
TargetHardwareRegistration<ExampleHardware> example_hardware(
"Example device", CreateExampleHardware);
} // namespace tac
} // namespace TFL
} // namespace mlir
tensorflow/compiler/mlir/lite/experimental/tac/examples/example_hardware.h
+45
View File
@@ -0,0 +1,45 @@
/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_TAC_EXAMPLES_EXAMPLE_HARDWARE_H_
#define TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_TAC_EXAMPLES_EXAMPLE_HARDWARE_H_
#include "tensorflow/compiler/mlir/lite/experimental/tac/hardwares/simple_hardware.h"
namespace mlir {
namespace TFL {
namespace tac {
class ExampleHardware : public SimpleHardware {
public:
static constexpr char kId[] = "ExampleHardware";
mlir::RewritePatternSet GetTransformations(
MLIRContext* context) const override;
mlir::TypeID GetTypeId() const override {
return mlir::TypeID::get<ExampleHardware>();
}
bool IsNotSupportedOp(mlir::Operation* op) const override { return false; }
float AdvantageOverCPU() const override { return 5.0; }
};
} // namespace tac
} // namespace TFL
} // namespace mlir
#endif // TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_TAC_EXAMPLES_EXAMPLE_HARDWARE_H_
tensorflow/compiler/mlir/lite/experimental/tac/execution_metadata_exporter.cc
+206
View File
@@ -0,0 +1,206 @@
// Copyright 2020 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#include "tensorflow/compiler/mlir/lite/experimental/tac/execution_metadata_exporter.h"
#include <cstdint>
#include <map>
#include <optional>
#include <string>
#include <utility>
#include <vector>
#include "flatbuffers/buffer.h" // from @flatbuffers
#include "flatbuffers/flatbuffer_builder.h" // from @flatbuffers
#include "flatbuffers/string.h" // from @flatbuffers
#include "flatbuffers/vector.h" // from @flatbuffers
#include "llvm/Support/Casting.h"
#include "mlir/Dialect/Arith/IR/Arith.h" // from @llvm-project
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/Attributes.h" // from @llvm-project
#include "mlir/IR/BuiltinAttributes.h" // from @llvm-project
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/IR/Operation.h" // from @llvm-project
#include "mlir/IR/Region.h" // from @llvm-project
#include "mlir/Support/LLVM.h" // from @llvm-project
#include "tensorflow/compiler/mlir/lite/experimental/tac/common/targets.h"
#include "tensorflow/compiler/mlir/lite/experimental/tac/hardwares/target_hardware.h"
#include "tensorflow/compiler/mlir/lite/experimental/tac/runtime_metadata_generated.h"
#include "tensorflow/compiler/mlir/lite/ir/tfl_ops.h"
#include "tensorflow/compiler/mlir/lite/quantization/ir/QuantOps.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
namespace tflite {
namespace {
bool IsConst(mlir::Operation* op) {
return llvm::isa<mlir::arith::ConstantOp, mlir::TF::ConstOp,
mlir::TFL::ConstOp, mlir::TFL::QConstOp>(op);
}
bool IsOpSupported(mlir::Operation* op, const std::string& hardware) {
auto* devce_hardware = mlir::TFL::tac::GetTargetHardware(hardware);
if (devce_hardware == nullptr) return {};
return devce_hardware->IsOpSupported(op);
}
bool HasValidHardwareTarget(mlir::Operation* op) {
// All TFLite ops has CPU interface, should be enough to check for cpu.
return IsOpSupported(op, "CPU");
}
std::optional<std::string> GetDeviceName(mlir::Operation* op) {
if (IsConst(op)) return std::nullopt;
// The model may contain quant stats op which is unrelevant to the
// execution.
if (llvm::isa<mlir::func::ReturnOp, mlir::quantfork::StatisticsOp>(op))
return std::nullopt;
if (!HasValidHardwareTarget(op)) return std::nullopt;
auto device = op->getAttrOfType<mlir::StringAttr>(mlir::TFL::tac::kDevice);
if (device == nullptr) return std::nullopt;
llvm::StringRef device_name_str = device.getValue();
return device_name_str.str();
}
std::optional<std::vector<float>> GetPerDeviceCosts(
const std::map<std::string, uint8_t>& hardware_map, mlir::Operation* op) {
auto device_costs_attr =
op->getAttrOfType<mlir::DictionaryAttr>("per_device_costs");
if (device_costs_attr == nullptr) return std::nullopt;
std::vector<float> device_costs(hardware_map.size(), -1.f);
for (const auto& kv : hardware_map) {
auto cost_attr = device_costs_attr.getNamed(kv.first);
if (!cost_attr.has_value()) return std::nullopt;
float cost = mlir::dyn_cast_or_null<mlir::FloatAttr>(cost_attr->getValue())
.getValueAsDouble();
device_costs[kv.second] = cost;
}
return device_costs;
}
flatbuffers::Offset<SubgraphMetadata> CreateSubgraphMetadata(
const std::map<std::string, uint8_t>& hardware_map, mlir::Region* Region,
flatbuffers::FlatBufferBuilder* builder) {
auto& block = Region->front();
int index = 0;
std::vector<flatbuffers::Offset<tflite::OpMetadata>> ops;
for (auto& inst : block) {
// Const nodes are mapped to const vectors in flatbuffer, so skip.
if (IsConst(&inst)) continue;
// The model may contain quant stats op which is unrelevant to the
// execution.
if (llvm::isa<mlir::func::ReturnOp, mlir::quantfork::StatisticsOp>(&inst))
continue;
// If an op doesn't implement any of the hardware interface we skip it.
// This can happen in cases like Flex when we have non TFLite ops.
auto device_name = GetDeviceName(&inst);
if (device_name.has_value()) {
// Add per device costs if present.
auto per_device_cost = GetPerDeviceCosts(hardware_map, &inst);
flatbuffers::Offset<flatbuffers::Vector<float>> per_device_cost_offset;
if (per_device_cost.has_value()) {
per_device_cost_offset = builder->CreateVector(*per_device_cost);
}
OpMetadataBuilder op_builder(*builder);
op_builder.add_index(index);
uint8_t hardware = hardware_map.at(*device_name);
op_builder.add_hardware(hardware);
if (per_device_cost.has_value()) {
op_builder.add_op_costs(per_device_cost_offset);
}
ops.push_back(op_builder.Finish());
}
index++;
}
return CreateSubgraphMetadata(*builder, builder->CreateVector(ops));
}
flatbuffers::Offset<tflite::HardwareMetadata>
CreateHardwareMetadataAndPopulateLookupTable(
std::vector<mlir::func::FuncOp>* funcs,
flatbuffers::FlatBufferBuilder* builder,
std::map<std::string, uint8_t>* hardware_names) {
uint8_t index = 0;
for (auto& func : *funcs) {
func.walk([&hardware_names, &index](mlir::Operation* op) {
auto device_name = GetDeviceName(op);
if (!device_name.has_value()) return;
auto iter = hardware_names->find(*device_name);
if (iter == hardware_names->end()) {
hardware_names->insert({*device_name, index++});
}
});
}
// Build the flatbuffer.
std::vector<flatbuffers::Offset<flatbuffers::String>> hardwares;
for (const auto& kv : *hardware_names) {
hardwares.push_back(builder->CreateString(kv.first));
}
return CreateHardwareMetadata(*builder, builder->CreateVector(hardwares));
}
} // namespace
std::optional<std::string> ExportRuntimeMetadata(mlir::ModuleOp module) {
mlir::func::FuncOp main_fn = module.lookupSymbol<mlir::func::FuncOp>("main");
if (!main_fn) return std::string("");
flatbuffers::FlatBufferBuilder fb_builder;
std::vector<mlir::func::FuncOp> funcs;
funcs.push_back(main_fn);
module.walk([&](mlir::func::FuncOp fn) {
if (fn != main_fn) {
funcs.push_back(fn);
}
});
// Populate the hardware metadata.
// And collect the hardwares used.
std::map<std::string, uint8_t> hardware_map;
flatbuffers::Offset<tflite::HardwareMetadata> hardware_metadata_offset =
CreateHardwareMetadataAndPopulateLookupTable(&funcs, &fb_builder,
&hardware_map);
// Populate the runtime metadata.
std::vector<flatbuffers::Offset<SubgraphMetadata>> subgraphs_metadata;
subgraphs_metadata.reserve(funcs.size());
for (auto& func : funcs) {
subgraphs_metadata.push_back(
CreateSubgraphMetadata(hardware_map, &func.getBody(), &fb_builder));
}
auto runtime_metadata =
CreateRuntimeMetadata(fb_builder, hardware_metadata_offset,
fb_builder.CreateVector(subgraphs_metadata));
fb_builder.Finish(runtime_metadata);
return std::string(
reinterpret_cast<const char*>(fb_builder.GetBufferPointer()),
fb_builder.GetSize());
}
} // namespace tflite
tensorflow/compiler/mlir/lite/experimental/tac/execution_metadata_exporter.h
+31
View File
@@ -0,0 +1,31 @@
// Copyright 2020 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#ifndef TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_TAC_EXECUTION_METADATA_EXPORTER_H_
#define TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_TAC_EXECUTION_METADATA_EXPORTER_H_
#include <optional>
#include <string>
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
namespace tflite {
// Returns serialized string for the generated flatbuffer.
std::optional<std::string> ExportRuntimeMetadata(mlir::ModuleOp module);
} // namespace tflite
#endif // TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_TAC_EXECUTION_METADATA_EXPORTER_H_
tensorflow/compiler/mlir/lite/experimental/tac/execution_metadata_exporter_test.cc
+128
View File
@@ -0,0 +1,128 @@
// Copyright 2020 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#include "tensorflow/compiler/mlir/lite/experimental/tac/execution_metadata_exporter.h"
#include <string>
#include <vector>
#include <gtest/gtest.h>
#include "flatbuffers/buffer.h" // from @flatbuffers
#include "flatbuffers/flatbuffer_builder.h" // from @flatbuffers
#include "flatbuffers/string.h" // from @flatbuffers
#include "mlir/Dialect/Arith/IR/Arith.h" // from @llvm-project
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/IR/DialectRegistry.h" // from @llvm-project
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "mlir/IR/OwningOpRef.h" // from @llvm-project
#include "mlir/Parser/Parser.h" // from @llvm-project
#include "tensorflow/compiler/mlir/lite/experimental/tac/runtime_metadata_generated.h"
#include "tensorflow/compiler/mlir/lite/ir/tfl_ops.h"
namespace tflite {
std::string CreateRuntimeMetadata() {
flatbuffers::FlatBufferBuilder fb_builder;
std::vector<flatbuffers::Offset<flatbuffers::String>> device_names = {
fb_builder.CreateString("GPU"), fb_builder.CreateString("CPU")};
const auto hardwares =
CreateHardwareMetadata(fb_builder, fb_builder.CreateVector(device_names));
const auto ops = {
CreateOpMetadata(fb_builder, 0, 0,
fb_builder.CreateVector(std::vector<float>({1.0, 5.0}))),
CreateOpMetadata(fb_builder, 1, 0,
fb_builder.CreateVector(std::vector<float>({1.0, 5.0}))),
CreateOpMetadata(fb_builder, 2, 0,
fb_builder.CreateVector(std::vector<float>({1.0, 5.0}))),
CreateOpMetadata(
fb_builder, 3, 1,
fb_builder.CreateVector(std::vector<float>({-1.0, 2.0}))),
};
const auto subgraphs = {CreateSubgraphMetadata(
fb_builder, fb_builder.CreateVector(ops.begin(), ops.size()))};
const auto metadata = CreateRuntimeMetadata(
fb_builder, hardwares,
fb_builder.CreateVector(subgraphs.begin(), subgraphs.size()));
fb_builder.Finish(metadata);
return std::string(
reinterpret_cast<const char*>(fb_builder.GetBufferPointer()),
fb_builder.GetSize());
}
void Verify(const RuntimeMetadata* result, const RuntimeMetadata* expected) {
EXPECT_EQ(result->subgraph_metadata()->size(),
expected->subgraph_metadata()->size());
for (int i = 0; i < result->subgraph_metadata()->size(); ++i) {
auto result_subgraph_metadata =
result->subgraph_metadata()->GetAs<SubgraphMetadata>(i);
auto expected_subgraph_metadata =
expected->subgraph_metadata()->GetAs<SubgraphMetadata>(i);
if (expected_subgraph_metadata->op_metadata() == nullptr &&
result_subgraph_metadata->op_metadata() == nullptr) {
return;
}
ASSERT_EQ(expected_subgraph_metadata->op_metadata()->size(),
result_subgraph_metadata->op_metadata()->size());
for (int j = 0; j < expected_subgraph_metadata->op_metadata()->size();
++j) {
auto result_op_metadata =
result_subgraph_metadata->op_metadata()->GetAs<OpMetadata>(j);
auto expected_op_metadata =
expected_subgraph_metadata->op_metadata()->GetAs<OpMetadata>(j);
EXPECT_EQ(result_op_metadata->index(), expected_op_metadata->index());
EXPECT_EQ(result_op_metadata->hardware(),
expected_op_metadata->hardware());
EXPECT_EQ(result_op_metadata->op_costs()->size(),
expected_op_metadata->op_costs()->size());
for (int i = 0; i < result_op_metadata->op_costs()->size(); ++i) {
EXPECT_FLOAT_EQ(result_op_metadata->op_costs()->Get(i),
expected_op_metadata->op_costs()->Get(i));
}
}
}
}
TEST(ExporterTest, Valid) {
const std::string kMLIR = R"(
func.func @main(%arg0: tensor<1xf32>, %arg1: tensor<1xf32>, %arg2: tensor<1xf32>, %arg3: tensor<1xf32>) -> tensor<2x1xf32> {
%0 = "tfl.add"(%arg0, %arg1) {fused_activation_function = "RELU6", per_device_costs = {CPU = 5.0 : f32, GPU = 1.0 : f32}, tac.device = "GPU"} : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32>
%1 = "tfl.mul"(%0, %arg2) {fused_activation_function = "RELU6", per_device_costs = {CPU = 5.0 : f32, GPU = 1.0 : f32}, tac.device = "GPU"} : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32>
%2 = "tfl.add"(%arg0, %arg3) {fused_activation_function = "RELU6", per_device_costs = {CPU = 5.0 : f32, GPU = 1.0 : f32}, tac.device = "GPU"} : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32>
%3 = "tfl.pack"(%1, %2) {axis = 0 : i32, per_device_costs = {CPU = 2.0 : f32, GPU = -1.0 : f32}, values_count = 2 : i32, tac.device = "CPU"} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2x1xf32>
func.return %3 : tensor<2x1xf32>
})";
const std::string kExpectedFB = CreateRuntimeMetadata();
mlir::DialectRegistry registry;
registry.insert<mlir::TFL::TensorFlowLiteDialect, mlir::arith::ArithDialect,
mlir::func::FuncDialect>();
mlir::MLIRContext context(registry);
auto module = mlir::OwningOpRef<mlir::ModuleOp>(
mlir::parseSourceString<mlir::ModuleOp>(kMLIR, &context));
auto module_op = module.get();
auto serialized_result_fb = ExportRuntimeMetadata(module_op);
const auto* result = GetRuntimeMetadata(serialized_result_fb.value().c_str());
const auto* expected = GetRuntimeMetadata(kExpectedFB.c_str());
ASSERT_TRUE(result != nullptr);
ASSERT_TRUE(result->subgraph_metadata() != nullptr);
ASSERT_TRUE(expected->subgraph_metadata() != nullptr);
Verify(result, expected);
}
} // namespace tflite
tensorflow/compiler/mlir/lite/experimental/tac/g3doc/images/compute_cost.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 38 KiB

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

After

Width:  |  Height:  |  Size: 29 KiB

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

After

Width:  |  Height:  |  Size: 41 KiB

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

After

Width:  |  Height:  |  Size: 38 KiB

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

After

Width:  |  Height:  |  Size: 38 KiB

tensorflow/compiler/mlir/lite/experimental/tac/hardwares/BUILD
+89
View File
@@ -0,0 +1,89 @@
load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:LICENSE"],
default_visibility = [
"//visibility:public",
],
licenses = ["notice"],
)
cc_library(
name = "gpu_hardware",
srcs = ["gpu_hardware.cc"],
hdrs = ["gpu_hardware.h"],
deps = [
":target_hardware",
"//tensorflow/compiler/mlir/lite:cost_estimators",
"//tensorflow/compiler/mlir/lite:tensorflow_lite",
"//tensorflow/compiler/mlir/lite/experimental/tac:common",
"//tensorflow/compiler/mlir/lite/experimental/tac:device_transform_patterns",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Support",
],
alwayslink = 1,
)
cc_library(
name = "cpu_hardware",
srcs = ["cpu_hardware.cc"],
deps = [
":target_hardware",
"//tensorflow/compiler/mlir/lite:cost_estimators",
"//tensorflow/compiler/mlir/lite:tensorflow_lite",
"//tensorflow/compiler/mlir/lite/experimental/tac:common",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Support",
],
alwayslink = 1,
)
cc_library(
name = "nnapi_hardware",
srcs = ["nnapi_hardware.cc"],
hdrs = ["nnapi_hardware.h"],
deps = [
":target_hardware",
"//tensorflow/compiler/mlir/lite:cost_estimators",
"//tensorflow/compiler/mlir/lite:tensorflow_lite",
"//tensorflow/compiler/mlir/lite/experimental/tac:device_transform_patterns",
"//tensorflow/compiler/mlir/lite/experimental/tac/hardwares:simple_hardware",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Support",
],
alwayslink = 1,
)
cc_library(
name = "target_hardware",
srcs = ["target_hardware.cc"],
hdrs = ["target_hardware.h"],
deps = [
"//tensorflow/compiler/mlir/lite:tensorflow_lite",
"//tensorflow/compiler/mlir/lite/experimental/tac:common",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Support",
],
)
cc_library(
name = "simple_hardware",
srcs = ["simple_hardware.cc"],
hdrs = ["simple_hardware.h"],
deps = [
":target_hardware",
"@llvm-project//mlir:IR",
],
)
cc_library(
name = "all-target-hardwares",
deps = [
":cpu_hardware",
":gpu_hardware",
":nnapi_hardware",
],
alwayslink = 1,
)
tensorflow/compiler/mlir/lite/experimental/tac/hardwares/cpu_hardware.cc
+181
View File
@@ -0,0 +1,181 @@
/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include <cstddef>
#include <cstdint>
#include <memory>
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "mlir/IR/Operation.h" // from @llvm-project
#include "mlir/IR/PatternMatch.h" // from @llvm-project
#include "mlir/Support/TypeID.h" // from @llvm-project
#include "tensorflow/compiler/mlir/lite/experimental/tac/common/targets.h"
#include "tensorflow/compiler/mlir/lite/experimental/tac/common/utils.h"
#include "tensorflow/compiler/mlir/lite/experimental/tac/hardwares/target_hardware.h"
#include "tensorflow/compiler/mlir/lite/ir/tfl_ops.h"
#include "tensorflow/compiler/mlir/lite/utils/arithmetic_count_util.h"
namespace mlir {
namespace TFL {
namespace tac {
namespace {
// CPU
constexpr float kCPUArithmeticUnitCost = 1.0;
// This basically assumes pure load/store. This is just fake data.
constexpr float kCPUCopyUnitCost = 0.5;
// Default values.
constexpr float kCPUDefaultFixedValuedCost = 10000.0;
// Quantized inference cost efficiency.
// For CPU, quantized inference is ~3x faster than the float alternative, this
// is just an estimation.
constexpr float kQuantizedInferenceEfficiency = 0.3;
inline float InferenceTypeEfficiency(InferenceType inference_type) {
if (inference_type == QUANTIZED_INT8 || inference_type == QUANTIZED_UINT8) {
return kQuantizedInferenceEfficiency;
}
return 1.0;
}
// CPU hardware class which handles CPU capabilities in TFLite.
// This is used by TAC to get op supported/ op cost estimates on CPU.
class CpuHardware : public TargetHardware {
public:
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(CpuHardware)
// String Identifier for CPU hardware.
static constexpr char kId[] = "CPU";
double GetHardwareSwitchingCost(const TargetHardware* from,
size_t buffer_size) const override {
auto from_type = from->GetTypeId();
auto to_type = GetTypeId();
if (from_type == to_type) return 0.0f;
// TODO(renjieliu): Implement a better version for different hardware cases.
return buffer_size * kCrossHardwareTransferPerByteCost / 8.0 +
kCrossHardwareTransferFixedCost;
}
mlir::RewritePatternSet GetTransformations(
MLIRContext* context) const override {
return {context};
}
mlir::TypeID GetTypeId() const override {
return mlir::TypeID::get<CpuHardware>();
}
bool IsOpSupported(mlir::Operation* op) const override {
// All ops in TFL dialect are supported on CPU.
if (op->getDialect() == nullptr) return false;
if (op->getDialect()->getNamespace() != "tfl") return false;
return true;
}
};
constexpr char CpuHardware::kId[]; // Define kId.
std::unique_ptr<TargetHardware> CreateCpuHardware() {
return std::make_unique<CpuHardware>();
}
TargetHardwareRegistration<CpuHardware> cpu_hardware("Target device for CPU",
CreateCpuHardware);
#define TAC_REGISTER_CPU_OP(Op, Create) \
TargetHardwareOpRegistration<CpuHardware, Op> Op##_CpuHardware_hardware( \
Create);
// Operation costs on CPU
// Currently used for these ops:
// tfl.conv_2d / tfl.depthwise_conv_2d / tfl.fully_connected
class CpuConvOp : public TargetHardwareOperation {
double GetOpCost(mlir::Operation* op) const override {
float cost = 0.0;
int64_t arithmetic_count;
if (ArithmeticCountUtilHelper::GetArithmeticCountForConvAndFullyconnectedOp(
op, &arithmetic_count)) {
cost = arithmetic_count * kCPUArithmeticUnitCost;
} else {
cost = kCPUDefaultFixedValuedCost;
}
return cost * InferenceTypeEfficiency(GetInferenceType(op));
}
bool IsOpSupported(mlir::Operation* op) const override { return true; }
};
std::unique_ptr<TargetHardwareOperation> CreateConvOp() {
return std::make_unique<CpuConvOp>();
}
// Currently used for these ops:
// tfl.Add / tfl.mul
class CpuArithmeticOp : public TargetHardwareOperation {
double GetOpCost(mlir::Operation* op) const override {
float cost = 0.0;
int64_t count;
if (ArithmeticCountUtilHelper::GetFirstOutputCount(op, &count)) {
cost = kCPUArithmeticUnitCost * count;
} else {
cost = kCPUDefaultFixedValuedCost;
}
return cost * InferenceTypeEfficiency(GetInferenceType(op));
}
bool IsOpSupported(mlir::Operation* op) const override { return true; }
};
std::unique_ptr<TargetHardwareOperation> CreateArithmeticOp() {
return std::make_unique<CpuArithmeticOp>();
}
// Currently used for these ops:
// tfl.concatenation / tfl.reshape / tfl.pack
class CpuConcatOp : public TargetHardwareOperation {
double GetOpCost(mlir::Operation* op) const override {
float cost = 0.0;
int64_t count;
if (ArithmeticCountUtilHelper::GetInputTensorTotalSize(op, &count)) {
cost = kCPUCopyUnitCost * count;
} else {
cost = kCPUDefaultFixedValuedCost;
}
return cost * InferenceTypeEfficiency(GetInferenceType(op));
}
bool IsOpSupported(mlir::Operation* op) const override { return true; }
};
std::unique_ptr<TargetHardwareOperation> CreateConcatOp() {
return std::make_unique<CpuConcatOp>();
}
TAC_REGISTER_CPU_OP(Conv2DOp, CreateConvOp);
TAC_REGISTER_CPU_OP(DepthwiseConv2DOp, CreateConvOp);
TAC_REGISTER_CPU_OP(FullyConnectedOp, CreateConvOp);
TAC_REGISTER_CPU_OP(AddOp, CreateArithmeticOp);
TAC_REGISTER_CPU_OP(MulOp, CreateArithmeticOp);
TAC_REGISTER_CPU_OP(ConcatenationOp, CreateConcatOp);
TAC_REGISTER_CPU_OP(ReshapeOp, CreateConcatOp);
TAC_REGISTER_CPU_OP(PackOp, CreateConcatOp);
#undef TAC_REGISTER_CPU_OP
} // namespace
} // namespace tac
} // namespace TFL
} // namespace mlir
tensorflow/compiler/mlir/lite/experimental/tac/hardwares/gpu_hardware.cc
+229
View File
@@ -0,0 +1,229 @@
/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/compiler/mlir/lite/experimental/tac/hardwares/gpu_hardware.h"
#include <cstddef>
#include <cstdint>
#include <memory>
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "mlir/IR/Operation.h" // from @llvm-project
#include "mlir/IR/PatternMatch.h" // from @llvm-project
#include "tensorflow/compiler/mlir/lite/experimental/tac/common/targets.h"
#include "tensorflow/compiler/mlir/lite/experimental/tac/common/utils.h"
#include "tensorflow/compiler/mlir/lite/experimental/tac/hardwares/target_hardware.h"
#include "tensorflow/compiler/mlir/lite/experimental/tac/transforms/device_transform_patterns.h"
#include "tensorflow/compiler/mlir/lite/ir/tfl_ops.h"
#include "tensorflow/compiler/mlir/lite/utils/arithmetic_count_util.h"
namespace mlir {
namespace TFL {
namespace tac {
namespace {
#include "tensorflow/compiler/mlir/lite/experimental/tac/transforms/generated_transform_patterns.inc"
} // namespace
constexpr char GpuHardware::kId[]; // Define kId.
mlir::RewritePatternSet GpuHardware::GetTransformations(
MLIRContext* context) const {
mlir::RewritePatternSet patterns(context);
patterns.add<LowerPackIntoConcatReshape, UnrollSplit, UnrollSplitV, SubToAdd,
EnsureBiasForConv2d, PadSlice, FullyConnectedToConv, PadConcat,
SquaredDifference>(context);
return patterns;
}
double GpuHardware::GetHardwareSwitchingCost(const TargetHardware* from,
size_t buffer_size) const {
auto from_type = from->GetTypeId();
auto to_type = GetTypeId();
if (from_type == to_type) return 0.0f;
// TODO(renjieliu): Implement a better version for different hardware cases.
return buffer_size * kCrossHardwareTransferPerByteCost / 8.0 +
kCrossHardwareTransferFixedCost;
}
bool GpuHardware::IsOpSupported(mlir::Operation* op) const {
if (TargetHardware::IsOpSupported(op)) {
return true;
}
// We also support quantized ops.
return !NotTFLQuantDequantizeOp(op);
}
namespace {
// GPU
constexpr float kGPUArithmeticUnitCost = 0.2;
// The copy can be non-consectutive copy. This is just fake data.
constexpr float kGPUCopyUnitCost = 0.2;
// Default values.
constexpr float kGPUDefaultFixedValuedCost = 10000.0;
std::unique_ptr<TargetHardware> CreateGpuHardware() {
return std::make_unique<GpuHardware>();
}
TargetHardwareRegistration<GpuHardware> gpu_hardware("Target device for GPU",
CreateGpuHardware);
#define TAC_REGISTER_GPU_OP(Op, Create) \
TargetHardwareOpRegistration<GpuHardware, Op> Op##_GpuHardware_hardware( \
Create);
// Currently used for these ops:
// tfl.Abs / tfl.Average_pool_2d / tfl.Cos / tfl.div / tfl.exp / tfl.hardswish /
// tfl.log / tfl.logistic / tfl.max_pool_2d / tfl.mirror_pad / tfl.maximum /
// tfl.custom / tfl.mean / tfl.minimum / tfl.pad / tfl.pow / tfl.prelu /
// tfl.relu / tfl.relu6 / tfl.rsqrt / tfl.sin / tfl.slice / tfl.softmax /
// tfl.space_to_depth / tfl.sqrt / tfl.square / tfl.squared_difference /
// tfl.strided_slice / tfl.tanh / tfl.transpose / tfl.transpose_conv
class GpuBasicSupportedOpNoCost : public TargetHardwareOperation {
double GetOpCost(mlir::Operation* op) const override { return 0; }
bool IsOpSupported(mlir::Operation* op) const override {
InferenceType inference_type = GetInferenceType(op);
if (inference_type != FLOAT) {
return false;
}
return true;
}
};
std::unique_ptr<TargetHardwareOperation> CreateBasicOpNoCost() {
return std::make_unique<GpuBasicSupportedOpNoCost>();
}
// Currently used for these ops:
// tfl.Add / tfl.mul
class GpuArithmeticOp : public TargetHardwareOperation {
double GetOpCost(mlir::Operation* op) const override {
int64_t count;
if (ArithmeticCountUtilHelper::GetFirstOutputCount(op, &count))
return kGPUArithmeticUnitCost * count;
return kGPUDefaultFixedValuedCost;
}
bool IsOpSupported(mlir::Operation* op) const override {
InferenceType inference_type = GetInferenceType(op);
if (inference_type != FLOAT) {
return false;
}
return true;
}
};
std::unique_ptr<TargetHardwareOperation> CreateArithmeticOp() {
return std::make_unique<GpuArithmeticOp>();
}
// Currently used for these ops:
// tfl.concatenation / tfl.reshape
class GpuConcatOp : public TargetHardwareOperation {
double GetOpCost(mlir::Operation* op) const override {
int64_t count;
if (ArithmeticCountUtilHelper::GetInputTensorTotalSize(op, &count))
return kGPUCopyUnitCost * count;
return kGPUDefaultFixedValuedCost;
}
bool IsOpSupported(mlir::Operation* op) const override {
InferenceType inference_type = GetInferenceType(op);
if (inference_type != FLOAT) {
return false;
}
return true;
}
};
std::unique_ptr<TargetHardwareOperation> CreateConcatOp() {
return std::make_unique<GpuConcatOp>();
}
// Currently used for these ops:
// tfl.conv_2d / tfl.depthwise_conv_2d / tfl.fully_connected
class GpuConvOp : public TargetHardwareOperation {
double GetOpCost(mlir::Operation* op) const override {
int64_t arithmetic_count;
if (ArithmeticCountUtilHelper::GetArithmeticCountForConvAndFullyconnectedOp(
op, &arithmetic_count)) {
return arithmetic_count * kGPUArithmeticUnitCost;
}
return kGPUDefaultFixedValuedCost;
}
bool IsOpSupported(mlir::Operation* op) const override {
InferenceType inference_type = GetInferenceType(op);
if (inference_type != FLOAT) {
return false;
}
return true;
}
};
std::unique_ptr<TargetHardwareOperation> CreateConvOp() {
return std::make_unique<GpuConvOp>();
}
// Op registrations
TAC_REGISTER_GPU_OP(AbsOp, CreateBasicOpNoCost);
TAC_REGISTER_GPU_OP(AveragePool2DOp, CreateBasicOpNoCost);
TAC_REGISTER_GPU_OP(CosOp, CreateBasicOpNoCost);
TAC_REGISTER_GPU_OP(DivOp, CreateBasicOpNoCost);
TAC_REGISTER_GPU_OP(ExpOp, CreateBasicOpNoCost);
TAC_REGISTER_GPU_OP(HardSwishOp, CreateBasicOpNoCost);
TAC_REGISTER_GPU_OP(LogOp, CreateBasicOpNoCost);
TAC_REGISTER_GPU_OP(LogisticOp, CreateBasicOpNoCost);
TAC_REGISTER_GPU_OP(MaxPool2DOp, CreateBasicOpNoCost);
TAC_REGISTER_GPU_OP(MirrorPadOp, CreateBasicOpNoCost);
TAC_REGISTER_GPU_OP(MaximumOp, CreateBasicOpNoCost);
TAC_REGISTER_GPU_OP(MinimumOp, CreateBasicOpNoCost);
TAC_REGISTER_GPU_OP(MeanOp, CreateBasicOpNoCost);
TAC_REGISTER_GPU_OP(CustomOp, CreateBasicOpNoCost);
TAC_REGISTER_GPU_OP(PadOp, CreateBasicOpNoCost);
TAC_REGISTER_GPU_OP(PowOp, CreateBasicOpNoCost);
TAC_REGISTER_GPU_OP(PReluOp, CreateBasicOpNoCost);
TAC_REGISTER_GPU_OP(ReluOp, CreateBasicOpNoCost);
TAC_REGISTER_GPU_OP(Relu6Op, CreateBasicOpNoCost);
TAC_REGISTER_GPU_OP(RsqrtOp, CreateBasicOpNoCost);
TAC_REGISTER_GPU_OP(SinOp, CreateBasicOpNoCost);
TAC_REGISTER_GPU_OP(SliceOp, CreateBasicOpNoCost);
TAC_REGISTER_GPU_OP(SoftmaxOp, CreateBasicOpNoCost);
TAC_REGISTER_GPU_OP(SpaceToDepthOp, CreateBasicOpNoCost);
TAC_REGISTER_GPU_OP(SqrtOp, CreateBasicOpNoCost);
TAC_REGISTER_GPU_OP(SquareOp, CreateBasicOpNoCost);
TAC_REGISTER_GPU_OP(SquaredDifferenceOp, CreateBasicOpNoCost);
TAC_REGISTER_GPU_OP(StridedSliceOp, CreateBasicOpNoCost);
TAC_REGISTER_GPU_OP(TanhOp, CreateBasicOpNoCost);
TAC_REGISTER_GPU_OP(TransposeOp, CreateBasicOpNoCost);
TAC_REGISTER_GPU_OP(TransposeConvOp, CreateBasicOpNoCost);
TAC_REGISTER_GPU_OP(ConcatenationOp, CreateConcatOp);
TAC_REGISTER_GPU_OP(ReshapeOp, CreateConcatOp);
TAC_REGISTER_GPU_OP(Conv2DOp, CreateConvOp);
TAC_REGISTER_GPU_OP(DepthwiseConv2DOp, CreateConvOp);
TAC_REGISTER_GPU_OP(FullyConnectedOp, CreateConvOp);
TAC_REGISTER_GPU_OP(AddOp, CreateArithmeticOp);
TAC_REGISTER_GPU_OP(MulOp, CreateArithmeticOp);
#undef TAC_REGISTER_GPU_OP
} // namespace
} // namespace tac
} // namespace TFL
} // namespace mlir
tensorflow/compiler/mlir/lite/experimental/tac/hardwares/gpu_hardware.h
+51
View File
@@ -0,0 +1,51 @@
/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_TAC_HARDWARES_GPU_HARDWARE_H_
#define TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_TAC_HARDWARES_GPU_HARDWARE_H_
#include <cstddef>
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "mlir/IR/PatternMatch.h" // from @llvm-project
#include "mlir/Support/TypeID.h" // from @llvm-project
#include "tensorflow/compiler/mlir/lite/experimental/tac/hardwares/target_hardware.h"
#include "tensorflow/compiler/mlir/lite/ir/tfl_ops.h"
namespace mlir {
namespace TFL {
namespace tac {
// Gpu hardware class which handles GPU capabilities in TFLite.
// This is used by TAC to get op supported/ op cost estimates on GPU.
class GpuHardware : public TargetHardware {
public:
static constexpr char kId[] = "GPU";
mlir::RewritePatternSet GetTransformations(
MLIRContext* context) const override;
mlir::TypeID GetTypeId() const override {
return mlir::TypeID::get<GpuHardware>();
}
double GetHardwareSwitchingCost(const TargetHardware* from,
size_t buffer_size) const override;
bool IsOpSupported(mlir::Operation* op) const override;
};
} // namespace tac
} // namespace TFL
} // namespace mlir
#endif // TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_TAC_HARDWARES_GPU_HARDWARE_H_
tensorflow/compiler/mlir/lite/experimental/tac/hardwares/nnapi_hardware.cc
+100
View File
@@ -0,0 +1,100 @@
/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/compiler/mlir/lite/experimental/tac/hardwares/nnapi_hardware.h"
#include <cstdint>
#include <memory>
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "mlir/IR/Operation.h" // from @llvm-project
#include "mlir/IR/PatternMatch.h" // from @llvm-project
#include "tensorflow/compiler/mlir/lite/experimental/tac/hardwares/target_hardware.h"
#include "tensorflow/compiler/mlir/lite/experimental/tac/transforms/device_transform_patterns.h"
#include "tensorflow/compiler/mlir/lite/ir/tfl_ops.h"
#include "tensorflow/compiler/mlir/lite/utils/arithmetic_count_util.h"
namespace mlir {
namespace TFL {
namespace tac {
// The copy can be non-consectutive copy. This is just fake data.
constexpr float kNNAPICopyUnitCost = 0.2;
// Default values.
constexpr float kNNAPIDefaultFixedValuedCost = 10000.0;
constexpr char NNAPIHardware::kId[]; // Define kId.
mlir::RewritePatternSet NNAPIHardware::GetTransformations(
MLIRContext* context) const {
mlir::RewritePatternSet patterns(context);
patterns.add<SquaredDifference, LowerPackIntoConcatReshape,
ReduceMeanToAvgPool, InsertRequantForReduceMean>(context);
return patterns;
}
std::unique_ptr<TargetHardware> CreateNNAPIHardware() {
return std::make_unique<NNAPIHardware>();
}
TargetHardwareRegistration<NNAPIHardware> nnapi_hardware(
"Target device for NNAPI", CreateNNAPIHardware);
// Currently used for these ops:
// tfl.squared_difference
class NNAPIBasicSupportedOpNoCost : public TargetHardwareOperation {
double GetOpCost(mlir::Operation* op) const override { return 0; }
bool IsOpSupported(mlir::Operation* op) const override {
return true;
}
};
std::unique_ptr<TargetHardwareOperation> CreateBasicOpNoCost() {
return std::make_unique<NNAPIBasicSupportedOpNoCost>();
}
// Currently used for these ops:
// tfl.concatenation / tfl.reshape / tfl.pack
class NNAPIConcatOp : public TargetHardwareOperation {
double GetOpCost(mlir::Operation* op) const override {
int64_t count;
if (ArithmeticCountUtilHelper::GetInputTensorTotalSize(op, &count))
return kNNAPICopyUnitCost * count;
return kNNAPIDefaultFixedValuedCost;
}
bool IsOpSupported(mlir::Operation* op) const override { return true; }
};
std::unique_ptr<TargetHardwareOperation> CreateConcatOp() {
return std::make_unique<NNAPIConcatOp>();
}
#define TAC_REGISTER_NNAPI_OP(Op, Create) \
TargetHardwareOpRegistration<NNAPIHardware, Op> Op##_NNAPIHardware_hardware( \
Create);
// Op registeration
TAC_REGISTER_NNAPI_OP(SquaredDifferenceOp, CreateBasicOpNoCost);
TAC_REGISTER_NNAPI_OP(ConcatenationOp, CreateConcatOp);
TAC_REGISTER_NNAPI_OP(ReshapeOp, CreateConcatOp);
TAC_REGISTER_NNAPI_OP(PackOp, CreateConcatOp);
#undef TAC_REGISTER_NNAPI_OP
} // namespace tac
} // namespace TFL
} // namespace mlir
tensorflow/compiler/mlir/lite/experimental/tac/hardwares/nnapi_hardware.h
+50
View File
@@ -0,0 +1,50 @@
/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
/* NNAPI Hardware Implementation */
#ifndef TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_TAC_HARDWARES_NNAPI_HARDWARE_H_
#define TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_TAC_HARDWARES_NNAPI_HARDWARE_H_
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "mlir/IR/Operation.h" // from @llvm-project
#include "mlir/IR/PatternMatch.h" // from @llvm-project
#include "mlir/Support/TypeID.h" // from @llvm-project
#include "tensorflow/compiler/mlir/lite/experimental/tac/hardwares/simple_hardware.h"
namespace mlir {
namespace TFL {
namespace tac {
class NNAPIHardware : public SimpleHardware {
public:
static constexpr char kId[] = "NNAPI";
mlir::RewritePatternSet GetTransformations(
MLIRContext* context) const override;
mlir::TypeID GetTypeId() const override {
return mlir::TypeID::get<NNAPIHardware>();
}
bool IsNotSupportedOp(mlir::Operation* op) const override { return false; }
float AdvantageOverCPU() const override { return 5.0; }
};
} // namespace tac
} // namespace TFL
} // namespace mlir
#endif // TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_TAC_HARDWARES_NNAPI_HARDWARE_H_
tensorflow/compiler/mlir/lite/experimental/tac/hardwares/simple_hardware.cc
+53
View File
@@ -0,0 +1,53 @@
/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/compiler/mlir/lite/experimental/tac/hardwares/simple_hardware.h"
#include <cstddef>
#include "mlir/IR/Operation.h" // from @llvm-project
#include "tensorflow/compiler/mlir/lite/experimental/tac/hardwares/target_hardware.h"
namespace mlir {
namespace TFL {
namespace tac {
bool SimpleHardware::IsOpSupported(mlir::Operation* op) const {
if (IsNotSupportedOp(op)) {
return false;
}
const TargetHardware* cpu = GetTargetHardware("CPU");
return cpu->IsOpSupported(op);
}
double SimpleHardware::GetHardwareSwitchingCost(const TargetHardware* from,
size_t buffer_size) const {
auto from_type = from->GetTypeId();
auto to_type = GetTypeId();
if (from_type == to_type) return 0.0f;
// TODO(renjieliu): Implement a better version for different hardware cases.
return buffer_size * kCrossHardwareTransferPerByteCost / 8.0 +
kCrossHardwareTransferFixedCost;
}
double SimpleHardware::GetOpCost(mlir::Operation* op) const {
const TargetHardware* cpu = GetTargetHardware("CPU");
return cpu->GetOpCost(op) / AdvantageOverCPU();
}
} // namespace tac
} // namespace TFL
} // namespace mlir
tensorflow/compiler/mlir/lite/experimental/tac/hardwares/simple_hardware.h
+67
View File
@@ -0,0 +1,67 @@
/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_TAC_HARDWARES_SIMPLE_HARDWARE_H_
#define TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_TAC_HARDWARES_SIMPLE_HARDWARE_H_
#include <cstddef>
#include "mlir/IR/Operation.h" // from @llvm-project
#include "tensorflow/compiler/mlir/lite/experimental/tac/hardwares/target_hardware.h"
namespace mlir {
namespace TFL {
namespace tac {
// A simple hardware is an interface makes you add a target backend easily if
// you don't want too much customization.
//
// It allows you to easily specify the ops capabilities (by
// specifying the denylist), the rest ops will be considered supported. Also you
// can also specify the advantage over CPU.
//
// If you need more customization, e.g., if you have your own hardware dialect,
// please consider use TargetHardware directly.
class SimpleHardware : public TargetHardware {
public:
// This is essentially a denylist.
// TODO(renjieliu): Consider whether we want an allowlist for custom op as
// well.
virtual bool IsNotSupportedOp(mlir::Operation* op) const = 0;
// The larger the value is, the more preferrable over CPU.
// If the value > 1, means the hardware has advantage over CPU.
// If the value < 1, means CPU is more preferred.
// If we specify 10.0, meaning the hardware is 10x faster than CPU.
// The value should be > 0.
// TODO(renjieliu): Consider add an interface for more detailed customization,
// for example, users should be able to specify some ops are preferred and
// some are less preferred.
virtual float AdvantageOverCPU() const = 0;
private:
bool IsOpSupported(mlir::Operation* op) const override;
double GetHardwareSwitchingCost(const TargetHardware* from,
size_t buffer_size) const override;
double GetOpCost(mlir::Operation* op) const override;
};
} // namespace tac
} // namespace TFL
} // namespace mlir
#endif // TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_TAC_HARDWARES_SIMPLE_HARDWARE_H_
tensorflow/compiler/mlir/lite/experimental/tac/hardwares/target_hardware.cc
+284
View File
@@ -0,0 +1,284 @@
/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/compiler/mlir/lite/experimental/tac/hardwares/target_hardware.h"
#include <functional>
#include <memory>
#include <string>
#include <vector>
#include "llvm/ADT/ArrayRef.h"
#include "llvm/ADT/DenseMap.h"
#include "llvm/Support/Casting.h"
#include "llvm/Support/raw_ostream.h"
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/OpDefinition.h" // from @llvm-project
#include "mlir/IR/Operation.h" // from @llvm-project
#include "mlir/Support/TypeID.h" // from @llvm-project
#include "tensorflow/compiler/mlir/lite/experimental/tac/common/targets.h"
#include "tensorflow/compiler/mlir/lite/experimental/tac/common/utils.h"
#include "tensorflow/compiler/mlir/lite/ir/tfl_ops.h"
#include "tensorflow/compiler/mlir/lite/utils/utils.h"
namespace mlir {
namespace TFL {
namespace tac {
namespace {
struct RegisteredTargetHardware {
RegisteredTargetHardware(
const std::string& name, const std::string& description,
mlir::TypeID type_id,
std::function<std::unique_ptr<TargetHardware>()> target_hardware_factory)
: unique_name(GetCanonicalHardwareName(name)),
description(description),
type_id(type_id),
target_hardware(target_hardware_factory()),
target_hardware_factory(target_hardware_factory) {}
std::string unique_name;
std::string description;
mlir::TypeID type_id;
std::unique_ptr<TargetHardware> target_hardware;
std::function<std::unique_ptr<TargetHardware>()> target_hardware_factory;
};
struct RegisteredTargetHardwareOps {
explicit RegisteredTargetHardwareOps(mlir::TypeID hardware_type)
: hardware_typeid(hardware_type) {}
// Key is the Operation TypeID
llvm::DenseMap<mlir::TypeID, std::unique_ptr<TargetHardwareOperation>>
target_hardware_ops;
// Key is the Operation TypeID
llvm::DenseMap<mlir::TypeID,
std::function<std::unique_ptr<TargetHardwareOperation>()>>
target_hardware_ops_factory;
mlir::TypeID hardware_typeid;
};
std::vector<std::unique_ptr<RegisteredTargetHardwareOps>>*
GetRegisteredTargetHardwareOps() {
static std::vector<std::unique_ptr<RegisteredTargetHardwareOps>>*
hardwares_ops =
[]() -> std::vector<std::unique_ptr<RegisteredTargetHardwareOps>>* {
return new std::vector<std::unique_ptr<RegisteredTargetHardwareOps>>();
}();
return hardwares_ops;
}
std::vector<RegisteredTargetHardware>* GetRegisteredHardwares() {
static std::vector<RegisteredTargetHardware>* hardwares =
[]() -> std::vector<RegisteredTargetHardware>* {
return new std::vector<RegisteredTargetHardware>();
}();
return hardwares;
}
llvm::DenseMap<mlir::TypeID, std::unique_ptr<TargetHardwareOperation>>*
getRegisteredOperationsForHardware(mlir::TypeID type_id) {
auto* hardwares = GetRegisteredTargetHardwareOps();
for (auto& hardware : *hardwares) {
if (hardware->hardware_typeid == type_id) {
return &hardware->target_hardware_ops;
}
}
return nullptr;
}
// A deny list for op cost computation since those ops are not arithemtic.
inline bool IsNonArithmeticOp(mlir::Operation* op) {
if (llvm::isa<func::ReturnOp, func::FuncOp>(op)) return true;
if (op->hasTrait<OpTrait::ConstantLike>()) return true;
if (llvm::isa<QConstOp, SparseQConstOp>(op)) return true;
if (!NotTFLQuantDequantizeOp(op)) return true;
return false;
}
} // namespace
bool TargetHardware::Init() {
auto* hardware_ops_factory = GetRegisteredTargetHardwareOps();
for (auto& hardware_ops : *hardware_ops_factory) {
if (hardware_ops->hardware_typeid != this->GetTypeId()) continue;
auto& op_factories = hardware_ops->target_hardware_ops_factory;
for (auto& op_factory : op_factories) {
hardware_ops_.emplace_back(op_factory.getSecond()());
}
break;
}
return true;
}
double TargetHardware::GetOpCost(mlir::Operation* op) const {
auto* registered_ops = getRegisteredOperationsForHardware(GetTypeId());
if (registered_ops == nullptr) {
return kDefaultFixedValuedCost;
}
auto abstract_op = op->getRegisteredInfo();
auto hardware_op = registered_ops->find(abstract_op->getTypeID());
if (hardware_op == registered_ops->end()) return kDefaultFixedValuedCost;
return hardware_op->second->GetOpCost(op);
}
bool TargetHardware::IsOpSupported(mlir::Operation* op) const {
auto* registered_ops = getRegisteredOperationsForHardware(GetTypeId());
if (registered_ops == nullptr) {
return false;
}
auto abstract_op = op->getRegisteredInfo();
auto hardware_op = registered_ops->find(abstract_op->getTypeID());
if (hardware_op == registered_ops->end()) return false;
return hardware_op->second->IsOpSupported(op);
}
double TargetHardware::GetFuncCost(func::FuncOp* func) const {
double total_cost = 0.0;
func->walk([&](Operation* op) {
if (IsNonArithmeticOp(op)) return;
// We will always defer to the hardware to decide the cost.
total_cost += GetOpCost(op);
});
return total_cost;
}
const TargetHardware* GetTargetHardware(const std::string& hardware_name) {
const std::string canonical_name = GetCanonicalHardwareName(hardware_name);
// Just loop for now, we don't expect number of hardwares to be huge.
// Revisit to have map if number of elements increased.
auto* registered_hardwares = GetRegisteredHardwares();
for (const auto& hardware : *registered_hardwares) {
if (hardware.unique_name == canonical_name) {
return hardware.target_hardware.get();
}
}
return nullptr;
}
std::function<std::unique_ptr<TargetHardware>()> GetTargetHardwareFactory(
const std::string& hardware_name) {
const std::string canonical_name = GetCanonicalHardwareName(hardware_name);
// Just loop for now, we don't expect number of hardwares to be huge.
// Revisit to have map if number of elements increased.
auto* registered_hardwares = GetRegisteredHardwares();
for (const auto& hardware : *registered_hardwares) {
if (hardware.unique_name == canonical_name) {
return hardware.target_hardware_factory;
}
}
return nullptr;
}
namespace internal {
void RegisterTargetHardwareFactory(
const std::string& unique_name, const std::string& description,
mlir::TypeID type_id,
std::function<std::unique_ptr<TargetHardware>()> target_hardware_factory) {
auto* registered_hardwares = GetRegisteredHardwares();
for (auto& hardware : *registered_hardwares) {
if (hardware.unique_name == unique_name) {
llvm::errs() << "Ignoring duplicate hardware. Hardware " << unique_name
<< " already registered\n";
hardware.target_hardware_factory = target_hardware_factory;
return;
}
}
registered_hardwares->push_back(RegisteredTargetHardware(
unique_name, description, type_id, target_hardware_factory));
}
void RegisterTargetHardwareOp(
mlir::TypeID hardware_type, mlir::TypeID op_type,
std::function<std::unique_ptr<TargetHardwareOperation>()>
target_hardware_op_factory) {
auto* registered_hardware_ops = GetRegisteredTargetHardwareOps();
for (auto& hardware : *registered_hardware_ops) {
if (hardware->hardware_typeid == hardware_type) {
if (hardware->target_hardware_ops.count(op_type)) {
llvm::errs() << "Trying to register duplicate Op";
return;
}
hardware->target_hardware_ops[op_type] = target_hardware_op_factory();
return;
}
}
registered_hardware_ops->push_back(
std::make_unique<RegisteredTargetHardwareOps>(
RegisteredTargetHardwareOps(hardware_type)));
registered_hardware_ops->back()->target_hardware_ops[op_type] =
target_hardware_op_factory();
}
void RegisterTargetHardwareOpFactory(
mlir::TypeID hardware_type, mlir::TypeID op_type,
std::function<std::unique_ptr<TargetHardwareOperation>()>
target_hardware_op_factory) {
auto* registered_hardware_ops = GetRegisteredTargetHardwareOps();
for (auto& hardware : *registered_hardware_ops) {
if (hardware->hardware_typeid == hardware_type) {
if (hardware->target_hardware_ops_factory.count(op_type)) {
llvm::errs() << "Trying to register duplicate Op";
return;
}
hardware->target_hardware_ops_factory[op_type] =
target_hardware_op_factory;
return;
}
}
registered_hardware_ops->push_back(
std::make_unique<RegisteredTargetHardwareOps>(
RegisteredTargetHardwareOps(hardware_type)));
registered_hardware_ops->back()->target_hardware_ops_factory[op_type] =
target_hardware_op_factory;
}
} // namespace internal
bool ProcessTargetDevices(llvm::ArrayRef<std::string> specified_device_specs,
std::vector<std::string>* device_specs) {
bool cpu_include = false;
for (auto& device_spec : specified_device_specs) {
auto device = GetCanonicalHardwareName(device_spec);
if (device == "CPU") cpu_include = true;
device_specs->push_back(device);
}
if (!cpu_include) {
device_specs->push_back("CPU");
}
// Make sure all the devices are registered.
for (const std::string& device : *device_specs) {
if (GetTargetHardware(device) == nullptr) {
llvm::errs() << "cannot get target hardware for device: " << device;
return false;
}
}
return true;
}
std::string GetHardwareName(const TargetHardware* hardware) {
const auto* registered_hardwares = GetRegisteredHardwares();
for (const auto& registered_hardware : *registered_hardwares) {
if (registered_hardware.type_id == hardware->GetTypeId())
return registered_hardware.unique_name;
}
return "";
}
} // namespace tac
} // namespace TFL
} // namespace mlir
tensorflow/compiler/mlir/lite/experimental/tac/hardwares/target_hardware.h
+196
View File
@@ -0,0 +1,196 @@
/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_TAC_HARDWARES_TARGET_HARDWARE_H_
#define TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_TAC_HARDWARES_TARGET_HARDWARE_H_
#include <cstddef>
#include <functional>
#include <memory>
#include <string>
#include <vector>
#include "llvm/ADT/ArrayRef.h"
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/DialectRegistry.h" // from @llvm-project
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "mlir/IR/Operation.h" // from @llvm-project
#include "mlir/IR/PatternMatch.h" // from @llvm-project
#include "mlir/Support/TypeID.h" // from @llvm-project
namespace mlir {
namespace TFL {
namespace tac {
// Default fixed values for ops.
constexpr static float kDefaultFixedValuedCost = 1000000.0;
// This is just fake data.
constexpr static float kCrossHardwareTransferPerByteCost = 5.0f;
// This is just fake data.
constexpr static float kCrossHardwareTransferFixedCost = 10.f;
// Interface for an Operation capabilities which should be tied to
// a specific hardware.
// Users should implement the interface and use TargetHardwareOpRegistration
// for registering the operation.
class TargetHardwareOperation {
public:
virtual ~TargetHardwareOperation() = default;
virtual double GetOpCost(mlir::Operation* op) const = 0;
virtual bool IsOpSupported(mlir::Operation* op) const = 0;
};
// Abstract base class for a hardware.
// To introduce new hardware
// users should implement the interface and use TargetHardwareRegistration
// for registering the hardware.
// Subclasses must implement the pure virtual function interface and
// define static member variable that retrieves string identifying the Target
// Hardware. Example,
// class MyType : public TargetHardware {
// public:
// static constexpr char kId[] = "MyHardware";
// };
class TargetHardware {
public:
virtual ~TargetHardware() = default;
// Initializes all TargetHardwareOperation registered for this hardware.
// Users overriding this function, should call the base class method to
// initialize the ops.
virtual bool Init();
// Returns the cost of running 'op' on this Hardware.
virtual double GetOpCost(mlir::Operation* op) const;
// Returns the cost of running the whole function on this hardware.
// By default this is the sum of the cost of individual cost for each op.
virtual double GetFuncCost(func::FuncOp* func) const;
// Returns true if 'op' can run on this Hardware.
virtual bool IsOpSupported(mlir::Operation* op) const;
// Switching cost between from hardware and this hardware.
// If both the hardwares are the same, the transfer cost is basically 0.
virtual double GetHardwareSwitchingCost(const TargetHardware* from,
size_t buffer_size) const = 0;
// Returns a list of all patterns to apply for this hardware.
virtual mlir::RewritePatternSet GetTransformations(
MLIRContext* context) const = 0;
// Returns TypeId for the provided hardware.
// Usually should be something like mlir::TypeID::get<MyType>()
virtual mlir::TypeID GetTypeId() const = 0;
virtual void GetDependentDialects(mlir::DialectRegistry& registry) const {}
protected:
// All registered hardware ops.
std::vector<std::unique_ptr<TargetHardwareOperation>> hardware_ops_;
};
// Returns pointer to the Hardware identified by 'hardware_name'.
// If not found nullptr is returned.
// DEPRECATED: Do not use, prefer GetTargetHardwareFactory instead.
const TargetHardware* GetTargetHardware(const std::string& hardware_name);
// Returns the factory method for the requested hardware if present.
std::function<std::unique_ptr<TargetHardware>()> GetTargetHardwareFactory(
const std::string& hardware_name);
namespace internal {
void RegisterTargetHardwareFactory(
const std::string& unique_name, const std::string& description,
mlir::TypeID type_id,
std::function<std::unique_ptr<TargetHardware>()> target_hardware_factory);
// Registers the provided target hardware factory.
template <typename T>
void RegisterTargetHardwareFactory(
const std::string& description,
std::function<std::unique_ptr<TargetHardware>()> target_hardware_factory) {
RegisterTargetHardwareFactory(T::kId, description, mlir::TypeID::get<T>(),
target_hardware_factory);
}
// DEPRECATED: Do not use, prefer RegisterTargetHardwareOpFactory intstead.
void RegisterTargetHardwareOp(
mlir::TypeID hardware_type, mlir::TypeID op_type,
std::function<std::unique_ptr<TargetHardwareOperation>()>
target_hardware_op_factory);
void RegisterTargetHardwareOpFactory(
mlir::TypeID hardware_type, mlir::TypeID op_type,
std::function<std::unique_ptr<TargetHardwareOperation>()>
target_hardware_op_factory);
} // namespace internal
// Register target hardware.
template <typename Hardware>
struct TargetHardwareRegistration {
TargetHardwareRegistration(const std::string& description,
std::function<std::unique_ptr<TargetHardware>()>
target_hardware_factory) {
internal::RegisterTargetHardwareFactory<Hardware>(description,
target_hardware_factory);
}
};
// Register Op capabilities for specific hardware.
template <typename Hardware, typename Op>
struct TargetHardwareOpRegistration {
explicit TargetHardwareOpRegistration(
std::function<std::unique_ptr<TargetHardwareOperation>()>
target_hardware_op_factory) {
// TODO(b/177376459): remove this.
internal::RegisterTargetHardwareOp(mlir::TypeID::get<Hardware>(),
mlir::TypeID::get<Op>(),
target_hardware_op_factory);
internal::RegisterTargetHardwareOpFactory(mlir::TypeID::get<Hardware>(),
mlir::TypeID::get<Op>(),
target_hardware_op_factory);
}
};
//======== util functions ==========
// Process user specified device specs, will always add CPU if it's not there.
// specified_device_specs: ',' separated, like "GPU,DSP,CPU".
// device_specs: processed device specs enum.
bool ProcessTargetDevices(llvm::ArrayRef<std::string> specified_device_specs,
std::vector<std::string>* device_specs);
// Check whether two hardwares are the same.
inline bool IsSameHardware(const TargetHardware* lhs,
const TargetHardware* rhs) {
return lhs->GetTypeId() == rhs->GetTypeId();
}
// Returns the ID identifying 'hardware'. This should match the ID defined
// in the hardware field ID.
// For example, if MyHardware is passed the value returned should match
// MyHardware::kId.
std::string GetHardwareName(const TargetHardware* hardware);
} // namespace tac
} // namespace TFL
} // namespace mlir
#endif // TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_TAC_HARDWARES_TARGET_HARDWARE_H_
tensorflow/compiler/mlir/lite/experimental/tac/py_wrapper/BUILD
+110
View File
@@ -0,0 +1,110 @@
load("@rules_cc//cc:cc_library.bzl", "cc_library")
load("//tensorflow:tensorflow.bzl", "VERSION")
load("//tensorflow:tensorflow.default.bzl", "pybind_extension")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:LICENSE"],
default_visibility = [
"//tensorflow/compiler/mlir/lite/experimental/tac:__subpackages__",
],
licenses = ["notice"],
)
cc_library(
name = "tac_wrapper_lib",
srcs = ["tac_wrapper.cc"],
hdrs = [
"tac_wrapper.h",
],
deps = [
"//tensorflow/compiler/mlir/lite/experimental/tac:tac_importer_exporter",
"//tensorflow/compiler/mlir/lite/experimental/tac:target_aware_conversion",
"//tensorflow/compiler/mlir/lite/experimental/tac:tflite_importer_exporter",
"//tensorflow/compiler/mlir/lite/experimental/tac/hardwares:all-target-hardwares",
"//tensorflow/compiler/mlir/tensorflow",
"@com_google_absl//absl/status",
"@llvm-project//mlir:IR",
"@xla//third_party/python_runtime:headers", # buildcleaner: keep
],
)
pybind_extension(
name = "_pywrap_tac_wrapper",
srcs = [
"tac_wrapper_pybind11.cc",
],
hdrs = ["tac_wrapper.h"],
dynamic_deps = select({
"//tensorflow:macos": ["//tensorflow:libtensorflow_framework.%s.dylib" % VERSION],
"//tensorflow:windows": [],
"//conditions:default": ["//tensorflow:libtensorflow_framework.so.%s" % VERSION],
}),
enable_stub_generation = True,
pytype_srcs = [
"_pywrap_tac_wrapper.pyi",
],
static_deps = [
"@arm_neon_2_x86_sse//:__subpackages__",
"@bazel_tools//:__subpackages__",
"@boringssl//:__subpackages__",
"@clog//:__subpackages__",
"@com_github_cares_cares//:__subpackages__",
"@com_github_googlecloudplatform_tensorflow_gcp_tools//:__subpackages__",
"@com_github_grpc_grpc//:__subpackages__",
"@com_google_absl//:__subpackages__",
"@com_google_googleapis//:__subpackages__",
"@com_google_protobuf//:__subpackages__",
"@com_googlesource_code_re2//:__subpackages__",
"@compute_library//:__subpackages__",
"@cpuinfo//:__subpackages__",
"@curl//:__subpackages__",
"@eigen_archive//:__subpackages__",
"@farmhash_archive//:__subpackages__",
"@farmhash_gpu_archive//:__subpackages__",
"@fft2d//:__subpackages__",
"@flatbuffers//:__subpackages__",
"@FP16//:__subpackages__",
"@FXdiv//:__subpackages__",
"@gemmlowp//:__subpackages__",
"@go_googleapis//google/api:annotations_proto",
"@gif//:__subpackages__",
"@highwayhash//:__subpackages__",
"@hwloc//:__subpackages__",
"@icu//:__subpackages__",
"@jsoncpp_git//:__subpackages__",
"@libjpeg_turbo//:__subpackages__",
"@llvm_openmp//:__subpackages__",
"@llvm-project//:__subpackages__",
"@llvm_terminfo//:__subpackages__",
"@llvm_zlib//:__subpackages__",
"@local_config_cuda//:__subpackages__",
"@local_config_git//:__subpackages__",
"@local_config_nccl//:__subpackages__",
"@local_config_python//:__subpackages__",
"@local_config_rocm//:__subpackages__",
"@local_config_tensorrt//:__subpackages__",
"@mkl_dnn_acl_compatible//:__subpackages__",
"@onednn//:__subpackages__",
"@org_sqlite//:__subpackages__",
"@platforms//:__subpackages__",
"@png//:__subpackages__",
"@pthreadpool//:__subpackages__",
"@pybind11//:__subpackages__",
"@ruy//:__subpackages__",
"@snappy//:__subpackages__",
"@sobol_data//:__subpackages__",
"@stablehlo//:__subpackages__",
"//:__subpackages__",
"@upb//:__subpackages__",
"@XNNPACK//:__subpackages__",
"@zlib//:__subpackages__",
"@tsl//tsl:__subpackages__",
"@xla//xla:__subpackages__",
],
deps = [
":tac_wrapper_lib",
"//tensorflow/python/lib/core:pybind11_lib",
"@pybind11",
"@xla//third_party/python_runtime:headers",
],
)
tensorflow/compiler/mlir/lite/experimental/tac/py_wrapper/_pywrap_tac_wrapper.pyi
+16
View File
@@ -0,0 +1,16 @@
# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
def run_tac(arg0: str, arg1: list[str], arg2: str) -> bool: ...
tensorflow/compiler/mlir/lite/experimental/tac/py_wrapper/tac_wrapper.cc
+72
View File
@@ -0,0 +1,72 @@
/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/compiler/mlir/lite/experimental/tac/py_wrapper/tac_wrapper.h"
#include <memory>
#include <string>
#include <vector>
#include "absl/status/status.h"
#include "mlir/IR/DialectRegistry.h" // from @llvm-project
#include "tensorflow/compiler/mlir/lite/experimental/tac/tac_importer_exporter.h"
#include "tensorflow/compiler/mlir/lite/experimental/tac/tac_module.h"
#include "tensorflow/compiler/mlir/lite/experimental/tac/tflite_import_export.h"
#include "tensorflow/compiler/mlir/lite/experimental/tac/transforms/passes.h"
#include "tensorflow/compiler/mlir/tensorflow/dialect_registration.h"
namespace tflite {
namespace {
std::unique_ptr<mlir::TFL::tac::TacImporter> CreateTfLiteImporter(
const std::string input_file_name) {
mlir::TFL::tac::TfLiteImporter::Options options;
options.file_name = input_file_name;
options.input_mlir = false;
return std::make_unique<mlir::TFL::tac::TfLiteImporter>(options);
}
std::unique_ptr<mlir::TFL::tac::TacExporter> CreateTfLiteExporter(
const std::string output_file_name,
const std::vector<std::string>& target_hardware_backends) {
mlir::TFL::tac::TfLiteExporter::Options options;
options.output_mlir = false;
options.output_file_name = output_file_name;
options.export_runtime_metadata = false;
options.target_hardware_backends = target_hardware_backends;
return std::make_unique<mlir::TFL::tac::TfLiteExporter>(options);
}
} // namespace
// Run target-aware-conversion for the given tflite model with the given device
// specs.
// Warning: The API is experimental and subject to changes.
bool run_tac(const std::string& model_file_path,
const std::vector<std::string>& device_specs,
const std::string& model_output_path) {
mlir::TFL::tac::TacModule::Options options;
options.hardware_backends = device_specs;
options.enable_inliner = true;
options.legalize_to_tflite_ops = true;
mlir::TFL::tac::TacModule tac_module(options);
mlir::DialectRegistry registry;
mlir::RegisterAllTensorFlowDialects(registry);
tac_module.RegisterExtraDialects(registry);
tac_module.SetImporter(CreateTfLiteImporter(model_file_path));
tac_module.SetExporter(
CreateTfLiteExporter(model_output_path, options.hardware_backends));
return tac_module.Run().ok();
}
} // namespace tflite
tensorflow/compiler/mlir/lite/experimental/tac/py_wrapper/tac_wrapper.h
+42
View File
@@ -0,0 +1,42 @@
/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_TAC_PY_WRAPPER_TAC_WRAPPER_H_
#define TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_TAC_PY_WRAPPER_TAC_WRAPPER_H_
#include <functional>
#include <memory>
#include <string>
#include <vector>
// Place `<locale>` before <Python.h> to avoid build failures in macOS.
#include <locale>
// The empty line above is on purpose as otherwise clang-format will
// automatically move <Python.h> before <locale>.
#include <Python.h>
namespace tflite {
// Run target-aware-conversion for the given tflite model with the given device
// specs.
// Warning: The API is experimental and subject to change.
bool run_tac(const std::string& model_file_path,
const std::vector<std::string>& device_specs,
const std::string& model_output_path);
} // namespace tflite
#endif // TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_TAC_PY_WRAPPER_TAC_WRAPPER_H_
tensorflow/compiler/mlir/lite/experimental/tac/py_wrapper/tac_wrapper_pybind11.cc
+37
View File
@@ -0,0 +1,37 @@
/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include <pybind11/stl.h>
#include <string>
#include <vector>
#include "pybind11/pybind11.h" // from @pybind11
#include "tensorflow/compiler/mlir/lite/experimental/tac/py_wrapper/tac_wrapper.h"
// Warning: The API is experimental and subject to change.
PYBIND11_MODULE(_pywrap_tac_wrapper, m) {
m.def(
"run_tac",
[](const std::string& model_file_path,
const std::vector<std::string>& device_specs,
const std::string& model_output_path) {
return ::tflite::run_tac(model_file_path, device_specs,
model_output_path);
},
R"pbdoc(
Run target-aware-conversion with the given device specs.
)pbdoc");
}
tensorflow/compiler/mlir/lite/experimental/tac/runtime_metadata.fbs
+89
View File
@@ -0,0 +1,89 @@
// Copyright 2020 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
namespace tflite;
// Here is an overview of the pipeline & runtime metadata looks like:
//
// ======== IRs after TAC ========
// op1 = .... {tac.device = "CPU"}
// op2 = .... {tac.device = "GPU"}
// op3 = .... {tac.device = "GPU"}
// op4 = .... {tac.device = "CPU"}
//
// We will run a separate cost estimation for each op with each different
// hardwares.
//
// ======== IRs after cost-estimation ========
// op1 = .... {tac.device = "CPU"} {"CPU": 10, "GPU": -1.0}
// op2 = .... {tac.device = "GPU"} {"CPU": 20, "GPU": 5.0}
// op3 = .... {tac.device = "GPU"} {"CPU": 30, "GPU": 10.0}
// op4 = .... {tac.device = "CPU"} {"CPU": 40, "GPU": -1.0}
//
// The exported runtime metadata will look like below:
// ======== Runtime Metadata ========
//
// RuntimeMetadata: {
// hardwares: {
// ["CPU", "GPU"]
// }
//
// subgraph_metadata: {
// [
// op1:
// hardware: 0
// op_costs: [10, -1],
// op2:
// hardware: 1
// op_costs: [20, 5],
// op3:
// hardware: 1
// op_costs: [30, 10],
// op4:
// hardware: 0
// op_costs: [40, -1],
// ]
// }
// }
table HardwareMetadata {
// List of different hardwares this model can use.
hardware_name:[string];
}
// Metadata for a single Op.
table OpMetadata {
// Node index in the corresponding subgraph.
index:int;
// The hardware suggested for inference.
// This will correspond to the index in the hardware_name field in the
// hardware metadata.
hardware:uint8;
// The estimated costs for the given op.
// The index matches the hardware array.
op_costs: [float];
}
table SubgraphMetadata {
// Metadata for all ops in the current subgraph.
op_metadata:[OpMetadata];
}
table RuntimeMetadata {
// Metadata about different hardwares used by the model.
hardwares: HardwareMetadata;
// Metadata for each subgraph in same order as in the TFLite model file.
subgraph_metadata:[SubgraphMetadata];
}
root_type RuntimeMetadata;
tensorflow/compiler/mlir/lite/experimental/tac/tac.py
+46
View File
@@ -0,0 +1,46 @@
# Copyright 2021 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Target aware conversion for TFLite model."""
from tensorflow.compiler.mlir.lite.experimental.tac.py_wrapper import _pywrap_tac_wrapper
def run_tac(model_path, targets, output_path):
"""Run target aware conversion for the given tflite model file.
Args:
model_path: Path to the tflite model file.
targets: A list of string of the desired targets. E.g., ['GPU', 'CPU'].
output_path: The output path.
Returns:
Whether the optimization succeeded.
Raises:
ValueError:
Invalid model_path.
Targets are not specified.
Invalid output_path.
"""
if not model_path:
raise ValueError("Invalid model_path.")
if not targets:
raise ValueError("Targets are not specified.")
if not output_path:
raise ValueError("Invalid output_path.")
return _pywrap_tac_wrapper.run_tac(model_path, targets, output_path)
tensorflow/compiler/mlir/lite/experimental/tac/tac_filter.proto
+87
View File
@@ -0,0 +1,87 @@
// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
syntax = "proto3";
package third_party.tensorflow.compiler.mlir.lite.experimental.tac;
import "google/protobuf/any.proto";
// A list of filters for TAC users to run ops/functions on ML hardwares. The
// intuition is that, for ops/functions that can be run on ML hardware (e.g.
// EdgeTPU) and TFLite CPU, TAC users give a hint that they're more performant
// to run on TFLite CPU. These filters give the TAC users freedom to specify the
// parts that they want to use other hardware to accelerate.
message TacFilters {
// A list of filters/rules to specify the parts that user wants to run on
// other hardware.
repeated TacFilter tac_filters = 1;
}
// A filter can be used for an op or function.
message TacFilter {
oneof filter {
OpFilter op_filter = 1;
FunctionFilter function_filter = 2;
}
}
// Function filter is to include/exclude a function in the target annotation
// pass in the TAC tool pipeline.
message FunctionFilter {
// Function filter types that are supported. If one function is matched for
// two rules with conflict, INCLUDE_TARGET_ANNOTATION has higher priority.
enum FunctionFilterType {
// To skip this function in the target annotation pass. This means all ops
// in this function run on TFLite CPU.
SKIP_TARGET_ANNOTATION = 0;
// To include this function in the target annotation pass. This has higher
// priority than `SKIP_TARGET_ANNOTATION`.
INCLUDE_TARGET_ANNOTATION = 1;
}
// This name corresponds to the TFLite subgraph name in the flatbuffer.
// `function_name_pattern` supports regex matching.
string function_name_pattern = 1;
FunctionFilterType filter_type = 2;
}
// Op filter is to filter out ops that user wants to run. Ops with this filter
// run on TFLite CPU.
message OpFilter {
// This name corresponds to the mlir::Location of the tensor.
// `op_name_pattern` supports regex matching.
string op_name_pattern = 1;
enum MatchType {
// To match the op name with the `op_name_pattern` directly.
MATCH = 0;
// To invert the match with the `op_name_pattern`, ie, the filter will
// select an op if its name does not match the pattern above.
INVERT_MATCH = 1;
}
MatchType match_type = 2;
enum DeviceType {
// To run the op on CPU.
CPU = 0;
// To run the op on an hardware unit other than the CPU.
OTHER = 1;
}
DeviceType device_type = 3;
// An arbitrary string that can be used to identify the filter or a set of
// filters.
string tag = 4;
// Custom options that can be used to pass additional information for the
// filter.
google.protobuf.Any custom_options = 5;
}
tensorflow/compiler/mlir/lite/experimental/tac/tac_importer_exporter.h
+52
View File
@@ -0,0 +1,52 @@
/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_TAC_TAC_IMPORTER_EXPORTER_H_
#define TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_TAC_TAC_IMPORTER_EXPORTER_H_
#include "absl/status/statusor.h"
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
namespace mlir {
namespace TFL {
namespace tac {
// Interface for Importing program to TAC (Target Aware Conversion) Module.
// This class is an interface for importing program in TAC.
// See TacModule in how to register it with the module and use it.
class TacImporter {
public:
virtual ~TacImporter() = default;
// Imports and returns the Module for the imported program.
virtual absl::StatusOr<mlir::OwningOpRef<mlir::ModuleOp>> Import() = 0;
};
// Interface for exporting a module.
// Users should implement the interface for exporting the result from TAC
// in their preferred way.
// See TacModule in how to register it with the module and use it.
class TacExporter {
public:
virtual ~TacExporter() = default;
// Imports and returns the Module for the imported program.
virtual absl::Status Export(mlir::ModuleOp module) = 0;
};
} // namespace tac
} // namespace TFL
} // namespace mlir
#endif // TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_TAC_TAC_IMPORTER_EXPORTER_H_
tensorflow/compiler/mlir/lite/experimental/tac/tac_module.cc
+149
View File
@@ -0,0 +1,149 @@
/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/compiler/mlir/lite/experimental/tac/tac_module.h"
#include <memory>
#include <string>
#include <vector>
#include "absl/status/status.h"
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/Transforms/Passes.h" // from @llvm-project
#include "tensorflow/compiler/mlir/lite/experimental/tac/hardwares/target_hardware.h"
#include "tensorflow/compiler/mlir/lite/experimental/tac/transforms/passes.h"
#include "tensorflow/compiler/mlir/lite/transforms/passes.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/error_util.h"
namespace mlir {
namespace TFL {
namespace tac {
namespace {
// TODO(b/177376459): We should make this configureable.
void AddExportTFLPass(mlir::OpPassManager* pass_manager, bool enable_inliner) {
if (enable_inliner) pass_manager->addPass(mlir::createInlinerPass());
pass_manager->addPass(mlir::createSymbolDCEPass());
pass_manager->addNestedPass<mlir::func::FuncOp>(
mlir::createCanonicalizerPass());
pass_manager->addNestedPass<mlir::func::FuncOp>(mlir::createCSEPass());
}
} // namespace
// TODO(b/177376459): We should make this configureable.
void TacModule::AddTACPass(mlir::OpPassManager* pass_manager,
llvm::ArrayRef<std::string> device_specs) {
pass_manager->addPass(mlir::TFL::tac::CreateTargetAnnotationPass(this));
pass_manager->addPass(mlir::TFL::tac::CreateRaiseTargetSubgraphsPass());
pass_manager->addPass(mlir::TFL::tac::CreateFoldConstantsToSubgraphPass(
/*fold_all_constants=*/false));
pass_manager->addPass(
mlir::TFL::tac::CreateAlternativeSubgraphPass(device_specs));
if (options_.legalize_to_tflite_ops) {
// After we creat the alternative subgraph, we can still do canonicalization
// legalization & other optimizations as long as we're not inlining the
// function.
// And in fact, we probably need to do the proper legalization, for the
// compute cost to work. (in case we added some TF ops)
pass_manager->addPass(mlir::TFL::CreatePrepareTFPass(
/*unfold_batch_matmul=*/true,
/*allow_bf16_and_f16_type_legalization=*/false));
pass_manager->addNestedPass<mlir::func::FuncOp>(
mlir::createCanonicalizerPass());
pass_manager->addPass(
mlir::TFL::CreateLegalizeTFPass(/*run_tfl_runtime_verification=*/true));
pass_manager->addPass(mlir::TFL::CreateOptimizePass());
}
pass_manager->addPass(mlir::TFL::tac::CreateComputeCostPass());
pass_manager->addPass(mlir::TFL::tac::CreatePickSubgraphsPass());
// After this pass, we may consider add a pass to merge small functions into
// large functions (and maybe other metadata as well).
}
const tac::TargetHardware* TacModule::GetTargetHardware(
const std::string& hardware_name) const {
for (auto& hardware : backends_) {
if (GetHardwareName(hardware.get()) == hardware_name) return hardware.get();
}
return nullptr;
}
absl::Status TacModule::RunTacPasses(mlir::ModuleOp* module, bool debug_mode) {
mlir::PassManager pm((*module)->getName(),
mlir::OpPassManager::Nesting::Implicit);
AddTACPass(&pm, options_.hardware_backends);
if (!debug_mode) {
AddExportTFLPass(&pm, options_.enable_inliner);
}
mlir::StatusScopedDiagnosticHandler statusHandler(module->getContext(),
/*propagate=*/true);
if (failed(pm.run(*module))) {
return absl::InternalError("conversion error");
}
return absl::OkStatus();
}
std::vector<std::unique_ptr<tac::TargetHardware>>
TacModule::InstantiateBackends() {
std::vector<std::unique_ptr<tac::TargetHardware>> backends;
for (const auto& hardware_name : options_.hardware_backends) {
auto factory = tac::GetTargetHardwareFactory(hardware_name);
backends.emplace_back(factory());
backends.back()->Init();
}
return backends;
}
absl::Status TacModule::Run() {
// Construct all backends.
backends_ = InstantiateBackends();
const_backends_.resize(backends_.size());
for (const auto& backend : backends_)
const_backends_.emplace_back(backend.get());
if (!importer_) {
return absl::Status(absl::StatusCode::kFailedPrecondition,
"Null Importer provided");
}
if (!exporter_) {
return absl::Status(absl::StatusCode::kFailedPrecondition,
"Null Exporter provided");
}
auto module_status = importer_->Import();
if (!module_status.ok()) {
return module_status.status();
}
auto module = module_status->get();
auto* context = module->getContext();
context->appendDialectRegistry(registry_);
context->loadAllAvailableDialects();
// Run TAC passes.
auto status = RunTacPasses(&module, options_.debug_mode);
if (!status.ok()) {
return status;
}
return exporter_->Export(module);
}
void TacModule::RegisterExtraDialects(mlir::DialectRegistry& registry) {
registry.appendTo(registry_);
}
} // namespace tac
} // namespace TFL
} // namespace mlir
tensorflow/compiler/mlir/lite/experimental/tac/tac_module.h
+123
View File
@@ -0,0 +1,123 @@
/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_TAC_TAC_MODULE_H_
#define TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_TAC_TAC_MODULE_H_
#include <memory>
#include <string>
#include <utility>
#include <vector>
#include "absl/status/status.h"
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/Pass/PassManager.h" // from @llvm-project
#include "tensorflow/compiler/mlir/lite/experimental/tac/hardwares/target_hardware.h"
#include "tensorflow/compiler/mlir/lite/experimental/tac/tac_importer_exporter.h"
namespace mlir {
namespace TFL {
namespace tac {
// Main class for using Target Aware Conversion (TAC).
// To run TAC:
// 1) users should create object form this class, with desired options
// (TacModule::Options).
// 2) Use SetImporter/SetExporter to the desired importer
// and exporter.
// 3) Call Run()
//
// The module fetches all TargetHardware backends registered in the binary
// and only create TargetHardware requested in Options.
//
// This class is not thread safe.
class TacModule {
public:
// TAC options. Contains knobs to configure TAC as needed.
struct Options {
// List of names for the requested Target hardware.
std::vector<std::string> hardware_backends;
// Debug mode.
// This will output different alternative subgraphs in mlir format for debug
// purpose.
bool debug_mode = false;
// Whether to enable inliner passes or not.
bool enable_inliner = false;
// Whether to legalize ops to TFLite ops before exporting.
bool legalize_to_tflite_ops = false;
};
virtual ~TacModule() = default;
explicit TacModule(const Options& options) : options_(options) {}
void SetImporter(std::unique_ptr<TacImporter> importer) {
importer_ = std::move(importer);
}
void SetExporter(std::unique_ptr<TacExporter> exporter) {
exporter_ = std::move(exporter);
}
// Returns pointer to the TargetHardware that is identified by 'hardware_name'
// Returns NULL If no hardware with this name found.
const tac::TargetHardware* GetTargetHardware(
const std::string& hardware_name) const;
// Runs the TAC workflow, configured as in the options provided during
// construction.
// SetImporter/SetExporter should be called prior to invoking `Run`.
// Returns Status of the Run.
virtual absl::Status Run();
// Returns all available hardware backends registered in this module
// instance.
const std::vector<const tac::TargetHardware*>& GetAvailableHardwares() const {
return const_backends_;
}
// Registers all dialects in 'registry' with the module.
// This to allow clients to register extra dialects required.
void RegisterExtraDialects(mlir::DialectRegistry& registry);
protected:
// Adds TAC passes to the 'pass_manager'.
virtual void AddTACPass(mlir::OpPassManager* pass_manager,
llvm::ArrayRef<std::string> device_specs);
private:
// Runs all TAC passes on the provided module.
absl::Status RunTacPasses(mlir::ModuleOp* module, bool debug_mode = false);
// Create instances of all registered hardwares.
std::vector<std::unique_ptr<tac::TargetHardware>> InstantiateBackends();
std::unique_ptr<TacImporter> importer_;
std::unique_ptr<TacExporter> exporter_;
// Owned list of all target hardware backends.
std::vector<std::unique_ptr<tac::TargetHardware>> backends_;
// Holder for const pointers for the data in 'backends_'
std::vector<const tac::TargetHardware*> const_backends_;
// Extra dialects requested by the user.
mlir::DialectRegistry registry_;
const Options options_;
};
} // namespace tac
} // namespace TFL
} // namespace mlir
#endif // TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_TAC_TAC_MODULE_H_
tensorflow/compiler/mlir/lite/experimental/tac/tac_translate.cc
+148
View File
@@ -0,0 +1,148 @@
/* 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 <memory>
#include <string>
#include <vector>
#include "absl/log/log.h"
#include "absl/status/status.h"
#include "absl/strings/str_split.h"
#include "llvm/ADT/SmallString.h"
#include "llvm/Support/CommandLine.h"
#include "llvm/Support/MemoryBuffer.h"
#include "llvm/Support/Path.h"
#include "llvm/Support/SourceMgr.h"
#include "llvm/Support/ToolOutputFile.h"
#include "llvm/Support/raw_ostream.h"
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/Attributes.h" // from @llvm-project
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "mlir/Parser/Parser.h" // from @llvm-project
#include "mlir/Pass/PassManager.h" // from @llvm-project
#include "mlir/Pass/PassRegistry.h" // from @llvm-project
#include "mlir/Support/FileUtilities.h" // from @llvm-project
#include "mlir/Support/LLVM.h" // from @llvm-project
#include "mlir/Support/LogicalResult.h" // from @llvm-project
#include "mlir/Tools/mlir-translate/Translation.h" // from @llvm-project
#include "mlir/Transforms/Passes.h" // from @llvm-project
#include "tensorflow/compiler/mlir/init_mlir.h"
#include "tensorflow/compiler/mlir/lite/experimental/tac/common/targets.h"
#include "tensorflow/compiler/mlir/lite/experimental/tac/common/utils.h"
#include "tensorflow/compiler/mlir/lite/experimental/tac/execution_metadata_exporter.h"
#include "tensorflow/compiler/mlir/lite/experimental/tac/hardwares/target_hardware.h"
#include "tensorflow/compiler/mlir/lite/experimental/tac/tac_module.h"
#include "tensorflow/compiler/mlir/lite/experimental/tac/tflite_import_export.h"
#include "tensorflow/compiler/mlir/lite/experimental/tac/transforms/passes.h"
#include "tensorflow/compiler/mlir/lite/experimental/tac/utils/utils.h"
#include "tensorflow/compiler/mlir/lite/transforms/passes.h"
#include "tensorflow/compiler/mlir/tensorflow/dialect_registration.h"
#include "tensorflow/compiler/mlir/tensorflow/transforms/passes.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/error_util.h"
#include "tensorflow/core/platform/init_main.h"
using llvm::cl::opt;
// NOLINTNEXTLINE
opt<std::string> input_file_name(llvm::cl::Positional,
llvm::cl::desc("<input file>"),
llvm::cl::init("-"));
// NOLINTNEXTLINE
opt<std::string> output_file_name("o", llvm::cl::desc("<output file>"),
llvm::cl::value_desc("filename"),
llvm::cl::init("-"));
// NOLINTNEXTLINE
opt<bool> input_mlir("input-mlir",
llvm::cl::desc("Input is MLIR rather than FlatBuffer"),
llvm::cl::init(false));
// NOLINTNEXTLINE
opt<bool> output_mlir(
"output-mlir",
llvm::cl::desc(
"Output MLIR rather than FlatBuffer for the generated TFLite model"),
llvm::cl::init(false));
// NOLINTNEXTLINE
opt<bool> inline_subgraphs(
"inline-subgraphs",
llvm::cl::desc("Whether or not to inline all the subgraphs"),
llvm::cl::init(true));
// NOLINTNEXTLINE
opt<std::string> device_specs(
"device-specs", llvm::cl::desc("comma separated list of device specs."),
llvm::cl::init(""));
// NOLINTNEXTLINE
opt<bool> export_runtime_metadata(
"export-runtime-metadata",
llvm::cl::desc("Whether or not to export metadata, if yes, the metadata "
"file will be exported along with the output model"),
llvm::cl::init(false));
namespace {
std::unique_ptr<mlir::TFL::tac::TacImporter> CreateTfLiteImporter() {
mlir::TFL::tac::TfLiteImporter::Options options;
options.file_name = input_file_name;
options.input_mlir = input_mlir;
return std::make_unique<mlir::TFL::tac::TfLiteImporter>(options);
}
std::unique_ptr<mlir::TFL::tac::TacExporter> CreateTfLiteExporter(
const std::vector<std::string>& target_hardware_backends) {
mlir::TFL::tac::TfLiteExporter::Options options;
options.output_mlir = output_mlir;
options.output_file_name = output_file_name;
options.export_runtime_metadata = export_runtime_metadata;
options.target_hardware_backends = target_hardware_backends;
return std::make_unique<mlir::TFL::tac::TfLiteExporter>(options);
}
absl::Status TargetAwareConversionMain() {
std::vector<std::string> device_specs_array =
absl::StrSplit(device_specs, ',', absl::SkipEmpty());
mlir::TFL::tac::TacModule::Options options;
options.hardware_backends = device_specs_array;
options.debug_mode = true;
if (!output_mlir || inline_subgraphs) {
options.debug_mode = false;
}
options.enable_inliner = true;
options.legalize_to_tflite_ops = true;
mlir::TFL::tac::TacModule tac_module(options);
mlir::DialectRegistry registry;
mlir::RegisterAllTensorFlowDialects(registry);
tac_module.RegisterExtraDialects(registry);
tac_module.SetImporter(CreateTfLiteImporter());
tac_module.SetExporter(CreateTfLiteExporter(options.hardware_backends));
return tac_module.Run();
}
} // namespace
int main(int argc, char** argv) {
tensorflow::InitMlir y(&argc, &argv);
llvm::cl::ParseCommandLineOptions(argc, argv, "Target aware conversion\n");
absl::Status status = TargetAwareConversionMain();
if (!status.ok()) {
LOG(ERROR) << status;
}
return 0;
}
tensorflow/compiler/mlir/lite/experimental/tac/tests/BUILD
+26
View File
@@ -0,0 +1,26 @@
load("//tensorflow:tensorflow.default.bzl", "filegroup")
load("//tensorflow/compiler/mlir:glob_lit_test.bzl", "glob_lit_tests")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:LICENSE"],
licenses = ["notice"],
)
glob_lit_tests(
name = "all_tests",
data = [":test_utilities"],
driver = "@llvm-project//mlir:run_lit.sh",
test_file_exts = ["mlir"],
)
# Bundle together all of the test utilities that are used by tests.
filegroup(
name = "test_utilities",
testonly = True,
data = [
"//tensorflow/compiler/mlir/lite/experimental/common:outline_operations",
"//tensorflow/compiler/mlir/lite/experimental/tac:tac-opt-all-backends",
"@llvm-project//llvm:FileCheck",
"@llvm-project//llvm:not",
],
)
tensorflow/compiler/mlir/lite/experimental/tac/tests/compute-cost.mlir
+65
View File
@@ -0,0 +1,65 @@
// 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: tac-opt-all-backends -tfl-compute-cost %s -split-input-file -verify-diagnostics | FileCheck %s
// CHECK: tac.cost = 7.864320e+05
func.func @func_0_CPU(%arg0: tensor<256x32x32x3xf32>, %arg1: tensor<256x32x32x3xf32>) -> tensor<256x32x32x3xf32> attributes {tac.device = "CPU", tac.interface_name = "func_0"} {
%0 = "tfl.add"(%arg0, %arg1) {fused_activation_function = "RELU", tac.device = "CPU"} : (tensor<256x32x32x3xf32>, tensor<256x32x32x3xf32>) -> tensor<256x32x32x3xf32>
func.return %0 : tensor<256x32x32x3xf32>
}
// CHECK: tac.cost = 157286.4
func.func @func_0_GPU(%arg0: tensor<256x32x32x3xf32>, %arg1: tensor<256x32x32x3xf32>) -> tensor<256x32x32x3xf32> attributes {tac.device = "GPU", tac.interface_name = "func_0"} {
%0 = "tfl.add"(%arg0, %arg1) {fused_activation_function = "RELU", tac.device = "GPU"} : (tensor<256x32x32x3xf32>, tensor<256x32x32x3xf32>) -> tensor<256x32x32x3xf32>
func.return %0 : tensor<256x32x32x3xf32>
}
// -----
// CHECK: tac.cost = 1.000000e+03
func.func @func_0_CPU(%arg0: tensor<10x10x10xf32>, %arg1: tensor<10xf32>) -> tensor<10x10x10xf32> attributes {tac.device = "CPU", tac.interface_name = "func_0"} {
%0 = "tfl.add"(%arg0, %arg1) {fused_activation_function = "RELU", tac.device = "CPU"} : (tensor<10x10x10xf32>, tensor<10xf32>) -> tensor<10x10x10xf32>
func.return %0 : tensor<10x10x10xf32>
}
// -----
// CHECK: tac.cost = 2.000000e+02
func.func @func_0_GPU(%arg0: tensor<10x10x10xf32>, %arg1: tensor<f32>) -> tensor<10x10x10xf32> attributes {tac.device = "GPU", tac.interface_name = "func_0"} {
%0 = "tfl.add"(%arg0, %arg1) {fused_activation_function = "RELU", tac.device = "GPU"} : (tensor<10x10x10xf32>, tensor<f32>) -> tensor<10x10x10xf32>
func.return %0 : tensor<10x10x10xf32>
}
// -----
// CHECK: tac.cost = 2.000000e+03
func.func @func_0_CPU(%arg0: tensor<10x10x10xf32>, %arg1: tensor<10xf32>) -> tensor<10x10x10xf32> attributes {tac.device = "CPU", tac.interface_name = "func_0"} {
%0 = "tfl.add"(%arg0, %arg1) {fused_activation_function = "RELU", tac.device = "CPU"} : (tensor<10x10x10xf32>, tensor<10xf32>) -> tensor<10x10x10xf32>
%1 = "tfl.mul"(%0, %arg1) {fused_activation_function = "RELU", tac.device = "CPU"} : (tensor<10x10x10xf32>, tensor<10xf32>) -> tensor<10x10x10xf32>
func.return %1 : tensor<10x10x10xf32>
}
// -----
// CHECK: tac.cost = 0x4B673001
func.func @quantize_ops_CPU_QUANTIZED_INT8(%arg0: tensor<384x512x!quant.uniform<i8:f32, 0.1>>, %arg1: tensor<128x512x!quant.uniform<i8<-127:127>:f32, 0.1>>, %arg2: tensor<128x!quant.uniform<i8:f32, 0.2:-128>>, %arg3: tensor<128x!quant.uniform<i8:f32, 0.2:-4>>) -> tensor<1x384x128x!quant.uniform<i8:f32, 0.3:-3>> attributes {tac.device = "CPU", tac.inference_type = "QUANTIZED_INT8", tac.interface_name = "quantize_ops"} {
%0 = "tfl.pseudo_qconst"() {qtype = tensor<128x!quant.uniform<i32:f32, 0.7>>, value = dense<0> : tensor<128xi32>} : () -> tensor<128x!quant.uniform<i32:f32, 0.7>>
%1 = "tfl.fully_connected"(%arg0, %arg1, %0) {tac.device = "CPU", tac.inference_type = "QUANTIZED_INT8", fused_activation_function = "NONE", keep_num_dims = false, weights_format = "DEFAULT"} : (tensor<384x512x!quant.uniform<i8:f32, 0.1>>, tensor<128x512x!quant.uniform<i8<-127:127>:f32, 0.1>>, tensor<128x!quant.uniform<i32:f32, 0.7>>) -> tensor<384x128x!quant.uniform<i8:f32, 0.9:-4>>
%2 = "tfl.pseudo_const"() {value = dense<[1, 384, 128]> : tensor<3xi32>} : () -> tensor<3xi32>
%3 = "tfl.reshape"(%1, %2) {tac.device = "CPU", tac.inference_type = "QUANTIZED_INT8"} : (tensor<384x128x!quant.uniform<i8:f32, 0.9:-4>>, tensor<3xi32>) -> tensor<1x384x128x!quant.uniform<i8:f32, 0.9:-4>>
%4 = "tfl.mul"(%3, %arg2) {tac.device = "CPU", tac.inference_type = "QUANTIZED_INT8", fused_activation_function = "NONE"} : (tensor<1x384x128x!quant.uniform<i8:f32, 0.9:-4>>, tensor<128x!quant.uniform<i8:f32, 0.2:-128>>) -> tensor<1x384x128x!quant.uniform<i8:f32, 0.3:3>>
%5 = "tfl.add"(%4, %arg3) {tac.device = "CPU", tac.inference_type = "QUANTIZED_INT8", fused_activation_function = "NONE"} : (tensor<1x384x128x!quant.uniform<i8:f32, 0.3:3>>, tensor<128x!quant.uniform<i8:f32, 0.2:-4>>) -> tensor<1x384x128x!quant.uniform<i8:f32, 0.3:-3>>
func.return %5 : tensor<1x384x128x!quant.uniform<i8:f32, 0.3:-3>>
}
tensorflow/compiler/mlir/lite/experimental/tac/tests/device-transform-gpu.mlir
+239
View File
@@ -0,0 +1,239 @@
// 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: tac-opt-all-backends -tfl-device-transform-gpu %s -split-input-file -verify-diagnostics | FileCheck %s
func.func @pack(%arg0: tensor<1xf32>, %arg1: tensor<1xf32>) -> tensor<2x1xf32> {
%0 = "tfl.pack"(%arg0, %arg1) {axis = 0 : i32, values_count = 2 : i32} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2x1xf32>
func.return %0 : tensor<2x1xf32>
}
// CHECK: func @pack(%[[VAL_0:.*]]: tensor<1xf32>, %[[VAL_1:.*]]: tensor<1xf32>) -> tensor<2x1xf32> {
// CHECK-DAG: %[[VAL_2:.*]] = "tfl.pseudo_const"{{.*}}dense<1> : tensor<4xi32>
// CHECK-DAG: %[[VAL_3:.*]] = "tfl.pseudo_const"{{.*}}dense<2> : tensor<1xi32>
// CHECK-DAG: %[[VAL_4:.*]] = "tfl.pseudo_const"{{.*}}dense<[2, 1]> : tensor<2xi32>
// CHECK: %[[VAL_5:.*]] = "tfl.reshape"(%[[VAL_0]], %[[VAL_2]]) : (tensor<1xf32>, tensor<4xi32>) -> tensor<1x1x1x1xf32>
// CHECK: %[[VAL_6:.*]] = "tfl.reshape"(%[[VAL_1]], %[[VAL_2]]) : (tensor<1xf32>, tensor<4xi32>) -> tensor<1x1x1x1xf32>
// CHECK: %[[VAL_7:.*]] = "tfl.concatenation"(%[[VAL_5]], %[[VAL_6]]) <{axis = 3 : i32, fused_activation_function = "NONE"}> : (tensor<1x1x1x1xf32>, tensor<1x1x1x1xf32>) -> tensor<1x1x1x2xf32>
// CHECK: %[[VAL_8:.*]] = "tfl.reshape"(%[[VAL_7]], %[[VAL_3]]) : (tensor<1x1x1x2xf32>, tensor<1xi32>) -> tensor<2xf32>
// CHECK: %[[VAL_9:.*]] = "tfl.reshape"(%[[VAL_8]], %[[VAL_4]]) : (tensor<2xf32>, tensor<2xi32>) -> tensor<2x1xf32>
// CHECK: return %[[VAL_9]] : tensor<2x1xf32>
// CHECK: }
// -----
// CHECK-LABEL: @avoidPackToConcatenationOnUnknownShapes
func.func @avoidPackToConcatenationOnUnknownShapes(%arg0: tensor<?x?xf32>, %arg1: tensor<?x?xf32>) -> tensor<?x?x1xf32> {
%0 = "tfl.pack"(%arg0, %arg1) {axis = 0 : i32, values_count = 2 : i32} : (tensor<?x?xf32>, tensor<?x?xf32>) -> tensor<?x?x1xf32>
func.return %0 : tensor<?x?x1xf32>
}
// CHECK-NOT: "tfl.reshape"
// CHECK: "tfl.pack"
// -----
// CHECK-LABEL: @avoidPackToConcatenationOnUnknownRank
func.func @avoidPackToConcatenationOnUnknownRank(%arg0: tensor<*xf32>, %arg1: tensor<*xf32>) -> tensor<*xf32> {
%0 = "tfl.pack"(%arg0, %arg1) {axis = 0 : i32, values_count = 2 : i32} : (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32>
func.return %0 : tensor<*xf32>
}
// CHECK-NOT: "tfl.reshape"
// CHECK: "tfl.pack"
// -----
func.func @squaredDifference(%arg0: tensor<4xf32>, %arg1: tensor<4xf32>) -> tensor<4xf32> {
%0 = "tfl.squared_difference"(%arg0, %arg1) : (tensor<4xf32>, tensor<4xf32>) -> tensor<4xf32>
func.return %0 : tensor<4xf32>
}
// CHECK: func @squaredDifference(%arg0: tensor<4xf32>, %arg1: tensor<4xf32>) -> tensor<4xf32> {
// CHECK: %0 = "tf.Sub"(%arg0, %arg1) : (tensor<4xf32>, tensor<4xf32>) -> tensor<4xf32>
// CHECK: %1 = "tf.Mul"(%0, %0) : (tensor<4xf32>, tensor<4xf32>) -> tensor<4xf32>
// CHECK: return %1 : tensor<4xf32>
// CHECK: }
// -----
func.func @unrollSplit(%arg0: tensor<i32>, %arg1: tensor<1x8x8x1024xf32>) -> (tensor<1x8x8x256xf32>, tensor<1x8x8x256xf32>, tensor<1x8x8x256xf32>) {
%0:4 = "tfl.split"(%arg0, %arg1) {num_splits = 4 : i32, tac.device = "CPU"} : (tensor<i32>, tensor<1x8x8x1024xf32>) -> (tensor<1x8x8x256xf32>, tensor<1x8x8x256xf32>, tensor<1x8x8x256xf32>, tensor<1x8x8x256xf32>)
func.return %0#0, %0#1, %0#3 : tensor<1x8x8x256xf32>, tensor<1x8x8x256xf32>, tensor<1x8x8x256xf32>
}
// CHECK: func @unrollSplit([[VAL_0:%.*]]: tensor<i32>, [[VAL_1:%.*]]: tensor<1x8x8x1024xf32>) -> (tensor<1x8x8x256xf32>, tensor<1x8x8x256xf32>, tensor<1x8x8x256xf32>) {
// CHECK-DAG: [[VAL_2:%.*]] = "tfl.pseudo_const"{{.*}}dense<0> : tensor<4xi32>
// CHECK-DAG: [[VAL_3:%.*]] = "tfl.pseudo_const"{{.*}}dense<[0, 0, 0, 256]> : tensor<4xi32>
// CHECK-DAG: [[VAL_4:%.*]] = "tfl.pseudo_const"{{.*}}dense<[0, 0, 0, 768]> : tensor<4xi32>
// CHECK-DAG: [[VAL_5:%.*]] = "tfl.pseudo_const"{{.*}}dense<[-1, -1, -1, 256]> : tensor<4xi32>
// CHECK: [[VAL_6:%.*]] = "tfl.slice"([[VAL_1]], [[VAL_2]], [[VAL_5]]) : (tensor<1x8x8x1024xf32>, tensor<4xi32>, tensor<4xi32>) -> tensor<1x8x8x256xf32>
// CHECK: [[VAL_7:%.*]] = "tfl.slice"([[VAL_1]], [[VAL_3]], [[VAL_5]]) : (tensor<1x8x8x1024xf32>, tensor<4xi32>, tensor<4xi32>) -> tensor<1x8x8x256xf32>
// CHECK: [[VAL_8:%.*]] = "tfl.slice"([[VAL_1]], [[VAL_4]], [[VAL_5]]) : (tensor<1x8x8x1024xf32>, tensor<4xi32>, tensor<4xi32>) -> tensor<1x8x8x256xf32>
// CHECK: return [[VAL_6]], [[VAL_7]], [[VAL_8]] : tensor<1x8x8x256xf32>, tensor<1x8x8x256xf32>, tensor<1x8x8x256xf32>
// CHECK: }
// -----
func.func @unrollSplitUnknownRankResults(%arg0: tensor<i32>, %arg1: tensor<1x8x8x1024xf32>) -> (tensor<*xf32>, tensor<*xf32>, tensor<*xf32>) {
%0:4 = "tfl.split"(%arg0, %arg1) {num_splits = 4 : i32, tac.device = "CPU"} : (tensor<i32>, tensor<1x8x8x1024xf32>) -> (tensor<*xf32>, tensor<*xf32>, tensor<*xf32>, tensor<*xf32>)
func.return %0#0, %0#1, %0#3 : tensor<*xf32>, tensor<*xf32>, tensor<*xf32>
}
// CHECK-LABEL: @unrollSplitUnknownRankResults
// CHECK-NOT: "tfl.slice"
// CHECK: "tfl.split"
// -----
func.func @unrollSplitV(%arg0: tensor<?x13x13x85xf32>) -> (tensor<?x13x13x2xf32>, tensor<?x13x13x2xf32>, tensor<?x13x13x81xf32>) {
%0 = "tfl.pseudo_const"() {value = dense<[2, 2, 81]> : tensor<3xi32>} : () -> tensor<3xi32>
%1 = "tfl.pseudo_const"() {value = dense<-1> : tensor<i32>} : () -> tensor<i32>
%2:3 = "tfl.split_v"(%arg0, %0, %1) {num_splits = 3 : i32} : (tensor<?x13x13x85xf32>, tensor<3xi32>, tensor<i32>) -> (tensor<?x13x13x2xf32>, tensor<?x13x13x2xf32>, tensor<?x13x13x81xf32>)
func.return %2#0, %2#1, %2#2 : tensor<?x13x13x2xf32>, tensor<?x13x13x2xf32>, tensor<?x13x13x81xf32>
}
// CHECK: func @unrollSplitV(%[[VAL_0:.*]]: tensor<?x13x13x85xf32>) -> (tensor<?x13x13x2xf32>, tensor<?x13x13x2xf32>, tensor<?x13x13x81xf32>) {
// CHECK-DAG: %[[VAL_1:.*]] = "tfl.pseudo_const"{{.*}}dense<0> : tensor<4xi32>
// CHECK-DAG: %[[VAL_2:.*]] = "tfl.pseudo_const"{{.*}}dense<[0, 0, 0, 2]> : tensor<4xi32>
// CHECK-DAG: %[[VAL_3:.*]] = "tfl.pseudo_const"{{.*}}dense<[-1, -1, -1, 2]> : tensor<4xi32>
// CHECK-DAG: %[[VAL_4:.*]] = "tfl.pseudo_const"{{.*}}dense<[0, 0, 0, 4]> : tensor<4xi32>
// CHECK-DAG: %[[VAL_5:.*]] = "tfl.pseudo_const"{{.*}}dense<[-1, -1, -1, 81]> : tensor<4xi32>
// CHECK: %[[VAL_6:.*]] = "tfl.slice"(%[[VAL_0]], %[[VAL_1]], %[[VAL_3]]) : (tensor<?x13x13x85xf32>, tensor<4xi32>, tensor<4xi32>) -> tensor<?x13x13x2xf32>
// CHECK: %[[VAL_7:.*]] = "tfl.slice"(%[[VAL_0]], %[[VAL_2]], %[[VAL_3]]) : (tensor<?x13x13x85xf32>, tensor<4xi32>, tensor<4xi32>) -> tensor<?x13x13x2xf32>
// CHECK: %[[VAL_8:.*]] = "tfl.slice"(%[[VAL_0]], %[[VAL_4]], %[[VAL_5]]) : (tensor<?x13x13x85xf32>, tensor<4xi32>, tensor<4xi32>) -> tensor<?x13x13x81xf32>
// CHECK: return %[[VAL_6]], %[[VAL_7]], %[[VAL_8]] : tensor<?x13x13x2xf32>, tensor<?x13x13x2xf32>, tensor<?x13x13x81xf32>
// CHECK: }
// -----
func.func @unrollSplitVUnknownRankResults(%arg0: tensor<?x13x13x85xf32>) -> (tensor<*xf32>, tensor<*xf32>, tensor<*xf32>) {
%0 = "tfl.pseudo_const"() {value = dense<[2, 2, 81]> : tensor<3xi32>} : () -> tensor<3xi32>
%1 = "tfl.pseudo_const"() {value = dense<-1> : tensor<i32>} : () -> tensor<i32>
%2:3 = "tfl.split_v"(%arg0, %0, %1) {num_splits = 3 : i32} : (tensor<?x13x13x85xf32>, tensor<3xi32>, tensor<i32>) -> (tensor<*xf32>, tensor<*xf32>, tensor<*xf32>)
func.return %2#0, %2#1, %2#2 : tensor<*xf32>, tensor<*xf32>, tensor<*xf32>
}
// CHECK-LABEL: @unrollSplitVUnknownRankResults
// CHECK-NOT: "tfl.slice"
// CHECK: "tfl.split_v"
// -----
func.func @sub(%arg0: tensor<1x384x384x3xf32>, %arg1: tensor<3xf32>) -> tensor<1x384x384x3xf32> {
%0 = "tfl.sub"(%arg0, %arg1) {fused_activation_function = "NONE"} : (tensor<1x384x384x3xf32>, tensor<3xf32>) -> tensor<1x384x384x3xf32>
func.return %0 : tensor<1x384x384x3xf32>
}
// CHECK: func @sub(%[[VAL_0:.*]]: tensor<1x384x384x3xf32>, %[[VAL_1:.*]]: tensor<3xf32>) -> tensor<1x384x384x3xf32> {
// CHECK: %[[VAL_2:.*]] = arith.constant dense<-1.000000e+00> : tensor<f32>
// CHECK: %[[VAL_3:.*]] = tfl.mul(%[[VAL_1]], %[[VAL_2]]) <{fused_activation_function = "NONE"}> : (tensor<3xf32>, tensor<f32>) -> tensor<3xf32>
// CHECK: %[[VAL_4:.*]] = tfl.add(%[[VAL_0]], %[[VAL_3]]) <{fused_activation_function = "NONE"}> : (tensor<1x384x384x3xf32>, tensor<3xf32>) -> tensor<1x384x384x3xf32>
// CHECK: return %[[VAL_4]] : tensor<1x384x384x3xf32>
// CHECK: }
// -----
func.func @ensureBiasForConv2d(%arg0: tensor<128x32x32x3xf32>, %arg1: tensor<32x1x1x3xf32>) -> tensor<128x32x32x32xf32> {
%cst = "tfl.no_value"() {value = unit} : () -> none
%0 = "tfl.conv_2d"(%arg0, %arg1, %cst) {dilation_h_factor = 1 : i32, dilation_w_factor = 1 : i32, fused_activation_function = "NONE", padding = "VALID", stride_h = 1 : i32, stride_w = 1 : i32} : (tensor<128x32x32x3xf32>, tensor<32x1x1x3xf32>, none) -> tensor<128x32x32x32xf32>
func.return %0 : tensor<128x32x32x32xf32>
}
// CHECK: func @ensureBiasForConv2d(%[[VAL_0:.*]]: tensor<128x32x32x3xf32>, %[[VAL_1:.*]]: tensor<32x1x1x3xf32>) -> tensor<128x32x32x32xf32> {
// CHECK: %[[VAL_2:.*]] = "tfl.pseudo_const"{{.*}}dense<0.000000e+00> : tensor<32xf32>
// CHECK: %[[VAL_3:.*]] = "tfl.conv_2d"(%[[VAL_0]], %[[VAL_1]], %[[VAL_2]]) <{dilation_h_factor = 1 : i32, dilation_w_factor = 1 : i32, fused_activation_function = "NONE", padding = "VALID", stride_h = 1 : i32, stride_w = 1 : i32}> : (tensor<128x32x32x3xf32>, tensor<32x1x1x3xf32>, tensor<32xf32>) -> tensor<128x32x32x32xf32>
// CHECK: return %[[VAL_3]] : tensor<128x32x32x32xf32>
// CHECK: }
// -----
func.func @padSliceTo4D(%arg0: tensor<4x384x32xf32>) -> tensor<1x384x32xf32> {
%0 = "tfl.pseudo_const"() {value = dense<0> : tensor<3xi32>} : () -> tensor<3xi32>
%1 = "tfl.pseudo_const"() {value = dense<[1, 384, 32]> : tensor<3xi32>} : () -> tensor<3xi32>
%2 = "tfl.slice"(%arg0, %0, %1) : (tensor<4x384x32xf32>, tensor<3xi32>, tensor<3xi32>) -> tensor<1x384x32xf32>
func.return %2 : tensor<1x384x32xf32>
}
// CHECK: func @padSliceTo4D(%[[VAL_0:.*]]: tensor<4x384x32xf32>) -> tensor<1x384x32xf32> {
// CHECK-DAG: %[[VAL_1:.*]] = "tf.Const"() <{value = dense<0> : tensor<4xi32>}> : () -> tensor<4xi32>
// CHECK-DAG: %[[VAL_2:.*]] = "tf.Const"() <{value = dense<[1, 1, 384, 32]> : tensor<4xi32>}> : () -> tensor<4xi32>
// CHECK-DAG: %[[VAL_3:.*]] = "tfl.pseudo_const"{{.*}}dense<[1, 4, 384, 32]> : tensor<4xi32>
// CHECK-DAG: %[[VAL_4:.*]] = "tfl.pseudo_const"() <{value = dense<[1, 384, 32]> : tensor<3xi32>
// CHECK: %[[VAL_5:.*]] = "tfl.reshape"(%[[VAL_0]], %[[VAL_3]]) : (tensor<4x384x32xf32>, tensor<4xi32>) -> tensor<1x4x384x32xf32>
// CHECK: %[[VAL_6:.*]] = "tfl.slice"(%[[VAL_5]], %[[VAL_1]], %[[VAL_2]]) : (tensor<1x4x384x32xf32>, tensor<4xi32>, tensor<4xi32>) -> tensor<1x1x384x32xf32>
// CHECK: %[[VAL_7:.*]] = "tfl.reshape"(%[[VAL_6]], %[[VAL_4]]) : (tensor<1x1x384x32xf32>, tensor<3xi32>) -> tensor<1x384x32xf32>
// CHECK: return %[[VAL_7]] : tensor<1x384x32xf32>
// CHECK: }
// -----
// CHECK-LABEL: @avoidPadSliceTo4DOnUnknownOutputShape
func.func @avoidPadSliceTo4DOnUnknownOutputShape(%arg0: tensor<4x384x32xf32>) -> tensor<1x?x?xf32> {
%0 = "tfl.pseudo_const"() {value = dense<0> : tensor<3xi32>} : () -> tensor<3xi32>
%1 = "tfl.pseudo_const"() {value = dense<[1, 384, 32]> : tensor<3xi32>} : () -> tensor<3xi32>
%2 = "tfl.slice"(%arg0, %0, %1) : (tensor<4x384x32xf32>, tensor<3xi32>, tensor<3xi32>) -> tensor<1x?x?xf32>
func.return %2 : tensor<1x?x?xf32>
}
// CHECK-NOT: "tfl.reshape"
// CHECK: "tfl.slice"
// -----
func.func @fullyConnectedToConv(%arg0: tensor<384x384xf32>, %arg1: tensor<512x384xf32>, %arg2: tensor<512xf32>) -> tensor<384x512xf32> {
%0 = "tfl.fully_connected"(%arg0, %arg1, %arg2) {fused_activation_function = "NONE", keep_num_dims = false, weights_format = "DEFAULT"} : (tensor<384x384xf32>, tensor<512x384xf32>, tensor<512xf32>) -> tensor<384x512xf32>
func.return %0: tensor<384x512xf32>
}
// CHECK: func @fullyConnectedToConv(%[[VAL_0:.*]]: tensor<384x384xf32>, %[[VAL_1:.*]]: tensor<512x384xf32>, %[[VAL_2:.*]]: tensor<512xf32>) -> tensor<384x512xf32> {
// CHECK-DAG: %[[VAL_3:.*]] = "tfl.pseudo_const"{{.*}}dense<[1, 1, 384, 384]> : tensor<4xi32>
// CHECK-DAG: %[[VAL_4:.*]] = "tfl.pseudo_const"{{.*}}dense<[512, 1, 1, 384]> : tensor<4xi32>
// CHECK-DAG: %[[VAL_5:.*]] = "tfl.pseudo_const"{{.*}}dense<[384, 512]> : tensor<2xi32>
// CHECK: %[[VAL_6:.*]] = "tfl.reshape"(%[[VAL_0]], %[[VAL_3]]) : (tensor<384x384xf32>, tensor<4xi32>) -> tensor<1x1x384x384xf32>
// CHECK: %[[VAL_7:.*]] = "tfl.reshape"(%[[VAL_1]], %[[VAL_4]]) : (tensor<512x384xf32>, tensor<4xi32>) -> tensor<512x1x1x384xf32>
// CHECK: %[[VAL_8:.*]] = "tfl.conv_2d"(%[[VAL_6]], %[[VAL_7]], %[[VAL_2]]) <{dilation_h_factor = 1 : i32, dilation_w_factor = 1 : i32, fused_activation_function = "NONE", padding = "VALID", stride_h = 1 : i32, stride_w = 1 : i32}> : (tensor<1x1x384x384xf32>, tensor<512x1x1x384xf32>, tensor<512xf32>) -> tensor<1x1x384x512xf32>
// CHECK: %[[VAL_9:.*]] = "tfl.reshape"(%[[VAL_8]], %[[VAL_5]]) : (tensor<1x1x384x512xf32>, tensor<2xi32>) -> tensor<384x512xf32>
// CHECK: return %[[VAL_9]] : tensor<384x512xf32>
// CHECK: }
// -----
func.func @padConcatTo4D(%arg0: tensor<384x384xf32>, %arg1: tensor<384x384xf32>, %arg2: tensor<384x384xf32>, %arg3: tensor<384x384xf32>) -> tensor<1536x384xf32> {
%0 = "tfl.concatenation"(%arg0, %arg1, %arg2, %arg3) {axis = 0 : i32, fused_activation_function = "NONE"} : (tensor<384x384xf32>, tensor<384x384xf32>, tensor<384x384xf32>, tensor<384x384xf32>) -> tensor<1536x384xf32>
func.return %0: tensor<1536x384xf32>
}
// CHECK: func @padConcatTo4D(%[[VAL_0:.*]]: tensor<384x384xf32>, %[[VAL_1:.*]]: tensor<384x384xf32>, %[[VAL_2:.*]]: tensor<384x384xf32>, %[[VAL_3:.*]]: tensor<384x384xf32>) -> tensor<1536x384xf32> {
// CHECK-DAG: %[[VAL_4:.*]] = "tfl.pseudo_const"{{.*}}dense<[1, 1, 384, 384]> : tensor<4xi32>
// CHECK-DAG: %[[VAL_5:.*]] = "tfl.pseudo_const"{{.*}}dense<[1536, 384]> : tensor<2xi32>
// CHECK: %[[VAL_6:.*]] = "tfl.reshape"(%[[VAL_0]], %[[VAL_4]]) : (tensor<384x384xf32>, tensor<4xi32>) -> tensor<1x1x384x384xf32>
// CHECK: %[[VAL_7:.*]] = "tfl.reshape"(%[[VAL_1]], %[[VAL_4]]) : (tensor<384x384xf32>, tensor<4xi32>) -> tensor<1x1x384x384xf32>
// CHECK: %[[VAL_8:.*]] = "tfl.reshape"(%[[VAL_2]], %[[VAL_4]]) : (tensor<384x384xf32>, tensor<4xi32>) -> tensor<1x1x384x384xf32>
// CHECK: %[[VAL_9:.*]] = "tfl.reshape"(%[[VAL_3]], %[[VAL_4]]) : (tensor<384x384xf32>, tensor<4xi32>) -> tensor<1x1x384x384xf32>
// CHECK: %[[VAL_10:.*]] = "tfl.concatenation"(%[[VAL_6]], %[[VAL_7]], %[[VAL_8]], %[[VAL_9]]) <{axis = 2 : i32, fused_activation_function = "NONE"}> : (tensor<1x1x384x384xf32>, tensor<1x1x384x384xf32>, tensor<1x1x384x384xf32>, tensor<1x1x384x384xf32>) -> tensor<1x1x1536x384xf32>
// CHECK: %[[VAL_11:.*]] = "tfl.reshape"(%[[VAL_10]], %[[VAL_5]]) : (tensor<1x1x1536x384xf32>, tensor<2xi32>) -> tensor<1536x384xf32>
// CHECK: return %[[VAL_11]] : tensor<1536x384xf32>
// CHECK: }
// -----
// CHECK-LABEL: @avoidPadConcatTo4DOnUnknownOutputShape
func.func @avoidPadConcatTo4DOnUnknownOutputShape(%arg0: tensor<384x384xf32>, %arg1: tensor<384x384xf32>, %arg2: tensor<384x384xf32>, %arg3: tensor<384x384xf32>) -> tensor<?x?xf32> {
%0 = "tfl.concatenation"(%arg0, %arg1, %arg2, %arg3) {axis = 0 : i32, fused_activation_function = "NONE"} : (tensor<384x384xf32>, tensor<384x384xf32>, tensor<384x384xf32>, tensor<384x384xf32>) -> tensor<?x?xf32>
func.return %0: tensor<?x?xf32>
}
// CHECK-NOT: "tfl.reshape"
// CHECK: "tfl.concatenation"
tensorflow/compiler/mlir/lite/experimental/tac/tests/device-transform-nnapi.mlir
+72
View File
@@ -0,0 +1,72 @@
// 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: tac-opt-all-backends -tfl-device-transform-nnapi %s -split-input-file -verify-diagnostics | FileCheck %s
func.func @mean_4d_keepdim(%arg0: tensor<1x48x48x512xf32>) -> tensor<1x1x1x512xf32> {
%cst = arith.constant dense<[1, 2]> : tensor<2xi32>
%0 = "tfl.mean"(%arg0, %cst) {keep_dims = true} : (tensor<1x48x48x512xf32>, tensor<2xi32>) -> tensor<1x1x1x512xf32>
func.return %0 : tensor<1x1x1x512xf32>
}
// CHECK: func @mean_4d_keepdim([[VAL_0:%.*]]: tensor<1x48x48x512xf32>) -> tensor<1x1x1x512xf32> {
// CHECK: [[VAL_1:%.*]] = "tfl.average_pool_2d"([[VAL_0]]) <{filter_height = 48 : i32, filter_width = 48 : i32, fused_activation_function = "NONE", padding = "VALID", stride_h = 1 : i32, stride_w = 1 : i32}> : (tensor<1x48x48x512xf32>) -> tensor<1x1x1x512xf32>
// CHECK: return [[VAL_1]] : tensor<1x1x1x512xf32>
// CHECK: }
// -----
func.func @mean_4d_no_keepdim(%arg0: tensor<1x48x48x512xf32>) -> tensor<1x512xf32> {
%cst = arith.constant dense<[1, 2]> : tensor<2xi32>
%0 = "tfl.mean"(%arg0, %cst) {keep_dims = false} : (tensor<1x48x48x512xf32>, tensor<2xi32>) -> tensor<1x512xf32>
func.return %0 : tensor<1x512xf32>
}
// CHECK: func @mean_4d_no_keepdim([[VAL_0:%.*]]: tensor<1x48x48x512xf32>) -> tensor<1x512xf32> {
// CHECK: [[VAL_1:%.*]] = "tfl.pseudo_const"(){{.*}}dense<[1, 512]> : tensor<2xi32>
// CHECK: [[VAL_2:%.*]] = "tfl.average_pool_2d"([[VAL_0]]) <{filter_height = 48 : i32, filter_width = 48 : i32, fused_activation_function = "NONE", padding = "VALID", stride_h = 1 : i32, stride_w = 1 : i32}> : (tensor<1x48x48x512xf32>) -> tensor<1x1x1x512xf32>
// CHECK: [[VAL_3:%.*]] = "tfl.reshape"([[VAL_2]], [[VAL_1]]) : (tensor<1x1x1x512xf32>, tensor<2xi32>) -> tensor<1x512xf32>
// CHECK: return [[VAL_3]] : tensor<1x512xf32>
// CHECK: }
// -----
func.func @mean_quant_same_scale(%arg0: tensor<?x7x7x2048x!quant.uniform<i8:f32, 0.6:-128>>) -> tensor<?x2048x!quant.uniform<i8:f32, 0.6:-128>> {
%0 = "tfl.pseudo_const"() {value = dense<[1, 2]> : tensor<2xi32>} : () -> tensor<2xi32>
%1 = "tfl.mean"(%arg0, %0) {keep_dims = false} : (tensor<?x7x7x2048x!quant.uniform<i8:f32, 0.6:-128>>, tensor<2xi32>) -> tensor<?x2048x!quant.uniform<i8:f32, 0.6:-128>>
func.return %1 : tensor<?x2048x!quant.uniform<i8:f32, 0.6:-128>>
}
// CHECK: func @mean_quant_same_scale(%[[VAL_0:.*]]: tensor<?x7x7x2048x!quant.uniform<i8:f32, 6.000000e-01:-128>>) -> tensor<?x2048x!quant.uniform<i8:f32, 6.000000e-01:-128>> {
// CHECK: %[[VAL_1:.*]] = "tfl.pseudo_const"(){{.*}}dense<[-1, 2048]> : tensor<2xi32>
// CHECK: %[[VAL_2:.*]] = "tfl.average_pool_2d"(%[[VAL_0]]) <{filter_height = 7 : i32, filter_width = 7 : i32, fused_activation_function = "NONE", padding = "VALID", stride_h = 1 : i32, stride_w = 1 : i32}> : (tensor<?x7x7x2048x!quant.uniform<i8:f32, 6.000000e-01:-128>>) -> tensor<?x1x1x2048x!quant.uniform<i8:f32, 6.000000e-01:-128>>
// CHECK: %[[VAL_3:.*]] = "tfl.reshape"(%[[VAL_2]], %[[VAL_1]]) : (tensor<?x1x1x2048x!quant.uniform<i8:f32, 6.000000e-01:-128>>, tensor<2xi32>) -> tensor<?x2048x!quant.uniform<i8:f32, 6.000000e-01:-128>>
// CHECK: return %[[VAL_3]] : tensor<?x2048x!quant.uniform<i8:f32, 6.000000e-01:-128>>
// CHECK: }
// -----
func.func @mean_quant_different_scales(%arg0: tensor<?x7x7x2048x!quant.uniform<i8:f32, 0.6:-128>>) -> tensor<?x2048x!quant.uniform<i8:f32, 0.9:-128>> {
%0 = "tfl.pseudo_const"() {value = dense<[1, 2]> : tensor<2xi32>} : () -> tensor<2xi32>
%1 = "tfl.mean"(%arg0, %0) {keep_dims = false} : (tensor<?x7x7x2048x!quant.uniform<i8:f32, 0.6:-128>>, tensor<2xi32>) -> tensor<?x2048x!quant.uniform<i8:f32, 0.9:-128>>
func.return %1 : tensor<?x2048x!quant.uniform<i8:f32, 0.9:-128>>
}
// CHECK: func @mean_quant_different_scales(%[[VAL_0:.*]]: tensor<?x7x7x2048x!quant.uniform<i8:f32, 6.000000e-01:-128>>) -> tensor<?x2048x!quant.uniform<i8:f32, 9.000000e-01:-128>> {
// CHECK: %[[VAL_1:.*]] = "tfl.pseudo_const"(){{.*}}dense<[-1, 2048]> : tensor<2xi32>
// CHECK: %[[VAL_2:.*]] = "tfl.average_pool_2d"(%[[VAL_0]]) <{filter_height = 7 : i32, filter_width = 7 : i32, fused_activation_function = "NONE", padding = "VALID", stride_h = 1 : i32, stride_w = 1 : i32}> : (tensor<?x7x7x2048x!quant.uniform<i8:f32, 6.000000e-01:-128>>) -> tensor<?x1x1x2048x!quant.uniform<i8:f32, 6.000000e-01:-128>>
// CHECK: %[[VAL_3:.*]] = "tfl.reshape"(%[[VAL_2]], %[[VAL_1]]) : (tensor<?x1x1x2048x!quant.uniform<i8:f32, 6.000000e-01:-128>>, tensor<2xi32>) -> tensor<?x2048x!quant.uniform<i8:f32, 6.000000e-01:-128>>
// CHECK: %[[VAL_4:.*]] = "tfl.quantize"(%[[VAL_3]]) <{qtype = tensor<?x2048x!quant.uniform<i8:f32, 9.000000e-01:-128>>}> : (tensor<?x2048x!quant.uniform<i8:f32, 6.000000e-01:-128>>) -> tensor<?x2048x!quant.uniform<i8:f32, 9.000000e-01:-128>>
// CHECK: return %[[VAL_4]] : tensor<?x2048x!quant.uniform<i8:f32, 9.000000e-01:-128>>
// CHECK: }
tensorflow/compiler/mlir/lite/experimental/tac/tests/e2e/BUILD
+30
View File
@@ -0,0 +1,30 @@
load("//tensorflow:tensorflow.default.bzl", "filegroup")
load("//tensorflow/compiler/mlir:glob_lit_test.bzl", "glob_lit_tests")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:LICENSE"],
default_visibility = [
"//visibility:public",
],
licenses = ["notice"],
)
glob_lit_tests(
name = "all_tests",
data = [":test_utilities"],
driver = "@llvm-project//mlir:run_lit.sh",
test_file_exts = [
"mlir",
"cc",
],
)
# Bundle together all of the test utilities that are used by tests.
filegroup(
name = "test_utilities",
testonly = True,
data = [
"//tensorflow/compiler/mlir/lite/experimental/tac:tac-translate",
"@llvm-project//llvm:FileCheck",
],
)
tensorflow/compiler/mlir/lite/experimental/tac/tests/e2e/device-transform-nnapi.mlir
+35
View File
@@ -0,0 +1,35 @@
// 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: tac-translate -input-mlir -output-mlir -device-specs=NNAPI %s -o - 2>&1 | FileCheck %s
module {
// CHECK-LABEL: main
func.func @main(%arg0: tensor<4xf32>, %arg1: tensor<4xf32>) -> tensor<4xf32> {
%0 = "tfl.squared_difference"(%arg0, %arg1) : (tensor<4xf32>, tensor<4xf32>) -> tensor<4xf32>
func.return %0 : tensor<4xf32>
// CHECK: [[VAL_0:%.*]] = tfl.sub %arg0, %arg1 {fused_activation_function = "NONE"} : tensor<4xf32>
// CHECK: [[VAL_1:%.*]] = tfl.mul [[VAL_0]], [[VAL_0]] {fused_activation_function = "NONE"} : tensor<4xf32
}
// CHECK-LABEL: pack
func.func @pack(%arg0: tensor<1xf32>, %arg1: tensor<1xf32>) -> tensor<2x1xf32> {
%0 = "tfl.pack"(%arg0, %arg1) {axis = 0 : i32, values_count = 2 : i32} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2x1xf32>
func.return %0 : tensor<2x1xf32>
// CHECK: %[[VAL_0:.*]] = arith.constant dense<[2, 1]> : tensor<2xi32>
// CHECK: %[[CONCAT:.*]] = "tfl.concatenation"(%arg0, %arg1) <{axis = 0 : i32, fused_activation_function = "NONE"}> : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>
// CHECK: %[[VAL_1:.*]] = "tfl.reshape"(%[[CONCAT]], %[[VAL_0]]) : (tensor<2xf32>, tensor<2xi32>) -> tensor<2x1xf32>
// CHECK: return %[[VAL_1]]
}
}
tensorflow/compiler/mlir/lite/experimental/tac/tests/e2e/simple-graph.mlir
+32
View File
@@ -0,0 +1,32 @@
// 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: tac-translate -input-mlir -output-mlir -device-specs=GPU %s -o - 2>&1 | FileCheck %s
module {
func.func @main(%arg0: tensor<1xf32>, %arg1: tensor<1xf32>, %arg2: tensor<1xf32>, %arg3: tensor<1xf32>) -> tensor<2x1xf32> attributes {tf.entry_function = {inputs = "input0,input1,input2,input3", outputs = "output"}} {
%0 = "tfl.add"(%arg0, %arg1) {fused_activation_function = "RELU6"} : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32>
%1 = "tfl.mul"(%0, %arg2) {fused_activation_function = "RELU6"} : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32>
%2 = "tfl.add"(%arg0, %arg3) {fused_activation_function = "RELU6"} : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32>
%3 = "tfl.pack"(%1, %2) {axis = 0 : i32, values_count = 2 : i32} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2x1xf32>
func.return %3 : tensor<2x1xf32>
}
// CHECK: %[[CST:.*]] = arith.constant dense<1> : tensor<4xi32>
// CHECK: [[VAL_0:%.*]] = "tfl.reshape"(%1, %[[CST]]) {tac.device = "GPU", tac.inference_type = "FLOAT"} : (tensor<1xf32>, tensor<4xi32>) -> tensor<1x1x1x1xf32>
// CHECK: [[VAL_1:%.*]] = "tfl.reshape"(%2, %[[CST]]) {tac.device = "GPU", tac.inference_type = "FLOAT"} : (tensor<1xf32>, tensor<4xi32>) -> tensor<1x1x1x1xf32>
// CHECK: [[VAL_2:%.*]] = "tfl.concatenation"([[VAL_0]], [[VAL_1]]) <{axis = 3 : i32, fused_activation_function = "NONE"}> {tac.device = "GPU", tac.inference_type = "FLOAT"} : (tensor<1x1x1x1xf32>, tensor<1x1x1x1xf32>) -> tensor<1x1x1x2xf32>
// CHECK: [[VAL_3:%.*]] = "tfl.reshape"([[VAL_2]], %{{.*}}) : (tensor<1x1x1x2xf32>, tensor<2xi32>) -> tensor<2x1xf32>
}
tensorflow/compiler/mlir/lite/experimental/tac/tests/fold-constants-to-subgraph.mlir
+205
View File
@@ -0,0 +1,205 @@
// 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.
// ==============================================================================
// Test for partial folding: only fold i32 constants.
// RUN: tac-opt-all-backends -tfl-fold-constants-to-subgraph='fold-all-constants=false' %s -split-input-file -verify-diagnostics | FileCheck --check-prefix=PARTIAL %s
// Test for fold all constants.
// RUN: tac-opt-all-backends -tfl-fold-constants-to-subgraph='fold-all-constants=true' %s -split-input-file -verify-diagnostics | FileCheck --check-prefix=ALL %s
module {
func.func @main(%arg0: tensor<4x384x32xf32>) -> tensor<1x384x32xf32> {
%0 = "tfl.pseudo_const"() {value = dense<0> : tensor<3xi32>} : () -> tensor<3xi32>
%1 = "tfl.pseudo_const"() {value = dense<[1, 384, 32]> : tensor<3xi32>} : () -> tensor<3xi32>
%2 = func.call @simple_test(%arg0, %0, %1) {tac.interface_name = "func1"} : (tensor<4x384x32xf32>, tensor<3xi32>, tensor<3xi32>) -> tensor<1x384x32xf32>
func.return %2 : tensor<1x384x32xf32>
}
// PARTIAL-LABEL: @simple_test
func.func @simple_test(%arg0: tensor<4x384x32xf32>, %arg1: tensor<3xi32>, %arg2: tensor<3xi32>) -> tensor<1x384x32xf32> attributes {tac.interface_name = "func1"} {
%0 = "tfl.slice"(%arg0, %arg1, %arg2) : (tensor<4x384x32xf32>, tensor<3xi32>, tensor<3xi32>) -> tensor<1x384x32xf32>
func.return %0 : tensor<1x384x32xf32>
}
// PARTIAL: func @simple_test(%[[VAL_0:.*]]: tensor<4x384x32xf32>, %[[VAL_1:.*]]: tensor<3xi32>, %[[VAL_2:.*]]: tensor<3xi32>) -> tensor<1x384x32xf32> attributes {tac.interface_name = "func1"} {
// PARTIAL: %[[VAL_3:.*]] = "tfl.pseudo_const"() <{value = dense<[1, 384, 32]> : tensor<3xi32>}> : () -> tensor<3xi32>
// PARTIAL: %[[VAL_4:.*]] = "tfl.pseudo_const"() <{value = dense<0> : tensor<3xi32>}> : () -> tensor<3xi32>
// PARTIAL: %[[VAL_5:.*]] = "tfl.slice"(%[[VAL_0]], %[[VAL_4]], %[[VAL_3]]) : (tensor<4x384x32xf32>, tensor<3xi32>, tensor<3xi32>) -> tensor<1x384x32xf32>
// PARTIAL: return %[[VAL_5]] : tensor<1x384x32xf32>
// PARTIAL: }
}
// -----
module {
func.func @main(%arg0: tensor<4x384x32xf32>) -> (tensor<1x384x32xf32>, tensor<1x384x32xf32>) {
%0 = "tfl.pseudo_const"() {value = dense<0> : tensor<3xi32>} : () -> tensor<3xi32>
%1 = "tfl.pseudo_const"() {value = dense<[1, 384, 32]> : tensor<3xi32>} : () -> tensor<3xi32>
%2 = func.call @arg_reuse_test_1(%arg0, %0, %1) {tac.interface_name = "func1"} : (tensor<4x384x32xf32>, tensor<3xi32>, tensor<3xi32>) -> tensor<1x384x32xf32>
%3 = func.call @arg_reuse_test_2(%arg0, %0, %1) {tac.interface_name = "func2"} : (tensor<4x384x32xf32>, tensor<3xi32>, tensor<3xi32>) -> tensor<1x384x32xf32>
func.return %2, %3 : tensor<1x384x32xf32>, tensor<1x384x32xf32>
}
// PARTIAL-LABEL: @arg_reuse_test_1
func.func @arg_reuse_test_1(%arg0: tensor<4x384x32xf32>, %arg1: tensor<3xi32>, %arg2: tensor<3xi32>) -> tensor<1x384x32xf32> attributes {tac.interface_name = "func1"} {
%0 = "tfl.slice"(%arg0, %arg1, %arg2) : (tensor<4x384x32xf32>, tensor<3xi32>, tensor<3xi32>) -> tensor<1x384x32xf32>
func.return %0 : tensor<1x384x32xf32>
}
// PARTIAL-LABEL: @arg_reuse_test_2
func.func @arg_reuse_test_2(%arg0: tensor<4x384x32xf32>, %arg1: tensor<3xi32>, %arg2: tensor<3xi32>) -> tensor<1x384x32xf32> attributes {tac.interface_name = "func2"} {
%0 = "tfl.slice"(%arg0, %arg1, %arg2) : (tensor<4x384x32xf32>, tensor<3xi32>, tensor<3xi32>) -> tensor<1x384x32xf32>
func.return %0 : tensor<1x384x32xf32>
}
// PARTIAL: func @arg_reuse_test_1(%[[VAL_0:.*]]: tensor<4x384x32xf32>, %[[VAL_1:.*]]: tensor<3xi32>, %[[VAL_2:.*]]: tensor<3xi32>) -> tensor<1x384x32xf32> attributes {tac.interface_name = "func1"} {
// PARTIAL: %[[VAL_3:.*]] = "tfl.pseudo_const"() <{value = dense<[1, 384, 32]> : tensor<3xi32>}> : () -> tensor<3xi32>
// PARTIAL: %[[VAL_4:.*]] = "tfl.pseudo_const"() <{value = dense<0> : tensor<3xi32>}> : () -> tensor<3xi32>
// PARTIAL: %[[VAL_5:.*]] = "tfl.slice"(%[[VAL_0]], %[[VAL_4]], %[[VAL_3]]) : (tensor<4x384x32xf32>, tensor<3xi32>, tensor<3xi32>) -> tensor<1x384x32xf32>
// PARTIAL: return %[[VAL_5]] : tensor<1x384x32xf32>
// PARTIAL: }
// PARTIAL: func @arg_reuse_test_2(%[[VAL_6:.*]]: tensor<4x384x32xf32>, %[[VAL_7:.*]]: tensor<3xi32>, %[[VAL_8:.*]]: tensor<3xi32>) -> tensor<1x384x32xf32> attributes {tac.interface_name = "func2"} {
// PARTIAL: %[[VAL_9:.*]] = "tfl.pseudo_const"() <{value = dense<[1, 384, 32]> : tensor<3xi32>}> : () -> tensor<3xi32>
// PARTIAL: %[[VAL_10:.*]] = "tfl.pseudo_const"() <{value = dense<0> : tensor<3xi32>}> : () -> tensor<3xi32>
// PARTIAL: %[[VAL_11:.*]] = "tfl.slice"(%[[VAL_6]], %[[VAL_10]], %[[VAL_9]]) : (tensor<4x384x32xf32>, tensor<3xi32>, tensor<3xi32>) -> tensor<1x384x32xf32>
// PARTIAL: return %[[VAL_11]] : tensor<1x384x32xf32>
// PARTIAL: }
}
// -----
module {
func.func @main(%arg0: tensor<384x512x!quant.uniform<i8:f32, 0.1>>) -> tensor<384x128x!quant.uniform<i8:f32, 0.09:-4>> {
%0 = "tfl.pseudo_qconst"() {qtype = tensor<128x512x!quant.uniform<i8<-127:127>:f32, 0.01>>, value = dense<0> : tensor<128x512xi8>} : () -> tensor<128x512x!quant.uniform<i8<-127:127>:f32, 0.01>>
%1 = "tfl.pseudo_qconst"() {qtype = tensor<128x!quant.uniform<i32:f32, 0.7>>, value = dense<0> : tensor<128xi32>} : () -> tensor<128x!quant.uniform<i32:f32, 0.7>>
%2 = func.call @quantization_test(%arg0, %0, %1) {tac.interface_name = "func1"} : (tensor<384x512x!quant.uniform<i8:f32, 0.1>>, tensor<128x512x!quant.uniform<i8<-127:127>:f32, 0.01>>, tensor<128x!quant.uniform<i32:f32, 0.7>>) -> tensor<384x128x!quant.uniform<i8:f32, 0.09:-4>>
func.return %2 : tensor<384x128x!quant.uniform<i8:f32, 0.09:-4>>
}
// PARTIAL-LABEL: @quantization_test
func.func @quantization_test(%arg0: tensor<384x512x!quant.uniform<i8:f32, 0.1>>, %arg1: tensor<128x512x!quant.uniform<i8<-127:127>:f32, 0.01>>, %arg2: tensor<128x!quant.uniform<i32:f32, 0.7>>) -> tensor<384x128x!quant.uniform<i8:f32, 0.09:-4>> {
%0 = "tfl.fully_connected"(%arg0, %arg1, %arg2) {fused_activation_function = "NONE", keep_num_dims = false, weights_format = "DEFAULT"} : (tensor<384x512x!quant.uniform<i8:f32, 0.1>>, tensor<128x512x!quant.uniform<i8<-127:127>:f32, 0.01>>, tensor<128x!quant.uniform<i32:f32, 0.7>>) -> tensor<384x128x!quant.uniform<i8:f32, 0.09:-4>>
func.return %0 : tensor<384x128x!quant.uniform<i8:f32, 0.09:-4>>
}
// PARTIAL: func @quantization_test(%[[VAL_0:.*]]: tensor<384x512x!quant.uniform<i8:f32, 1.000000e-01>>, %[[VAL_1:.*]]: tensor<128x512x!quant.uniform<i8<-127:127>:f32, 1.000000e-02>>, %[[VAL_2:.*]]: tensor<128x!quant.uniform<i32:f32, 0.69999999999999996>>) -> tensor<384x128x!quant.uniform<i8:f32, 0.089999999999999996:-4>> {
// PARTIAL: %[[VAL_3:.*]] = "tfl.pseudo_qconst"() <{qtype = tensor<128x!quant.uniform<i32:f32, 0.69999999999999996>>, value = dense<0> : tensor<128xi32>}> : () -> tensor<128x!quant.uniform<i32:f32, 0.69999999999999996>>
// PARTIAL: %[[VAL_4:.*]] = "tfl.fully_connected"(%[[VAL_0]], %[[VAL_1]], %[[VAL_3]]) <{fused_activation_function = "NONE", keep_num_dims = false, weights_format = "DEFAULT"}> : (tensor<384x512x!quant.uniform<i8:f32, 1.000000e-01>>, tensor<128x512x!quant.uniform<i8<-127:127>:f32, 1.000000e-02>>, tensor<128x!quant.uniform<i32:f32, 0.69999999999999996>>) -> tensor<384x128x!quant.uniform<i8:f32, 0.089999999999999996:-4>>
// PARTIAL: return %[[VAL_4]] : tensor<384x128x!quant.uniform<i8:f32, 0.089999999999999996:-4>>
// PARTIAL: }
}
// -----
module {
func.func @main(%arg0: tensor<256x32x32x3xf32>) -> tensor<256x30x30x16xf32> {
%0 = "tfl.pseudo_const"() {value = dense<1.000000e+00> : tensor<16x3x3x3xf32>} : () -> tensor<16x3x3x3xf32>
%1 = "tfl.pseudo_const"() {value = dense<1.000000e+00> : tensor<16xf32>} : () -> tensor<16xf32>
%2 = func.call @fold_all_test(%arg0, %0, %1) : (tensor<256x32x32x3xf32>, tensor<16x3x3x3xf32>, tensor<16xf32>) -> tensor<256x30x30x16xf32>
func.return %2 : tensor<256x30x30x16xf32>
}
// ALL-LABEL: @fold_all_test
func.func @fold_all_test(%arg0: tensor<256x32x32x3xf32>, %arg1: tensor<16x3x3x3xf32>, %arg2: tensor<16xf32>) -> tensor<256x30x30x16xf32> {
%0 = "tfl.conv_2d"(%arg0, %arg1, %arg2) {tac.device = "GPU", tac.inference_type = "FLOAT", dilation_h_factor = 1 : i32, dilation_w_factor = 1 : i32, fused_activation_function = "NONE", padding = "VALID", stride_h = 1 : i32, stride_w = 1 : i32} : (tensor<256x32x32x3xf32>, tensor<16x3x3x3xf32>, tensor<16xf32>) -> tensor<256x30x30x16xf32>
func.return %0 : tensor<256x30x30x16xf32>
}
// ALL: func @fold_all_test(%[[VAL_0:.*]]: tensor<256x32x32x3xf32>, %[[VAL_1:.*]]: tensor<16x3x3x3xf32>, %[[VAL_2:.*]]: tensor<16xf32>) -> tensor<256x30x30x16xf32> {
// ALL: %[[VAL_3:.*]] = "tfl.pseudo_const"() <{value = dense<1.000000e+00> : tensor<16xf32>}> : () -> tensor<16xf32>
// ALL: %[[VAL_4:.*]] = "tfl.pseudo_const"() <{value = dense<1.000000e+00> : tensor<16x3x3x3xf32>}> : () -> tensor<16x3x3x3xf32>
// ALL: %[[VAL_5:.*]] = "tfl.conv_2d"(%[[VAL_0]], %[[VAL_4]], %[[VAL_3]]) <{dilation_h_factor = 1 : i32, dilation_w_factor = 1 : i32, fused_activation_function = "NONE", padding = "VALID", stride_h = 1 : i32, stride_w = 1 : i32}> {tac.device = "GPU", tac.inference_type = "FLOAT"} : (tensor<256x32x32x3xf32>, tensor<16x3x3x3xf32>, tensor<16xf32>) -> tensor<256x30x30x16xf32>
// ALL: return %[[VAL_5]] : tensor<256x30x30x16xf32>
// ALL: }
}
// -----
module {
func.func @main(%arg0: tensor<256x32x32x3xf32>) -> tensor<256x30x30x16xf32> {
%0 = stablehlo.constant dense<1.000000e+00> : tensor<16x3x3x3xf32>
%1 = stablehlo.constant dense<1.000000e+00> : tensor<16xf32>
%2 = func.call @fold_stablehlo_constant_test(%arg0, %0, %1) : (tensor<256x32x32x3xf32>, tensor<16x3x3x3xf32>, tensor<16xf32>) -> tensor<256x30x30x16xf32>
func.return %2 : tensor<256x30x30x16xf32>
}
func.func @fold_stablehlo_constant_test(%arg0: tensor<256x32x32x3xf32>, %arg1: tensor<16x3x3x3xf32>, %arg2: tensor<16xf32>) -> tensor<256x30x30x16xf32> {
%0 = "tfl.conv_2d"(%arg0, %arg1, %arg2) {tac.device = "GPU", tac.inference_type = "FLOAT", dilation_h_factor = 1 : i32, dilation_w_factor = 1 : i32, fused_activation_function = "NONE", padding = "VALID", stride_h = 1 : i32, stride_w = 1 : i32} : (tensor<256x32x32x3xf32>, tensor<16x3x3x3xf32>, tensor<16xf32>) -> tensor<256x30x30x16xf32>
func.return %0 : tensor<256x30x30x16xf32>
}
// ALL-LABEL: func @fold_stablehlo_constant_test
// ALL-SAME: (%[[VAL_0:.*]]: tensor<256x32x32x3xf32>, %[[VAL_1:.*]]: tensor<16x3x3x3xf32>, %[[VAL_2:.*]]: tensor<16xf32>) -> tensor<256x30x30x16xf32> {
// ALL-DAG: %[[VAL_3:.*]] = stablehlo.constant dense<1.000000e+00> : tensor<16xf32>
// ALL-DAG: %[[VAL_4:.*]] = stablehlo.constant dense<1.000000e+00> : tensor<16x3x3x3xf32>
// ALL: %[[VAL_5:.*]] = "tfl.conv_2d"(%[[VAL_0]], %[[VAL_4]], %[[VAL_3]]) <{dilation_h_factor = 1 : i32, dilation_w_factor = 1 : i32, fused_activation_function = "NONE", padding = "VALID", stride_h = 1 : i32, stride_w = 1 : i32}> {tac.device = "GPU", tac.inference_type = "FLOAT"} : (tensor<256x32x32x3xf32>, tensor<16x3x3x3xf32>, tensor<16xf32>) -> tensor<256x30x30x16xf32>
// ALL: return %[[VAL_5]] : tensor<256x30x30x16xf32>
// ALL: }
}
// -----
module {
func.func @main(%arg0: tensor<4x384x32xf32>) -> tensor<1x384x32xf32> {
%0 = arith.constant dense<0> : tensor<3xi32>
%1 = arith.constant dense<[1, 384, 32]> : tensor<3xi32>
%2 = func.call @simple_test(%arg0, %0, %1) {tac.interface_name = "func1"} : (tensor<4x384x32xf32>, tensor<3xi32>, tensor<3xi32>) -> tensor<1x384x32xf32>
func.return %2 : tensor<1x384x32xf32>
}
// PARTIAL-LABEL: @simple_test
func.func @simple_test(%arg0: tensor<4x384x32xf32>, %arg1: tensor<3xi32>, %arg2: tensor<3xi32>) -> tensor<1x384x32xf32> attributes {tac.interface_name = "func1"} {
%0 = "tfl.slice"(%arg0, %arg1, %arg2) : (tensor<4x384x32xf32>, tensor<3xi32>, tensor<3xi32>) -> tensor<1x384x32xf32>
func.return %0 : tensor<1x384x32xf32>
}
// PARTIAL: func @simple_test(%[[VAL_0:.*]]: tensor<4x384x32xf32>, %[[VAL_1:.*]]: tensor<3xi32>, %[[VAL_2:.*]]: tensor<3xi32>) -> tensor<1x384x32xf32> attributes {tac.interface_name = "func1"} {
// PARTIAL: %[[VAL_3:.*]] = arith.constant dense<[1, 384, 32]> : tensor<3xi32>
// PARTIAL: %[[VAL_4:.*]] = arith.constant dense<0> : tensor<3xi32>
// PARTIAL: %[[VAL_5:.*]] = "tfl.slice"(%[[VAL_0]], %[[VAL_4]], %[[VAL_3]]) : (tensor<4x384x32xf32>, tensor<3xi32>, tensor<3xi32>) -> tensor<1x384x32xf32>
// PARTIAL: return %[[VAL_5]] : tensor<1x384x32xf32>
// PARTIAL: }
}
// -----
module {
func.func @main(%arg0: tensor<4x384x32xf32>) -> tensor<1x384x32xf32> {
%0 = stablehlo.constant dense<0> : tensor<3xi32>
%1 = stablehlo.constant dense<[1, 384, 32]> : tensor<3xi32>
%2 = func.call @partial_stablehlo_constant_test(%arg0, %0, %1) {tac.interface_name = "func1"} : (tensor<4x384x32xf32>, tensor<3xi32>, tensor<3xi32>) -> tensor<1x384x32xf32>
func.return %2 : tensor<1x384x32xf32>
}
func.func @partial_stablehlo_constant_test(%arg0: tensor<4x384x32xf32>, %arg1: tensor<3xi32>, %arg2: tensor<3xi32>) -> tensor<1x384x32xf32> attributes {tac.interface_name = "func1"} {
%0 = "tfl.slice"(%arg0, %arg1, %arg2) : (tensor<4x384x32xf32>, tensor<3xi32>, tensor<3xi32>) -> tensor<1x384x32xf32>
func.return %0 : tensor<1x384x32xf32>
}
// PARTIAL-LABEL: func @partial_stablehlo_constant_test
// PARTIAL-SAME: (%[[VAL_0:.*]]: tensor<4x384x32xf32>, %[[VAL_1:.*]]: tensor<3xi32>, %[[VAL_2:.*]]: tensor<3xi32>) -> tensor<1x384x32xf32> attributes {tac.interface_name = "func1"} {
// PARTIAL-DAG: %[[VAL_3:.*]] = stablehlo.constant dense<[1, 384, 32]> : tensor<3xi32>
// PARTIAL-DAG: %[[VAL_4:.*]] = stablehlo.constant dense<0> : tensor<3xi32>
// PARTIAL: %[[VAL_5:.*]] = "tfl.slice"(%[[VAL_0]], %[[VAL_4]], %[[VAL_3]]) : (tensor<4x384x32xf32>, tensor<3xi32>, tensor<3xi32>) -> tensor<1x384x32xf32>
// PARTIAL: return %[[VAL_5]] : tensor<1x384x32xf32>
// PARTIAL: }
}
tensorflow/compiler/mlir/lite/experimental/tac/tests/get-alternative-subgraph.mlir
+182
View File
@@ -0,0 +1,182 @@
// 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: tac-opt-all-backends -tfl-get-alternative-subgraph='device-specs=GPU' %s -split-input-file -verify-diagnostics | FileCheck %s
module {
func.func @simpleTest(%arg0: tensor<1xf32>, %arg1: tensor<1xf32>, %arg2: tensor<1xf32>, %arg3: tensor<1xf32>) -> tensor<2x1xf32> {
%0 = func.call @func_0_GPU_FLOAT(%arg0, %arg1, %arg2) {tac.interface_name = "func_0"} : (tensor<1xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32>
%1 = func.call @func_1_GPU_FLOAT(%arg0, %arg3) {tac.interface_name = "func_1"} : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32>
%2 = func.call @func_2_CPU_FLOAT(%0, %1) {tac.interface_name = "func_2"} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2x1xf32>
func.return %2 : tensor<2x1xf32>
}
func.func private @func_2_CPU_FLOAT(%arg0: tensor<1xf32>, %arg1: tensor<1xf32>) -> tensor<2x1xf32> attributes {tac.device = "CPU", tac.inference_type = "FLOAT", tac.interface_name = "func_2"} {
%0 = "tfl.pack"(%arg0, %arg1) {axis = 0 : i32, tac.device = "CPU", tac.inference_type = "FLOAT", values_count = 2 : i32} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2x1xf32>
func.return %0 : tensor<2x1xf32>
}
func.func private @func_0_GPU_FLOAT(%arg0: tensor<1xf32>, %arg1: tensor<1xf32>, %arg2: tensor<1xf32>) -> tensor<1xf32> attributes {tac.device = "GPU", tac.inference_type = "FLOAT", tac.interface_name = "func_0"} {
%0 = tfl.add %arg0, %arg1 {fused_activation_function = "RELU6", tac.device = "GPU", tac.inference_type = "FLOAT"} : tensor<1xf32>
%1 = tfl.mul %0, %arg2 {fused_activation_function = "RELU6", tac.device = "GPU", tac.inference_type = "FLOAT"} : tensor<1xf32>
func.return %1 : tensor<1xf32>
}
func.func private @func_1_GPU_FLOAT(%arg0: tensor<1xf32>, %arg1: tensor<1xf32>) -> tensor<1xf32> attributes {tac.device = "GPU", tac.inference_type = "FLOAT", tac.interface_name = "func_1"} {
%0 = tfl.add %arg0, %arg1 {fused_activation_function = "RELU6", tac.device = "GPU", tac.inference_type = "FLOAT"} : tensor<1xf32>
func.return %0 : tensor<1xf32>
}
// CHECK: func @simpleTest(%[[VAL_0:.*]]: tensor<1xf32>, %[[VAL_1:.*]]: tensor<1xf32>, %[[VAL_2:.*]]: tensor<1xf32>, %[[VAL_3:.*]]: tensor<1xf32>) -> tensor<2x1xf32> {
// CHECK: %[[VAL_4:.*]] = call @func_0_GPU_FLOAT(%[[VAL_0]], %[[VAL_1]], %[[VAL_2]]) {tac.interface_name = "func_0"} : (tensor<1xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32>
// CHECK: %[[VAL_5:.*]] = call @func_1_GPU_FLOAT(%[[VAL_0]], %[[VAL_3]]) {tac.interface_name = "func_1"} : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32>
// CHECK: %[[VAL_6:.*]] = call @func_2_CPU_FLOAT(%[[VAL_4]], %[[VAL_5]]) {tac.interface_name = "func_2"} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2x1xf32>
// CHECK: return %[[VAL_6]] : tensor<2x1xf32>
// CHECK: }
// CHECK: func private @func_2_CPU_FLOAT(%[[VAL_0:.*]]: tensor<1xf32>, %[[VAL_1:.*]]: tensor<1xf32>) -> tensor<2x1xf32> attributes {tac.device = "CPU", tac.inference_type = "FLOAT", tac.interface_name = "func_2"} {
// CHECK: %[[VAL_2:.*]] = "tfl.pack"(%[[VAL_0]], %[[VAL_1]]) <{axis = 0 : i32, values_count = 2 : i32}> {tac.device = "CPU", tac.inference_type = "FLOAT"} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2x1xf32>
// CHECK: return %[[VAL_2]] : tensor<2x1xf32>
// CHECK: }
// CHECK: func private @func_0_GPU_FLOAT(%[[VAL_0:.*]]: tensor<1xf32>, %[[VAL_1:.*]]: tensor<1xf32>, %[[VAL_2:.*]]: tensor<1xf32>) -> tensor<1xf32> attributes {tac.device = "GPU", tac.inference_type = "FLOAT", tac.interface_name = "func_0"} {
// CHECK: %[[VAL_3:.*]] = tfl.add %[[VAL_0]], %[[VAL_1]] {fused_activation_function = "RELU6", tac.device = "GPU", tac.inference_type = "FLOAT"} : tensor<1xf32>
// CHECK: %[[VAL_4:.*]] = tfl.mul %[[VAL_3]], %[[VAL_2]] {fused_activation_function = "RELU6", tac.device = "GPU", tac.inference_type = "FLOAT"} : tensor<1xf32>
// CHECK: return %[[VAL_4]] : tensor<1xf32>
// CHECK: }
// CHECK: func private @func_1_GPU_FLOAT(%[[VAL_0:.*]]: tensor<1xf32>, %[[VAL_1:.*]]: tensor<1xf32>) -> tensor<1xf32> attributes {tac.device = "GPU", tac.inference_type = "FLOAT", tac.interface_name = "func_1"} {
// CHECK: %[[VAL_2:.*]] = tfl.add %[[VAL_0]], %[[VAL_1]] {fused_activation_function = "RELU6", tac.device = "GPU", tac.inference_type = "FLOAT"} : tensor<1xf32>
// CHECK: return %[[VAL_2]] : tensor<1xf32>
// CHECK: }
// CHECK: func private @func_2_GPU_FLOAT(%[[VAL_0:.*]]: tensor<1xf32>, %[[VAL_1:.*]]: tensor<1xf32>) -> tensor<2x1xf32> attributes {tac.device = "GPU", tac.inference_type = "FLOAT", tac.interface_name = "func_2"} {
// CHECK-DAG: %[[VAL_2:.*]] = "tfl.pseudo_const"(){{.*}} dense<1> : tensor<4xi32>
// CHECK-DAG: %[[VAL_3:.*]] = "tfl.pseudo_const"(){{.*}}dense<2> : tensor<1xi32>
// CHECK-DAG: %[[VAL_4:.*]] = "tfl.pseudo_const"(){{.*}}dense<[2, 1]> : tensor<2xi32>
// CHECK: %[[VAL_5:.*]] = "tfl.reshape"(%[[VAL_0]], %[[VAL_2]]) {tac.device = "GPU", tac.inference_type = "FLOAT"} : (tensor<1xf32>, tensor<4xi32>) -> tensor<1x1x1x1xf32>
// CHECK: %[[VAL_6:.*]] = "tfl.reshape"(%[[VAL_1]], %[[VAL_2]]) {tac.device = "GPU", tac.inference_type = "FLOAT"} : (tensor<1xf32>, tensor<4xi32>) -> tensor<1x1x1x1xf32>
// CHECK: %[[VAL_7:.*]] = "tfl.concatenation"(%[[VAL_5]], %[[VAL_6]]) <{axis = 3 : i32, fused_activation_function = "NONE"}> {tac.device = "GPU", tac.inference_type = "FLOAT"} : (tensor<1x1x1x1xf32>, tensor<1x1x1x1xf32>) -> tensor<1x1x1x2xf32>
// CHECK: %[[VAL_8:.*]] = "tfl.reshape"(%[[VAL_7]], %[[VAL_3]]) {tac.device = "GPU", tac.inference_type = "FLOAT"} : (tensor<1x1x1x2xf32>, tensor<1xi32>) -> tensor<2xf32>
// CHECK: %[[VAL_9:.*]] = "tfl.reshape"(%[[VAL_8]], %[[VAL_4]]) {tac.device = "GPU", tac.inference_type = "FLOAT"} : (tensor<2xf32>, tensor<2xi32>) -> tensor<2x1xf32>
// CHECK: return %[[VAL_9]] : tensor<2x1xf32>
// CHECK: }
// CHECK: func private @func_0_CPU_FLOAT(%[[VAL_0:.*]]: tensor<1xf32>, %[[VAL_1:.*]]: tensor<1xf32>, %[[VAL_2:.*]]: tensor<1xf32>) -> tensor<1xf32> attributes {tac.device = "CPU", tac.inference_type = "FLOAT", tac.interface_name = "func_0"} {
// CHECK: %[[VAL_3:.*]] = tfl.add %[[VAL_0]], %[[VAL_1]] {fused_activation_function = "RELU6", tac.device = "CPU", tac.inference_type = "FLOAT"} : tensor<1xf32>
// CHECK: %[[VAL_4:.*]] = tfl.mul %[[VAL_3]], %[[VAL_2]] {fused_activation_function = "RELU6", tac.device = "CPU", tac.inference_type = "FLOAT"} : tensor<1xf32>
// CHECK: return %[[VAL_4]] : tensor<1xf32>
// CHECK: }
// CHECK: func private @func_1_CPU_FLOAT(%[[VAL_0:.*]]: tensor<1xf32>, %[[VAL_1:.*]]: tensor<1xf32>) -> tensor<1xf32> attributes {tac.device = "CPU", tac.inference_type = "FLOAT", tac.interface_name = "func_1"} {
// CHECK: %[[VAL_2:.*]] = tfl.add %[[VAL_0]], %[[VAL_1]] {fused_activation_function = "RELU6", tac.device = "CPU", tac.inference_type = "FLOAT"} : tensor<1xf32>
// CHECK: return %[[VAL_2]] : tensor<1xf32>
// CHECK: }
}
// -----
module {
func.func private @func_10_CPU_FLOAT(%arg0: tensor<3xi32>, %arg1: tensor<i32>, %arg2: tensor<f32>, %arg3: tensor<f32>) -> tensor<*xf32> attributes {tac.device = "CPU", tac.inference_type = "FLOAT", tac.interface_name = "func_10"} {
%0 = "tfl.one_hot"(%arg0, %arg1, %arg2, %arg3) {axis = -1 : i32, tac.device = "CPU", tac.inference_type = "FLOAT"} : (tensor<3xi32>, tensor<i32>, tensor<f32>, tensor<f32>) -> tensor<*xf32>
func.return %0 : tensor<*xf32>
}
// CHECK: func private @func_10_CPU_FLOAT(%[[VAL_0:.*]]: tensor<3xi32>, %[[VAL_1:.*]]: tensor<i32>, %[[VAL_2:.*]]: tensor<f32>, %[[VAL_3:.*]]: tensor<f32>) -> tensor<*xf32> attributes {tac.device = "CPU", tac.inference_type = "FLOAT", tac.interface_name = "func_10"} {
// CHECK: %[[VAL_4:.*]] = "tfl.one_hot"(%[[VAL_0]], %[[VAL_1]], %[[VAL_2]], %[[VAL_3]]) <{axis = -1 : i32}> {tac.device = "CPU", tac.inference_type = "FLOAT"} : (tensor<3xi32>, tensor<i32>, tensor<f32>, tensor<f32>) -> tensor<*xf32>
// CHECK: return %[[VAL_4]] : tensor<*xf32>
// CHECK: }
}
// -----
module {
func.func private @func_20_GPU_FLOAT(%arg0: tensor<128x128xf32>, %arg1: tensor<3xi32>) -> tensor<1x128x128xf32> attributes {tac.device = "GPU", tac.inference_type = "FLOAT", tac.interface_name = "func_20"} {
%0 = "tfl.reshape"(%arg0, %arg1) {tac.device = "GPU", tac.inference_type = "FLOAT"} : (tensor<128x128xf32>, tensor<3xi32>) -> tensor<1x128x128xf32>
func.return %0 : tensor<1x128x128xf32>
}
// CHECK: func private @func_20_GPU_FLOAT(%[[VAL_0:.*]]: tensor<128x128xf32>, %[[VAL_1:.*]]: tensor<3xi32>) -> tensor<1x128x128xf32> attributes {tac.device = "GPU", tac.inference_type = "FLOAT", tac.interface_name = "func_20"} {
// CHECK: %[[VAL_2:.*]] = "tfl.reshape"(%[[VAL_0]], %[[VAL_1]]) {tac.device = "GPU", tac.inference_type = "FLOAT"} : (tensor<128x128xf32>, tensor<3xi32>) -> tensor<1x128x128xf32>
// CHECK: return %[[VAL_2]] : tensor<1x128x128xf32>
// CHECK: }
// CHECK: func private @func_20_CPU_FLOAT(%[[VAL_0:.*]]: tensor<128x128xf32>, %[[VAL_1:.*]]: tensor<3xi32>) -> tensor<1x128x128xf32> attributes {tac.device = "CPU", tac.inference_type = "FLOAT", tac.interface_name = "func_20"} {
// CHECK: %[[VAL_2:.*]] = "tfl.reshape"(%[[VAL_0]], %[[VAL_1]]) {tac.device = "CPU", tac.inference_type = "FLOAT"} : (tensor<128x128xf32>, tensor<3xi32>) -> tensor<1x128x128xf32>
// CHECK: return %[[VAL_2]] : tensor<1x128x128xf32>
// CHECK: }
}
// -----
module {
func.func private @quantize_ops_CPU_QUANTIZED_INT8(%arg0: tensor<384x512x!quant.uniform<i8:f32, 0.1>>, %arg1: tensor<128x512x!quant.uniform<i8<-127:127>:f32, 0.1>>, %arg2: tensor<128x!quant.uniform<i8:f32, 0.2:-128>>, %arg3: tensor<128x!quant.uniform<i8:f32, 0.2:-4>>) -> tensor<1x384x128x!quant.uniform<i8:f32, 0.3:-3>> attributes {tac.device = "CPU", tac.inference_type = "QUANTIZED_INT8", tac.interface_name = "quantize_ops"} {
%0 = "tfl.pseudo_qconst"() {qtype = tensor<128x!quant.uniform<i32:f32, 0.7>>, value = dense<0> : tensor<128xi32>} : () -> tensor<128x!quant.uniform<i32:f32, 0.7>>
%1 = "tfl.fully_connected"(%arg0, %arg1, %0) {tac.device = "CPU", tac.inference_type = "QUANTIZED_INT8", fused_activation_function = "NONE", keep_num_dims = false, weights_format = "DEFAULT"} : (tensor<384x512x!quant.uniform<i8:f32, 0.1>>, tensor<128x512x!quant.uniform<i8<-127:127>:f32, 0.1>>, tensor<128x!quant.uniform<i32:f32, 0.7>>) -> tensor<384x128x!quant.uniform<i8:f32, 0.9:-4>>
%2 = "arith.constant"() {value = dense<[1, 384, 128]> : tensor<3xi32>} : () -> tensor<3xi32>
%3 = "tfl.reshape"(%1, %2) {tac.device = "CPU", tac.inference_type = "QUANTIZED_INT8"} : (tensor<384x128x!quant.uniform<i8:f32, 0.9:-4>>, tensor<3xi32>) -> tensor<1x384x128x!quant.uniform<i8:f32, 0.9:-4>>
%4 = "tfl.mul"(%3, %arg2) {tac.device = "CPU", tac.inference_type = "QUANTIZED_INT8", fused_activation_function = "NONE"} : (tensor<1x384x128x!quant.uniform<i8:f32, 0.9:-4>>, tensor<128x!quant.uniform<i8:f32, 0.2:-128>>) -> tensor<1x384x128x!quant.uniform<i8:f32, 0.3:3>>
%5 = "tfl.add"(%4, %arg3) {tac.device = "CPU", tac.inference_type = "QUANTIZED_INT8", fused_activation_function = "NONE"} : (tensor<1x384x128x!quant.uniform<i8:f32, 0.3:3>>, tensor<128x!quant.uniform<i8:f32, 0.2:-4>>) -> tensor<1x384x128x!quant.uniform<i8:f32, 0.3:-3>>
func.return %5 : tensor<1x384x128x!quant.uniform<i8:f32, 0.3:-3>>
}
// CHECK: func private @quantize_ops_CPU_QUANTIZED_INT8(%[[VAL_0:.*]]: tensor<384x512x!quant.uniform<i8:f32, 1.000000e-01>>, %[[VAL_1:.*]]: tensor<128x512x!quant.uniform<i8<-127:127>:f32, 1.000000e-01>>, %[[VAL_2:.*]]: tensor<128x!quant.uniform<i8:f32, 2.000000e-01:-128>>, %[[VAL_3:.*]]: tensor<128x!quant.uniform<i8:f32, 2.000000e-01:-4>>) -> tensor<1x384x128x!quant.uniform<i8:f32, 3.000000e-01:-3>> attributes {tac.device = "CPU", tac.inference_type = "QUANTIZED_INT8", tac.interface_name = "quantize_ops"} {
// CHECK-DAG: %[[VAL_4:.*]] = arith.constant dense<[1, 384, 128]> : tensor<3xi32>
// CHECK-DAG: %[[VAL_5:.*]] = "tfl.pseudo_qconst"() <{qtype = tensor<128x!quant.uniform<i32:f32, 0.69999999999999996>>, value = dense<0> : tensor<128xi32>}> : () -> tensor<128x!quant.uniform<i32:f32, 0.69999999999999996>>
// CHECK: %[[VAL_6:.*]] = "tfl.fully_connected"(%[[VAL_0]], %[[VAL_1]], %[[VAL_5]]) <{fused_activation_function = "NONE", keep_num_dims = false, weights_format = "DEFAULT"}> {tac.device = "CPU", tac.inference_type = "QUANTIZED_INT8"} : (tensor<384x512x!quant.uniform<i8:f32, 1.000000e-01>>, tensor<128x512x!quant.uniform<i8<-127:127>:f32, 1.000000e-01>>, tensor<128x!quant.uniform<i32:f32, 0.69999999999999996>>) -> tensor<384x128x!quant.uniform<i8:f32, 9.000000e-01:-4>>
// CHECK: %[[VAL_7:.*]] = "tfl.reshape"(%[[VAL_6]], %[[VAL_4]]) {tac.device = "CPU", tac.inference_type = "QUANTIZED_INT8"} : (tensor<384x128x!quant.uniform<i8:f32, 9.000000e-01:-4>>, tensor<3xi32>) -> tensor<1x384x128x!quant.uniform<i8:f32, 9.000000e-01:-4>>
// CHECK: %[[VAL_8:.*]] = tfl.mul(%[[VAL_7]], %[[VAL_2]]) <{fused_activation_function = "NONE"}> {tac.device = "CPU", tac.inference_type = "QUANTIZED_INT8"} : (tensor<1x384x128x!quant.uniform<i8:f32, 9.000000e-01:-4>>, tensor<128x!quant.uniform<i8:f32, 2.000000e-01:-128>>) -> tensor<1x384x128x!quant.uniform<i8:f32, 3.000000e-01:3>>
// CHECK: %[[VAL_9:.*]] = tfl.add(%[[VAL_8]], %[[VAL_3]]) <{fused_activation_function = "NONE"}> {tac.device = "CPU", tac.inference_type = "QUANTIZED_INT8"} : (tensor<1x384x128x!quant.uniform<i8:f32, 3.000000e-01:3>>, tensor<128x!quant.uniform<i8:f32, 2.000000e-01:-4>>) -> tensor<1x384x128x!quant.uniform<i8:f32, 3.000000e-01:-3>>
// CHECK: return %[[VAL_9]] : tensor<1x384x128x!quant.uniform<i8:f32, 3.000000e-01:-3>>
// CHECK: }
// CHECK: func private @quantize_ops_GPU_FLOAT(%[[VAL_0:.*]]: tensor<384x512x!quant.uniform<i8:f32, 1.000000e-01>>, %[[VAL_1:.*]]: tensor<128x512x!quant.uniform<i8<-127:127>:f32, 1.000000e-01>>, %[[VAL_2:.*]]: tensor<128x!quant.uniform<i8:f32, 2.000000e-01:-128>>, %[[VAL_3:.*]]: tensor<128x!quant.uniform<i8:f32, 2.000000e-01:-4>>) -> tensor<1x384x128x!quant.uniform<i8:f32, 3.000000e-01:-3>> attributes {tac.device = "GPU", tac.inference_type = "FLOAT", tac.interface_name = "quantize_ops"} {
// CHECK-DAG: %[[VAL_4:.*]] = "tfl.pseudo_const"(){{.*}}dense<0.000000e+00> : tensor<128xf32>
// CHECK-DAG: %[[VAL_5:.*]] = arith.constant dense<[1, 384, 128]> : tensor<3xi32>
// CHECK-DAG: %[[VAL_6:.*]] = "tfl.pseudo_const"(){{.*}}dense<[1, 1, 384, 512]> : tensor<4xi32>
// CHECK-DAG: %[[VAL_7:.*]] = "tfl.pseudo_const"(){{.*}}dense<[128, 1, 1, 512]> : tensor<4xi32>
// CHECK-DAG: %[[VAL_8:.*]] = "tfl.pseudo_const"(){{.*}}dense<[384, 128]> : tensor<2xi32>
// CHECK: %[[VAL_9:.*]] = "tfl.dequantize"(%[[VAL_0]]) {tac.device = "GPU", tac.inference_type = "FLOAT"} : (tensor<384x512x!quant.uniform<i8:f32, 1.000000e-01>>) -> tensor<384x512xf32>
// CHECK: %[[VAL_10:.*]] = "tfl.dequantize"(%[[VAL_1]]) {tac.device = "GPU", tac.inference_type = "FLOAT"} : (tensor<128x512x!quant.uniform<i8<-127:127>:f32, 1.000000e-01>>) -> tensor<128x512xf32>
// CHECK: %[[VAL_11:.*]] = "tfl.reshape"(%[[VAL_9]], %[[VAL_6]]) {tac.device = "GPU", tac.inference_type = "FLOAT"} : (tensor<384x512xf32>, tensor<4xi32>) -> tensor<1x1x384x512xf32>
// CHECK: %[[VAL_12:.*]] = "tfl.reshape"(%[[VAL_10]], %[[VAL_7]]) {tac.device = "GPU", tac.inference_type = "FLOAT"} : (tensor<128x512xf32>, tensor<4xi32>) -> tensor<128x1x1x512xf32>
// CHECK: %[[VAL_13:.*]] = "tfl.conv_2d"(%[[VAL_11]], %[[VAL_12]], %[[VAL_4]]) <{dilation_h_factor = 1 : i32, dilation_w_factor = 1 : i32, fused_activation_function = "NONE", padding = "VALID", stride_h = 1 : i32, stride_w = 1 : i32}> {tac.device = "GPU", tac.inference_type = "FLOAT"} : (tensor<1x1x384x512xf32>, tensor<128x1x1x512xf32>, tensor<128xf32>) -> tensor<1x1x384x128xf32>
// CHECK: %[[VAL_14:.*]] = "tfl.reshape"(%[[VAL_13]], %[[VAL_8]]) {tac.device = "GPU", tac.inference_type = "FLOAT"} : (tensor<1x1x384x128xf32>, tensor<2xi32>) -> tensor<384x128xf32>
// CHECK: %[[VAL_15:.*]] = "tfl.reshape"(%[[VAL_14]], %[[VAL_5]]) {tac.device = "GPU", tac.inference_type = "FLOAT"} : (tensor<384x128xf32>, tensor<3xi32>) -> tensor<1x384x128xf32>
// CHECK: %[[VAL_16:.*]] = "tfl.dequantize"(%[[VAL_2]]) {tac.device = "GPU", tac.inference_type = "FLOAT"} : (tensor<128x!quant.uniform<i8:f32, 2.000000e-01:-128>>) -> tensor<128xf32>
// CHECK: %[[VAL_17:.*]] = tfl.mul(%[[VAL_15]], %[[VAL_16]]) <{fused_activation_function = "NONE"}> {tac.device = "GPU", tac.inference_type = "FLOAT"} : (tensor<1x384x128xf32>, tensor<128xf32>) -> tensor<1x384x128xf32>
// CHECK: %[[VAL_18:.*]] = "tfl.dequantize"(%[[VAL_3]]) {tac.device = "GPU", tac.inference_type = "FLOAT"} : (tensor<128x!quant.uniform<i8:f32, 2.000000e-01:-4>>) -> tensor<128xf32>
// CHECK: %[[VAL_19:.*]] = tfl.add(%[[VAL_17]], %[[VAL_18]]) <{fused_activation_function = "NONE"}> {tac.device = "GPU", tac.inference_type = "FLOAT"} : (tensor<1x384x128xf32>, tensor<128xf32>) -> tensor<1x384x128xf32>
// CHECK: %[[VAL_20:.*]] = "tfl.quantize"(%[[VAL_19]]) <{qtype = tensor<1x384x128x!quant.uniform<i8:f32, 3.000000e-01:-3>>}> {tac.device = "GPU", tac.inference_type = "FLOAT"} : (tensor<1x384x128xf32>) -> tensor<1x384x128x!quant.uniform<i8:f32, 3.000000e-01:-3>>
// CHECK: return %[[VAL_20]] : tensor<1x384x128x!quant.uniform<i8:f32, 3.000000e-01:-3>>
// CHECK: }
// CHECK: func private @quantize_ops_CPU_FLOAT(%[[VAL_0:.*]]: tensor<384x512x!quant.uniform<i8:f32, 1.000000e-01>>, %[[VAL_1:.*]]: tensor<128x512x!quant.uniform<i8<-127:127>:f32, 1.000000e-01>>, %[[VAL_2:.*]]: tensor<128x!quant.uniform<i8:f32, 2.000000e-01:-128>>, %[[VAL_3:.*]]: tensor<128x!quant.uniform<i8:f32, 2.000000e-01:-4>>) -> tensor<1x384x128x!quant.uniform<i8:f32, 3.000000e-01:-3>> attributes {tac.device = "CPU", tac.inference_type = "FLOAT", tac.interface_name = "quantize_ops"} {
// CHECK-DAG: %[[VAL_4:.*]] = "tfl.pseudo_const"(){{.*}}dense<0.000000e+00> : tensor<128xf32>
// CHECK-DAG: %[[VAL_5:.*]] = arith.constant dense<[1, 384, 128]> : tensor<3xi32>
// CHECK: %[[VAL_6:.*]] = "tfl.dequantize"(%[[VAL_0]]) {tac.device = "CPU", tac.inference_type = "FLOAT"} : (tensor<384x512x!quant.uniform<i8:f32, 1.000000e-01>>) -> tensor<384x512xf32>
// CHECK: %[[VAL_7:.*]] = "tfl.dequantize"(%[[VAL_1]]) {tac.device = "CPU", tac.inference_type = "FLOAT"} : (tensor<128x512x!quant.uniform<i8<-127:127>:f32, 1.000000e-01>>) -> tensor<128x512xf32>
// CHECK: %[[VAL_8:.*]] = "tfl.fully_connected"(%[[VAL_6]], %[[VAL_7]], %[[VAL_4]]) <{fused_activation_function = "NONE", keep_num_dims = false, weights_format = "DEFAULT"}> {tac.device = "CPU", tac.inference_type = "FLOAT"} : (tensor<384x512xf32>, tensor<128x512xf32>, tensor<128xf32>) -> tensor<384x128xf32>
// CHECK: %[[VAL_9:.*]] = "tfl.reshape"(%[[VAL_8]], %[[VAL_5]]) {tac.device = "CPU", tac.inference_type = "FLOAT"} : (tensor<384x128xf32>, tensor<3xi32>) -> tensor<1x384x128xf32>
// CHECK: %[[VAL_10:.*]] = "tfl.dequantize"(%[[VAL_2]]) {tac.device = "CPU", tac.inference_type = "FLOAT"} : (tensor<128x!quant.uniform<i8:f32, 2.000000e-01:-128>>) -> tensor<128xf32>
// CHECK: %[[VAL_11:.*]] = tfl.mul(%[[VAL_9]], %[[VAL_10]]) <{fused_activation_function = "NONE"}> {tac.device = "CPU", tac.inference_type = "FLOAT"} : (tensor<1x384x128xf32>, tensor<128xf32>) -> tensor<1x384x128xf32>
// CHECK: %[[VAL_12:.*]] = "tfl.dequantize"(%[[VAL_3]]) {tac.device = "CPU", tac.inference_type = "FLOAT"} : (tensor<128x!quant.uniform<i8:f32, 2.000000e-01:-4>>) -> tensor<128xf32>
// CHECK: %[[VAL_13:.*]] = tfl.add(%[[VAL_11]], %[[VAL_12]]) <{fused_activation_function = "NONE"}> {tac.device = "CPU", tac.inference_type = "FLOAT"} : (tensor<1x384x128xf32>, tensor<128xf32>) -> tensor<1x384x128xf32>
// CHECK: %[[VAL_14:.*]] = "tfl.quantize"(%[[VAL_13]]) <{qtype = tensor<1x384x128x!quant.uniform<i8:f32, 3.000000e-01:-3>>}> {tac.device = "CPU", tac.inference_type = "FLOAT"} : (tensor<1x384x128xf32>) -> tensor<1x384x128x!quant.uniform<i8:f32, 3.000000e-01:-3>>
// CHECK: return %[[VAL_14]] : tensor<1x384x128x!quant.uniform<i8:f32, 3.000000e-01:-3>>
// CHECK: }
}
tensorflow/compiler/mlir/lite/experimental/tac/tests/get-op-cost.mlir
+101
View File
@@ -0,0 +1,101 @@
// 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: tac-opt-all-backends -tfl-get-op-cost %s -split-input-file -verify-diagnostics | FileCheck %s
func.func @func_0_CPU(%arg0: tensor<256x32x32x3xf32>, %arg1: tensor<256x32x32x3xf32>) -> tensor<256x32x32x3xf32> attributes {tac.device = "CPU", tac.interface_name = "func_0"} {
// CHECK: tac.cost = 7.864320e+05
%0 = "tfl.add"(%arg0, %arg1) {fused_activation_function = "RELU", tac.device = "CPU"} : (tensor<256x32x32x3xf32>, tensor<256x32x32x3xf32>) -> tensor<256x32x32x3xf32>
func.return %0 : tensor<256x32x32x3xf32>
}
func.func @func_0_GPU(%arg0: tensor<256x32x32x3xf32>, %arg1: tensor<256x32x32x3xf32>) -> tensor<256x32x32x3xf32> attributes {tac.device = "GPU", tac.interface_name = "func_0"} {
// CHECK: tac.cost = 157286.40
%0 = "tfl.add"(%arg0, %arg1) {fused_activation_function = "RELU", tac.device = "GPU"} : (tensor<256x32x32x3xf32>, tensor<256x32x32x3xf32>) -> tensor<256x32x32x3xf32>
func.return %0 : tensor<256x32x32x3xf32>
}
// -----
func.func @func_0_CPU(%arg0: tensor<10x10x10xf32>, %arg1: tensor<10xf32>) -> tensor<10x10x10xf32> attributes {tac.device = "CPU", tac.interface_name = "func_0"} {
// CHECK: tac.cost = 1.000000e+03
%0 = "tfl.add"(%arg0, %arg1) {fused_activation_function = "RELU", tac.device = "CPU"} : (tensor<10x10x10xf32>, tensor<10xf32>) -> tensor<10x10x10xf32>
func.return %0 : tensor<10x10x10xf32>
}
// -----
func.func @func_0_CPU(%arg0: tensor<10x10x10xf32>, %arg1: tensor<10xf32>) -> tensor<10x10x10xf32> attributes {tac.device = "CPU", tac.interface_name = "func_0"} {
// CHECK: tac.cost = 1.000000e+03
%0 = "tfl.add"(%arg0, %arg1) {fused_activation_function = "RELU", tac.device = "CPU"} : (tensor<10x10x10xf32>, tensor<10xf32>) -> tensor<10x10x10xf32>
// CHECK: tac.cost = 1.000000e+03
%1 = "tfl.mul"(%0, %arg1) {fused_activation_function = "RELU", tac.device = "CPU"} : (tensor<10x10x10xf32>, tensor<10xf32>) -> tensor<10x10x10xf32>
func.return %1 : tensor<10x10x10xf32>
}
// -----
func.func @pack_CPU(%arg0: tensor<100xf32>, %arg1: tensor<100xf32>) -> tensor<2x100xf32> attributes {tac.device = "CPU", tac.interface_name = "func_2"} {
// CHECK: tac.cost = 1.000000e+02
%0 = "tfl.pack"(%arg0, %arg1) {axis = 0 : i32, tac.device = "CPU", values_count = 2 : i32} : (tensor<100xf32>, tensor<100xf32>) -> tensor<2x100xf32>
func.return %0 : tensor<2x100xf32>
}
func.func @concat_reshape_GPU(%arg0: tensor<100xf32>, %arg1: tensor<100xf32>) -> tensor<2x100xf32> attributes {tac.device = "GPU", tac.interface_name = "func_2"} {
%cst = arith.constant dense<[2, 100]> : tensor<2xi64>
// CHECK: tac.cost = 4.000000e+01
%0 = "tfl.concatenation"(%arg0, %arg1) {axis = 0 : i32, fused_activation_function = "NONE", tac.device = "GPU"} : (tensor<100xf32>, tensor<100xf32>) -> tensor<200xf32>
// CHECK: tac.cost = 4.040000e+01
%1 = "tfl.reshape"(%0, %cst) {tac.device = "GPU"} : (tensor<200xf32>, tensor<2xi64>) -> tensor<2x100xf32>
func.return %1 : tensor<2x100xf32>
}
func.func @concat_reshape_CPU(%arg0: tensor<100xf32>, %arg1: tensor<100xf32>) -> tensor<2x100xf32> attributes {tac.device = "CPU", tac.interface_name = "func_2"} {
%cst = arith.constant dense<[2, 100]> : tensor<2xi64>
// CHECK: tac.cost = 1.000000e+02
%0 = "tfl.concatenation"(%arg0, %arg1) {axis = 0 : i32, fused_activation_function = "NONE", tac.device = "CPU"} : (tensor<100xf32>, tensor<100xf32>) -> tensor<200xf32>
// CHECK: tac.cost = 1.010000e+02
%1 = "tfl.reshape"(%0, %cst) {tac.device = "CPU"} : (tensor<200xf32>, tensor<2xi64>) -> tensor<2x100xf32>
func.return %1 : tensor<2x100xf32>
}
// -----
func.func @testConv2DCPU(tensor<256x32x32x3xf32>, tensor<16x3x3x3xf32>, tensor<16xf32>) -> tensor<256x32x32x16xf32> {
^bb0(%arg0: tensor<256x32x32x3xf32>, %arg1: tensor<16x3x3x3xf32>, %arg2: tensor<16xf32>):
// CHECK: tac.cost = 0x4D5C0000
%0 = "tfl.conv_2d"(%arg0, %arg1, %arg2) {dilation_h_factor = 1 : i32, dilation_w_factor = 1 : i32, padding = "SAME", stride_h = 1 : i32, stride_w = 1 : i32, fused_activation_function = "RELU6", tac.device = "CPU"} : (tensor<256x32x32x3xf32>, tensor<16x3x3x3xf32>, tensor<16xf32>) -> tensor<256x32x32x16xf32>
func.return %0 : tensor<256x32x32x16xf32>
}
func.func @testConv2DGPU(tensor<256x32x32x3xf32>, tensor<16x3x3x3xf32>, tensor<16xf32>) -> tensor<256x32x32x16xf32> {
^bb0(%arg0: tensor<256x32x32x3xf32>, %arg1: tensor<16x3x3x3xf32>, %arg2: tensor<16xf32>):
// CHECK: tac.cost = 0x4C300000
%0 = "tfl.conv_2d"(%arg0, %arg1, %arg2) {dilation_h_factor = 1 : i32, dilation_w_factor = 1 : i32, padding = "SAME", stride_h = 1 : i32, stride_w = 1 : i32, fused_activation_function = "RELU6", tac.device = "GPU"} : (tensor<256x32x32x3xf32>, tensor<16x3x3x3xf32>, tensor<16xf32>) -> tensor<256x32x32x16xf32>
func.return %0 : tensor<256x32x32x16xf32>
}
// -----
func.func @testConv2DNoBiasCPU(%arg0: tensor<128x32x32x3xf32>, %arg1: tensor<64x3x3x3xf32>, %arg2: none) -> tensor<128x32x32x64xf32> {
// CHECK: tac.cost = 0x4DD80000
%0 = "tfl.conv_2d"(%arg0, %arg1, %arg2) {dilation_h_factor = 1 : i32, dilation_w_factor = 1 : i32, padding = "SAME", stride_h = 1 : i32, stride_w = 1 : i32, fused_activation_function = "RELU6", tac.device = "CPU"} : (tensor<128x32x32x3xf32>, tensor<64x3x3x3xf32>, none) -> tensor<128x32x32x64xf32>
func.return %0 : tensor<128x32x32x64xf32>
}
func.func @testConv2DNoBiasGPU(%arg0: tensor<128x32x32x3xf32>, %arg1: tensor<64x3x3x3xf32>, %arg2: none) -> tensor<128x32x32x64xf32> {
// CHECK: tac.cost = 0x4CACCCCD
%0 = "tfl.conv_2d"(%arg0, %arg1, %arg2) {dilation_h_factor = 1 : i32, dilation_w_factor = 1 : i32, padding = "SAME", stride_h = 1 : i32, stride_w = 1 : i32, fused_activation_function = "RELU6", tac.device = "GPU"} : (tensor<128x32x32x3xf32>, tensor<64x3x3x3xf32>, none) -> tensor<128x32x32x64xf32>
func.return %0 : tensor<128x32x32x64xf32>
}
tensorflow/compiler/mlir/lite/experimental/tac/tests/pick-subgraphs.mlir
+204
View File
@@ -0,0 +1,204 @@
// 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: tac-opt-all-backends -tfl-pick-subgraphs %s -split-input-file -verify-diagnostics | FileCheck %s
module {
func.func @main(%arg0: tensor<100xf32>, %arg1: tensor<100xf32>, %arg2: tensor<100xf32>, %arg3: tensor<100xf32>) -> tensor<2x100xf32> {
%0 = func.call @func_0_GPU_FLOAT(%arg0, %arg1, %arg2) {tac.device = "GPU", tac.inference_type = "FLOAT", tac.interface_name = "func_0"} : (tensor<100xf32>, tensor<100xf32>, tensor<100xf32>) -> tensor<100xf32>
%1 = func.call @func_1_GPU_FLOAT(%arg0, %arg3) {tac.device = "GPU", tac.inference_type = "FLOAT", tac.interface_name = "func_1"} : (tensor<100xf32>, tensor<100xf32>) -> tensor<100xf32>
%2 = func.call @func_2_CPU_FLOAT(%0, %1) {tac.device = "CPU", tac.inference_type = "FLOAT", tac.interface_name = "func_2"} : (tensor<100xf32>, tensor<100xf32>) -> tensor<2x100xf32>
func.return %2 : tensor<2x100xf32>
}
func.func @func_2_CPU_FLOAT(%arg0: tensor<100xf32>, %arg1: tensor<100xf32>) -> tensor<2x100xf32> attributes {tac.cost = 2.000000e+01 : f32, tac.device = "CPU", tac.inference_type = "FLOAT", tac.interface_name = "func_2"} {
%0 = "tfl.pack"(%arg0, %arg1) {axis = 0 : i32, tac.device = "CPU", values_count = 2 : i32} : (tensor<100xf32>, tensor<100xf32>) -> tensor<2x100xf32>
func.return %0 : tensor<2x100xf32>
}
func.func @func_0_GPU_FLOAT(%arg0: tensor<100xf32>, %arg1: tensor<100xf32>, %arg2: tensor<100xf32>) -> tensor<100xf32> attributes {tac.cost = 4.000000e+01 : f32, tac.device = "GPU", tac.inference_type = "FLOAT", tac.interface_name = "func_0"} {
%0 = tfl.add %arg0, %arg1 {fused_activation_function = "RELU6", tac.device = "GPU"} : tensor<100xf32>
%1 = tfl.mul %0, %arg2 {fused_activation_function = "RELU6", tac.device = "GPU"} : tensor<100xf32>
func.return %1 : tensor<100xf32>
}
func.func @func_1_GPU_FLOAT(%arg0: tensor<100xf32>, %arg1: tensor<100xf32>) -> tensor<100xf32> attributes {tac.cost = 2.000000e+01 : f32, tac.device = "GPU", tac.inference_type = "FLOAT", tac.interface_name = "func_1"} {
%0 = tfl.add %arg0, %arg1 {fused_activation_function = "RELU6", tac.device = "GPU"} : tensor<100xf32>
func.return %0 : tensor<100xf32>
}
func.func @func_2_GPU_FLOAT(%arg0: tensor<100xf32>, %arg1: tensor<100xf32>) -> tensor<2x100xf32> attributes {tac.cost = 8.040000e+01 : f32, tac.device = "GPU", tac.inference_type = "FLOAT", tac.interface_name = "func_2"} {
%cst = arith.constant dense<[2, 100]> : tensor<2xi64>
%0 = "tfl.concatenation"(%arg0, %arg1) {axis = 0 : i32, fused_activation_function = "NONE"} : (tensor<100xf32>, tensor<100xf32>) -> tensor<200xf32>
%1 = "tfl.reshape"(%0, %cst) : (tensor<200xf32>, tensor<2xi64>) -> tensor<2x100xf32>
func.return %1 : tensor<2x100xf32>
}
func.func @func_0_CPU_FLOAT(%arg0: tensor<100xf32>, %arg1: tensor<100xf32>, %arg2: tensor<100xf32>) -> tensor<100xf32> attributes {tac.cost = 2.000000e+02 : f32, tac.device = "CPU", tac.inference_type = "FLOAT", tac.interface_name = "func_0"} {
%0 = tfl.add %arg0, %arg1 {fused_activation_function = "RELU6", tac.device = "GPU"} : tensor<100xf32>
%1 = tfl.mul %0, %arg2 {fused_activation_function = "RELU6", tac.device = "GPU"} : tensor<100xf32>
func.return %1 : tensor<100xf32>
}
func.func @func_1_CPU_FLOAT(%arg0: tensor<100xf32>, %arg1: tensor<100xf32>) -> tensor<100xf32> attributes {tac.cost = 1.000000e+02 : f32, tac.device = "CPU", tac.inference_type = "FLOAT", tac.interface_name = "func_1"} {
%0 = tfl.add %arg0, %arg1 {fused_activation_function = "RELU6", tac.device = "GPU"} : tensor<100xf32>
func.return %0 : tensor<100xf32>
}
// CHECK: func @main([[VAL_0:%.*]]: tensor<100xf32>, [[VAL_1:%.*]]: tensor<100xf32>, [[VAL_2:%.*]]: tensor<100xf32>, [[VAL_3:%.*]]: tensor<100xf32>) -> tensor<2x100xf32> {
// CHECK: [[VAL_4:%.*]] = call @func_0_GPU_FLOAT([[VAL_0]], [[VAL_1]], [[VAL_2]]) {tac.device = "GPU", tac.inference_type = "FLOAT", tac.interface_name = "func_0"} : (tensor<100xf32>, tensor<100xf32>, tensor<100xf32>) -> tensor<100xf32>
// CHECK: [[VAL_5:%.*]] = call @func_1_GPU_FLOAT([[VAL_0]], [[VAL_3]]) {tac.device = "GPU", tac.inference_type = "FLOAT", tac.interface_name = "func_1"} : (tensor<100xf32>, tensor<100xf32>) -> tensor<100xf32>
// CHECK: [[VAL_6:%.*]] = call @func_2_GPU_FLOAT([[VAL_4]], [[VAL_5]]) {tac.device = "GPU", tac.inference_type = "FLOAT", tac.interface_name = "func_2"} : (tensor<100xf32>, tensor<100xf32>) -> tensor<2x100xf32>
// CHECK: return [[VAL_6]] : tensor<2x100xf32>
// CHECK: }
}
// -----
module {
func.func @main(%arg0: tensor<1x200x200x200xf32>) -> tensor<2x1x200x200x200xf32> attributes {tf.entry_function = {inputs = "Placeholder", outputs = "mul_1"}} {
%0 = "tfl.pseudo_const"() {value = dense<0.962260901> : tensor<1xf32>} : () -> tensor<1xf32>
%1 = func.call @func_0_GPU_FLOAT(%arg0, %0) {tac.device = "GPU", tac.inference_type = "FLOAT", tac.interface_name = "func_0"} : (tensor<1x200x200x200xf32>, tensor<1xf32>) -> tensor<1x200x200x200xf32>
%2 = "tfl.pseudo_const"() {value = dense<0.895973444> : tensor<1xf32>} : () -> tensor<1xf32>
%3 = func.call @func_1_GPU_FLOAT(%arg0, %2) {tac.device = "GPU", tac.inference_type = "FLOAT", tac.interface_name = "func_1"} : (tensor<1x200x200x200xf32>, tensor<1xf32>) -> tensor<1x200x200x200xf32>
%4 = func.call @func_2_CPU_FLOAT(%3, %1) {tac.device = "CPU", tac.inference_type = "FLOAT", tac.interface_name = "func_2"} : (tensor<1x200x200x200xf32>, tensor<1x200x200x200xf32>) -> tensor<2x1x200x200x200xf32>
%5 = "tfl.pseudo_const"() {value = dense<0.0778453499> : tensor<1xf32>} : () -> tensor<1xf32>
%6 = func.call @func_3_GPU_FLOAT(%4, %5) {tac.device = "GPU", tac.inference_type = "FLOAT", tac.interface_name = "func_3"} : (tensor<2x1x200x200x200xf32>, tensor<1xf32>) -> tensor<2x1x200x200x200xf32>
func.return %6 : tensor<2x1x200x200x200xf32>
}
func.func @func_2_CPU_FLOAT(%arg0: tensor<1x200x200x200xf32>, %arg1: tensor<1x200x200x200xf32>) -> tensor<2x1x200x200x200xf32> attributes {tac.cost = 1.600000e+06 : f32, tac.device = "CPU", tac.inference_type = "FLOAT", tac.interface_name = "func_2"} {
%0 = "tfl.pack"(%arg0, %arg1) {axis = 0 : i32, tac.device = "CPU", values_count = 2 : i32} : (tensor<1x200x200x200xf32>, tensor<1x200x200x200xf32>) -> tensor<2x1x200x200x200xf32>
func.return %0 : tensor<2x1x200x200x200xf32>
}
func.func @func_0_GPU_FLOAT(%arg0: tensor<1x200x200x200xf32>, %arg1: tensor<1xf32>) -> tensor<1x200x200x200xf32> attributes {tac.cost = 1.600000e+06 : f32, tac.device = "GPU", tac.inference_type = "FLOAT", tac.interface_name = "func_0"} {
%0 = "tfl.add"(%arg0, %arg1) {fused_activation_function = "NONE", tac.device = "GPU"} : (tensor<1x200x200x200xf32>, tensor<1xf32>) -> tensor<1x200x200x200xf32>
func.return %0 : tensor<1x200x200x200xf32>
}
func.func @func_1_GPU_FLOAT(%arg0: tensor<1x200x200x200xf32>, %arg1: tensor<1xf32>) -> tensor<1x200x200x200xf32> attributes {tac.cost = 1.600000e+06 : f32, tac.device = "GPU", tac.inference_type = "FLOAT", tac.interface_name = "func_1"} {
%0 = "tfl.mul"(%arg0, %arg1) {fused_activation_function = "NONE", tac.device = "GPU"} : (tensor<1x200x200x200xf32>, tensor<1xf32>) -> tensor<1x200x200x200xf32>
func.return %0 : tensor<1x200x200x200xf32>
}
func.func @func_3_GPU_FLOAT(%arg0: tensor<2x1x200x200x200xf32>, %arg1: tensor<1xf32>) -> tensor<2x1x200x200x200xf32> attributes {tac.cost = 3.200000e+06 : f32, tac.device = "GPU", tac.inference_type = "FLOAT", tac.interface_name = "func_3"} {
%0 = "tfl.mul"(%arg0, %arg1) {fused_activation_function = "NONE", tac.device = "GPU"} : (tensor<2x1x200x200x200xf32>, tensor<1xf32>) -> tensor<2x1x200x200x200xf32>
func.return %0 : tensor<2x1x200x200x200xf32>
}
func.func @func_2_GPU_FLOAT(%arg0: tensor<1x200x200x200xf32>, %arg1: tensor<1x200x200x200xf32>) -> tensor<2x1x200x200x200xf32> attributes {tac.cost = 0x4AC35002 : f32, tac.device = "GPU", tac.inference_type = "FLOAT", tac.interface_name = "func_2"} {
%cst = arith.constant dense<[2, 1, 200, 200, 200]> : tensor<5xi64>
%0 = "tfl.concatenation"(%arg0, %arg1) {axis = 0 : i32, fused_activation_function = "NONE"} : (tensor<1x200x200x200xf32>, tensor<1x200x200x200xf32>) -> tensor<2x200x200x200xf32>
%1 = "tfl.reshape"(%0, %cst) : (tensor<2x200x200x200xf32>, tensor<5xi64>) -> tensor<2x1x200x200x200xf32>
func.return %1 : tensor<2x1x200x200x200xf32>
}
func.func @func_0_CPU_FLOAT(%arg0: tensor<1x200x200x200xf32>, %arg1: tensor<1xf32>) -> tensor<1x200x200x200xf32> attributes {tac.cost = 8.000000e+06 : f32, tac.device = "CPU", tac.inference_type = "FLOAT", tac.interface_name = "func_0"} {
%0 = "tfl.add"(%arg0, %arg1) {fused_activation_function = "NONE", tac.device = "GPU"} : (tensor<1x200x200x200xf32>, tensor<1xf32>) -> tensor<1x200x200x200xf32>
func.return %0 : tensor<1x200x200x200xf32>
}
func.func @func_1_CPU_FLOAT(%arg0: tensor<1x200x200x200xf32>, %arg1: tensor<1xf32>) -> tensor<1x200x200x200xf32> attributes {tac.cost = 8.000000e+06 : f32, tac.device = "CPU", tac.inference_type = "FLOAT", tac.interface_name= "func_1"} {
%0 = "tfl.mul"(%arg0, %arg1) {fused_activation_function = "NONE", tac.device = "GPU"} : (tensor<1x200x200x200xf32>, tensor<1xf32>) -> tensor<1x200x200x200xf32>
func.return %0 : tensor<1x200x200x200xf32>
}
func.func @func_3_CPU_FLOAT(%arg0: tensor<2x1x200x200x200xf32>, %arg1: tensor<1xf32>) -> tensor<2x1x200x200x200xf32> attributes {tac.cost = 1.600000e+07 : f32, tac.device = "CPU", tac.inference_type = "FLOAT", tac.interface_name = "func_3"} {
%0 = "tfl.mul"(%arg0, %arg1) {fused_activation_function = "NONE", tac.device = "GPU"} : (tensor<2x1x200x200x200xf32>, tensor<1xf32>) -> tensor<2x1x200x200x200xf32>
func.return %0 : tensor<2x1x200x200x200xf32>
}
// CHECK: func @main([[VAL_0:%.*]]: tensor<1x200x200x200xf32>) -> tensor<2x1x200x200x200xf32> attributes {tf.entry_function = {inputs = "Placeholder", outputs = "mul_1"}} {
// CHECK: [[VAL_1:%.*]] = "tfl.pseudo_const"() <{value = dense<0.962260901> : tensor<1xf32>}> : () -> tensor<1xf32>
// CHECK: [[VAL_2:%.*]] = call @func_0_GPU_FLOAT([[VAL_0]], [[VAL_1]]) {tac.device = "GPU", tac.inference_type = "FLOAT", tac.interface_name = "func_0"} : (tensor<1x200x200x200xf32>, tensor<1xf32>) -> tensor<1x200x200x200xf32>
// CHECK: [[VAL_3:%.*]] = "tfl.pseudo_const"() <{value = dense<0.895973444> : tensor<1xf32>}> : () -> tensor<1xf32>
// CHECK: [[VAL_4:%.*]] = call @func_1_GPU_FLOAT([[VAL_0]], [[VAL_3]]) {tac.device = "GPU", tac.inference_type = "FLOAT", tac.interface_name = "func_1"} : (tensor<1x200x200x200xf32>, tensor<1xf32>) -> tensor<1x200x200x200xf32>
// CHECK: [[VAL_5:%.*]] = call @func_2_GPU_FLOAT([[VAL_4]], [[VAL_2]]) {tac.device = "GPU", tac.inference_type = "FLOAT", tac.interface_name = "func_2"} : (tensor<1x200x200x200xf32>, tensor<1x200x200x200xf32>) -> tensor<2x1x200x200x200xf32>
// CHECK: [[VAL_6:%.*]] = "tfl.pseudo_const"() <{value = dense<0.0778453499> : tensor<1xf32>}> : () -> tensor<1xf32>
// CHECK: [[VAL_7:%.*]] = call @func_3_GPU_FLOAT([[VAL_5]], [[VAL_6]]) {tac.device = "GPU", tac.inference_type = "FLOAT", tac.interface_name = "func_3"} : (tensor<2x1x200x200x200xf32>, tensor<1xf32>) -> tensor<2x1x200x200x200xf32>
// CHECK: return [[VAL_7]] : tensor<2x1x200x200x200xf32>
// CHECK: }
}
// -----
module {
func.func @main(%arg0: tensor<100x!quant.uniform<i8:f32, 2.000000e-01:-3>>, %arg1: tensor<100x!quant.uniform<i8:f32, 2.000000e-01:-3>>, %arg2: tensor<100x!quant.uniform<i8:f32, 2.000000e-01:-3>>, %arg3: tensor<100x!quant.uniform<i8:f32, 2.000000e-01:-3>>) -> tensor<2x100x!quant.uniform<i8:f32, 2.000000e-01:-3>> attributes {tf.entry_function = {inputs = "input0,input1,input2,input3", outputs = "output"}} {
%0 = func.call @func_0_CPU_QUANTIZED_INT8(%arg0, %arg1, %arg2) {tac.device = "CPU", tac.inference_type = "QUANTIZED_INT8", tac.interface_name = "func_0"} : (tensor<100x!quant.uniform<i8:f32, 2.000000e-01:-3>>, tensor<100x!quant.uniform<i8:f32, 2.000000e-01:-3>>, tensor<100x!quant.uniform<i8:f32, 2.000000e-01:-3>>) -> tensor<100x!quant.uniform<i8:f32, 2.000000e-01:-3>>
%1 = func.call @func_1_CPU_QUANTIZED_INT8(%arg0, %arg3) {tac.device = "CPU", tac.inference_type = "QUANTIZED_INT8", tac.interface_name = "func_1"} : (tensor<100x!quant.uniform<i8:f32, 2.000000e-01:-3>>, tensor<100x!quant.uniform<i8:f32, 2.000000e-01:-3>>) -> tensor<100x!quant.uniform<i8:f32, 2.000000e-01:-3>>
%2 = func.call @func_2_CPU_QUANTIZED_INT8(%0, %1) {tac.device = "CPU", tac.inference_type = "QUANTIZED_INT8", tac.interface_name = "func_2"} : (tensor<100x!quant.uniform<i8:f32, 2.000000e-01:-3>>, tensor<100x!quant.uniform<i8:f32, 2.000000e-01:-3>>) -> tensor<2x100x!quant.uniform<i8:f32, 2.000000e-01:-3>>
func.return %2 : tensor<2x100x!quant.uniform<i8:f32, 2.000000e-01:-3>>
}
func.func private @func_0_CPU_QUANTIZED_INT8(%arg0: tensor<100x!quant.uniform<i8:f32, 2.000000e-01:-3>>, %arg1: tensor<100x!quant.uniform<i8:f32, 2.000000e-01:-3>>, %arg2: tensor<100x!quant.uniform<i8:f32, 2.000000e-01:-3>>) -> tensor<100x!quant.uniform<i8:f32, 2.000000e-01:-3>> attributes {tac.cost = 2.000000e+02 : f32, tac.device = "CPU", tac.inference_type = "QUANTIZED_INT8", tac.interface_name = "func_0"} {
%0 = tfl.add %arg0, %arg1 {fused_activation_function = "RELU6", tac.device = "CPU", tac.inference_type = "QUANTIZED_INT8"} : tensor<100x!quant.uniform<i8:f32, 2.000000e-01:-3>>
%1 = tfl.mul %0, %arg2 {fused_activation_function = "RELU6", tac.device = "CPU", tac.inference_type = "QUANTIZED_INT8"} : tensor<100x!quant.uniform<i8:f32, 2.000000e-01:-3>>
func.return %1 : tensor<100x!quant.uniform<i8:f32, 2.000000e-01:-3>>
}
func.func private @func_2_CPU_QUANTIZED_INT8(%arg0: tensor<100x!quant.uniform<i8:f32, 2.000000e-01:-3>>, %arg1: tensor<100x!quant.uniform<i8:f32, 2.000000e-01:-3>>) -> tensor<2x100x!quant.uniform<i8:f32, 2.000000e-01:-3>> attributes {tac.cost = 1.000000e+02 : f32, tac.device = "CPU", tac.inference_type = "QUANTIZED_INT8", tac.interface_name = "func_2"} {
%0 = "tfl.pack"(%arg0, %arg1) {axis = 0 : i32, tac.device = "CPU", tac.inference_type = "QUANTIZED_INT8", values_count = 2 : i32} : (tensor<100x!quant.uniform<i8:f32, 2.000000e-01:-3>>, tensor<100x!quant.uniform<i8:f32, 2.000000e-01:-3>>) -> tensor<2x100x!quant.uniform<i8:f32, 2.000000e-01:-3>>
func.return %0 : tensor<2x100x!quant.uniform<i8:f32, 2.000000e-01:-3>>
}
func.func private @func_1_CPU_QUANTIZED_INT8(%arg0: tensor<100x!quant.uniform<i8:f32, 2.000000e-01:-3>>, %arg1: tensor<100x!quant.uniform<i8:f32, 2.000000e-01:-3>>) -> tensor<100x!quant.uniform<i8:f32, 2.000000e-01:-3>> attributes {tac.cost = 1.000000e+02 : f32, tac.device = "CPU", tac.inference_type = "QUANTIZED_INT8", tac.interface_name = "func_1"} {
%0 = tfl.add %arg0, %arg1 {fused_activation_function = "RELU6", tac.device = "CPU", tac.inference_type = "QUANTIZED_INT8"} : tensor<100x!quant.uniform<i8:f32, 2.000000e-01:-3>>
func.return %0 : tensor<100x!quant.uniform<i8:f32, 2.000000e-01:-3>>
}
func.func private @func_0_GPU_FLOAT(%arg0: tensor<100x!quant.uniform<i8:f32, 2.000000e-01:-3>>, %arg1: tensor<100x!quant.uniform<i8:f32, 2.000000e-01:-3>>, %arg2: tensor<100x!quant.uniform<i8:f32, 2.000000e-01:-3>>) -> tensor<100x!quant.uniform<i8:f32, 2.000000e-01:-3>> attributes {tac.cost = 4.000000e+01 : f32, tac.device = "GPU", tac.inference_type = "FLOAT", tac.interface_name = "func_0"} {
%0 = "tfl.dequantize"(%arg0) {tac.device = "GPU", tac.inference_type = "FLOAT"} : (tensor<100x!quant.uniform<i8:f32, 2.000000e-01:-3>>) -> tensor<100xf32>
%1 = "tfl.dequantize"(%arg1) {tac.device = "GPU", tac.inference_type = "FLOAT"} : (tensor<100x!quant.uniform<i8:f32, 2.000000e-01:-3>>) -> tensor<100xf32>
%2 = tfl.add %0, %1 {fused_activation_function = "RELU6", tac.device = "GPU", tac.inference_type = "FLOAT"} : tensor<100xf32>
%3 = "tfl.dequantize"(%arg2) {tac.device = "GPU", tac.inference_type = "FLOAT"} : (tensor<100x!quant.uniform<i8:f32, 2.000000e-01:-3>>) -> tensor<100xf32>
%4 = tfl.mul %2, %3 {fused_activation_function = "RELU6", tac.device = "GPU", tac.inference_type = "FLOAT"} : tensor<100xf32>
%5 = "tfl.quantize"(%4) {qtype = tensor<100x!quant.uniform<i8:f32, 2.000000e-01:-3>>, tac.device = "GPU", tac.inference_type = "FLOAT"} : (tensor<100xf32>) -> tensor<100x!quant.uniform<i8:f32, 2.000000e-01:-3>>
func.return %5 : tensor<100x!quant.uniform<i8:f32, 2.000000e-01:-3>>
}
func.func private @func_0_CPU_FLOAT(%arg0: tensor<100x!quant.uniform<i8:f32, 2.000000e-01:-3>>, %arg1: tensor<100x!quant.uniform<i8:f32, 2.000000e-01:-3>>, %arg2: tensor<100x!quant.uniform<i8:f32, 2.000000e-01:-3>>) -> tensor<100x!quant.uniform<i8:f32, 2.000000e-01:-3>> attributes {tac.cost = 2.000000e+02 : f32, tac.device = "CPU", tac.inference_type = "FLOAT", tac.interface_name = "func_0"} {
%0 = "tfl.dequantize"(%arg0) {tac.device = "CPU", tac.inference_type = "FLOAT"} : (tensor<100x!quant.uniform<i8:f32, 2.000000e-01:-3>>) -> tensor<100xf32>
%1 = "tfl.dequantize"(%arg1) {tac.device = "CPU", tac.inference_type = "FLOAT"} : (tensor<100x!quant.uniform<i8:f32, 2.000000e-01:-3>>) -> tensor<100xf32>
%2 = tfl.add %0, %1 {fused_activation_function = "RELU6", tac.device = "CPU", tac.inference_type = "FLOAT"} : tensor<100xf32>
%3 = "tfl.dequantize"(%arg2) {tac.device = "CPU", tac.inference_type = "FLOAT"} : (tensor<100x!quant.uniform<i8:f32, 2.000000e-01:-3>>) -> tensor<100xf32>
%4 = tfl.mul %2, %3 {fused_activation_function = "RELU6", tac.device = "CPU", tac.inference_type = "FLOAT"} : tensor<100xf32>
%5 = "tfl.quantize"(%4) {qtype = tensor<100x!quant.uniform<i8:f32, 2.000000e-01:-3>>, tac.device = "CPU", tac.inference_type = "FLOAT"} : (tensor<100xf32>) -> tensor<100x!quant.uniform<i8:f32, 2.000000e-01:-3>>
func.return %5 : tensor<100x!quant.uniform<i8:f32, 2.000000e-01:-3>>
}
func.func private @func_2_GPU_FLOAT(%arg0: tensor<100x!quant.uniform<i8:f32, 2.000000e-01:-3>>, %arg1: tensor<100x!quant.uniform<i8:f32, 2.000000e-01:-3>>) -> tensor<2x100x!quant.uniform<i8:f32, 2.000000e-01:-3>> attributes {tac.cost = 162.200012 : f32, tac.device = "GPU", tac.inference_type = "FLOAT", tac.interface_name = "func_2"} {
%cst = arith.constant dense<[1, 1, 1, 100]> : tensor<4xi32>
%cst_0 = arith.constant dense<200> : tensor<1xi32>
%cst_1 = arith.constant dense<[2, 100]> : tensor<2xi32>
%0 = "tfl.dequantize"(%arg0) {tac.device = "GPU", tac.inference_type = "FLOAT"} : (tensor<100x!quant.uniform<i8:f32, 2.000000e-01:-3>>) -> tensor<100xf32>
%1 = "tfl.dequantize"(%arg1) {tac.device = "GPU", tac.inference_type = "FLOAT"} : (tensor<100x!quant.uniform<i8:f32, 2.000000e-01:-3>>) -> tensor<100xf32>
%2 = "tfl.reshape"(%0, %cst) {tac.device = "GPU", tac.inference_type = "FLOAT"} : (tensor<100xf32>, tensor<4xi32>) -> tensor<1x1x1x100xf32>
%3 = "tfl.reshape"(%1, %cst) {tac.device = "GPU", tac.inference_type = "FLOAT"} : (tensor<100xf32>, tensor<4xi32>) -> tensor<1x1x1x100xf32>
%4 = "tfl.concatenation"(%2, %3) {axis = 3 : i32, fused_activation_function = "NONE", tac.device = "GPU", tac.inference_type = "FLOAT"} : (tensor<1x1x1x100xf32>, tensor<1x1x1x100xf32>) -> tensor<1x1x1x200xf32>
%5 = "tfl.reshape"(%4, %cst_0) {tac.device = "GPU", tac.inference_type = "FLOAT"} : (tensor<1x1x1x200xf32>, tensor<1xi32>) -> tensor<200xf32>
%6 = "tfl.reshape"(%5, %cst_1) {tac.device = "GPU", tac.inference_type = "FLOAT"} : (tensor<200xf32>, tensor<2xi32>) -> tensor<2x100xf32>
%7 = "tfl.quantize"(%6) {qtype = tensor<2x100x!quant.uniform<i8:f32, 2.000000e-01:-3>>, tac.device = "GPU", tac.inference_type = "FLOAT"} : (tensor<2x100xf32>) -> tensor<2x100x!quant.uniform<i8:f32, 2.000000e-01:-3>>
func.return %7 : tensor<2x100x!quant.uniform<i8:f32, 2.000000e-01:-3>>
}
func.func private @func_2_CPU_FLOAT(%arg0: tensor<100x!quant.uniform<i8:f32, 2.000000e-01:-3>>, %arg1: tensor<100x!quant.uniform<i8:f32, 2.000000e-01:-3>>) -> tensor<2x100x!quant.uniform<i8:f32, 2.000000e-01:-3>> attributes {tac.cost = 1.000000e+02 : f32, tac.device = "CPU", tac.inference_type = "FLOAT", tac.interface_name = "func_2"} {
%0 = "tfl.dequantize"(%arg0) {tac.device = "CPU", tac.inference_type = "FLOAT"} : (tensor<100x!quant.uniform<i8:f32, 2.000000e-01:-3>>) -> tensor<100xf32>
%1 = "tfl.dequantize"(%arg1) {tac.device = "CPU", tac.inference_type = "FLOAT"} : (tensor<100x!quant.uniform<i8:f32, 2.000000e-01:-3>>) -> tensor<100xf32>
%2 = "tfl.pack"(%0, %1) {axis = 0 : i32, tac.device = "CPU", tac.inference_type = "FLOAT", values_count = 2 : i32} : (tensor<100xf32>, tensor<100xf32>) -> tensor<2x100xf32>
%3 = "tfl.quantize"(%2) {qtype = tensor<2x100x!quant.uniform<i8:f32, 2.000000e-01:-3>>, tac.device = "CPU", tac.inference_type = "FLOAT"} : (tensor<2x100xf32>) -> tensor<2x100x!quant.uniform<i8:f32, 2.000000e-01:-3>>
func.return %3 : tensor<2x100x!quant.uniform<i8:f32, 2.000000e-01:-3>>
}
func.func private @func_1_GPU_FLOAT(%arg0: tensor<100x!quant.uniform<i8:f32, 2.000000e-01:-3>>, %arg1: tensor<100x!quant.uniform<i8:f32, 2.000000e-01:-3>>) -> tensor<100x!quant.uniform<i8:f32, 2.000000e-01:-3>> attributes {tac.cost = 2.000000e+01 : f32, tac.device = "GPU", tac.inference_type = "FLOAT", tac.interface_name = "func_1"} {
%0 = "tfl.dequantize"(%arg0) {tac.device = "GPU", tac.inference_type = "FLOAT"} : (tensor<100x!quant.uniform<i8:f32, 2.000000e-01:-3>>) -> tensor<100xf32>
%1 = "tfl.dequantize"(%arg1) {tac.device = "GPU", tac.inference_type = "FLOAT"} : (tensor<100x!quant.uniform<i8:f32, 2.000000e-01:-3>>) -> tensor<100xf32>
%2 = tfl.add %0, %1 {fused_activation_function = "RELU6", tac.device = "GPU", tac.inference_type = "FLOAT"} : tensor<100xf32>
%3 = "tfl.quantize"(%2) {qtype = tensor<100x!quant.uniform<i8:f32, 2.000000e-01:-3>>, tac.device = "GPU", tac.inference_type = "FLOAT"} : (tensor<100xf32>) -> tensor<100x!quant.uniform<i8:f32, 2.000000e-01:-3>>
func.return %3 : tensor<100x!quant.uniform<i8:f32, 2.000000e-01:-3>>
}
func.func private @func_1_CPU_FLOAT(%arg0: tensor<100x!quant.uniform<i8:f32, 2.000000e-01:-3>>, %arg1: tensor<100x!quant.uniform<i8:f32, 2.000000e-01:-3>>) -> tensor<100x!quant.uniform<i8:f32, 2.000000e-01:-3>> attributes {tac.cost = 1.000000e+02 : f32, tac.device = "CPU", tac.inference_type = "FLOAT", tac.interface_name = "func_1"} {
%0 = "tfl.dequantize"(%arg0) {tac.device = "CPU", tac.inference_type = "FLOAT"} : (tensor<100x!quant.uniform<i8:f32, 2.000000e-01:-3>>) -> tensor<100xf32>
%1 = "tfl.dequantize"(%arg1) {tac.device = "CPU", tac.inference_type = "FLOAT"} : (tensor<100x!quant.uniform<i8:f32, 2.000000e-01:-3>>) -> tensor<100xf32>
%2 = tfl.add %0, %1 {fused_activation_function = "RELU6", tac.device = "CPU", tac.inference_type = "FLOAT"} : tensor<100xf32>
%3 = "tfl.quantize"(%2) {qtype = tensor<100x!quant.uniform<i8:f32, 2.000000e-01:-3>>, tac.device = "CPU", tac.inference_type = "FLOAT"} : (tensor<100xf32>) -> tensor<100x!quant.uniform<i8:f32, 2.000000e-01:-3>>
func.return %3 : tensor<100x!quant.uniform<i8:f32, 2.000000e-01:-3>>
}
// CHECK: @main(%[[VAL_0:.*]]: tensor<100x!quant.uniform<i8:f32, 2.000000e-01:-3>>, %[[VAL_1:.*]]: tensor<100x!quant.uniform<i8:f32, 2.000000e-01:-3>>, %[[VAL_2:.*]]: tensor<100x!quant.uniform<i8:f32, 2.000000e-01:-3>>, %[[VAL_3:.*]]: tensor<100x!quant.uniform<i8:f32, 2.000000e-01:-3>>) -> tensor<2x100x!quant.uniform<i8:f32, 2.000000e-01:-3>> attributes {tf.entry_function = {inputs = "input0,input1,input2,input3", outputs = "output"}} {
// CHECK: %[[VAL_4:.*]] = call @func_0_GPU_FLOAT(%[[VAL_0]], %[[VAL_1]], %[[VAL_2]]) {tac.device = "GPU", tac.inference_type = "FLOAT", tac.interface_name = "func_0"} : (tensor<100x!quant.uniform<i8:f32, 2.000000e-01:-3>>, tensor<100x!quant.uniform<i8:f32, 2.000000e-01:-3>>, tensor<100x!quant.uniform<i8:f32, 2.000000e-01:-3>>) -> tensor<100x!quant.uniform<i8:f32, 2.000000e-01:-3>>
// CHECK: %[[VAL_5:.*]] = call @func_1_GPU_FLOAT(%[[VAL_0]], %[[VAL_3]]) {tac.device = "GPU", tac.inference_type = "FLOAT", tac.interface_name = "func_1"} : (tensor<100x!quant.uniform<i8:f32, 2.000000e-01:-3>>, tensor<100x!quant.uniform<i8:f32, 2.000000e-01:-3>>) -> tensor<100x!quant.uniform<i8:f32, 2.000000e-01:-3>>
// CHECK: %[[VAL_6:.*]] = call @func_2_GPU_FLOAT(%[[VAL_4]], %[[VAL_5]]) {tac.device = "GPU", tac.inference_type = "FLOAT", tac.interface_name = "func_2"} : (tensor<100x!quant.uniform<i8:f32, 2.000000e-01:-3>>, tensor<100x!quant.uniform<i8:f32, 2.000000e-01:-3>>) -> tensor<2x100x!quant.uniform<i8:f32, 2.000000e-01:-3>>
// CHECK: return %[[VAL_6]] : tensor<2x100x!quant.uniform<i8:f32, 2.000000e-01:-3>>
// CHECK: }
}
tensorflow/compiler/mlir/lite/experimental/tac/tests/raise-target-subgraphs.mlir
+604
View File
@@ -0,0 +1,604 @@
// 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: tac-opt-all-backends -tfl-raise-target-subgraphs %s -split-input-file | FileCheck %s
// RUN: tac-opt-all-backends -tfl-raise-target-subgraphs="skip-raise-cpu-ops=true" %s -split-input-file | FileCheck %s --check-prefixes=CHECK-SKIP-CPU
// RUN: tac-opt-all-backends -tfl-raise-target-subgraphs="ignore-inference-type=true" %s -split-input-file | FileCheck %s --check-prefixes=CHECK-IGNORE-INFERENCE-TYPE
module {
func.func @simpleWhile(%arg0: tensor<i32>) -> tensor<i32> {
%0 = "tfl.while"(%arg0) ({
^bb0(%block: tensor<i32>):
"tfl.yield"(%block) : (tensor<i32>) -> ()
},{
^bb0(%block: tensor<i32>):
"tfl.yield"(%block) : (tensor<i32>) -> ()
}) {tac.device = "CPU", fused_activation_function = "RELU6", tac.inference_type = "FLOAT"} : (tensor<i32>) -> tensor<i32>
func.return %0 : tensor<i32>
}
}
// CHECK: func.func @simpleWhile(%arg0: tensor<i32>) -> tensor<i32> {
// CHECK: %0 = call @func_0_CPU_FLOAT(%arg0) {tac.device = "CPU", tac.inference_type = "FLOAT", tac.interface_name = "func_0"} : (tensor<i32>) -> tensor<i32>
// CHECK: return %0 : tensor<i32>
// CHECK: }
// CHECK: func.func private @func_0_CPU_FLOAT(%arg0: tensor<i32>) -> tensor<i32> attributes {tac.device = "CPU", tac.inference_type = "FLOAT", tac.interface_name = "func_0"} {
// CHECK: %0 = "tfl.while"(%arg0) ({
// CHECK: ^bb0(%arg1: tensor<i32>):
// CHECK: "tfl.yield"(%arg1) : (tensor<i32>) -> ()
// CHECK: }, {
// CHECK: ^bb0(%arg1: tensor<i32>):
// CHECK: "tfl.yield"(%arg1) : (tensor<i32>) -> ()
// CHECK: }) {fused_activation_function = "RELU6", tac.device = "CPU", tac.inference_type = "FLOAT"} : (tensor<i32>) -> tensor<i32>
// CHECK: return %0 : tensor<i32>
// CHECK: }
// -----
module {
func.func @whileWithNested(%arg0: tensor<i32>) -> tensor<i32> {
%0 = "tfl.while"(%arg0) ({
^bb0(%block: tensor<i32>):
%1 = "tfl.add"(%arg0, %arg0) { fused_activation_function = "NONE", tac.device = "CPU", tac.inference_type = "FLOAT"} : (tensor<i32>, tensor<i32>) -> (tensor<i32>)
%2 = "tfl.add"(%1, %1) { fused_activation_function = "NONE", tac.device = "GPU", tac.inference_type = "FLOAT"} : (tensor<i32>, tensor<i32>) -> (tensor<i32>)
"tfl.yield"(%2) : (tensor<i32>) -> ()
},{
^bb0(%block: tensor<i32>):
"tfl.yield"(%block) : (tensor<i32>) -> ()
}) {tac.device = "CPU", fused_activation_function = "RELU6", tac.inference_type = "FLOAT"} : (tensor<i32>) -> tensor<i32>
func.return %0 : tensor<i32>
}
}
// CHECK: func.func @whileWithNested(%arg0: tensor<i32>) -> tensor<i32> {
// CHECK: %0 = call @func_0_CPU_FLOAT(%arg0) {tac.device = "CPU", tac.inference_type = "FLOAT", tac.interface_name = "func_0"} : (tensor<i32>) -> tensor<i32>
// CHECK: return %0 : tensor<i32>
// CHECK: }
// CHECK: func.func private @func_0_CPU_FLOAT(%arg0: tensor<i32>) -> tensor<i32> attributes {tac.device = "CPU", tac.inference_type = "FLOAT", tac.interface_name = "func_0"} {
// CHECK: %0 = "tfl.while"(%arg0) ({
// CHECK: ^bb0(%arg1: tensor<i32>):
// CHECK: %1 = tfl.add %arg0, %arg0 {fused_activation_function = "NONE", tac.device = "CPU", tac.inference_type = "FLOAT"} : tensor<i32>
// CHECK: %2 = func.call @func_1_GPU_FLOAT(%1) {tac.device = "GPU", tac.inference_type = "FLOAT", tac.interface_name = "func_1"} : (tensor<i32>) -> tensor<i32>
// CHECK: "tfl.yield"(%2) : (tensor<i32>) -> ()
// CHECK: }, {
// CHECK: ^bb0(%arg1: tensor<i32>):
// CHECK: "tfl.yield"(%arg1) : (tensor<i32>) -> ()
// CHECK: }) {fused_activation_function = "RELU6", tac.device = "CPU", tac.inference_type = "FLOAT"} : (tensor<i32>) -> tensor<i32>
// CHECK: return %0 : tensor<i32>
// CHECK: }
// CHECK: func.func private @func_1_GPU_FLOAT(%arg0: tensor<i32>) -> tensor<i32> attributes {tac.device = "GPU", tac.inference_type = "FLOAT", tac.interface_name = "func_1"} {
// CHECK: %0 = tfl.add %arg0, %arg0 {fused_activation_function = "NONE", tac.device = "GPU", tac.inference_type = "FLOAT"} : tensor<i32>
// CHECK: return %0 : tensor<i32>
// CHECK: }
// -----
module {
func.func @degenerateCase(%arg0: tensor<1xf32>) -> tensor<1xf32> {
%0 = "tfl.add"(%arg0, %arg0) {tac.device = "GPU", fused_activation_function = "RELU6", tac.inference_type = "FLOAT"} : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32>
func.return %0 : tensor<1xf32>
}
}
// CHECK: func.func @degenerateCase(%arg0: tensor<1xf32>) -> tensor<1xf32> {
// CHECK: %0 = call @func_0_GPU_FLOAT(%arg0) {tac.device = "GPU", tac.inference_type = "FLOAT", tac.interface_name = "func_0"} : (tensor<1xf32>) -> tensor<1xf32>
// CHECK: return %0 : tensor<1xf32>
// CHECK: }
// CHECK: func.func private @func_0_GPU_FLOAT(%arg0: tensor<1xf32>) -> tensor<1xf32> attributes {tac.device = "GPU", tac.inference_type = "FLOAT", tac.interface_name = "func_0"} {
// CHECK: %0 = tfl.add %arg0, %arg0 {fused_activation_function = "RELU6", tac.device = "GPU", tac.inference_type = "FLOAT"} : tensor<1xf32>
// CHECK: return %0 : tensor<1xf32>
// CHECK: }
// -----
module {
func.func @simpleTest(%arg0: tensor<1xf32>, %arg1: tensor<1xf32>, %arg2: tensor<1xf32>, %arg3: tensor<1xf32>) -> tensor<2x1xf32> {
%0 = "tfl.add"(%arg0, %arg1) {tac.device = "GPU", fused_activation_function = "RELU6", tac.inference_type = "FLOAT"} : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32>
%1 = "tfl.mul"(%0, %arg2) {tac.device = "GPU", fused_activation_function = "RELU6", tac.inference_type = "FLOAT"} : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32>
%2 = "tfl.add"(%arg0, %arg3) {tac.device = "GPU", fused_activation_function = "RELU6", tac.inference_type = "FLOAT"} : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32>
%3 = "tfl.pack"(%1, %2) {tac.device = "CPU", tac.inference_type = "FLOAT", axis = 0 : i32, values_count = 2 : i32} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2x1xf32>
func.return %3 : tensor<2x1xf32>
}
}
// CHECK: func @simpleTest(%[[VAL_0:.*]]: tensor<1xf32>, %[[VAL_1:.*]]: tensor<1xf32>, %[[VAL_2:.*]]: tensor<1xf32>, %[[VAL_3:.*]]: tensor<1xf32>) -> tensor<2x1xf32> {
// CHECK: %[[VAL_4:.*]]:2 = call @func_0_GPU_FLOAT(%[[VAL_0]], %[[VAL_1]], %[[VAL_2]], %[[VAL_3]]) {tac.device = "GPU", tac.inference_type = "FLOAT", tac.interface_name = "func_0"} : (tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>) -> (tensor<1xf32>, tensor<1xf32>)
// CHECK: %[[VAL_5:.*]] = call @func_1_CPU_FLOAT(%[[VAL_4]]#0, %[[VAL_4]]#1) {tac.device = "CPU", tac.inference_type = "FLOAT", tac.interface_name = "func_1"} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2x1xf32>
// CHECK: return %[[VAL_5]] : tensor<2x1xf32>
// CHECK: }
// CHECK: func private @func_0_GPU_FLOAT(%[[VAL_0:.*]]: tensor<1xf32>, %[[VAL_1:.*]]: tensor<1xf32>, %[[VAL_2:.*]]: tensor<1xf32>, %[[VAL_4:.*]]: tensor<1xf32>) -> (tensor<1xf32>, tensor<1xf32>) attributes {tac.device = "GPU", tac.inference_type = "FLOAT", tac.interface_name = "func_0"} {
// CHECK: %[[VAL_5:.*]] = tfl.add %[[VAL_0]], %[[VAL_1]] {fused_activation_function = "RELU6", tac.device = "GPU", tac.inference_type = "FLOAT"} : tensor<1xf32>
// CHECK: %[[VAL_6:.*]] = tfl.mul %[[VAL_5]], %[[VAL_2]] {fused_activation_function = "RELU6", tac.device = "GPU", tac.inference_type = "FLOAT"} : tensor<1xf32>
// CHECK: %[[VAL_7:.*]] = tfl.add %[[VAL_0]], %[[VAL_4]] {fused_activation_function = "RELU6", tac.device = "GPU", tac.inference_type = "FLOAT"} : tensor<1xf32>
// CHECK: return %[[VAL_6]], %[[VAL_7]] : tensor<1xf32>, tensor<1xf32>
// CHECK: }
// CHECK: func private @func_1_CPU_FLOAT(%[[VAL_0:.*]]: tensor<1xf32>, %[[VAL_1:.*]]: tensor<1xf32>) -> tensor<2x1xf32> attributes {tac.device = "CPU", tac.inference_type = "FLOAT", tac.interface_name = "func_1"} {
// CHECK: %[[VAL_2:.*]] = "tfl.pack"(%[[VAL_0]], %[[VAL_1]]) <{axis = 0 : i32, values_count = 2 : i32}> {tac.device = "CPU", tac.inference_type = "FLOAT"} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2x1xf32>
// CHECK: return %[[VAL_2]] : tensor<2x1xf32>
// CHECK: }
// -----
module {
func.func @constWeight(%arg0: tensor<256x32x32x3xf32>) -> tensor<256x30x30x16xf32> {
%0 = "tfl.pseudo_const"() {value = dense<1.000000e+00> : tensor<16x3x3x3xf32>} : () -> tensor<16x3x3x3xf32>
%1 = "tfl.pseudo_const"() {value = dense<1.000000e+00> : tensor<16xf32>} : () -> tensor<16xf32>
%2 = "tfl.conv_2d"(%arg0, %0, %1) {tac.device = "GPU", tac.inference_type = "FLOAT", dilation_h_factor = 1 : i32, dilation_w_factor = 1 : i32, fused_activation_function = "NONE", padding = "VALID", stride_h = 1 : i32, stride_w = 1 : i32} : (tensor<256x32x32x3xf32>, tensor<16x3x3x3xf32>, tensor<16xf32>) -> tensor<256x30x30x16xf32>
%3 = "tfl.pseudo_const"() {value = dense<1.000000e+00> : tensor<16x3x3x16xf32>} : () -> tensor<16x3x3x16xf32>
%4 = "tfl.pseudo_const"() {value = dense<1.000000e+00> : tensor<16xf32>} : () -> tensor<16xf32>
%5 = "tfl.conv_2d"(%2, %3, %4) {tac.device = "GPU", tac.inference_type = "FLOAT", dilation_h_factor = 1 : i32, dilation_w_factor = 1 : i32, fused_activation_function = "NONE", padding = "SAME", stride_h = 1 : i32, stride_w = 1 : i32} : (tensor<256x30x30x16xf32>, tensor<16x3x3x16xf32>, tensor<16xf32>) -> tensor<256x30x30x16xf32>
func.return %5 : tensor<256x30x30x16xf32>
}
// CHECK: func @constWeight(%[[VAL_0:.*]]: tensor<256x32x32x3xf32>) -> tensor<256x30x30x16xf32> {
// CHECK-DAG: %[[VAL_1:.*]] = "tfl.pseudo_const"() <{value = dense<1.000000e+00> : tensor<16x3x3x3xf32>}> : () -> tensor<16x3x3x3xf32>
// CHECK-DAG: %[[VAL_2:.*]] = "tfl.pseudo_const"() <{value = dense<1.000000e+00> : tensor<16xf32>}> : () -> tensor<16xf32>
// CHECK-DAG: %[[VAL_3:.*]] = "tfl.pseudo_const"() <{value = dense<1.000000e+00> : tensor<16x3x3x16xf32>}> : () -> tensor<16x3x3x16xf32>
// CHECK-DAG: %[[VAL_4:.*]] = "tfl.pseudo_const"() <{value = dense<1.000000e+00> : tensor<16xf32>}> : () -> tensor<16xf32>
// CHECK: %[[VAL_5:.*]] = call @func_0_GPU_FLOAT(%[[VAL_0]], %[[VAL_1]], %[[VAL_2]], %[[VAL_3]], %[[VAL_4]]) {tac.device = "GPU", tac.inference_type = "FLOAT", tac.interface_name = "func_0"} : (tensor<256x32x32x3xf32>, tensor<16x3x3x3xf32>, tensor<16xf32>, tensor<16x3x3x16xf32>, tensor<16xf32>) -> tensor<256x30x30x16xf32>
// CHECK: return %[[VAL_5]] : tensor<256x30x30x16xf32>
// CHECK: }
// CHECK: func private @func_0_GPU_FLOAT(%[[VAL_0:.*]]: tensor<256x32x32x3xf32>, %[[VAL_1:.*]]: tensor<16x3x3x3xf32>, %[[VAL_2:.*]]: tensor<16xf32>, %[[VAL_3:.*]]: tensor<16x3x3x16xf32>, %[[VAL_4:.*]]: tensor<16xf32>) -> tensor<256x30x30x16xf32> attributes {tac.device = "GPU", tac.inference_type = "FLOAT", tac.interface_name = "func_0"} {
// CHECK: %[[VAL_5:.*]] = "tfl.conv_2d"(%[[VAL_0]], %[[VAL_1]], %[[VAL_2]]) <{dilation_h_factor = 1 : i32, dilation_w_factor = 1 : i32, fused_activation_function = "NONE", padding = "VALID", stride_h = 1 : i32, stride_w = 1 : i32}> {tac.device = "GPU", tac.inference_type = "FLOAT"} : (tensor<256x32x32x3xf32>, tensor<16x3x3x3xf32>, tensor<16xf32>) -> tensor<256x30x30x16xf32>
// CHECK: %[[VAL_6:.*]] = "tfl.conv_2d"(%[[VAL_5]], %[[VAL_3]], %[[VAL_4]]) <{dilation_h_factor = 1 : i32, dilation_w_factor = 1 : i32, fused_activation_function = "NONE", padding = "SAME", stride_h = 1 : i32, stride_w = 1 : i32}> {tac.device = "GPU", tac.inference_type = "FLOAT"} : (tensor<256x30x30x16xf32>, tensor<16x3x3x16xf32>, tensor<16xf32>) -> tensor<256x30x30x16xf32>
// CHECK: return %[[VAL_6]] : tensor<256x30x30x16xf32>
// CHECK: }
}
// -----
module {
func.func @norm1(%arg0: tensor<1x128x128xf32>) -> tensor<1x128x128xf32> {
%0 = "tfl.pseudo_const"() {value = dense<1.000000e+00> : tensor<128xf32>} : () -> tensor<128xf32>
%1 = "tfl.add"(%arg0, %0) {tac.device = "GPU", tac.inference_type = "FLOAT", fused_activation_function = "NONE"} : (tensor<1x128x128xf32>, tensor<128xf32>) -> tensor<1x128x128xf32>
%2 = "tfl.pseudo_const"() {value = dense<[128, 128]> : tensor<2xi32>} : () -> tensor<2xi32>
%3 = "tfl.reshape"(%1, %2) {tac.device = "GPU", tac.inference_type = "FLOAT"} : (tensor<1x128x128xf32>, tensor<2xi32>) -> tensor<128x128xf32>
%4 = "tfl.relu"(%3) {tac.device = "GPU", tac.inference_type = "FLOAT"} : (tensor<128x128xf32>) -> tensor<128x128xf32>
%5 = "tfl.pseudo_const"() {value = dense<[1, 128, 128]> : tensor<3xi32>} : () -> tensor<3xi32>
%6 = "tfl.reshape"(%4, %5) {tac.device = "GPU", tac.inference_type = "FLOAT"} : (tensor<128x128xf32>, tensor<3xi32>) -> tensor<1x128x128xf32>
%7 = "tfl.add"(%1, %6) {tac.device = "GPU", tac.inference_type = "FLOAT", fused_activation_function = "NONE"} : (tensor<1x128x128xf32>, tensor<1x128x128xf32>) -> tensor<1x128x128xf32>
func.return %7 : tensor<1x128x128xf32>
}
// CHECK: func @norm1(%[[VAL_0:.*]]: tensor<1x128x128xf32>) -> tensor<1x128x128xf32> {
// CHECK-DAG: %[[VAL_1:.*]] = "tfl.pseudo_const"() <{value = dense<1.000000e+00> : tensor<128xf32>}> : () -> tensor<128xf32>
// CHECK-DAG: %[[VAL_2:.*]] = "tfl.pseudo_const"() <{value = dense<128> : tensor<2xi32>}> : () -> tensor<2xi32>
// CHECK-DAG: %[[VAL_3:.*]] = "tfl.pseudo_const"() <{value = dense<[1, 128, 128]> : tensor<3xi32>}> : () -> tensor<3xi32>
// CHECK: %[[VAL_4:.*]] = call @func_0_GPU_FLOAT(%[[VAL_0]], %[[VAL_1]], %[[VAL_2]], %[[VAL_3]]) {tac.device = "GPU", tac.inference_type = "FLOAT", tac.interface_name = "func_0"} : (tensor<1x128x128xf32>, tensor<128xf32>, tensor<2xi32>, tensor<3xi32>) -> tensor<1x128x128xf32>
// CHECK: return %[[VAL_4]] : tensor<1x128x128xf32>
// CHECK: }
// CHECK: func private @func_0_GPU_FLOAT(%[[VAL_0:.*]]: tensor<1x128x128xf32>, %[[VAL_1:.*]]: tensor<128xf32>, %[[VAL_2:.*]]: tensor<2xi32>, %[[VAL_3:.*]]: tensor<3xi32>) -> tensor<1x128x128xf32> attributes {tac.device = "GPU", tac.inference_type = "FLOAT", tac.interface_name = "func_0"} {
// CHECK: %[[VAL_4:.*]] = tfl.add(%[[VAL_0]], %[[VAL_1]]) <{fused_activation_function = "NONE"}> {tac.device = "GPU", tac.inference_type = "FLOAT"} : (tensor<1x128x128xf32>, tensor<128xf32>) -> tensor<1x128x128xf32>
// CHECK: %[[VAL_5:.*]] = "tfl.reshape"(%[[VAL_4]], %[[VAL_2]]) {tac.device = "GPU", tac.inference_type = "FLOAT"} : (tensor<1x128x128xf32>, tensor<2xi32>) -> tensor<128x128xf32>
// CHECK: %[[VAL_6:.*]] = "tfl.relu"(%[[VAL_5]]) {tac.device = "GPU", tac.inference_type = "FLOAT"} : (tensor<128x128xf32>) -> tensor<128x128xf32>
// CHECK: %[[VAL_7:.*]] = "tfl.reshape"(%[[VAL_6]], %[[VAL_3]]) {tac.device = "GPU", tac.inference_type = "FLOAT"} : (tensor<128x128xf32>, tensor<3xi32>) -> tensor<1x128x128xf32>
// CHECK: %[[VAL_8:.*]] = tfl.add %[[VAL_4]], %[[VAL_7]] {fused_activation_function = "NONE", tac.device = "GPU", tac.inference_type = "FLOAT"} : tensor<1x128x128xf32>
// CHECK: return %[[VAL_8]] : tensor<1x128x128xf32>
// CHECK: }
}
// -----
module {
func.func @norm2(%arg0: tensor<1x128x128xf32>) -> tensor<1x128x128xf32> {
%0 = "tfl.pseudo_const"() {value = dense<1.000000e+00> : tensor<128xf32>} : () -> tensor<128xf32>
%1 = "tfl.add"(%arg0, %0) {tac.device = "GPU", tac.inference_type = "FLOAT", fused_activation_function = "NONE"} : (tensor<1x128x128xf32>, tensor<128xf32>) -> tensor<1x128x128xf32>
%2 = "tfl.pseudo_const"() {value = dense<[128, 128]> : tensor<2xi32>} : () -> tensor<2xi32>
%3 = "tfl.reshape"(%1, %2) {tac.device = "GPU", tac.inference_type = "FLOAT"} : (tensor<1x128x128xf32>, tensor<2xi32>) -> tensor<128x128xf32>
%4 = "tfl.relu"(%3) {tac.device = "GPU", tac.inference_type = "FLOAT"} : (tensor<128x128xf32>) -> tensor<128x128xf32>
%5 = "tfl.pseudo_const"() {value = dense<1.000000e+00> : tensor<128x128xf32>} : () -> tensor<128x128xf32>
%6 = "tfl.pseudo_const"() {value = dense<1.000000e+00> : tensor<128xf32>} : () -> tensor<128xf32>
%7 = "tfl.fully_connected"(%4, %5, %6) {tac.device = "CPU", tac.inference_type = "FLOAT", fused_activation_function = "NONE", keep_num_dims = false, weights_format = "DEFAULT"} : (tensor<128x128xf32>, tensor<128x128xf32>, tensor<128xf32>) -> tensor<128x128xf32>
%8 = "tfl.pseudo_const"() {value = dense<[1, 128, 128]> : tensor<3xi32>} : () -> tensor<3xi32>
%9 = "tfl.reshape"(%7, %8) {tac.device = "GPU", tac.inference_type = "FLOAT"} : (tensor<128x128xf32>, tensor<3xi32>) -> tensor<1x128x128xf32>
%10 = "tfl.add"(%1, %9) {tac.device = "GPU", tac.inference_type = "FLOAT", fused_activation_function = "NONE"} : (tensor<1x128x128xf32>, tensor<1x128x128xf32>) -> tensor<1x128x128xf32>
func.return %10 : tensor<1x128x128xf32>
}
}
// CHECK: func @norm2(%[[VAL_0:.*]]: tensor<1x128x128xf32>) -> tensor<1x128x128xf32> {
// CHECK-DAG: %[[VAL_1:.*]] = "tfl.pseudo_const"() <{value = dense<1.000000e+00> : tensor<128xf32>}> : () -> tensor<128xf32>
// CHECK-DAG: %[[VAL_2:.*]] = "tfl.pseudo_const"() <{value = dense<128> : tensor<2xi32>}> : () -> tensor<2xi32>
// CHECK: %[[VAL_3:.*]]:2 = call @func_0_GPU_FLOAT(%[[VAL_0]], %[[VAL_1]], %[[VAL_2]]) {tac.device = "GPU", tac.inference_type = "FLOAT", tac.interface_name = "func_0"} : (tensor<1x128x128xf32>, tensor<128xf32>, tensor<2xi32>) -> (tensor<1x128x128xf32>, tensor<128x128xf32>)
// CHECK-DAG: %[[VAL_4:.*]] = "tfl.pseudo_const"() <{value = dense<1.000000e+00> : tensor<128x128xf32>}> : () -> tensor<128x128xf32>
// CHECK-DAG: %[[VAL_5:.*]] = "tfl.pseudo_const"() <{value = dense<1.000000e+00> : tensor<128xf32>}> : () -> tensor<128xf32>
// CHECK: %[[VAL_6:.*]] = call @func_2_CPU_FLOAT(%[[VAL_3]]#1, %[[VAL_4]], %[[VAL_5]]) {tac.device = "CPU", tac.inference_type = "FLOAT", tac.interface_name = "func_2"} : (tensor<128x128xf32>, tensor<128x128xf32>, tensor<128xf32>) -> tensor<128x128xf32>
// CHECK: %[[VAL_7:.*]] = "tfl.pseudo_const"() <{value = dense<[1, 128, 128]> : tensor<3xi32>}> : () -> tensor<3xi32>
// CHECK: %[[VAL_8:.*]] = call @func_1_GPU_FLOAT(%[[VAL_6]], %[[VAL_7]], %[[VAL_3]]#0) {tac.device = "GPU", tac.inference_type = "FLOAT", tac.interface_name = "func_1"} : (tensor<128x128xf32>, tensor<3xi32>, tensor<1x128x128xf32>) -> tensor<1x128x128xf32>
// CHECK: return %[[VAL_8]] : tensor<1x128x128xf32>
// CHECK: }
// CHECK: func.func private @func_0_GPU_FLOAT(%[[VAL_0:.*]]: tensor<1x128x128xf32>, %[[VAL_1:.*]]: tensor<128xf32>, %[[VAL_2:.*]]: tensor<2xi32>) -> (tensor<1x128x128xf32>, tensor<128x128xf32>) attributes {tac.device = "GPU", tac.inference_type = "FLOAT", tac.interface_name = "func_0"} {
// CHECK: %[[VAL_3:.*]] = tfl.add(%[[VAL_0]], %[[VAL_1]]) <{fused_activation_function = "NONE"}> {tac.device = "GPU", tac.inference_type = "FLOAT"} : (tensor<1x128x128xf32>, tensor<128xf32>) -> tensor<1x128x128xf32>
// CHECK: %[[VAL_4:.*]] = "tfl.reshape"(%[[VAL_3]], %[[VAL_2]]) {tac.device = "GPU", tac.inference_type = "FLOAT"} : (tensor<1x128x128xf32>, tensor<2xi32>) -> tensor<128x128xf32>
// CHECK: %[[VAL_5:.*]] = "tfl.relu"(%[[VAL_4]]) {tac.device = "GPU", tac.inference_type = "FLOAT"} : (tensor<128x128xf32>) -> tensor<128x128xf32>
// CHECK: return %[[VAL_3]], %[[VAL_5]] : tensor<1x128x128xf32>, tensor<128x128xf32>
// CHECK: }
// CHECK: func.func private @func_1_GPU_FLOAT(%[[VAL_0:.*]]: tensor<128x128xf32>, %[[VAL_1:.*]]: tensor<3xi32>, %[[VAL_2:.*]]: tensor<1x128x128xf32>) -> tensor<1x128x128xf32> attributes {tac.device = "GPU", tac.inference_type = "FLOAT", tac.interface_name = "func_1"} {
// CHECK: %[[VAL_3:.*]] = "tfl.reshape"(%[[VAL_0]], %[[VAL_1]]) {tac.device = "GPU", tac.inference_type = "FLOAT"} : (tensor<128x128xf32>, tensor<3xi32>) -> tensor<1x128x128xf32>
// CHECK: %[[VAL_4:.*]] = tfl.add %[[VAL_2]], %[[VAL_3]] {fused_activation_function = "NONE", tac.device = "GPU", tac.inference_type = "FLOAT"} : tensor<1x128x128xf32>
// CHECK: return %[[VAL_4]] : tensor<1x128x128xf32>
// CHECK: }
// CHECK: func.func private @func_2_CPU_FLOAT(%[[VAL_0:.*]]: tensor<128x128xf32>, %[[VAL_1:.*]]: tensor<128x128xf32>, %[[VAL_2:.*]]: tensor<128xf32>) -> tensor<128x128xf32> attributes {tac.device = "CPU", tac.inference_type = "FLOAT", tac.interface_name = "func_2"} {
// CHECK: %[[VAL_3:.*]] = "tfl.fully_connected"(%[[VAL_0]], %[[VAL_1]], %[[VAL_2]]) <{fused_activation_function = "NONE", keep_num_dims = false, weights_format = "DEFAULT"}> {tac.device = "CPU", tac.inference_type = "FLOAT"} : (tensor<128x128xf32>, tensor<128x128xf32>, tensor<128xf32>) -> tensor<128x128xf32>
// CHECK: return %[[VAL_3]] : tensor<128x128xf32>
// CHECK: }
// -----
module {
func.func @quantizedOpOnly(%arg0: tensor<1x!quant.uniform<i8:f32, 0.003:-128>>, %arg1: tensor<1x!quant.uniform<i8:f32, 0.003:-128>>) -> tensor<2x1x!quant.uniform<i8:f32, 0.003:-128>> {
%0 = "tfl.pseudo_qconst"() {qtype = tensor<1x!quant.uniform<i8:f32, 0.003:-128>>, value = dense<127> : tensor<1xi8>} : () -> tensor<1x!quant.uniform<i8:f32, 0.003:-128>>
%1 = "tfl.pseudo_qconst"() {qtype = tensor<1x!quant.uniform<i8:f32, 0.003:-128>>, value = dense<127> : tensor<1xi8>} : () -> tensor<1x!quant.uniform<i8:f32, 0.003:-128>>
%2 = "tfl.mul"(%arg0, %0) {tac.device = "CPU", tac.inference_type = "QUANTIZED_INT8", fused_activation_function = "NONE"} : (tensor<1x!quant.uniform<i8:f32, 0.003:-128>>, tensor<1x!quant.uniform<i8:f32, 0.003:-128>>) -> tensor<1x!quant.uniform<i8:f32, 0.003:-128>>
%3 = "tfl.add"(%2, %1) {tac.device = "CPU", tac.inference_type = "QUANTIZED_INT8", fused_activation_function = "NONE"} : (tensor<1x!quant.uniform<i8:f32, 0.003:-128>>, tensor<1x!quant.uniform<i8:f32, 0.003:-128>>) -> tensor<1x!quant.uniform<i8:f32, 0.003:-128>>
%4 = "tfl.add"(%arg1, %0) {tac.device = "CPU", tac.inference_type = "QUANTIZED_INT8", fused_activation_function = "NONE"} : (tensor<1x!quant.uniform<i8:f32, 0.003:-128>>, tensor<1x!quant.uniform<i8:f32, 0.003:-128>>) -> tensor<1x!quant.uniform<i8:f32, 0.003:-128>>
%5 = "tfl.pack"(%3, %4) {tac.device = "CPU", tac.inference_type = "QUANTIZED_INT8", axis = 0 : i32, values_count = 2 : i32} : (tensor<1x!quant.uniform<i8:f32, 0.003:-128>>, tensor<1x!quant.uniform<i8:f32, 0.003:-128>>) -> tensor<2x1x!quant.uniform<i8:f32, 0.003:-128>>
func.return %5: tensor<2x1x!quant.uniform<i8:f32, 0.003:-128>>
}
// CHECK: func @quantizedOpOnly(%[[VAL_0:.*]]: tensor<1x!quant.uniform<i8:f32, 3.000000e-03:-128>>, %[[VAL_1:.*]]: tensor<1x!quant.uniform<i8:f32, 3.000000e-03:-128>>) -> tensor<2x1x!quant.uniform<i8:f32, 3.000000e-03:-128>> {
// CHECK: %[[VAL_2:.*]] = "tfl.pseudo_qconst"() <{qtype = tensor<1x!quant.uniform<i8:f32, 3.000000e-03:-128>>, value = dense<127> : tensor<1xi8>}> : () -> tensor<1x!quant.uniform<i8:f32, 3.000000e-03:-128>>
// CHECK: %[[VAL_3:.*]] = "tfl.pseudo_qconst"() <{qtype = tensor<1x!quant.uniform<i8:f32, 3.000000e-03:-128>>, value = dense<127> : tensor<1xi8>}> : () -> tensor<1x!quant.uniform<i8:f32, 3.000000e-03:-128>>
// CHECK: %[[VAL_4:.*]] = call @func_0_CPU_QUANTIZED_INT8(%[[VAL_0]], %[[VAL_2]], %[[VAL_3]], %[[VAL_1]]) {tac.device = "CPU", tac.inference_type = "QUANTIZED_INT8", tac.interface_name = "func_0"} : (tensor<1x!quant.uniform<i8:f32, 3.000000e-03:-128>>, tensor<1x!quant.uniform<i8:f32, 3.000000e-03:-128>>, tensor<1x!quant.uniform<i8:f32, 3.000000e-03:-128>>, tensor<1x!quant.uniform<i8:f32, 3.000000e-03:-128>>) -> tensor<2x1x!quant.uniform<i8:f32, 3.000000e-03:-128>>
// CHECK: return %[[VAL_4]] : tensor<2x1x!quant.uniform<i8:f32, 3.000000e-03:-128>>
// CHECK: }
// CHECK: func private @func_0_CPU_QUANTIZED_INT8(%[[VAL_0:.*]]: tensor<1x!quant.uniform<i8:f32, 3.000000e-03:-128>>, %[[VAL_1:.*]]: tensor<1x!quant.uniform<i8:f32, 3.000000e-03:-128>>, %[[VAL_2:.*]]: tensor<1x!quant.uniform<i8:f32, 3.000000e-03:-128>>, %[[VAL_3:.*]]: tensor<1x!quant.uniform<i8:f32, 3.000000e-03:-128>>) -> tensor<2x1x!quant.uniform<i8:f32, 3.000000e-03:-128>> attributes {tac.device = "CPU", tac.inference_type = "QUANTIZED_INT8", tac.interface_name = "func_0"} {
// CHECK: %[[VAL_5:.*]] = tfl.mul %[[VAL_0]], %[[VAL_1]] {fused_activation_function = "NONE", tac.device = "CPU", tac.inference_type = "QUANTIZED_INT8"} : tensor<1x!quant.uniform<i8:f32, 3.000000e-03:-128>>
// CHECK: %[[VAL_6:.*]] = tfl.add %[[VAL_5]], %[[VAL_2]] {fused_activation_function = "NONE", tac.device = "CPU", tac.inference_type = "QUANTIZED_INT8"} : tensor<1x!quant.uniform<i8:f32, 3.000000e-03:-128>>
// CHECK: %[[VAL_7:.*]] = tfl.add %[[VAL_3]], %[[VAL_1]] {fused_activation_function = "NONE", tac.device = "CPU", tac.inference_type = "QUANTIZED_INT8"} : tensor<1x!quant.uniform<i8:f32, 3.000000e-03:-128>>
// CHECK: %[[VAL_8:.*]] = "tfl.pack"(%[[VAL_6]], %[[VAL_7]]) <{axis = 0 : i32, values_count = 2 : i32}> {tac.device = "CPU", tac.inference_type = "QUANTIZED_INT8"} : (tensor<1x!quant.uniform<i8:f32, 3.000000e-03:-128>>, tensor<1x!quant.uniform<i8:f32, 3.000000e-03:-128>>) -> tensor<2x1x!quant.uniform<i8:f32, 3.000000e-03:-128>>
// CHECK: return %[[VAL_8]] : tensor<2x1x!quant.uniform<i8:f32, 3.000000e-03:-128>>
// CHECK: }
}
// -----
module {
func.func @quantizationWithFloat(%arg0: tensor<1x1x384x!quant.uniform<i8:f32, 0.003:-128>>, %arg1: tensor<1x1x384x!quant.uniform<i8:f32, 0.003:-128>>) -> tensor<1x384x384x!quant.uniform<i8:f32, 0.003:-128>> {
%0 = "tfl.pseudo_qconst"() {qtype = tensor<1x384x1x!quant.uniform<i8:f32, 0.003:-128>>, value = dense<127> : tensor<1x384x1xi8>} : () -> tensor<1x384x1x!quant.uniform<i8:f32, 0.003:-128>>
%1 = "tfl.mul"(%arg0, %0) {tac.device = "CPU", tac.inference_type = "QUANTIZED_INT8", fused_activation_function = "NONE"} : (tensor<1x1x384x!quant.uniform<i8:f32, 0.003:-128>>, tensor<1x384x1x!quant.uniform<i8:f32, 0.003:-128>>) -> tensor<1x384x384x!quant.uniform<i8:f32, 0.003:-128>>
%2 = "tfl.dequantize"(%1) : (tensor<1x384x384x!quant.uniform<i8:f32, 0.003:-128>>) -> tensor<1x384x384xf32>
%3 = "tfl.pseudo_const"() {value = dense<1.000000e+00> : tensor<1x1x384xf32>} : () -> tensor<1x384x384xf32>
%4 = "tfl.add"(%2, %3) {tac.device = "GPU", tac.inference_type = "FLOAT", fused_activation_function = "NONE"} : (tensor<1x384x384xf32>, tensor<1x384x384xf32>) -> tensor<1x384x384xf32>
%5 = "tfl.quantize"(%4) {qtype = tensor<1x384x1x!quant.uniform<i8:f32, 0.003:-128>>} : (tensor<1x384x384xf32>) -> tensor<1x384x384x!quant.uniform<i8:f32, 0.003:-128>>
%6 = "tfl.mul"(%arg1, %5) {tac.device = "CPU", tac.inference_type = "QUANTIZED_INT8", fused_activation_function = "NONE"} : (tensor<1x1x384x!quant.uniform<i8:f32, 0.003:-128>>, tensor<1x384x384x!quant.uniform<i8:f32, 0.003:-128>>) -> tensor<1x384x384x!quant.uniform<i8:f32, 0.003:-128>>
func.return %6: tensor<1x384x384x!quant.uniform<i8:f32, 0.003:-128>>
}
}
// CHECK: func @quantizationWithFloat(%[[VAL_0:.*]]: tensor<1x1x384x!quant.uniform<i8:f32, 3.000000e-03:-128>>, %[[VAL_1:.*]]: tensor<1x1x384x!quant.uniform<i8:f32, 3.000000e-03:-128>>) -> tensor<1x384x384x!quant.uniform<i8:f32, 3.000000e-03:-128>> {
// CHECK: %[[VAL_2:.*]] = "tfl.pseudo_qconst"() <{qtype = tensor<1x384x1x!quant.uniform<i8:f32, 3.000000e-03:-128>>, value = dense<127> : tensor<1x384x1xi8>}> : () -> tensor<1x384x1x!quant.uniform<i8:f32, 3.000000e-03:-128>>
// CHECK: %[[VAL_3:.*]] = call @func_1_CPU_QUANTIZED_INT8(%[[VAL_0]], %[[VAL_2]]) {tac.device = "CPU", tac.inference_type = "QUANTIZED_INT8", tac.interface_name = "func_1"} : (tensor<1x1x384x!quant.uniform<i8:f32, 3.000000e-03:-128>>, tensor<1x384x1x!quant.uniform<i8:f32, 3.000000e-03:-128>>) -> tensor<1x384x384x!quant.uniform<i8:f32, 3.000000e-03:-128>>
// CHECK: %[[VAL_4:.*]] = "tfl.dequantize"(%[[VAL_3]]) : (tensor<1x384x384x!quant.uniform<i8:f32, 3.000000e-03:-128>>) -> tensor<1x384x384xf32>
// CHECK: %[[VAL_5:.*]] = "tfl.pseudo_const"() <{value = dense<1.000000e+00> : tensor<1x1x384xf32>}> : () -> tensor<1x384x384xf32>
// CHECK: %[[VAL_6:.*]] = call @func_0_GPU_FLOAT(%[[VAL_4]], %[[VAL_5]]) {tac.device = "GPU", tac.inference_type = "FLOAT", tac.interface_name = "func_0"} : (tensor<1x384x384xf32>, tensor<1x384x384xf32>) -> tensor<1x384x384xf32>
// CHECK: %[[VAL_7:.*]] = "tfl.quantize"(%[[VAL_6]]) <{qtype = tensor<1x384x1x!quant.uniform<i8:f32, 3.000000e-03:-128>>}> : (tensor<1x384x384xf32>) -> tensor<1x384x384x!quant.uniform<i8:f32, 3.000000e-03:-128>>
// CHECK: %[[VAL_8:.*]] = call @func_2_CPU_QUANTIZED_INT8(%[[VAL_1]], %[[VAL_7]]) {tac.device = "CPU", tac.inference_type = "QUANTIZED_INT8", tac.interface_name = "func_2"} : (tensor<1x1x384x!quant.uniform<i8:f32, 3.000000e-03:-128>>, tensor<1x384x384x!quant.uniform<i8:f32, 3.000000e-03:-128>>) -> tensor<1x384x384x!quant.uniform<i8:f32, 3.000000e-03:-128>>
// CHECK: return %[[VAL_8]] : tensor<1x384x384x!quant.uniform<i8:f32, 3.000000e-03:-128>>
// CHECK: }
// CHECK: func private @func_0_GPU_FLOAT(%[[VAL_0:.*]]: tensor<1x384x384xf32>, %[[VAL_1:.*]]: tensor<1x384x384xf32>) -> tensor<1x384x384xf32> attributes {tac.device = "GPU", tac.inference_type = "FLOAT", tac.interface_name = "func_0"} {
// CHECK: %[[VAL_2:.*]] = tfl.add %[[VAL_0]], %[[VAL_1]] {fused_activation_function = "NONE", tac.device = "GPU", tac.inference_type = "FLOAT"} : tensor<1x384x384xf32>
// CHECK: return %[[VAL_2]] : tensor<1x384x384xf32>
// CHECK: }
// CHECK: func private @func_1_CPU_QUANTIZED_INT8(%[[VAL_0:.*]]: tensor<1x1x384x!quant.uniform<i8:f32, 3.000000e-03:-128>>, %[[VAL_1:.*]]: tensor<1x384x1x!quant.uniform<i8:f32, 3.000000e-03:-128>>) -> tensor<1x384x384x!quant.uniform<i8:f32, 3.000000e-03:-128>> attributes {tac.device = "CPU", tac.inference_type = "QUANTIZED_INT8", tac.interface_name = "func_1"} {
// CHECK: %[[VAL_2:.*]] = tfl.mul(%[[VAL_0]], %[[VAL_1]]) <{fused_activation_function = "NONE"}> {tac.device = "CPU", tac.inference_type = "QUANTIZED_INT8"} : (tensor<1x1x384x!quant.uniform<i8:f32, 3.000000e-03:-128>>, tensor<1x384x1x!quant.uniform<i8:f32, 3.000000e-03:-128>>) -> tensor<1x384x384x!quant.uniform<i8:f32, 3.000000e-03:-128>>
// CHECK: return %[[VAL_2]] : tensor<1x384x384x!quant.uniform<i8:f32, 3.000000e-03:-128>>
// CHECK: }
// CHECK: func private @func_2_CPU_QUANTIZED_INT8(%[[VAL_0:.*]]: tensor<1x1x384x!quant.uniform<i8:f32, 3.000000e-03:-128>>, %[[VAL_1:.*]]: tensor<1x384x384x!quant.uniform<i8:f32, 3.000000e-03:-128>>) -> tensor<1x384x384x!quant.uniform<i8:f32, 3.000000e-03:-128>> attributes {tac.device = "CPU", tac.inference_type = "QUANTIZED_INT8", tac.interface_name = "func_2"} {
// CHECK: %[[VAL_2:.*]] = tfl.mul(%[[VAL_0]], %[[VAL_1]]) <{fused_activation_function = "NONE"}> {tac.device = "CPU", tac.inference_type = "QUANTIZED_INT8"} : (tensor<1x1x384x!quant.uniform<i8:f32, 3.000000e-03:-128>>, tensor<1x384x384x!quant.uniform<i8:f32, 3.000000e-03:-128>>) -> tensor<1x384x384x!quant.uniform<i8:f32, 3.000000e-03:-128>>
// CHECK: return %[[VAL_2]] : tensor<1x384x384x!quant.uniform<i8:f32, 3.000000e-03:-128>>
// CHECK: }
// -----
func.func @cond_false_72730(%arg0: tensor<?x?x!tf_type.string>, %arg1: tensor<?x?x!tf_type.string>, %arg2: tensor<?x?xi32>, %arg3: tensor<?x!tf_type.string>, %arg4: tensor<?x!tf_type.string>, %arg5: tensor<?xi32>, %arg6: tensor<?xi32>) -> (tensor<?x?x!tf_type.string>, tensor<?x?x!tf_type.string>) {
%cst = arith.constant dense<[1, 0]> : tensor<2xi32>
%cst_0 = arith.constant dense<0> : tensor<i32>
%cst_1 = arith.constant dense<-1> : tensor<1xi32>
%cst_2 = arith.constant dense<1> : tensor<1xi32>
%cst_3 = arith.constant dense<0> : tensor<1xi32>
%0 = "tfl.shape"(%arg2) {tac.device = "DARWINN", tac.inference_type = "FLOAT"} : (tensor<?x?xi32>) -> tensor<2xi32>
%1 = "tfl.strided_slice"(%0, %cst_3, %cst_2, %cst_2) {begin_mask = 0 : i32, ellipsis_mask = 0 : i32, end_mask = 0 : i32, new_axis_mask = 0 : i32, shrink_axis_mask = 1 : i32, offset = false, tac.device = "DARWINN", tac.inference_type = "FLOAT"} : (tensor<2xi32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<i32>
%2 = "tfl.custom"(%cst_1, %1) {custom_code = "FlexTensorListReserve", custom_option = #tfl<const_bytes : "0x1154656E736F724C697374526573657276650040121154656E736F724C697374526573657276651A001A002A130A0D656C656D656E745F6474797065120230072A100A0A73686170655F74797065120230033200000255431414042801">, tac.device = "CPU", tac.inference_type = "FLOAT"} : (tensor<1xi32>, tensor<i32>) -> tensor<!tf_type.variant<tensor<?x!tf_type.string>>>
%3 = "tfl.custom"(%cst_1, %1) {custom_code = "FlexTensorListReserve", custom_option = #tfl<const_bytes : "0x1154656E736F724C697374526573657276650040121154656E736F724C697374526573657276651A001A002A130A0D656C656D656E745F6474797065120230032A100A0A73686170655F74797065120230033200000255431414042801">, tac.device = "CPU", tac.inference_type = "FLOAT"} : (tensor<1xi32>, tensor<i32>) -> tensor<!tf_type.variant<tensor<?xi32>>>
%4:8 = "tfl.while"(%cst_0, %cst_0, %arg5, %arg6, %2, %2, %3, %3) ({
^bb0(%arg7: tensor<i32>, %arg8: tensor<i32>, %arg9: tensor<?xi32>, %arg10: tensor<?xi32>, %arg11: tensor<!tf_type.variant<tensor<?x!tf_type.string>>>, %arg12: tensor<!tf_type.variant<tensor<?x!tf_type.string>>>, %arg13: tensor<!tf_type.variant<tensor<?xi32>>>, %arg14: tensor<!tf_type.variant<tensor<?xi32>>>):
%9 = tfl.less(%arg8, %1) {tac.device = "DARWINN", tac.inference_type = "FLOAT"} : (tensor<i32>, tensor<i32>) -> tensor<i1>
%10 = tfl.less(%arg7, %1) {tac.device = "DARWINN", tac.inference_type = "FLOAT"} : (tensor<i32>, tensor<i32>) -> tensor<i1>
%11 = tfl.logical_and %10, %9 {tac.device = "DARWINN", tac.inference_type = "FLOAT"} : tensor<i1>
"tfl.yield"(%11) : (tensor<i1>) -> ()
}, {
^bb0(%arg7: tensor<i32>, %arg8: tensor<i32>, %arg9: tensor<?xi32>, %arg10: tensor<?xi32>, %arg11: tensor<!tf_type.variant<tensor<?x!tf_type.string>>>, %arg12: tensor<!tf_type.variant<tensor<?x!tf_type.string>>>, %arg13: tensor<!tf_type.variant<tensor<?xi32>>>, %arg14: tensor<!tf_type.variant<tensor<?xi32>>>):
%cst_4 = arith.constant dense<[0, 0, 1, 1, 1]> : tensor<5xi32>
%cst_5 = arith.constant dense<[0, 1, 0, 1, 1]> : tensor<5xi32>
%cst_6 = arith.constant dense<2> : tensor<i32>
%cst_7 = arith.constant dense<"*"> : tensor<!tf_type.string>
%cst_8 = arith.constant dense<-1> : tensor<1xi32>
%cst_9 = arith.constant dense<-1> : tensor<i32>
%cst_10 = arith.constant dense<2> : tensor<1xi32>
%cst_11 = arith.constant dense<1> : tensor<i32>
%cst_12 = arith.constant dense<0> : tensor<i32>
%cst_13 = arith.constant dense<1> : tensor<1xi32>
%cst_14 = arith.constant dense<0> : tensor<1xi32>
%9 = "tfl.shape"(%arg1) {tac.device = "DARWINN", tac.inference_type = "FLOAT"} : (tensor<?x?x!tf_type.string>) -> tensor<2xi32>
%10 = "tfl.strided_slice"(%9, %cst_14, %cst_13, %cst_13) {begin_mask = 0 : i32, ellipsis_mask = 0 : i32, end_mask = 0 : i32, new_axis_mask = 0 : i32, shrink_axis_mask = 1 : i32, offset = false, tac.device = "DARWINN", tac.inference_type = "FLOAT"} : (tensor<2xi32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<i32>
%11 = "tfl.range"(%cst_12, %10, %cst_11) {tac.device = "DARWINN", tac.inference_type = "FLOAT"} : (tensor<i32>, tensor<i32>, tensor<i32>) -> tensor<?xi32>
%12 = "tfl.pack"(%10, %cst_11) {axis = 0 : i32, tac.device = "DARWINN", tac.inference_type = "FLOAT", values_count = 2 : i32} : (tensor<i32>, tensor<i32>) -> tensor<2xi32>
%13 = "tfl.strided_slice"(%9, %cst_13, %cst_10, %cst_13) {begin_mask = 0 : i32, ellipsis_mask = 0 : i32, end_mask = 0 : i32, new_axis_mask = 0 : i32, shrink_axis_mask = 1 : i32, offset = false, tac.device = "DARWINN", tac.inference_type = "FLOAT"} : (tensor<2xi32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<i32>
%14 = tfl.mul(%11, %13) {fused_activation_function = "NONE", tac.device = "DARWINN", tac.inference_type = "FLOAT"} : (tensor<?xi32>, tensor<i32>) -> tensor<?xi32>
%15 = "tfl.reshape"(%14, %12) {tac.device = "DARWINN", tac.inference_type = "FLOAT"} : (tensor<?xi32>, tensor<2xi32>) -> tensor<?x1xi32>
%16 = "tfl.strided_slice"(%9, %cst_14, %cst_10, %cst_13) {begin_mask = 1 : i32, ellipsis_mask = 0 : i32, end_mask = 0 : i32, new_axis_mask = 0 : i32, shrink_axis_mask = 0 : i32, offset = false, tac.device = "DARWINN", tac.inference_type = "FLOAT"} : (tensor<2xi32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<2xi32>
%17 = "tfl.reduce_prod"(%16, %cst_14) {keep_dims = true, tac.device = "CPU", tac.inference_type = "FLOAT"} : (tensor<2xi32>, tensor<1xi32>) -> tensor<1xi32>
%18 = "tfl.reshape"(%arg1, %17) {tac.device = "DARWINN", tac.inference_type = "FLOAT"} : (tensor<?x?x!tf_type.string>, tensor<1xi32>) -> tensor<?x!tf_type.string>
%19 = "tfl.shape"(%arg0) {tac.device = "DARWINN", tac.inference_type = "FLOAT"} : (tensor<?x?x!tf_type.string>) -> tensor<2xi32>
%20 = "tfl.strided_slice"(%19, %cst_14, %cst_13, %cst_13) {begin_mask = 0 : i32, ellipsis_mask = 0 : i32, end_mask = 0 : i32, new_axis_mask = 0 : i32, shrink_axis_mask = 1 : i32, offset = false, tac.device = "DARWINN", tac.inference_type = "FLOAT"} : (tensor<2xi32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<i32>
%21 = "tfl.range"(%cst_12, %20, %cst_11) {tac.device = "DARWINN", tac.inference_type = "FLOAT"} : (tensor<i32>, tensor<i32>, tensor<i32>) -> tensor<?xi32>
%22 = "tfl.pack"(%20, %cst_11) {axis = 0 : i32, tac.device = "DARWINN", tac.inference_type = "FLOAT", values_count = 2 : i32} : (tensor<i32>, tensor<i32>) -> tensor<2xi32>
%23 = "tfl.strided_slice"(%19, %cst_13, %cst_10, %cst_13) {begin_mask = 0 : i32, ellipsis_mask = 0 : i32, end_mask = 0 : i32, new_axis_mask = 0 : i32, shrink_axis_mask = 1 : i32, offset = false, tac.device = "DARWINN", tac.inference_type = "FLOAT"} : (tensor<2xi32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<i32>
%24 = tfl.mul(%21, %23) {fused_activation_function = "NONE", tac.device = "DARWINN", tac.inference_type = "FLOAT"} : (tensor<?xi32>, tensor<i32>) -> tensor<?xi32>
%25 = "tfl.reshape"(%24, %22) {tac.device = "DARWINN", tac.inference_type = "FLOAT"} : (tensor<?xi32>, tensor<2xi32>) -> tensor<?x1xi32>
%26 = "tfl.strided_slice"(%19, %cst_14, %cst_10, %cst_13) {begin_mask = 1 : i32, ellipsis_mask = 0 : i32, end_mask = 0 : i32, new_axis_mask = 0 : i32, shrink_axis_mask = 0 : i32, offset = false, tac.device = "DARWINN", tac.inference_type = "FLOAT"} : (tensor<2xi32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<2xi32>
%27 = "tfl.reduce_prod"(%26, %cst_14) {keep_dims = true, tac.device = "CPU", tac.inference_type = "FLOAT"} : (tensor<2xi32>, tensor<1xi32>) -> tensor<1xi32>
%28 = "tfl.reshape"(%arg0, %27) {tac.device = "DARWINN", tac.inference_type = "FLOAT"} : (tensor<?x?x!tf_type.string>, tensor<1xi32>) -> tensor<?x!tf_type.string>
%29 = tfl.add %arg8, %cst_11 {fused_activation_function = "NONE", tac.device = "DARWINN", tac.inference_type = "FLOAT"} : tensor<i32>
%30 = "tfl.expand_dims"(%arg9, %cst_9) {tac.device = "CPU", tac.inference_type = "FLOAT"} : (tensor<?xi32>, tensor<i32>) -> tensor<?x1xi32>
%31 = tfl.add %30, %25 {fused_activation_function = "NONE", tac.device = "DARWINN", tac.inference_type = "FLOAT"} : tensor<?x1xi32>
%32 = "tfl.reshape"(%31, %cst_8) {tac.device = "DARWINN", tac.inference_type = "FLOAT"} : (tensor<?x1xi32>, tensor<1xi32>) -> tensor<?xi32>
%33 = "tfl.gather"(%28, %32) <{axis = 0 : i32, batch_dims = 0 : i32}> {tac.device = "DARWINN", tac.inference_type = "FLOAT"} : (tensor<?x!tf_type.string>, tensor<?xi32>) -> tensor<?x!tf_type.string>
%34 = "tfl.reshape"(%33, %cst_8) {tac.device = "DARWINN", tac.inference_type = "FLOAT"} : (tensor<?x!tf_type.string>, tensor<1xi32>) -> tensor<?x!tf_type.string>
%35 = "tfl.shape"(%34) {tac.device = "DARWINN", tac.inference_type = "FLOAT"} : (tensor<?x!tf_type.string>) -> tensor<1xi32>
%36 = "tfl.fill"(%35, %cst_7) {tac.device = "DARWINN", tac.inference_type = "FLOAT"} : (tensor<1xi32>, tensor<!tf_type.string>) -> tensor<?x!tf_type.string>
%37 = "tfl.expand_dims"(%arg10, %cst_9) {tac.device = "CPU", tac.inference_type = "FLOAT"} : (tensor<?xi32>, tensor<i32>) -> tensor<?x1xi32>
%38 = tfl.add %37, %15 {fused_activation_function = "NONE", tac.device = "DARWINN", tac.inference_type = "FLOAT"} : tensor<?x1xi32>
%39 = "tfl.reshape"(%38, %cst_8) {tac.device = "DARWINN", tac.inference_type = "FLOAT"} : (tensor<?x1xi32>, tensor<1xi32>) -> tensor<?xi32>
%40 = "tfl.gather"(%18, %39) <{axis = 0 : i32, batch_dims = 0 : i32}> {tac.device = "DARWINN", tac.inference_type = "FLOAT"} : (tensor<?x!tf_type.string>, tensor<?xi32>) -> tensor<?x!tf_type.string>
%41 = "tfl.reshape"(%40, %cst_8) {tac.device = "DARWINN", tac.inference_type = "FLOAT"} : (tensor<?x!tf_type.string>, tensor<1xi32>) -> tensor<?x!tf_type.string>
%42 = "tfl.shape"(%41) {tac.device = "DARWINN", tac.inference_type = "FLOAT"} : (tensor<?x!tf_type.string>) -> tensor<1xi32>
%43 = "tfl.fill"(%42, %cst_7) {tac.device = "DARWINN", tac.inference_type = "FLOAT"} : (tensor<1xi32>, tensor<!tf_type.string>) -> tensor<?x!tf_type.string>
%44 = "tfl.gather"(%arg2, %arg8) <{axis = 0 : i32, batch_dims = 0 : i32}> {tac.device = "DARWINN", tac.inference_type = "FLOAT"} : (tensor<?x?xi32>, tensor<i32>) -> tensor<?xi32>
%45 = "tfl.equal"(%44, %cst_6) {tac.device = "DARWINN", tac.inference_type = "FLOAT"} : (tensor<?xi32>, tensor<i32>) -> tensor<?xi1>
%46 = "tfl.custom"(%45, %36, %34) {custom_code = "FlexSelect", custom_option = #tfl<const_bytes : "0x0653656C6563740031120653656C6563741A001A001A002A070A01541202300732180A052E31323131120F7768696C655F626F64795F3733303200023B341414042801">, tac.device = "CPU", tac.inference_type = "FLOAT"} : (tensor<?xi1>, tensor<?x!tf_type.string>, tensor<?x!tf_type.string>) -> tensor<?x!tf_type.string>
%47 = "tfl.custom"(%arg11, %arg8, %46) {custom_code = "FlexTensorListSetItem", custom_option = #tfl<const_bytes : "0x1154656E736F724C6973745365744974656D0047121154656E736F724C6973745365744974656D1A001A001A002A130A0D656C656D656E745F64747970651202300732170A042E326333120F7768696C655F626F64795F3733303200025C4A1414042801">, tac.device = "CPU", tac.inference_type = "FLOAT"} : (tensor<!tf_type.variant<tensor<?x!tf_type.string>>>, tensor<i32>, tensor<?x!tf_type.string>) -> tensor<!tf_type.variant<tensor<?x!tf_type.string>>>
%48 = "tfl.equal"(%44, %cst_11) {tac.device = "DARWINN", tac.inference_type = "FLOAT"} : (tensor<?xi32>, tensor<i32>) -> tensor<?xi1>
%49 = "tfl.custom"(%48, %43, %41) {custom_code = "FlexSelect", custom_option = #tfl<const_bytes : "0x0653656C6563740031120653656C6563741A001A001A002A070A01541202300732180A052E31323166120F7768696C655F626F64795F3733303200023B341414042801">, tac.device = "CPU", tac.inference_type = "FLOAT"} : (tensor<?xi1>, tensor<?x!tf_type.string>, tensor<?x!tf_type.string>) -> tensor<?x!tf_type.string>
%50 = "tfl.custom"(%arg12, %arg8, %49) {custom_code = "FlexTensorListSetItem", custom_option = #tfl<const_bytes : "0x1154656E736F724C6973745365744974656D0047121154656E736F724C6973745365744974656D1A001A001A002A130A0D656C656D656E745F64747970651202300732170A042E326335120F7768696C655F626F64795F3733303200025C4A1414042801">, tac.device = "CPU", tac.inference_type = "FLOAT"} : (tensor<!tf_type.variant<tensor<?x!tf_type.string>>>, tensor<i32>, tensor<?x!tf_type.string>) -> tensor<!tf_type.variant<tensor<?x!tf_type.string>>>
%51 = "tfl.gather"(%cst_5, %44) <{axis = 0 : i32, batch_dims = 0 : i32}> {tac.device = "DARWINN", tac.inference_type = "FLOAT"} : (tensor<5xi32>, tensor<?xi32>) -> tensor<?xi32>
%52 = tfl.add %arg9, %51 {fused_activation_function = "NONE", tac.device = "DARWINN", tac.inference_type = "FLOAT"} : tensor<?xi32>
%53 = "tfl.custom"(%arg13, %arg8, %52) {custom_code = "FlexTensorListSetItem", custom_option = #tfl<const_bytes : "0x1154656E736F724C6973745365744974656D0047121154656E736F724C6973745365744974656D1A001A001A002A130A0D656C656D656E745F64747970651202300332170A042E326337120F7768696C655F626F64795F3733303200025C4A1414042801">, tac.device = "CPU", tac.inference_type = "FLOAT"} : (tensor<!tf_type.variant<tensor<?xi32>>>, tensor<i32>, tensor<?xi32>) -> tensor<!tf_type.variant<tensor<?xi32>>>
%54 = "tfl.gather"(%cst_4, %44) <{axis = 0 : i32, batch_dims = 0 : i32}> {tac.device = "DARWINN", tac.inference_type = "FLOAT"} : (tensor<5xi32>, tensor<?xi32>) -> tensor<?xi32>
%55 = tfl.add %arg10, %54 {fused_activation_function = "NONE", tac.device = "DARWINN", tac.inference_type = "FLOAT"} : tensor<?xi32>
%56 = "tfl.custom"(%arg14, %arg8, %55) {custom_code = "FlexTensorListSetItem", custom_option = #tfl<const_bytes : "0x1154656E736F724C6973745365744974656D0047121154656E736F724C6973745365744974656D1A001A001A002A130A0D656C656D656E745F64747970651202300332170A042E343139120F7768696C655F626F64795F3733303200025C4A1414042801">, tac.device = "CPU", tac.inference_type = "FLOAT"} : (tensor<!tf_type.variant<tensor<?xi32>>>, tensor<i32>, tensor<?xi32>) -> tensor<!tf_type.variant<tensor<?xi32>>>
%57 = tfl.add %arg7, %cst_11 {fused_activation_function = "NONE", tac.device = "DARWINN", tac.inference_type = "FLOAT"} : tensor<i32>
"tfl.yield"(%57, %29, %52, %55, %47, %50, %53, %56) : (tensor<i32>, tensor<i32>, tensor<?xi32>, tensor<?xi32>, tensor<!tf_type.variant<tensor<?x!tf_type.string>>>, tensor<!tf_type.variant<tensor<?x!tf_type.string>>>, tensor<!tf_type.variant<tensor<?xi32>>>, tensor<!tf_type.variant<tensor<?xi32>>>) -> ()
}) {tac.device = "CPU", tac.inference_type = "FLOAT"} : (tensor<i32>, tensor<i32>, tensor<?xi32>, tensor<?xi32>, tensor<!tf_type.variant<tensor<?x!tf_type.string>>>, tensor<!tf_type.variant<tensor<?x!tf_type.string>>>, tensor<!tf_type.variant<tensor<?xi32>>>, tensor<!tf_type.variant<tensor<?xi32>>>) -> (tensor<i32>, tensor<i32>, tensor<?xi32>, tensor<?xi32>, tensor<!tf_type.variant<tensor<?x!tf_type.string>>>, tensor<!tf_type.variant<tensor<?x!tf_type.string>>>, tensor<!tf_type.variant<tensor<?xi32>>>, tensor<!tf_type.variant<tensor<?xi32>>>)
%5 = "tfl.custom"(%4#4, %cst_1) {custom_code = "FlexTensorListStack", custom_option = #tfl<const_bytes : "0x0F54656E736F724C697374537461636B0049120F54656E736F724C697374537461636B1A001A002A130A0D656C656D656E745F6474797065120230072A1B0A0C6E756D5F656C656D656E7473120B18FFFFFFFFFFFFFFFFFF01320000025C4C1414042801">, tac.device = "CPU", tac.inference_type = "FLOAT"} : (tensor<!tf_type.variant<tensor<?x!tf_type.string>>>, tensor<1xi32>) -> tensor<?x?x!tf_type.string>
%6 = "tfl.custom"(%5, %cst) {custom_code = "FlexTranspose", custom_option = #tfl<const_bytes : "0x095472616E73706F7365002712095472616E73706F73651A001A002A0B0A05547065726D120230032A070A01541202300732000002342A1414042801">, tac.device = "CPU", tac.inference_type = "FLOAT"} : (tensor<?x?x!tf_type.string>, tensor<2xi32>) -> tensor<?x?x!tf_type.string>
%7 = "tfl.custom"(%4#5, %cst_1) {custom_code = "FlexTensorListStack", custom_option = #tfl<const_bytes : "0x0F54656E736F724C697374537461636B0049120F54656E736F724C697374537461636B1A001A002A130A0D656C656D656E745F6474797065120230072A1B0A0C6E756D5F656C656D656E7473120B18FFFFFFFFFFFFFFFFFF01320000025C4C1414042801">, tac.device = "CPU", tac.inference_type = "FLOAT"} : (tensor<!tf_type.variant<tensor<?x!tf_type.string>>>, tensor<1xi32>) -> tensor<?x?x!tf_type.string>
%8 = "tfl.custom"(%7, %cst) {custom_code = "FlexTranspose", custom_option = #tfl<const_bytes : "0x095472616E73706F7365002712095472616E73706F73651A001A002A070A0154120230072A0B0A05547065726D1202300332000002342A1414042801">, tac.device = "CPU", tac.inference_type = "FLOAT"} : (tensor<?x?x!tf_type.string>, tensor<2xi32>) -> tensor<?x?x!tf_type.string>
return %6, %8 : tensor<?x?x!tf_type.string>, tensor<?x?x!tf_type.string>
}
// CHECK: func.func @cond_false_72730(%arg0: tensor<?x?x!tf_type.string>, %arg1: tensor<?x?x!tf_type.string>, %arg2: tensor<?x?xi32>, %arg3: tensor<?x!tf_type.string>, %arg4: tensor<?x!tf_type.string>, %arg5: tensor<?xi32>, %arg6: tensor<?xi32>) -> (tensor<?x?x!tf_type.string>, tensor<?x?x!tf_type.string>) {
// CHECK: %cst = arith.constant dense<[1, 0]> : tensor<2xi32>
// CHECK: %cst_0 = arith.constant dense<0> : tensor<i32>
// CHECK: %cst_1 = arith.constant dense<-1> : tensor<1xi32>
// CHECK: %cst_2 = arith.constant dense<1> : tensor<1xi32>
// CHECK: %cst_3 = arith.constant dense<0> : tensor<1xi32>
// CHECK: %0 = call @func_0_DARWINN_FLOAT(%arg2, %cst_3, %cst_2) {tac.device = "DARWINN", tac.inference_type = "FLOAT", tac.interface_name = "func_0"} : (tensor<?x?xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<i32>
// CHECK: %1:2 = call @func_1_CPU_FLOAT(%cst_1, %0, %cst_0, %arg5, %arg6, %arg1, %arg0, %arg2, %cst) {tac.device = "CPU", tac.inference_type = "FLOAT", tac.interface_name = "func_1"} : (tensor<1xi32>, tensor<i32>, tensor<i32>, tensor<?xi32>, tensor<?xi32>, tensor<?x?x!tf_type.string>, tensor<?x?x!tf_type.string>, tensor<?x?xi32>, tensor<2xi32>) -> (tensor<?x?x!tf_type.string>, tensor<?x?x!tf_type.string>)
// CHECK: return %1#0, %1#1 : tensor<?x?x!tf_type.string>, tensor<?x?x!tf_type.string>
// CHECK: }
// CHECK: func.func private @func_0_DARWINN_FLOAT(%arg0: tensor<?x?xi32>, %arg1: tensor<1xi32>, %arg2: tensor<1xi32>) -> tensor<i32> attributes {tac.device = "DARWINN", tac.inference_type = "FLOAT", tac.interface_name = "func_0"} {
// CHECK: %0 = "tfl.shape"(%arg0) {tac.device = "DARWINN", tac.inference_type = "FLOAT"} : (tensor<?x?xi32>) -> tensor<2xi32>
// CHECK: %1 = "tfl.strided_slice"(%0, %arg1, %arg2, %arg2) <{begin_mask = 0 : i32, ellipsis_mask = 0 : i32, end_mask = 0 : i32, new_axis_mask = 0 : i32, offset = false, shrink_axis_mask = 1 : i32}> {tac.device = "DARWINN", tac.inference_type = "FLOAT"} : (tensor<2xi32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<i32>
// CHECK: return %1 : tensor<i32>
// CHECK: }
// CHECK: func.func private @func_1_CPU_FLOAT(%arg0: tensor<1xi32>, %arg1: tensor<i32>, %arg2: tensor<i32>, %arg3: tensor<?xi32>, %arg4: tensor<?xi32>, %arg5: tensor<?x?x!tf_type.string>, %arg6: tensor<?x?x!tf_type.string>, %arg7: tensor<?x?xi32>, %arg8: tensor<2xi32>) -> (tensor<?x?x!tf_type.string>, tensor<?x?x!tf_type.string>) attributes {tac.device = "CPU", tac.inference_type = "FLOAT", tac.interface_name = "func_1"} {
// CHECK: %0 = "tfl.custom"(%arg0, %arg1) <{custom_code = "FlexTensorListReserve", custom_option = #tfl<const_bytes : "0x1154656E736F724C697374526573657276650040121154656E736F724C697374526573657276651A001A002A130A0D656C656D656E745F6474797065120230072A100A0A73686170655F74797065120230033200000255431414042801">}> {tac.device = "CPU", tac.inference_type = "FLOAT"} : (tensor<1xi32>, tensor<i32>) -> tensor<!tf_type.variant<tensor<?x!tf_type.string>>>
// CHECK: %1 = "tfl.custom"(%arg0, %arg1) <{custom_code = "FlexTensorListReserve", custom_option = #tfl<const_bytes : "0x1154656E736F724C697374526573657276650040121154656E736F724C697374526573657276651A001A002A130A0D656C656D656E745F6474797065120230032A100A0A73686170655F74797065120230033200000255431414042801">}> {tac.device = "CPU", tac.inference_type = "FLOAT"} : (tensor<1xi32>, tensor<i32>) -> tensor<!tf_type.variant<tensor<?xi32>>>
// CHECK: %2:8 = "tfl.while"(%arg2, %arg2, %arg3, %arg4, %0, %0, %1, %1) ({
// CHECK: ^bb0(%arg9: tensor<i32>, %arg10: tensor<i32>, %arg11: tensor<?xi32>, %arg12: tensor<?xi32>, %arg13: tensor<!tf_type.variant<tensor<?x!tf_type.string>>>, %arg14: tensor<!tf_type.variant<tensor<?x!tf_type.string>>>, %arg15: tensor<!tf_type.variant<tensor<?xi32>>>, %arg16: tensor<!tf_type.variant<tensor<?xi32>>>):
// CHECK: %7 = func.call @func_2_DARWINN_FLOAT(%arg10, %arg1, %arg9) {tac.device = "DARWINN", tac.inference_type = "FLOAT", tac.interface_name = "func_2"} : (tensor<i32>, tensor<i32>, tensor<i32>) -> tensor<i1>
// CHECK: "tfl.yield"(%7) : (tensor<i1>) -> ()
// CHECK: }, {
// CHECK: ^bb0(%arg9: tensor<i32>, %arg10: tensor<i32>, %arg11: tensor<?xi32>, %arg12: tensor<?xi32>, %arg13: tensor<!tf_type.variant<tensor<?x!tf_type.string>>>, %arg14: tensor<!tf_type.variant<tensor<?x!tf_type.string>>>, %arg15: tensor<!tf_type.variant<tensor<?xi32>>>, %arg16: tensor<!tf_type.variant<tensor<?xi32>>>):
// CHECK: %cst = arith.constant dense<[0, 0, 1, 1, 1]> : tensor<5xi32>
// CHECK: %cst_0 = arith.constant dense<[0, 1, 0, 1, 1]> : tensor<5xi32>
// CHECK: %cst_1 = arith.constant dense<2> : tensor<i32>
// CHECK: %cst_2 = arith.constant dense<"*"> : tensor<!tf_type.string>
// CHECK: %cst_3 = arith.constant dense<-1> : tensor<1xi32>
// CHECK: %cst_4 = arith.constant dense<-1> : tensor<i32>
// CHECK: %cst_5 = arith.constant dense<2> : tensor<1xi32>
// CHECK: %cst_6 = arith.constant dense<1> : tensor<i32>
// CHECK: %cst_7 = arith.constant dense<0> : tensor<i32>
// CHECK: %cst_8 = arith.constant dense<1> : tensor<1xi32>
// CHECK: %cst_9 = arith.constant dense<0> : tensor<1xi32>
// CHECK: %7:2 = func.call @func_3_DARWINN_FLOAT(%arg5, %cst_9, %cst_8, %cst_7, %cst_6, %cst_5) {tac.device = "DARWINN", tac.inference_type = "FLOAT", tac.interface_name = "func_3"} : (tensor<?x?x!tf_type.string>, tensor<1xi32>, tensor<1xi32>, tensor<i32>, tensor<i32>, tensor<1xi32>) -> (tensor<?x1xi32>, tensor<2xi32>)
// CHECK: %8 = "tfl.reduce_prod"(%7#1, %cst_9) <{keep_dims = true}> {tac.device = "CPU", tac.inference_type = "FLOAT"} : (tensor<2xi32>, tensor<1xi32>) -> tensor<1xi32>
// CHECK: %9:3 = func.call @func_4_DARWINN_FLOAT(%arg5, %8, %arg6, %cst_9, %cst_8, %cst_7, %cst_6, %cst_5) {tac.device = "DARWINN", tac.inference_type = "FLOAT", tac.interface_name = "func_4"} : (tensor<?x?x!tf_type.string>, tensor<1xi32>, tensor<?x?x!tf_type.string>, tensor<1xi32>, tensor<1xi32>, tensor<i32>, tensor<i32>, tensor<1xi32>) -> (tensor<?x!tf_type.string>, tensor<?x1xi32>, tensor<2xi32>)
// CHECK: %10 = "tfl.reduce_prod"(%9#2, %cst_9) <{keep_dims = true}> {tac.device = "CPU", tac.inference_type = "FLOAT"} : (tensor<2xi32>, tensor<1xi32>) -> tensor<1xi32>
// CHECK: %11 = "tfl.expand_dims"(%arg11, %cst_4) {tac.device = "CPU", tac.inference_type = "FLOAT"} : (tensor<?xi32>, tensor<i32>) -> tensor<?x1xi32>
// CHECK: %12 = "tfl.expand_dims"(%arg12, %cst_4) {tac.device = "CPU", tac.inference_type = "FLOAT"} : (tensor<?xi32>, tensor<i32>) -> tensor<?x1xi32>
// CHECK: %13:10 = func.call @func_5_DARWINN_FLOAT(%arg6, %10, %arg10, %cst_6, %11, %9#1, %cst_3, %cst_2, %12, %7#0, %9#0, %arg7, %cst_1, %cst_0, %arg11, %cst, %arg12, %arg9) {tac.device = "DARWINN", tac.inference_type = "FLOAT", tac.interface_name = "func_5"} : (tensor<?x?x!tf_type.string>, tensor<1xi32>, tensor<i32>, tensor<i32>, tensor<?x1xi32>, tensor<?x1xi32>, tensor<1xi32>, tensor<!tf_type.string>, tensor<?x1xi32>, tensor<?x1xi32>, tensor<?x!tf_type.string>, tensor<?x?xi32>, tensor<i32>, tensor<5xi32>, tensor<?xi32>, tensor<5xi32>, tensor<?xi32>, tensor<i32>) -> (tensor<i32>, tensor<?x!tf_type.string>, tensor<?x!tf_type.string>, tensor<?x!tf_type.string>, tensor<?x!tf_type.string>, tensor<?xi1>, tensor<?xi1>, tensor<?xi32>, tensor<?xi32>, tensor<i32>)
// CHECK: %14 = "tfl.custom"(%13#5, %13#2, %13#1) <{custom_code = "FlexSelect", custom_option = #tfl<const_bytes : "0x0653656C6563740031120653656C6563741A001A001A002A070A01541202300732180A052E31323131120F7768696C655F626F64795F3733303200023B341414042801">}> {tac.device = "CPU", tac.inference_type = "FLOAT"} : (tensor<?xi1>, tensor<?x!tf_type.string>, tensor<?x!tf_type.string>) -> tensor<?x!tf_type.string>
// CHECK: %15 = "tfl.custom"(%arg13, %arg10, %14) <{custom_code = "FlexTensorListSetItem", custom_option = #tfl<const_bytes : "0x1154656E736F724C6973745365744974656D0047121154656E736F724C6973745365744974656D1A001A001A002A130A0D656C656D656E745F64747970651202300732170A042E326333120F7768696C655F626F64795F3733303200025C4A1414042801">}> {tac.device = "CPU", tac.inference_type = "FLOAT"} : (tensor<!tf_type.variant<tensor<?x!tf_type.string>>>, tensor<i32>, tensor<?x!tf_type.string>) -> tensor<!tf_type.variant<tensor<?x!tf_type.string>>>
// CHECK: %16 = "tfl.custom"(%13#6, %13#4, %13#3) <{custom_code = "FlexSelect", custom_option = #tfl<const_bytes : "0x0653656C6563740031120653656C6563741A001A001A002A070A01541202300732180A052E31323166120F7768696C655F626F64795F3733303200023B341414042801">}> {tac.device = "CPU", tac.inference_type = "FLOAT"} : (tensor<?xi1>, tensor<?x!tf_type.string>, tensor<?x!tf_type.string>) -> tensor<?x!tf_type.string>
// CHECK: %17 = "tfl.custom"(%arg14, %arg10, %16) <{custom_code = "FlexTensorListSetItem", custom_option = #tfl<const_bytes : "0x1154656E736F724C6973745365744974656D0047121154656E736F724C6973745365744974656D1A001A001A002A130A0D656C656D656E745F64747970651202300732170A042E326335120F7768696C655F626F64795F3733303200025C4A1414042801">}> {tac.device = "CPU", tac.inference_type = "FLOAT"} : (tensor<!tf_type.variant<tensor<?x!tf_type.string>>>, tensor<i32>, tensor<?x!tf_type.string>) -> tensor<!tf_type.variant<tensor<?x!tf_type.string>>>
// CHECK: %18 = "tfl.custom"(%arg15, %arg10, %13#7) <{custom_code = "FlexTensorListSetItem", custom_option = #tfl<const_bytes : "0x1154656E736F724C6973745365744974656D0047121154656E736F724C6973745365744974656D1A001A001A002A130A0D656C656D656E745F64747970651202300332170A042E326337120F7768696C655F626F64795F3733303200025C4A1414042801">}> {tac.device = "CPU", tac.inference_type = "FLOAT"} : (tensor<!tf_type.variant<tensor<?xi32>>>, tensor<i32>, tensor<?xi32>) -> tensor<!tf_type.variant<tensor<?xi32>>>
// CHECK: %19 = "tfl.custom"(%arg16, %arg10, %13#8) <{custom_code = "FlexTensorListSetItem", custom_option = #tfl<const_bytes : "0x1154656E736F724C6973745365744974656D0047121154656E736F724C6973745365744974656D1A001A001A002A130A0D656C656D656E745F64747970651202300332170A042E343139120F7768696C655F626F64795F3733303200025C4A1414042801">}> {tac.device = "CPU", tac.inference_type = "FLOAT"} : (tensor<!tf_type.variant<tensor<?xi32>>>, tensor<i32>, tensor<?xi32>) -> tensor<!tf_type.variant<tensor<?xi32>>>
// CHECK: "tfl.yield"(%13#9, %13#0, %13#7, %13#8, %15, %17, %18, %19) : (tensor<i32>, tensor<i32>, tensor<?xi32>, tensor<?xi32>, tensor<!tf_type.variant<tensor<?x!tf_type.string>>>, tensor<!tf_type.variant<tensor<?x!tf_type.string>>>, tensor<!tf_type.variant<tensor<?xi32>>>, tensor<!tf_type.variant<tensor<?xi32>>>) -> ()
// CHECK: }) {tac.device = "CPU", tac.inference_type = "FLOAT"} : (tensor<i32>, tensor<i32>, tensor<?xi32>, tensor<?xi32>, tensor<!tf_type.variant<tensor<?x!tf_type.string>>>, tensor<!tf_type.variant<tensor<?x!tf_type.string>>>, tensor<!tf_type.variant<tensor<?xi32>>>, tensor<!tf_type.variant<tensor<?xi32>>>) -> (tensor<i32>, tensor<i32>, tensor<?xi32>, tensor<?xi32>, tensor<!tf_type.variant<tensor<?x!tf_type.string>>>, tensor<!tf_type.variant<tensor<?x!tf_type.string>>>, tensor<!tf_type.variant<tensor<?xi32>>>, tensor<!tf_type.variant<tensor<?xi32>>>)
// CHECK: %3 = "tfl.custom"(%2#4, %arg0) <{custom_code = "FlexTensorListStack", custom_option = #tfl<const_bytes : "0x0F54656E736F724C697374537461636B0049120F54656E736F724C697374537461636B1A001A002A130A0D656C656D656E745F6474797065120230072A1B0A0C6E756D5F656C656D656E7473120B18FFFFFFFFFFFFFFFFFF01320000025C4C1414042801">}> {tac.device = "CPU", tac.inference_type = "FLOAT"} : (tensor<!tf_type.variant<tensor<?x!tf_type.string>>>, tensor<1xi32>) -> tensor<?x?x!tf_type.string>
// CHECK: %4 = "tfl.custom"(%3, %arg8) <{custom_code = "FlexTranspose", custom_option = #tfl<const_bytes : "0x095472616E73706F7365002712095472616E73706F73651A001A002A0B0A05547065726D120230032A070A01541202300732000002342A1414042801">}> {tac.device = "CPU", tac.inference_type = "FLOAT"} : (tensor<?x?x!tf_type.string>, tensor<2xi32>) -> tensor<?x?x!tf_type.string>
// CHECK: %5 = "tfl.custom"(%2#5, %arg0) <{custom_code = "FlexTensorListStack", custom_option = #tfl<const_bytes : "0x0F54656E736F724C697374537461636B0049120F54656E736F724C697374537461636B1A001A002A130A0D656C656D656E745F6474797065120230072A1B0A0C6E756D5F656C656D656E7473120B18FFFFFFFFFFFFFFFFFF01320000025C4C1414042801">}> {tac.device = "CPU", tac.inference_type = "FLOAT"} : (tensor<!tf_type.variant<tensor<?x!tf_type.string>>>, tensor<1xi32>) -> tensor<?x?x!tf_type.string>
// CHECK: %6 = "tfl.custom"(%5, %arg8) <{custom_code = "FlexTranspose", custom_option = #tfl<const_bytes : "0x095472616E73706F7365002712095472616E73706F73651A001A002A070A0154120230072A0B0A05547065726D1202300332000002342A1414042801">}> {tac.device = "CPU", tac.inference_type = "FLOAT"} : (tensor<?x?x!tf_type.string>, tensor<2xi32>) -> tensor<?x?x!tf_type.string>
// CHECK: return %4, %6 : tensor<?x?x!tf_type.string>, tensor<?x?x!tf_type.string>
// CHECK: }
// CHECK: func.func private @func_2_DARWINN_FLOAT(%arg0: tensor<i32>, %arg1: tensor<i32>, %arg2: tensor<i32>) -> tensor<i1> attributes {tac.device = "DARWINN", tac.inference_type = "FLOAT", tac.interface_name = "func_2"} {
// CHECK: %0 = tfl.less(%arg0, %arg1) {tac.device = "DARWINN", tac.inference_type = "FLOAT"} : (tensor<i32>, tensor<i32>) -> tensor<i1>
// CHECK: %1 = tfl.less(%arg2, %arg1) {tac.device = "DARWINN", tac.inference_type = "FLOAT"} : (tensor<i32>, tensor<i32>) -> tensor<i1>
// CHECK: %2 = tfl.logical_and %1, %0 {tac.device = "DARWINN", tac.inference_type = "FLOAT"} : tensor<i1>
// CHECK: return %2 : tensor<i1>
// CHECK: }
// CHECK: func.func private @func_3_DARWINN_FLOAT(%arg0: tensor<?x?x!tf_type.string>, %arg1: tensor<1xi32>, %arg2: tensor<1xi32>, %arg3: tensor<i32>, %arg4: tensor<i32>, %arg5: tensor<1xi32>) -> (tensor<?x1xi32>, tensor<2xi32>) attributes {tac.device = "DARWINN", tac.inference_type = "FLOAT", tac.interface_name = "func_3"} {
// CHECK: %0 = "tfl.shape"(%arg0) {tac.device = "DARWINN", tac.inference_type = "FLOAT"} : (tensor<?x?x!tf_type.string>) -> tensor<2xi32>
// CHECK: %1 = "tfl.strided_slice"(%0, %arg1, %arg2, %arg2) <{begin_mask = 0 : i32, ellipsis_mask = 0 : i32, end_mask = 0 : i32, new_axis_mask = 0 : i32, offset = false, shrink_axis_mask = 1 : i32}> {tac.device = "DARWINN", tac.inference_type = "FLOAT"} : (tensor<2xi32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<i32>
// CHECK: %2 = "tfl.range"(%arg3, %1, %arg4) {tac.device = "DARWINN", tac.inference_type = "FLOAT"} : (tensor<i32>, tensor<i32>, tensor<i32>) -> tensor<?xi32>
// CHECK: %3 = "tfl.pack"(%1, %arg4) <{axis = 0 : i32, values_count = 2 : i32}> {tac.device = "DARWINN", tac.inference_type = "FLOAT"} : (tensor<i32>, tensor<i32>) -> tensor<2xi32>
// CHECK: %4 = "tfl.strided_slice"(%0, %arg2, %arg5, %arg2) <{begin_mask = 0 : i32, ellipsis_mask = 0 : i32, end_mask = 0 : i32, new_axis_mask = 0 : i32, offset = false, shrink_axis_mask = 1 : i32}> {tac.device = "DARWINN", tac.inference_type = "FLOAT"} : (tensor<2xi32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<i32>
// CHECK: %5 = tfl.mul(%2, %4) <{fused_activation_function = "NONE"}> {tac.device = "DARWINN", tac.inference_type = "FLOAT"} : (tensor<?xi32>, tensor<i32>) -> tensor<?xi32>
// CHECK: %6 = "tfl.reshape"(%5, %3) {tac.device = "DARWINN", tac.inference_type = "FLOAT"} : (tensor<?xi32>, tensor<2xi32>) -> tensor<?x1xi32>
// CHECK: %7 = "tfl.strided_slice"(%0, %arg1, %arg5, %arg2) <{begin_mask = 1 : i32, ellipsis_mask = 0 : i32, end_mask = 0 : i32, new_axis_mask = 0 : i32, offset = false, shrink_axis_mask = 0 : i32}> {tac.device = "DARWINN", tac.inference_type = "FLOAT"} : (tensor<2xi32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<2xi32>
// CHECK: return %6, %7 : tensor<?x1xi32>, tensor<2xi32>
// CHECK: }
// CHECK: func.func private @func_4_DARWINN_FLOAT(%arg0: tensor<?x?x!tf_type.string>, %arg1: tensor<1xi32>, %arg2: tensor<?x?x!tf_type.string>, %arg3: tensor<1xi32>, %arg4: tensor<1xi32>, %arg5: tensor<i32>, %arg6: tensor<i32>, %arg7: tensor<1xi32>) -> (tensor<?x!tf_type.string>, tensor<?x1xi32>, tensor<2xi32>) attributes {tac.device = "DARWINN", tac.inference_type = "FLOAT", tac.interface_name = "func_4"} {
// CHECK: %0 = "tfl.reshape"(%arg0, %arg1) {tac.device = "DARWINN", tac.inference_type = "FLOAT"} : (tensor<?x?x!tf_type.string>, tensor<1xi32>) -> tensor<?x!tf_type.string>
// CHECK: %1 = "tfl.shape"(%arg2) {tac.device = "DARWINN", tac.inference_type = "FLOAT"} : (tensor<?x?x!tf_type.string>) -> tensor<2xi32>
// CHECK: %2 = "tfl.strided_slice"(%1, %arg3, %arg4, %arg4) <{begin_mask = 0 : i32, ellipsis_mask = 0 : i32, end_mask = 0 : i32, new_axis_mask = 0 : i32, offset = false, shrink_axis_mask = 1 : i32}> {tac.device = "DARWINN", tac.inference_type = "FLOAT"} : (tensor<2xi32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<i32>
// CHECK: %3 = "tfl.range"(%arg5, %2, %arg6) {tac.device = "DARWINN", tac.inference_type = "FLOAT"} : (tensor<i32>, tensor<i32>, tensor<i32>) -> tensor<?xi32>
// CHECK: %4 = "tfl.pack"(%2, %arg6) <{axis = 0 : i32, values_count = 2 : i32}> {tac.device = "DARWINN", tac.inference_type = "FLOAT"} : (tensor<i32>, tensor<i32>) -> tensor<2xi32>
// CHECK: %5 = "tfl.strided_slice"(%1, %arg4, %arg7, %arg4) <{begin_mask = 0 : i32, ellipsis_mask = 0 : i32, end_mask = 0 : i32, new_axis_mask = 0 : i32, offset = false, shrink_axis_mask = 1 : i32}> {tac.device = "DARWINN", tac.inference_type = "FLOAT"} : (tensor<2xi32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<i32>
// CHECK: %6 = tfl.mul(%3, %5) <{fused_activation_function = "NONE"}> {tac.device = "DARWINN", tac.inference_type = "FLOAT"} : (tensor<?xi32>, tensor<i32>) -> tensor<?xi32>
// CHECK: %7 = "tfl.reshape"(%6, %4) {tac.device = "DARWINN", tac.inference_type = "FLOAT"} : (tensor<?xi32>, tensor<2xi32>) -> tensor<?x1xi32>
// CHECK: %8 = "tfl.strided_slice"(%1, %arg3, %arg7, %arg4) <{begin_mask = 1 : i32, ellipsis_mask = 0 : i32, end_mask = 0 : i32, new_axis_mask = 0 : i32, offset = false, shrink_axis_mask = 0 : i32}> {tac.device = "DARWINN", tac.inference_type = "FLOAT"} : (tensor<2xi32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<2xi32>
// CHECK: return %0, %7, %8 : tensor<?x!tf_type.string>, tensor<?x1xi32>, tensor<2xi32>
// CHECK: }
// CHECK: func.func private @func_5_DARWINN_FLOAT(%arg0: tensor<?x?x!tf_type.string>, %arg1: tensor<1xi32>, %arg2: tensor<i32>, %arg3: tensor<i32>, %arg4: tensor<?x1xi32>, %arg5: tensor<?x1xi32>, %arg6: tensor<1xi32>, %arg7: tensor<!tf_type.string>, %arg8: tensor<?x1xi32>, %arg9: tensor<?x1xi32>, %arg10: tensor<?x!tf_type.string>, %arg11: tensor<?x?xi32>, %arg12: tensor<i32>, %arg13: tensor<5xi32>, %arg14: tensor<?xi32>, %arg15: tensor<5xi32>, %arg16: tensor<?xi32>, %arg17: tensor<i32>) -> (tensor<i32>, tensor<?x!tf_type.string>, tensor<?x!tf_type.string>, tensor<?x!tf_type.string>, tensor<?x!tf_type.string>, tensor<?xi1>, tensor<?xi1>, tensor<?xi32>, tensor<?xi32>, tensor<i32>) attributes {tac.device = "DARWINN", tac.inference_type = "FLOAT", tac.interface_name = "func_5"} {
// CHECK: %0 = "tfl.reshape"(%arg0, %arg1) {tac.device = "DARWINN", tac.inference_type = "FLOAT"} : (tensor<?x?x!tf_type.string>, tensor<1xi32>) -> tensor<?x!tf_type.string>
// CHECK: %1 = tfl.add %arg2, %arg3 {fused_activation_function = "NONE", tac.device = "DARWINN", tac.inference_type = "FLOAT"} : tensor<i32>
// CHECK: %2 = tfl.add %arg4, %arg5 {fused_activation_function = "NONE", tac.device = "DARWINN", tac.inference_type = "FLOAT"} : tensor<?x1xi32>
// CHECK: %3 = "tfl.reshape"(%2, %arg6) {tac.device = "DARWINN", tac.inference_type = "FLOAT"} : (tensor<?x1xi32>, tensor<1xi32>) -> tensor<?xi32>
// CHECK: %4 = "tfl.gather"(%0, %3) <{axis = 0 : i32, batch_dims = 0 : i32}> {tac.device = "DARWINN", tac.inference_type = "FLOAT"} : (tensor<?x!tf_type.string>, tensor<?xi32>) -> tensor<?x!tf_type.string>
// CHECK: %5 = "tfl.reshape"(%4, %arg6) {tac.device = "DARWINN", tac.inference_type = "FLOAT"} : (tensor<?x!tf_type.string>, tensor<1xi32>) -> tensor<?x!tf_type.string>
// CHECK: %6 = "tfl.shape"(%5) {tac.device = "DARWINN", tac.inference_type = "FLOAT"} : (tensor<?x!tf_type.string>) -> tensor<1xi32>
// CHECK: %7 = "tfl.fill"(%6, %arg7) {tac.device = "DARWINN", tac.inference_type = "FLOAT"} : (tensor<1xi32>, tensor<!tf_type.string>) -> tensor<?x!tf_type.string>
// CHECK: %8 = tfl.add %arg8, %arg9 {fused_activation_function = "NONE", tac.device = "DARWINN", tac.inference_type = "FLOAT"} : tensor<?x1xi32>
// CHECK: %9 = "tfl.reshape"(%8, %arg6) {tac.device = "DARWINN", tac.inference_type = "FLOAT"} : (tensor<?x1xi32>, tensor<1xi32>) -> tensor<?xi32>
// CHECK: %10 = "tfl.gather"(%arg10, %9) <{axis = 0 : i32, batch_dims = 0 : i32}> {tac.device = "DARWINN", tac.inference_type = "FLOAT"} : (tensor<?x!tf_type.string>, tensor<?xi32>) -> tensor<?x!tf_type.string>
// CHECK: %11 = "tfl.reshape"(%10, %arg6) {tac.device = "DARWINN", tac.inference_type = "FLOAT"} : (tensor<?x!tf_type.string>, tensor<1xi32>) -> tensor<?x!tf_type.string>
// CHECK: %12 = "tfl.shape"(%11) {tac.device = "DARWINN", tac.inference_type = "FLOAT"} : (tensor<?x!tf_type.string>) -> tensor<1xi32>
// CHECK: %13 = "tfl.fill"(%12, %arg7) {tac.device = "DARWINN", tac.inference_type = "FLOAT"} : (tensor<1xi32>, tensor<!tf_type.string>) -> tensor<?x!tf_type.string>
// CHECK: %14 = "tfl.gather"(%arg11, %arg2) <{axis = 0 : i32, batch_dims = 0 : i32}> {tac.device = "DARWINN", tac.inference_type = "FLOAT"} : (tensor<?x?xi32>, tensor<i32>) -> tensor<?xi32>
// CHECK: %15 = "tfl.equal"(%14, %arg12) {tac.device = "DARWINN", tac.inference_type = "FLOAT"} : (tensor<?xi32>, tensor<i32>) -> tensor<?xi1>
// CHECK: %16 = "tfl.equal"(%14, %arg3) {tac.device = "DARWINN", tac.inference_type = "FLOAT"} : (tensor<?xi32>, tensor<i32>) -> tensor<?xi1>
// CHECK: %17 = "tfl.gather"(%arg13, %14) <{axis = 0 : i32, batch_dims = 0 : i32}> {tac.device = "DARWINN", tac.inference_type = "FLOAT"} : (tensor<5xi32>, tensor<?xi32>) -> tensor<?xi32>
// CHECK: %18 = tfl.add %arg14, %17 {fused_activation_function = "NONE", tac.device = "DARWINN", tac.inference_type = "FLOAT"} : tensor<?xi32>
// CHECK: %19 = "tfl.gather"(%arg15, %14) <{axis = 0 : i32, batch_dims = 0 : i32}> {tac.device = "DARWINN", tac.inference_type = "FLOAT"} : (tensor<5xi32>, tensor<?xi32>) -> tensor<?xi32>
// CHECK: %20 = tfl.add %arg16, %19 {fused_activation_function = "NONE", tac.device = "DARWINN", tac.inference_type = "FLOAT"} : tensor<?xi32>
// CHECK: %21 = tfl.add %arg17, %arg3 {fused_activation_function = "NONE", tac.device = "DARWINN", tac.inference_type = "FLOAT"} : tensor<i32>
// CHECK: return %1, %5, %7, %11, %13, %15, %16, %18, %20, %21 : tensor<i32>, tensor<?x!tf_type.string>, tensor<?x!tf_type.string>, tensor<?x!tf_type.string>, tensor<?x!tf_type.string>, tensor<?xi1>, tensor<?xi1>, tensor<?xi32>, tensor<?xi32>, tensor<i32>
// CHECK: }
// -----
// CHECK-SKIP-CPU-LABEL: testSkipCpuOps
func.func @testSkipCpuOps(%arg0: tensor<1xf32>) -> (tensor<1xf32>, tensor<1xf32>) {
%0 = "tfl.add"(%arg0, %arg0) {tac.device = "GPU", fused_activation_function = "RELU6", tac.inference_type = "FLOAT"} : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32>
%1 = "tfl.add"(%arg0, %0) {tac.device = "CPU", fused_activation_function = "RELU6", tac.inference_type = "FLOAT"} : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32>
func.return %0, %1 : tensor<1xf32>, tensor<1xf32>
}
// CHECK-SKIP-CPU: %[[RES0:.*]] = call @func_0_GPU_FLOAT(%arg0) {tac.device = "GPU", tac.inference_type = "FLOAT", tac.interface_name = "func_0"} : (tensor<1xf32>) -> tensor<1xf32>
// CHECK-SKIP-CPU: %[[RES1:.*]] = tfl.add %arg0, %[[RES0]] {fused_activation_function = "RELU6", tac.device = "CPU", tac.inference_type = "FLOAT"} : tensor<1xf32>
// CHECK-SKIP-CPU: return %[[RES0]], %[[RES1]] : tensor<1xf32>, tensor<1xf32>
// CHECK-SKIP-CPU: }
// CHECK-SKIP-CPU: func.func private @func_0_GPU_FLOAT(%arg0: tensor<1xf32>) -> tensor<1xf32> attributes {tac.device = "GPU", tac.inference_type = "FLOAT", tac.interface_name = "func_0"} {
// CHECK-SKIP-CPU: %[[RES2:.*]] = tfl.add %arg0, %arg0 {fused_activation_function = "RELU6", tac.device = "GPU", tac.inference_type = "FLOAT"} : tensor<1xf32>
// CHECK-SKIP-CPU: return %[[RES2]] : tensor<1xf32>
// CHECK-SKIP-CPU: }
// -----
// CHECK-SKIP-CPU-LABEL: testSkipCpuOpsWithinLoop
func.func @testSkipCpuOpsWithinLoop(%arg0: tensor<i32>) -> tensor<i32> {
%0 = "tfl.while"(%arg0) ({
^bb0(%block: tensor<i32>):
"tfl.yield"(%block) : (tensor<i32>) -> ()
},{
^bb0(%block: tensor<i32>):
%0 = "tfl.add"(%arg0, %block) {tac.device = "GPU", fused_activation_function = "RELU6", tac.inference_type = "FLOAT"} : (tensor<i32>, tensor<i32>) -> tensor<i32>
"tfl.yield"(%0) : (tensor<i32>) -> ()
}) {tac.device = "CPU", fused_activation_function = "RELU6", tac.inference_type = "FLOAT"} : (tensor<i32>) -> tensor<i32>
func.return %0 : tensor<i32>
}
// CHECK-SKIP-CPU: "tfl.while"
// CHECK-SKIP-CPU: ^bb0(%[[ARG0:.*]]: tensor<i32>):
// CHECK-SKIP-CPU: "tfl.yield"(%[[ARG0]]) : (tensor<i32>) -> ()
// CHECK-SKIP-CPU: }, {
// CHECK-SKIP-CPU: ^bb0(%[[ARG1:.*]]: tensor<i32>):
// CHECK-SKIP-CPU: %[[RES0:.*]] = func.call @func_0_GPU_FLOAT(%{{.*}}, %[[ARG1]]) {tac.device = "GPU", tac.inference_type = "FLOAT", tac.interface_name = "func_0"} : (tensor<i32>, tensor<i32>) -> tensor<i32>
// CHECK-SKIP-CPU: "tfl.yield"(%[[RES0]]) : (tensor<i32>) -> ()
// CHECK-SKIP-CPU: }) {fused_activation_function = "RELU6", tac.device = "CPU", tac.inference_type = "FLOAT"} : (tensor<i32>) -> tensor<i32>
// CHECK-SKIP-CPU: func.func private @func_0_GPU_FLOAT(%[[ARG2:.*]]: tensor<i32>, %[[ARG3:.*]]: tensor<i32>) -> tensor<i32> attributes {tac.device = "GPU", tac.inference_type = "FLOAT", tac.interface_name = "func_0"} {
// CHECK-SKIP-CPU: %[[RES1:.*]] = tfl.add %[[ARG2]], %[[ARG3]] {fused_activation_function = "RELU6", tac.device = "GPU", tac.inference_type = "FLOAT"} : tensor<i32>
// CHECK-SKIP-CPU: return %[[RES1]] : tensor<i32>
// CHECK-SKIP-CPU: }
// -----
// CHECK-IGNORE-INFERENCE-TYPE-LABEL: testIgnoreInferenceType
func.func @testIgnoreInferenceType(%arg0: tensor<1x384x384xf32>, %arg1: tensor<1x1x384x!quant.uniform<i8:f32, 0.003:-128>>) -> (tensor<1x384x384xf32>, tensor<1x1x384x!quant.uniform<i8:f32, 0.003:-128>>) {
// These 2 ops are clustered together when `ignore-inference-type` sets to true.
%0 = "tfl.add"(%arg0, %arg0) {tac.device = "GPU", tac.inference_type = "FLOAT", fused_activation_function = "NONE"} : (tensor<1x384x384xf32>, tensor<1x384x384xf32>) -> tensor<1x384x384xf32>
%1 = "tfl.mul"(%arg1, %arg1) {tac.device = "GPU", tac.inference_type = "QUANTIZED_INT8", fused_activation_function = "NONE"} : (tensor<1x1x384x!quant.uniform<i8:f32, 0.003:-128>>, tensor<1x1x384x!quant.uniform<i8:f32, 0.003:-128>>) -> tensor<1x1x384x!quant.uniform<i8:f32, 0.003:-128>>
func.return %0, %1: tensor<1x384x384xf32>, tensor<1x1x384x!quant.uniform<i8:f32, 0.003:-128>>
}
// CHECK-IGNORE-INFERENCE-TYPE: %[[RES0:.*]]:2 = call @[[FUNC_NAME:.*]](%arg0, %arg1) {tac.device = "GPU", tac.inference_type = "FLOAT", tac.interface_name = "func_0"} : (tensor<1x384x384xf32>, tensor<1x1x384x!quant.uniform<i8:f32, 3.000000e-03:-128>>) -> (tensor<1x384x384xf32>, tensor<1x1x384x!quant.uniform<i8:f32, 3.000000e-03:-128>>)
// CHECK-IGNORE-INFERENCE-TYPE: return %[[RES0]]#0, %[[RES0]]#1 : tensor<1x384x384xf32>, tensor<1x1x384x!quant.uniform<i8:f32, 3.000000e-03:-128>>
// CHECK-IGNORE-INFERENCE-TYPE: }
// CHECK-IGNORE-INFERENCE-TYPE: func.func private @[[FUNC_NAME]](%arg0: tensor<1x384x384xf32>, %arg1: tensor<1x1x384x!quant.uniform<i8:f32, 3.000000e-03:-128>>) -> (tensor<1x384x384xf32>, tensor<1x1x384x!quant.uniform<i8:f32, 3.000000e-03:-128>>) attributes {tac.device = "GPU", tac.inference_type = "FLOAT", tac.interface_name = "func_0"} {
// CHECK-IGNORE-INFERENCE-TYPE: %[[RES1:.*]] = tfl.add %arg0, %arg0 {fused_activation_function = "NONE", tac.device = "GPU", tac.inference_type = "FLOAT"} : tensor<1x384x384xf32>
// CHECK-IGNORE-INFERENCE-TYPE: %[[RES2:.*]] = tfl.mul %arg1, %arg1 {fused_activation_function = "NONE", tac.device = "GPU", tac.inference_type = "QUANTIZED_INT8"} : tensor<1x1x384x!quant.uniform<i8:f32, 3.000000e-03:-128>>
// CHECK-IGNORE-INFERENCE-TYPE: return %[[RES1]], %[[RES2]] : tensor<1x384x384xf32>, tensor<1x1x384x!quant.uniform<i8:f32, 3.000000e-03:-128>>
// CHECK-IGNORE-INFERENCE-TYPE: }
// -----
func.func @testStackVersionPropagation(%arg0: tensor<1x4x4x1920xf32>) -> tensor<1x4x4x64xf32> {
%cst = arith.constant dense<1.200000e+00> : tensor<64xf32>
%0 = "tfl.pseudo_qconst"() <{qtype = tensor<64x1x1x1920x!quant.uniform<i8<-127:127>:f32, 0.002>>, value = dense<1> : tensor<64x1x1x1920xi8>}> : () -> tensor<64x1x1x1920x!quant.uniform<i8<-127:127>:f32, 0.002>>
%1 = "tfl.conv_2d"(%arg0, %0, %cst) <{dilation_h_factor = 1 : i32, dilation_w_factor = 1 : i32, fused_activation_function = "NONE", padding = "VALID", stride_h = 1 : i32, stride_w = 1 : i32}> {tac.delegate_compiler_version = "V_1_0", tac.device = "DARWINN", tac.inference_type = "HYBRID"} : (tensor<1x4x4x1920xf32>, tensor<64x1x1x1920x!quant.uniform<i8<-127:127>:f32, 0.002>>, tensor<64xf32>) -> tensor<1x4x4x64xf32>
%2 = tfl.add %1, %1 {fused_activation_function = "NONE", tac.delegate_compiler_version = "V_1_5", tac.device = "DARWINN", tac.inference_type = "FLOAT"} : tensor<1x4x4x64xf32>
func.return %2 : tensor<1x4x4x64xf32>
}
// CHECK-LABEL: func @testStackVersionPropagation
// CHECK: %[[RES0:.*]] = call @func_0_DARWINN_HYBRID(%arg0, %{{.*}}, %{{.*}}) {tac.device = "DARWINN", tac.inference_type = "HYBRID", tac.interface_name = "func_0"} : (tensor<1x4x4x1920xf32>, tensor<64x1x1x1920x!quant.uniform<i8<-127:127>:f32, 2.000000e-03>>, tensor<64xf32>) -> tensor<1x4x4x64xf32>
// CHECK: %[[RES1:.*]] = call @func_1_DARWINN_FLOAT(%[[RES0]]) {tac.device = "DARWINN", tac.inference_type = "FLOAT", tac.interface_name = "func_1"} : (tensor<1x4x4x64xf32>) -> tensor<1x4x4x64xf32>
// CHECK: return %[[RES1]]
// CHECK: }
// CHECK: func.func private @func_0_DARWINN_HYBRID(%arg0: tensor<1x4x4x1920xf32>, %arg1: tensor<64x1x1x1920x!quant.uniform<i8<-127:127>:f32, 2.000000e-03>>, %arg2: tensor<64xf32>) -> tensor<1x4x4x64xf32> attributes {tac.delegate_compiler_version = "V_1_0", tac.device = "DARWINN", tac.inference_type = "HYBRID", tac.interface_name = "func_0"}
// CHECK: func.func private @func_1_DARWINN_FLOAT(%arg0: tensor<1x4x4x64xf32>) -> tensor<1x4x4x64xf32> attributes {tac.delegate_compiler_version = "V_1_5", tac.device = "DARWINN", tac.inference_type = "FLOAT", tac.interface_name = "func_1"}
tensorflow/compiler/mlir/lite/experimental/tac/tests/tac-filter.mlir
+98
View File
@@ -0,0 +1,98 @@
// 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: tac-opt-all-backends -tfl-tac-filter='use-test-setting=true' %s -split-input-file -verify-diagnostics | FileCheck %s
// expected-remark@below {{Tac filter (0): filter type: function filter SKIP_TARGET_ANNOTATION, filter_pattern: "^testFunction"}}
// expected-remark@below {{Tac filter (1): filter type: function filter INCLUDE_TARGET_ANNOTATION, filter_pattern: "testFunctionInclude"}}
// expected-remark@below {{Tac filter (1) specified but not applied to any op}}
// expected-remark@below {{Tac filter (2): filter type: op filter, filter_pattern: "^test_op"}}
// expected-remark@below {{Tac filter (2) specified but not applied to any op}}
module {
// CHECK-LABEL: testFunctionSkiped
// expected-remark@+1 {{filtered by tac filter (0)}}
func.func @testFunctionSkiped(%arg0: tensor<1xf32>, %arg1: tensor<1xf32>) {
// CHECK: tfl.add
// CHECK-SAME: tac.skip_target_annotation
%0 = "tfl.add"(%arg0, %arg1) {fused_activation_function = "RELU6"} : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32>
// CHECK: tfl.add
// CHECK-SAME: tac.skip_target_annotation
%1 = "tfl.add"(%arg0, %0) {fused_activation_function = "RELU"} : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32>
// CHECK: tfl.relu
// CHECK-SAME: tac.skip_target_annotation
%2 = "tfl.relu"(%arg0) : (tensor<1xf32>) -> tensor<1xf32>
func.return
}
}
// -----
// expected-remark@below {{Tac filter (0): filter type: function filter SKIP_TARGET_ANNOTATION, filter_pattern: "^testFunction"}}
// expected-remark@below {{Tac filter (1): filter type: function filter INCLUDE_TARGET_ANNOTATION, filter_pattern: "testFunctionInclude"}}
// expected-remark@below {{Tac filter (2): filter type: op filter, filter_pattern: "^test_op"}}
// expected-remark@below {{Tac filter (2) specified but not applied to any op}}
module {
// CHECK-LABEL: testFunctionInclude
// CHECK-NOT: tac.skip_target_annotation
// expected-remark@+2 {{filtered by tac filter (0)}}
// expected-remark@+1 {{filtered by tac filter (1)}}
func.func @testFunctionInclude(%arg0: tensor<1xf32>, %arg1: tensor<1xf32>) {
%0 = "tfl.add"(%arg0, %arg1) {fused_activation_function = "RELU6"} : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32>
func.return
}
}
// -----
// expected-remark@below {{Tac filter (0): filter type: function filter SKIP_TARGET_ANNOTATION, filter_pattern: "^testFunction"}}
// expected-remark@below {{Tac filter (0) specified but not applied to any op}}
// expected-remark@below {{Tac filter (1): filter type: function filter INCLUDE_TARGET_ANNOTATION, filter_pattern: "testFunctionInclude"}}
// expected-remark@below {{Tac filter (1) specified but not applied to any op}}
// expected-remark@below {{Tac filter (2): filter type: op filter, filter_pattern: "^test_op"}}
module {
// CHECK-LABEL: testOpFilter
// expected-remark@+1 {{all ops filtered by tac filter (2): "tfl.add", "tfl.relu"}}
func.func @testOpFilter(%arg0: tensor<1xf32>, %arg1: tensor<1xf32>) {
// CHECK: tfl.add
// CHECK-SAME: tac.skip_target_annotation
%0 = "tfl.add"(%arg0, %arg1) {fused_activation_function = "RELU6"} : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32> loc("test_op_0")
// CHECK: tfl.add
// CHECK-NOT: tac.skip_target_annotation
%1 = "tfl.add"(%arg0, %0) {fused_activation_function = "RELU"} : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32> loc("non_test_op")
// CHECK: tfl.relu
// CHECK-SAME: tac.skip_target_annotation
%2 = "tfl.relu"(%arg0) : (tensor<1xf32>) -> tensor<1xf32> loc("test_op_1")
func.return
}
}
// -----
// expected-remark@below {{Tac filter (0): filter type: function filter SKIP_TARGET_ANNOTATION, filter_pattern: "^testFunction"}}
// expected-remark@below {{Tac filter (0) specified but not applied to any op}}
// expected-remark@below {{Tac filter (1): filter type: function filter INCLUDE_TARGET_ANNOTATION, filter_pattern: "testFunctionInclude"}}
// expected-remark@below {{Tac filter (1) specified but not applied to any op}}
// expected-remark@below {{Tac filter (2): filter type: op filter, filter_pattern: "^test_op"}}
module {
// CHECK-LABEL: testOpMultipleResults
// expected-remark@+1 {{all ops filtered by tac filter (2): "tfl.split_v"}}
func.func @testOpMultipleResults(%arg0: tensor<16x4x4xf32>) -> (tensor<7x4x4xf32>, tensor<3x4x4xf32>, tensor<6x4x4xf32>) {
%size_splits = arith.constant dense<[7, 3, 6]> : tensor<3xi32>
%split_dim = arith.constant dense<0> : tensor<i32>
// CHECK: tfl.split_v
// CHECK-SAME: tac.skip_target_annotation
%0, %1, %2 = "tfl.split_v"(%arg0, %size_splits, %split_dim) {num_splits = 3 : i32} : (tensor<16x4x4xf32>, tensor<3xi32>, tensor<i32>) -> (tensor<7x4x4xf32>, tensor<3x4x4xf32>, tensor<6x4x4xf32>) loc("test_op_split"("/tmp/test_model.tflite":0:0))
func.return %0, %1, %2 : tensor<7x4x4xf32>, tensor<3x4x4xf32>, tensor<6x4x4xf32>
}
}
tensorflow/compiler/mlir/lite/experimental/tac/tests/target-annotation.mlir
+105
View File
@@ -0,0 +1,105 @@
// 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: tac-opt-all-backends -tfl-target-annotation='device-specs=GPU' %s -split-input-file -verify-diagnostics | FileCheck %s
func.func @testConv(%arg0: tensor<256x32x32x3xf32>, %arg1: tensor<16x3x3x3xf32>, %arg2: tensor<16xf32>) -> tensor<256x30x30x16xf32> {
// CHECK: tac.device = "GPU", tac.inference_type = "FLOAT"
%0 = "tfl.conv_2d"(%arg0, %arg1, %arg2) {dilation_h_factor = 1 : i32, dilation_w_factor = 1 : i32, fused_activation_function = "NONE", padding = "VALID", stride_h = 1 : i32, stride_w = 1 : i32} : (tensor<256x32x32x3xf32>, tensor<16x3x3x3xf32>, tensor<16xf32>) -> tensor<256x30x30x16xf32>
func.return %0 : tensor<256x30x30x16xf32>
}
// -----
func.func @testDepthwiseConv(%arg0: tensor<1x112x112x32xf32>, %arg1: tensor<1x3x3x32xf32>, %arg2: tensor<32xf32>) -> tensor<1x112x112x32xf32> {
// CHECK: tac.device = "GPU", tac.inference_type = "FLOAT"
%0 = "tfl.depthwise_conv_2d"(%arg0, %arg1, %arg2) {depth_multiplier = 1 : i32, dilation_h_factor = 1 : i32, dilation_w_factor = 1 : i32, fused_activation_function = "NONE", padding = "VALID", stride_h = 1 : i32, stride_w = 1 : i32} : (tensor<1x112x112x32xf32>, tensor<1x3x3x32xf32>, tensor<32xf32>) -> tensor<1x112x112x32xf32>
func.return %0 : tensor<1x112x112x32xf32>
}
// -----
func.func @testAvgPool(%arg0: tensor<256x32x32x3xf32>) -> tensor<256x30x30x16xf32> {
// CHECK: tac.device = "GPU", tac.inference_type = "FLOAT"
%0 = "tfl.average_pool_2d"(%arg0) {filter_height = 3 : i32, filter_width = 3 : i32, fused_activation_function = "NONE", padding = "VALID", stride_h = 1 : i32, stride_w = 1 : i32} : (tensor<256x32x32x3xf32>) -> tensor<256x30x30x16xf32>
func.return %0 : tensor<256x30x30x16xf32>
}
// -----
func.func @testMaxPool(%arg0: tensor<256x32x32x3xf32>) -> tensor<256x30x30x16xf32> {
// CHECK: tac.device = "GPU", tac.inference_type = "FLOAT"
%0 = "tfl.max_pool_2d"(%arg0) {filter_height = 3 : i32, filter_width = 3 : i32, fused_activation_function = "NONE", padding = "VALID", stride_h = 1 : i32, stride_w = 1 : i32} : (tensor<256x32x32x3xf32>) -> tensor<256x30x30x16xf32>
func.return %0 : tensor<256x30x30x16xf32>
}
// -----
func.func @testAddReluPack(%arg0: tensor<1xf32>, %arg1: tensor<1xf32>) {
// CHECK: tac.device = "GPU", tac.inference_type = "FLOAT"
%0 = "tfl.add"(%arg0, %arg1) {fused_activation_function = "RELU6"} : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32>
// CHECK: tac.device = "GPU", tac.inference_type = "FLOAT"
%1 = "tfl.add"(%arg0, %0) {fused_activation_function = "RELU"} : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32>
// CHECK: tac.device = "GPU", tac.inference_type = "FLOAT"
%2 = "tfl.relu"(%arg0) : (tensor<1xf32>) -> tensor<1xf32>
// CHECK: tac.device = "CPU", tac.inference_type = "FLOAT"
%3 = "tfl.pack"(%arg0, %arg1) {axis = 0 : i32, values_count = 2 : i32} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2x1xf32>
func.return
}
func.func @notAnnotateConst(%arg0: tensor<256x32x32x3xf32>) -> tensor<256x30x30x16xf32> {
// CHECK-NOT: tac.device tac.inference_type
%0 = "tfl.pseudo_const"() {value = dense<1.000000e+00> : tensor<16x3x3x3xf32>} : () -> tensor<16x3x3x3xf32>
// CHECK-NOT: tac.device tac.inference_type
%1 = "tfl.pseudo_const"() {value = dense<1.000000e+00> : tensor<16xf32>} : () -> tensor<16xf32>
// CHECK: tac.device = "GPU", tac.inference_type = "FLOAT"
%2 = "tfl.conv_2d"(%arg0, %0, %1) {dilation_h_factor = 1 : i32, dilation_w_factor = 1 : i32, fused_activation_function = "NONE", padding = "VALID", stride_h = 1 : i32, stride_w = 1 : i32} : (tensor<256x32x32x3xf32>, tensor<16x3x3x3xf32>, tensor<16xf32>) -> tensor<256x30x30x16xf32>
func.return %2 : tensor<256x30x30x16xf32>
}
// -----
func.func @notAnnotateQuantizeDequantize(%arg0: tensor<4x384x32x!quant.uniform<i8:f32, 0.2:-3>>) -> tensor<4x384x32xf32> {
// CHECK-NOT: tac.device tac.inference_type
%0 = "tfl.pseudo_const"() {value = dense<[1, 4, 384, 32]> : tensor<4xi32>} : () -> tensor<4xi32>
// CHECK-NOT: tac.device tac.inference_type
%1 = "tfl.pseudo_const"() {value = dense<[4, 384, 32]> : tensor<3xi32>} : () -> tensor<3xi32>
// CHECK: tac.device = "CPU", tac.inference_type = "QUANTIZED_INT8"
%2 = "tfl.reshape"(%arg0, %0) : (tensor<4x384x32x!quant.uniform<i8:f32, 0.2:-3>>, tensor<4xi32>) -> tensor<1x4x384x32x!quant.uniform<i8:f32, 0.2:-3>>
// CHECK-NOT: tac.device tac.inference_type
%3 = "tfl.quantize"(%2) {qtype = tensor<1x4x384x32x!quant.uniform<i8:f32, 0.19:1>>} : (tensor<1x4x384x32x!quant.uniform<i8:f32, 0.2:-3>>) -> tensor<1x4x384x32x!quant.uniform<i8:f32, 0.19:1>>
// CHECK: tac.device = "CPU", tac.inference_type = "QUANTIZED_INT8"
%4 = "tfl.reshape"(%3, %1) : (tensor<1x4x384x32x!quant.uniform<i8:f32, 0.19:1>>, tensor<3xi32>) -> tensor<4x384x32x!quant.uniform<i8:f32, 0.19:1>>
// CHECK-NOT: tac.device tac.inference_type
%5 = "tfl.dequantize"(%4) : (tensor<4x384x32x!quant.uniform<i8:f32, 0.19:1>>) -> tensor<4x384x32xf32>
func.return %5 : tensor<4x384x32xf32>
}
func.func @annotateInferenceType(%arg0: tensor<1x1x384x!quant.uniform<i8:f32, 0.003:-128>>) -> tensor<1x384x384x!quant.uniform<i8:f32, 0.003:-128>>{
// CHECK-NOT: tac.device tac.inference_type
%0 = "tfl.pseudo_qconst"() {qtype = tensor<1x384x1x!quant.uniform<i8:f32, 0.003:-128>>, value = dense<127> : tensor<1x384x1xi8>} : () -> tensor<1x384x1x!quant.uniform<i8:f32, 0.003:-128>>
// CHECK: tac.device = "CPU", tac.inference_type = "QUANTIZED_INT8"
%1 = "tfl.mul"(%arg0, %0) {fused_activation_function = "NONE"} : (tensor<1x1x384x!quant.uniform<i8:f32, 0.003:-128>>, tensor<1x384x1x!quant.uniform<i8:f32, 0.003:-128>>) -> tensor<1x384x384x!quant.uniform<i8:f32, 0.003:-128>>
func.return %1 : tensor<1x384x384x!quant.uniform<i8:f32, 0.003:-128>>
}
// -----
// CHECK-LABEL: testSkipAnnotation
func.func @testSkipAnnotation(%arg0: tensor<256x32x32x3xf32>, %arg1: tensor<16x3x3x3xf32>, %arg2: tensor<16xf32>) -> tensor<256x30x30x16xf32> {
// CHECK-NOT: tac.device
%0 = "tfl.conv_2d"(%arg0, %arg1, %arg2) {dilation_h_factor = 1 : i32, dilation_w_factor = 1 : i32, fused_activation_function = "NONE", padding = "VALID", stride_h = 1 : i32, stride_w = 1 : i32, tac.skip_target_annotation } : (tensor<256x32x32x3xf32>, tensor<16x3x3x3xf32>, tensor<16xf32>) -> tensor<256x30x30x16xf32>
func.return %0 : tensor<256x30x30x16xf32>
}
tensorflow/compiler/mlir/lite/experimental/tac/tflite_import_export.cc
+120
View File
@@ -0,0 +1,120 @@
/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/compiler/mlir/lite/experimental/tac/tflite_import_export.h"
#include <memory>
#include <set>
#include <string>
#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "absl/strings/str_cat.h"
#include "llvm/Support/Path.h"
#include "llvm/Support/ToolOutputFile.h"
#include "mlir/Support/FileUtilities.h" // from @llvm-project
#include "tensorflow/compiler/mlir/lite/experimental/tac/common/targets.h"
#include "tensorflow/compiler/mlir/lite/experimental/tac/common/utils.h"
#include "tensorflow/compiler/mlir/lite/experimental/tac/execution_metadata_exporter.h"
#include "tensorflow/compiler/mlir/lite/experimental/tac/hardwares/target_hardware.h"
#include "tensorflow/compiler/mlir/lite/experimental/tac/utils/utils.h"
namespace mlir {
namespace TFL {
namespace tac {
namespace {
void AttachCostPerDevice(mlir::ModuleOp module,
llvm::ArrayRef<std::string> device_specs) {
std::set<std::string> processed_device_specs;
for (const auto& device_spec : device_specs) {
processed_device_specs.insert(
mlir::TFL::tac::GetCanonicalHardwareName(device_spec));
}
processed_device_specs.insert("CPU");
module.walk([&](mlir::Operation* op) {
if (!mlir::TFL::tac::IsNonConstOp(op) &&
!llvm::isa<func::ReturnOp, func::FuncOp, CallOpInterface>(op))
return;
// Attach cost per target.
// Unsupported op will have negative values.
mlir::SmallVector<mlir::NamedAttribute, 4> device_costs;
for (const auto& device : processed_device_specs) {
auto* target_hardware = mlir::TFL::tac::GetTargetHardware(device);
float cost = -1;
if (target_hardware->IsOpSupported(op)) {
cost = target_hardware->GetOpCost(op);
}
mlir::StringAttr device_identifier =
mlir::StringAttr::get(module.getContext(), device);
auto float_type = mlir::Float32Type::get(module.getContext());
auto float_attr =
mlir::FloatAttr::get(float_type, static_cast<float>(cost));
device_costs.push_back({device_identifier, float_attr});
}
op->setAttr("per_device_costs",
mlir::DictionaryAttr::get(module.getContext(), device_costs));
});
}
} // namespace
//////////// Importer ////////////
absl::StatusOr<OwningOpRef<mlir::ModuleOp>> TfLiteImporter::Import() {
source_mgr_handler_ = std::make_unique<mlir::SourceMgrDiagnosticHandler>(
source_mgr_, &context_);
return ImportFlatbufferOrMlir(
options_.file_name, options_.input_mlir,
/*experimental_prune_unreachable_nodes_unconditionally=*/true,
&source_mgr_, &context_);
}
//////////// Exporter ////////////
absl::Status TfLiteExporter::Export(mlir::ModuleOp module) {
// return absl::OkStatus();
if (options_.export_runtime_metadata) {
// Run the cost model for each device/op.
AttachCostPerDevice(module, options_.target_hardware_backends);
// We will export the runtime metadata with the same name under the same
// directory except with a different extention ".rtmeta".
llvm::SmallString<128> metadata_filename(options_.output_file_name);
const char kRuntimeMetadataName[] = "rtmeta";
llvm::sys::path::replace_extension(metadata_filename, kRuntimeMetadataName);
std::string error_msg;
auto output = mlir::openOutputFile(metadata_filename, &error_msg);
if (output == nullptr) {
return absl::InvalidArgumentError(
absl::StrCat("cannot open output file: ", error_msg));
}
auto result = tflite::ExportRuntimeMetadata(module);
if (!result) {
return absl::InvalidArgumentError("Cannot export runtime metadata.");
}
output->os() << result;
output->keep();
}
return mlir::TFL::tac::ExportFlatbufferOrMlir(options_.output_file_name,
options_.output_mlir, module,
/*enable_select_tf_ops=*/false);
}
} // namespace tac
} // namespace TFL
} // namespace mlir
tensorflow/compiler/mlir/lite/experimental/tac/tflite_import_export.h
+74
View File
@@ -0,0 +1,74 @@
/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_TAC_TFLITE_IMPORT_EXPORT_H_
#define TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_TAC_TFLITE_IMPORT_EXPORT_H_
#include <memory>
#include <string>
#include <vector>
#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "llvm/Support/SourceMgr.h"
#include "tensorflow/compiler/mlir/lite/experimental/tac/tac_importer_exporter.h"
namespace mlir {
namespace TFL {
namespace tac {
// TAC Importer for TFLite.
// This import to MLIR from tflite file or MLIR
class TfLiteImporter : public mlir::TFL::tac::TacImporter {
public:
// Options for configuring the importer.
struct Options {
std::string file_name;
// Whether the input file is an MLIR not tflite file.
bool input_mlir = false;
};
explicit TfLiteImporter(const Options& options) : options_(options) {}
absl::StatusOr<mlir::OwningOpRef<mlir::ModuleOp>> Import() override;
private:
Options options_;
mlir::MLIRContext context_;
llvm::SourceMgr source_mgr_;
std::unique_ptr<mlir::SourceMgrDiagnosticHandler> source_mgr_handler_;
};
// TAC Exporter. It exports the provided Module to a tflite file.
class TfLiteExporter : public mlir::TFL::tac::TacExporter {
public:
// Exporter configuration options.
struct Options {
bool export_runtime_metadata = false;
bool output_mlir = false;
std::string output_file_name;
std::vector<std::string> target_hardware_backends;
};
explicit TfLiteExporter(const Options& options) : options_(options) {}
absl::Status Export(mlir::ModuleOp module) override;
private:
Options options_;
};
} // namespace tac
} // namespace TFL
} // namespace mlir
#endif // TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_TAC_TFLITE_IMPORT_EXPORT_H_
tensorflow/compiler/mlir/lite/experimental/tac/transforms/compute_cost.cc
+115
View File
@@ -0,0 +1,115 @@
/* 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 <memory>
#include <string>
#include "llvm/ADT/ArrayRef.h"
#include "llvm/ADT/DenseSet.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/Support/Casting.h"
#include "llvm/Support/CommandLine.h"
#include "mlir/IR/Builders.h" // from @llvm-project
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "mlir/IR/OpDefinition.h" // from @llvm-project
#include "mlir/IR/Operation.h" // from @llvm-project
#include "mlir/Pass/Pass.h" // from @llvm-project
#include "mlir/Pass/PassRegistry.h" // from @llvm-project
#include "mlir/Support/LLVM.h" // from @llvm-project
#include "tensorflow/compiler/mlir/lite/experimental/tac/common/cost.h"
#include "tensorflow/compiler/mlir/lite/experimental/tac/common/subgraph.h"
#include "tensorflow/compiler/mlir/lite/experimental/tac/common/targets.h"
#include "tensorflow/compiler/mlir/lite/experimental/tac/common/utils.h"
#include "tensorflow/compiler/mlir/lite/experimental/tac/transforms/cost_model.h"
#include "tensorflow/compiler/mlir/lite/experimental/tac/transforms/passes.h"
#include "tensorflow/compiler/mlir/lite/ir/tfl_ops.h"
namespace mlir {
namespace TFL {
namespace tac {
namespace {
// We will caculate the total compute cost for each Func Op.
//
// The compute cost is simply an add-up of the costs of all the operations
// within the FuncOp. (Excluding const ops since they're just "data".)
// We will ignore quant/dequant/requant costs within the Func Op as well,
// intuition:
//
// The assumpution is that quant/dequant/requant will only happen at the begin
// and the end of the FuncOp (basically the "boundaries" of the subgraph).
// So we can imagine if multiple "same-inference-typed" graph are presented at
// the same time, the quant/dequant ops pair can be squashed:
//
// dequant ------------
// |
// ops... FuncOp1
// |
// quant -------------
// | <--- can be squashed
// dequant -------------
// |
// ops... FuncOp2
// |
// quant ---------------
//
// But it's true quant & dequant ops can happen "within" the FuncOp as well,
// normally as "quantization params" adjust. We should check more careful to
// include those as those ops wouldn't be "squashed".
class ComputeCostPass
: public mlir::PassWrapper<ComputeCostPass, mlir::OperationPass<ModuleOp>> {
public:
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(ComputeCostPass)
private:
llvm::StringRef getArgument() const final { return "tfl-compute-cost"; }
llvm::StringRef getDescription() const final {
return "Compute the total cost for each available subgraph.";
}
void runOnOperation() override;
};
void ComputeCostPass::runOnOperation() {
auto module = getOperation();
for (auto func : module.getOps<func::FuncOp>()) {
// We only care about those functions annotated with "tac.interface_name".
auto interface_name = GetInterFaceName(func);
if (!interface_name.has_value()) continue;
auto target = GetTargetAnnotation(func);
if (!target.has_value()) {
func.emitError("we cannot get hardware info for this function.");
signalPassFailure();
}
float total_cost = GetCostForFunc(&func, *target);
OpBuilder builder(func);
UpdateCost(func, total_cost, &builder);
}
}
} // namespace
std::unique_ptr<OperationPass<ModuleOp>> CreateComputeCostPass() {
return std::make_unique<ComputeCostPass>();
}
static PassRegistration<ComputeCostPass> pass;
} // namespace tac
} // namespace TFL
} // namespace mlir
tensorflow/compiler/mlir/lite/experimental/tac/transforms/cost_model.cc
+204
View File
@@ -0,0 +1,204 @@
/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/compiler/mlir/lite/experimental/tac/transforms/cost_model.h"
#include <algorithm>
#include <cstdint>
#include <memory>
#include <string>
#include "absl/strings/str_cat.h"
#include "llvm/Support/Casting.h"
#include "llvm/Support/raw_ostream.h"
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/Builders.h" // from @llvm-project
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/IR/BuiltinTypes.h" // from @llvm-project
#include "mlir/IR/PatternMatch.h" // from @llvm-project
#include "mlir/IR/Visitors.h" // from @llvm-project
#include "mlir/Interfaces/CallInterfaces.h" // from @llvm-project
#include "mlir/Pass/Pass.h" // from @llvm-project
#include "mlir/Pass/PassRegistry.h" // from @llvm-project
#include "mlir/Support/LLVM.h" // from @llvm-project
#include "tensorflow/compiler/mlir/lite/experimental/tac/common/cost.h"
#include "tensorflow/compiler/mlir/lite/experimental/tac/common/targets.h"
#include "tensorflow/compiler/mlir/lite/experimental/tac/common/utils.h"
#include "tensorflow/compiler/mlir/lite/experimental/tac/hardwares/target_hardware.h"
#include "tensorflow/compiler/mlir/lite/ir/tfl_ops.h"
namespace mlir {
namespace TFL {
namespace tac {
namespace {
// These are just fake costs.
constexpr float kDequantCost = 2.0;
constexpr float kQuantCost = 2.0;
constexpr float kRequantCost = 2.0;
// TODO(renjieliu): Ideally this should consider different kinds of SOCs as
// well.
// Get total bytes transferred.
int64_t GetTransferredTensorBytes(func::CallOp from_graph,
func::CallOp to_graph) {
int64_t total_size_transferred = 0;
for (auto input : to_graph.getOperands()) {
Operation* input_op = input.getDefiningOp();
if (input_op && input_op == from_graph.getOperation()) {
auto input_type =
mlir::dyn_cast_or_null<RankedTensorType>(input.getType());
if (input_type == nullptr || !input_type.hasStaticShape()) continue;
// Quantized type does not support getSizeInBits.
if (IsQUI8Type(input_type) || IsQI8Type(input_type)) {
total_size_transferred += input_type.getNumElements() * 8;
} else {
auto s_type = mlir::cast<ShapedType>(input_type);
total_size_transferred +=
s_type.getNumElements() * s_type.getElementTypeBitWidth();
}
}
}
return total_size_transferred;
}
// Get total tensor element size transferred.
int64_t GetTransferredElementCount(func::CallOp from_graph,
func::CallOp to_graph) {
int64_t total_element_count = 0;
for (auto input : to_graph.getOperands()) {
Operation* input_op = input.getDefiningOp();
if (input_op && input_op == from_graph.getOperation()) {
auto input_type =
mlir::dyn_cast_or_null<RankedTensorType>(input.getType());
if (input_type == nullptr || !input_type.hasStaticShape()) continue;
total_element_count += input_type.getNumElements();
}
}
return total_element_count;
}
struct GetOpCostPass
: mlir::PassWrapper<GetOpCostPass, OperationPass<func::FuncOp>> {
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(GetOpCostPass)
llvm::StringRef getArgument() const final { return "tfl-get-op-cost"; }
llvm::StringRef getDescription() const final {
return "Get cost for every op";
}
void runOnOperation() override;
};
void GetOpCostPass::runOnOperation() {
auto func = getOperation();
OpBuilder builder(func);
func.walk([&](Operation* op) {
if (IsNonConstOp(op) && !IsTerminatorOp(op) &&
!llvm::isa<func::ReturnOp, func::FuncOp, CallOpInterface>(op)) {
auto hardware = GetTargetAnnotation(op);
if (!hardware) return;
float cost = GetCostForOp(op, *hardware);
UpdateCost(op, cost, &builder);
}
});
}
} // namespace
float GetCostForOp(Operation* op, const std::string& hardware) {
auto* device_hardware = GetTargetHardware(hardware);
if (device_hardware == nullptr) {
return kDefaultFixedValuedCost;
}
return device_hardware->GetOpCost(op);
}
float GetCostForFunc(func::FuncOp* func, const std::string& hardware) {
auto* device_hardware = GetTargetHardware(hardware);
if (device_hardware == nullptr) {
return kDefaultFixedValuedCost;
}
return device_hardware->GetFuncCost(func);
}
float GetTransferCost(const std::string& from_hardware_str,
const std::string& to_hardware_str,
func::CallOp from_graph, func::CallOp to_graph) {
auto from_hardware = GetTargetHardware(from_hardware_str);
auto to_hardware = GetTargetHardware(to_hardware_str);
if (from_hardware == nullptr) {
from_graph.emitError(absl::StrCat(
"we cannot find the registered hardware: ", from_hardware_str));
}
if (to_hardware == nullptr) {
to_graph.emitError(absl::StrCat("we cannot find the registered hardware: ",
to_hardware_str));
}
const int64_t total_size_transferred =
GetTransferredTensorBytes(from_graph, to_graph);
return to_hardware->GetHardwareSwitchingCost(from_hardware,
total_size_transferred);
}
float GetQuantDequantCost(InferenceType from_inference_type,
InferenceType to_inference_type,
func::CallOp from_graph, func::CallOp to_graph) {
// Same inference type, no dequant/quant happens.
if (from_inference_type == to_inference_type) return 0;
const int64_t total_element_count_transferred =
GetTransferredElementCount(from_graph, to_graph);
if (from_inference_type == FLOAT || from_inference_type == HYBRID) {
// FLOAT <-> HYBRID will have no quant/dequant as well.
if (to_inference_type == FLOAT || to_inference_type == HYBRID) {
return 0;
} else if (to_inference_type == QUANTIZED_INT8 ||
to_inference_type == QUANTIZED_UINT8) {
// QUANT path.
return kQuantCost * total_element_count_transferred;
}
}
if (from_inference_type == QUANTIZED_INT8 ||
from_inference_type == QUANTIZED_UINT8) {
// Dequant path.
if (to_inference_type == FLOAT || to_inference_type == HYBRID) {
return kDequantCost * total_element_count_transferred;
} else if (to_inference_type == QUANTIZED_INT8 ||
to_inference_type == QUANTIZED_UINT8) {
// Requant path.
return kRequantCost * total_element_count_transferred;
}
}
// Default quant/dequant/requant cost.
return kDefaultFixedValuedCost;
}
std::unique_ptr<OperationPass<func::FuncOp>> CreateGetOpCostPass() {
return std::make_unique<GetOpCostPass>();
}
static PassRegistration<GetOpCostPass> pass;
} // namespace tac
} // namespace TFL
} // namespace mlir
tensorflow/compiler/mlir/lite/experimental/tac/transforms/cost_model.h
+64
View File
@@ -0,0 +1,64 @@
/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_TAC_TRANSFORMS_COST_MODEL_H_
#define TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_TAC_TRANSFORMS_COST_MODEL_H_
#include <string>
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "tensorflow/compiler/mlir/lite/experimental/tac/common/targets.h"
namespace mlir {
namespace TFL {
namespace tac {
// TODO(renjieliu): We need to come up with a better strategy to do cost
// estimatation. Maybe build a big lookup table for all the ops.
// TODO(renjieliu): We need to consider what's the default value if we cannot
// analyze the cost.
// ================== Interface ========================
// Get the estimated cost for the op under the given hardware spec senario.
float GetCostForOp(Operation* op, const std::string& hardware);
// Get the estimated cost for the whole function under the given hardware.
float GetCostForFunc(func::FuncOp* func, const std::string& hardware);
// Get the transfer cost given from & to hardware info.
// We will only calculate for the "necessary" tensor transferred.
// from_graph & to_graph are used to compute the "necessary" tensors.
// from_graph
// / \ \
// out1 out2 out3
// \ /
// to_graph
// So only out2 & out3 are counted.
float GetTransferCost(const std::string& from_hardware_str,
const std::string& to_hardware_str,
func::CallOp from_graph, func::CallOp to_graph);
// Get the cross quantization/dequantization boundary cost.
float GetQuantDequantCost(InferenceType from_inference_type,
InferenceType to_inference_type,
func::CallOp from_graph, func::CallOp to_graph);
} // namespace tac
} // namespace TFL
} // namespace mlir
#endif // TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_TAC_TRANSFORMS_COST_MODEL_H_
tensorflow/compiler/mlir/lite/experimental/tac/transforms/device_transform.cc
+220
View File
@@ -0,0 +1,220 @@
/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/compiler/mlir/lite/experimental/tac/transforms/device_transform.h"
#include <string>
#include <utility>
#include "llvm/ADT/SmallVector.h"
#include "llvm/Support/Casting.h"
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/Dialect/Quant/IR/QuantTypes.h" // from @llvm-project
#include "mlir/IR/BuiltinTypes.h" // from @llvm-project
#include "mlir/IR/PatternMatch.h" // from @llvm-project
#include "mlir/Support/LLVM.h" // from @llvm-project
#include "mlir/Transforms/GreedyPatternRewriteDriver.h" // from @llvm-project
#include "tensorflow/compiler/mlir/lite/experimental/tac/common/targets.h"
#include "tensorflow/compiler/mlir/lite/experimental/tac/common/utils.h"
#include "tensorflow/compiler/mlir/lite/experimental/tac/hardwares/target_hardware.h"
#include "tensorflow/compiler/mlir/lite/experimental/tac/transforms/device_transform_gpu.h"
#include "tensorflow/compiler/mlir/lite/ir/tfl_ops.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
namespace mlir {
namespace TFL {
namespace tac {
namespace {
#include "tensorflow/compiler/mlir/lite/experimental/tac/transforms/generated_transform_patterns.inc"
} // namespace
RewritePatternSet GetHardwareRewritePatterns(MLIRContext* context,
const std::string& hardware) {
auto* devce_hardware = GetTargetHardware(hardware);
if (devce_hardware == nullptr) return {context};
return devce_hardware->GetTransformations(context);
}
bool IsSupported(Operation* op, const std::string& hardware) {
auto* devce_hardware = GetTargetHardware(hardware);
if (devce_hardware == nullptr) return {};
return devce_hardware->IsOpSupported(op);
}
// ================== Convert Quantized Op ============================
// Walk through the func and convert the quantize ops to their float version.
void ConvertQuantizedOpToFloat(mlir::func::FuncOp func, OpBuilder* builder) {
func.walk([&](Operation* op) {
// TODO(renjieliu): Find a generic way to deal with const ops.
if (op->hasTrait<OpTrait::IsTerminator>() ||
llvm::isa<TFL::QConstOp, TFL::ConstOp>(op) ||
llvm::isa<TFL::QConstOp, TFL::ConstOp, TF::ConstOp, ConstOp>(op))
return;
bool int8_type_observed = false;
bool uint8_type_observed = false;
for (auto& input : op->getOpOperands()) {
auto input_type = input.get().getType();
if (IsQI8Type(input_type)) {
int8_type_observed = true;
} else if (IsQUI8Type(input_type)) {
uint8_type_observed = true;
}
}
// TODO(renjieliu): We probably should check whether the op supports float
// execution to be safe. Although normally they should support float
// execution. Not Quantized ops.
if (!int8_type_observed && !uint8_type_observed) return;
// Insert dequantize ops for every quantized input.
SmallVector<Value, 4> dequantized_inputs;
for (auto& input : op->getOpOperands()) {
auto input_type = input.get().getType();
if (IsQI8Type(input_type) || IsQUI8Type(input_type) ||
IsQI32Type(input_type)) {
auto dequantized_input_type =
mlir::quant::QuantizedType::castToExpressedType(input_type);
builder->setInsertionPoint(op);
auto dequantize_op = TFL::DequantizeOp::create(
*builder, op->getLoc(), dequantized_input_type, input.get());
dequantized_inputs.push_back(dequantize_op);
} else {
dequantized_inputs.push_back(input.get());
}
}
// Result types.
SmallVector<Type, 4> result_types;
for (auto result_type : op->getResultTypes()) {
if (IsQI8Type(result_type) || IsQUI8Type(result_type)) {
auto dequantized_result_type =
mlir::quant::QuantizedType::castToExpressedType(result_type);
result_types.push_back(dequantized_result_type);
} else {
result_types.push_back(result_type);
}
}
// Build the new float-versioned op.
OperationState state(op->getLoc(), op->getName());
state.operands = dequantized_inputs;
state.types = result_types;
state.attributes = op->getAttrs();
state.successors = op->getSuccessors();
builder->setInsertionPoint(op);
Operation* new_op = builder->create(state);
// Insert quantize ops for every outputs and rewrite.
for (int i = 0; i < op->getNumResults(); ++i) {
auto result = op->getResult(i);
auto result_type = result.getType();
Value new_result = new_op->getResult(i);
if (IsQI8Type(result_type) || IsQUI8Type(result_type)) {
builder->setInsertionPoint(op);
TFL::QuantizeOp quant_op =
TFL::QuantizeOp::create(*builder, op->getLoc(), result_type,
new_result, TypeAttr::get(result_type));
new_result = quant_op.getResult();
}
// Rewire the outputs.
result.replaceAllUsesWith(new_result);
}
// Remove the old op.
op->erase();
});
}
// Fold quantized i32 (normally bias) into their float values.
struct FoldQuantizedI32ToFloat : public OpRewritePattern<TFL::DequantizeOp> {
using OpRewritePattern<TFL::DequantizeOp>::OpRewritePattern;
LogicalResult matchAndRewrite(TFL::DequantizeOp dequant_op,
PatternRewriter& rewriter) const override {
// We only fold i32 -> float pattern.
auto input = dequant_op.getInput().getDefiningOp();
if (!input) return failure();
auto input_dequant = llvm::dyn_cast_or_null<TFL::QConstOp>(input);
if (!input_dequant) return failure();
if (!IsQI32Type(input_dequant.getType())) return failure();
auto output_type =
mlir::dyn_cast_or_null<ShapedType>(dequant_op.getOutput().getType());
if (!output_type || !output_type.getElementType().isF32()) return failure();
auto input_type = mlir::dyn_cast<ShapedType>(input_dequant.getType());
// TODO(renjieliu): support UniformQuantizedPerAxisType.
auto q_type = mlir::dyn_cast_or_null<quant::UniformQuantizedType>(
input_type.getElementType());
if (!q_type) return failure();
const float scale = q_type.getScale();
const float zp = q_type.getZeroPoint();
auto input_values = input_dequant.getValue();
// mapValues always takes a function returning APInt, even when the output
// is actually float.
using DequantizeFuncType = llvm::APInt(const llvm::APInt&);
auto dequantize_func = [&](const APInt& ap_int_value) -> APInt {
const int64_t int_value = ap_int_value.getSExtValue();
const float real = (int_value - zp) * scale;
auto real_int = absl::bit_cast<uint32_t>(real);
return APInt(/*numBits=*/32, real_int);
};
auto dequant_values =
mlir::cast<DenseIntOrFPElementsAttr>(input_values)
.mapValues(Float32Type::get(rewriter.getContext()),
llvm::function_ref<DequantizeFuncType>(dequantize_func));
rewriter.replaceOpWithNewOp<TFL::ConstOp>(dequant_op, dequant_op.getType(),
dequant_values);
return success();
}
};
// If the quant op has no consumer, we will remove them.
struct RemoveUnusedQuant : public OpRewritePattern<TFL::QuantizeOp> {
using OpRewritePattern<TFL::QuantizeOp>::OpRewritePattern;
LogicalResult matchAndRewrite(TFL::QuantizeOp quant_op,
PatternRewriter& rewriter) const override {
if (!quant_op.getResult().use_empty()) return failure();
rewriter.eraseOp(quant_op);
return success();
}
};
void OptimizeQuantizedOpToFloat(func::FuncOp func, MLIRContext* context) {
RewritePatternSet patterns(func.getContext());
patterns
.add<FoldQuantizedI32ToFloat, FoldQuantizeDequantize, RemoveUnusedQuant>(
context);
(void)applyPatternsGreedily(func, std::move(patterns));
}
} // namespace tac
} // namespace TFL
} // namespace mlir
tensorflow/compiler/mlir/lite/experimental/tac/transforms/device_transform.h
+49
View File
@@ -0,0 +1,49 @@
/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_TAC_TRANSFORMS_DEVICE_TRANSFORM_H_
#define TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_TAC_TRANSFORMS_DEVICE_TRANSFORM_H_
#include <string>
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "mlir/IR/PatternMatch.h" // from @llvm-project
#include "tensorflow/compiler/mlir/lite/experimental/tac/common/targets.h"
#include "tensorflow/compiler/mlir/lite/ir/tfl_ops.h"
namespace mlir {
namespace TFL {
namespace tac {
// Returns true if 'op' is supported to run on 'hardware'.
bool IsSupported(Operation* op, const std::string& hardware);
// Return proper rewriter patterns for different hardwares.
RewritePatternSet GetHardwareRewritePatterns(MLIRContext* context,
const std::string& hardware);
// Convert quantized ops to float, this will essentially insert dequantize &
// quantize pair around the op.
void ConvertQuantizedOpToFloat(func::FuncOp func, OpBuilder* builder);
// This will optimize the quantized ops -> float graph.
void OptimizeQuantizedOpToFloat(func::FuncOp func, MLIRContext* context);
} // namespace tac
} // namespace TFL
} // namespace mlir
#endif // TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_TAC_TRANSFORMS_DEVICE_TRANSFORM_H_
tensorflow/compiler/mlir/lite/experimental/tac/transforms/device_transform_gpu.cc
+85
View File
@@ -0,0 +1,85 @@
/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/compiler/mlir/lite/experimental/tac/transforms/device_transform_gpu.h"
#include <memory>
#include <utility>
#include "llvm/ADT/ArrayRef.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/Support/Casting.h"
#include "llvm/Support/raw_ostream.h"
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/Attributes.h" // from @llvm-project
#include "mlir/IR/Builders.h" // from @llvm-project
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/IR/BuiltinTypes.h" // from @llvm-project
#include "mlir/IR/Value.h" // from @llvm-project
#include "mlir/Pass/Pass.h" // from @llvm-project
#include "mlir/Pass/PassRegistry.h" // from @llvm-project
#include "mlir/Support/LLVM.h" // from @llvm-project
#include "mlir/Support/LogicalResult.h" // from @llvm-project
#include "mlir/Transforms/GreedyPatternRewriteDriver.h" // from @llvm-project
#include "tensorflow/compiler/mlir/lite/experimental/tac/common/targets.h"
#include "tensorflow/compiler/mlir/lite/experimental/tac/hardwares/gpu_hardware.h"
#include "tensorflow/compiler/mlir/lite/ir/tfl_ops.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
namespace mlir {
namespace TFL {
namespace tac {
namespace {
struct DeviceTransformGPUPass
: public mlir::PassWrapper<DeviceTransformGPUPass,
OperationPass<func::FuncOp>> {
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(DeviceTransformGPUPass)
llvm::StringRef getArgument() const final {
return "tfl-device-transform-gpu";
}
llvm::StringRef getDescription() const final {
return "Suitable transformation for gpu only.";
}
void getDependentDialects(DialectRegistry& registry) const override {
registry.insert<TF::TensorFlowDialect>();
}
void runOnOperation() override;
};
void DeviceTransformGPUPass::runOnOperation() {
auto func = getOperation();
auto* ctx = &getContext();
RewritePatternSet patterns = GetHardwareRewritePatternsGPU(ctx);
(void)applyPatternsGreedily(func, std::move(patterns));
}
} // namespace
RewritePatternSet GetHardwareRewritePatternsGPU(MLIRContext* context) {
GpuHardware gpu_hardware;
return gpu_hardware.GetTransformations(context);
}
std::unique_ptr<OperationPass<func::FuncOp>> CreateDeviceTransformGPUPass() {
return std::make_unique<DeviceTransformGPUPass>();
}
static PassRegistration<DeviceTransformGPUPass> pass;
} // namespace tac
} // namespace TFL
} // namespace mlir
tensorflow/compiler/mlir/lite/experimental/tac/transforms/device_transform_gpu.h
+34
View File
@@ -0,0 +1,34 @@
/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_TAC_TRANSFORMS_DEVICE_TRANSFORM_GPU_H_
#define TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_TAC_TRANSFORMS_DEVICE_TRANSFORM_GPU_H_
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "mlir/IR/PatternMatch.h" // from @llvm-project
namespace mlir {
namespace TFL {
namespace tac {
// nit: Returns all the gpu suitable transformation patterns.
RewritePatternSet GetHardwareRewritePatternsGPU(MLIRContext* context);
} // namespace tac
} // namespace TFL
} // namespace mlir
#endif // TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_TAC_TRANSFORMS_DEVICE_TRANSFORM_GPU_H_
tensorflow/compiler/mlir/lite/experimental/tac/transforms/device_transform_nnapi.cc
+75
View File
@@ -0,0 +1,75 @@
/* 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 <memory>
#include <utility>
#include "llvm/ADT/ArrayRef.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/Support/Casting.h"
#include "llvm/Support/raw_ostream.h"
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/Attributes.h" // from @llvm-project
#include "mlir/IR/Builders.h" // from @llvm-project
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/IR/BuiltinTypes.h" // from @llvm-project
#include "mlir/IR/Value.h" // from @llvm-project
#include "mlir/Pass/Pass.h" // from @llvm-project
#include "mlir/Pass/PassRegistry.h" // from @llvm-project
#include "mlir/Support/LLVM.h" // from @llvm-project
#include "mlir/Support/LogicalResult.h" // from @llvm-project
#include "mlir/Transforms/GreedyPatternRewriteDriver.h" // from @llvm-project
#include "tensorflow/compiler/mlir/lite/experimental/tac/common/targets.h"
#include "tensorflow/compiler/mlir/lite/experimental/tac/hardwares/nnapi_hardware.h"
#include "tensorflow/compiler/mlir/lite/ir/tfl_ops.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
namespace mlir {
namespace TFL {
namespace tac {
namespace {
struct DeviceTransformNNAPIPass
: public mlir::PassWrapper<DeviceTransformNNAPIPass,
OperationPass<func::FuncOp>> {
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(DeviceTransformNNAPIPass)
llvm::StringRef getArgument() const final {
return "tfl-device-transform-nnapi";
}
llvm::StringRef getDescription() const final {
return "Suitable transformation for nnapi only.";
}
void getDependentDialects(DialectRegistry& registry) const override {
registry.insert<TF::TensorFlowDialect>();
}
void runOnOperation() override;
};
void DeviceTransformNNAPIPass::runOnOperation() {
auto func = getOperation();
auto* ctx = &getContext();
NNAPIHardware nnapi_hardware;
RewritePatternSet patterns = nnapi_hardware.GetTransformations(ctx);
(void)applyPatternsGreedily(func, std::move(patterns));
}
} // namespace
static PassRegistration<DeviceTransformNNAPIPass> pass;
} // namespace tac
} // namespace TFL
} // namespace mlir
tensorflow/compiler/mlir/lite/experimental/tac/transforms/device_transform_patterns.cc
+706
View File
@@ -0,0 +1,706 @@
/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/compiler/mlir/lite/experimental/tac/transforms/device_transform_patterns.h"
#include <limits>
#include <memory>
#include "llvm/ADT/ArrayRef.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/Support/Casting.h"
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/Dialect/Quant/IR/QuantTypes.h" // from @llvm-project
#include "mlir/IR/Attributes.h" // from @llvm-project
#include "mlir/IR/Builders.h" // from @llvm-project
#include "mlir/IR/BuiltinAttributes.h" // from @llvm-project
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/IR/BuiltinTypes.h" // from @llvm-project
#include "mlir/IR/Matchers.h" // from @llvm-project
#include "mlir/IR/PatternMatch.h" // from @llvm-project
#include "mlir/IR/Value.h" // from @llvm-project
#include "mlir/Pass/Pass.h" // from @llvm-project
#include "mlir/Pass/PassRegistry.h" // from @llvm-project
#include "mlir/Support/LLVM.h" // from @llvm-project
#include "mlir/Support/LogicalResult.h" // from @llvm-project
#include "tensorflow/compiler/mlir/lite/experimental/tac/common/targets.h"
#include "tensorflow/compiler/mlir/lite/experimental/tac/common/utils.h"
#include "tensorflow/compiler/mlir/lite/ir/tfl_ops.h"
#include "tensorflow/compiler/mlir/lite/utils/attribute_utils.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/verification_utils.h"
namespace mlir {
namespace TFL {
namespace tac {
namespace {
// ================== Common ========================
// Converts any IntegerAttr to an IntegerAttr of an i32 type.
// The value won't change in the new attribute, but if the value is out of
// the bound of i32, the function returns a failure.
LogicalResult ConvertToI32Attr(IntegerAttr attr, IntegerAttr* attr_i32) {
if (attr.getType().isInteger(/*width=*/32)) {
*attr_i32 = attr;
return success();
}
int64_t value = attr.getInt();
if (value > std::numeric_limits<int>::max() ||
value < std::numeric_limits<int>::min()) {
return failure();
}
*attr_i32 = IntegerAttr::get(
IntegerType::get(attr.getContext(), /*width=*/32), value);
return success();
}
TFL::ReshapeOp InsertReshapeOp(Location loc, Value input, Type element_type,
llvm::ArrayRef<int64_t> new_shape_array,
OpBuilder* builder) {
auto reshape_shape_type = mlir::RankedTensorType::get(
new_shape_array.size(), builder->getIntegerType(32));
// This is to workaround the unnecessary cast i64 -> i32. :(
// TODO(renjieliu): Revisit this later.
SmallVector<int32_t, 4> new_shape_array_i32;
for (auto size : new_shape_array) {
new_shape_array_i32.push_back(
ShapedType::isDynamic(size) ? -1 : static_cast<int32_t>(size));
}
auto new_shape_attr =
mlir::DenseIntElementsAttr::get(reshape_shape_type, new_shape_array_i32);
auto new_shape = TFL::ConstOp::create(*builder, loc, new_shape_attr);
auto reshape_out_type = RankedTensorType::get(new_shape_array, element_type);
return TFL::ReshapeOp::create(*builder, loc, reshape_out_type, input,
new_shape);
}
LogicalResult EnsureBias(Operation* op, int bias_idx,
PatternRewriter& rewriter) {
auto bias = op->getOperand(bias_idx);
if (!mlir::isa<NoneType>(bias.getType())) return failure();
// Proceed to create a zero bias.
auto output = op->getResult(0);
auto output_type = mlir::dyn_cast_or_null<RankedTensorType>(output.getType());
if (!output_type) return failure();
// bias should be a vector sized of the last output dim.
int64_t num_units = output_type.getDimSize(output_type.getRank() - 1);
auto bias_type =
mlir::RankedTensorType::get({num_units}, output_type.getElementType());
mlir::DenseElementsAttr bias_attr;
if (output_type.getElementType().isF32()) {
float val = 0.0;
bias_attr = mlir::DenseFPElementsAttr::get(bias_type, val);
} else {
// TODO(renjieliu): Refactor this and share the logic with
// CreateConstOpWithSingleValue. Also, make sure it works with QConst.
return failure();
}
auto zero_bias = TFL::ConstOp::create(rewriter, op->getLoc(), bias_attr);
op->setOperand(bias_idx, zero_bias);
return success();
}
TF::ConstOp PadConstValues(Operation* input_op, int value_to_pad,
int pad_dimensions, Location loc,
OpBuilder* builder) {
if (input_op == nullptr) return nullptr;
mlir::DenseIntElementsAttr attr;
if (!matchPattern(input_op, m_Constant(&attr))) {
return nullptr;
}
auto value_shape_type = mlir::RankedTensorType::get(
{pad_dimensions}, builder->getIntegerType(32));
SmallVector<int32_t, 4> value_i32;
value_i32.reserve(pad_dimensions);
for (int i = 0; i < pad_dimensions - attr.getNumElements(); ++i) {
value_i32.push_back(value_to_pad);
}
for (const auto& size : attr) {
value_i32.push_back(static_cast<int32_t>(size.getSExtValue()));
}
auto new_value_i32_attr =
mlir::DenseIntElementsAttr::get(value_shape_type, value_i32);
return TF::ConstOp::create(*builder, loc, new_value_i32_attr);
}
SmallVector<Value, 4> SliceOutputs(Operation* split_op, Value input,
RankedTensorType input_type, int split_dim,
int num_splits, PatternRewriter* rewriter) {
SmallVector<Value, 4> slice_outputs;
int begin = 0;
for (int i = 0; i < num_splits; ++i) {
// Create slice op.
// Populate begin & size.
SmallVector<int32_t, 4> slice_begin;
SmallVector<int32_t, 4> slice_size;
auto current_output = split_op->getResult(i);
auto current_output_type =
mlir::cast<RankedTensorType>(current_output.getType());
for (int d = 0; d < input_type.getRank(); ++d) {
if (d == split_dim) {
// Split dimension.
slice_begin.push_back(begin);
int64_t size = current_output_type.getDimSize(d);
slice_size.push_back(size);
begin += size;
} else {
slice_begin.push_back(0);
// -1 means every elements.
slice_size.push_back(-1);
}
}
auto slice_type = mlir::RankedTensorType::get(slice_begin.size(),
rewriter->getIntegerType(32));
auto slice_begin_attr =
mlir::DenseIntElementsAttr::get(slice_type, slice_begin);
auto slice_size_attr =
mlir::DenseIntElementsAttr::get(slice_type, slice_size);
auto slice_begin_const =
TFL::ConstOp::create(*rewriter, split_op->getLoc(), slice_begin_attr);
auto slice_size_const =
TFL::ConstOp::create(*rewriter, split_op->getLoc(), slice_size_attr);
auto slice_op =
TFL::SliceOp::create(*rewriter, split_op->getLoc(), current_output_type,
input, slice_begin_const, slice_size_const);
// Rewire output.
slice_outputs.push_back(slice_op.getResult());
}
return slice_outputs;
}
} // namespace
// ================== Pack ========================
LogicalResult LowerPackIntoConcatReshape::matchAndRewrite(
TFL::PackOp pack_op, PatternRewriter& rewriter) const {
// Pack op should have same shape type.
SmallVector<Value, 5> pack_inputs(pack_op.getValues());
auto input_type = mlir::dyn_cast<RankedTensorType>(pack_inputs[0].getType());
if (!input_type) return failure();
// Figure out output shapes.
SmallVector<int64_t, 4> concat_out_shape;
SmallVector<int64_t, 4> pack_out_shape;
const int64_t rank = input_type.getRank();
int64_t pack_axis = pack_op.getAxis();
size_t count = pack_inputs.size();
if (pack_axis < 0) {
pack_axis += rank;
}
// Concat out shape.
for (int i = 0; i < rank; ++i) {
int64_t dim_size = input_type.getDimSize(i);
if (i == pack_axis) {
dim_size *= count;
}
concat_out_shape.push_back(dim_size);
}
// Pack out shape.
int j = 0;
for (int i = 0; i < rank + 1; ++i) {
if (i == pack_axis) {
pack_out_shape.push_back(count);
} else {
pack_out_shape.push_back(input_type.getDimSize(j));
j++;
}
}
if (failed(TF::VerifyShapeOfReshapeOp(pack_out_shape))) return failure();
// Insert the concat op.
auto concat_out_type =
RankedTensorType::get(concat_out_shape, input_type.getElementType());
auto concat_op =
TFL::ConcatenationOp::create(rewriter, pack_op.getLoc(), concat_out_type,
pack_inputs, pack_op.getAxis(), "NONE");
auto reshape_op =
InsertReshapeOp(pack_op.getLoc(), concat_op, input_type.getElementType(),
pack_out_shape, &rewriter);
// Rewire output & get rid of the pack op.
rewriter.replaceOp(pack_op, reshape_op.getResult());
return success();
}
// ================== squared_difference ========================
LogicalResult SquaredDifference::matchAndRewrite(
TFL::SquaredDifferenceOp squared_diff_op, PatternRewriter& rewriter) const {
auto x = squared_diff_op.getLhs();
auto y = squared_diff_op.getRhs();
auto x_type = mlir::dyn_cast<RankedTensorType>(x.getType());
auto y_type = mlir::dyn_cast<RankedTensorType>(y.getType());
if (!x_type || !y_type) return failure();
if (x_type.getShape() != y_type.getShape()) return failure();
auto result_type = squared_diff_op.getType();
if (!result_type) return failure();
auto sub_op =
TF::SubOp::create(rewriter, squared_diff_op.getLoc(), result_type, x, y);
auto mul_op =
TF::MulOp::create(rewriter, squared_diff_op.getLoc(), result_type,
sub_op.getResult(), sub_op.getResult());
rewriter.replaceOp(squared_diff_op, mul_op.getResult());
return success();
}
// ================== split ========================
LogicalResult UnrollSplit::matchAndRewrite(TFL::SplitOp split_op,
PatternRewriter& rewriter) const {
auto num_splits = split_op.getNumSplits();
auto input = split_op.getValue();
auto input_type = mlir::dyn_cast<RankedTensorType>(input.getType());
if (input_type == nullptr || !input_type.hasStaticShape()) return failure();
for (auto result : split_op.getResults()) {
auto result_type = mlir::dyn_cast<RankedTensorType>(result.getType());
if (result_type == nullptr) return failure();
}
auto output = split_op.getResult(0);
auto output_type = mlir::cast<RankedTensorType>(output.getType());
// TODO(renjieliu): change to use split_dim when we raise the constants
// as well.
int64_t split_dim = -1;
for (int64_t d = 0; d < input_type.getRank(); ++d) {
if (input_type.getDimSize(d) != output_type.getDimSize(d)) split_dim = d;
}
const SmallVector<Value, 4>& slice_outputs = SliceOutputs(
split_op, input, input_type, split_dim, num_splits, &rewriter);
rewriter.replaceOp(split_op, slice_outputs);
return success();
}
// ================== splitV ========================
LogicalResult UnrollSplitV::matchAndRewrite(TFL::SplitVOp splitv_op,
PatternRewriter& rewriter) const {
// We need to make sure both splits & split dim are constants.
auto splits = splitv_op.getSizeSplits().getDefiningOp();
mlir::DenseIntElementsAttr splits_attr;
if (!splits || !matchPattern(splits, m_Constant(&splits_attr)))
return failure();
auto split_dim = splitv_op.getSplitDim().getDefiningOp();
mlir::ElementsAttr split_dim_attr;
if (!split_dim || !matchPattern(split_dim, m_Constant(&split_dim_attr)))
return failure();
auto input = splitv_op.getValue();
auto input_type = mlir::dyn_cast_or_null<RankedTensorType>(input.getType());
if (!input_type || !input_type.hasRank()) return failure();
for (auto result : splitv_op.getResults()) {
auto result_type = mlir::dyn_cast<RankedTensorType>(result.getType());
if (result_type == nullptr) return failure();
}
const int64_t rank = input_type.getRank();
IntegerAttr dim_int = ExtractSingleElementAsInteger(split_dim_attr);
// "axis" operand could be a i64 tensor. Resolve it here.
IntegerAttr dim_i32;
if (failed(ConvertToI32Attr(dim_int, &dim_i32))) return failure();
int dim = dim_i32.getInt();
if (dim < 0) dim += rank;
const SmallVector<Value, 4>& slice_outputs = SliceOutputs(
splitv_op, input, input_type, dim, splitv_op.getNumSplits(), &rewriter);
rewriter.replaceOp(splitv_op, slice_outputs);
return success();
}
// ================== conv_2d ========================
LogicalResult EnsureBiasForConv2d::matchAndRewrite(
TFL::Conv2DOp conv_op, PatternRewriter& rewriter) const {
return EnsureBias(conv_op, 2, rewriter);
}
// ================== slice ============================
// If a slice op has < 4d dimension, will pad it to 4d.
LogicalResult PadSlice::matchAndRewrite(TFL::SliceOp slice_op,
PatternRewriter& rewriter) const {
// We have to know the shape of the input, as well as the begin/size.
// also, begin and size have to be constants.
auto input = slice_op.getInput();
auto input_type = mlir::dyn_cast_or_null<RankedTensorType>(input.getType());
if (!input_type || !input_type.hasStaticShape()) return failure();
if (input_type.getRank() >= 4) return failure();
auto begin = slice_op.getBegin();
auto begin_type = mlir::dyn_cast_or_null<RankedTensorType>(begin.getType());
if (!begin_type || !begin_type.hasStaticShape()) return failure();
auto size = slice_op.getSize();
auto size_type = mlir::dyn_cast_or_null<RankedTensorType>(size.getType());
if (!size_type || !size_type.hasStaticShape()) return failure();
auto output_type =
mlir::dyn_cast_or_null<RankedTensorType>(slice_op.getType());
if (!output_type || !output_type.hasStaticShape()) return failure();
// Pad 0s in front of the begin.
TF::ConstOp new_begin =
PadConstValues(begin.getDefiningOp(), 0, 4, slice_op.getLoc(), &rewriter);
if (!new_begin) return failure();
// Pad 1s in front of the size.
TF::ConstOp new_size =
PadConstValues(size.getDefiningOp(), 1, 4, slice_op.getLoc(), &rewriter);
if (!new_size) return failure();
// Reshape the input to 4d.
SmallVector<int64_t, 4> new_shape;
const int current_rank = input_type.getRank();
for (int i = 0; i < 4 - current_rank; ++i) {
new_shape.push_back(1);
}
for (auto size : input_type.getShape()) {
new_shape.push_back(size);
}
auto reshape_op =
InsertReshapeOp(slice_op.getLoc(), input, input_type.getElementType(),
new_shape, &rewriter);
// Replace with the new slice op.
SmallVector<int64_t, 4> new_output_shape;
for (int i = 0; i < 4 - current_rank; ++i) {
new_output_shape.push_back(1);
}
for (auto size : output_type.getShape()) {
new_output_shape.push_back(size);
}
RankedTensorType new_output_type =
RankedTensorType::get(new_output_shape, output_type.getElementType());
auto new_slice =
TFL::SliceOp::create(rewriter, slice_op.getLoc(), new_output_type,
reshape_op, new_begin, new_size);
// Append a reshape at the bottom.
auto output_reshape_op = InsertReshapeOp(slice_op.getLoc(), new_slice,
output_type.getElementType(),
output_type.getShape(), &rewriter);
rewriter.replaceOp(slice_op, output_reshape_op.getResult());
return success();
}
// ================== fully_connected ========================
// TFL fully_connected basically does:
// Weight * Input + bias.
// Input layout is : [..., depth]
// Weight layout is : [output, depth]
// Bias is [output].
//
// While conv2d is:
// Filter: [NHWC]
// Input is also: [NHWC]
// Bias is [N]
//
// So to perform the transform, we need to insert a few reshape ops:
//
// Input weight bias
// \ / /
// FC
// |
// output
//
// |
// \/
//
// Input weight
// | |
// Reshape Reshape bias
// | | /
// conv
// |
// reshape
// |
// output
LogicalResult FullyConnectedToConv::matchAndRewrite(
TFL::FullyConnectedOp fc_op, PatternRewriter& rewriter) const {
// We have to know the shape of the input.
auto input = fc_op.getInput();
auto input_type = mlir::dyn_cast_or_null<RankedTensorType>(input.getType());
if (!input_type || !input_type.hasStaticShape()) return failure();
// We have to know the shape of the weight.
auto weight = fc_op.getFilter();
auto weight_type = mlir::dyn_cast_or_null<RankedTensorType>(weight.getType());
if (!weight_type || !weight_type.hasStaticShape()) return failure();
// We have to know the shape of the output as well.
auto output = fc_op.getResult(0);
auto output_type = mlir::dyn_cast_or_null<RankedTensorType>(output.getType());
if (!output_type || !output_type.hasStaticShape()) return failure();
// Insert a reshape after the input.
// Since the input maybe more than 2-d, we may collect the flat size of the
// input then reshape into [1, 1, flat_size / depth, depth].
const int64_t depth = input_type.getDimSize(input_type.getRank() - 1);
const int64_t flat_size = input_type.getNumElements();
const int64_t width = flat_size / depth;
SmallVector<int64_t, 4> input_new_shape({1, 1, width, depth});
auto reshaped_input =
InsertReshapeOp(fc_op.getLoc(), input, input_type.getElementType(),
input_new_shape, &rewriter);
// Insert a reshape after the weight.
// We will reshape the weight into [output, 1, 1, depth]
const int64_t output_size = weight_type.getDimSize(0);
SmallVector<int64_t, 2> weight_new_shape({output_size, 1, 1, depth});
auto reshaped_weight =
InsertReshapeOp(fc_op.getLoc(), weight, weight_type.getElementType(),
weight_new_shape, &rewriter);
// Replace the fc with conv.
// The output would be [1, 1, width, output].
auto conv_output_type = RankedTensorType::get({1, 1, width, output_size},
output_type.getElementType());
auto conv = TFL::Conv2DOp::create(
rewriter, fc_op.getLoc(), conv_output_type, reshaped_input,
reshaped_weight, fc_op.getBias(), rewriter.getI32IntegerAttr(1),
rewriter.getI32IntegerAttr(1), fc_op.getFusedActivationFunctionAttr(),
rewriter.getStringAttr("VALID"), rewriter.getI32IntegerAttr(1),
rewriter.getI32IntegerAttr(1));
// Insert a shape after the conv.
auto reshaped_conv =
InsertReshapeOp(fc_op.getLoc(), conv, output_type.getElementType(),
output_type.getShape(), &rewriter);
rewriter.replaceOp(fc_op, reshaped_conv.getResult());
return success();
}
// ================== concat ============================
// If a concat op has < 4d dimension, will pad it to 4d.
LogicalResult PadConcat::matchAndRewrite(TFL::ConcatenationOp concat_op,
PatternRewriter& rewriter) const {
int rank = -1;
for (auto input : concat_op.getValues()) {
auto input_type = mlir::dyn_cast_or_null<RankedTensorType>(input.getType());
if (!input_type || !input_type.hasStaticShape()) return failure();
rank = input_type.getRank();
}
auto output_type =
mlir::dyn_cast_or_null<RankedTensorType>(concat_op.getType());
if (!output_type || !output_type.hasStaticShape()) return failure();
if (rank >= 4) return failure();
// All values should have the same rank.
// We will insert a reshape op after every input.
SmallVector<Value, 4> reshape_ops;
for (auto input : concat_op.getValues()) {
auto input_type = mlir::cast<RankedTensorType>(input.getType());
// Get the new shape.
SmallVector<int64_t, 4> new_shape;
for (int i = 0; i < 4 - rank; ++i) {
new_shape.push_back(1);
}
for (auto size : input_type.getShape()) {
new_shape.push_back(size);
}
auto reshape_op =
InsertReshapeOp(concat_op.getLoc(), input, input_type.getElementType(),
new_shape, &rewriter);
reshape_ops.push_back(reshape_op.getResult());
}
// Deal with the axis.
// We don't need to handle axis < 0, since it's counting reversely.
int32_t axis = concat_op.getAxis();
if (axis >= 0) {
axis += (4 - rank);
}
// Replace with the new concat op.
SmallVector<int64_t, 4> new_output_shape;
for (int i = 0; i < 4 - rank; ++i) {
new_output_shape.push_back(1);
}
for (auto size : output_type.getShape()) {
new_output_shape.push_back(size);
}
RankedTensorType new_output_type =
RankedTensorType::get(new_output_shape, output_type.getElementType());
auto new_concat = TFL::ConcatenationOp::create(
rewriter, concat_op.getLoc(), new_output_type, reshape_ops, axis,
concat_op.getFusedActivationFunction());
// Append a reshape at the bottom.
auto output_reshape_op = InsertReshapeOp(concat_op.getLoc(), new_concat,
output_type.getElementType(),
output_type.getShape(), &rewriter);
rewriter.replaceOp(concat_op, output_reshape_op.getResult());
return success();
}
// ================== mean ========================
// Currently NNAPI does not support mean op with different scales (quantization
// cases), and in TFLite avg_pool will ensure the input & output has the same
// scales.
LogicalResult ReduceMeanToAvgPool::matchAndRewrite(
TFL::MeanOp mean_op, PatternRewriter& rewriter) const {
auto input = mean_op.getInput();
auto input_type = mlir::dyn_cast_or_null<RankedTensorType>(input.getType());
// Only 4d is supported here.
if (!input_type || input_type.getRank() != 4) return failure();
// The axes has to be [1, 2].
DenseElementsAttr axis_const;
if (!matchPattern(mean_op.getAxis(), m_Constant(&axis_const)))
return failure();
if (axis_const.size() != 2) return failure();
auto axis_values = axis_const.getValues<APInt>();
int i = 1;
for (auto axis_value : axis_values) {
if (axis_value != i++) return failure();
}
auto output = mean_op.getOutput();
auto output_type = mlir::dyn_cast_or_null<RankedTensorType>(output.getType());
if (!output_type) return failure();
auto input_quantized_type =
quant::QuantizedType::getQuantizedElementType(input_type);
auto output_quantized_type =
quant::QuantizedType::getQuantizedElementType(output_type);
// If both the input & output types are non-quantized, they will be both
// nullptrs.
if (input_quantized_type != output_quantized_type) {
return failure();
}
int64_t batch = input_type.getDimSize(0);
int64_t height = input_type.getDimSize(1);
int64_t width = input_type.getDimSize(2);
int64_t channel = input_type.getDimSize(3);
auto avg_pool_output_type = RankedTensorType::get(
{batch, 1, 1, channel}, input_type.getElementType());
auto avg_pool = TFL::AveragePool2DOp::create(
rewriter, mean_op.getLoc(), avg_pool_output_type, input,
rewriter.getI32IntegerAttr(height), rewriter.getI32IntegerAttr(width),
rewriter.getStringAttr("VALID"), rewriter.getI32IntegerAttr(1),
rewriter.getI32IntegerAttr(1), rewriter.getStringAttr("NONE"));
auto value_to_replace = avg_pool.getResult();
// If it's not keep dim, we need to insert a reshape after the average
// pool.
if (!mean_op.getKeepDims()) {
// Insert the reshape.
SmallVector<int64_t, 2> new_shape({batch, channel});
auto reshape_op =
InsertReshapeOp(mean_op.getLoc(), avg_pool.getResult(),
input_type.getElementType(), new_shape, &rewriter);
value_to_replace = reshape_op.getResult();
}
rewriter.replaceOp(mean_op, value_to_replace);
return success();
}
// Insert a "requant" op after the mean op if the mean has different scales for
// input & output.
// Please note: THIS IS NOT a mathmetically-equivalent transformation and it may
// loose accuracy, so we need to use this very very carefully.
LogicalResult InsertRequantForReduceMean::matchAndRewrite(
TFL::MeanOp mean_op, PatternRewriter& rewriter) const {
auto input = mean_op.getInput();
auto input_type = mlir::dyn_cast_or_null<ShapedType>(input.getType());
if (!input_type) return failure();
// Only need to do this for quantized input.
auto input_quantized_type =
quant::QuantizedType::getQuantizedElementType(input_type);
if (!input_quantized_type) return failure();
auto output = mean_op.getOutput();
auto output_type = mlir::dyn_cast_or_null<ShapedType>(output.getType());
if (!output_type) return failure();
auto output_quantized_type =
quant::QuantizedType::getQuantizedElementType(output_type);
// If the quantized type is the same, we don't need to do anything.
if (input_quantized_type == output_quantized_type) return failure();
auto new_output_type =
RankedTensorType::get(output_type.getShape(), input_quantized_type);
auto new_mean_op =
TFL::MeanOp::create(rewriter, mean_op->getLoc(), new_output_type, input,
mean_op.getAxis(), mean_op.getKeepDims());
// Insert a requant op.
rewriter.replaceOpWithNewOp<TFL::QuantizeOp>(
mean_op, output_type, new_mean_op, mlir::TypeAttr::get(output_type));
return success();
}
} // namespace tac
} // namespace TFL
} // namespace mlir
tensorflow/compiler/mlir/lite/experimental/tac/transforms/device_transform_patterns.h
+115
View File
@@ -0,0 +1,115 @@
/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_TAC_TRANSFORMS_DEVICE_TRANSFORM_PATTERNS_H_
#define TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_TAC_TRANSFORMS_DEVICE_TRANSFORM_PATTERNS_H_
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/PatternMatch.h" // from @llvm-project
#include "mlir/Support/LogicalResult.h" // from @llvm-project
#include "tensorflow/compiler/mlir/lite/ir/tfl_ops.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
namespace mlir {
namespace TFL {
namespace tac {
// TODO(renjieliu): add more patterns.
// This basically:
// Pack => (Concat -> Reshape)
struct LowerPackIntoConcatReshape : public OpRewritePattern<TFL::PackOp> {
using OpRewritePattern<TFL::PackOp>::OpRewritePattern;
LogicalResult matchAndRewrite(TFL::PackOp pack_op,
PatternRewriter& rewriter) const override;
};
struct SquaredDifference : public OpRewritePattern<TFL::SquaredDifferenceOp> {
using OpRewritePattern<TFL::SquaredDifferenceOp>::OpRewritePattern;
LogicalResult matchAndRewrite(TFL::SquaredDifferenceOp squared_diff_op,
PatternRewriter& rewriter) const override;
};
// Unroll split into a bunch of slice ops.
struct UnrollSplit : public OpRewritePattern<TFL::SplitOp> {
using OpRewritePattern<TFL::SplitOp>::OpRewritePattern;
LogicalResult matchAndRewrite(TFL::SplitOp split_op,
PatternRewriter& rewriter) const override;
};
// Unroll splitv into a bunch of slice ops.
struct UnrollSplitV : public OpRewritePattern<TFL::SplitVOp> {
using OpRewritePattern<TFL::SplitVOp>::OpRewritePattern;
LogicalResult matchAndRewrite(TFL::SplitVOp splitv_op,
PatternRewriter& rewriter) const override;
};
// Ensure bias for conv2d op.
struct EnsureBiasForConv2d : public OpRewritePattern<TFL::Conv2DOp> {
using OpRewritePattern<TFL::Conv2DOp>::OpRewritePattern;
LogicalResult matchAndRewrite(TFL::Conv2DOp conv_op,
PatternRewriter& rewriter) const override;
};
// Pad slice to 4d.
struct PadSlice : public OpRewritePattern<TFL::SliceOp> {
using OpRewritePattern<TFL::SliceOp>::OpRewritePattern;
LogicalResult matchAndRewrite(TFL::SliceOp slice_op,
PatternRewriter& rewriter) const override;
};
// Fully connected to conv2d.
struct FullyConnectedToConv : public OpRewritePattern<TFL::FullyConnectedOp> {
using OpRewritePattern<TFL::FullyConnectedOp>::OpRewritePattern;
LogicalResult matchAndRewrite(TFL::FullyConnectedOp fc_op,
PatternRewriter& rewriter) const override;
};
// Pad concat to 4d.
struct PadConcat : public OpRewritePattern<TFL::ConcatenationOp> {
using OpRewritePattern<TFL::ConcatenationOp>::OpRewritePattern;
LogicalResult matchAndRewrite(TFL::ConcatenationOp concat_op,
PatternRewriter& rewriter) const override;
};
// Convert reduce mean 4d to avg pool.
struct ReduceMeanToAvgPool : public OpRewritePattern<TFL::MeanOp> {
using OpRewritePattern<TFL::MeanOp>::OpRewritePattern;
LogicalResult matchAndRewrite(TFL::MeanOp mean_op,
PatternRewriter& rewriter) const override;
};
// Insert Requant ops for reduce_mean.
struct InsertRequantForReduceMean : public OpRewritePattern<TFL::MeanOp> {
using OpRewritePattern<TFL::MeanOp>::OpRewritePattern;
LogicalResult matchAndRewrite(TFL::MeanOp mean_op,
PatternRewriter& rewriter) const override;
};
} // namespace tac
} // namespace TFL
} // namespace mlir
#endif // TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_TAC_TRANSFORMS_DEVICE_TRANSFORM_PATTERNS_H_
tensorflow/compiler/mlir/lite/experimental/tac/transforms/fold_constants_to_subgraph.cc
+189
View File
@@ -0,0 +1,189 @@
/* 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 <cassert>
#include <memory>
#include <string>
#include "absl/strings/str_cat.h"
#include "llvm/ADT/ArrayRef.h"
#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/DenseSet.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/ADT/StringRef.h"
#include "llvm/Support/Casting.h"
#include "mlir/Dialect/Arith/IR/Arith.h" // from @llvm-project
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/Builders.h" // from @llvm-project
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/IR/BuiltinTypeInterfaces.h" // from @llvm-project
#include "mlir/IR/BuiltinTypes.h" // from @llvm-project
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "mlir/Pass/Pass.h" // from @llvm-project
#include "mlir/Pass/PassRegistry.h" // from @llvm-project
#include "mlir/Support/LLVM.h" // from @llvm-project
#include "stablehlo/dialect/StablehloOps.h" // from @stablehlo
#include "tensorflow/compiler/mlir/lite/experimental/tac/common/subgraph.h"
#include "tensorflow/compiler/mlir/lite/experimental/tac/transforms/passes.h"
#include "tensorflow/compiler/mlir/lite/ir/tfl_ops.h"
namespace mlir {
namespace TFL {
namespace tac {
namespace {
// This pass is used to fold tfl.const ops to each subgraph (func::FuncOp):
// See the example below:
//
// In main:
// %0 = tfl.const...
// %1 = tfl.const...
// %2 = call func_1(..., %0,...)
// %3 = call func_2(..., %0, ..., %1...)
// ...
//
// Then those consts will be copied into each function and replace their usage.
// func_1:
// %0 = tfl.const...
// func_2:
// %0 = tfl.const...
// %1 = tfl.const...
class FoldConstantsToSubgraphPass
: public mlir::PassWrapper<FoldConstantsToSubgraphPass,
mlir::OperationPass<ModuleOp>> {
public:
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(FoldConstantsToSubgraphPass)
llvm::StringRef getArgument() const final {
return "tfl-fold-constants-to-subgraph";
}
llvm::StringRef getDescription() const final {
return "Fold constants into each subgraph.";
}
FoldConstantsToSubgraphPass() = default;
FoldConstantsToSubgraphPass(const FoldConstantsToSubgraphPass& other) {
this->fold_all_constants_flag_ = other.fold_all_constants_flag_;
}
explicit FoldConstantsToSubgraphPass(bool fold_all_constants) {
fold_all_constants_flag_ = fold_all_constants;
}
private:
void runOnOperation() override;
Option<bool> fold_all_constants_flag_{
*this, "fold-all-constants",
llvm::cl::desc("Whether to fold all constants or just i32."),
llvm::cl::init(false)};
};
void CopyConstantIntoFunc(int argument_index, Operation* const_op,
func::FuncOp func) {
assert((llvm::isa<TFL::ConstOp, TFL::QConstOp, arith::ConstantOp,
stablehlo::ConstantOp>(const_op)) &&
"Expect QConst or Const op.");
OpBuilder builder(func.getBody());
auto cloned_const_op = const_op->clone();
cloned_const_op->setLoc(func.getBody().getLoc());
builder.insert(cloned_const_op);
// Rewire the usage.
func.getArgument(argument_index)
.replaceAllUsesWith(cloned_const_op->getResult(0));
}
bool IsConstOrQConstInt(Operation* op) {
if (!llvm::isa<TFL::ConstOp, TFL::QConstOp, arith::ConstantOp,
stablehlo::ConstantOp>(op))
return false;
if (auto arith_const_op = dyn_cast_or_null<arith::ConstantOp>(op)) {
// arith ConstOp path.
auto type =
mlir::cast<ShapedType>(arith_const_op.getType()).getElementType();
if (!type.isInteger(32) && !type.isInteger(64)) return false;
} else if (auto shlo_const_op = dyn_cast_or_null<stablehlo::ConstantOp>(op)) {
auto type =
mlir::cast<ShapedType>(shlo_const_op.getType()).getElementType();
if (!type.isInteger(32) && !type.isInteger(64)) return false;
} else if (auto const_op = dyn_cast_or_null<TFL::ConstOp>(op)) {
// ConstOp path.
auto type = mlir::cast<ShapedType>(const_op.getType()).getElementType();
if (!type.isInteger(32) && !type.isInteger(64)) return false;
} else {
// QConstOp path.
auto qconst_op = dyn_cast<TFL::QConstOp>(op);
auto type =
quant::QuantizedType::getQuantizedElementType(qconst_op.getType());
if (type.getStorageTypeIntegralWidth() != 32) {
return false;
}
}
return true;
}
void FoldConstantsToSubgraphPass::runOnOperation() {
auto module = getOperation();
for (auto fn : module.getOps<func::FuncOp>()) {
fn.walk([&](Operation* op) {
if (!llvm::isa<TFL::ConstOp, TFL::QConstOp, arith::ConstantOp,
stablehlo::ConstantOp>(op))
return;
// We only fold int32/int64 for Const and i32 for QConst if not specify
// all constants flag. (Since they're more like "configs" or i32 biases.)
// We will fold every const ops (and q_const ops) if we speicfy the
// fold_all_constants_flag.
if (!fold_all_constants_flag_) {
if (!IsConstOrQConstInt(op)) return;
}
for (auto consumer : op->getResult(0).getUsers()) {
auto consumer_call = llvm::dyn_cast_or_null<func::CallOp>(consumer);
if (!consumer_call) continue;
auto function_name = consumer_call.getCallee();
// Locate the argument position of the use.
int argument_index = -1;
for (int i = 0; i < consumer_call.getNumOperands(); ++i) {
if (consumer_call.getOperand(i) == op->getResult(0)) {
argument_index = i;
break;
}
}
// Copy the const into the consumer func and replace their usages.
func::FuncOp func = module.lookupSymbol<func::FuncOp>(function_name);
CopyConstantIntoFunc(argument_index, op, func);
}
});
}
}
} // namespace
std::unique_ptr<OperationPass<ModuleOp>> CreateFoldConstantsToSubgraphPass(
bool fold_all_constants) {
return std::make_unique<FoldConstantsToSubgraphPass>(fold_all_constants);
}
static PassRegistration<FoldConstantsToSubgraphPass> pass;
} // namespace tac
} // namespace TFL
} // namespace mlir
tensorflow/compiler/mlir/lite/experimental/tac/transforms/get_alternative_subgraph.cc
+314
View File
@@ -0,0 +1,314 @@
/* 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 <memory>
#include <string>
#include <utility>
#include <vector>
#include "absl/strings/str_cat.h"
#include "llvm/ADT/ArrayRef.h"
#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/DenseSet.h"
#include "llvm/ADT/SetVector.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/ADT/StringRef.h"
#include "llvm/Support/Casting.h"
#include "llvm/Support/CommandLine.h"
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/Attributes.h" // from @llvm-project
#include "mlir/IR/Block.h" // from @llvm-project
#include "mlir/IR/Builders.h" // from @llvm-project
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "mlir/IR/OperationSupport.h" // from @llvm-project
#include "mlir/IR/Value.h" // from @llvm-project
#include "mlir/Interfaces/CallInterfaces.h" // from @llvm-project
#include "mlir/Pass/Pass.h" // from @llvm-project
#include "mlir/Pass/PassRegistry.h" // from @llvm-project
#include "mlir/Support/LLVM.h" // from @llvm-project
#include "mlir/Transforms/GreedyPatternRewriteDriver.h" // from @llvm-project
#include "tensorflow/compiler/mlir/lite/experimental/tac/common/subgraph.h"
#include "tensorflow/compiler/mlir/lite/experimental/tac/common/targets.h"
#include "tensorflow/compiler/mlir/lite/experimental/tac/common/utils.h"
#include "tensorflow/compiler/mlir/lite/experimental/tac/hardwares/target_hardware.h"
#include "tensorflow/compiler/mlir/lite/experimental/tac/transforms/device_transform.h"
#include "tensorflow/compiler/mlir/lite/experimental/tac/transforms/passes.h"
#include "tensorflow/compiler/mlir/lite/ir/tfl_ops.h"
namespace mlir {
namespace TFL {
namespace tac {
namespace {
// Given the function interface name and the InferenceDeviceType, return the
// new function name.
std::string GetFunctionImplName(
std::string interface_name,
const InferenceDeviceType& device_inference_type) {
return absl::StrCat(interface_name, "_", device_inference_type.hardware, "_",
GetInferenceString(device_inference_type.inference_type));
}
// For every device, we will do the following:
// If the inference type is quantized, we will try the float alternative.
// If it's float, we will just keep it as it is.
std::vector<InferenceDeviceType> GetAllAlternativeInferenceDeviceType(
InferenceType inference_type, ArrayRef<std::string> devices) {
std::vector<InferenceDeviceType> all_device_inference_types;
for (const auto& device : devices) {
if (inference_type == QUANTIZED_INT8) {
all_device_inference_types.push_back({device, QUANTIZED_INT8});
} else if (inference_type == QUANTIZED_UINT8) {
all_device_inference_types.push_back({device, QUANTIZED_UINT8});
}
// We will alway enable float.
all_device_inference_types.push_back({device, FLOAT});
}
return all_device_inference_types;
}
// This pass will try to get alternative subgraph:
// Say a subgraph is annotated with CPU (it probably means the ops it contains
// cannot be run on other deviecs):
//
// We will try:
// 1) If we can do some mathmatically equaivalent transformation so this
// subgraph can be run on other devices.
// 2) We will other apply device-specifics optimizations as well, that includes
// maybe tensor layout transformation, device specific fusion, etc.
class AlternativeSubgraphPass
: public mlir::PassWrapper<AlternativeSubgraphPass,
mlir::OperationPass<ModuleOp>> {
public:
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(AlternativeSubgraphPass)
llvm::StringRef getArgument() const final {
return "tfl-get-alternative-subgraph";
}
llvm::StringRef getDescription() const final {
return "Get alternative subgraph representation (if appliable) for all the "
"given devices, will by default include the cpu implementation.";
}
AlternativeSubgraphPass() = default;
AlternativeSubgraphPass(const AlternativeSubgraphPass&) {}
explicit AlternativeSubgraphPass(llvm::ArrayRef<std::string> device_specs) {
device_specs_flag_ = device_specs;
}
private:
void runOnOperation() override;
// Given a func and targeted devices, we will try to clonse the func &
// transform/optimize for those devices.
// This will only happen if the whole subgraph can be supported by the target
// or can be supported after some transformations.
void GetAlternativeGraphForFunc(ArrayRef<std::string> devices,
func::FuncOp func, ModuleOp module,
OpBuilder* builder);
// If all ops in the func op is able to be represented in the hardware, we
// will return true, else will be false.
// This is basically all or nothing.
bool IsAllSupportedbySpec(func::FuncOp func,
const InferenceDeviceType& inference_type);
// Given a func and a targeted device, we will try to clonse the func &
// transform/optimize for that device.
// It's simply clone the FuncOp and hardware specific transformations.
func::FuncOp GetAlternativeViewForSpec(
func::FuncOp func,
const InferenceDeviceType& current_device_inference_type,
const InferenceDeviceType& target_device_inference_type, ModuleOp module,
OpBuilder* builder);
// Apply any device-specific optimizations.
void Optimize(func::FuncOp func, const std::string& hardware);
ListOption<std::string> device_specs_flag_{
*this, "device-specs",
llvm::cl::desc(
"comma separated list of device specs, like CPU, GPU, DPS."),
llvm::cl::ZeroOrMore};
};
void AlternativeSubgraphPass::GetAlternativeGraphForFunc(
ArrayRef<std::string> devices, func::FuncOp func, ModuleOp module,
OpBuilder* builder) {
auto current_device = GetTargetAnnotation(func);
if (current_device->empty()) {
func.emitError(
"cannot find target annotation or unknown device specified for current "
"function");
return;
}
auto current_inference_type = GetInferenceTypeAnnotation(func);
if (!current_inference_type.has_value() ||
current_inference_type == UNKNOWN) {
func.emitError(
"cannot find inference type annotation or unknown inference type "
"specified for current "
"function");
return;
}
const InferenceDeviceType current_device_type(
{*current_device, *current_inference_type});
const std::vector<InferenceDeviceType>& all_inference_device_type =
GetAllAlternativeInferenceDeviceType(*current_inference_type, devices);
for (const auto& device_inference_type : all_inference_device_type) {
if (device_inference_type != current_device_type) {
func::FuncOp cloned_func = GetAlternativeViewForSpec(
func, current_device_type, device_inference_type, module, builder);
// If we found unsupported ops, we will just go ahead and remove this
// function.
// TODO(b/160284136): currently we check if the ops are supported then
// see if we need to erase the func op.
// Ideally it would be nice if we can utilize dynamic illegal op to do
// the job.
if (device_inference_type.hardware != "CPU" &&
!IsAllSupportedbySpec(cloned_func, device_inference_type)) {
cloned_func.erase();
}
}
}
// Perform the device-specific optimization last.
// We need to run the optimization for the current device last because we
// need to avoid any changes made the current graph polluting other
// alternative graph views.
Optimize(func, *current_device);
}
bool AlternativeSubgraphPass::IsAllSupportedbySpec(
func::FuncOp func, const InferenceDeviceType& device_inference_type) {
bool found_unsupported = false;
func.walk([&](Operation* op) {
if (IsNonConstOp(op) && !IsTerminatorOp(op) &&
!llvm::isa<func::ReturnOp, func::FuncOp, CallOpInterface>(op) &&
!IsSupported(op, device_inference_type.hardware)) {
found_unsupported = true;
}
});
return !found_unsupported;
}
void AlternativeSubgraphPass::Optimize(func::FuncOp func,
const std::string& hardware) {
auto* ctx = &getContext();
RewritePatternSet patterns = GetHardwareRewritePatterns(ctx, hardware);
(void)applyPatternsGreedily(func, std::move(patterns));
}
// Get the alternative view of the func for the given device_inference_type.
// It's possible the transformed func can still contain unsupported ops for the
// given device_inference_type.
func::FuncOp AlternativeSubgraphPass::GetAlternativeViewForSpec(
func::FuncOp func, const InferenceDeviceType& current_device_inference_type,
const InferenceDeviceType& target_device_inference_type, ModuleOp module,
OpBuilder* builder) {
func::FuncOp cloned_func = func.clone();
cloned_func.setPrivate();
auto interface_name = GetInterFaceName(func);
if (!interface_name.has_value()) {
func.emitError("the func op does not have interface_name");
return nullptr;
}
cloned_func->setAttr(
kDevice, builder->getStringAttr(target_device_inference_type.hardware));
cloned_func->setAttr(kInferenceType,
builder->getStringAttr(GetInferenceString(
target_device_inference_type.inference_type)));
std::string new_function_name =
GetFunctionImplName(*interface_name, target_device_inference_type);
cloned_func.setName(new_function_name);
// If it's quantized -> float, we need to wrap all the ops around with dequant
// and quant.
if ((current_device_inference_type.inference_type == QUANTIZED_UINT8 ||
current_device_inference_type.inference_type == QUANTIZED_INT8) &&
target_device_inference_type.inference_type == FLOAT) {
OpBuilder cloned_func_builder(cloned_func);
ConvertQuantizedOpToFloat(cloned_func, &cloned_func_builder);
OptimizeQuantizedOpToFloat(cloned_func, &getContext());
}
Optimize(cloned_func, target_device_inference_type.hardware);
// Set device for each op.
cloned_func.walk([&](Operation* op) {
if (IsNonConstOp(op) && !IsTerminatorOp(op) &&
!llvm::isa<func::ReturnOp, func::FuncOp, CallableOpInterface>(op)) {
op->setAttr(kDevice, builder->getStringAttr(
target_device_inference_type.hardware));
op->setAttr(kInferenceType,
builder->getStringAttr(GetInferenceString(
target_device_inference_type.inference_type)));
}
});
module.push_back(cloned_func);
return cloned_func;
}
void AlternativeSubgraphPass::runOnOperation() {
auto module = getOperation();
// Process devices specs.
if (device_specs_flag_.empty()) {
module.emitError("no device specs specified");
signalPassFailure();
}
std::vector<std::string> device_specs;
if (!ProcessTargetDevices(device_specs_flag_, &device_specs)) {
module.emitError("unknown devices specified");
signalPassFailure();
}
SmallVector<func::FuncOp, 25> funcs_to_be_processed;
// We only process if func has device annotations.
for (auto func : module.getOps<func::FuncOp>()) {
auto device_attr = func->getAttrOfType<StringAttr>(kDevice);
if (device_attr != nullptr) funcs_to_be_processed.push_back(func);
}
OpBuilder builder(module);
// Go head to process those funcs.
// We don't process in the previous loop is we're adding new funcs,
// this is to avoid unnecessary processing.
for (auto func : funcs_to_be_processed) {
GetAlternativeGraphForFunc(device_specs, func, module, &builder);
}
}
} // namespace
std::unique_ptr<OperationPass<ModuleOp>> CreateAlternativeSubgraphPass(
llvm::ArrayRef<std::string> device_specs) {
return std::make_unique<AlternativeSubgraphPass>(device_specs);
}
static PassRegistration<AlternativeSubgraphPass> pass;
} // namespace tac
} // namespace TFL
} // namespace mlir
tensorflow/compiler/mlir/lite/experimental/tac/transforms/passes.h
+83
View File
@@ -0,0 +1,83 @@
/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_TAC_TRANSFORMS_PASSES_H_
#define TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_TAC_TRANSFORMS_PASSES_H_
#include <functional>
#include <memory>
#include <string>
#include "llvm/ADT/ArrayRef.h"
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/IR/Operation.h" // from @llvm-project
#include "mlir/Pass/Pass.h" // from @llvm-project
#include "tensorflow/compiler/mlir/lite/experimental/tac/tac_filter.pb.h"
namespace mlir {
namespace TFL {
namespace tac {
class TacModule;
// Create an instance of the TargetAnnotationPass.
// TODO(b/177376459): Remove in favor of the one below.
std::unique_ptr<OperationPass<func::FuncOp>> CreateTargetAnnotationPass(
llvm::ArrayRef<std::string> device_specs);
// Create and instance of TargetAnnotationPass.
std::unique_ptr<OperationPass<func::FuncOp>> CreateTargetAnnotationPass(
const TacModule* module);
// Create an instance of the RaiseTargetSubgraphsPass. If `skip_raise_cpu_ops`,
// we skip clustering for CPU ops for better clustering of ops running on other
// ML accelerators. When `ignore_inference_type` is set to true, the inference
// types are set to "NOT_CARE" for better clustering.
std::unique_ptr<OperationPass<ModuleOp>> CreateRaiseTargetSubgraphsPass(
bool skip_raise_cpu_ops = false, bool ignore_inference_type = false);
// Create an instance of the AlternativeSubgraphPass.
std::unique_ptr<OperationPass<ModuleOp>> CreateAlternativeSubgraphPass(
llvm::ArrayRef<std::string> device_specs);
// Create an instance of ComputeCostPass.
std::unique_ptr<OperationPass<ModuleOp>> CreateComputeCostPass();
// Create an instance of PickSubgraphsPass.
std::unique_ptr<OperationPass<ModuleOp>> CreatePickSubgraphsPass();
// Create an instance of DeviceTransformGPUPass.
std::unique_ptr<OperationPass<func::FuncOp>> CreateDeviceTransformGPUPass();
// Create an instance of GetOpCostPass.
std::unique_ptr<OperationPass<func::FuncOp>> CreateGetOpCostPass();
// Create an instance of FoldConstantsToSubgraphPass.
std::unique_ptr<OperationPass<ModuleOp>> CreateFoldConstantsToSubgraphPass(
bool fold_all_constants);
// Create an instance of TacFilterPass.
std::unique_ptr<OperationPass<ModuleOp>> CreateTacFilterPass(
::third_party::tensorflow::compiler::mlir::lite::experimental::tac::
TacFilters* tac_filters,
std::function<void(mlir::Operation* op,
const google::protobuf::Any& custom_options)>
custom_options_callback);
} // namespace tac
} // namespace TFL
} // namespace mlir
#endif // TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_TAC_TRANSFORMS_PASSES_H_
tensorflow/compiler/mlir/lite/experimental/tac/transforms/pick_subgraphs.cc
+526
View File
@@ -0,0 +1,526 @@
/* 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 <limits>
#include <memory>
#include <queue>
#include <string>
#include <unordered_map>
#include <vector>
#include "llvm/ADT/ArrayRef.h"
#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/DenseSet.h"
#include "llvm/ADT/SetVector.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/ADT/StringRef.h"
#include "llvm/Support/Casting.h"
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/Attributes.h" // from @llvm-project
#include "mlir/IR/Block.h" // from @llvm-project
#include "mlir/IR/Builders.h" // from @llvm-project
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/IR/BuiltinTypes.h" // from @llvm-project
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "mlir/IR/PatternMatch.h" // from @llvm-project
#include "mlir/IR/Value.h" // from @llvm-project
#include "mlir/Pass/Pass.h" // from @llvm-project
#include "mlir/Pass/PassRegistry.h" // from @llvm-project
#include "mlir/Support/LLVM.h" // from @llvm-project
#include "tensorflow/compiler/mlir/lite/experimental/tac/common/cost.h"
#include "tensorflow/compiler/mlir/lite/experimental/tac/common/subgraph.h"
#include "tensorflow/compiler/mlir/lite/experimental/tac/common/targets.h"
#include "tensorflow/compiler/mlir/lite/experimental/tac/common/utils.h"
#include "tensorflow/compiler/mlir/lite/experimental/tac/transforms/cost_model.h"
#include "tensorflow/compiler/mlir/lite/experimental/tac/transforms/passes.h"
#include "tensorflow/compiler/mlir/lite/ir/tfl_ops.h"
// This pass is used to "determine" the best combinations of the whole "graph".
//
// Assume we have the graph looks like below:
// subgraph1 (CPU/GPU) subgraph2 (CPU)
// \ /
// subgraph3 (CPU/GPU) subgraph4 (CPU/GPU)
// | /
// subgraph5 (CPU/GPU)
// |
// subgraph6 (CPU)
//
// We want to evaluate the possible options and minize the overall costs to
// produce a graph like below:
//
// subgraph1 (GPU) subgraph2(CPU)
// \ /
// subgraph3 (GPU) subgraph4(GPU)
// | /
// subgraph5 (GPU)
// |
// subgraph6 (CPU)
//
// The overall workflow of the pick subgraphs pass:
// 1) Build subgraphs
// 1.1) Collect output subgraphs.
// 1.2) Build `Subgraph` and their "alternative view" from FuncOp.
// 2) Pick subgraphs
// 2.1) Populate the "dp table" for (subgraph, hardware).
// 2.2) Make decisions based on the populated dp table.
// 2.3) Rewire the whole graph based on the desicions.
//
namespace mlir {
namespace TFL {
namespace tac {
namespace {
// GrapView is used to hold the aggregated cost for the given hardware
// view.
struct GraphView {
float total_cost;
std::unordered_map<Operation*, InferenceDeviceType> input_subgraph_plans;
};
// Subgraph is to hold the "conceptual" subgraph.
// A subgraph may associate with 1...n FuncOp, and each FuncOp may correspond
// with different hardwares.
struct Subgraph {
// The call can be thought as an "API".
func::CallOp call;
// available_choces can be viewed as "real implementation" assosicated with
// the hardware.
std::unordered_map<InferenceDeviceType, func::FuncOp,
InferenceDeviceType::inference_device_type_hash>
available_choices;
// This will include self (the subgraph itself).
// subgraphn
// |
// current_subgraph <- aggregated cost
std::unordered_map<InferenceDeviceType, GraphView,
InferenceDeviceType::inference_device_type_hash>
aggregated_cost_with_decisions;
};
// If the output is produced by a callop, will return the callop, otherwise,
// will return nullptr.
inline func::CallOp GetProducerCallOpOrNull(Value output) {
Operation* output_op = output.getDefiningOp();
if (output_op != nullptr && llvm::isa<func::CallOp>(output_op)) {
return llvm::cast<func::CallOp>(output_op);
}
return nullptr;
}
class PickSubgraphsPass
: public mlir::PassWrapper<PickSubgraphsPass,
mlir::OperationPass<ModuleOp>> {
public:
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(PickSubgraphsPass)
private:
llvm::StringRef getArgument() const final { return "tfl-pick-subgraphs"; }
llvm::StringRef getDescription() const final {
return "Pick the best subgraphs to minimize the overall total costs.";
}
void runOnOperation() override;
std::unordered_map<std::string, std::vector<func::FuncOp>>
CollectSubgraphFuncs(ModuleOp module);
void BuildSubgraphs(
func::FuncOp main_fn,
const std::unordered_map<std::string, std::vector<func::FuncOp>>&
func_impls,
llvm::SetVector<Operation*>* unprocessed_subgraphs,
SmallVector<func::CallOp, 4>* output_subgraphs);
void ProcessSubgraph(func::CallOp current_graph,
llvm::SetVector<Operation*>* unprocessed_subgraphs);
bool PickSubgraphs(
llvm::SetVector<Operation*>* all_subgraphs,
ArrayRef<func::CallOp> output_subgraphs,
const std::unordered_map<std::string, std::vector<func::FuncOp>>&
collected_impl_funcs,
OpBuilder* builder);
// Make the decisions based on the subgraphs.
// It may be the case we cannot decide the best scenarios for the user,
// in this case, we just return false.
bool MakeDecisions(ArrayRef<func::CallOp> output_subgraphs);
// Rewire the subgraphs based on the decisions made.
// If we cannot make a decisions, we just don't do anything.
// TODO(renjieliu): we may change the vector to a map of hardware with
// corresponding ipml.
void RewireSubgraphs(
const std::unordered_map<std::string, std::vector<func::FuncOp>>&
collected_impl_funcs,
OpBuilder* builder);
float GetCostOrFail(func::FuncOp func);
llvm::DenseMap<Operation*, Subgraph> subgraphs_;
llvm::DenseMap<Operation*, InferenceDeviceType> decisions_;
};
float PickSubgraphsPass::GetCostOrFail(func::FuncOp func) {
float self_cost;
if (!GetCostOnOp(func, &self_cost)) {
func.emitError("we cannot find cost for this func");
signalPassFailure();
}
return self_cost;
}
// Here we choose to do a greedy dynamic programming based algorithm for
// simplicity.
//
// See the following graph:
//
// input_subgraph_1 .... input_subgraph_n
// \ /
// \ /
// current_subgraph
// / | \
//
// Assume all the input subgraphs of the current subgraph are independent.
// If we already got optimal results for all the input subgraphs.
// Then the current_subgraph's aggregated optimal costs with regards to target
// perspective is simply:
// for target in current_subgraph.supported_targets:
// total_cost = 0
// for input_subgraph in current_subgraph.input_subgraphs:
// input_cost = kInfinity
// for input_target in input_subgraphs.upported_targets:
// # cost = aggregated cost for input_subgraph with transfer cost.
// input_cost = min(input_cost, cost)
// total_cost += input_cost
// total_cost += current_subgraph.get_computation_cost(target)
//
// Note: for input subgraphs are not independent case, the dp case it a little
// bit complicated to handle. A potential thought is resolve only where
// conflict "happened".
//
// The above mentioned thought should probably be revisited for better thought
// or expanded more for more careful design.
// TODO(renjieliu): We may revisit this later.
void PickSubgraphsPass::ProcessSubgraph(
func::CallOp current_graph_call,
llvm::SetVector<Operation*>* unprocessed_subgraphs) {
Subgraph& current_subgraph = subgraphs_.find(current_graph_call)->second;
std::vector<Subgraph*> input_subgraphs;
for (auto input : current_graph_call.getOperands()) {
func::CallOp input_call = GetProducerCallOpOrNull(input);
// If the input subgraph is not processed yet, we just go ahead and process
// that one first.
if (input_call == nullptr) continue;
if (unprocessed_subgraphs->count(input_call) > 0) {
unprocessed_subgraphs->remove(input_call);
ProcessSubgraph(input_call, unprocessed_subgraphs);
}
Subgraph& input_subgraph = subgraphs_.find(input_call)->second;
input_subgraphs.push_back(&input_subgraph);
}
// Find the best plan for the current subgraph.
for (const auto& kv : current_subgraph.available_choices) {
const auto& current_inference_device_type = kv.first;
func::FuncOp impl_target = kv.second;
float self_compute_cost = GetCostOrFail(impl_target);
GraphView current_graph_view;
auto& input_subgraph_plans = current_graph_view.input_subgraph_plans;
float inputs_total_costs = 0.0;
for (Subgraph* input_subgraph : input_subgraphs) {
float input_total_cost = std::numeric_limits<float>::max();
for (const auto& input_kv : input_subgraph->available_choices) {
const auto& input_inference_device_type = input_kv.first;
func::FuncOp input_impl_target = input_kv.second;
float input_compute_cost = GetCostOrFail(input_impl_target);
float transfer_cost =
GetTransferCost(input_inference_device_type.hardware,
current_inference_device_type.hardware,
input_subgraph->call, current_graph_call);
float quant_dequant_cost =
GetQuantDequantCost(input_inference_device_type.inference_type,
current_inference_device_type.inference_type,
input_subgraph->call, current_graph_call);
float summed_cost =
transfer_cost + quant_dequant_cost + input_compute_cost;
if (summed_cost < input_total_cost) {
// Looks this hardware is better for this input_subgraph, let's change
// it.
input_total_cost = summed_cost;
input_subgraph_plans[input_subgraph->call] =
input_inference_device_type;
}
} // for every hardware of input_subgraph
inputs_total_costs += input_total_cost;
} // for every input_subgraph
current_graph_view.total_cost = inputs_total_costs + self_compute_cost;
current_subgraph
.aggregated_cost_with_decisions[current_inference_device_type] =
current_graph_view;
} // for every subgraph
}
void PickSubgraphsPass::BuildSubgraphs(
func::FuncOp fn,
const std::unordered_map<std::string, std::vector<func::FuncOp>>&
func_impls,
llvm::SetVector<Operation*>* unprocessed_subgraphs,
SmallVector<func::CallOp, 4>* output_subgraphs) {
llvm::DenseSet<Operation*> returned_call_op_set;
// Collect all returns first from the main function.
// all the outputs of the main function are actually the final outputs.
// main_func:
// %output1 = call @subgraph_1...
// ...
// %output2 = call @subgraph_m...
// ...
// %outputn = call @subgraph_k...
// return %output1, output2, ..., outputn.
fn.walk([&](func::ReturnOp return_op) {
for (auto output : return_op.getOperands()) {
func::CallOp output_call = GetProducerCallOpOrNull(output);
if (output_call != nullptr) {
returned_call_op_set.insert(output_call);
}
}
});
// Each call op actually is the entry of the subgraph.
fn.walk([&](func::CallOp call_op) {
auto interface_name = GetInterFaceName(call_op);
// we only need to care about the call ops those have interface_name.
if (!interface_name.has_value()) return;
unprocessed_subgraphs->insert(call_op);
// Build the subgraph.
Subgraph subgraph;
subgraph.call = call_op;
auto impl_iter = func_impls.find(*interface_name);
if (impl_iter == func_impls.end()) {
call_op.emitError(
"we cannot find corresponding implementation for this call op");
signalPassFailure();
}
for (auto impl : impl_iter->second) {
auto inference_device_type = GetInferenceDeviceTypeForOp(impl);
if (!inference_device_type.has_value()) {
impl.emitError("we cannot find inference device type for this func");
signalPassFailure();
}
subgraph.available_choices.emplace(*inference_device_type, impl);
}
// Insert in the subgraphs.
subgraphs_.try_emplace(call_op, subgraph);
// If it's an output subgraph, we will add to the output_subgraphs.
if (returned_call_op_set.find(call_op) != returned_call_op_set.end()) {
output_subgraphs->push_back(call_op);
}
});
}
// Collect all the subgraphs (and their alternatives) in the module.
std::unordered_map<std::string, std::vector<func::FuncOp>>
PickSubgraphsPass::CollectSubgraphFuncs(ModuleOp module) {
std::unordered_map<std::string, std::vector<func::FuncOp>> func_impls;
for (auto func : module.getOps<func::FuncOp>()) {
auto interface_name = GetInterFaceName(func);
if (interface_name.has_value()) {
auto impls_iter = func_impls.find(*interface_name);
if (impls_iter == func_impls.end())
impls_iter =
func_impls.emplace(*interface_name, std::vector<func::FuncOp>())
.first;
impls_iter->second.push_back(func);
}
}
return func_impls;
}
// Given the final outputs, evaluate on the overall costs and pick the best
// plan, if we cannot make a decision, nothing would change, just fallback
// to the original plan.
bool PickSubgraphsPass::MakeDecisions(ArrayRef<func::CallOp> output_subgraphs) {
// BFS to make decisions.
std::queue<const GraphView*> processing_queue;
for (func::CallOp output : output_subgraphs) {
const GraphView* preferred_graph_view;
float minimum_cost = std::numeric_limits<float>::max();
const Subgraph& subgraph = subgraphs_.find(output)->second;
for (const auto& kv : subgraph.aggregated_cost_with_decisions) {
if (minimum_cost > kv.second.total_cost) {
minimum_cost = kv.second.total_cost;
preferred_graph_view = &kv.second;
decisions_[output] = kv.first;
}
}
processing_queue.push(preferred_graph_view);
}
// If we see conflict, we will just abort.
while (!processing_queue.empty()) {
const GraphView* current = processing_queue.front();
processing_queue.pop();
for (const auto& input_with_plans : current->input_subgraph_plans) {
func::CallOp input = llvm::cast<func::CallOp>(input_with_plans.first);
const InferenceDeviceType& input_decision = input_with_plans.second;
auto made_input_decision_it = decisions_.find(input);
if (made_input_decision_it == decisions_.end()) {
// Input is not processed.
// Let's process it, also push it to the queue.
decisions_[input] = input_decision;
const Subgraph& input_subgraph = subgraphs_.find(input)->second;
const GraphView& input_subgraph_view =
input_subgraph.aggregated_cost_with_decisions.find(input_decision)
->second;
processing_queue.push(&input_subgraph_view);
} else if (made_input_decision_it->second != input_decision) {
// We see confliction, we need to abort.
return false;
}
}
}
return true;
}
// This rewire subgraph is essentially "hook" the call op with the "best" choice
// (subgraph).
void PickSubgraphsPass::RewireSubgraphs(
const std::unordered_map<std::string, std::vector<func::FuncOp>>&
collected_impl_funcs,
OpBuilder* builder) {
for (auto& kv : decisions_) {
func::CallOp call = llvm::cast<func::CallOp>(kv.first);
const InferenceDeviceType& preferred_inference_device_type = kv.second;
// We need to rewire the call.
std::string interface_name = *GetInterFaceName(call);
for (auto impl : collected_impl_funcs.find(interface_name)->second) {
const auto& impl_inference_device_type =
GetInferenceDeviceTypeForOp(impl);
if (*impl_inference_device_type == preferred_inference_device_type) {
if (call.getCallee() != impl.getName()) {
// We need to rebuild the call op. :(
builder->setInsertionPoint(call);
auto new_call = func::CallOp::create(*builder, call.getLoc(), impl,
call.getOperands());
// Set interface_name & target to the call_op as well.
new_call->setAttr(kInterfaceNameAttr,
builder->getStringAttr(interface_name));
new_call->setAttr(
kDevice,
builder->getStringAttr(preferred_inference_device_type.hardware));
new_call->setAttr(
kInferenceType,
builder->getStringAttr(GetInferenceString(
preferred_inference_device_type.inference_type)));
call.replaceAllUsesWith(new_call.getResults());
call.erase();
}
}
}
}
}
bool PickSubgraphsPass::PickSubgraphs(
llvm::SetVector<Operation*>* all_subgraphs,
ArrayRef<func::CallOp> output_subgraphs,
const std::unordered_map<std::string, std::vector<func::FuncOp>>&
collected_impl_funcs,
OpBuilder* builder) {
// Process those collected unprocessed subgraphs.
//
// Algorithm complexity for this:
// This Complexity should be O(edge * specs ^ 2).
// We should expect the specs to be a small number.
// In future, the spesc can be Hardwares x inference_types
// The Hardware can be {CPU, GPU, DSP, EDGE_TPU}
// The inference_types can be {float, Q_INT8, float16}.
// But still, we should expect the specs to be a small number.
//
// The process is essentially evaluating the accumulated cost for the dp table
// for all the subgraphs (and their alternatives).
while (!all_subgraphs->empty()) {
func::CallOp current_subgraph =
llvm::cast<func::CallOp>(all_subgraphs->front());
all_subgraphs->remove(current_subgraph);
ProcessSubgraph(current_subgraph, all_subgraphs);
}
// Make decisions given the "outputs" and the populated dp table.
// This is hoping to achieve a global minimum.
if (!MakeDecisions(output_subgraphs)) {
return false;
}
// Once the design has been made.
// Start from the outputs and go back and checkout the plan.
RewireSubgraphs(collected_impl_funcs, builder);
return true;
}
void PickSubgraphsPass::runOnOperation() {
auto module = getOperation();
// Collect & build the subgraphs.
// Also collect the output subgraphs.
// Output subgraphs are essentially those subgraphs pointed by the return
// op.
const std::unordered_map<std::string, std::vector<func::FuncOp>> func_impls =
CollectSubgraphFuncs(module);
llvm::SetVector<Operation*> unprocessed_subgraphs;
SmallVector<func::CallOp, 4> output_subgraphs;
for (auto fn : module.getOps<func::FuncOp>()) {
BuildSubgraphs(fn, func_impls, &unprocessed_subgraphs, &output_subgraphs);
}
OpBuilder builder(module);
if (!PickSubgraphs(&unprocessed_subgraphs, output_subgraphs, func_impls,
&builder)) {
module.emitWarning(
"we cannot find the best scenarios for your case, so we just use "
"your original model plans");
}
}
} // namespace
std::unique_ptr<OperationPass<ModuleOp>> CreatePickSubgraphsPass() {
return std::make_unique<PickSubgraphsPass>();
}
static PassRegistration<PickSubgraphsPass> pass;
} // namespace tac
} // namespace TFL
} // namespace mlir
tensorflow/compiler/mlir/lite/experimental/tac/transforms/raise_target_subgraphs.cc
+341
View File
@@ -0,0 +1,341 @@
/* 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 <memory>
#include <optional>
#include <string>
#include <unordered_map>
#include <utility>
#include "absl/strings/str_cat.h"
#include "llvm/ADT/ArrayRef.h"
#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/DenseSet.h"
#include "llvm/ADT/SetVector.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/ADT/StringMap.h"
#include "llvm/ADT/StringRef.h"
#include "llvm/Support/Casting.h"
#include "llvm/Support/raw_ostream.h"
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/Attributes.h" // from @llvm-project
#include "mlir/IR/Block.h" // from @llvm-project
#include "mlir/IR/Builders.h" // from @llvm-project
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "mlir/IR/OpDefinition.h" // from @llvm-project
#include "mlir/IR/Value.h" // from @llvm-project
#include "mlir/IR/Visitors.h" // from @llvm-project
#include "mlir/Pass/Pass.h" // from @llvm-project
#include "mlir/Pass/PassRegistry.h" // from @llvm-project
#include "mlir/Support/LLVM.h" // from @llvm-project
#include "mlir/Transforms/RegionUtils.h" // from @llvm-project
#include "tensorflow/compiler/mlir/lite/experimental/common/outline_operations.h"
#include "tensorflow/compiler/mlir/lite/experimental/tac/common/subgraph.h"
#include "tensorflow/compiler/mlir/lite/experimental/tac/common/targets.h"
#include "tensorflow/compiler/mlir/lite/experimental/tac/common/utils.h"
#include "tensorflow/compiler/mlir/lite/experimental/tac/transforms/passes.h"
#include "tensorflow/compiler/mlir/lite/ir/tfl_ops.h"
#include "tensorflow/compiler/mlir/lite/utils/utils.h"
#include "tensorflow/compiler/mlir/tensorflow/analysis/side_effect_analysis.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/cluster_util.h"
namespace mlir {
namespace TFL {
namespace tac {
namespace {
constexpr StringRef kCpuDeviceName = "CPU";
using ::mlir::TFL::common::OpsAdded;
using ::mlir::TFL::common::Subgraph;
class RaiseTargetSubgraphsPass
: public PassWrapper<RaiseTargetSubgraphsPass, OperationPass<ModuleOp>> {
public:
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(RaiseTargetSubgraphsPass)
RaiseTargetSubgraphsPass() = default;
RaiseTargetSubgraphsPass(const RaiseTargetSubgraphsPass& other) {
this->skip_raise_cpu_ops_ = other.skip_raise_cpu_ops_;
this->ignore_inference_type_ = other.ignore_inference_type_;
}
explicit RaiseTargetSubgraphsPass(bool skip_raise_cpu_ops,
bool ignore_inference_type) {
skip_raise_cpu_ops_ = skip_raise_cpu_ops;
ignore_inference_type_ = ignore_inference_type;
}
private:
Option<bool> skip_raise_cpu_ops_{
*this, "skip-raise-cpu-ops",
llvm::cl::desc("Whether to cluster and raise CPU ops."),
llvm::cl::init(false)};
Option<bool> ignore_inference_type_{
*this, "ignore-inference-type",
llvm::cl::desc("Whether to ignore the inference type in clustering."),
llvm::cl::init(false)};
llvm::StringRef getArgument() const final {
return "tfl-raise-target-subgraphs";
}
llvm::StringRef getDescription() const final {
return "This pass will merge those have target-annotated TFL IRs together "
"& raise them as a function.";
}
void runOnOperation() override;
bool RaiseTargetSubgraphsForBlock(
Block& block, OpBuilder& builder, ModuleOp module, bool skip_cpu,
int& func_count, const TF::SideEffectAnalysis::Info& side_effect_info);
};
// Returns the value of the string attribute `attr_name` for an added function
// op, by walking through the ops in the function and collating all the
// attribute values. If the attribute values are not all the same, returns
// std::nullopt.
std::optional<StringAttr> GetConsolidatedStringAttr(func::FuncOp func,
StringRef attr_name) {
StringAttr attr_value;
WalkResult result = func.walk([&](Operation* op) {
if (op->hasAttr(attr_name)) {
StringAttr current_attr = mlir::cast<StringAttr>(op->getAttr(attr_name));
if (attr_value && attr_value != current_attr) {
// If the attribute value is not null, and it is different
// from the current op's attribute value, then it is an
// error.
return WalkResult::interrupt();
}
attr_value = current_attr;
}
return WalkResult::advance();
});
if (result.wasInterrupted()) {
return std::nullopt;
}
return attr_value;
}
// After raising ops and adding the Func & Call op, call this function
// to set attributes specific to this pass.
bool AddAttrs(OpsAdded& ops_added, OpBuilder& builder, int func_count) {
func::FuncOp& added_func_op = ops_added.func_op;
func::CallOp& added_call_op = ops_added.call_op;
StringAttr interface_name =
builder.getStringAttr(absl::StrCat("func_", func_count));
added_func_op->setAttr(kInterfaceNameAttr, interface_name);
added_call_op->setAttr(kInterfaceNameAttr, interface_name);
auto custom_options_fingerprint =
GetConsolidatedStringAttr(added_func_op, kCustomOptionsFingerprint);
if (!custom_options_fingerprint.has_value()) {
return false;
}
if (custom_options_fingerprint.value() != nullptr) {
added_func_op->setAttr(kCustomOptionsFingerprint,
custom_options_fingerprint.value());
}
auto delegate_compiler_version =
GetConsolidatedStringAttr(added_func_op, kDelegateCompilerVersion);
if (!delegate_compiler_version.has_value()) {
return false;
}
if (delegate_compiler_version.value() != nullptr) {
added_func_op->setAttr(kDelegateCompilerVersion,
delegate_compiler_version.value());
}
StringAttr device = mlir::cast<StringAttr>(
added_func_op->getRegion(0).getBlocks().front().front().getAttr(kDevice));
StringAttr inference_type = mlir::cast<StringAttr>(
added_func_op->getRegion(0).getBlocks().front().front().getAttr(
kInferenceType));
added_call_op->setAttr(kDevice, device);
added_call_op->setAttr(kInferenceType, inference_type);
added_func_op->setAttr(kDevice, device);
added_func_op->setAttr(kInferenceType, inference_type);
std::string function_name = absl::StrCat(interface_name.getValue().str(), "_",
device.getValue().str(), "_",
inference_type.getValue().str());
added_func_op.setName(builder.getStringAttr(function_name));
added_call_op.setCallee(builder.getStringAttr(function_name));
return true;
}
// Raises partitioned sequential `Operations` from a block to a new function
// definition. `Operations` are partitioned into classes from the cartesian
// product of possible devices and inference datatypes. For example, we might
// raise a chunk of sequential operations from a block all having attributes
// `{ tac.device = "GPU", tac.inference_type = "FLOAT"}` to a function
// with the matching attributes. Assumed is that device type "CPU"
// is the only device that is allowed to call other devices. I.e. ancestors of a
// "CPU" `Operation` may only `Operations` without a device or other "CPU"
// `Operations`. Implied is that "CPU" ops may contain subgraphs of different
// device types which also need to be raised. The `side_effect_info` is used in
// the cluster algorithm for ops with side effect.
bool RaiseTargetSubgraphsPass::RaiseTargetSubgraphsForBlock(
Block& block, OpBuilder& builder, ModuleOp module, bool skip_cpu,
int& func_count, const TF::SideEffectAnalysis::Info& side_effect_info) {
llvm::SetVector<Operation*> partition_ops;
auto device_is = [&](InferenceDeviceType t, llvm::StringRef device) -> bool {
return t.hardware == device;
};
auto op_has_device = [&](Operation& op, InferenceDeviceType& device) -> bool {
std::optional<InferenceDeviceType> op_device =
GetInferenceDeviceTypeForOp(&op);
if (!op_device.has_value()) return false;
device = op_device.value();
return true;
};
auto op_device_is = [&](Operation& op, llvm::StringRef device) -> bool {
InferenceDeviceType device_type;
if (!op_has_device(op, device_type)) return false;
return device_is(device_type, device);
};
// Given a list of `Operation`s to partition, raise them to a new
// function. If the partitons is of type "CPU" then it may contain
// other device subgraphs that need to be raised. We recur on
// any nested blocks of "CPU" ops and skip raising "CPU" ops for the
// remainder of that recursive call.
auto extract = [&](const llvm::SetVector<Operation*>& partition_ops) -> bool {
if (partition_ops.empty()) return true;
InferenceDeviceType device =
GetInferenceDeviceTypeForOp(partition_ops.front()).value();
Subgraph old_subgraph(partition_ops, ++func_count);
OpsAdded ops_added;
ExtractSubgraphToFunc(old_subgraph, builder, module, ops_added);
if (!AddAttrs(ops_added, builder, func_count)) {
return false;
}
// Ops in "CPU" subgraphs may nested regions with other device subgraphs.
// We recur into these nested blocks to raise those as well. We don't raise
// "CPU" ops who are themselves nested within a "CPU" op, so set
// `skip_cpu` to true.
if (device_is(device, kCpuDeviceName)) {
for (auto& block : ops_added.func_op->getRegion(0).getBlocks())
for (auto& op : block) {
auto op_device = GetInferenceDeviceTypeForOp(&op);
if (op_device_is(op, kCpuDeviceName)) {
// The recently raised func is device type cpu & `op` is a "CPU".
// Recursivley call again to raise any non-"CPU" subgraphs contained
// within nested region of `op`.
for (auto& region : op.getRegions())
for (auto& block : region.getBlocks())
if (!RaiseTargetSubgraphsForBlock(block, builder, module,
/*skip_cpu=*/true, func_count,
side_effect_info)) {
return false;
}
}
}
}
return true;
};
auto get_inference_device_type_string = [&](Operation* op) {
auto device_type = GetInferenceDeviceTypeForOp(op);
if (!device_type.has_value()) {
return std::string("");
}
std::string concat_inference_device_type_string =
ignore_inference_type_
? device_type.value().hardware
: absl::StrCat(
device_type.value().hardware, "_",
GetInferenceString(device_type.value().inference_type));
// Append the custom options fingerprint to the inference device type
// string, to compile functions with different fingerprints separately.
auto custom_options_fingerprint = op->getAttr(kCustomOptionsFingerprint);
if (custom_options_fingerprint != nullptr) {
StringAttr stack_str = mlir::cast<StringAttr>(custom_options_fingerprint);
absl::StrAppend(&concat_inference_device_type_string, "_",
stack_str.getValue().str());
}
return concat_inference_device_type_string;
};
auto op_can_be_ignored = [&](Operation* op) {
auto device_type = GetInferenceDeviceTypeForOp(op);
return !device_type.has_value() ||
(skip_cpu && device_is(device_type.value(), kCpuDeviceName));
};
const llvm::StringMap<SmallVector<TF::Cluster>>& all_clusters =
TF::BuildAllClusters(block, side_effect_info,
get_inference_device_type_string, op_can_be_ignored);
for (const auto& [device, clusters] : all_clusters) {
for (const TF::Cluster& cluster : clusters) {
extract(cluster.ops);
}
}
if (skip_cpu) {
for (auto& op : block) {
auto op_device = GetInferenceDeviceTypeForOp(&op);
if (op_device_is(op, kCpuDeviceName)) {
// The recently raised func is device type cpu & `op` is a "CPU".
// Recursivley call again to raise any non-"CPU" subgraphs contained
// within nested region of `op`.
for (auto& region : op.getRegions())
for (auto& block : region.getBlocks())
if (!RaiseTargetSubgraphsForBlock(block, builder, module,
/*skip_cpu=*/true, func_count,
side_effect_info)) {
return false;
}
}
}
}
return true;
}
void RaiseTargetSubgraphsPass::runOnOperation() {
ModuleOp module = getOperation();
auto& side_effect_analysis = getAnalysis<TF::SideEffectAnalysis>();
SmallVector<func::FuncOp> funcs(module.getOps<func::FuncOp>());
int func_count = -1;
for (auto func : funcs) {
const auto& info = side_effect_analysis.GetAnalysisForFunc(func);
for (auto& block : func) {
OpBuilder builder = OpBuilder::atBlockBegin(&block);
if (!RaiseTargetSubgraphsForBlock(block, builder, module,
/*skip_cpu=*/skip_raise_cpu_ops_,
func_count, info)) {
module.emitError("Raising target subgraph for block failed.");
signalPassFailure();
}
}
}
}
} // namespace
std::unique_ptr<OperationPass<ModuleOp>> CreateRaiseTargetSubgraphsPass(
bool skip_raise_cpu_ops, bool ignore_inference_type) {
return std::make_unique<RaiseTargetSubgraphsPass>(skip_raise_cpu_ops,
ignore_inference_type);
}
static PassRegistration<RaiseTargetSubgraphsPass> pass;
} // namespace tac
} // namespace TFL
} // namespace mlir
tensorflow/compiler/mlir/lite/experimental/tac/transforms/tac_filter.cc
+343
View File
@@ -0,0 +1,343 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include <algorithm>
#include <functional>
#include <memory>
#include <string>
#include "google/protobuf/text_format.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/StringRef.h"
#include "llvm/ADT/Twine.h"
#include "llvm/Support/Casting.h"
#include "llvm/Support/CommandLine.h"
#include "llvm/Support/FormatVariadic.h"
#include "llvm/Support/Regex.h"
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/Builders.h" // from @llvm-project
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/IR/Diagnostics.h" // from @llvm-project
#include "mlir/IR/Location.h" // from @llvm-project
#include "mlir/Interfaces/CallInterfaces.h" // from @llvm-project
#include "mlir/Pass/Pass.h" // from @llvm-project
#include "mlir/Pass/PassRegistry.h" // from @llvm-project
#include "mlir/Support/LLVM.h" // from @llvm-project
#include "tensorflow/compiler/mlir/lite/experimental/tac/common/targets.h"
#include "tensorflow/compiler/mlir/lite/experimental/tac/common/utils.h"
#include "tensorflow/compiler/mlir/lite/experimental/tac/tac_filter.pb.h"
#include "tensorflow/compiler/mlir/lite/experimental/tac/transforms/passes.h"
#include "tensorflow/compiler/mlir/lite/utils/utils.h"
namespace mlir {
namespace TFL {
namespace tac {
namespace {
using ::third_party::tensorflow::compiler::mlir::lite::experimental::tac::
FunctionFilter;
using ::third_party::tensorflow::compiler::mlir::lite::experimental::tac::
OpFilter;
using ::third_party::tensorflow::compiler::mlir::lite::experimental::tac::
TacFilter;
using ::third_party::tensorflow::compiler::mlir::lite::experimental::tac::
TacFilters;
class TacFilterPass
: public PassWrapper<TacFilterPass, OperationPass<ModuleOp>> {
public:
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TacFilterPass)
TacFilterPass() = default;
TacFilterPass(const TacFilterPass& other) {
this->tac_filters_ = other.tac_filters_;
}
explicit TacFilterPass(
TacFilters* tac_filters,
std::function<void(Operation*, const google::protobuf::Any&)>
custom_options_callback) {
tac_filters_ = tac_filters;
custom_options_callback_ = custom_options_callback;
}
private:
TacFilters* tac_filters_ = nullptr;
llvm::StringRef getArgument() const final { return "tfl-tac-filter"; }
llvm::StringRef getDescription() const final {
return "This pass marks the ops to skip target annotation by inserting "
"`tac.skip_target_annotation` attribute to them based on user "
"provided config.";
}
Option<bool> use_test_setting_{
*this, "use-test-setting",
llvm::cl::desc(
"Whether to use the test config for the tac filter protobuf."),
llvm::cl::init(false)};
std::function<void(Operation*, const google::protobuf::Any&)>
custom_options_callback_;
void runOnOperation() override;
};
void ApplyFunctionTacFilter(func::FuncOp func,
FunctionFilter::FunctionFilterType type,
OpBuilder& builder) {
for (Operation& op : func.front()) {
if (type == FunctionFilter::SKIP_TARGET_ANNOTATION) {
op.setAttr(kSkipTargetAnnotation, builder.getUnitAttr());
} else if (type == FunctionFilter::INCLUDE_TARGET_ANNOTATION) {
op.removeAttr(kSkipTargetAnnotation);
}
}
}
void ApplyTacFilter(
ModuleOp module, const TacFilter& tac_filter,
SmallVector<Operation*>& filtered_ops, OpBuilder& builder,
std::function<void(Operation*, const google::protobuf::Any&)>
custom_options_callback) {
if (tac_filter.has_function_filter()) {
llvm::Regex func_regex(
tac_filter.function_filter().function_name_pattern());
for (auto func : module.getOps<func::FuncOp>()) {
if (!func_regex.match(func.getName())) {
continue;
}
ApplyFunctionTacFilter(func, tac_filter.function_filter().filter_type(),
builder);
filtered_ops.push_back(func);
}
return;
}
auto should_filter_op = [](mlir::Operation* op) {
return IsNonConstOp(op) && !IsTerminatorOp(op) &&
!llvm::isa<func::ReturnOp, func::FuncOp, CallOpInterface>(op);
};
auto map_op_to_cpu = [&](mlir::Operation* op) {
if (!should_filter_op(op)) {
return;
}
op->setAttr(kSkipTargetAnnotation, builder.getUnitAttr());
filtered_ops.push_back(op);
};
auto map_op_to_custom_device = [&](mlir::Operation* op) {
if (!should_filter_op(op)) {
return;
}
if (!tac_filter.op_filter().has_custom_options()) {
return;
}
const google::protobuf::Any& custom_options =
tac_filter.op_filter().custom_options();
custom_options_callback(op, custom_options);
};
llvm::Regex op_regex(tac_filter.op_filter().op_name_pattern());
OpFilter::MatchType match_type = tac_filter.op_filter().match_type();
OpFilter::DeviceType device_type = tac_filter.op_filter().device_type();
module.walk([&](Operation* op) {
NameLoc loc;
if (auto name_loc = mlir::dyn_cast<NameLoc>(op->getLoc())) {
loc = name_loc;
} else if (auto fused_loc = mlir::dyn_cast<FusedLoc>(op->getLoc())) {
loc = dyn_cast<NameLoc>(fused_loc.getLocations().front());
}
if (!loc) {
return;
}
// There can be two kinds of `match_type`:
// 1. MATCH: the op name matches the pattern.
// 2. INVERT_MATCH: the op name does not match the pattern.
//
// On each of the above, the op filter can specify if it should be run on
// the CPU or on a custom device, with the `device_type` field.
// Running on CPU and the match type MATCH is the default if not specified.
//
// The code below maps an op to the appropriate device based on the above
// fields.
if (op_regex.match(loc.getName())) {
switch (match_type) {
case OpFilter::MATCH:
if (device_type == OpFilter::CPU) {
map_op_to_cpu(op);
return;
}
map_op_to_custom_device(op);
break;
default:
break;
}
} else {
switch (match_type) {
case OpFilter::INVERT_MATCH:
if (device_type == OpFilter::CPU) {
map_op_to_cpu(op);
return;
}
map_op_to_custom_device(op);
break;
default:
break;
}
}
});
}
// A custom string for tac filter.
std::string TacFilterToString(const TacFilter& tac_filter) {
std::string tac_filter_type_str;
std::string tac_filter_name_pattern_str;
if (tac_filter.has_function_filter()) {
tac_filter_type_str = (llvm::Twine("function filter ") +
FunctionFilter::FunctionFilterType_Name(
tac_filter.function_filter().filter_type()))
.str();
tac_filter_name_pattern_str =
tac_filter.function_filter().function_name_pattern();
} else {
tac_filter_type_str = "op filter";
tac_filter_name_pattern_str = tac_filter.op_filter().op_name_pattern();
}
return (llvm::Twine("filter type: ") + tac_filter_type_str +
", filter_pattern: \"" + tac_filter_name_pattern_str + "\"")
.str();
}
void PrintTacFilterResult(Location module_loc, const TacFilter& tac_filter,
int count,
const SmallVector<Operation*>& filtered_ops) {
emitRemark(module_loc) << llvm::formatv("Tac filter ({0}): {1}", count,
TacFilterToString(tac_filter));
if (filtered_ops.empty()) {
emitRemark(module_loc) << llvm::formatv(
"Tac filter ({0}) specified but not applied to any op", count);
return;
}
if (tac_filter.has_function_filter()) {
for (Operation* op : filtered_ops) {
auto func = cast<func::FuncOp>(op);
func.emitRemark() << llvm::formatv("filtered by tac filter ({0})", count);
}
return;
}
DenseMap<func::FuncOp, SmallVector<Operation*>> func_to_filtered_ops_map;
for (Operation* op : filtered_ops) {
auto func = op->getParentOfType<func::FuncOp>();
func_to_filtered_ops_map[func].push_back(op);
}
for (auto& [func, ops] : func_to_filtered_ops_map) {
std::string interleaved_op_name;
llvm::raw_string_ostream os(interleaved_op_name);
llvm::interleaveComma(
ops, os, [&](Operation* op) { os << "\"" << op->getName() << "\""; });
os.flush();
func.emitRemark() << llvm::formatv(
"all ops filtered by tac filter ({0}): {1}", count,
interleaved_op_name);
}
}
void TacFilterPass::runOnOperation() {
TacFilters test_tac_filters;
if (use_test_setting_) {
// Sets up the test config used in the mlir LIT test.
google::protobuf::TextFormat::ParseFromString(R"(
tac_filters {
function_filter {
function_name_pattern: "^testFunction"
}
}
tac_filters {
function_filter {
function_name_pattern: "testFunctionInclude"
filter_type: INCLUDE_TARGET_ANNOTATION
}
}
tac_filters {
op_filter {
op_name_pattern: "^test_op"
}
}
)",
&test_tac_filters);
tac_filters_ = &test_tac_filters;
}
if (!tac_filters_) {
return;
}
ModuleOp module = getOperation();
OpBuilder builder(module);
std::sort(tac_filters_->mutable_tac_filters()->pointer_begin(),
tac_filters_->mutable_tac_filters()->pointer_end(),
[](const TacFilter* a, const TacFilter* b) {
const bool a_is_function_filter = a->has_function_filter();
const bool b_is_function_filter = b->has_function_filter();
if (a_is_function_filter != b_is_function_filter) {
// Function filter is applied before op filter.
return a_is_function_filter > b_is_function_filter;
}
if (!a_is_function_filter && !b_is_function_filter) {
// The order of 2 op filters don't matter.
return false;
}
const bool a_is_function_exclude =
(a->function_filter().filter_type() ==
FunctionFilter::SKIP_TARGET_ANNOTATION);
const bool b_is_function_exclude =
(b->function_filter().filter_type() ==
FunctionFilter::SKIP_TARGET_ANNOTATION);
// Function exclude filter is applied before function include
// filter.
return a_is_function_exclude > b_is_function_exclude;
});
for (const auto& tac_filter : llvm::enumerate(tac_filters_->tac_filters())) {
SmallVector<Operation*> filtered_ops;
ApplyTacFilter(module, tac_filter.value(), filtered_ops, builder,
custom_options_callback_);
PrintTacFilterResult(module.getLoc(), tac_filter.value(),
tac_filter.index(), filtered_ops);
}
}
} // namespace
std::unique_ptr<OperationPass<ModuleOp>> CreateTacFilterPass(
TacFilters* tac_filters,
std::function<void(Operation*, const google::protobuf::Any&)>
custom_options_callback) {
return std::make_unique<TacFilterPass>(tac_filters, custom_options_callback);
}
static PassRegistration<TacFilterPass> pass;
} // namespace tac
} // namespace TFL
} // namespace mlir
tensorflow/compiler/mlir/lite/experimental/tac/transforms/tac_pass.h
+90
View File
@@ -0,0 +1,90 @@
/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_TAC_TRANSFORMS_TAC_PASS_H_
#define TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_TAC_TRANSFORMS_TAC_PASS_H_
#include <memory>
#include <string>
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/Pass/Pass.h" // from @llvm-project
#include "tensorflow/compiler/mlir/lite/experimental/tac/hardwares/target_hardware.h"
#include "tensorflow/compiler/mlir/lite/experimental/tac/tac_module.h"
namespace mlir {
namespace TFL {
namespace tac {
// An OperationPass<> with access to the TAC module instance that the
// pass is running part of.
// See OperationPass<> comments for all details/restrictions of OperationPass.
//
// When adding new Pass to TAC, users should use this class as the base class
// as it provides access to the TAC module.
template <typename T>
class TacPass : public OperationPass<T> {
public:
using OperationPass<T>::OperationPass;
explicit TacPass(const TacModule* module)
: OperationPass<T>::OperationPass(mlir::TypeID::get<T>()),
module_(module) {}
~TacPass() override = default;
const TargetHardware* GetTargetHardware(
const std::string& hardware_name) const {
return module_ != nullptr
? module_->GetTargetHardware(hardware_name)
: mlir::TFL::tac::GetTargetHardware(hardware_name);
}
protected:
const TacModule* module_ = nullptr; // Not owned.
};
// A FunctionPass but with access to TAC module.
// See FunctionPass comments for all details/restrictions of FunctionPass.
//
// When adding new Pass to TAC, users should use this class as the base class
// as it provides access to the TAC module.
template <typename T>
class TacFunctionPass : public TacPass<func::FuncOp> {
public:
using TacPass<func::FuncOp>::TacPass;
~TacFunctionPass() override = default;
mlir::func::FuncOp getFunction() { return getOperation(); }
virtual void runOnFunction() = 0;
void runOnOperation() final {
if (!getFunction().isExternal()) runOnFunction();
}
protected:
// Returns the derived pass name.
StringRef getName() const override { return llvm::getTypeName<T>(); }
// A clone method to create a copy of this pass.
std::unique_ptr<Pass> clonePass() const override {
return std::make_unique<T>(*static_cast<const T*>(this));
}
};
} // namespace tac
} // namespace TFL
} // namespace mlir
#endif // TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_TAC_TRANSFORMS_TAC_PASS_H_
tensorflow/compiler/mlir/lite/experimental/tac/transforms/target_annotation.cc
+166
View File
@@ -0,0 +1,166 @@
/* 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 <memory>
#include <string>
#include "llvm/ADT/ArrayRef.h"
#include "llvm/ADT/DenseSet.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/Support/Casting.h"
#include "llvm/Support/CommandLine.h"
#include "mlir/IR/Builders.h" // from @llvm-project
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "mlir/Pass/Pass.h" // from @llvm-project
#include "mlir/Pass/PassRegistry.h" // from @llvm-project
#include "mlir/Support/LLVM.h" // from @llvm-project
#include "tensorflow/compiler/mlir/lite/experimental/tac/common/targets.h"
#include "tensorflow/compiler/mlir/lite/experimental/tac/common/utils.h"
#include "tensorflow/compiler/mlir/lite/experimental/tac/hardwares/target_hardware.h"
#include "tensorflow/compiler/mlir/lite/experimental/tac/transforms/device_transform.h"
#include "tensorflow/compiler/mlir/lite/experimental/tac/transforms/passes.h"
#include "tensorflow/compiler/mlir/lite/experimental/tac/transforms/tac_pass.h"
#include "tensorflow/compiler/mlir/lite/ir/tfl_ops.h"
namespace mlir {
namespace TFL {
namespace tac {
namespace {
class TargetAnnotationPass : public TacFunctionPass<TargetAnnotationPass> {
public:
llvm::StringRef getArgument() const final { return "tfl-target-annotation"; }
llvm::StringRef getDescription() const final {
return "Add user specified target annotations to the TFL operations given "
"operation capabilities, will default to CPU.";
}
// using TacFunctionPass::TacFunctionPass;
TargetAnnotationPass() : TacFunctionPass(nullptr) {}
TargetAnnotationPass(const TargetAnnotationPass& copy)
: TacFunctionPass(copy.module_) {}
explicit TargetAnnotationPass(llvm::ArrayRef<std::string> device_specs)
: TacFunctionPass(nullptr) {
device_specs_flag_ = device_specs;
}
explicit TargetAnnotationPass(const TacModule* module)
: TacFunctionPass(module) {}
private:
void runOnFunction() override;
void SetTargetAnnotation(Operation* op,
llvm::ArrayRef<std::string> device_specs,
OpBuilder* builder);
ListOption<std::string> device_specs_flag_{
*this, "device-specs",
llvm::cl::desc(
"comma separated list of device specs, like CPU, GPU, Hexagon."),
llvm::cl::ZeroOrMore};
void getDependentDialects(mlir::DialectRegistry& registry) const override {
if (!module_) {
for (const auto& device : device_specs_flag_) {
auto* hardware = this->GetTargetHardware(device);
if (hardware == nullptr) continue;
hardware->GetDependentDialects(registry);
}
}
}
};
void SetAnnotation(Operation* op, std::string attribute, std::string annotation,
OpBuilder* builder) {
// TODO(karimnosseir): Maybe set device capabilities to allow us to have
// more flexbility when raise the subgraphs.
auto default_target = builder->getStringAttr(annotation);
op->setAttr(attribute, default_target);
}
void TargetAnnotationPass::SetTargetAnnotation(
Operation* op, llvm::ArrayRef<std::string> device_specs,
OpBuilder* builder) {
if (op->hasAttr(kSkipTargetAnnotation)) {
return;
}
const InferenceType inference_type = GetInferenceType(op);
const std::string inference_type_str = GetInferenceString(inference_type);
SetAnnotation(op, kInferenceType, inference_type_str, builder);
bool device_is_set = false;
// TODO(b/177376459): Remove the usage of device_specs.
// TODO(b/177376459): Update if needed to make testing easy.
if (!module_) {
for (const auto& device : device_specs) {
auto* hardware = this->GetTargetHardware(device);
if (hardware == nullptr) continue;
if (hardware->IsOpSupported(op)) {
SetAnnotation(op, kDevice, device, builder);
device_is_set = true;
break;
}
}
} else {
for (const auto* hardware : module_->GetAvailableHardwares()) {
if (hardware == nullptr) continue;
if (hardware->IsOpSupported(op)) {
SetAnnotation(op, kDevice, GetHardwareName(hardware), builder);
device_is_set = true;
break;
}
}
}
// default to CPU
if (!device_is_set) {
if (IsNonConstOp(op) && !IsTerminatorOp(op) &&
!llvm::isa<func::ReturnOp, func::FuncOp, CallableOpInterface>(op)) {
SetAnnotation(op, kDevice, "CPU", builder);
device_is_set = true;
}
}
if (!device_is_set) {
op->emitError("cannot set target device for this ops");
}
}
void TargetAnnotationPass::runOnFunction() {
auto func = getFunction();
OpBuilder builder(func);
func.walk([&](Operation* op) {
// We only care about TFL dialect.
if (IsNonConstOp(op) && !IsTerminatorOp(op) &&
!llvm::isa<func::ReturnOp, func::FuncOp, CallOpInterface>(op)) {
SetTargetAnnotation(op, device_specs_flag_, &builder);
}
});
}
} // namespace
std::unique_ptr<OperationPass<func::FuncOp>> CreateTargetAnnotationPass(
llvm::ArrayRef<std::string> device_specs) {
return std::make_unique<TargetAnnotationPass>(device_specs);
}
std::unique_ptr<OperationPass<func::FuncOp>> CreateTargetAnnotationPass(
const TacModule* module) {
return std::make_unique<TargetAnnotationPass>(module);
}
static PassRegistration<TargetAnnotationPass> pass;
} // namespace tac
} // namespace TFL
} // namespace mlir
tensorflow/compiler/mlir/lite/experimental/tac/transforms/transform_patterns.td
+32
View File
@@ -0,0 +1,32 @@
/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
include "mlir/IR/OpBase.td"
include "mlir/IR/PatternBase.td"
include "mlir/Dialect/Arith/IR/ArithOps.td"
include "mlir/Dialect/Func/IR/FuncOps.td"
include "tensorflow/compiler/mlir/lite/ir/tfl_ops.td"
include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.td"
// Basically x - y => x + y * -1
def SubToAdd : Pat<(TFL_SubOp $lhs, $rhs, $act),
(TFL_AddOp $lhs, (TFL_MulOp $rhs,
(Arith_ConstantOp ConstantAttr<RankedF32ElementsAttr<[]>,
"-1.0f">), TFL_AF_None), $act)>;
// Squash tfl.dequantize and tfl.quantize pairs.
// TODO(b/185915462): Compare the scale of input and output. This can also be
// squashed to a requantize op if the scales are different.
def FoldQuantizeDequantize: Pat<(TFL_DequantizeOp (TFL_QuantizeOp $arg, $qt)), (replaceWithValue $arg)>;
tensorflow/compiler/mlir/lite/experimental/tac/utils/BUILD
+38
View File
@@ -0,0 +1,38 @@
load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:LICENSE"],
default_visibility = [
"//visibility:public",
],
licenses = ["notice"],
)
cc_library(
name = "utils",
srcs = [
"utils.cc",
],
hdrs = [
"utils.h",
],
deps = [
"//tensorflow/compiler/mlir/lite:flatbuffer_translate_lib",
"//tensorflow/compiler/mlir/lite:tensorflow_lite",
"//tensorflow/compiler/mlir/lite/experimental/tac:common",
"//tensorflow/compiler/mlir/lite/stablehlo:prepare_hlo",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings:str_format",
"@com_google_absl//absl/strings:string_view",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:ArithDialect",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:FuncExtensions",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Parser",
"@llvm-project//mlir:Pass",
"@llvm-project//mlir:ReconcileUnrealizedCasts",
"@llvm-project//mlir:Support",
],
)
tensorflow/compiler/mlir/lite/experimental/tac/utils/utils.cc
+146
View File
@@ -0,0 +1,146 @@
/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/compiler/mlir/lite/experimental/tac/utils/utils.h"
#include <memory>
#include <optional>
#include <string>
#include <utility>
#include <vector>
#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "absl/strings/str_format.h"
#include "absl/strings/string_view.h"
#include "llvm/Support/MemoryBuffer.h"
#include "llvm/Support/SMLoc.h"
#include "llvm/Support/SourceMgr.h"
#include "llvm/Support/ToolOutputFile.h"
#include "llvm/Support/raw_ostream.h"
#include "mlir/Conversion/ReconcileUnrealizedCasts/ReconcileUnrealizedCasts.h" // from @llvm-project
#include "mlir/Dialect/Arith/IR/Arith.h" // from @llvm-project
#include "mlir/Dialect/Func/Extensions/AllExtensions.h" // from @llvm-project
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/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/OwningOpRef.h" // from @llvm-project
#include "mlir/Parser/Parser.h" // from @llvm-project
#include "mlir/Pass/PassManager.h" // from @llvm-project
#include "mlir/Support/FileUtilities.h" // from @llvm-project
#include "mlir/Support/LLVM.h" // from @llvm-project
#include "tensorflow/compiler/mlir/lite/experimental/tac/common/targets.h"
#include "tensorflow/compiler/mlir/lite/flatbuffer_export.h"
#include "tensorflow/compiler/mlir/lite/flatbuffer_import.h"
#include "tensorflow/compiler/mlir/lite/ir/tfl_ops.h"
#include "tensorflow/compiler/mlir/lite/stablehlo/transforms/stablehlo_passes.h"
namespace mlir {
namespace TFL {
namespace tac {
absl::StatusOr<mlir::OwningOpRef<mlir::ModuleOp>> ImportFlatbufferOrMlir(
const std::string& input_filename, bool input_mlir,
bool experimental_prune_unreachable_nodes_unconditionally,
llvm::SourceMgr* source_mgr, mlir::MLIRContext* context) {
std::string error;
std::unique_ptr<llvm::MemoryBuffer> buffer =
mlir::openInputFile(input_filename, &error);
if (buffer == nullptr) {
return absl::InvalidArgumentError(absl::StrFormat(
"Cannot open input file: %s. %s", input_filename, error));
}
if (input_mlir) {
mlir::DialectRegistry registry;
registry.insert<mlir::TFL::TensorFlowLiteDialect, mlir::arith::ArithDialect,
mlir::func::FuncDialect>();
mlir::func::registerAllExtensions(registry);
context->appendDialectRegistry(registry);
source_mgr->AddNewSourceBuffer(std::move(buffer), llvm::SMLoc());
return mlir::OwningOpRef<mlir::ModuleOp>(
mlir::parseSourceFile<mlir::ModuleOp>(*source_mgr, context));
}
mlir::Location loc =
mlir::FileLineColLoc::get(context, input_filename, /*line=*/0,
/*column=*/0);
std::vector<std::string> inputs;
std::vector<std::string> outputs;
return tflite::FlatBufferToMlir(
absl::string_view(buffer->getBufferStart(), buffer->getBufferSize()),
context, loc, /*use_external_constant=*/false, inputs, outputs,
experimental_prune_unreachable_nodes_unconditionally);
}
absl::Status ExportFlatbufferOrMlir(
const std::string& output_filename, bool output_mlir, mlir::ModuleOp module,
bool enable_select_tf_ops, std::optional<int> custom_option_alignment) {
std::string error_msg;
auto output = mlir::openOutputFile(output_filename, &error_msg);
if (output == nullptr) {
llvm::errs() << error_msg << '\n';
return absl::InvalidArgumentError("cannot open output file.");
}
std::string result;
if (output_mlir) {
llvm::raw_string_ostream os(result);
module.print(os);
os.flush();
} else {
// This extra attribute is added by TAC pass. We need to remove it before
// converting to VHLO.
module.walk([&](mlir::Operation* op) {
if (op->hasAttr(mlir::TFL::tac::kSkipTargetAnnotation)) {
op->removeAttr(mlir::TFL::tac::kSkipTargetAnnotation);
}
});
// Converts stablehlo to vhlo so that flatbuffer export can handle it.
auto pass_manager =
std::make_unique<mlir::PassManager>(module.getContext());
pass_manager->addPass(mlir::odml::createLegalizeStablehloToVhloPass());
pass_manager->addPass(mlir::createReconcileUnrealizedCastsPass());
if (failed(pass_manager->run(module))) {
return absl::UnknownError("Failed to legalize stablehlo to vhlo.");
}
tflite::FlatbufferExportOptions options;
options.converter_flags.set_force_select_tf_ops(false);
options.converter_flags.set_allow_custom_ops(true);
if (enable_select_tf_ops) {
options.converter_flags.set_enable_select_tf_ops(true);
options.converter_flags.set_allow_all_select_tf_ops(true);
} else {
options.converter_flags.set_enable_select_tf_ops(false);
}
if (custom_option_alignment.has_value()) {
options.custom_option_alignment = *custom_option_alignment;
}
if (!tflite::MlirToFlatBufferTranslateFunction(
module, options, &result, /*serialize_stablehlo_ops=*/true)) {
return absl::UnknownError("Failed to export tflite file.");
}
}
output->os() << result;
output->keep();
return absl::OkStatus();
}
} // namespace tac
} // namespace TFL
} // namespace mlir
tensorflow/compiler/mlir/lite/experimental/tac/utils/utils.h
+50
View File
@@ -0,0 +1,50 @@
/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_TAC_UTILS_UTILS_H_
#define TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_TAC_UTILS_UTILS_H_
#include <optional>
#include <string>
#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "mlir/IR/OwningOpRef.h" // from @llvm-project
#include "mlir/Parser/Parser.h" // from @llvm-project
namespace mlir {
namespace TFL {
namespace tac {
// Import the file as mlir module, the input maybe flatbuffer or mlir file.
absl::StatusOr<mlir::OwningOpRef<mlir::ModuleOp>> ImportFlatbufferOrMlir(
const std::string& input_filename, bool input_mlir,
bool experimental_prune_unreachable_nodes_unconditionally,
llvm::SourceMgr* source_mgr, mlir::MLIRContext* context);
// Export the module to file, can be either mlir or flatbuffer.
absl::Status ExportFlatbufferOrMlir(
const std::string& output_filename, bool output_mlir, mlir::ModuleOp module,
bool enable_select_tf_ops,
std::optional<int> custom_option_alignment = std::nullopt);
} // namespace tac
} // namespace TFL
} // namespace mlir
#endif // TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_TAC_UTILS_UTILS_H_