[mlir][SCF][NFC] Add helper functions to get body of scf.while

[springerm]

Summary

Add two new helper functions getBeforeBody and getAfterBody to be consistent with "scf.for" (getBody) and to show in the API that both regions have exactly one block. Also simplify some code that assumed that there can be more than one block in a region.

[guraypp]

Nice, it simplifies some codes.

mlir/lib/Conversion/SCFToControlFlow/SCFToControlFlow.cpp

578

why do we loose notifyMatchFailure here?

mlir/lib/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.cpp

762–763

again, why do we loose asserts?

[springerm]

mlir/lib/Conversion/SCFToControlFlow/SCFToControlFlow.cpp

578

scf.while always has one block in the "after" region.

mlir/lib/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.cpp

762–763

both regions always have exactly one block

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

def WhileOp : SCF_Op<"while",
   let extraClassDeclaration = [{
     using BodyBuilderFn =
         function_ref<void(OpBuilder &, Location, ValueRange)>;
 
     ConditionOp getConditionOp();
     YieldOp getYieldOp();
     Block::BlockArgListType getBeforeArguments();
     Block::BlockArgListType getAfterArguments();
+    Block *getBeforeBody() { return &getBefore().front(); }
+    Block *getAfterBody() { return &getAfter().front(); }
   }];
 
   let hasCanonicalizer = 1;
   let hasCustomAssemblyFormat = 1;
   let hasVerifier = 1;
 }
 
 //===----------------------------------------------------------------------===//

mlir/lib/Conversion/SCFToControlFlow/SCFToControlFlow.cpp

