Propagate the allowed unregistered attributes to the expanded operations

Since the expansion is by two passes, the propagation of the attributes is by
the chain of "tf composite op -> tfr func call op -> inlined tfr func call op -> inlined tf ops". In case the
propagation is broken by the canonicalization of TF ops, these
canonicalizations are excluded.

PiperOrigin-RevId: 346146558
Change-Id: Iecf633b511e4c15e72394c26f313a98ff19206a3

Author: liufengdb

tensorflow/compiler/mlir/tfr/BUILD

@@ -104,7 +104,9 @@ cc_library(
         "utils/utils.h",
     ],
     deps = [
+        ":tfr",
         "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:IR",
         "@llvm-project//mlir:Support",
     ],
 )

tensorflow/compiler/mlir/tfr/integration/tfr_decompose_ctx_test.cc

@@ -82,6 +82,7 @@ tfr.func @tf__my_add_n(%values: !tfr.tensor_list,
   tfr.return %res : !tfr.tensor
 }
 
+tfr.func @tf__my_add_n_(!tfr.tensor_list<N,T>, i64 {tfr.name="N"}) -> !tfr.tensor attributes{N,T}
 tfr.func @tf__risc_add_dummy_(!tfr.tensor<T>, !tfr.tensor<T>) -> !tfr.tensor<T> attributes{T}
 )";
 

tensorflow/compiler/mlir/tfr/ir/tfr_ops.h

@@ -16,6 +16,7 @@ limitations under the License.
 #ifndef TENSORFLOW_COMPILER_MLIR_TFR_IR_TFR_OPS_H_
 #define TENSORFLOW_COMPILER_MLIR_TFR_IR_TFR_OPS_H_
 
+#include "llvm/ADT/StringSet.h"
 #include "mlir/IR/BuiltinOps.h"  // from @llvm-project
 #include "mlir/IR/Dialect.h"  // from @llvm-project
 #include "mlir/IR/DialectImplementation.h"  // from @llvm-project

tensorflow/compiler/mlir/tfr/ir/tfr_ops.td

@@ -411,6 +411,22 @@ def TFR_TFRFuncOp : TFR_Op<"func", [HasParent<"ModuleOp">,
     // Hooks for the input/output type enumeration in FunctionLike .
     unsigned getNumFuncArguments() { return getType().getNumInputs(); }
     unsigned getNumFuncResults() { return getType().getNumResults(); }
+
+    // Get the names of all defined attributes, including both derived and
+    // non-derived ones.
+    llvm::StringSet<> getDefinedAttributeNames() {
+      llvm::StringSet<> all_attrs;
+      for (auto& attr : getAttrs()) {
+        all_attrs.insert(attr.first.strref());
+      }
+      for (const auto& operand : llvm::enumerate(getType().getInputs())) {
+        if (auto attr_name = getArgAttrOfType<StringAttr>(
+            operand.index(), kAttrArgumentNameAttr)) {
+          all_attrs.insert(attr_name.getValue());
+        }
+      }
+      return all_attrs;
+    }
   }];
 
   let verifier = [{ return Verify(*this); }];

tensorflow/compiler/mlir/tfr/passes/canonicalize.cc

@@ -151,9 +151,21 @@ LogicalResult SimplifySCFIfOp::InlineRegion(Location loc,
 
 }  // namespace
 
-void populateSCFOpsCanonicalizationPatterns(OwningRewritePatternList &results,
-                                            MLIRContext *context) {
-  results.insert<UnrollSCFForOp, SimplifySCFIfOp>(context);
+void populateCanonicalizationPatterns(FuncOp func,
+                                      OwningRewritePatternList &patterns) {
+  MLIRContext *context = func.getContext();
+  mlir::Dialect *tf = context->getLoadedDialect<mlir::TF::TensorFlowDialect>();
+  // Load all official canonicalization patterns. Here we skip the
+  // canonicalization of the ops in the tf dialect, because they couldn't
+  // propagate the attributes correctly. These optimization will be played by
+  // bridge.
+  func->walk([&](Operation *op) {
+    if (op->getDialect() != tf) {
+      op->getAbstractOperation()->getCanonicalizationPatterns(patterns,
+                                                              context);
+    }
+  });
+  patterns.insert<UnrollSCFForOp, SimplifySCFIfOp>(context);
 }
 
 }  // namespace TFR

tensorflow/compiler/mlir/tfr/passes/decompose.cc

