experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Use tf._implements to fuse Tf.Text APIs

Browse Source 
PiperOrigin-RevId: 321068074
Change-Id: Ia43c6fe5b29f41a59461a458fb97e986602a2a63
This commit is contained in:
A. Unique TensorFlower 2020-07-13 17:33:04 -07:00 committed by TensorFlower Gardener
parent a7e6b483d3
commit 1e1bcbbf80
6 changed files with 60 additions and 32 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
tensorflow/compiler/mlir/lite/BUILD
View File

@ -273,6 +273,7 @@ cc_library(
deps = [
":tensorflow_lite",
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_attributes",
"//tensorflow/core:framework",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:IR",
@ -360,6 +361,7 @@ cc_library(
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/compiler/mlir/tensorflow:convert_tensor",
"//tensorflow/compiler/mlir/tensorflow:mangling_util",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_attributes",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_types",
"//tensorflow/compiler/mlir/tensorflow:unroll_batch_matmul_pass",
"//tensorflow/compiler/xla:status",

13
tensorflow/compiler/mlir/lite/tests/fuse-tftext.mlir
View File

@ -1,7 +1,7 @@
// RUN: tf-opt -tfl-prepare-composite-funcs-tf -tfl-fuse-tftext=true %s -split-input-file | FileCheck %s
module {
func @whitespace_tokenizer_rank1(%arg0: tensor<1x!tf.string> {tf._user_specified_name = "input"}) -> (tensor<?x!tf.string>, tensor<?xi64>) attributes {tf._input_shapes = [#tf.shape<1>], tf.api_implements = "tftext:WhitespaceTokenizer", tf.signature.is_stateful} {
func @whitespace_tokenizer_rank1(%arg0: tensor<1x!tf.string> {tf._user_specified_name = "input"}) -> (tensor<?x!tf.string>, tensor<?xi64>) attributes {sym_visibility = "private", tf._input_shapes = [#tf.shape<1>], tf._implements = #tf.func<@"tftext:WhitespaceTokenizer", {}>, tf.signature.is_stateful} {
%0 = "tf.Const"() {value = dense<[0, 1]> : tensor<2xi64>} : () -> tensor<2xi64>
%1 = "tf.Const"() {value = dense<[]> : tensor<0xi64>} : () -> tensor<0xi64>
%2 = "tf.Const"() {value = dense<true> : tensor<i1>} : () -> tensor<i1>
@ -1027,11 +1027,11 @@ module {
return %1 : tensor<i1>
}
// CHECK: func @whitespace_tokenizer_rank1(%arg0: tensor<1x!tf.string> {tf._user_specified_name = "input"}) -> (tensor<?x!tf.string>, tensor<?xi64>) attributes {tf._input_shapes = [#tf.shape<1>], tf.api_implements = "tftext:WhitespaceTokenizer", tf.signature.is_stateful} {
// CHECK: func @whitespace_tokenizer_rank1(%arg0: tensor<1x!tf.string> {tf._user_specified_name = "input"}) -> (tensor<?x!tf.string>, tensor<?xi64>) attributes {sym_visibility = "private", tf._implements = #tf.func<@"tftext:WhitespaceTokenizer", {}>, tf._input_shapes = [#tf.shape<1>], tf.signature.is_stateful} {
// CHECK: %0:2 = "tfl.custom"(%arg0) {custom_code = "tftext:WhitespaceTokenizer", custom_option = opaque<"tfl", "0x"> : tensor<0xi8>} : (tensor<1x!tf.string>) -> (tensor<?x!tf.string>, tensor<?xi64>)
// CHECK: return %0#0, %0#1 : tensor<?x!tf.string>, tensor<?xi64>
func @whitespace_tokenizer_rank2(%arg0: tensor<?x1x!tf.string> {tf._user_specified_name = "input"}) -> (tensor<?x!tf.string>, tensor<?xi64>, tensor<?xi64>) attributes {tf._input_shapes = [#tf.shape<?x1>], tf.api_implements = "tftext:WhitespaceTokenizer", tf.signature.is_stateful} {
func @whitespace_tokenizer_rank2(%arg0: tensor<?x1x!tf.string> {tf._user_specified_name = "input"}) -> (tensor<?x!tf.string>, tensor<?xi64>, tensor<?xi64>) attributes {sym_visibility = "private", tf._input_shapes = [#tf.shape<?x1>], tf._implements = #tf.func<@"tftext:WhitespaceTokenizer", {}>, tf.signature.is_stateful} {
%0 = "tf.Const"() {value = dense<[]> : tensor<0xi64>} : () -> tensor<0xi64>
%1 = "tf.Const"() {value = dense<true> : tensor<i1>} : () -> tensor<i1>
%2 = "tf.Const"() {value = dense<-1> : tensor<i32>} : () -> tensor<i32>
@ -2160,11 +2160,12 @@ module {
}
// CHECK: func @whitespace_tokenizer_rank2(%arg0: tensor<?x1x!tf.string> {tf._user_specified_name = "input"}) -> (tensor<?x!tf.string>, tensor<?xi64>, tensor<?xi64>) attributes {tf._input_shapes = [#tf.shape<?x1>], tf.api_implements = "tftext:WhitespaceTokenizer", tf.signature.is_stateful} {
// CHECK: func @whitespace_tokenizer_rank2(%arg0: tensor<?x1x!tf.string> {tf._user_specified_name = "input"}) -> (tensor<?x!tf.string>, tensor<?xi64>, tensor<?xi64>) attributes {sym_visibility = "private", tf._implements = #tf.func<@"tftext:WhitespaceTokenizer", {}>, tf._input_shapes = [#tf.shape<?x1>], tf.signature.is_stateful} {
// CHECK: %0:3 = "tfl.custom"(%arg0) {custom_code = "tftext:WhitespaceTokenizer", custom_option = opaque<"tfl", "0x"> : tensor<0xi8>} : (tensor<?x1x!tf.string>) -> (tensor<?x!tf.string>, tensor<?xi64>, tensor<?xi64>)
// CHECK: return %0#0, %0#1, %0#2 : tensor<?x!tf.string>, tensor<?xi64>, tensor<?xi64>
func @whitespace_tokenizer_rank0(%arg0: tensor<!tf.string> {tf._user_specified_name = "input"}) -> tensor<?x!tf.string> attributes {tf._input_shapes = [#tf.shape<>], tf.api_implements = "tftext:WhitespaceTokenizer", tf.signature.is_stateful} {
func @whitespace_tokenizer_rank0(%arg0: tensor<!tf.string> {tf._user_specified_name = "input"}) -> tensor<?x!tf.string> attributes {sym_visibility = "private", tf._input_shapes = [#tf.shape<>], tf._implements = #tf.func<@"tftext:WhitespaceTokenizer", {}>, tf.signature.is_stateful} {
%0 = "tf.Const"() {value = dense<[0, 1]> : tensor<2xi64>} : () -> tensor<2xi64>
%1 = "tf.Const"() {value = dense<[]> : tensor<0xi64>} : () -> tensor<0xi64>
%2 = "tf.Const"() {value = dense<true> : tensor<i1>} : () -> tensor<i1>
@ -3190,7 +3191,7 @@ module {
return %1 : tensor<i1>
}
// CHECK: func @whitespace_tokenizer_rank0(%arg0: tensor<!tf.string> {tf._user_specified_name = "input"}) -> tensor<?x!tf.string> attributes {tf._input_shapes = [#tf.shape<>], tf.api_implements = "tftext:WhitespaceTokenizer", tf.signature.is_stateful} {
// CHECK: func @whitespace_tokenizer_rank0(%arg0: tensor<!tf.string> {tf._user_specified_name = "input"}) -> tensor<?x!tf.string> attributes {sym_visibility = "private", tf._implements = #tf.func<@"tftext:WhitespaceTokenizer", {}>, tf._input_shapes = [#tf.shape<>], tf.signature.is_stateful} {
// CHECK: %0 = "tfl.custom"(%arg0) {custom_code = "tftext:WhitespaceTokenizer", custom_option = opaque<"tfl", "0x"> : tensor<0xi8>} : (tensor<!tf.string>) -> tensor<?x!tf.string>
// CHECK: return %0 : tensor<?x!tf.string>
}

44
tensorflow/compiler/mlir/lite/transforms/prepare_composite_functions_tf.cc
View File

@ -42,6 +42,7 @@ limitations under the License.
#include "tensorflow/compiler/mlir/lite/transforms/passes.h"
#include "tensorflow/compiler/mlir/lite/utils/lstm_utils.h"
#include "tensorflow/compiler/mlir/lite/utils/tftext_utils.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_attributes.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
// The cmd line flag to turn on/off Tf.Text API fusion.
@ -56,9 +57,11 @@ namespace TFL {
namespace {
constexpr char kTFAPIImplements[] = "tf.api_implements";
constexpr char kTfTextAPIPRefix[] = "tftext:";
constexpr char kTFTextAPIPrefix[] = "tftext:";
constexpr char kTfNMSPadded[] = "non_max_suppression_padded_v2";
using mlir::TF::FuncAttr;
// Abstracts the conversion of the embedded lookup composite function.
class ConvertEmbeddedLookupFunc {
public:
@ -161,7 +164,9 @@ class PrepareCompositeFunctionsPass
explicit PrepareCompositeFunctionsPass() {}
private:
// TODO(b/160915525): Consolidate FuncAttr and StringAttr into one.
void ConvertTFImplements(FuncOp func, StringAttr attr);
void ConvertTFImplementsWithAttributes(FuncOp func, FuncAttr attr);
void ConvertTFAPIImplements(FuncOp func, StringAttr attr, ModuleOp module);
void runOnOperation() override;
};
@ -204,6 +209,18 @@ void PrepareCompositeFunctionsPass::ConvertTFImplements(FuncOp func,
}
}
void PrepareCompositeFunctionsPass::ConvertTFImplementsWithAttributes(
FuncOp func, FuncAttr attr) {
auto api_name = attr.GetName().getLeafReference();
bool enable_fuse_tftext =
fuse_tftext_flag || IsTFTextRegistered(tensorflow::OpRegistry::Global());
if (api_name.startswith(kTFTextAPIPrefix) && enable_fuse_tftext) {
if (failed(ConvertTFTextAPI(func, api_name, attr))) {
return signalPassFailure();
}
}
}
LogicalResult CheckOutputConsumer(
Operation* call_op, int expected_num_outputs,
llvm::DenseSet<int> expected_consumer_indices) {
@ -256,26 +273,27 @@ void PrepareCompositeFunctionsPass::ConvertTFAPIImplements(FuncOp func,
OpBuilder builder(func.getBody());
if (failed(ConvertKerasLSTMLayer(func, &builder)))
return signalPassFailure();
} else if (fuse_tftext_flag ||
IsTfTextRegistered(tensorflow::OpRegistry::Global())) {
if (attr.getValue().startswith(kTfTextAPIPRefix)) {
if (failed(ConvertTFTextAPI(func, attr.getValue()))) {
return signalPassFailure();
}
}
}
}
void PrepareCompositeFunctionsPass::runOnOperation() {
auto module = getOperation();
for (auto func : module.getOps<FuncOp>()) {
// We have two kinds of implements:
// 1) tf._implements.
// 2) tf.api_implements.
// We have three kinds of implements:
// 1) tf._implements, with string attributes.
// 2) tf._implements, with proto attributes.
// 3) tf.api_implements.
// We need to handle them separately.
auto tf_implements_attr = func.getAttrOfType<StringAttr>(kTFImplements);
auto tf_implements_attr_str = func.getAttrOfType<StringAttr>(kTFImplements);
if (tf_implements_attr_str) {
ConvertTFImplements(func, tf_implements_attr_str);
continue;
}
auto tf_implements_attr = func.getAttrOfType<FuncAttr>(kTFImplements);
if (tf_implements_attr) {
ConvertTFImplements(func, tf_implements_attr);
ConvertTFImplementsWithAttributes(func, tf_implements_attr);
continue;
}
auto tf_api_implements_attr =

21
tensorflow/compiler/mlir/lite/utils/tftext_utils.cc
View File

@ -44,7 +44,9 @@ namespace TFL {
namespace {
constexpr char kWhitespaceTokenizer[] = "tftext:WhitespaceTokenizer";
constexpr char kTFAPIImplements[] = "tf.api_implements";
constexpr char kTFImplements[] = "tf._implements";
using mlir::TF::FuncAttr;
inline OpaqueElementsAttr emptyCustomOption(OpBuilder* builder) {
std::string content = "";
@ -121,11 +123,11 @@ LogicalResult VerifyWhitespaceTokenizer(mlir::FuncOp func) {
return success();
}
LogicalResult ConvertWhitespaceTokenizer(mlir::FuncOp func,
llvm::StringRef api) {
LogicalResult ConvertWhitespaceTokenizer(mlir::FuncOp func, llvm::StringRef api,
FuncAttr attr) {
func.eraseBody();
func.addEntryBlock();
func.setAttr(kTFAPIImplements, StringAttr::get(api, func.getContext()));
func.setAttr(kTFImplements, attr);
Value text = func.getArgument(0);
OpBuilder builder(func.getBody());
@ -137,20 +139,21 @@ LogicalResult ConvertWhitespaceTokenizer(mlir::FuncOp func,
}
} // namespace
LogicalResult ConvertTFTextAPI(mlir::FuncOp func, llvm::StringRef api) {
LogicalResult ConvertTFTextAPI(mlir::FuncOp func, llvm::StringRef api,
FuncAttr attr) {
if (api.str() == kWhitespaceTokenizer) {
if (succeeded(VerifyWhitespaceTokenizer(func))) {
return ConvertWhitespaceTokenizer(func, api);
return ConvertWhitespaceTokenizer(func, api, attr);
}
}
return failure();
}
bool IsTfTextRegistered(const tensorflow::OpRegistry* op_registery) {
const std::vector<std::string> kTfTextOps = {
bool IsTFTextRegistered(const tensorflow::OpRegistry* op_registery) {
const std::vector<std::string> kTFTextOps = {
"WhitespaceTokenizeWithOffsets",
};
for (const auto& iter : kTfTextOps) {
for (const auto& iter : kTFTextOps) {
if (op_registery->LookUp(iter)) {
return true;
}

8
tensorflow/compiler/mlir/lite/utils/tftext_utils.h
View File

@ -27,14 +27,18 @@ limitations under the License.
#include "mlir/IR/Value.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_attributes.h"
#include "tensorflow/core/framework/op.h"
namespace mlir {
namespace TFL {
LogicalResult ConvertTFTextAPI(mlir::FuncOp func, llvm::StringRef api);
// Fuse TF.Text APIs annotated by tf.function to a TFLite custom op.
LogicalResult ConvertTFTextAPI(mlir::FuncOp func, llvm::StringRef api,
mlir::TF::FuncAttr attr);
bool IsTfTextRegistered(const tensorflow::OpRegistry* op_registery);
// Check if TF.Text Tensorflow ops are registered.
bool IsTFTextRegistered(const tensorflow::OpRegistry* op_registery);
} // end namespace TFL
} // end namespace mlir

4
tensorflow/compiler/mlir/lite/utils/tftext_utils_test.cc
View File

@ -41,13 +41,13 @@ void Register(const std::string& op_name, OpRegistry* registry) {
TEST(TfTextUtilsTest, TestTfTextRegistered) {
std::unique_ptr<OpRegistry> registry(new OpRegistry);
Register("WhitespaceTokenizeWithOffsets", registry.get());
EXPECT_TRUE(IsTfTextRegistered(registry.get()));
EXPECT_TRUE(IsTFTextRegistered(registry.get()));
}
TEST(TfTextUtilsTest, TestTfTextNotRegistered) {
std::unique_ptr<OpRegistry> registry(new OpRegistry);
Register("Test", registry.get());
EXPECT_FALSE(IsTfTextRegistered(registry.get()));
EXPECT_FALSE(IsTFTextRegistered(registry.get()));
}
} // namespace TFL
} // namespace mlir