[mlir] lower loop.if operations that yield a result

[gysit]

Summary

The patch extends the loop dialect lowering to support loop.if operations that yield one or more results. In particular, the patch ensures the results yielded by the then and else blocks are forwarded to the continuation block.

[ftynse]

Unfortunately, this is harder than it looks. We need to make sure all IR changes in the pattern go through rewriter and we currently don't have a way of adding new block arguments there. If you are blocked by this, feel free to temporarily loop.if lowering from the pattern rewriter infra and run it as a plain IR-mutating pass.

mlir/lib/Conversion/LoopToStandard/ConvertLoopToStandard.cpp

243

We should not be calling replaceAllUsesWith inside matchAndRewrite, it messes up with the pattern rewriter (it tracks all transformations and may roll some changes back, going outside the rewrite makes changes untrackable and ultimately leads to tricky crashes).

[gysit]

In D76350#1928810, @ftynse wrote:

Unfortunately, this is harder than it looks. We need to make sure all IR changes in the pattern go through rewriter and we currently don't have a way of adding new block arguments there. If you are blocked by this, feel free to temporarily loop.if lowering from the pattern rewriter infra and run it as a plain IR-mutating pass.

Thanks for that hint!

I guess I have multiple patterns where I use replaceAllUsesWith(). Can a rollback happen even if there is no code path that later on returns failure()? I ask since my stuff seems to work so far.

Let's postpone the patch then!

[ftynse]

In D76350#1928935, @gysit wrote:

I guess I have multiple patterns where I use replaceAllUsesWith(). Can a rollback happen even if there is no code path that later on returns failure()? I ask since my stuff seems to work so far.

The only case where it is actually necessary is replacing block arguments. Otherwise, you can do rewriter.replaceOp(op-that-defines-the-value, new-values).

Currently, we don't rollback if patterns don't fail. However, you don't necessarily have a guarantee that your patterns won't be combined with other patterns that may fail...

mlir/lib/Conversion/LoopToStandard/ConvertLoopToStandard.cpp

 // operations following the loop.if are split into a continuation (subgraph
 // exit) block. The condition is lowered to a chain of blocks that implement the
 // short-circuit scheme.  Condition blocks are created by splitting out an empty
 // block from the block that contains the loop.if operation.  They
 // conditionally branch to either the first block of the "then" region, or to
 // the first block of the "else" region.  If the latter is absent, they branch
 // to the continuation block instead.  The last blocks of "then" and "else"
 // regions (which are known to be exit blocks thanks to the invariant we
-// maintain).
+// maintain). The terminator operands of the then and else blocks are the
+// results of the loop.if operation and thus forwarded to the continuation 
+// block.
 //
 //      +--------------------------------+
 //      | <code before the IfOp>         |
 //      | cond_br %cond, %then, %else    |
 //      +--------------------------------+
 //             |              |
 //             |              --------------|
 //             v                            |
 //      +--------------------------------+  |
 //      | then:                          |  |
 //      |   <then contents>              |  |
-//      |   br continue                  |  |
+//      |   br continue(%yields...)      |  |
 //      +--------------------------------+  |
 //             |                            |
 //   |----------               |-------------
 //   |                         V
 //   |  +--------------------------------+
 //   |  | else:                          |
 //   |  |   <else contents>              |
-//   |  |   br continue                  |
+//   |  |   br continue(%yields...)      |
 //   |  +--------------------------------+
 //   |         |
 //   ------|   |
 //         v   v
 //      +--------------------------------+