LogicalResult WhileLowering::matchAndRewrite(WhileOp whileOp,
   Location loc = whileOp.getLoc();
 
   // Split the current block before the WhileOp to create the inlining point.
   Block *currentBlock = rewriter.getInsertionBlock();
   Block *continuation =
       rewriter.splitBlock(currentBlock, rewriter.getInsertionPoint());
 
   // Inline both regions.
-  Block *after = &whileOp.getAfter().front();
-  Block *afterLast = &whileOp.getAfter().back();
-  Block *before = &whileOp.getBefore().front();
-  Block *beforeLast = &whileOp.getBefore().back();
+  Block *after = whileOp.getAfterBody();
+  Block *before = whileOp.getBeforeBody();
   rewriter.inlineRegionBefore(whileOp.getAfter(), continuation);
   rewriter.inlineRegionBefore(whileOp.getBefore(), after);
 
   // Branch to the "before" region.
   rewriter.setInsertionPointToEnd(currentBlock);
   rewriter.create<cf::BranchOp>(loc, before, whileOp.getInits());
 
   // Replace terminators with branches. Assuming bodies are SESE, which holds
   // given only the patterns from this file, we only need to look at the last
   // block. This should be reconsidered if we allow break/continue in SCF.
-  rewriter.setInsertionPointToEnd(beforeLast);
-  auto condOp = cast<ConditionOp>(beforeLast->getTerminator());
+  rewriter.setInsertionPointToEnd(before);
+  auto condOp = cast<ConditionOp>(before->getTerminator());
   rewriter.replaceOpWithNewOp<cf::CondBranchOp>(condOp, condOp.getCondition(),
                                                 after, condOp.getArgs(),
                                                 continuation, ValueRange());
 
-  rewriter.setInsertionPointToEnd(afterLast);
-  auto yieldOp = cast<scf::YieldOp>(afterLast->getTerminator());
+  rewriter.setInsertionPointToEnd(after);
+  auto yieldOp = cast<scf::YieldOp>(after->getTerminator());
   rewriter.replaceOpWithNewOp<cf::BranchOp>(yieldOp, before,
                                             yieldOp.getResults());
 
   // Replace the op with values "yielded" from the "before" region, which are
   // visible by dominance.
   rewriter.replaceOp(whileOp, condOp.getArgs());
 
   return success();
 }
 
 LogicalResult
 DoWhileLowering::matchAndRewrite(WhileOp whileOp,
                                  PatternRewriter &rewriter) const {
-  if (!llvm::hasSingleElement(whileOp.getAfter()))
-    return rewriter.notifyMatchFailure(whileOp,
-                                       "do-while simplification applicable to "
-                                       "single-block 'after' region only");
-
-  Block &afterBlock = whileOp.getAfter().front();
+  Block &afterBlock = *whileOp.getAfterBody();
   if (!llvm::hasSingleElement(afterBlock))
     return rewriter.notifyMatchFailure(whileOp,
                                        "do-while simplification applicable "
                                        "only if 'after' region has no payload");
 
   auto yield = dyn_cast<scf::YieldOp>(&afterBlock.front());
   if (!yield || yield.getResults() != afterBlock.getArguments())
     return rewriter.notifyMatchFailure(whileOp,
                                        "do-while simplification applicable "
                                        "only to forwarding 'after' regions");
 
   // Split the current block before the WhileOp to create the inlining point.
   OpBuilder::InsertionGuard guard(rewriter);
   Block *currentBlock = rewriter.getInsertionBlock();
   Block *continuation =
       rewriter.splitBlock(currentBlock, rewriter.getInsertionPoint());
 
   // Only the "before" region should be inlined.
-  Block *before = &whileOp.getBefore().front();
-  Block *beforeLast = &whileOp.getBefore().back();
+  Block *before = whileOp.getBeforeBody();
   rewriter.inlineRegionBefore(whileOp.getBefore(), continuation);
 
   // Branch to the "before" region.
   rewriter.setInsertionPointToEnd(currentBlock);
   rewriter.create<cf::BranchOp>(whileOp.getLoc(), before, whileOp.getInits());
 
   // Loop around the "before" region based on condition.
-  rewriter.setInsertionPointToEnd(beforeLast);
-  auto condOp = cast<ConditionOp>(beforeLast->getTerminator());
+  rewriter.setInsertionPointToEnd(before);
+  auto condOp = cast<ConditionOp>(before->getTerminator());
   rewriter.replaceOpWithNewOp<cf::CondBranchOp>(condOp, condOp.getCondition(),
                                                 before, condOp.getArgs(),
                                                 continuation, ValueRange());
 
   // Replace the op with values "yielded" from the "before" region, which are
   // visible by dominance.
   rewriter.replaceOp(whileOp, condOp.getArgs());
 

mlir/lib/Dialect/SCF/IR/SCF.cpp

 OperandRange WhileOp::getEntrySuccessorOperands(std::optional<unsigned> index) {
   assert(index && *index == 0 &&
          "WhileOp is expected to branch only to the first region");
 
   return getInits();
 }
 
 ConditionOp WhileOp::getConditionOp() {
-  return cast<ConditionOp>(getBefore().front().getTerminator());
+  return cast<ConditionOp>(getBeforeBody()->getTerminator());
 }
 
 YieldOp WhileOp::getYieldOp() {
-  return cast<YieldOp>(getAfter().front().getTerminator());
+  return cast<YieldOp>(getAfterBody()->getTerminator());
 }
 
 Block::BlockArgListType WhileOp::getBeforeArguments() {
-  return getBefore().front().getArguments();
+  return getBeforeBody()->getArguments();
 }
 
 Block::BlockArgListType WhileOp::getAfterArguments() {
-  return getAfter().front().getArguments();
+  return getAfterBody()->getArguments();
 }
 
 void WhileOp::getSuccessorRegions(std::optional<unsigned> index,
                                   SmallVectorImpl<RegionSuccessor> &regions) {
   // The parent op always branches to the condition region.
   if (!index) {
     regions.emplace_back(&getBefore(), getBefore().getArguments());
     return;
ParseResult scf::WhileOp::parse(OpAsmParser &parser, OperationState &result) {
 
   return failure(parser.parseRegion(*before, regionArgs) ||
                  parser.parseKeyword("do") || parser.parseRegion(*after) ||
                  parser.parseOptionalAttrDictWithKeyword(result.attributes));
 }
 
 /// Prints a `while` op.
 void scf::WhileOp::print(OpAsmPrinter &p) {
-  printInitializationList(p, getBefore().front().getArguments(), getInits(),
-                          " ");
+  printInitializationList(p, getBeforeArguments(), getInits(), " ");
   p << " : ";
   p.printFunctionalType(getInits().getTypes(), getResults().getTypes());
   p << ' ';
   p.printRegion(getBefore(), /*printEntryBlockArgs=*/false);
   p << " do ";
   p.printRegion(getAfter());
   p.printOptionalAttrDictWithKeyword((*this)->getAttrs());
 }

 ///           block argument and replace all uses of 1-th before block argument
 ///           with %b.
 struct RemoveLoopInvariantArgsFromBeforeBlock
     : public OpRewritePattern<WhileOp> {
   using OpRewritePattern<WhileOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(WhileOp op,
                                 PatternRewriter &rewriter) const override {
-    Block &afterBlock = op.getAfter().front();
+    Block &afterBlock = *op.getAfterBody();
     Block::BlockArgListType beforeBlockArgs = op.getBeforeArguments();
     ConditionOp condOp = op.getConditionOp();
     OperandRange condOpArgs = condOp.getArgs();
     Operation *yieldOp = afterBlock.getTerminator();
     ValueRange yieldOpArgs = yieldOp->getOperands();
 
     bool canSimplify = false;
     for (const auto &it :
LogicalResult matchAndRewrite(WhileOp op,
 
     auto newWhile =
         rewriter.create<WhileOp>(op.getLoc(), op.getResultTypes(), newInitArgs);
 
     Block &newBeforeBlock = *rewriter.createBlock(
         &newWhile.getBefore(), /*insertPt*/ {},
         ValueRange(newYieldOpArgs).getTypes(), newBeforeBlockArgLocs);
 
-    Block &beforeBlock = op.getBefore().front();
+    Block &beforeBlock = *op.getBeforeBody();
     SmallVector<Value> newBeforeBlockArgs(beforeBlock.getNumArguments());
     // For each i-th before block argument we find it's replacement value as :-
     //   1. If i-th before block argument is a loop invariant, we fetch it's
     //      initial value from `beforeBlockInitValMap` by querying for key `i`.
     //   2. Else we fetch j-th new before block argument as the replacement
     //      value of i-th before block argument.
     for (unsigned i = 0, j = 0, n = beforeBlock.getNumArguments(); i < n; i++) {
       // If the index 'i' argument was a loop invariant we fetch it's initial

 ///          replaced by %b.
 ///       3. The corresponding after block argument %arg1_after's uses are
 ///          replaced by %a and %arg2_after's uses are replaced by %b.
 struct RemoveLoopInvariantValueYielded : public OpRewritePattern<WhileOp> {
   using OpRewritePattern<WhileOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(WhileOp op,
                                 PatternRewriter &rewriter) const override {
-    Block &beforeBlock = op.getBefore().front();
+    Block &beforeBlock = *op.getBeforeBody();
     ConditionOp condOp = op.getConditionOp();
     OperandRange condOpArgs = condOp.getArgs();
 
     bool canSimplify = false;
     for (Value condOpArg : condOpArgs) {
       // Those values not defined within `before` block will be considered as
       // loop invariant values. We map the corresponding `index` with their
       // value.
LogicalResult matchAndRewrite(WhileOp op,
 
     auto newWhile = rewriter.create<WhileOp>(op.getLoc(), newAfterBlockType,
                                              op.getOperands());
 
     Block &newAfterBlock =
         *rewriter.createBlock(&newWhile.getAfter(), /*insertPt*/ {},
                               newAfterBlockType, newAfterBlockArgLocs);
 
-    Block &afterBlock = op.getAfter().front();
+    Block &afterBlock = *op.getAfterBody();
     // Since a new scf.condition op was created, we need to fetch the new
     // `after` block arguments which will be used while replacing operations of
     // previous scf.while's `after` blocks. We'd also be fetching new result
     // values too.
     SmallVector<Value> newAfterBlockArgs(afterBlock.getNumArguments());
     SmallVector<Value> newWhileResults(afterBlock.getNumArguments());
     for (unsigned i = 0, j = 0, n = afterBlock.getNumArguments(); i < n; i++) {
       Value afterBlockArg, result;
LogicalResult matchAndRewrite(WhileOp op,
     for (const auto &it : llvm::enumerate(newResultsIndices)) {
       newResults[it.value()] = newWhile.getResult(it.index());
       newAfterBlockArgs[it.value()] = newAfterBlock.getArgument(it.index());
     }
 
     rewriter.inlineRegionBefore(op.getBefore(), newWhile.getBefore(),
                                 newWhile.getBefore().begin());
 
-    Block &afterBlock = op.getAfter().front();
+    Block &afterBlock = *op.getAfterBody();
     rewriter.mergeBlocks(&afterBlock, &newAfterBlock, newAfterBlockArgs);
 
     rewriter.replaceOp(op, newResults);
     return success();
   }
 };
 
 /// Replace operations equivalent to the condition in the do block with true,
struct WhileCmpCond : public OpRewritePattern<scf::WhileOp> {
   LogicalResult matchAndRewrite(scf::WhileOp op,
                                 PatternRewriter &rewriter) const override {
     using namespace scf;
     auto cond = op.getConditionOp();
     auto cmp = cond.getCondition().getDefiningOp<arith::CmpIOp>();
     if (!cmp)
       return failure();
     bool changed = false;
-    for (auto tup :
-         llvm::zip(cond.getArgs(), op.getAfter().front().getArguments())) {
+    for (auto tup : llvm::zip(cond.getArgs(), op.getAfterArguments())) {
       for (size_t opIdx = 0; opIdx < 2; opIdx++) {
         if (std::get<0>(tup) != cmp.getOperand(opIdx))
           continue;
         for (OpOperand &u :
              llvm::make_early_inc_range(std::get<1>(tup).getUses())) {
           auto cmp2 = dyn_cast<arith::CmpIOp>(u.getOwner());
           if (!cmp2)
             continue;
for (auto &&[beforeArg, yieldValue, initValue] : llvm::zip(
         argsToErase.push_back(true);
       } else {
         argsToErase.push_back(false);
         newYields.emplace_back(yieldValue);
         newInits.emplace_back(initValue);
       }
     }
 
-    Block &beforeBlock = op.getBefore().front();
-    Block &afterBlock = op.getAfter().front();
+    Block &beforeBlock = *op.getBeforeBody();
+    Block &afterBlock = *op.getAfterBody();
 
     beforeBlock.eraseArguments(argsToErase);
 
     Location loc = op.getLoc();
     auto newWhileOp =
         rewriter.create<WhileOp>(loc, op.getResultTypes(), newInits,
                                  /*beforeBody*/ nullptr, /*afterBody*/ nullptr);
-    Block &newBeforeBlock = newWhileOp.getBefore().front();
-    Block &newAfterBlock = newWhileOp.getAfter().front();
+    Block &newBeforeBlock = *newWhileOp.getBeforeBody();
+    Block &newAfterBlock = *newWhileOp.getAfterBody();
 
     OpBuilder::InsertionGuard g(rewriter);
     rewriter.setInsertionPoint(yield);
     rewriter.replaceOpWithNewOp<YieldOp>(yield, newYields);
 
     rewriter.mergeBlocks(&beforeBlock, &newBeforeBlock,
                          newBeforeBlock.getArguments());
     rewriter.mergeBlocks(&afterBlock, &newAfterBlock,
LogicalResult matchAndRewrite(WhileOp op,
     }
 
     ValueRange argsRange(newArgs);
 
     Location loc = op.getLoc();
     auto newWhileOp = rewriter.create<scf::WhileOp>(
         loc, argsRange.getTypes(), op.getInits(), /*beforeBody*/ nullptr,
         /*afterBody*/ nullptr);
-    Block &newBeforeBlock = newWhileOp.getBefore().front();
-    Block &newAfterBlock = newWhileOp.getAfter().front();
+    Block &newBeforeBlock = *newWhileOp.getBeforeBody();
+    Block &newAfterBlock = *newWhileOp.getAfterBody();
 
     SmallVector<Value> afterArgsMapping;
     SmallVector<Value> resultsMapping;
     for (auto &&[i, arg] : llvm::enumerate(condOpArgs)) {
       auto it = argsMap.find(arg);
       assert(it != argsMap.end());
       auto pos = it->second;
       afterArgsMapping.emplace_back(newAfterBlock.getArgument(pos));
       resultsMapping.emplace_back(newWhileOp->getResult(pos));
     }
 
     OpBuilder::InsertionGuard g(rewriter);
     rewriter.setInsertionPoint(condOp);
     rewriter.replaceOpWithNewOp<ConditionOp>(condOp, condOp.getCondition(),
                                              argsRange);
 
-    Block &beforeBlock = op.getBefore().front();
-    Block &afterBlock = op.getAfter().front();
+    Block &beforeBlock = *op.getBeforeBody();
+    Block &afterBlock = *op.getAfterBody();
 
     rewriter.mergeBlocks(&beforeBlock, &newBeforeBlock,
                          newBeforeBlock.getArguments());
     rewriter.mergeBlocks(&afterBlock, &newAfterBlock, afterArgsMapping);
     rewriter.replaceOp(op, resultsMapping);
     return success();
   }
 };

mlir/lib/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.cpp

LogicalResult resolveConflicts(Operation *op, RewriterBase &rewriter,
     });
 
     return success();
   }
 
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
                           const BufferizationOptions &options) const {
     auto whileOp = cast<scf::WhileOp>(op);
 
-    assert(whileOp.getBefore().getBlocks().size() == 1 &&
-           "regions with multiple blocks not supported");
-    Block *beforeBody = &whileOp.getBefore().front();
-    assert(whileOp.getAfter().getBlocks().size() == 1 &&
-           "regions with multiple blocks not supported");
-    Block *afterBody = &whileOp.getAfter().front();
-
     // Indices of all bbArgs that have tensor type. These are the ones that
     // are bufferized. The "before" and "after" regions may have different args.
     DenseSet<int64_t> indicesBefore = getTensorIndices(whileOp.getInits());
     DenseSet<int64_t> indicesAfter =
         getTensorIndices(whileOp.getAfterArguments());
 
     // The new memref init_args of the loop.
     FailureOr<SmallVector<Value>> maybeInitArgs =
         getBuffers(rewriter, whileOp->getOpOperands(), options);
LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
     newWhileOp.getAfter().addArguments(argsTypesAfter, bbArgLocsAfter);
 
     // Set up new iter_args and move the loop condition block to the new op.
     // The old block uses tensors, so wrap the (memref) bbArgs of the new block
     // in ToTensorOps.
     rewriter.setInsertionPointToStart(newBeforeBody);
     SmallVector<Value> newBeforeArgs = getBbArgReplacements(
         rewriter, newWhileOp.getBeforeArguments(), indicesBefore);
-    rewriter.mergeBlocks(beforeBody, newBeforeBody, newBeforeArgs);
+    rewriter.mergeBlocks(whileOp.getBeforeBody(), newBeforeBody, newBeforeArgs);
 
     // Set up new iter_args and move the loop body block to the new op.
     // The old block uses tensors, so wrap the (memref) bbArgs of the new block
     // in ToTensorOps.
     rewriter.setInsertionPointToStart(newAfterBody);
     SmallVector<Value> newAfterArgs = getBbArgReplacements(
         rewriter, newWhileOp.getAfterArguments(), indicesAfter);
-    rewriter.mergeBlocks(afterBody, newAfterBody, newAfterArgs);
+    rewriter.mergeBlocks(whileOp.getAfterBody(), newAfterBody, newAfterArgs);
 
     // Replace loop results.
     replaceOpWithBufferizedValues(rewriter, op, newWhileOp->getResults());
 
     return success();
   }
 
   FailureOr<BaseMemRefType>

mlir/lib/Dialect/SCF/Transforms/ForToWhile.cpp

LogicalResult matchAndRewrite(ForOp forOp,
     llvm::append_range(initArgs, forOp.getInitArgs());
     auto whileOp = rewriter.create<WhileOp>(forOp.getLoc(), lcvTypes, initArgs,
                                             forOp->getAttrs());
 
     // 'before' region contains the loop condition and forwarding of iteration
     // arguments to the 'after' region.
     auto *beforeBlock = rewriter.createBlock(
         &whileOp.getBefore(), whileOp.getBefore().begin(), lcvTypes, lcvLocs);
-    rewriter.setInsertionPointToStart(&whileOp.getBefore().front());
+    rewriter.setInsertionPointToStart(whileOp.getBeforeBody());
     auto cmpOp = rewriter.create<arith::CmpIOp>(
         whileOp.getLoc(), arith::CmpIPredicate::slt,
         beforeBlock->getArgument(0), forOp.getUpperBound());
     rewriter.create<scf::ConditionOp>(whileOp.getLoc(), cmpOp.getResult(),
                                       beforeBlock->getArguments());
 
     // Inline for-loop body into an executeRegion operation in the "after"
     // region. The return type of the execRegionOp does not contain the