[mlir][transform] Add region to ApplyPatternsOp

[springerm]

Summary

Patterns should be selected by adding ops that implement PatternDescriptorOpInterface to the region of apply_pattern ops. Such ops can have operands, allowing for pattern parameterization. The existing way of selecting patterns from the PatternRegistry is deprecated.

mlir/include/mlir/Dialect/Transform/IR/TransformInterfaces.td

InterfaceMethod<
       }],
       /*returnType=*/"::llvm::SmallVector<::mlir::Value>",
       /*name=*/"getNextOperands",
       /*arguments=*/(ins)
     >,
   ];
 }
 
+def PatternDescriptorOpInterface : OpInterface<"PatternDescriptorOpInterface"> {
+  let description = [{
+    This interface should be implemented by ops that select patterns of a
+    `transform.apply_patterns` op. It provides a method to populate a rewrite
+    pattern set with patterns.
+  }];
+
+  let cppNamespace = "::mlir::transform";
+
+  let methods = [
+    InterfaceMethod<
+      /*desc=*/[{
+        Populate patterns into the given pattern set.
+      }],
+      /*returnType=*/"void",
+      /*name=*/"populatePatterns",
+      /*arguments=*/(ins "RewritePatternSet &":$patterns)
+    >,
+  ];
+}
+
 #endif // MLIR_DIALECT_TRANSFORM_IR_TRANSFORM_INTERFACES_TD

mlir/include/mlir/Dialect/Transform/IR/TransformOps.td

 
 include "mlir/Interfaces/CallInterfaces.td"
 include "mlir/Interfaces/CastInterfaces.td"
 include "mlir/Interfaces/ControlFlowInterfaces.td"
 include "mlir/Interfaces/InferTypeOpInterface.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
 include "mlir/IR/FunctionInterfaces.td"
 include "mlir/IR/OpAsmInterface.td"
+include "mlir/IR/RegionKindInterface.td"
 include "mlir/IR/SymbolInterfaces.td"
 include "mlir/Dialect/Transform/IR/MatchInterfaces.td"
 include "mlir/Dialect/Transform/IR/TransformAttrs.td"
 include "mlir/Dialect/Transform/IR/TransformDialect.td"
 include "mlir/Dialect/Transform/IR/TransformInterfaces.td"
 
 def AlternativesOp : TransformDialectOp<"alternatives",
     [DeclareOpInterfaceMethods<RegionBranchOpInterface,
def AnnotateOp : TransformDialectOp<"annotate",
 
   let assemblyFormat =
     "$target $name attr-dict (`=` $param^)?"
     "`:` type($target) (`,` type($param)^)?";
 }
 
 def ApplyPatternsOp : TransformDialectOp<"apply_patterns",
     [TransformOpInterface, TransformEachOpTrait,
-     DeclareOpInterfaceMethods<MemoryEffectsOpInterface>]> {
+     DeclareOpInterfaceMethods<MemoryEffectsOpInterface>]
+        # GraphRegionNoTerminator.traits> {
   let summary = "Greedily applies patterns to the body of the targeted op";
   let description = [{
     This transform greedily applies the specified patterns to the body of the
     targeted op until a fixpoint was reached. Patterns are not applied to the
     targeted op itself.
 
-    Only patterns that were registered in the transform dialect's
-    `PatternRegistry` are available. Additional patterns can be registered as
-    part of transform dialect extensions. "canonicalization" is a special set
-    of patterns that refers to all canonicalization patterns of all loaded
+    The patterns that should be applied are specified in the graph region of
+    this op. They must implement the `PatternDescriptorOpInterface`.
+
+    (Deprecated) In addition, patterns that were registered in the transform
+    dialect's `PatternRegistry` are available. "canonicalization" is a special
+    set of patterns that refers to all canonicalization patterns of all loaded
     dialects.
 
     This transform only reads the target handle and modifies the payload. If a
     pattern erases or replaces a tracked op, the mapping is updated accordingly.
 
     Only replacements via `RewriterBase::replaceOp` or `replaceOpWithNewOp` are
     considered "payload op replacements". Furthermore, only if the replacement
     values are defined by the same op and that op has the same type as the
     original op, the mapping is updated. Otherwise, this transform fails
     silently unless `fail_on_payload_replacement_not_found` is set to "false".
     More details can be found at the documentation site of `TrackingListener`.
 
     This transform also fails silently if the pattern application did not
     converge within the default number of iterations/rewrites of the greedy
     pattern rewrite driver.
   }];
 
   let arguments = (ins
     TransformHandleTypeInterface:$target, ArrayAttr:$patterns,
     DefaultValuedAttr<BoolAttr, "true">:$fail_on_payload_replacement_not_found);
   let results = (outs);