@@ -99,15 +99,12 @@ struct DecomposeTFOpsPass
 };
 
 void DecomposeTFOpsPass::ApplyCanonicalization() {
+  FuncOp func = getFunction();
   OwningRewritePatternList patterns;
 
-  auto* context = &getContext();
-  for (auto* op : context->getRegisteredOperations()) {
-    op->getCanonicalizationPatterns(patterns, context);
-  }
-  populateSCFOpsCanonicalizationPatterns(patterns, context);
+  populateCanonicalizationPatterns(func, patterns);
 
-  applyPatternsAndFoldGreedily(getFunction(), std::move(patterns));
+  applyPatternsAndFoldGreedily(func, std::move(patterns));
 }
 
 LogicalResult DecomposeTFOpsPass::RewriteUnregisteredTFOps() {
@@ -122,15 +119,25 @@ LogicalResult DecomposeTFOpsPass::RewriteUnregisteredTFOps() {
     // either will be constant folded or lowered by the rules defined in the
     // bridge.
     if (op->isRegistered()) {
-      return;
+      return WalkResult::advance();
     }
 
     // Find out the compose function
     auto compose_func_name = GetComposeFuncName(op->getName().getStringRef());
     auto compose_func = table.lookup<TFRFuncOp>(compose_func_name);
     if (!compose_func || compose_func.isExternal()) {
       // There are no decomposition methods defined for this op, skip.
-      return;
+      return WalkResult::advance();
+    }
+
+    // Make sure all the attributes are valid. An attribute is valid when it is
+    // in the signature or it is allowed explicitly.
+    auto compose_func_signature =
+        table.lookup<TFRFuncOp>(compose_func_name + "_");
+    if (!compose_func_signature) compose_func_signature = compose_func;
+    auto defined_attrs = compose_func_signature.getDefinedAttributeNames();
+    if (failed(ValidateAttrs(op, defined_attrs))) {
+      return WalkResult::interrupt();
     }
 
     tensorflow::IncreaseOpExpansionExecuteCounterByOne(
@@ -215,8 +222,15 @@ LogicalResult DecomposeTFOpsPass::RewriteUnregisteredTFOps() {
           op->getLoc(), std::get<0>(res).getType(), std::get<1>(res));
       std::get<0>(res).replaceAllUsesWith(casted.out());
     }
+
+    // Copy all the unregisted attributes to the new op.
+    if (failed(CopyAllowedUnregisteredAttrs(op, new_op, defined_attrs))) {
+      return WalkResult::interrupt();
+    }
+
     op->erase();
     changed |= true;
+    return WalkResult::advance();
   });
 
   // If `changed` is false, it is considered as a failure, so the recursive