-//      | continue:                      |
+//      | continue(%yields...):          |
 //      |   <code after the IfOp>        |
 //      +--------------------------------+
 //
 struct IfLowering : public OpRewritePattern<IfOp> {
   using OpRewritePattern<IfOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(IfOp ifOp,
                                 PatternRewriter &rewriter) const override;
LogicalResult IfLowering::matchAndRewrite(IfOp ifOp,
 
   // Start by splitting the block containing the 'loop.if' into two parts.
   // The part before will contain the condition, the part after will be the
   // continuation point.
   auto *condBlock = rewriter.getInsertionBlock();
   auto opPosition = rewriter.getInsertionPoint();
   auto *continueBlock = rewriter.splitBlock(condBlock, opPosition);
 
+  // Add the results of the if loop as block argument to the continuation block
+  for (auto result : ifOp.results())
+    result.replaceAllUsesWith(continueBlock->addArgument(result.getType()));
+
   // Move blocks from the "then" region to the region containing 'loop.if',
-  // place it before the continuation block, and branch to it.
+  // place it before the continuation block, and branch to it. The results of
+  // the thenBlock are taken from the terminator and passed to the branch op.
   auto &thenRegion = ifOp.thenRegion();
   auto *thenBlock = &thenRegion.front();
-  rewriter.eraseOp(thenRegion.back().getTerminator());
+  Operation *thenTerminator = thenRegion.back().getTerminator();
+  SmallVector<Value, 8> thenResults;
+  thenResults.append(thenTerminator->operand_begin(),
+                     thenTerminator->operand_end());
+  rewriter.eraseOp(thenTerminator);
   rewriter.setInsertionPointToEnd(&thenRegion.back());
-  rewriter.create<BranchOp>(loc, continueBlock);
+  rewriter.create<BranchOp>(loc, continueBlock, thenResults);
   rewriter.inlineRegionBefore(thenRegion, continueBlock);
 
   // Move blocks from the "else" region (if present) to the region containing
   // 'loop.if', place it before the continuation block and branch to it.  It
-  // will be placed after the "then" regions.
+  // will be placed after the "then" regions. The results of the elseBlock are 
+  // taken from the terminator and passed to the branch op.
   auto *elseBlock = continueBlock;
   auto &elseRegion = ifOp.elseRegion();
   if (!elseRegion.empty()) {
     elseBlock = &elseRegion.front();
-    rewriter.eraseOp(elseRegion.back().getTerminator());
+    Operation *elseTerminator = elseRegion.back().getTerminator();
+    SmallVector<Value, 8> elseResults;
+    elseResults.append(elseTerminator->operand_begin(),
+                       elseTerminator->operand_end());
+    rewriter.eraseOp(elseTerminator);
     rewriter.setInsertionPointToEnd(&elseRegion.back());
-    rewriter.create<BranchOp>(loc, continueBlock);
+    rewriter.create<BranchOp>(loc, continueBlock, elseResults);
     rewriter.inlineRegionBefore(elseRegion, continueBlock);
   }
 
   rewriter.setInsertionPointToEnd(condBlock);
   rewriter.create<CondBranchOp>(loc, ifOp.condition(), thenBlock,
                                 /*trueArgs=*/ArrayRef<Value>(), elseBlock,
                                 /*falseArgs=*/ArrayRef<Value>());
 

mlir/test/Conversion/convert-to-cfg.mlir

func @simple_std_if_else(%arg0: i1) {
   loop.if %arg0 {
     %c1 = constant 1 : index
   } else {
     %c1_0 = constant 1 : index
   }
   return
 }
 
+// CHECK-LABEL: func @simple_yield_if_else(%{{.*}}: i1) -> (index, index) {
+//  CHECK-NEXT:   cond_br %{{.*}}, ^bb1, ^bb2
+//  CHECK-NEXT:   ^bb1:   // pred: ^bb0
+//  CHECK-NEXT:     [[C1:%.*]] = constant 1 : index
+//  CHECK-NEXT:     br ^bb3([[C1]], [[C1]] : index, index) 
+//  CHECK-NEXT:   ^bb2:   // pred: ^bb0
+//  CHECK-NEXT:     [[C2:%.*]] = constant 2 : index
+//  CHECK-NEXT:     br ^bb3([[C2]], [[C2]] : index, index) 
+//  CHECK-NEXT:   ^bb3([[R1:%.*]]: index, [[R2:%.*]]: index):   // 2 preds: ^bb1, ^bb2
+//  CHECK-NEXT:     return [[R1]], [[R2]] : index, index
+func @simple_yield_if_else(%arg0 : i1) -> (index, index) {
+  %result:2 = loop.if %arg0 -> (index, index) {
+    %c1 = constant 1 : index
+    loop.yield %c1, %c1 : index, index
+  } else {
+    %c2 = constant 2 : index
+    loop.yield %c2, %c2 : index, index
+  }
+  return %result#0, %result#1 : index, index
+}
+
 // CHECK-LABEL: func @simple_std_2_ifs(%{{.*}}: i1) {
 //  CHECK-NEXT:   cond_br %{{.*}}, ^bb1, ^bb5
 //  CHECK-NEXT: ^bb1:   // pred: ^bb0
 //  CHECK-NEXT:   %{{.*}} = constant 1 : index
 //  CHECK-NEXT:   cond_br %{{.*}}, ^bb2, ^bb3
 //  CHECK-NEXT: ^bb2:   // pred: ^bb1
 //  CHECK-NEXT:   %{{.*}} = constant 1 : index
 //  CHECK-NEXT:   br ^bb4