Use tf._implements to fuse Tf.Text APIs

PiperOrigin-RevId: 321068074
Change-Id: Ia43c6fe5b29f41a59461a458fb97e986602a2a63

tensorflow/compiler/mlir/lite/BUILD

@@ -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",

tensorflow/compiler/mlir/lite/tests/fuse-tftext.mlir

@@ -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>
 }

tensorflow/compiler/mlir/lite/transforms/prepare_composite_functions_tf.cc

@@ -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:
+    // We have three kinds of implements:
-    // 1) tf._implements.
+    // 1) tf._implements, with string attributes.
-    // 2) tf.api_implements.
+    // 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 =

tensorflow/compiler/mlir/lite/utils/tftext_utils.cc

@@ -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,
+LogicalResult ConvertWhitespaceTokenizer(mlir::FuncOp func, llvm::StringRef api,
-                                         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) {
+bool IsTFTextRegistered(const tensorflow::OpRegistry* op_registery) {
-  const std::vector<std::string> kTfTextOps = {
+  const std::vector<std::string> kTFTextOps = {
       "WhitespaceTokenizeWithOffsets",
   };
-  for (const auto& iter : kTfTextOps) {
+  for (const auto& iter : kTFTextOps) {
     if (op_registery->LookUp(iter)) {
       return true;
     }

tensorflow/compiler/mlir/lite/utils/tftext_utils.h

@@ -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

tensorflow/compiler/mlir/lite/utils/tftext_utils_test.cc

@@ -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