[MLIR][SCF] Inline ExecuteRegion if parent can contain multiple blocks

[wsmoses]

Summary

The executeregionop is used to allow multiple blocks within SCF constructs. If the container allows multiple blocks, inline the region

[wsmoses]

Add newline

[bondhugula]

The execution region op can be inlined even if its parent can only hold a single block (like scf.for, scf.if) when the execute_region only contains a single block. I think the rewrite pattern you added can be transparently generalized for that case without really any addition in code?

[wsmoses]

The more generic single block simplification you're suggesting was added here: https://reviews.llvm.org/D104865 .

[bondhugula]

In D104960#2842230, @wsmoses wrote:

The more generic single block simplification you're suggesting was added here: https://reviews.llvm.org/D104865 .

I see - thanks. I had missed that.

[kiranchandramohan]

mlir/lib/Dialect/SCF/SCF.cpp

214

Is there an assumption here that the yield terminator will be the terminator of the last block? Is it possible that there are multiple terminators? I don't see anything in the verifier of the scf.execute_region op that prevents it.

[bondhugula]

mlir/lib/Dialect/SCF/SCF.cpp

214

scf.execute_region can have multiple blocks terminated with yield - it's no different from a function's region in that respect.

[wsmoses]

Handle regions with multiple yields

[ftynse]

mlir/include/mlir/Dialect/SCF/SCFOps.td

112–113

I'm not sure it's feasible. Folders are not supposed to create new operations, and inlining most likely will (the inlined operations will be "new" in the block).

mlir/lib/Dialect/SCF/SCF.cpp

223

Nit: just use SmallVector<Value> with the default number of stack elements

227

Not: drop trivial braces

mlir/test/Dialect/SCF/canonicalize.mlir

973

Could we relax this test a little? In particular,