-  let assemblyFormat = "$patterns `to` $target attr-dict `:` type($target)";
+  let regions = (region MaxSizedRegion<1>:$region);
+
+  let assemblyFormat = "$patterns `to` $target $region attr-dict `:` type($target)";
   let hasVerifier = 1;
 
   let extraClassDeclaration = [{
     ::mlir::DiagnosedSilenceableFailure applyToOne(
       ::mlir::Operation *target,
       ::mlir::transform::ApplyToEachResultList &results,
       ::mlir::transform::TransformState &state);
   }];
 }
 
+def ApplyCanonicalizationPatternsOp
+    : TransformDialectOp<"apply_patterns.canonicalization",
+        [DeclareOpInterfaceMethods<PatternDescriptorOpInterface>]> {
+  let summary = "Populates canonicalization patterns";
+  let description = [{
+    This op populates all canonicalization patterns of all loaded dialects in
+    an `apply_patterns` transform.
+  }];
+  let assemblyFormat = "attr-dict";
+}
+
 def CastOp : TransformDialectOp<"cast",
     [TransformOpInterface, TransformEachOpTrait,
      DeclareOpInterfaceMethods<CastOpInterface>,
      DeclareOpInterfaceMethods<MemoryEffectsOpInterface>]> {
   let arguments = (ins TransformHandleTypeInterface:$input);
   let results = (outs TransformHandleTypeInterface:$output);
   let assemblyFormat = "$input attr-dict `:` type($input) `to` type($output)";
 

mlir/lib/Dialect/Transform/IR/TransformDialect.cpp

 void transform::detail::checkImplementsTransformOpInterface(
     StringRef name, MLIRContext *context) {
   // Since the operation is being inserted into the Transform dialect and the
   // dialect does not implement the interface fallback, only check for the op
   // itself having the interface implementation.
   RegisteredOperationName opName =
       *RegisteredOperationName::lookup(name, context);
   assert((opName.hasInterface<TransformOpInterface>() ||
+          opName.hasInterface<PatternDescriptorOpInterface>() ||
           opName.hasTrait<OpTrait::IsTerminator>()) &&
          "non-terminator ops injected into the transform dialect must "
-         "implement TransformOpInterface");
+         "implement TransformOpInterface or PatternDescriptorOpInterface");
+  if (!opName.hasInterface<PatternDescriptorOpInterface>()) {
     assert(opName.hasInterface<MemoryEffectOpInterface>() &&
            "ops injected into the transform dialect must implement "
            "MemoryEffectsOpInterface");
   }
+}
 
 void transform::detail::checkImplementsTransformHandleTypeInterface(
     TypeID typeID, MLIRContext *context) {
   const auto &abstractType = AbstractType::lookup(typeID, context);
   assert((abstractType.hasInterface(
               TransformHandleTypeInterface::getInterfaceID()) ||
           abstractType.hasInterface(
               TransformParamTypeInterface::getInterfaceID()) ||
           abstractType.hasInterface(
               TransformValueHandleTypeInterface::getInterfaceID())) &&
          "expected Transform dialect type to implement one of the three "
          "interfaces");
 }
 #endif // NDEBUG
 
 void transform::TransformDialect::initialize() {
   // Using the checked versions to enable the same assertions as for the ops
   // from extensions.
   addOperationsChecked<
 #define GET_OP_LIST
 #include "mlir/Dialect/Transform/IR/TransformOps.cpp.inc"
       >();
   initializeTypes();
-
-  // Register all canonicalization patterns.
-  getOrCreateExtraData<transform::PatternRegistry>().registerPatterns(
-      "canonicalization", [](RewritePatternSet &patterns) {
-        MLIRContext *ctx = patterns.getContext();
-        for (Dialect *dialect : ctx->getLoadedDialects())
-          dialect->getCanonicalizationPatterns(patterns);
-        for (RegisteredOperationName op : ctx->getRegisteredOperations())
-          op.getCanonicalizationPatterns(patterns, ctx);
-      });
 }
 
 Type transform::TransformDialect::parseType(DialectAsmParser &parser) const {
   StringRef keyword;
   SMLoc loc = parser.getCurrentLocation();
   if (failed(parser.parseKeyword(&keyword)))
     return nullptr;
 

mlir/lib/Dialect/Transform/IR/TransformOps.cpp

transform::ApplyPatternsOp::applyToOne(Operation *target,
   // Gather all specified patterns.
   MLIRContext *ctx = target->getContext();
   RewritePatternSet patterns(ctx);
   const auto &registry = getContext()
                              ->getLoadedDialect<transform::TransformDialect>()
                              ->getExtraData<transform::PatternRegistry>();
   for (Attribute attr : getPatterns())
     registry.populatePatterns(attr.cast<StringAttr>(), patterns);
+  if (!getRegion().empty()) {
+    for (Operation &op : getRegion().front()) {
+      cast<transform::PatternDescriptorOpInterface>(&op).populatePatterns(
+          patterns);
+    }
+  }
 
   // Configure the GreedyPatternRewriteDriver.
   ErrorCheckingTrackingListener listener(state, *this);
   GreedyRewriteConfig config;
   config.listener = &listener;
 
   // Manually gather list of ops because the other GreedyPatternRewriteDriver
   // overloads only accepts ops that are isolated from above. This way, patterns
LogicalResult transform::ApplyPatternsOp::verify() {
   for (Attribute attr : getPatterns()) {
     auto strAttr = attr.dyn_cast<StringAttr>();
     if (!strAttr)
       return emitOpError() << "expected " << getPatternsAttrName()
                            << " to be an array of strings";
     if (!registry.hasPatterns(strAttr))
       return emitOpError() << "patterns not registered: " << strAttr.strref();
   }
+  if (!getRegion().empty()) {
+    for (Operation &op : getRegion().front()) {
+      if (!isa<transform::PatternDescriptorOpInterface>(&op)) {
+        InFlightDiagnostic diag = emitOpError()
+                                  << "expected children ops to implement "
+                                     "PatternDescriptorOpInterface";
+        diag.attachNote(op.getLoc()) << "op without interface";
+        return diag;
+      }
+    }
+  }
   return success();
 }
 
 void transform::ApplyPatternsOp::getEffects(
     SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
   transform::onlyReadsHandle(getTarget(), effects);
   transform::modifiesPayload(effects);
 }
 
 //===----------------------------------------------------------------------===//
+// ApplyCanonicalizationPatternsOp
+//===----------------------------------------------------------------------===//
+
+void transform::ApplyCanonicalizationPatternsOp::populatePatterns(
+    RewritePatternSet &patterns) {
+  MLIRContext *ctx = patterns.getContext();
+  for (Dialect *dialect : ctx->getLoadedDialects())
+    dialect->getCanonicalizationPatterns(patterns);
+  for (RegisteredOperationName op : ctx->getRegisteredOperations())
+    op.getCanonicalizationPatterns(patterns, ctx);
+}
+
+//===----------------------------------------------------------------------===//
 // CastOp
 //===----------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure
 transform::CastOp::applyToOne(Operation *target, ApplyToEachResultList &results,
                               transform::TransformState &state) {
   results.push_back(target);
   return DiagnosedSilenceableFailure::success();

mlir/test/Dialect/Transform/ops-invalid.mlir

module attributes { transform.with_named_sequence } {
   }
 }
 
 // -----
 
 transform.sequence failures(propagate) {
 ^bb0(%arg0: !transform.any_op):
   // expected-error @below {{patterns not registered: transform.invalid_pattern_identifier}}
-  transform.apply_patterns ["transform.invalid_pattern_identifier"] to %arg0 : !transform.any_op
+  transform.apply_patterns ["transform.invalid_pattern_identifier"] to %arg0 {} : !transform.any_op
 }
 
 // -----
 
 transform.sequence failures(propagate) {
 ^bb0(%arg0: !transform.any_op):
   // expected-error @below {{expected "patterns" to be an array of strings}}
-  transform.apply_patterns [3, 9] to %arg0 : !transform.any_op
+  transform.apply_patterns [3, 9] to %arg0 {} : !transform.any_op
+}
+
+// -----
+transform.sequence failures(propagate) {
+^bb0(%arg0: !transform.any_op):
+  // expected-error @below {{expected children ops to implement PatternDescriptorOpInterface}}
+  transform.apply_patterns [] to %arg0 {
+    // expected-note @below {{op without interface}}
+    transform.named_sequence @foo()
+  } : !transform.any_op
 }

mlir/test/Dialect/Transform/test-pattern-application.mlir

func.func @update_tracked_op_mapping() {
   }) : () -> ()
   return
 }
 
 transform.sequence failures(propagate) {
 ^bb1(%arg1: !transform.any_op):
   %0 = transform.structured.match ops{["test.container"]} in %arg1 : (!transform.any_op) -> !transform.any_op
   %1 = transform.structured.match ops{["test.foo"]} in %arg1 : (!transform.any_op) -> !transform.any_op
-  transform.apply_patterns ["transform.test"] to %0 : !transform.any_op
+  transform.apply_patterns ["transform.test"] to %0 {} : !transform.any_op
+  // Add an attribute to %1, which is now mapped to a new op.
+  transform.annotate %1 "annotated" : !transform.any_op
+}
+
+// -----
+
+// CHECK-LABEL: func @update_tracked_op_mapping_region()
+//       CHECK:   "test.container"() ({
+//       CHECK:     %0 = "test.foo"() {annotated} : () -> i32
+//       CHECK:   }) : () -> ()
+func.func @update_tracked_op_mapping_region() {
+  "test.container"() ({
+    %0 = "test.foo"() {replace_with_new_op = "test.foo"} : () -> (i32)
+  }) : () -> ()
+  return
+}
+
+transform.sequence failures(propagate) {
+^bb1(%arg1: !transform.any_op):
+  %0 = transform.structured.match ops{["test.container"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+  %1 = transform.structured.match ops{["test.foo"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+  transform.apply_patterns [] to %0 {
+    transform.apply_patterns.transform.test_patterns
+  } : !transform.any_op
   // Add an attribute to %1, which is now mapped to a new op.
   transform.annotate %1 "annotated" : !transform.any_op
 }
 
 // -----
 
 func.func @replacement_op_not_found() {
   "test.container"() ({
     // expected-note @below {{[0] replaced op}}
     // expected-note @below {{[0] replacement value 0}}
     %0 = "test.foo"() {replace_with_new_op = "test.bar"} : () -> (i32)
   }) : () -> ()
   return
 }
 
 transform.sequence failures(propagate) {
 ^bb1(%arg1: !transform.any_op):
   %0 = transform.structured.match ops{["test.container"]} in %arg1 : (!transform.any_op) -> !transform.any_op
   %1 = transform.structured.match ops{["test.foo"]} in %arg1 : (!transform.any_op) -> !transform.any_op
   // expected-error @below {{tracking listener failed to find replacement op}}
-  transform.apply_patterns ["transform.test"] to %0 : !transform.any_op
+  transform.apply_patterns ["transform.test"] to %0 {} : !transform.any_op
   // %1 must be used in some way. If no replacement payload op could be found,
   // an error is thrown only if the handle is not dead.
   transform.annotate %1 "annotated" : !transform.any_op
 }
 
 // -----
 
 // CHECK-LABEL: func @replacement_op_for_dead_handle_not_found()
 //       CHECK:   "test.container"() ({
 //       CHECK:     %0 = "test.bar"() : () -> i32
 //       CHECK:   }) : () -> ()
 func.func @replacement_op_for_dead_handle_not_found() {
   "test.container"() ({
     %0 = "test.foo"() {replace_with_new_op = "test.bar"} : () -> (i32)
   }) : () -> ()
   return
 }
 
 transform.sequence failures(propagate) {
 ^bb1(%arg1: !transform.any_op):
   %0 = transform.structured.match ops{["test.container"]} in %arg1 : (!transform.any_op) -> !transform.any_op
   %1 = transform.structured.match ops{["test.foo"]} in %arg1 : (!transform.any_op) -> !transform.any_op
   // No error because %1 is dead.
-  transform.apply_patterns ["transform.test"] to %0 : !transform.any_op
+  transform.apply_patterns ["transform.test"] to %0 {} : !transform.any_op
 }
 
 // -----
 
 // CHECK-LABEL: func @replacement_op_not_found_silenced()
 //       CHECK:   "test.container"() ({
 //       CHECK:     %0 = "test.bar"() : () -> i32
 //       CHECK:   }) : () -> ()
 func.func @replacement_op_not_found_silenced() {
   "test.container"() ({
     %0 = "test.foo"() {replace_with_new_op = "test.bar"} : () -> (i32)
   }) : () -> ()
   return
 }
 
 transform.sequence failures(propagate) {
 ^bb1(%arg1: !transform.any_op):
   %0 = transform.structured.match ops{["test.container"]} in %arg1 : (!transform.any_op) -> !transform.any_op
   %1 = transform.structured.match ops{["test.foo"]} in %arg1 : (!transform.any_op) -> !transform.any_op
-  transform.apply_patterns ["transform.test"] to %0 {fail_on_payload_replacement_not_found = false}: !transform.any_op
+  transform.apply_patterns ["transform.test"] to %0 {} {fail_on_payload_replacement_not_found = false}: !transform.any_op
   transform.annotate %1 "annotated" : !transform.any_op
 }
 
 // -----
 
 // CHECK-LABEL: func @patterns_apply_only_to_target_body()
 //       CHECK:   %0 = "test.foo"() {replace_with_new_op = "test.bar"} : () -> i32
 func.func @patterns_apply_only_to_target_body() {
   %0 = "test.foo"() {replace_with_new_op = "test.bar"} : () -> (i32)
   return
 }
 
 transform.sequence failures(propagate) {
 ^bb1(%arg1: !transform.any_op):
-  %0 = transform.structured.match ops{["test.foo"]} in %arg1 : (!transform.any_op) -> !transform.any_op
-  transform.apply_patterns ["transform.test"] to %0 : !transform.any_op
+  transform.apply_patterns ["transform.test"] to %0 {} : !transform.any_op
 }
 
 // -----
 
 // CHECK-LABEL: func @erase_tracked_op()
 //       CHECK:   "test.container"() ({
 //  CHECK-NEXT:   ^bb0:
 //  CHECK-NEXT:   }) : () -> ()
 func.func @erase_tracked_op() {
   "test.container"() ({
     // expected-remark @below {{matched op}}
     %0 = "test.erase_op"() {replace_with_new_op = "test.foo"} : () -> (i32)
   }) : () -> ()
   return
 }
 
 transform.sequence failures(propagate) {
 ^bb1(%arg1: !transform.any_op):
   %0 = transform.structured.match ops{["test.container"]} in %arg1 : (!transform.any_op) -> !transform.any_op
   %1 = transform.structured.match ops{["test.erase_op"]} in %arg1 : (!transform.any_op) -> !transform.any_op
   transform.test_print_remark_at_operand %1, "matched op" : !transform.any_op
-  transform.apply_patterns ["transform.test"] to %0 : !transform.any_op
+  transform.apply_patterns ["transform.test"] to %0 {} : !transform.any_op
   transform.test_print_remark_at_operand %1, "op was deleted" : !transform.any_op
 }
 
 // -----
 
 // CHECK-LABEL: func @canonicalization(
 //       CHECK:   %[[c5:.*]] = arith.constant 5 : index
 //       CHECK:   return %[[c5]]
 func.func @canonicalization(%t: tensor<5xf32>) -> index {
   %c0 = arith.constant 0 : index
   // expected-remark @below {{op was replaced}}
   %dim = tensor.dim %t, %c0 : tensor<5xf32>
   return %dim : index
 }
 
 transform.sequence failures(propagate) {
 ^bb1(%arg1: !transform.any_op):
   %0 = transform.structured.match ops{["tensor.dim"]} in %arg1 : (!transform.any_op) -> !transform.any_op
   %1 = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
-  transform.apply_patterns ["canonicalization"] to %1 : !transform.any_op
+  transform.apply_patterns [] to %1 {
+    transform.apply_patterns.canonicalization
+  } : !transform.any_op
   transform.test_print_remark_at_operand %0, "op was replaced" : !transform.any_op
 }

mlir/test/lib/Dialect/Transform/TestTransformDialectExtension.cpp

public:
   EraseOp(MLIRContext *context)
       : RewritePattern("test.erase_op", /*benefit=*/1, context) {}
   LogicalResult matchAndRewrite(Operation *op,
                                 PatternRewriter &rewriter) const override {
     rewriter.eraseOp(op);
     return success();
   }
 };
+} // namespace
 
+void mlir::test::ApplyTestPatternsOp::populatePatterns(
+    RewritePatternSet &patterns) {
+  patterns.insert<ReplaceWithNewOp, EraseOp>(patterns.getContext());
+}
+
+namespace {
 void populateTestPatterns(RewritePatternSet &patterns) {
   patterns.insert<ReplaceWithNewOp, EraseOp>(patterns.getContext());
 }
 
 /// Test extension of the Transform dialect. Registers additional ops and
 /// declares PDL as dependent dialect since the additional ops are using PDL
 /// types for operands and results.
 class TestTransformDialectExtension

mlir/test/lib/Dialect/Transform/TestTransformDialectExtension.td

 def TestProduceValueHandleToSelfOperand
   : Op<Transform_Dialect, "test_produce_value_handle_to_self_operand",
        [DeclareOpInterfaceMethods<TransformOpInterface>,
         DeclareOpInterfaceMethods<MemoryEffectsOpInterface>]> {
   let arguments = (ins TransformHandleTypeInterface:$in);
   let results = (outs TransformValueHandleTypeInterface:$out);
   let assemblyFormat = "$in attr-dict `:` functional-type(operands, results)";
   let cppNamespace = "::mlir::test";
-  
 }
 
 def TestProduceValueHandleToResult
   : Op<Transform_Dialect, "test_produce_value_handle_to_result",
        [TransformEachOpTrait, TransformOpInterface,
         DeclareOpInterfaceMethods<MemoryEffectsOpInterface>]> {
   let arguments = (ins TransformHandleTypeInterface:$in,
                        I64Attr:$number);
: Op<Transform_Dialect, "test_tracked_rewrite",
       [DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
        DeclareOpInterfaceMethods<TransformOpInterface>]> {
   let arguments = (ins TransformHandleTypeInterface:$in);
   let results = (outs);
   let assemblyFormat = "$in attr-dict `:` functional-type(operands, results)";
   let cppNamespace = "::mlir::test";
 }
 
+def ApplyTestPatternsOp
+  : Op<Transform_Dialect, "apply_patterns.transform.test_patterns",
+      [DeclareOpInterfaceMethods<PatternDescriptorOpInterface>]> {
+  let arguments = (ins);
+  let results = (outs);
+  let assemblyFormat = "attr-dict";
+  let cppNamespace = "::mlir::test";
+}
+
 #endif // MLIR_TESTTRANSFORMDIALECTEXTENSION_TD