@@ -237,6 +251,15 @@ LogicalResult DecomposeTFOpsPass::InlineTFRFuncCalls() {
   auto walk_result = func.walk([&](CallOp call_op) {
     auto callee = table.lookup<TFRFuncOp>(call_op.callee());
     if (!callee || callee.isExternal()) return WalkResult::advance();
+
+    // Record the boundary of the inlined operations. The inlined operation will
+    // be inserted between these two operations.
+    Operation* inlined_point = call_op.getOperation();
+    Operation* after_inlined_point =
+        &*std::next(Block::iterator(call_op.getOperation()));
+
+    // Use the inliner to replace all the uses of the call_op by its
+    // composition.
     if (failed(inlineCall(inliner,
                           cast<CallOpInterface>(call_op.getOperation()),
                           cast<CallableOpInterface>(callee.getOperation()),
@@ -246,6 +269,13 @@ LogicalResult DecomposeTFOpsPass::InlineTFRFuncCalls() {
       // This call will be raised to TF ops.
       return WalkResult::interrupt();
     }
+
+    // Propagate all the attributes to the inlined operations, which are defined
+    // by the two boundary operations.
+    PropagateAttrsToOperations(call_op, Block::iterator(inlined_point),
+                               Block::iterator(after_inlined_point));
+
+    // Remove the call_op to finish the op expansion.
     call_op.erase();
     changed |= true;
     return WalkResult::advance();

tensorflow/compiler/mlir/tfr/passes/passes.h

@@ -25,8 +25,10 @@ limitations under the License.
 namespace mlir {
 namespace TFR {
 
-void populateSCFOpsCanonicalizationPatterns(OwningRewritePatternList &results,
-                                            MLIRContext *context);
+// Scans the func op and adds all the canonicalization patterns of the ops
+// except the tf ops, inside the function.
+void populateCanonicalizationPatterns(FuncOp func,
+                                      OwningRewritePatternList &patterns);
 
 // Decompose ops.
 std::unique_ptr<OperationPass<FuncOp>> CreateDecomposeTFOpsPass(

tensorflow/compiler/mlir/tfr/passes/raise_to_tf.cc

@@ -377,6 +377,10 @@ LogicalResult RewriteTFRCallOp::CreateAndReplaceOp(
       new_results.push_back(list_op.out());
     }
   }
+
+  // Copy all the allowed attributes to the new op.
+  if (failed(CopyNonSymbolRefAttrs(call_op, new_op))) return failure();
+
   rewriter.replaceOp(call_op, new_results);
   return success();
 }
@@ -450,9 +454,8 @@ void RaiseToTFOpsPass::runOnFunction() {
 
   OwningRewritePatternList patterns;
   patterns.insert<RewriteTFRCallOp>(ctx, table, materialize_derived_attrs);
-  for (auto* op : ctx->getRegisteredOperations()) {
-    op->getCanonicalizationPatterns(patterns, ctx);
-  }
+
+  populateCanonicalizationPatterns(func, patterns);
 
   applyPatternsAndFoldGreedily(func, std::move(patterns));
 }

tensorflow/compiler/mlir/tfr/resources/decomposition_lib.mlir

@@ -28,6 +28,8 @@ tfr.func @tf__my_add_n(%values: !tfr.tensor_list,
   tfr.return %res : !tfr.tensor
 }
 
+tfr.func @tf__my_add_n_(!tfr.tensor_list<N,T>, i64 {tfr.name="N"}) -> !tfr.tensor attributes {N,T}
+
 // Translated from tf.compose Python function.
 tfr.func @tf__my_biased_dense(%input: !tfr.tensor, %weight: !tfr.tensor,
                               %bias: !tfr.tensor,
@@ -55,6 +57,9 @@ tfr.func @tf__my_biased_dense(%input: !tfr.tensor, %weight: !tfr.tensor,
   tfr.return %res : !tfr.tensor
 }
 
+tfr.func @tf__my_biased_dense_(!tfr.tensor<T>, !tfr.tensor<T>, !tfr.tensor<T>,
+    !tfr.attr{tfr.name="act", tfr.default=""}) -> !tfr.tensor attributes {T}
+
 // This is a wong decomposition and used to verify that tf.Elu isn't decomposed
 // since its kernel has been registered.
 tfr.func @tf__elu_(%input: !tfr.tensor) -> !tfr.tensor {

tensorflow/compiler/mlir/tfr/tests/decompose.mlir

@@ -82,3 +82,43 @@ func @decompose_fused_n(%arg0: tensor<1x2x3x4x!tf.string>, %arg1: tensor<f32>, %
 // CHECK-NEXT: return %[[back]] : tensor<f32>
 }
 
+// CHECK-LABEL: attribute_propagate_direct
+func @attribute_propagate_direct(%arg0: tensor<1x2x3x4x!tf.string>) -> tensor<1x2x3x4x!tf.string> {
+  %0 = "tf.Intermediate"(%arg0) {_tpu_replicate, device="hello"} : (tensor<1x2x3x4x!tf.string>) -> tensor<1x2x3x4x!tf.string>
+  return %0 : tensor<1x2x3x4x!tf.string>
+
+// CHECK-NEXT: %[[casted:.*]] = "tfr.cast"(%arg0) : (tensor<1x2x3x4x!tf.string>) -> !tfr.tensor
+// CHECK-NEXT: %[[id:.*]] = tfr.call @tf__risc(%[[casted]]) {_tpu_replicate, device = "hello"}
+// CHECK-NEXT: %[[back:.*]] = "tfr.cast"(%[[id]]) : (!tfr.tensor) -> tensor<1x2x3x4x!tf.string>
+// CHECK-NEXT: return %[[back]]
+}
+
+// CHECK-LABEL: attribute_propagate
+func @attribute_propagate(%arg0: tensor<1x2x3x4x!tf.string>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
+  %0:2 = "tf.FusedN"(%arg0, %arg1, %arg2) {A=0:index, _tpu_replicate, device="hello"} : (tensor<1x2x3x4x!tf.string>, tensor<f32>, tensor<f32>) -> (tensor<1x2x3x4x!tf.string>, tensor<f32>)
+  return %0#1 : tensor<f32>
+
+// CHECK-NEXT: %[[in0:.*]] = "tfr.cast"(%arg0) : (tensor<1x2x3x4x!tf.string>) -> !tfr.tensor
+// CHECK-NEXT: %[[in1:.*]] = "tfr.cast"(%arg1) : (tensor<f32>) -> !tfr.tensor
+// CHECK-NEXT: %[[id0:.*]] = tfr.call @tf__risc(%[[in0]]) {_tpu_replicate, device = "hello"}
+// CHECK-NEXT: %[[id1:.*]] = tfr.call @tf__risc(%[[in1]]) {_tpu_replicate, device = "hello"}
+// CHECK-NEXT: %[[back:.*]] = "tfr.cast"(%[[id1]]) : (!tfr.tensor) -> tensor<f32>
+// CHECK-NEXT: return %[[back]] : tensor<f32>
+}
+
+// CHECK-LABEL: no_tf_canonicalization
+func @no_tf_canonicalization(%arg0: tensor<8xi1>, %arg1: tensor<8x3xf32>, %arg2: tensor<8x3xf32>) -> tensor<8x3xf32> {
+  %0 = "tf.Select"(%arg0, %arg1, %arg2) : (tensor<8xi1>, tensor<8x3xf32>, tensor<8x3xf32>) -> tensor<8x3xf32>
+  return %0: tensor<8x3xf32>
+
+// CHECK:   "tf.Select"
+}
+
+// CHECK-LABEL: denied_attribute
+func @denied_attribute(%arg0: tensor<1x2x3x4x!tf.string>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
+  // expected-error@+1 {{Denied unregistered attribute was found: denied_attr}}
+  %0:2 = "tf.FusedN"(%arg0, %arg1, %arg2) {A=0:index, denied_attr} : (tensor<1x2x3x4x!tf.string>, tensor<f32>, tensor<f32>) -> (tensor<1x2x3x4x!tf.string>, tensor<f32>)
+  return %0#1 : tensor<f32>
+
+// CHECK-NEXT:   "tf.FusedN"(%arg0, %arg1, %arg2) {A = 0 : index, denied_attr}
+}