• do we care that operations are on subsequent lines? (I think we don't, use CHECK)
• do we care about type signatures of these operations? (they are not affected by the pattern, so drop them)
• do we care about the comment // n preds:?
• do we care about the closing brace?
• do we care about block names that are not guaranteed to be fixed? (we don't, pattern-match them similarly to variables).

On the contrary, we do care about execute-region being removed, so I would add CHECK-NOT: execute_region.

990

Nit: could you change the name of this function to be able to tell the "test cases" apart?

[bondhugula]

This revision was not accepted when it landed; it landed in state Needs Review.

Was this committed unintentionally? I don't see an approval from @ftynse post comments nor from anyone else.

[wsmoses]

In D104960#2844778, @bondhugula wrote:

This revision was not accepted when it landed; it landed in state Needs Review.

Was this committed unintentionally? I don't see an approval from @ftynse post comments nor from anyone else.

Oh shoot, yes completely unintentional. I'll revert now and resubmit the PR.

[wsmoses]

Reopen changes

[ftynse]

Would you mind to mark comments as done next time? There's a checkbox in the top right corner of each comment box. This makes it easier to re-review the patch.

mlir/test/Dialect/SCF/canonicalize.mlir

977

980

1011

1014

mlir/include/mlir/Dialect/SCF/SCFOps.td

let description = [{
     }
     ```
   }];
 
   let results = (outs Variadic<AnyType>);
 
   let regions = (region AnyRegion:$region);
 
-  // TODO: If the parent is a func like op (which would be the case if all other
-  // ops are from the std dialect), the inliner logic could be readily used to
-  // inline.
   let hasCanonicalizer = 1;
 
-  // TODO: can fold if it returns a constant.
-  // TODO: Single block execute_region ops can be readily inlined irrespective
-  // of which op is a parent. Add a fold for this.
   let hasFolder = 0;
 }
 
 def ForOp : SCF_Op<"for",
       [DeclareOpInterfaceMethods<LoopLikeOpInterface>,
        DeclareOpInterfaceMethods<RegionBranchOpInterface>,
        SingleBlockImplicitTerminator<"scf::YieldOp">,
        RecursiveSideEffects]> {
   let summary = "for operation";

mlir/lib/Dialect/SCF/SCF.cpp

 //        scf.yield %x : i64
 //      }
 //      "test.bar"(%v) : (i64) -> ()
 //
 //  becomes
 //
 //     "test.foo"() : () -> ()
 //     %x = "test.val"() : () -> i64
-//     "test.bar"(%v) : (i64) -> ()
+//     "test.bar"(%x) : (i64) -> ()
 //
 struct SingleBlockExecuteInliner : public OpRewritePattern<ExecuteRegionOp> {
   using OpRewritePattern<ExecuteRegionOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(ExecuteRegionOp op,
                                 PatternRewriter &rewriter) const override {
-    if (op.region().getBlocks().size() != 1)
+    if (!llvm::hasSingleElement(op.region()))
       return failure();
     replaceOpWithRegion(rewriter, op, op.region());
     return success();
   }
 };
 
+// Inline an ExecuteRegionOp if its parent can contain multiple blocks.
+// TODO generalize the conditions for operations which can be inlined into.
+// func @func_execute_region_elim() {
+//     "test.foo"() : () -> ()
+//     %v = scf.execute_region -> i64 {
+//       %c = "test.cmp"() : () -> i1
+//       cond_br %c, ^bb2, ^bb3
+//     ^bb2:
+//       %x = "test.val1"() : () -> i64
+//       br ^bb4(%x : i64)
+//     ^bb3:
+//       %y = "test.val2"() : () -> i64
+//       br ^bb4(%y : i64)
+//     ^bb4(%z : i64):
+//       scf.yield %z : i64
+//     }
+//     "test.bar"(%v) : (i64) -> ()
+//   return
+// }
+//
+//  becomes
+//
+// func @func_execute_region_elim() {
+//    "test.foo"() : () -> ()
+//    %c = "test.cmp"() : () -> i1
+//    cond_br %c, ^bb1, ^bb2
+//  ^bb1:  // pred: ^bb0
+//    %x = "test.val1"() : () -> i64
+//    br ^bb3(%x : i64)
+//  ^bb2:  // pred: ^bb0
+//    %y = "test.val2"() : () -> i64
+//    br ^bb3(%y : i64)
+//  ^bb3(%z: i64):  // 2 preds: ^bb1, ^bb2
+//    "test.bar"(%z) : (i64) -> ()
+//    return
+//  }
+//
+struct MultiBlockExecuteInliner : public OpRewritePattern<ExecuteRegionOp> {
+  using OpRewritePattern<ExecuteRegionOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(ExecuteRegionOp op,
+                                PatternRewriter &rewriter) const override {
+    if (!isa<FuncOp, ExecuteRegionOp>(op->getParentOp()))
+      return failure();
+
+    Block *prevBlock = op->getBlock();
+    Block *postBlock = rewriter.splitBlock(prevBlock, op->getIterator());
+    rewriter.setInsertionPointToEnd(prevBlock);
+
+    rewriter.create<BranchOp>(op.getLoc(), &op.region().front());
+
+    for (Block &blk : op.region()) {
+      if (YieldOp yieldOp = dyn_cast<YieldOp>(blk.getTerminator())) {
+        rewriter.setInsertionPoint(yieldOp);
+        rewriter.create<BranchOp>(yieldOp.getLoc(), postBlock,
+                                  yieldOp.results());
+        rewriter.eraseOp(yieldOp);
+      }
+    }
+
+    rewriter.inlineRegionBefore(op.region(), postBlock);
+    SmallVector<Value> blockArgs;
+
+    for (auto res : op.getResults())
+      blockArgs.push_back(postBlock->addArgument(res.getType()));
+
+    rewriter.replaceOp(op, blockArgs);
+    return success();
+  }
+};
+
 void ExecuteRegionOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                                   MLIRContext *context) {
-  results.add<SingleBlockExecuteInliner>(context);
+  results.add<SingleBlockExecuteInliner, MultiBlockExecuteInliner>(context);
 }
 
 //===----------------------------------------------------------------------===//
 // ForOp
 //===----------------------------------------------------------------------===//
 
 void ForOp::build(OpBuilder &builder, OperationState &result, Value lb,
                   Value ub, Value step, ValueRange iterArgs,

mlir/test/Dialect/SCF/canonicalize.mlir

 }
 
 // CHECK-NEXT:     affine.for %arg0 = 0 to 100 {
 // CHECK-NEXT:       "test.foo"() : () -> ()
 // CHECK-NEXT:       %[[VAL:.*]] = "test.val"() : () -> i64
 // CHECK-NEXT:       "test.bar"(%[[VAL]]) : (i64) -> ()
 // CHECK-NEXT:     }
 
+
+// -----
+
+// CHECK-LABEL: func @func_execute_region_elim
+func @func_execute_region_elim() {
+    "test.foo"() : () -> ()
+    %v = scf.execute_region -> i64 {
+      %c = "test.cmp"() : () -> i1
+      cond_br %c, ^bb2, ^bb3
+    ^bb2:
+      %x = "test.val1"() : () -> i64
+      br ^bb4(%x : i64)
+    ^bb3:
+      %y = "test.val2"() : () -> i64
+      br ^bb4(%y : i64)
+    ^bb4(%z : i64):
+      scf.yield %z : i64
+    }
+    "test.bar"(%v) : (i64) -> ()
+  return
+}
+
+// CHECK-NOT: execute_region
+// CHECK:     "test.foo"
+// CHECK:     %[[cmp:.+]] = "test.cmp"
+// CHECK:     cond_br %[[cmp]], ^[[bb1:.+]], ^[[bb2:.+]]
+// CHECK:   ^[[bb1]]:
+// CHECK:     %[[x:.+]] = "test.val1"
+// CHECK:     br ^[[bb3:.+]](%[[x]] : i64)
+// CHECK:   ^[[bb2]]:
+// CHECK:     %[[y:.+]] = "test.val2"
+// CHECK:     br ^[[bb3]](%[[y:.+]] : i64)
+// CHECK:   ^[[bb3]](%[[z:.+]]: i64):
+// CHECK:     "test.bar"(%[[z]])
+// CHECK:     return
+
+
+// -----
+
+// CHECK-LABEL: func @func_execute_region_elim_multi_yield
+func @func_execute_region_elim_multi_yield() {
+    "test.foo"() : () -> ()
+    %v = scf.execute_region -> i64 {
+      %c = "test.cmp"() : () -> i1
+      cond_br %c, ^bb2, ^bb3
+    ^bb2:
+      %x = "test.val1"() : () -> i64
+      scf.yield %x : i64
+    ^bb3:
+      %y = "test.val2"() : () -> i64
+      scf.yield %y : i64
+    }
+    "test.bar"(%v) : (i64) -> ()
+  return
+}
+
+// CHECK-NOT: execute_region
+// CHECK:     "test.foo"
+// CHECK:     %[[cmp:.+]] = "test.cmp"
+// CHECK:     cond_br %[[cmp]], ^[[bb1:.+]], ^[[bb2:.+]]
+// CHECK:   ^[[bb1]]:
+// CHECK:     %[[x:.+]] = "test.val1"
+// CHECK:     br ^[[bb3:.+]](%[[x]] : i64)
+// CHECK:   ^[[bb2]]:
+// CHECK:     %[[y:.+]] = "test.val2"
+// CHECK:     br ^[[bb3]](%[[y:.+]] : i64)
+// CHECK:   ^[[bb3]](%[[z:.+]]: i64):
+// CHECK:     "test.bar"(%[[z]])
+// CHECK:     return