[mlir][PatternRewrite] Allow pattern application for regions of ops not isolated from above.

[mravishankar]

Summary

Right now applying patterns for ops not isolated from above hits an
assert. Trying this on a region not isolated from above also works,
but probably has broader impact. For now adding an option to enable
this as an opt-in.

[mehdi_amini]

mlir/include/mlir/Transforms/GreedyPatternRewriteDriver.h

54

It's not clear to me that this belongs to a config flag, why?

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

722

What about this comment?

[nicolasvasilache]

This discussion: https://discourse.llvm.org/t/why-does-greedypatternrewritedriver-requires-region-isolated-from-above/63724 and this PR: https://reviews.llvm.org/D141366 seem relevant.

mlir/include/mlir/Transforms/GreedyPatternRewriteDriver.h

public:
   bool useTopDownTraversal = false;
 
   /// Perform control flow optimizations to the region tree after applying all
   /// patterns.
   ///
   /// Note: Only applicable when simplifying entire regions.
   bool enableRegionSimplification = true;
 
+  /// Allow pattern rewrites for operations that are not isolated from above.
+  ///
+  /// Note: This is by default off.
+  bool enableRewritesForNonIsolatedFromAboveOps = false;
+
   /// This specifies the maximum number of times the rewriter will iterate
   /// between applying patterns and simplifying regions. Use `kNoLimit` to
   /// disable this iteration limit.
   ///
   /// Note: Only applicable when simplifying entire regions.
   int64_t maxIterations = 10;
 
   /// This specifies the maximum number of rewrites within an iteration. Use

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

LogicalResult RegionPatternRewriteDriver::simplify(bool *changed) && {
   // Whether the rewrite converges, i.e. wasn't changed in the last iteration.
   return success(!continueRewrites);
 }
 
 LogicalResult
 mlir::applyPatternsAndFoldGreedily(Region &region,
                                    const FrozenRewritePatternSet &patterns,
                                    GreedyRewriteConfig config, bool *changed) {
-  // The top-level operation must be known to be isolated from above to
-  // prevent performing canonicalizations on operations defined at or above
-  // the region containing 'op'.
-  assert(region.getParentOp()->hasTrait<OpTrait::IsIsolatedFromAbove>() &&
-         "patterns can only be applied to operations IsolatedFromAbove");
+  if (!region.getParentOp()->hasTrait<OpTrait::IsIsolatedFromAbove>() &&
+      !config.enableRewritesForNonIsolatedFromAboveOps) {
+    return region.getParentOp()->emitOpError(
+        "invalid application of patterns to operation not IsolatedFromAbove. "
+        "To enable set `enableRewritesForNonIsolatedFromAboveOps` to `true`");
+  }
 
   // Set scope if not specified.
   if (!config.scope)
     config.scope = &region;
 
   // Start the pattern driver.
   RegionPatternRewriteDriver driver(region.getContext(), patterns, config,
                                     region);

mlir/test/IR/pattern-rewriter-non-isolated-from-above.mlir

+// RUN: mlir-opt --test-pattern-non-isolated-from-above %s | FileCheck %s
+
+func.func @test(%arg0 : i32) {
+  %0 = "test.cast"(%arg0) : (i32) -> f32
+  "test.one_region_op"()({
+    %1 = "test.cast"(%0) : (f32) -> i32
+    "test.region_yield"(%1) : (i32) -> ()
+  }) : () -> ()
+  return
+}
+// CHECK-LABEL: func @test(
+//  CHECK-SAME:     %[[ARG0:.+]]: i32)
+//       CHECK:   test.one_region_op
+//       CHECK:     test.region_yield
+//  CHECK-SAME:         %[[ARG0]]

mlir/test/lib/Dialect/Test/TestPatterns.cpp

void runOnOperation() override {
     mlir::RewritePatternSet patterns(context);
     patterns.add<TestSelectiveOpReplacementPattern>(context);
     (void)applyPatternsAndFoldGreedily(getOperation(), std::move(patterns));
   }
 };
 } // namespace
 
 //===----------------------------------------------------------------------===//
+// Test Not isolated from above
+//===----------------------------------------------------------------------===//
+
+namespace {
+struct TestFoldCastOpPattern : public OpRewritePattern<TestCastOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(TestCastOp op,
+                                PatternRewriter &rewriter) const final {
+    if (op.getNumOperands() != 1)
+      return failure();
+    auto definingCastOp = op->getOperand(0).getDefiningOp<TestCastOp>();
+    if (!definingCastOp || definingCastOp->getNumOperands() != 1)
+      return failure();
+    Value source = definingCastOp.getOperand(0);
+    if (source.getType() != op.getResult().getType())
+      return failure();
+    rewriter.replaceOp(op, definingCastOp.getOperand(0));
+    return success();
+  }
+};
+
+struct TestNonIsolatedFromAbovePatternDriver
+    : public PassWrapper<TestNonIsolatedFromAbovePatternDriver,
+                         OperationPass<>> {
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(
+      TestNonIsolatedFromAbovePatternDriver)
+
+  StringRef getArgument() const final {
+    return "test-pattern-non-isolated-from-above";
+  }
+  StringRef getDescription() const final {
+    return "Test pattern application within operation that is not isolated "
+           "from above";
+  }
+  void runOnOperation() override {
+    MLIRContext *context = &getContext();
+    mlir::RewritePatternSet patterns(context);
+    patterns.insert<TestFoldCastOpPattern>(context);
+    SmallVector<OneRegionOp> oneRegionOps;
+    getOperation()->walk(
+        [&](OneRegionOp oneRegionOp) { oneRegionOps.push_back(oneRegionOp); });
+    GreedyRewriteConfig config;
+    config.enableRewritesForNonIsolatedFromAboveOps = true;
+    for (auto oneRegionOp : oneRegionOps) {
+      if (failed(applyPatternsAndFoldGreedily(oneRegionOp, std::move(patterns),
+                                              config))) {
+        return signalPassFailure();
+      }
+    }
+  }
+};
+} // namespace
+
+//===----------------------------------------------------------------------===//
 // PassRegistration
 //===----------------------------------------------------------------------===//
 
 namespace mlir {
 namespace test {
 void registerPatternsTestPass() {
   PassRegistration<TestReturnTypeDriver>();
 
void registerPatternsTestPass() {
 
   PassRegistration<TestTypeConversionDriver>();
   PassRegistration<TestTargetMaterializationWithNoUses>();
 
   PassRegistration<TestRewriteDynamicOpDriver>();
 
   PassRegistration<TestMergeBlocksPatternDriver>();
   PassRegistration<TestSelectiveReplacementPatternDriver>();
+
+  PassRegistration<TestNonIsolatedFromAbovePatternDriver>();
 }
 } // namespace test
 } // namespace mlir