brintos / llvm-project-archived public Read only

1//===- SCF.cpp - Structured Control Flow Operations -----------------------===//2//3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.4// See https://llvm.org/LICENSE.txt for license information.5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception6//7//===----------------------------------------------------------------------===//8 9#include "mlir/Dialect/SCF/IR/SCF.h"10#include "mlir/Conversion/ConvertToEmitC/ToEmitCInterface.h"11#include "mlir/Dialect/Arith/IR/Arith.h"12#include "mlir/Dialect/Arith/Utils/Utils.h"13#include "mlir/Dialect/Bufferization/IR/BufferDeallocationOpInterface.h"14#include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"15#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h"16#include "mlir/Dialect/MemRef/IR/MemRef.h"17#include "mlir/Dialect/SCF/IR/DeviceMappingInterface.h"18#include "mlir/Dialect/Tensor/IR/Tensor.h"19#include "mlir/IR/BuiltinAttributes.h"20#include "mlir/IR/IRMapping.h"21#include "mlir/IR/Matchers.h"22#include "mlir/IR/Operation.h"23#include "mlir/IR/OperationSupport.h"24#include "mlir/IR/PatternMatch.h"25#include "mlir/Interfaces/FunctionInterfaces.h"26#include "mlir/Interfaces/ParallelCombiningOpInterface.h"27#include "mlir/Interfaces/ValueBoundsOpInterface.h"28#include "mlir/Transforms/InliningUtils.h"29#include "mlir/Transforms/RegionUtils.h"30#include "llvm/ADT/MapVector.h"31#include "llvm/ADT/STLExtras.h"32#include "llvm/ADT/SmallPtrSet.h"33#include "llvm/Support/Casting.h"34#include "llvm/Support/DebugLog.h"35#include <optional>36 37using namespace mlir;38using namespace mlir::scf;39 40#include "mlir/Dialect/SCF/IR/SCFOpsDialect.cpp.inc"41 42//===----------------------------------------------------------------------===//43// SCFDialect Dialect Interfaces44//===----------------------------------------------------------------------===//45 46namespace {47struct SCFInlinerInterface : public DialectInlinerInterface {48  using DialectInlinerInterface::DialectInlinerInterface;49  // We don't have any special restrictions on what can be inlined into50  // destination regions (e.g. while/conditional bodies). Always allow it.51  bool isLegalToInline(Region *dest, Region *src, bool wouldBeCloned,52                       IRMapping &valueMapping) const final {53    return true;54  }55  // Operations in scf dialect are always legal to inline since they are56  // pure.57  bool isLegalToInline(Operation *, Region *, bool, IRMapping &) const final {58    return true;59  }60  // Handle the given inlined terminator by replacing it with a new operation61  // as necessary. Required when the region has only one block.62  void handleTerminator(Operation *op, ValueRange valuesToRepl) const final {63    auto retValOp = dyn_cast<scf::YieldOp>(op);64    if (!retValOp)65      return;66 67    for (auto retValue : llvm::zip(valuesToRepl, retValOp.getOperands())) {68      std::get<0>(retValue).replaceAllUsesWith(std::get<1>(retValue));69    }70  }71};72} // namespace73 74//===----------------------------------------------------------------------===//75// SCFDialect76//===----------------------------------------------------------------------===//77 78void SCFDialect::initialize() {79  addOperations<80#define GET_OP_LIST81#include "mlir/Dialect/SCF/IR/SCFOps.cpp.inc"82      >();83  addInterfaces<SCFInlinerInterface>();84  declarePromisedInterface<ConvertToEmitCPatternInterface, SCFDialect>();85  declarePromisedInterfaces<bufferization::BufferDeallocationOpInterface,86                            InParallelOp, ReduceReturnOp>();87  declarePromisedInterfaces<bufferization::BufferizableOpInterface, ConditionOp,88                            ExecuteRegionOp, ForOp, IfOp, IndexSwitchOp,89                            ForallOp, InParallelOp, WhileOp, YieldOp>();90  declarePromisedInterface<ValueBoundsOpInterface, ForOp>();91}92 93/// Default callback for IfOp builders. Inserts a yield without arguments.94void mlir::scf::buildTerminatedBody(OpBuilder &builder, Location loc) {95  scf::YieldOp::create(builder, loc);96}97 98/// Verifies that the first block of the given `region` is terminated by a99/// TerminatorTy. Reports errors on the given operation if it is not the case.100template <typename TerminatorTy>101static TerminatorTy verifyAndGetTerminator(Operation *op, Region &region,102                                           StringRef errorMessage) {103  Operation *terminatorOperation = nullptr;104  if (!region.empty() && !region.front().empty()) {105    terminatorOperation = &region.front().back();106    if (auto yield = dyn_cast_or_null<TerminatorTy>(terminatorOperation))107      return yield;108  }109  auto diag = op->emitOpError(errorMessage);110  if (terminatorOperation)111    diag.attachNote(terminatorOperation->getLoc()) << "terminator here";112  return nullptr;113}114 115std::optional<llvm::APSInt> mlir::scf::computeUbMinusLb(Value lb, Value ub,116                                                        bool isSigned) {117  llvm::APSInt diff;118  auto addOp = ub.getDefiningOp<arith::AddIOp>();119  if (!addOp)120    return std::nullopt;121  if ((isSigned && !addOp.hasNoSignedWrap()) ||122      (!isSigned && !addOp.hasNoUnsignedWrap()))123    return std::nullopt;124 125  if (addOp.getLhs() != lb ||126      !matchPattern(addOp.getRhs(), m_ConstantInt(&diff)))127    return std::nullopt;128  return diff;129}130 131//===----------------------------------------------------------------------===//132// ExecuteRegionOp133//===----------------------------------------------------------------------===//134 135/// Replaces the given op with the contents of the given single-block region,136/// using the operands of the block terminator to replace operation results.137static void replaceOpWithRegion(PatternRewriter &rewriter, Operation *op,138                                Region &region, ValueRange blockArgs = {}) {139  assert(region.hasOneBlock() && "expected single-block region");140  Block *block = &region.front();141  Operation *terminator = block->getTerminator();142  ValueRange results = terminator->getOperands();143  rewriter.inlineBlockBefore(block, op, blockArgs);144  rewriter.replaceOp(op, results);145  rewriter.eraseOp(terminator);146}147 148///149/// (ssa-id `=`)? `execute_region` `->` function-result-type `{`150///    block+151/// `}`152///153/// Example:154///   scf.execute_region -> i32 {155///     %idx = load %rI[%i] : memref<128xi32>156///     return %idx : i32157///   }158///159ParseResult ExecuteRegionOp::parse(OpAsmParser &parser,160                                   OperationState &result) {161  if (parser.parseOptionalArrowTypeList(result.types))162    return failure();163 164  if (succeeded(parser.parseOptionalKeyword("no_inline")))165    result.addAttribute("no_inline", parser.getBuilder().getUnitAttr());166 167  // Introduce the body region and parse it.168  Region *body = result.addRegion();169  if (parser.parseRegion(*body, /*arguments=*/{}, /*argTypes=*/{}) ||170      parser.parseOptionalAttrDict(result.attributes))171    return failure();172 173  return success();174}175 176void ExecuteRegionOp::print(OpAsmPrinter &p) {177  p.printOptionalArrowTypeList(getResultTypes());178  p << ' ';179  if (getNoInline())180    p << "no_inline ";181  p.printRegion(getRegion(),182                /*printEntryBlockArgs=*/false,183                /*printBlockTerminators=*/true);184  p.printOptionalAttrDict((*this)->getAttrs(), /*elidedAttrs=*/{"no_inline"});185}186 187LogicalResult ExecuteRegionOp::verify() {188  if (getRegion().empty())189    return emitOpError("region needs to have at least one block");190  if (getRegion().front().getNumArguments() > 0)191    return emitOpError("region cannot have any arguments");192  return success();193}194 195// Inline an ExecuteRegionOp if it only contains one block.196//     "test.foo"() : () -> ()197//      %v = scf.execute_region -> i64 {198//        %x = "test.val"() : () -> i64199//        scf.yield %x : i64200//      }201//      "test.bar"(%v) : (i64) -> ()202//203//  becomes204//205//     "test.foo"() : () -> ()206//     %x = "test.val"() : () -> i64207//     "test.bar"(%x) : (i64) -> ()208//209struct SingleBlockExecuteInliner : public OpRewritePattern<ExecuteRegionOp> {210  using OpRewritePattern<ExecuteRegionOp>::OpRewritePattern;211 212  LogicalResult matchAndRewrite(ExecuteRegionOp op,213                                PatternRewriter &rewriter) const override {214    if (!op.getRegion().hasOneBlock() || op.getNoInline())215      return failure();216    replaceOpWithRegion(rewriter, op, op.getRegion());217    return success();218  }219};220 221// Inline an ExecuteRegionOp if its parent can contain multiple blocks.222// TODO generalize the conditions for operations which can be inlined into.223// func @func_execute_region_elim() {224//     "test.foo"() : () -> ()225//     %v = scf.execute_region -> i64 {226//       %c = "test.cmp"() : () -> i1227//       cf.cond_br %c, ^bb2, ^bb3228//     ^bb2:229//       %x = "test.val1"() : () -> i64230//       cf.br ^bb4(%x : i64)231//     ^bb3:232//       %y = "test.val2"() : () -> i64233//       cf.br ^bb4(%y : i64)234//     ^bb4(%z : i64):235//       scf.yield %z : i64236//     }237//     "test.bar"(%v) : (i64) -> ()238//   return239// }240//241//  becomes242//243// func @func_execute_region_elim() {244//    "test.foo"() : () -> ()245//    %c = "test.cmp"() : () -> i1246//    cf.cond_br %c, ^bb1, ^bb2247//  ^bb1:  // pred: ^bb0248//    %x = "test.val1"() : () -> i64249//    cf.br ^bb3(%x : i64)250//  ^bb2:  // pred: ^bb0251//    %y = "test.val2"() : () -> i64252//    cf.br ^bb3(%y : i64)253//  ^bb3(%z: i64):  // 2 preds: ^bb1, ^bb2254//    "test.bar"(%z) : (i64) -> ()255//    return256//  }257//258struct MultiBlockExecuteInliner : public OpRewritePattern<ExecuteRegionOp> {259  using OpRewritePattern<ExecuteRegionOp>::OpRewritePattern;260 261  LogicalResult matchAndRewrite(ExecuteRegionOp op,262                                PatternRewriter &rewriter) const override {263    if (op.getNoInline())264      return failure();265    if (!isa<FunctionOpInterface, ExecuteRegionOp>(op->getParentOp()))266      return failure();267 268    Block *prevBlock = op->getBlock();269    Block *postBlock = rewriter.splitBlock(prevBlock, op->getIterator());270    rewriter.setInsertionPointToEnd(prevBlock);271 272    cf::BranchOp::create(rewriter, op.getLoc(), &op.getRegion().front());273 274    for (Block &blk : op.getRegion()) {275      if (YieldOp yieldOp = dyn_cast<YieldOp>(blk.getTerminator())) {276        rewriter.setInsertionPoint(yieldOp);277        cf::BranchOp::create(rewriter, yieldOp.getLoc(), postBlock,278                             yieldOp.getResults());279        rewriter.eraseOp(yieldOp);280      }281    }282 283    rewriter.inlineRegionBefore(op.getRegion(), postBlock);284    SmallVector<Value> blockArgs;285 286    for (auto res : op.getResults())287      blockArgs.push_back(postBlock->addArgument(res.getType(), res.getLoc()));288 289    rewriter.replaceOp(op, blockArgs);290    return success();291  }292};293 294// Pattern to eliminate ExecuteRegionOp results which forward external295// values from the region. In case there are multiple yield operations,296// all of them must have the same operands in order for the pattern to be297// applicable.298struct ExecuteRegionForwardingEliminator299    : public OpRewritePattern<ExecuteRegionOp> {300  using OpRewritePattern<ExecuteRegionOp>::OpRewritePattern;301 302  LogicalResult matchAndRewrite(ExecuteRegionOp op,303                                PatternRewriter &rewriter) const override {304    if (op.getNumResults() == 0)305      return failure();306 307    SmallVector<Operation *> yieldOps;308    for (Block &block : op.getRegion()) {309      if (auto yield = dyn_cast<scf::YieldOp>(block.getTerminator()))310        yieldOps.push_back(yield.getOperation());311    }312 313    if (yieldOps.empty())314      return failure();315 316    // Check if all yield operations have the same operands.317    auto yieldOpsOperands = yieldOps[0]->getOperands();318    for (auto *yieldOp : yieldOps) {319      if (yieldOp->getOperands() != yieldOpsOperands)320        return failure();321    }322 323    SmallVector<Value> externalValues;324    SmallVector<Value> internalValues;325    SmallVector<Value> opResultsToReplaceWithExternalValues;326    SmallVector<Value> opResultsToKeep;327    for (auto [index, yieldedValue] : llvm::enumerate(yieldOpsOperands)) {328      if (isValueFromInsideRegion(yieldedValue, op)) {329        internalValues.push_back(yieldedValue);330        opResultsToKeep.push_back(op.getResult(index));331      } else {332        externalValues.push_back(yieldedValue);333        opResultsToReplaceWithExternalValues.push_back(op.getResult(index));334      }335    }336    // No yielded external values - nothing to do.337    if (externalValues.empty())338      return failure();339 340    // There are yielded external values - create a new execute_region returning341    // just the internal values.342    SmallVector<Type> resultTypes;343    for (Value value : internalValues)344      resultTypes.push_back(value.getType());345    auto newOp =346        ExecuteRegionOp::create(rewriter, op.getLoc(), TypeRange(resultTypes));347    newOp->setAttrs(op->getAttrs());348 349    // Move old op's region to the new operation.350    rewriter.inlineRegionBefore(op.getRegion(), newOp.getRegion(),351                                newOp.getRegion().end());352 353    // Replace all yield operations with a new yield operation with updated354    // results. scf.execute_region must have at least one yield operation.355    for (auto *yieldOp : yieldOps) {356      rewriter.setInsertionPoint(yieldOp);357      rewriter.replaceOpWithNewOp<scf::YieldOp>(yieldOp,358                                                ValueRange(internalValues));359    }360 361    // Replace the old operation with the external values directly.362    rewriter.replaceAllUsesWith(opResultsToReplaceWithExternalValues,363                                externalValues);364    // Replace the old operation's remaining results with the new operation's365    // results.366    rewriter.replaceAllUsesWith(opResultsToKeep, newOp.getResults());367    rewriter.eraseOp(op);368    return success();369  }370 371private:372  bool isValueFromInsideRegion(Value value,373                               ExecuteRegionOp executeRegionOp) const {374    // Check if the value is defined within the execute_region375    if (Operation *defOp = value.getDefiningOp())376      return &executeRegionOp.getRegion() == defOp->getParentRegion();377 378    // If it's a block argument, check if it's from within the region379    if (BlockArgument blockArg = dyn_cast<BlockArgument>(value))380      return &executeRegionOp.getRegion() == blockArg.getParentRegion();381 382    return false; // Value is from outside the region383  }384};385 386void ExecuteRegionOp::getCanonicalizationPatterns(RewritePatternSet &results,387                                                  MLIRContext *context) {388  results.add<SingleBlockExecuteInliner, MultiBlockExecuteInliner,389              ExecuteRegionForwardingEliminator>(context);390}391 392void ExecuteRegionOp::getSuccessorRegions(393    RegionBranchPoint point, SmallVectorImpl<RegionSuccessor> &regions) {394  // If the predecessor is the ExecuteRegionOp, branch into the body.395  if (point.isParent()) {396    regions.push_back(RegionSuccessor(&getRegion()));397    return;398  }399 400  // Otherwise, the region branches back to the parent operation.401  regions.push_back(RegionSuccessor(getOperation(), getResults()));402}403 404//===----------------------------------------------------------------------===//405// ConditionOp406//===----------------------------------------------------------------------===//407 408MutableOperandRange409ConditionOp::getMutableSuccessorOperands(RegionSuccessor point) {410  assert(411      (point.isParent() || point.getSuccessor() == &getParentOp().getAfter()) &&412      "condition op can only exit the loop or branch to the after"413      "region");414  // Pass all operands except the condition to the successor region.415  return getArgsMutable();416}417 418void ConditionOp::getSuccessorRegions(419    ArrayRef<Attribute> operands, SmallVectorImpl<RegionSuccessor> &regions) {420  FoldAdaptor adaptor(operands, *this);421 422  WhileOp whileOp = getParentOp();423 424  // Condition can either lead to the after region or back to the parent op425  // depending on whether the condition is true or not.426  auto boolAttr = dyn_cast_or_null<BoolAttr>(adaptor.getCondition());427  if (!boolAttr || boolAttr.getValue())428    regions.emplace_back(&whileOp.getAfter(),429                         whileOp.getAfter().getArguments());430  if (!boolAttr || !boolAttr.getValue())431    regions.emplace_back(whileOp.getOperation(), whileOp.getResults());432}433 434//===----------------------------------------------------------------------===//435// ForOp436//===----------------------------------------------------------------------===//437 438void ForOp::build(OpBuilder &builder, OperationState &result, Value lb,439                  Value ub, Value step, ValueRange initArgs,440                  BodyBuilderFn bodyBuilder, bool unsignedCmp) {441  OpBuilder::InsertionGuard guard(builder);442 443  if (unsignedCmp)444    result.addAttribute(getUnsignedCmpAttrName(result.name),445                        builder.getUnitAttr());446  result.addOperands({lb, ub, step});447  result.addOperands(initArgs);448  for (Value v : initArgs)449    result.addTypes(v.getType());450  Type t = lb.getType();451  Region *bodyRegion = result.addRegion();452  Block *bodyBlock = builder.createBlock(bodyRegion);453  bodyBlock->addArgument(t, result.location);454  for (Value v : initArgs)455    bodyBlock->addArgument(v.getType(), v.getLoc());456 457  // Create the default terminator if the builder is not provided and if the458  // iteration arguments are not provided. Otherwise, leave this to the caller459  // because we don't know which values to return from the loop.460  if (initArgs.empty() && !bodyBuilder) {461    ForOp::ensureTerminator(*bodyRegion, builder, result.location);462  } else if (bodyBuilder) {463    OpBuilder::InsertionGuard guard(builder);464    builder.setInsertionPointToStart(bodyBlock);465    bodyBuilder(builder, result.location, bodyBlock->getArgument(0),466                bodyBlock->getArguments().drop_front());467  }468}469 470LogicalResult ForOp::verify() {471  // Check that the number of init args and op results is the same.472  if (getInitArgs().size() != getNumResults())473    return emitOpError(474        "mismatch in number of loop-carried values and defined values");475 476  return success();477}478 479LogicalResult ForOp::verifyRegions() {480  // Check that the body defines as single block argument for the induction481  // variable.482  if (getInductionVar().getType() != getLowerBound().getType())483    return emitOpError(484        "expected induction variable to be same type as bounds and step");485 486  if (getNumRegionIterArgs() != getNumResults())487    return emitOpError(488        "mismatch in number of basic block args and defined values");489 490  auto initArgs = getInitArgs();491  auto iterArgs = getRegionIterArgs();492  auto opResults = getResults();493  unsigned i = 0;494  for (auto e : llvm::zip(initArgs, iterArgs, opResults)) {495    if (std::get<0>(e).getType() != std::get<2>(e).getType())496      return emitOpError() << "types mismatch between " << i497                           << "th iter operand and defined value";498    if (std::get<1>(e).getType() != std::get<2>(e).getType())499      return emitOpError() << "types mismatch between " << i500                           << "th iter region arg and defined value";501 502    ++i;503  }504  return success();505}506 507std::optional<SmallVector<Value>> ForOp::getLoopInductionVars() {508  return SmallVector<Value>{getInductionVar()};509}510 511std::optional<SmallVector<OpFoldResult>> ForOp::getLoopLowerBounds() {512  return SmallVector<OpFoldResult>{OpFoldResult(getLowerBound())};513}514 515std::optional<SmallVector<OpFoldResult>> ForOp::getLoopSteps() {516  return SmallVector<OpFoldResult>{OpFoldResult(getStep())};517}518 519std::optional<SmallVector<OpFoldResult>> ForOp::getLoopUpperBounds() {520  return SmallVector<OpFoldResult>{OpFoldResult(getUpperBound())};521}522 523std::optional<ResultRange> ForOp::getLoopResults() { return getResults(); }524 525/// Promotes the loop body of a forOp to its containing block if the forOp526/// it can be determined that the loop has a single iteration.527LogicalResult ForOp::promoteIfSingleIteration(RewriterBase &rewriter) {528  std::optional<APInt> tripCount = getStaticTripCount();529  LDBG() << "promoteIfSingleIteration tripCount is " << tripCount530         << " for loop "531         << OpWithFlags(getOperation(), OpPrintingFlags().skipRegions());532  if (!tripCount.has_value() || tripCount->getSExtValue() > 1)533    return failure();534 535  if (*tripCount == 0) {536    rewriter.replaceAllUsesWith(getResults(), getInitArgs());537    rewriter.eraseOp(*this);538    return success();539  }540 541  // Replace all results with the yielded values.542  auto yieldOp = cast<scf::YieldOp>(getBody()->getTerminator());543  rewriter.replaceAllUsesWith(getResults(), getYieldedValues());544 545  // Replace block arguments with lower bound (replacement for IV) and546  // iter_args.547  SmallVector<Value> bbArgReplacements;548  bbArgReplacements.push_back(getLowerBound());549  llvm::append_range(bbArgReplacements, getInitArgs());550 551  // Move the loop body operations to the loop's containing block.552  rewriter.inlineBlockBefore(getBody(), getOperation()->getBlock(),553                             getOperation()->getIterator(), bbArgReplacements);554 555  // Erase the old terminator and the loop.556  rewriter.eraseOp(yieldOp);557  rewriter.eraseOp(*this);558 559  return success();560}561 562/// Prints the initialization list in the form of563///   <prefix>(%inner = %outer, %inner2 = %outer2, <...>)564/// where 'inner' values are assumed to be region arguments and 'outer' values565/// are regular SSA values.566static void printInitializationList(OpAsmPrinter &p,567                                    Block::BlockArgListType blocksArgs,568                                    ValueRange initializers,569                                    StringRef prefix = "") {570  assert(blocksArgs.size() == initializers.size() &&571         "expected same length of arguments and initializers");572  if (initializers.empty())573    return;574 575  p << prefix << '(';576  llvm::interleaveComma(llvm::zip(blocksArgs, initializers), p, [&](auto it) {577    p << std::get<0>(it) << " = " << std::get<1>(it);578  });579  p << ")";580}581 582void ForOp::print(OpAsmPrinter &p) {583  if (getUnsignedCmp())584    p << " unsigned";585 586  p << " " << getInductionVar() << " = " << getLowerBound() << " to "587    << getUpperBound() << " step " << getStep();588 589  printInitializationList(p, getRegionIterArgs(), getInitArgs(), " iter_args");590  if (!getInitArgs().empty())591    p << " -> (" << getInitArgs().getTypes() << ')';592  p << ' ';593  if (Type t = getInductionVar().getType(); !t.isIndex())594    p << " : " << t << ' ';595  p.printRegion(getRegion(),596                /*printEntryBlockArgs=*/false,597                /*printBlockTerminators=*/!getInitArgs().empty());598  p.printOptionalAttrDict((*this)->getAttrs(),599                          /*elidedAttrs=*/getUnsignedCmpAttrName().strref());600}601 602ParseResult ForOp::parse(OpAsmParser &parser, OperationState &result) {603  auto &builder = parser.getBuilder();604  Type type;605 606  OpAsmParser::Argument inductionVariable;607  OpAsmParser::UnresolvedOperand lb, ub, step;608 609  if (succeeded(parser.parseOptionalKeyword("unsigned")))610    result.addAttribute(getUnsignedCmpAttrName(result.name),611                        builder.getUnitAttr());612 613  // Parse the induction variable followed by '='.614  if (parser.parseOperand(inductionVariable.ssaName) || parser.parseEqual() ||615      // Parse loop bounds.616      parser.parseOperand(lb) || parser.parseKeyword("to") ||617      parser.parseOperand(ub) || parser.parseKeyword("step") ||618      parser.parseOperand(step))619    return failure();620 621  // Parse the optional initial iteration arguments.622  SmallVector<OpAsmParser::Argument, 4> regionArgs;623  SmallVector<OpAsmParser::UnresolvedOperand, 4> operands;624  regionArgs.push_back(inductionVariable);625 626  bool hasIterArgs = succeeded(parser.parseOptionalKeyword("iter_args"));627  if (hasIterArgs) {628    // Parse assignment list and results type list.629    if (parser.parseAssignmentList(regionArgs, operands) ||630        parser.parseArrowTypeList(result.types))631      return failure();632  }633 634  if (regionArgs.size() != result.types.size() + 1)635    return parser.emitError(636        parser.getNameLoc(),637        "mismatch in number of loop-carried values and defined values");638 639  // Parse optional type, else assume Index.640  if (parser.parseOptionalColon())641    type = builder.getIndexType();642  else if (parser.parseType(type))643    return failure();644 645  // Set block argument types, so that they are known when parsing the region.646  regionArgs.front().type = type;647  for (auto [iterArg, type] :648       llvm::zip_equal(llvm::drop_begin(regionArgs), result.types))649    iterArg.type = type;650 651  // Parse the body region.652  Region *body = result.addRegion();653  if (parser.parseRegion(*body, regionArgs))654    return failure();655  ForOp::ensureTerminator(*body, builder, result.location);656 657  // Resolve input operands. This should be done after parsing the region to658  // catch invalid IR where operands were defined inside of the region.659  if (parser.resolveOperand(lb, type, result.operands) ||660      parser.resolveOperand(ub, type, result.operands) ||661      parser.resolveOperand(step, type, result.operands))662    return failure();663  if (hasIterArgs) {664    for (auto argOperandType : llvm::zip_equal(llvm::drop_begin(regionArgs),665                                               operands, result.types)) {666      Type type = std::get<2>(argOperandType);667      std::get<0>(argOperandType).type = type;668      if (parser.resolveOperand(std::get<1>(argOperandType), type,669                                result.operands))670        return failure();671    }672  }673 674  // Parse the optional attribute list.675  if (parser.parseOptionalAttrDict(result.attributes))676    return failure();677 678  return success();679}680 681SmallVector<Region *> ForOp::getLoopRegions() { return {&getRegion()}; }682 683Block::BlockArgListType ForOp::getRegionIterArgs() {684  return getBody()->getArguments().drop_front(getNumInductionVars());685}686 687MutableArrayRef<OpOperand> ForOp::getInitsMutable() {688  return getInitArgsMutable();689}690 691FailureOr<LoopLikeOpInterface>692ForOp::replaceWithAdditionalYields(RewriterBase &rewriter,693                                   ValueRange newInitOperands,694                                   bool replaceInitOperandUsesInLoop,695                                   const NewYieldValuesFn &newYieldValuesFn) {696  // Create a new loop before the existing one, with the extra operands.697  OpBuilder::InsertionGuard g(rewriter);698  rewriter.setInsertionPoint(getOperation());699  auto inits = llvm::to_vector(getInitArgs());700  inits.append(newInitOperands.begin(), newInitOperands.end());701  scf::ForOp newLoop = scf::ForOp::create(702      rewriter, getLoc(), getLowerBound(), getUpperBound(), getStep(), inits,703      [](OpBuilder &, Location, Value, ValueRange) {}, getUnsignedCmp());704  newLoop->setAttrs(getPrunedAttributeList(getOperation(), {}));705 706  // Generate the new yield values and append them to the scf.yield operation.707  auto yieldOp = cast<scf::YieldOp>(getBody()->getTerminator());708  ArrayRef<BlockArgument> newIterArgs =709      newLoop.getBody()->getArguments().take_back(newInitOperands.size());710  {711    OpBuilder::InsertionGuard g(rewriter);712    rewriter.setInsertionPoint(yieldOp);713    SmallVector<Value> newYieldedValues =714        newYieldValuesFn(rewriter, getLoc(), newIterArgs);715    assert(newInitOperands.size() == newYieldedValues.size() &&716           "expected as many new yield values as new iter operands");717    rewriter.modifyOpInPlace(yieldOp, [&]() {718      yieldOp.getResultsMutable().append(newYieldedValues);719    });720  }721 722  // Move the loop body to the new op.723  rewriter.mergeBlocks(getBody(), newLoop.getBody(),724                       newLoop.getBody()->getArguments().take_front(725                           getBody()->getNumArguments()));726 727  if (replaceInitOperandUsesInLoop) {728    // Replace all uses of `newInitOperands` with the corresponding basic block729    // arguments.730    for (auto it : llvm::zip(newInitOperands, newIterArgs)) {731      rewriter.replaceUsesWithIf(std::get<0>(it), std::get<1>(it),732                                 [&](OpOperand &use) {733                                   Operation *user = use.getOwner();734                                   return newLoop->isProperAncestor(user);735                                 });736    }737  }738 739  // Replace the old loop.740  rewriter.replaceOp(getOperation(),741                     newLoop->getResults().take_front(getNumResults()));742  return cast<LoopLikeOpInterface>(newLoop.getOperation());743}744 745ForOp mlir::scf::getForInductionVarOwner(Value val) {746  auto ivArg = llvm::dyn_cast<BlockArgument>(val);747  if (!ivArg)748    return ForOp();749  assert(ivArg.getOwner() && "unlinked block argument");750  auto *containingOp = ivArg.getOwner()->getParentOp();751  return dyn_cast_or_null<ForOp>(containingOp);752}753 754OperandRange ForOp::getEntrySuccessorOperands(RegionSuccessor successor) {755  return getInitArgs();756}757 758void ForOp::getSuccessorRegions(RegionBranchPoint point,759                                SmallVectorImpl<RegionSuccessor> &regions) {760  // Both the operation itself and the region may be branching into the body or761  // back into the operation itself. It is possible for loop not to enter the762  // body.763  regions.push_back(RegionSuccessor(&getRegion(), getRegionIterArgs()));764  regions.push_back(RegionSuccessor(getOperation(), getResults()));765}766 767SmallVector<Region *> ForallOp::getLoopRegions() { return {&getRegion()}; }768 769/// Promotes the loop body of a forallOp to its containing block if it can be770/// determined that the loop has a single iteration.771LogicalResult scf::ForallOp::promoteIfSingleIteration(RewriterBase &rewriter) {772  for (auto [lb, ub, step] :773       llvm::zip(getMixedLowerBound(), getMixedUpperBound(), getMixedStep())) {774    auto tripCount =775        constantTripCount(lb, ub, step, /*isSigned=*/true, computeUbMinusLb);776    if (!tripCount.has_value() || *tripCount != 1)777      return failure();778  }779 780  promote(rewriter, *this);781  return success();782}783 784Block::BlockArgListType ForallOp::getRegionIterArgs() {785  return getBody()->getArguments().drop_front(getRank());786}787 788MutableArrayRef<OpOperand> ForallOp::getInitsMutable() {789  return getOutputsMutable();790}791 792/// Promotes the loop body of a scf::ForallOp to its containing block.793void mlir::scf::promote(RewriterBase &rewriter, scf::ForallOp forallOp) {794  OpBuilder::InsertionGuard g(rewriter);795  scf::InParallelOp terminator = forallOp.getTerminator();796 797  // Replace block arguments with lower bounds (replacements for IVs) and798  // outputs.799  SmallVector<Value> bbArgReplacements = forallOp.getLowerBound(rewriter);800  bbArgReplacements.append(forallOp.getOutputs().begin(),801                           forallOp.getOutputs().end());802 803  // Move the loop body operations to the loop's containing block.804  rewriter.inlineBlockBefore(forallOp.getBody(), forallOp->getBlock(),805                             forallOp->getIterator(), bbArgReplacements);806 807  // Replace the terminator with tensor.insert_slice ops.808  rewriter.setInsertionPointAfter(forallOp);809  SmallVector<Value> results;810  results.reserve(forallOp.getResults().size());811  for (auto &yieldingOp : terminator.getYieldingOps()) {812    auto parallelInsertSliceOp =813        dyn_cast<tensor::ParallelInsertSliceOp>(yieldingOp);814    if (!parallelInsertSliceOp)815      continue;816 817    Value dst = parallelInsertSliceOp.getDest();818    Value src = parallelInsertSliceOp.getSource();819    if (llvm::isa<TensorType>(src.getType())) {820      results.push_back(tensor::InsertSliceOp::create(821          rewriter, forallOp.getLoc(), dst.getType(), src, dst,822          parallelInsertSliceOp.getOffsets(), parallelInsertSliceOp.getSizes(),823          parallelInsertSliceOp.getStrides(),824          parallelInsertSliceOp.getStaticOffsets(),825          parallelInsertSliceOp.getStaticSizes(),826          parallelInsertSliceOp.getStaticStrides()));827    } else {828      llvm_unreachable("unsupported terminator");829    }830  }831  rewriter.replaceAllUsesWith(forallOp.getResults(), results);832 833  // Erase the old terminator and the loop.834  rewriter.eraseOp(terminator);835  rewriter.eraseOp(forallOp);836}837 838LoopNest mlir::scf::buildLoopNest(839    OpBuilder &builder, Location loc, ValueRange lbs, ValueRange ubs,840    ValueRange steps, ValueRange iterArgs,841    function_ref<ValueVector(OpBuilder &, Location, ValueRange, ValueRange)>842        bodyBuilder) {843  assert(lbs.size() == ubs.size() &&844         "expected the same number of lower and upper bounds");845  assert(lbs.size() == steps.size() &&846         "expected the same number of lower bounds and steps");847 848  // If there are no bounds, call the body-building function and return early.849  if (lbs.empty()) {850    ValueVector results =851        bodyBuilder ? bodyBuilder(builder, loc, ValueRange(), iterArgs)852                    : ValueVector();853    assert(results.size() == iterArgs.size() &&854           "loop nest body must return as many values as loop has iteration "855           "arguments");856    return LoopNest{{}, std::move(results)};857  }858 859  // First, create the loop structure iteratively using the body-builder860  // callback of `ForOp::build`. Do not create `YieldOp`s yet.861  OpBuilder::InsertionGuard guard(builder);862  SmallVector<scf::ForOp, 4> loops;863  SmallVector<Value, 4> ivs;864  loops.reserve(lbs.size());865  ivs.reserve(lbs.size());866  ValueRange currentIterArgs = iterArgs;867  Location currentLoc = loc;868  for (unsigned i = 0, e = lbs.size(); i < e; ++i) {869    auto loop = scf::ForOp::create(870        builder, currentLoc, lbs[i], ubs[i], steps[i], currentIterArgs,871        [&](OpBuilder &nestedBuilder, Location nestedLoc, Value iv,872            ValueRange args) {873          ivs.push_back(iv);874          // It is safe to store ValueRange args because it points to block875          // arguments of a loop operation that we also own.876          currentIterArgs = args;877          currentLoc = nestedLoc;878        });879    // Set the builder to point to the body of the newly created loop. We don't880    // do this in the callback because the builder is reset when the callback881    // returns.882    builder.setInsertionPointToStart(loop.getBody());883    loops.push_back(loop);884  }885 886  // For all loops but the innermost, yield the results of the nested loop.887  for (unsigned i = 0, e = loops.size() - 1; i < e; ++i) {888    builder.setInsertionPointToEnd(loops[i].getBody());889    scf::YieldOp::create(builder, loc, loops[i + 1].getResults());890  }891 892  // In the body of the innermost loop, call the body building function if any893  // and yield its results.894  builder.setInsertionPointToStart(loops.back().getBody());895  ValueVector results = bodyBuilder896                            ? bodyBuilder(builder, currentLoc, ivs,897                                          loops.back().getRegionIterArgs())898                            : ValueVector();899  assert(results.size() == iterArgs.size() &&900         "loop nest body must return as many values as loop has iteration "901         "arguments");902  builder.setInsertionPointToEnd(loops.back().getBody());903  scf::YieldOp::create(builder, loc, results);904 905  // Return the loops.906  ValueVector nestResults;907  llvm::append_range(nestResults, loops.front().getResults());908  return LoopNest{std::move(loops), std::move(nestResults)};909}910 911LoopNest mlir::scf::buildLoopNest(912    OpBuilder &builder, Location loc, ValueRange lbs, ValueRange ubs,913    ValueRange steps,914    function_ref<void(OpBuilder &, Location, ValueRange)> bodyBuilder) {915  // Delegate to the main function by wrapping the body builder.916  return buildLoopNest(builder, loc, lbs, ubs, steps, {},917                       [&bodyBuilder](OpBuilder &nestedBuilder,918                                      Location nestedLoc, ValueRange ivs,919                                      ValueRange) -> ValueVector {920                         if (bodyBuilder)921                           bodyBuilder(nestedBuilder, nestedLoc, ivs);922                         return {};923                       });924}925 926SmallVector<Value>927mlir::scf::replaceAndCastForOpIterArg(RewriterBase &rewriter, scf::ForOp forOp,928                                      OpOperand &operand, Value replacement,929                                      const ValueTypeCastFnTy &castFn) {930  assert(operand.getOwner() == forOp);931  Type oldType = operand.get().getType(), newType = replacement.getType();932 933  // 1. Create new iter operands, exactly 1 is replaced.934  assert(operand.getOperandNumber() >= forOp.getNumControlOperands() &&935         "expected an iter OpOperand");936  assert(operand.get().getType() != replacement.getType() &&937         "Expected a different type");938  SmallVector<Value> newIterOperands;939  for (OpOperand &opOperand : forOp.getInitArgsMutable()) {940    if (opOperand.getOperandNumber() == operand.getOperandNumber()) {941      newIterOperands.push_back(replacement);942      continue;943    }944    newIterOperands.push_back(opOperand.get());945  }946 947  // 2. Create the new forOp shell.948  scf::ForOp newForOp = scf::ForOp::create(949      rewriter, forOp.getLoc(), forOp.getLowerBound(), forOp.getUpperBound(),950      forOp.getStep(), newIterOperands, /*bodyBuilder=*/nullptr,951      forOp.getUnsignedCmp());952  newForOp->setAttrs(forOp->getAttrs());953  Block &newBlock = newForOp.getRegion().front();954  SmallVector<Value, 4> newBlockTransferArgs(newBlock.getArguments().begin(),955                                             newBlock.getArguments().end());956 957  // 3. Inject an incoming cast op at the beginning of the block for the bbArg958  // corresponding to the `replacement` value.959  OpBuilder::InsertionGuard g(rewriter);960  rewriter.setInsertionPointToStart(&newBlock);961  BlockArgument newRegionIterArg = newForOp.getTiedLoopRegionIterArg(962      &newForOp->getOpOperand(operand.getOperandNumber()));963  Value castIn = castFn(rewriter, newForOp.getLoc(), oldType, newRegionIterArg);964  newBlockTransferArgs[newRegionIterArg.getArgNumber()] = castIn;965 966  // 4. Steal the old block ops, mapping to the newBlockTransferArgs.967  Block &oldBlock = forOp.getRegion().front();968  rewriter.mergeBlocks(&oldBlock, &newBlock, newBlockTransferArgs);969 970  // 5. Inject an outgoing cast op at the end of the block and yield it instead.971  auto clonedYieldOp = cast<scf::YieldOp>(newBlock.getTerminator());972  rewriter.setInsertionPoint(clonedYieldOp);973  unsigned yieldIdx =974      newRegionIterArg.getArgNumber() - forOp.getNumInductionVars();975  Value castOut = castFn(rewriter, newForOp.getLoc(), newType,976                         clonedYieldOp.getOperand(yieldIdx));977  SmallVector<Value> newYieldOperands = clonedYieldOp.getOperands();978  newYieldOperands[yieldIdx] = castOut;979  scf::YieldOp::create(rewriter, newForOp.getLoc(), newYieldOperands);980  rewriter.eraseOp(clonedYieldOp);981 982  // 6. Inject an outgoing cast op after the forOp.983  rewriter.setInsertionPointAfter(newForOp);984  SmallVector<Value> newResults = newForOp.getResults();985  newResults[yieldIdx] =986      castFn(rewriter, newForOp.getLoc(), oldType, newResults[yieldIdx]);987 988  return newResults;989}990 991namespace {992// Fold away ForOp iter arguments when:993// 1) The op yields the iter arguments.994// 2) The argument's corresponding outer region iterators (inputs) are yielded.995// 3) The iter arguments have no use and the corresponding (operation) results996// have no use.997//998// These arguments must be defined outside of the ForOp region and can just be999// forwarded after simplifying the op inits, yields and returns.1000//1001// The implementation uses `inlineBlockBefore` to steal the content of the1002// original ForOp and avoid cloning.1003struct ForOpIterArgsFolder : public OpRewritePattern<scf::ForOp> {1004  using OpRewritePattern<scf::ForOp>::OpRewritePattern;1005 1006  LogicalResult matchAndRewrite(scf::ForOp forOp,1007                                PatternRewriter &rewriter) const final {1008    bool canonicalize = false;1009 1010    // An internal flat vector of block transfer1011    // arguments `newBlockTransferArgs` keeps the 1-1 mapping of original to1012    // transformed block argument mappings. This plays the role of a1013    // IRMapping for the particular use case of calling into1014    // `inlineBlockBefore`.1015    int64_t numResults = forOp.getNumResults();1016    SmallVector<bool, 4> keepMask;1017    keepMask.reserve(numResults);1018    SmallVector<Value, 4> newBlockTransferArgs, newIterArgs, newYieldValues,1019        newResultValues;1020    newBlockTransferArgs.reserve(1 + numResults);1021    newBlockTransferArgs.push_back(Value()); // iv placeholder with null value1022    newIterArgs.reserve(forOp.getInitArgs().size());1023    newYieldValues.reserve(numResults);1024    newResultValues.reserve(numResults);1025    DenseMap<std::pair<Value, Value>, std::pair<Value, Value>> initYieldToArg;1026    for (auto [init, arg, result, yielded] :1027         llvm::zip(forOp.getInitArgs(),       // iter from outside1028                   forOp.getRegionIterArgs(), // iter inside region1029                   forOp.getResults(),        // op results1030                   forOp.getYieldedValues()   // iter yield1031                   )) {1032      // Forwarded is `true` when:1033      // 1) The region `iter` argument is yielded.1034      // 2) The region `iter` argument the corresponding input is yielded.1035      // 3) The region `iter` argument has no use, and the corresponding op1036      // result has no use.1037      bool forwarded = (arg == yielded) || (init == yielded) ||1038                       (arg.use_empty() && result.use_empty());1039      if (forwarded) {1040        canonicalize = true;1041        keepMask.push_back(false);1042        newBlockTransferArgs.push_back(init);1043        newResultValues.push_back(init);1044        continue;1045      }1046 1047      // Check if a previous kept argument always has the same values for init1048      // and yielded values.1049      if (auto it = initYieldToArg.find({init, yielded});1050          it != initYieldToArg.end()) {1051        canonicalize = true;1052        keepMask.push_back(false);1053        auto [sameArg, sameResult] = it->second;1054        rewriter.replaceAllUsesWith(arg, sameArg);1055        rewriter.replaceAllUsesWith(result, sameResult);1056        // The replacement value doesn't matter because there are no uses.1057        newBlockTransferArgs.push_back(init);1058        newResultValues.push_back(init);1059        continue;1060      }1061 1062      // This value is kept.1063      initYieldToArg.insert({{init, yielded}, {arg, result}});1064      keepMask.push_back(true);1065      newIterArgs.push_back(init);1066      newYieldValues.push_back(yielded);1067      newBlockTransferArgs.push_back(Value()); // placeholder with null value1068      newResultValues.push_back(Value());      // placeholder with null value1069    }1070 1071    if (!canonicalize)1072      return failure();1073 1074    scf::ForOp newForOp =1075        scf::ForOp::create(rewriter, forOp.getLoc(), forOp.getLowerBound(),1076                           forOp.getUpperBound(), forOp.getStep(), newIterArgs,1077                           /*bodyBuilder=*/nullptr, forOp.getUnsignedCmp());1078    newForOp->setAttrs(forOp->getAttrs());1079    Block &newBlock = newForOp.getRegion().front();1080 1081    // Replace the null placeholders with newly constructed values.1082    newBlockTransferArgs[0] = newBlock.getArgument(0); // iv1083    for (unsigned idx = 0, collapsedIdx = 0, e = newResultValues.size();1084         idx != e; ++idx) {1085      Value &blockTransferArg = newBlockTransferArgs[1 + idx];1086      Value &newResultVal = newResultValues[idx];1087      assert((blockTransferArg && newResultVal) ||1088             (!blockTransferArg && !newResultVal));1089      if (!blockTransferArg) {1090        blockTransferArg = newForOp.getRegionIterArgs()[collapsedIdx];1091        newResultVal = newForOp.getResult(collapsedIdx++);1092      }1093    }1094 1095    Block &oldBlock = forOp.getRegion().front();1096    assert(oldBlock.getNumArguments() == newBlockTransferArgs.size() &&1097           "unexpected argument size mismatch");1098 1099    // No results case: the scf::ForOp builder already created a zero1100    // result terminator. Merge before this terminator and just get rid of the1101    // original terminator that has been merged in.1102    if (newIterArgs.empty()) {1103      auto newYieldOp = cast<scf::YieldOp>(newBlock.getTerminator());1104      rewriter.inlineBlockBefore(&oldBlock, newYieldOp, newBlockTransferArgs);1105      rewriter.eraseOp(newBlock.getTerminator()->getPrevNode());1106      rewriter.replaceOp(forOp, newResultValues);1107      return success();1108    }1109 1110    // No terminator case: merge and rewrite the merged terminator.1111    auto cloneFilteredTerminator = [&](scf::YieldOp mergedTerminator) {1112      OpBuilder::InsertionGuard g(rewriter);1113      rewriter.setInsertionPoint(mergedTerminator);1114      SmallVector<Value, 4> filteredOperands;1115      filteredOperands.reserve(newResultValues.size());1116      for (unsigned idx = 0, e = keepMask.size(); idx < e; ++idx)1117        if (keepMask[idx])1118          filteredOperands.push_back(mergedTerminator.getOperand(idx));1119      scf::YieldOp::create(rewriter, mergedTerminator.getLoc(),1120                           filteredOperands);1121    };1122 1123    rewriter.mergeBlocks(&oldBlock, &newBlock, newBlockTransferArgs);1124    auto mergedYieldOp = cast<scf::YieldOp>(newBlock.getTerminator());1125    cloneFilteredTerminator(mergedYieldOp);1126    rewriter.eraseOp(mergedYieldOp);1127    rewriter.replaceOp(forOp, newResultValues);1128    return success();1129  }1130};1131 1132/// Rewriting pattern that erases loops that are known not to iterate, replaces1133/// single-iteration loops with their bodies, and removes empty loops that1134/// iterate at least once and only return values defined outside of the loop.1135struct SimplifyTrivialLoops : public OpRewritePattern<ForOp> {1136  using OpRewritePattern<ForOp>::OpRewritePattern;1137 1138  LogicalResult matchAndRewrite(ForOp op,1139                                PatternRewriter &rewriter) const override {1140    std::optional<APInt> tripCount = op.getStaticTripCount();1141    if (!tripCount.has_value())1142      return rewriter.notifyMatchFailure(op,1143                                         "can't compute constant trip count");1144 1145    if (tripCount->isZero()) {1146      LDBG() << "SimplifyTrivialLoops tripCount is 0 for loop "1147             << OpWithFlags(op, OpPrintingFlags().skipRegions());1148      rewriter.replaceOp(op, op.getInitArgs());1149      return success();1150    }1151 1152    if (tripCount->getSExtValue() == 1) {1153      LDBG() << "SimplifyTrivialLoops tripCount is 1 for loop "1154             << OpWithFlags(op, OpPrintingFlags().skipRegions());1155      SmallVector<Value, 4> blockArgs;1156      blockArgs.reserve(op.getInitArgs().size() + 1);1157      blockArgs.push_back(op.getLowerBound());1158      llvm::append_range(blockArgs, op.getInitArgs());1159      replaceOpWithRegion(rewriter, op, op.getRegion(), blockArgs);1160      return success();1161    }1162 1163    // Now we are left with loops that have more than 1 iterations.1164    Block &block = op.getRegion().front();1165    if (!llvm::hasSingleElement(block))1166      return failure();1167    // The loop is empty and iterates at least once, if it only returns values1168    // defined outside of the loop, remove it and replace it with yield values.1169    if (llvm::any_of(op.getYieldedValues(),1170                     [&](Value v) { return !op.isDefinedOutsideOfLoop(v); }))1171      return failure();1172    LDBG() << "SimplifyTrivialLoops empty body loop allows replacement with "1173              "yield operands for loop "1174           << OpWithFlags(op, OpPrintingFlags().skipRegions());1175    rewriter.replaceOp(op, op.getYieldedValues());1176    return success();1177  }1178};1179 1180/// Fold scf.for iter_arg/result pairs that go through incoming/ougoing1181/// a tensor.cast op pair so as to pull the tensor.cast inside the scf.for:1182///1183/// ```1184///   %0 = tensor.cast %t0 : tensor<32x1024xf32> to tensor<?x?xf32>1185///   %1 = scf.for %i = %c0 to %c1024 step %c32 iter_args(%iter_t0 = %0)1186///      -> (tensor<?x?xf32>) {1187///     %2 = call @do(%iter_t0) : (tensor<?x?xf32>) -> tensor<?x?xf32>1188///     scf.yield %2 : tensor<?x?xf32>1189///   }1190///   use_of(%1)1191/// ```1192///1193/// folds into:1194///1195/// ```1196///   %0 = scf.for %arg2 = %c0 to %c1024 step %c32 iter_args(%arg3 = %arg0)1197///       -> (tensor<32x1024xf32>) {1198///     %2 = tensor.cast %arg3 : tensor<32x1024xf32> to tensor<?x?xf32>1199///     %3 = call @do(%2) : (tensor<?x?xf32>) -> tensor<?x?xf32>1200///     %4 = tensor.cast %3 : tensor<?x?xf32> to tensor<32x1024xf32>1201///     scf.yield %4 : tensor<32x1024xf32>1202///   }1203///   %1 = tensor.cast %0 : tensor<32x1024xf32> to tensor<?x?xf32>1204///   use_of(%1)1205/// ```1206struct ForOpTensorCastFolder : public OpRewritePattern<ForOp> {1207  using OpRewritePattern<ForOp>::OpRewritePattern;1208 1209  LogicalResult matchAndRewrite(ForOp op,1210                                PatternRewriter &rewriter) const override {1211    for (auto it : llvm::zip(op.getInitArgsMutable(), op.getResults())) {1212      OpOperand &iterOpOperand = std::get<0>(it);1213      auto incomingCast = iterOpOperand.get().getDefiningOp<tensor::CastOp>();1214      if (!incomingCast ||1215          incomingCast.getSource().getType() == incomingCast.getType())1216        continue;1217      // If the dest type of the cast does not preserve static information in1218      // the source type.1219      if (!tensor::preservesStaticInformation(1220              incomingCast.getDest().getType(),1221              incomingCast.getSource().getType()))1222        continue;1223      if (!std::get<1>(it).hasOneUse())1224        continue;1225 1226      // Create a new ForOp with that iter operand replaced.1227      rewriter.replaceOp(1228          op, replaceAndCastForOpIterArg(1229                  rewriter, op, iterOpOperand, incomingCast.getSource(),1230                  [](OpBuilder &b, Location loc, Type type, Value source) {1231                    return tensor::CastOp::create(b, loc, type, source);1232                  }));1233      return success();1234    }1235    return failure();1236  }1237};1238 1239} // namespace1240 1241void ForOp::getCanonicalizationPatterns(RewritePatternSet &results,1242                                        MLIRContext *context) {1243  results.add<ForOpIterArgsFolder, SimplifyTrivialLoops, ForOpTensorCastFolder>(1244      context);1245}1246 1247std::optional<APInt> ForOp::getConstantStep() {1248  IntegerAttr step;1249  if (matchPattern(getStep(), m_Constant(&step)))1250    return step.getValue();1251  return {};1252}1253 1254std::optional<MutableArrayRef<OpOperand>> ForOp::getYieldedValuesMutable() {1255  return cast<scf::YieldOp>(getBody()->getTerminator()).getResultsMutable();1256}1257 1258Speculation::Speculatability ForOp::getSpeculatability() {1259  // `scf.for (I = Start; I < End; I += 1)` terminates for all values of Start1260  // and End.1261  if (auto constantStep = getConstantStep())1262    if (*constantStep == 1)1263      return Speculation::RecursivelySpeculatable;1264 1265  // For Step != 1, the loop may not terminate.  We can add more smarts here if1266  // needed.1267  return Speculation::NotSpeculatable;1268}1269 1270std::optional<APInt> ForOp::getStaticTripCount() {1271  return constantTripCount(getLowerBound(), getUpperBound(), getStep(),1272                           /*isSigned=*/!getUnsignedCmp(), computeUbMinusLb);1273}1274 1275//===----------------------------------------------------------------------===//1276// ForallOp1277//===----------------------------------------------------------------------===//1278 1279LogicalResult ForallOp::verify() {1280  unsigned numLoops = getRank();1281  // Check number of outputs.1282  if (getNumResults() != getOutputs().size())1283    return emitOpError("produces ")1284           << getNumResults() << " results, but has only "1285           << getOutputs().size() << " outputs";1286 1287  // Check that the body defines block arguments for thread indices and outputs.1288  auto *body = getBody();1289  if (body->getNumArguments() != numLoops + getOutputs().size())1290    return emitOpError("region expects ") << numLoops << " arguments";1291  for (int64_t i = 0; i < numLoops; ++i)1292    if (!body->getArgument(i).getType().isIndex())1293      return emitOpError("expects ")1294             << i << "-th block argument to be an index";1295  for (unsigned i = 0; i < getOutputs().size(); ++i)1296    if (body->getArgument(i + numLoops).getType() != getOutputs()[i].getType())1297      return emitOpError("type mismatch between ")1298             << i << "-th output and corresponding block argument";1299  if (getMapping().has_value() && !getMapping()->empty()) {1300    if (getDeviceMappingAttrs().size() != numLoops)1301      return emitOpError() << "mapping attribute size must match op rank";1302    if (failed(getDeviceMaskingAttr()))1303      return emitOpError() << getMappingAttrName()1304                           << " supports at most one device masking attribute";1305  }1306 1307  // Verify mixed static/dynamic control variables.1308  Operation *op = getOperation();1309  if (failed(verifyListOfOperandsOrIntegers(op, "lower bound", numLoops,1310                                            getStaticLowerBound(),1311                                            getDynamicLowerBound())))1312    return failure();1313  if (failed(verifyListOfOperandsOrIntegers(op, "upper bound", numLoops,1314                                            getStaticUpperBound(),1315                                            getDynamicUpperBound())))1316    return failure();1317  if (failed(verifyListOfOperandsOrIntegers(op, "step", numLoops,1318                                            getStaticStep(), getDynamicStep())))1319    return failure();1320 1321  return success();1322}1323 1324void ForallOp::print(OpAsmPrinter &p) {1325  Operation *op = getOperation();1326  p << " (" << getInductionVars();1327  if (isNormalized()) {1328    p << ") in ";1329    printDynamicIndexList(p, op, getDynamicUpperBound(), getStaticUpperBound(),1330                          /*valueTypes=*/{}, /*scalables=*/{},1331                          OpAsmParser::Delimiter::Paren);1332  } else {1333    p << ") = ";1334    printDynamicIndexList(p, op, getDynamicLowerBound(), getStaticLowerBound(),1335                          /*valueTypes=*/{}, /*scalables=*/{},1336                          OpAsmParser::Delimiter::Paren);1337    p << " to ";1338    printDynamicIndexList(p, op, getDynamicUpperBound(), getStaticUpperBound(),1339                          /*valueTypes=*/{}, /*scalables=*/{},1340                          OpAsmParser::Delimiter::Paren);1341    p << " step ";1342    printDynamicIndexList(p, op, getDynamicStep(), getStaticStep(),1343                          /*valueTypes=*/{}, /*scalables=*/{},1344                          OpAsmParser::Delimiter::Paren);1345  }1346  printInitializationList(p, getRegionOutArgs(), getOutputs(), " shared_outs");1347  p << " ";1348  if (!getRegionOutArgs().empty())1349    p << "-> (" << getResultTypes() << ") ";1350  p.printRegion(getRegion(),1351                /*printEntryBlockArgs=*/false,1352                /*printBlockTerminators=*/getNumResults() > 0);1353  p.printOptionalAttrDict(op->getAttrs(), {getOperandSegmentSizesAttrName(),1354                                           getStaticLowerBoundAttrName(),1355                                           getStaticUpperBoundAttrName(),1356                                           getStaticStepAttrName()});1357}1358 1359ParseResult ForallOp::parse(OpAsmParser &parser, OperationState &result) {1360  OpBuilder b(parser.getContext());1361  auto indexType = b.getIndexType();1362 1363  // Parse an opening `(` followed by thread index variables followed by `)`1364  // TODO: when we can refer to such "induction variable"-like handles from the1365  // declarative assembly format, we can implement the parser as a custom hook.1366  SmallVector<OpAsmParser::Argument, 4> ivs;1367  if (parser.parseArgumentList(ivs, OpAsmParser::Delimiter::Paren))1368    return failure();1369 1370  DenseI64ArrayAttr staticLbs, staticUbs, staticSteps;1371  SmallVector<OpAsmParser::UnresolvedOperand> dynamicLbs, dynamicUbs,1372      dynamicSteps;1373  if (succeeded(parser.parseOptionalKeyword("in"))) {1374    // Parse upper bounds.1375    if (parseDynamicIndexList(parser, dynamicUbs, staticUbs,1376                              /*valueTypes=*/nullptr,1377                              OpAsmParser::Delimiter::Paren) ||1378        parser.resolveOperands(dynamicUbs, indexType, result.operands))1379      return failure();1380 1381    unsigned numLoops = ivs.size();1382    staticLbs = b.getDenseI64ArrayAttr(SmallVector<int64_t>(numLoops, 0));1383    staticSteps = b.getDenseI64ArrayAttr(SmallVector<int64_t>(numLoops, 1));1384  } else {1385    // Parse lower bounds.1386    if (parser.parseEqual() ||1387        parseDynamicIndexList(parser, dynamicLbs, staticLbs,1388                              /*valueTypes=*/nullptr,1389                              OpAsmParser::Delimiter::Paren) ||1390 1391        parser.resolveOperands(dynamicLbs, indexType, result.operands))1392      return failure();1393 1394    // Parse upper bounds.1395    if (parser.parseKeyword("to") ||1396        parseDynamicIndexList(parser, dynamicUbs, staticUbs,1397                              /*valueTypes=*/nullptr,1398                              OpAsmParser::Delimiter::Paren) ||1399        parser.resolveOperands(dynamicUbs, indexType, result.operands))1400      return failure();1401 1402    // Parse step values.1403    if (parser.parseKeyword("step") ||1404        parseDynamicIndexList(parser, dynamicSteps, staticSteps,1405                              /*valueTypes=*/nullptr,1406                              OpAsmParser::Delimiter::Paren) ||1407        parser.resolveOperands(dynamicSteps, indexType, result.operands))1408      return failure();1409  }1410 1411  // Parse out operands and results.1412  SmallVector<OpAsmParser::Argument, 4> regionOutArgs;1413  SmallVector<OpAsmParser::UnresolvedOperand, 4> outOperands;1414  SMLoc outOperandsLoc = parser.getCurrentLocation();1415  if (succeeded(parser.parseOptionalKeyword("shared_outs"))) {1416    if (outOperands.size() != result.types.size())1417      return parser.emitError(outOperandsLoc,1418                              "mismatch between out operands and types");1419    if (parser.parseAssignmentList(regionOutArgs, outOperands) ||1420        parser.parseOptionalArrowTypeList(result.types) ||1421        parser.resolveOperands(outOperands, result.types, outOperandsLoc,1422                               result.operands))1423      return failure();1424  }1425 1426  // Parse region.1427  SmallVector<OpAsmParser::Argument, 4> regionArgs;1428  std::unique_ptr<Region> region = std::make_unique<Region>();1429  for (auto &iv : ivs) {1430    iv.type = b.getIndexType();1431    regionArgs.push_back(iv);1432  }1433  for (const auto &it : llvm::enumerate(regionOutArgs)) {1434    auto &out = it.value();1435    out.type = result.types[it.index()];1436    regionArgs.push_back(out);1437  }1438  if (parser.parseRegion(*region, regionArgs))1439    return failure();1440 1441  // Ensure terminator and move region.1442  ForallOp::ensureTerminator(*region, b, result.location);1443  result.addRegion(std::move(region));1444 1445  // Parse the optional attribute list.1446  if (parser.parseOptionalAttrDict(result.attributes))1447    return failure();1448 1449  result.addAttribute("staticLowerBound", staticLbs);1450  result.addAttribute("staticUpperBound", staticUbs);1451  result.addAttribute("staticStep", staticSteps);1452  result.addAttribute("operandSegmentSizes",1453                      parser.getBuilder().getDenseI32ArrayAttr(1454                          {static_cast<int32_t>(dynamicLbs.size()),1455                           static_cast<int32_t>(dynamicUbs.size()),1456                           static_cast<int32_t>(dynamicSteps.size()),1457                           static_cast<int32_t>(outOperands.size())}));1458  return success();1459}1460 1461// Builder that takes loop bounds.1462void ForallOp::build(1463    mlir::OpBuilder &b, mlir::OperationState &result,1464    ArrayRef<OpFoldResult> lbs, ArrayRef<OpFoldResult> ubs,1465    ArrayRef<OpFoldResult> steps, ValueRange outputs,1466    std::optional<ArrayAttr> mapping,1467    function_ref<void(OpBuilder &, Location, ValueRange)> bodyBuilderFn) {1468  SmallVector<int64_t> staticLbs, staticUbs, staticSteps;1469  SmallVector<Value> dynamicLbs, dynamicUbs, dynamicSteps;1470  dispatchIndexOpFoldResults(lbs, dynamicLbs, staticLbs);1471  dispatchIndexOpFoldResults(ubs, dynamicUbs, staticUbs);1472  dispatchIndexOpFoldResults(steps, dynamicSteps, staticSteps);1473 1474  result.addOperands(dynamicLbs);1475  result.addOperands(dynamicUbs);1476  result.addOperands(dynamicSteps);1477  result.addOperands(outputs);1478  result.addTypes(TypeRange(outputs));1479 1480  result.addAttribute(getStaticLowerBoundAttrName(result.name),1481                      b.getDenseI64ArrayAttr(staticLbs));1482  result.addAttribute(getStaticUpperBoundAttrName(result.name),1483                      b.getDenseI64ArrayAttr(staticUbs));1484  result.addAttribute(getStaticStepAttrName(result.name),1485                      b.getDenseI64ArrayAttr(staticSteps));1486  result.addAttribute(1487      "operandSegmentSizes",1488      b.getDenseI32ArrayAttr({static_cast<int32_t>(dynamicLbs.size()),1489                              static_cast<int32_t>(dynamicUbs.size()),1490                              static_cast<int32_t>(dynamicSteps.size()),1491                              static_cast<int32_t>(outputs.size())}));1492  if (mapping.has_value()) {1493    result.addAttribute(ForallOp::getMappingAttrName(result.name),1494                        mapping.value());1495  }1496 1497  Region *bodyRegion = result.addRegion();1498  OpBuilder::InsertionGuard g(b);1499  b.createBlock(bodyRegion);1500  Block &bodyBlock = bodyRegion->front();1501 1502  // Add block arguments for indices and outputs.1503  bodyBlock.addArguments(1504      SmallVector<Type>(lbs.size(), b.getIndexType()),1505      SmallVector<Location>(staticLbs.size(), result.location));1506  bodyBlock.addArguments(1507      TypeRange(outputs),1508      SmallVector<Location>(outputs.size(), result.location));1509 1510  b.setInsertionPointToStart(&bodyBlock);1511  if (!bodyBuilderFn) {1512    ForallOp::ensureTerminator(*bodyRegion, b, result.location);1513    return;1514  }1515  bodyBuilderFn(b, result.location, bodyBlock.getArguments());1516}1517 1518// Builder that takes loop bounds.1519void ForallOp::build(1520    mlir::OpBuilder &b, mlir::OperationState &result,1521    ArrayRef<OpFoldResult> ubs, ValueRange outputs,1522    std::optional<ArrayAttr> mapping,1523    function_ref<void(OpBuilder &, Location, ValueRange)> bodyBuilderFn) {1524  unsigned numLoops = ubs.size();1525  SmallVector<OpFoldResult> lbs(numLoops, b.getIndexAttr(0));1526  SmallVector<OpFoldResult> steps(numLoops, b.getIndexAttr(1));1527  build(b, result, lbs, ubs, steps, outputs, mapping, bodyBuilderFn);1528}1529 1530// Checks if the lbs are zeros and steps are ones.1531bool ForallOp::isNormalized() {1532  auto allEqual = [](ArrayRef<OpFoldResult> results, int64_t val) {1533    return llvm::all_of(results, [&](OpFoldResult ofr) {1534      auto intValue = getConstantIntValue(ofr);1535      return intValue.has_value() && intValue == val;1536    });1537  };1538  return allEqual(getMixedLowerBound(), 0) && allEqual(getMixedStep(), 1);1539}1540 1541InParallelOp ForallOp::getTerminator() {1542  return cast<InParallelOp>(getBody()->getTerminator());1543}1544 1545SmallVector<Operation *> ForallOp::getCombiningOps(BlockArgument bbArg) {1546  SmallVector<Operation *> storeOps;1547  for (Operation *user : bbArg.getUsers()) {1548    if (auto parallelOp = dyn_cast<ParallelCombiningOpInterface>(user)) {1549      storeOps.push_back(parallelOp);1550    }1551  }1552  return storeOps;1553}1554 1555SmallVector<DeviceMappingAttrInterface> ForallOp::getDeviceMappingAttrs() {1556  SmallVector<DeviceMappingAttrInterface> res;1557  if (!getMapping())1558    return res;1559  for (auto attr : getMapping()->getValue()) {1560    auto m = dyn_cast<DeviceMappingAttrInterface>(attr);1561    if (m)1562      res.push_back(m);1563  }1564  return res;1565}1566 1567FailureOr<DeviceMaskingAttrInterface> ForallOp::getDeviceMaskingAttr() {1568  DeviceMaskingAttrInterface res;1569  if (!getMapping())1570    return res;1571  for (auto attr : getMapping()->getValue()) {1572    auto m = dyn_cast<DeviceMaskingAttrInterface>(attr);1573    if (m && res)1574      return failure();1575    if (m)1576      res = m;1577  }1578  return res;1579}1580 1581bool ForallOp::usesLinearMapping() {1582  SmallVector<DeviceMappingAttrInterface> ifaces = getDeviceMappingAttrs();1583  if (ifaces.empty())1584    return false;1585  return ifaces.front().isLinearMapping();1586}1587 1588std::optional<SmallVector<Value>> ForallOp::getLoopInductionVars() {1589  return SmallVector<Value>{getBody()->getArguments().take_front(getRank())};1590}1591 1592// Get lower bounds as OpFoldResult.1593std::optional<SmallVector<OpFoldResult>> ForallOp::getLoopLowerBounds() {1594  Builder b(getOperation()->getContext());1595  return getMixedValues(getStaticLowerBound(), getDynamicLowerBound(), b);1596}1597 1598// Get upper bounds as OpFoldResult.1599std::optional<SmallVector<OpFoldResult>> ForallOp::getLoopUpperBounds() {1600  Builder b(getOperation()->getContext());1601  return getMixedValues(getStaticUpperBound(), getDynamicUpperBound(), b);1602}1603 1604// Get steps as OpFoldResult.1605std::optional<SmallVector<OpFoldResult>> ForallOp::getLoopSteps() {1606  Builder b(getOperation()->getContext());1607  return getMixedValues(getStaticStep(), getDynamicStep(), b);1608}1609 1610ForallOp mlir::scf::getForallOpThreadIndexOwner(Value val) {1611  auto tidxArg = llvm::dyn_cast<BlockArgument>(val);1612  if (!tidxArg)1613    return ForallOp();1614  assert(tidxArg.getOwner() && "unlinked block argument");1615  auto *containingOp = tidxArg.getOwner()->getParentOp();1616  return dyn_cast<ForallOp>(containingOp);1617}1618 1619namespace {1620/// Fold tensor.dim(forall shared_outs(... = %t)) to tensor.dim(%t).1621struct DimOfForallOp : public OpRewritePattern<tensor::DimOp> {1622  using OpRewritePattern<tensor::DimOp>::OpRewritePattern;1623 1624  LogicalResult matchAndRewrite(tensor::DimOp dimOp,1625                                PatternRewriter &rewriter) const final {1626    auto forallOp = dimOp.getSource().getDefiningOp<ForallOp>();1627    if (!forallOp)1628      return failure();1629    Value sharedOut =1630        forallOp.getTiedOpOperand(llvm::cast<OpResult>(dimOp.getSource()))1631            ->get();1632    rewriter.modifyOpInPlace(1633        dimOp, [&]() { dimOp.getSourceMutable().assign(sharedOut); });1634    return success();1635  }1636};1637 1638class ForallOpControlOperandsFolder : public OpRewritePattern<ForallOp> {1639public:1640  using OpRewritePattern<ForallOp>::OpRewritePattern;1641 1642  LogicalResult matchAndRewrite(ForallOp op,1643                                PatternRewriter &rewriter) const override {1644    SmallVector<OpFoldResult> mixedLowerBound(op.getMixedLowerBound());1645    SmallVector<OpFoldResult> mixedUpperBound(op.getMixedUpperBound());1646    SmallVector<OpFoldResult> mixedStep(op.getMixedStep());1647    if (failed(foldDynamicIndexList(mixedLowerBound)) &&1648        failed(foldDynamicIndexList(mixedUpperBound)) &&1649        failed(foldDynamicIndexList(mixedStep)))1650      return failure();1651 1652    rewriter.modifyOpInPlace(op, [&]() {1653      SmallVector<Value> dynamicLowerBound, dynamicUpperBound, dynamicStep;1654      SmallVector<int64_t> staticLowerBound, staticUpperBound, staticStep;1655      dispatchIndexOpFoldResults(mixedLowerBound, dynamicLowerBound,1656                                 staticLowerBound);1657      op.getDynamicLowerBoundMutable().assign(dynamicLowerBound);1658      op.setStaticLowerBound(staticLowerBound);1659 1660      dispatchIndexOpFoldResults(mixedUpperBound, dynamicUpperBound,1661                                 staticUpperBound);1662      op.getDynamicUpperBoundMutable().assign(dynamicUpperBound);1663      op.setStaticUpperBound(staticUpperBound);1664 1665      dispatchIndexOpFoldResults(mixedStep, dynamicStep, staticStep);1666      op.getDynamicStepMutable().assign(dynamicStep);1667      op.setStaticStep(staticStep);1668 1669      op->setAttr(ForallOp::getOperandSegmentSizeAttr(),1670                  rewriter.getDenseI32ArrayAttr(1671                      {static_cast<int32_t>(dynamicLowerBound.size()),1672                       static_cast<int32_t>(dynamicUpperBound.size()),1673                       static_cast<int32_t>(dynamicStep.size()),1674                       static_cast<int32_t>(op.getNumResults())}));1675    });1676    return success();1677  }1678};1679 1680/// The following canonicalization pattern folds the iter arguments of1681/// scf.forall op if :-1682/// 1. The corresponding result has zero uses.1683/// 2. The iter argument is NOT being modified within the loop body.1684/// uses.1685///1686/// Example of first case :-1687///  INPUT:1688///   %res:3 = scf.forall ... shared_outs(%arg0 = %a, %arg1 = %b, %arg2 = %c)1689///            {1690///                ...1691///                <SOME USE OF %arg0>1692///                <SOME USE OF %arg1>1693///                <SOME USE OF %arg2>1694///                ...1695///                scf.forall.in_parallel {1696///                    <STORE OP WITH DESTINATION %arg1>1697///                    <STORE OP WITH DESTINATION %arg0>1698///                    <STORE OP WITH DESTINATION %arg2>1699///                }1700///             }1701///   return %res#11702///1703///  OUTPUT:1704///   %res:3 = scf.forall ... shared_outs(%new_arg0 = %b)1705///            {1706///                ...1707///                <SOME USE OF %a>1708///                <SOME USE OF %new_arg0>1709///                <SOME USE OF %c>1710///                ...1711///                scf.forall.in_parallel {1712///                    <STORE OP WITH DESTINATION %new_arg0>1713///                }1714///             }1715///   return %res1716///1717/// NOTE: 1. All uses of the folded shared_outs (iter argument) within the1718///          scf.forall is replaced by their corresponding operands.1719///       2. Even if there are <STORE OP WITH DESTINATION *> ops within the body1720///          of the scf.forall besides within scf.forall.in_parallel terminator,1721///          this canonicalization remains valid. For more details, please refer1722///          to :1723///          https://github.com/llvm/llvm-project/pull/90189#discussion_r15890111241724///       3. TODO(avarma): Generalize it for other store ops. Currently it1725///          handles tensor.parallel_insert_slice ops only.1726///1727/// Example of second case :-1728///  INPUT:1729///   %res:2 = scf.forall ... shared_outs(%arg0 = %a, %arg1 = %b)1730///            {1731///                ...1732///                <SOME USE OF %arg0>1733///                <SOME USE OF %arg1>1734///                ...1735///                scf.forall.in_parallel {1736///                    <STORE OP WITH DESTINATION %arg1>1737///                }1738///             }1739///   return %res#0, %res#11740///1741///  OUTPUT:1742///   %res = scf.forall ... shared_outs(%new_arg0 = %b)1743///            {1744///                ...1745///                <SOME USE OF %a>1746///                <SOME USE OF %new_arg0>1747///                ...1748///                scf.forall.in_parallel {1749///                    <STORE OP WITH DESTINATION %new_arg0>1750///                }1751///             }1752///   return %a, %res1753struct ForallOpIterArgsFolder : public OpRewritePattern<ForallOp> {1754  using OpRewritePattern<ForallOp>::OpRewritePattern;1755 1756  LogicalResult matchAndRewrite(ForallOp forallOp,1757                                PatternRewriter &rewriter) const final {1758    // Step 1: For a given i-th result of scf.forall, check the following :-1759    //         a. If it has any use.1760    //         b. If the corresponding iter argument is being modified within1761    //            the loop, i.e. has at least one store op with the iter arg as1762    //            its destination operand. For this we use1763    //            ForallOp::getCombiningOps(iter_arg).1764    //1765    //         Based on the check we maintain the following :-1766    //         a. `resultToDelete` - i-th result of scf.forall that'll be1767    //            deleted.1768    //         b. `resultToReplace` - i-th result of the old scf.forall1769    //            whose uses will be replaced by the new scf.forall.1770    //         c. `newOuts` - the shared_outs' operand of the new scf.forall1771    //            corresponding to the i-th result with at least one use.1772    SetVector<OpResult> resultToDelete;1773    SmallVector<Value> resultToReplace;1774    SmallVector<Value> newOuts;1775    for (OpResult result : forallOp.getResults()) {1776      OpOperand *opOperand = forallOp.getTiedOpOperand(result);1777      BlockArgument blockArg = forallOp.getTiedBlockArgument(opOperand);1778      if (result.use_empty() || forallOp.getCombiningOps(blockArg).empty()) {1779        resultToDelete.insert(result);1780      } else {1781        resultToReplace.push_back(result);1782        newOuts.push_back(opOperand->get());1783      }1784    }1785 1786    // Return early if all results of scf.forall have at least one use and being1787    // modified within the loop.1788    if (resultToDelete.empty())1789      return failure();1790 1791    // Step 2: For the the i-th result, do the following :-1792    //         a. Fetch the corresponding BlockArgument.1793    //         b. Look for store ops (currently tensor.parallel_insert_slice)1794    //            with the BlockArgument as its destination operand.1795    //         c. Remove the operations fetched in b.1796    for (OpResult result : resultToDelete) {1797      OpOperand *opOperand = forallOp.getTiedOpOperand(result);1798      BlockArgument blockArg = forallOp.getTiedBlockArgument(opOperand);1799      SmallVector<Operation *> combiningOps =1800          forallOp.getCombiningOps(blockArg);1801      for (Operation *combiningOp : combiningOps)1802        rewriter.eraseOp(combiningOp);1803    }1804 1805    // Step 3. Create a new scf.forall op with the new shared_outs' operands1806    //         fetched earlier1807    auto newForallOp = scf::ForallOp::create(1808        rewriter, forallOp.getLoc(), forallOp.getMixedLowerBound(),1809        forallOp.getMixedUpperBound(), forallOp.getMixedStep(), newOuts,1810        forallOp.getMapping(),1811        /*bodyBuilderFn =*/[](OpBuilder &, Location, ValueRange) {});1812 1813    // Step 4. Merge the block of the old scf.forall into the newly created1814    //         scf.forall using the new set of arguments.1815    Block *loopBody = forallOp.getBody();1816    Block *newLoopBody = newForallOp.getBody();1817    ArrayRef<BlockArgument> newBbArgs = newLoopBody->getArguments();1818    // Form initial new bbArg list with just the control operands of the new1819    // scf.forall op.1820    SmallVector<Value> newBlockArgs =1821        llvm::map_to_vector(newBbArgs.take_front(forallOp.getRank()),1822                            [](BlockArgument b) -> Value { return b; });1823    Block::BlockArgListType newSharedOutsArgs = newForallOp.getRegionOutArgs();1824    unsigned index = 0;1825    // Take the new corresponding bbArg if the old bbArg was used as a1826    // destination in the in_parallel op. For all other bbArgs, use the1827    // corresponding init_arg from the old scf.forall op.1828    for (OpResult result : forallOp.getResults()) {1829      if (resultToDelete.count(result)) {1830        newBlockArgs.push_back(forallOp.getTiedOpOperand(result)->get());1831      } else {1832        newBlockArgs.push_back(newSharedOutsArgs[index++]);1833      }1834    }1835    rewriter.mergeBlocks(loopBody, newLoopBody, newBlockArgs);1836 1837    // Step 5. Replace the uses of result of old scf.forall with that of the new1838    //         scf.forall.1839    for (auto &&[oldResult, newResult] :1840         llvm::zip(resultToReplace, newForallOp->getResults()))1841      rewriter.replaceAllUsesWith(oldResult, newResult);1842 1843    // Step 6. Replace the uses of those values that either has no use or are1844    //         not being modified within the loop with the corresponding1845    //         OpOperand.1846    for (OpResult oldResult : resultToDelete)1847      rewriter.replaceAllUsesWith(oldResult,1848                                  forallOp.getTiedOpOperand(oldResult)->get());1849    return success();1850  }1851};1852 1853struct ForallOpSingleOrZeroIterationDimsFolder1854    : public OpRewritePattern<ForallOp> {1855  using OpRewritePattern<ForallOp>::OpRewritePattern;1856 1857  LogicalResult matchAndRewrite(ForallOp op,1858                                PatternRewriter &rewriter) const override {1859    // Do not fold dimensions if they are mapped to processing units.1860    if (op.getMapping().has_value() && !op.getMapping()->empty())1861      return failure();1862    Location loc = op.getLoc();1863 1864    // Compute new loop bounds that omit all single-iteration loop dimensions.1865    SmallVector<OpFoldResult> newMixedLowerBounds, newMixedUpperBounds,1866        newMixedSteps;1867    IRMapping mapping;1868    for (auto [lb, ub, step, iv] :1869         llvm::zip(op.getMixedLowerBound(), op.getMixedUpperBound(),1870                   op.getMixedStep(), op.getInductionVars())) {1871      auto numIterations =1872          constantTripCount(lb, ub, step, /*isSigned=*/true, computeUbMinusLb);1873      if (numIterations.has_value()) {1874        // Remove the loop if it performs zero iterations.1875        if (*numIterations == 0) {1876          rewriter.replaceOp(op, op.getOutputs());1877          return success();1878        }1879        // Replace the loop induction variable by the lower bound if the loop1880        // performs a single iteration. Otherwise, copy the loop bounds.1881        if (*numIterations == 1) {1882          mapping.map(iv, getValueOrCreateConstantIndexOp(rewriter, loc, lb));1883          continue;1884        }1885      }1886      newMixedLowerBounds.push_back(lb);1887      newMixedUpperBounds.push_back(ub);1888      newMixedSteps.push_back(step);1889    }1890 1891    // All of the loop dimensions perform a single iteration. Inline loop body.1892    if (newMixedLowerBounds.empty()) {1893      promote(rewriter, op);1894      return success();1895    }1896 1897    // Exit if none of the loop dimensions perform a single iteration.1898    if (newMixedLowerBounds.size() == static_cast<unsigned>(op.getRank())) {1899      return rewriter.notifyMatchFailure(1900          op, "no dimensions have 0 or 1 iterations");1901    }1902 1903    // Replace the loop by a lower-dimensional loop.1904    ForallOp newOp;1905    newOp = ForallOp::create(rewriter, loc, newMixedLowerBounds,1906                             newMixedUpperBounds, newMixedSteps,1907                             op.getOutputs(), std::nullopt, nullptr);1908    newOp.getBodyRegion().getBlocks().clear();1909    // The new loop needs to keep all attributes from the old one, except for1910    // "operandSegmentSizes" and static loop bound attributes which capture1911    // the outdated information of the old iteration domain.1912    SmallVector<StringAttr> elidedAttrs{newOp.getOperandSegmentSizesAttrName(),1913                                        newOp.getStaticLowerBoundAttrName(),1914                                        newOp.getStaticUpperBoundAttrName(),1915                                        newOp.getStaticStepAttrName()};1916    for (const auto &namedAttr : op->getAttrs()) {1917      if (llvm::is_contained(elidedAttrs, namedAttr.getName()))1918        continue;1919      rewriter.modifyOpInPlace(newOp, [&]() {1920        newOp->setAttr(namedAttr.getName(), namedAttr.getValue());1921      });1922    }1923    rewriter.cloneRegionBefore(op.getRegion(), newOp.getRegion(),1924                               newOp.getRegion().begin(), mapping);1925    rewriter.replaceOp(op, newOp.getResults());1926    return success();1927  }1928};1929 1930/// Replace all induction vars with a single trip count with their lower bound.1931struct ForallOpReplaceConstantInductionVar : public OpRewritePattern<ForallOp> {1932  using OpRewritePattern<ForallOp>::OpRewritePattern;1933 1934  LogicalResult matchAndRewrite(ForallOp op,1935                                PatternRewriter &rewriter) const override {1936    Location loc = op.getLoc();1937    bool changed = false;1938    for (auto [lb, ub, step, iv] :1939         llvm::zip(op.getMixedLowerBound(), op.getMixedUpperBound(),1940                   op.getMixedStep(), op.getInductionVars())) {1941      if (iv.hasNUses(0))1942        continue;1943      auto numIterations =1944          constantTripCount(lb, ub, step, /*isSigned=*/true, computeUbMinusLb);1945      if (!numIterations.has_value() || numIterations.value() != 1) {1946        continue;1947      }1948      rewriter.replaceAllUsesWith(1949          iv, getValueOrCreateConstantIndexOp(rewriter, loc, lb));1950      changed = true;1951    }1952    return success(changed);1953  }1954};1955 1956struct FoldTensorCastOfOutputIntoForallOp1957    : public OpRewritePattern<scf::ForallOp> {1958  using OpRewritePattern<scf::ForallOp>::OpRewritePattern;1959 1960  struct TypeCast {1961    Type srcType;1962    Type dstType;1963  };1964 1965  LogicalResult matchAndRewrite(scf::ForallOp forallOp,1966                                PatternRewriter &rewriter) const final {1967    llvm::SmallMapVector<unsigned, TypeCast, 2> tensorCastProducers;1968    llvm::SmallVector<Value> newOutputTensors = forallOp.getOutputs();1969    for (auto en : llvm::enumerate(newOutputTensors)) {1970      auto castOp = en.value().getDefiningOp<tensor::CastOp>();1971      if (!castOp)1972        continue;1973 1974      // Only casts that that preserve static information, i.e. will make the1975      // loop result type "more" static than before, will be folded.1976      if (!tensor::preservesStaticInformation(castOp.getDest().getType(),1977                                              castOp.getSource().getType())) {1978        continue;1979      }1980 1981      tensorCastProducers[en.index()] =1982          TypeCast{castOp.getSource().getType(), castOp.getType()};1983      newOutputTensors[en.index()] = castOp.getSource();1984    }1985 1986    if (tensorCastProducers.empty())1987      return failure();1988 1989    // Create new loop.1990    Location loc = forallOp.getLoc();1991    auto newForallOp = ForallOp::create(1992        rewriter, loc, forallOp.getMixedLowerBound(),1993        forallOp.getMixedUpperBound(), forallOp.getMixedStep(),1994        newOutputTensors, forallOp.getMapping(),1995        [&](OpBuilder nestedBuilder, Location nestedLoc, ValueRange bbArgs) {1996          auto castBlockArgs =1997              llvm::to_vector(bbArgs.take_back(forallOp->getNumResults()));1998          for (auto [index, cast] : tensorCastProducers) {1999            Value &oldTypeBBArg = castBlockArgs[index];2000            oldTypeBBArg = tensor::CastOp::create(nestedBuilder, nestedLoc,2001                                                  cast.dstType, oldTypeBBArg);2002          }2003 2004          // Move old body into new parallel loop.2005          SmallVector<Value> ivsBlockArgs =2006              llvm::to_vector(bbArgs.take_front(forallOp.getRank()));2007          ivsBlockArgs.append(castBlockArgs);2008          rewriter.mergeBlocks(forallOp.getBody(),2009                               bbArgs.front().getParentBlock(), ivsBlockArgs);2010        });2011 2012    // After `mergeBlocks` happened, the destinations in the terminator were2013    // mapped to the tensor.cast old-typed results of the output bbArgs. The2014    // destination have to be updated to point to the output bbArgs directly.2015    auto terminator = newForallOp.getTerminator();2016    for (auto [yieldingOp, outputBlockArg] : llvm::zip(2017             terminator.getYieldingOps(), newForallOp.getRegionIterArgs())) {2018      if (auto parallelCombingingOp =2019              dyn_cast<ParallelCombiningOpInterface>(yieldingOp)) {2020        parallelCombingingOp.getUpdatedDestinations().assign(outputBlockArg);2021      }2022    }2023 2024    // Cast results back to the original types.2025    rewriter.setInsertionPointAfter(newForallOp);2026    SmallVector<Value> castResults = newForallOp.getResults();2027    for (auto &item : tensorCastProducers) {2028      Value &oldTypeResult = castResults[item.first];2029      oldTypeResult = tensor::CastOp::create(rewriter, loc, item.second.dstType,2030                                             oldTypeResult);2031    }2032    rewriter.replaceOp(forallOp, castResults);2033    return success();2034  }2035};2036 2037} // namespace2038 2039void ForallOp::getCanonicalizationPatterns(RewritePatternSet &results,2040                                           MLIRContext *context) {2041  results.add<DimOfForallOp, FoldTensorCastOfOutputIntoForallOp,2042              ForallOpControlOperandsFolder, ForallOpIterArgsFolder,2043              ForallOpSingleOrZeroIterationDimsFolder,2044              ForallOpReplaceConstantInductionVar>(context);2045}2046 2047/// Given the region at `index`, or the parent operation if `index` is None,2048/// return the successor regions. These are the regions that may be selected2049/// during the flow of control. `operands` is a set of optional attributes that2050/// correspond to a constant value for each operand, or null if that operand is2051/// not a constant.2052void ForallOp::getSuccessorRegions(RegionBranchPoint point,2053                                   SmallVectorImpl<RegionSuccessor> &regions) {2054  // In accordance with the semantics of forall, its body is executed in2055  // parallel by multiple threads. We should not expect to branch back into2056  // the forall body after the region's execution is complete.2057  if (point.isParent())2058    regions.push_back(RegionSuccessor(&getRegion(), getRegionIterArgs()));2059  else2060    regions.push_back(2061        RegionSuccessor(getOperation(), getOperation()->getResults()));2062}2063 2064//===----------------------------------------------------------------------===//2065// InParallelOp2066//===----------------------------------------------------------------------===//2067 2068// Build a InParallelOp with mixed static and dynamic entries.2069void InParallelOp::build(OpBuilder &b, OperationState &result) {2070  OpBuilder::InsertionGuard g(b);2071  Region *bodyRegion = result.addRegion();2072  b.createBlock(bodyRegion);2073}2074 2075LogicalResult InParallelOp::verify() {2076  scf::ForallOp forallOp =2077      dyn_cast<scf::ForallOp>(getOperation()->getParentOp());2078  if (!forallOp)2079    return this->emitOpError("expected forall op parent");2080 2081  for (Operation &op : getRegion().front().getOperations()) {2082    auto parallelCombiningOp = dyn_cast<ParallelCombiningOpInterface>(&op);2083    if (!parallelCombiningOp) {2084      return this->emitOpError("expected only ParallelCombiningOpInterface")2085             << " ops";2086    }2087 2088    // Verify that inserts are into out block arguments.2089    MutableOperandRange dests = parallelCombiningOp.getUpdatedDestinations();2090    ArrayRef<BlockArgument> regionOutArgs = forallOp.getRegionOutArgs();2091    for (OpOperand &dest : dests) {2092      if (!llvm::is_contained(regionOutArgs, dest.get()))2093        return op.emitOpError("may only insert into an output block argument");2094    }2095  }2096 2097  return success();2098}2099 2100void InParallelOp::print(OpAsmPrinter &p) {2101  p << " ";2102  p.printRegion(getRegion(),2103                /*printEntryBlockArgs=*/false,2104                /*printBlockTerminators=*/false);2105  p.printOptionalAttrDict(getOperation()->getAttrs());2106}2107 2108ParseResult InParallelOp::parse(OpAsmParser &parser, OperationState &result) {2109  auto &builder = parser.getBuilder();2110 2111  SmallVector<OpAsmParser::Argument, 8> regionOperands;2112  std::unique_ptr<Region> region = std::make_unique<Region>();2113  if (parser.parseRegion(*region, regionOperands))2114    return failure();2115 2116  if (region->empty())2117    OpBuilder(builder.getContext()).createBlock(region.get());2118  result.addRegion(std::move(region));2119 2120  // Parse the optional attribute list.2121  if (parser.parseOptionalAttrDict(result.attributes))2122    return failure();2123  return success();2124}2125 2126OpResult InParallelOp::getParentResult(int64_t idx) {2127  return getOperation()->getParentOp()->getResult(idx);2128}2129 2130SmallVector<BlockArgument> InParallelOp::getDests() {2131  SmallVector<BlockArgument> updatedDests;2132  for (Operation &yieldingOp : getYieldingOps()) {2133    auto parallelCombiningOp =2134        dyn_cast<ParallelCombiningOpInterface>(&yieldingOp);2135    if (!parallelCombiningOp)2136      continue;2137    for (OpOperand &updatedOperand :2138         parallelCombiningOp.getUpdatedDestinations())2139      updatedDests.push_back(cast<BlockArgument>(updatedOperand.get()));2140  }2141  return updatedDests;2142}2143 2144llvm::iterator_range<Block::iterator> InParallelOp::getYieldingOps() {2145  return getRegion().front().getOperations();2146}2147 2148//===----------------------------------------------------------------------===//2149// IfOp2150//===----------------------------------------------------------------------===//2151 2152bool mlir::scf::insideMutuallyExclusiveBranches(Operation *a, Operation *b) {2153  assert(a && "expected non-empty operation");2154  assert(b && "expected non-empty operation");2155 2156  IfOp ifOp = a->getParentOfType<IfOp>();2157  while (ifOp) {2158    // Check if b is inside ifOp. (We already know that a is.)2159    if (ifOp->isProperAncestor(b))2160      // b is contained in ifOp. a and b are in mutually exclusive branches if2161      // they are in different blocks of ifOp.2162      return static_cast<bool>(ifOp.thenBlock()->findAncestorOpInBlock(*a)) !=2163             static_cast<bool>(ifOp.thenBlock()->findAncestorOpInBlock(*b));2164    // Check next enclosing IfOp.2165    ifOp = ifOp->getParentOfType<IfOp>();2166  }2167 2168  // Could not find a common IfOp among a's and b's ancestors.2169  return false;2170}2171 2172LogicalResult2173IfOp::inferReturnTypes(MLIRContext *ctx, std::optional<Location> loc,2174                       IfOp::Adaptor adaptor,2175                       SmallVectorImpl<Type> &inferredReturnTypes) {2176  if (adaptor.getRegions().empty())2177    return failure();2178  Region *r = &adaptor.getThenRegion();2179  if (r->empty())2180    return failure();2181  Block &b = r->front();2182  if (b.empty())2183    return failure();2184  auto yieldOp = llvm::dyn_cast<YieldOp>(b.back());2185  if (!yieldOp)2186    return failure();2187  TypeRange types = yieldOp.getOperandTypes();2188  llvm::append_range(inferredReturnTypes, types);2189  return success();2190}2191 2192void IfOp::build(OpBuilder &builder, OperationState &result,2193                 TypeRange resultTypes, Value cond) {2194  return build(builder, result, resultTypes, cond, /*addThenBlock=*/false,2195               /*addElseBlock=*/false);2196}2197 2198void IfOp::build(OpBuilder &builder, OperationState &result,2199                 TypeRange resultTypes, Value cond, bool addThenBlock,2200                 bool addElseBlock) {2201  assert((!addElseBlock || addThenBlock) &&2202         "must not create else block w/o then block");2203  result.addTypes(resultTypes);2204  result.addOperands(cond);2205 2206  // Add regions and blocks.2207  OpBuilder::InsertionGuard guard(builder);2208  Region *thenRegion = result.addRegion();2209  if (addThenBlock)2210    builder.createBlock(thenRegion);2211  Region *elseRegion = result.addRegion();2212  if (addElseBlock)2213    builder.createBlock(elseRegion);2214}2215 2216void IfOp::build(OpBuilder &builder, OperationState &result, Value cond,2217                 bool withElseRegion) {2218  build(builder, result, TypeRange{}, cond, withElseRegion);2219}2220 2221void IfOp::build(OpBuilder &builder, OperationState &result,2222                 TypeRange resultTypes, Value cond, bool withElseRegion) {2223  result.addTypes(resultTypes);2224  result.addOperands(cond);2225 2226  // Build then region.2227  OpBuilder::InsertionGuard guard(builder);2228  Region *thenRegion = result.addRegion();2229  builder.createBlock(thenRegion);2230  if (resultTypes.empty())2231    IfOp::ensureTerminator(*thenRegion, builder, result.location);2232 2233  // Build else region.2234  Region *elseRegion = result.addRegion();2235  if (withElseRegion) {2236    builder.createBlock(elseRegion);2237    if (resultTypes.empty())2238      IfOp::ensureTerminator(*elseRegion, builder, result.location);2239  }2240}2241 2242void IfOp::build(OpBuilder &builder, OperationState &result, Value cond,2243                 function_ref<void(OpBuilder &, Location)> thenBuilder,2244                 function_ref<void(OpBuilder &, Location)> elseBuilder) {2245  assert(thenBuilder && "the builder callback for 'then' must be present");2246  result.addOperands(cond);2247 2248  // Build then region.2249  OpBuilder::InsertionGuard guard(builder);2250  Region *thenRegion = result.addRegion();2251  builder.createBlock(thenRegion);2252  thenBuilder(builder, result.location);2253 2254  // Build else region.2255  Region *elseRegion = result.addRegion();2256  if (elseBuilder) {2257    builder.createBlock(elseRegion);2258    elseBuilder(builder, result.location);2259  }2260 2261  // Infer result types.2262  SmallVector<Type> inferredReturnTypes;2263  MLIRContext *ctx = builder.getContext();2264  auto attrDict = DictionaryAttr::get(ctx, result.attributes);2265  if (succeeded(inferReturnTypes(ctx, std::nullopt, result.operands, attrDict,2266                                 /*properties=*/nullptr, result.regions,2267                                 inferredReturnTypes))) {2268    result.addTypes(inferredReturnTypes);2269  }2270}2271 2272LogicalResult IfOp::verify() {2273  if (getNumResults() != 0 && getElseRegion().empty())2274    return emitOpError("must have an else block if defining values");2275  return success();2276}2277 2278ParseResult IfOp::parse(OpAsmParser &parser, OperationState &result) {2279  // Create the regions for 'then'.2280  result.regions.reserve(2);2281  Region *thenRegion = result.addRegion();2282  Region *elseRegion = result.addRegion();2283 2284  auto &builder = parser.getBuilder();2285  OpAsmParser::UnresolvedOperand cond;2286  Type i1Type = builder.getIntegerType(1);2287  if (parser.parseOperand(cond) ||2288      parser.resolveOperand(cond, i1Type, result.operands))2289    return failure();2290  // Parse optional results type list.2291  if (parser.parseOptionalArrowTypeList(result.types))2292    return failure();2293  // Parse the 'then' region.2294  if (parser.parseRegion(*thenRegion, /*arguments=*/{}, /*argTypes=*/{}))2295    return failure();2296  IfOp::ensureTerminator(*thenRegion, parser.getBuilder(), result.location);2297 2298  // If we find an 'else' keyword then parse the 'else' region.2299  if (!parser.parseOptionalKeyword("else")) {2300    if (parser.parseRegion(*elseRegion, /*arguments=*/{}, /*argTypes=*/{}))2301      return failure();2302    IfOp::ensureTerminator(*elseRegion, parser.getBuilder(), result.location);2303  }2304 2305  // Parse the optional attribute list.2306  if (parser.parseOptionalAttrDict(result.attributes))2307    return failure();2308  return success();2309}2310 2311void IfOp::print(OpAsmPrinter &p) {2312  bool printBlockTerminators = false;2313 2314  p << " " << getCondition();2315  if (!getResults().empty()) {2316    p << " -> (" << getResultTypes() << ")";2317    // Print yield explicitly if the op defines values.2318    printBlockTerminators = true;2319  }2320  p << ' ';2321  p.printRegion(getThenRegion(),2322                /*printEntryBlockArgs=*/false,2323                /*printBlockTerminators=*/printBlockTerminators);2324 2325  // Print the 'else' regions if it exists and has a block.2326  auto &elseRegion = getElseRegion();2327  if (!elseRegion.empty()) {2328    p << " else ";2329    p.printRegion(elseRegion,2330                  /*printEntryBlockArgs=*/false,2331                  /*printBlockTerminators=*/printBlockTerminators);2332  }2333 2334  p.printOptionalAttrDict((*this)->getAttrs());2335}2336 2337void IfOp::getSuccessorRegions(RegionBranchPoint point,2338                               SmallVectorImpl<RegionSuccessor> &regions) {2339  // The `then` and the `else` region branch back to the parent operation or one2340  // of the recursive parent operations (early exit case).2341  if (!point.isParent()) {2342    regions.push_back(RegionSuccessor(getOperation(), getResults()));2343    return;2344  }2345 2346  regions.push_back(RegionSuccessor(&getThenRegion()));2347 2348  // Don't consider the else region if it is empty.2349  Region *elseRegion = &this->getElseRegion();2350  if (elseRegion->empty())2351    regions.push_back(2352        RegionSuccessor(getOperation(), getOperation()->getResults()));2353  else2354    regions.push_back(RegionSuccessor(elseRegion));2355}2356 2357void IfOp::getEntrySuccessorRegions(ArrayRef<Attribute> operands,2358                                    SmallVectorImpl<RegionSuccessor> &regions) {2359  FoldAdaptor adaptor(operands, *this);2360  auto boolAttr = dyn_cast_or_null<BoolAttr>(adaptor.getCondition());2361  if (!boolAttr || boolAttr.getValue())2362    regions.emplace_back(&getThenRegion());2363 2364  // If the else region is empty, execution continues after the parent op.2365  if (!boolAttr || !boolAttr.getValue()) {2366    if (!getElseRegion().empty())2367      regions.emplace_back(&getElseRegion());2368    else2369      regions.emplace_back(getOperation(), getResults());2370  }2371}2372 2373LogicalResult IfOp::fold(FoldAdaptor adaptor,2374                         SmallVectorImpl<OpFoldResult> &results) {2375  // if (!c) then A() else B() -> if c then B() else A()2376  if (getElseRegion().empty())2377    return failure();2378 2379  arith::XOrIOp xorStmt = getCondition().getDefiningOp<arith::XOrIOp>();2380  if (!xorStmt)2381    return failure();2382 2383  if (!matchPattern(xorStmt.getRhs(), m_One()))2384    return failure();2385 2386  getConditionMutable().assign(xorStmt.getLhs());2387  Block *thenBlock = &getThenRegion().front();2388  // It would be nicer to use iplist::swap, but that has no implemented2389  // callbacks See: https://llvm.org/doxygen/ilist_8h_source.html#l002242390  getThenRegion().getBlocks().splice(getThenRegion().getBlocks().begin(),2391                                     getElseRegion().getBlocks());2392  getElseRegion().getBlocks().splice(getElseRegion().getBlocks().begin(),2393                                     getThenRegion().getBlocks(), thenBlock);2394  return success();2395}2396 2397void IfOp::getRegionInvocationBounds(2398    ArrayRef<Attribute> operands,2399    SmallVectorImpl<InvocationBounds> &invocationBounds) {2400  if (auto cond = llvm::dyn_cast_or_null<BoolAttr>(operands[0])) {2401    // If the condition is known, then one region is known to be executed once2402    // and the other zero times.2403    invocationBounds.emplace_back(0, cond.getValue() ? 1 : 0);2404    invocationBounds.emplace_back(0, cond.getValue() ? 0 : 1);2405  } else {2406    // Non-constant condition. Each region may be executed 0 or 1 times.2407    invocationBounds.assign(2, {0, 1});2408  }2409}2410 2411namespace {2412// Pattern to remove unused IfOp results.2413struct RemoveUnusedResults : public OpRewritePattern<IfOp> {2414  using OpRewritePattern<IfOp>::OpRewritePattern;2415 2416  void transferBody(Block *source, Block *dest, ArrayRef<OpResult> usedResults,2417                    PatternRewriter &rewriter) const {2418    // Move all operations to the destination block.2419    rewriter.mergeBlocks(source, dest);2420    // Replace the yield op by one that returns only the used values.2421    auto yieldOp = cast<scf::YieldOp>(dest->getTerminator());2422    SmallVector<Value, 4> usedOperands;2423    llvm::transform(usedResults, std::back_inserter(usedOperands),2424                    [&](OpResult result) {2425                      return yieldOp.getOperand(result.getResultNumber());2426                    });2427    rewriter.modifyOpInPlace(yieldOp,2428                             [&]() { yieldOp->setOperands(usedOperands); });2429  }2430 2431  LogicalResult matchAndRewrite(IfOp op,2432                                PatternRewriter &rewriter) const override {2433    // Compute the list of used results.2434    SmallVector<OpResult, 4> usedResults;2435    llvm::copy_if(op.getResults(), std::back_inserter(usedResults),2436                  [](OpResult result) { return !result.use_empty(); });2437 2438    // Replace the operation if only a subset of its results have uses.2439    if (usedResults.size() == op.getNumResults())2440      return failure();2441 2442    // Compute the result types of the replacement operation.2443    SmallVector<Type, 4> newTypes;2444    llvm::transform(usedResults, std::back_inserter(newTypes),2445                    [](OpResult result) { return result.getType(); });2446 2447    // Create a replacement operation with empty then and else regions.2448    auto newOp =2449        IfOp::create(rewriter, op.getLoc(), newTypes, op.getCondition());2450    rewriter.createBlock(&newOp.getThenRegion());2451    rewriter.createBlock(&newOp.getElseRegion());2452 2453    // Move the bodies and replace the terminators (note there is a then and2454    // an else region since the operation returns results).2455    transferBody(op.getBody(0), newOp.getBody(0), usedResults, rewriter);2456    transferBody(op.getBody(1), newOp.getBody(1), usedResults, rewriter);2457 2458    // Replace the operation by the new one.2459    SmallVector<Value, 4> repResults(op.getNumResults());2460    for (const auto &en : llvm::enumerate(usedResults))2461      repResults[en.value().getResultNumber()] = newOp.getResult(en.index());2462    rewriter.replaceOp(op, repResults);2463    return success();2464  }2465};2466 2467struct RemoveStaticCondition : public OpRewritePattern<IfOp> {2468  using OpRewritePattern<IfOp>::OpRewritePattern;2469 2470  LogicalResult matchAndRewrite(IfOp op,2471                                PatternRewriter &rewriter) const override {2472    BoolAttr condition;2473    if (!matchPattern(op.getCondition(), m_Constant(&condition)))2474      return failure();2475 2476    if (condition.getValue())2477      replaceOpWithRegion(rewriter, op, op.getThenRegion());2478    else if (!op.getElseRegion().empty())2479      replaceOpWithRegion(rewriter, op, op.getElseRegion());2480    else2481      rewriter.eraseOp(op);2482 2483    return success();2484  }2485};2486 2487/// Hoist any yielded results whose operands are defined outside2488/// the if, to a select instruction.2489struct ConvertTrivialIfToSelect : public OpRewritePattern<IfOp> {2490  using OpRewritePattern<IfOp>::OpRewritePattern;2491 2492  LogicalResult matchAndRewrite(IfOp op,2493                                PatternRewriter &rewriter) const override {2494    if (op->getNumResults() == 0)2495      return failure();2496 2497    auto cond = op.getCondition();2498    auto thenYieldArgs = op.thenYield().getOperands();2499    auto elseYieldArgs = op.elseYield().getOperands();2500 2501    SmallVector<Type> nonHoistable;2502    for (auto [trueVal, falseVal] : llvm::zip(thenYieldArgs, elseYieldArgs)) {2503      if (&op.getThenRegion() == trueVal.getParentRegion() ||2504          &op.getElseRegion() == falseVal.getParentRegion())2505        nonHoistable.push_back(trueVal.getType());2506    }2507    // Early exit if there aren't any yielded values we can2508    // hoist outside the if.2509    if (nonHoistable.size() == op->getNumResults())2510      return failure();2511 2512    IfOp replacement = IfOp::create(rewriter, op.getLoc(), nonHoistable, cond,2513                                    /*withElseRegion=*/false);2514    if (replacement.thenBlock())2515      rewriter.eraseBlock(replacement.thenBlock());2516    replacement.getThenRegion().takeBody(op.getThenRegion());2517    replacement.getElseRegion().takeBody(op.getElseRegion());2518 2519    SmallVector<Value> results(op->getNumResults());2520    assert(thenYieldArgs.size() == results.size());2521    assert(elseYieldArgs.size() == results.size());2522 2523    SmallVector<Value> trueYields;2524    SmallVector<Value> falseYields;2525    rewriter.setInsertionPoint(replacement);2526    for (const auto &it :2527         llvm::enumerate(llvm::zip(thenYieldArgs, elseYieldArgs))) {2528      Value trueVal = std::get<0>(it.value());2529      Value falseVal = std::get<1>(it.value());2530      if (&replacement.getThenRegion() == trueVal.getParentRegion() ||2531          &replacement.getElseRegion() == falseVal.getParentRegion()) {2532        results[it.index()] = replacement.getResult(trueYields.size());2533        trueYields.push_back(trueVal);2534        falseYields.push_back(falseVal);2535      } else if (trueVal == falseVal)2536        results[it.index()] = trueVal;2537      else2538        results[it.index()] = arith::SelectOp::create(rewriter, op.getLoc(),2539                                                      cond, trueVal, falseVal);2540    }2541 2542    rewriter.setInsertionPointToEnd(replacement.thenBlock());2543    rewriter.replaceOpWithNewOp<YieldOp>(replacement.thenYield(), trueYields);2544 2545    rewriter.setInsertionPointToEnd(replacement.elseBlock());2546    rewriter.replaceOpWithNewOp<YieldOp>(replacement.elseYield(), falseYields);2547 2548    rewriter.replaceOp(op, results);2549    return success();2550  }2551};2552 2553/// Allow the true region of an if to assume the condition is true2554/// and vice versa. For example:2555///2556///   scf.if %cmp {2557///      print(%cmp)2558///   }2559///2560///  becomes2561///2562///   scf.if %cmp {2563///      print(true)2564///   }2565///2566struct ConditionPropagation : public OpRewritePattern<IfOp> {2567  using OpRewritePattern<IfOp>::OpRewritePattern;2568 2569  /// Kind of parent region in the ancestor cache.2570  enum class Parent { Then, Else, None };2571 2572  /// Returns the kind of region ("then", "else", or "none") of the2573  /// IfOp that the given region is transitively nested in. Updates2574  /// the cache accordingly.2575  static Parent getParentType(Region *toCheck, IfOp op,2576                              DenseMap<Region *, Parent> &cache,2577                              Region *endRegion) {2578    SmallVector<Region *> seen;2579    while (toCheck != endRegion) {2580      auto found = cache.find(toCheck);2581      if (found != cache.end())2582        return found->second;2583      seen.push_back(toCheck);2584      if (&op.getThenRegion() == toCheck) {2585        for (Region *region : seen)2586          cache[region] = Parent::Then;2587        return Parent::Then;2588      }2589      if (&op.getElseRegion() == toCheck) {2590        for (Region *region : seen)2591          cache[region] = Parent::Else;2592        return Parent::Else;2593      }2594      toCheck = toCheck->getParentRegion();2595    }2596 2597    for (Region *region : seen)2598      cache[region] = Parent::None;2599    return Parent::None;2600  }2601 2602  LogicalResult matchAndRewrite(IfOp op,2603                                PatternRewriter &rewriter) const override {2604    // Early exit if the condition is constant since replacing a constant2605    // in the body with another constant isn't a simplification.2606    if (matchPattern(op.getCondition(), m_Constant()))2607      return failure();2608 2609    bool changed = false;2610    mlir::Type i1Ty = rewriter.getI1Type();2611 2612    // These variables serve to prevent creating duplicate constants2613    // and hold constant true or false values.2614    Value constantTrue = nullptr;2615    Value constantFalse = nullptr;2616 2617    DenseMap<Region *, Parent> cache;2618    for (OpOperand &use :2619         llvm::make_early_inc_range(op.getCondition().getUses())) {2620      switch (getParentType(use.getOwner()->getParentRegion(), op, cache,2621                            op.getCondition().getParentRegion())) {2622      case Parent::Then: {2623        changed = true;2624 2625        if (!constantTrue)2626          constantTrue = arith::ConstantOp::create(2627              rewriter, op.getLoc(), i1Ty, rewriter.getIntegerAttr(i1Ty, 1));2628 2629        rewriter.modifyOpInPlace(use.getOwner(),2630                                 [&]() { use.set(constantTrue); });2631        break;2632      }2633      case Parent::Else: {2634        changed = true;2635 2636        if (!constantFalse)2637          constantFalse = arith::ConstantOp::create(2638              rewriter, op.getLoc(), i1Ty, rewriter.getIntegerAttr(i1Ty, 0));2639 2640        rewriter.modifyOpInPlace(use.getOwner(),2641                                 [&]() { use.set(constantFalse); });2642        break;2643      }2644      case Parent::None:2645        break;2646      }2647    }2648 2649    return success(changed);2650  }2651};2652 2653/// Remove any statements from an if that are equivalent to the condition2654/// or its negation. For example:2655///2656///    %res:2 = scf.if %cmp {2657///       yield something(), true2658///    } else {2659///       yield something2(), false2660///    }2661///    print(%res#1)2662///2663///  becomes2664///    %res = scf.if %cmp {2665///       yield something()2666///    } else {2667///       yield something2()2668///    }2669///    print(%cmp)2670///2671/// Additionally if both branches yield the same value, replace all uses2672/// of the result with the yielded value.2673///2674///    %res:2 = scf.if %cmp {2675///       yield something(), %arg12676///    } else {2677///       yield something2(), %arg12678///    }2679///    print(%res#1)2680///2681///  becomes2682///    %res = scf.if %cmp {2683///       yield something()2684///    } else {2685///       yield something2()2686///    }2687///    print(%arg1)2688///2689struct ReplaceIfYieldWithConditionOrValue : public OpRewritePattern<IfOp> {2690  using OpRewritePattern<IfOp>::OpRewritePattern;2691 2692  LogicalResult matchAndRewrite(IfOp op,2693                                PatternRewriter &rewriter) const override {2694    // Early exit if there are no results that could be replaced.2695    if (op.getNumResults() == 0)2696      return failure();2697 2698    auto trueYield =2699        cast<scf::YieldOp>(op.getThenRegion().back().getTerminator());2700    auto falseYield =2701        cast<scf::YieldOp>(op.getElseRegion().back().getTerminator());2702 2703    rewriter.setInsertionPoint(op->getBlock(),2704                               op.getOperation()->getIterator());2705    bool changed = false;2706    Type i1Ty = rewriter.getI1Type();2707    for (auto [trueResult, falseResult, opResult] :2708         llvm::zip(trueYield.getResults(), falseYield.getResults(),2709                   op.getResults())) {2710      if (trueResult == falseResult) {2711        if (!opResult.use_empty()) {2712          opResult.replaceAllUsesWith(trueResult);2713          changed = true;2714        }2715        continue;2716      }2717 2718      BoolAttr trueYield, falseYield;2719      if (!matchPattern(trueResult, m_Constant(&trueYield)) ||2720          !matchPattern(falseResult, m_Constant(&falseYield)))2721        continue;2722 2723      bool trueVal = trueYield.getValue();2724      bool falseVal = falseYield.getValue();2725      if (!trueVal && falseVal) {2726        if (!opResult.use_empty()) {2727          Dialect *constDialect = trueResult.getDefiningOp()->getDialect();2728          Value notCond = arith::XOrIOp::create(2729              rewriter, op.getLoc(), op.getCondition(),2730              constDialect2731                  ->materializeConstant(rewriter,2732                                        rewriter.getIntegerAttr(i1Ty, 1), i1Ty,2733                                        op.getLoc())2734                  ->getResult(0));2735          opResult.replaceAllUsesWith(notCond);2736          changed = true;2737        }2738      }2739      if (trueVal && !falseVal) {2740        if (!opResult.use_empty()) {2741          opResult.replaceAllUsesWith(op.getCondition());2742          changed = true;2743        }2744      }2745    }2746    return success(changed);2747  }2748};2749 2750/// Merge any consecutive scf.if's with the same condition.2751///2752///    scf.if %cond {2753///       firstCodeTrue();...2754///    } else {2755///       firstCodeFalse();...2756///    }2757///    %res = scf.if %cond {2758///       secondCodeTrue();...2759///    } else {2760///       secondCodeFalse();...2761///    }2762///2763///  becomes2764///    %res = scf.if %cmp {2765///       firstCodeTrue();...2766///       secondCodeTrue();...2767///    } else {2768///       firstCodeFalse();...2769///       secondCodeFalse();...2770///    }2771struct CombineIfs : public OpRewritePattern<IfOp> {2772  using OpRewritePattern<IfOp>::OpRewritePattern;2773 2774  LogicalResult matchAndRewrite(IfOp nextIf,2775                                PatternRewriter &rewriter) const override {2776    Block *parent = nextIf->getBlock();2777    if (nextIf == &parent->front())2778      return failure();2779 2780    auto prevIf = dyn_cast<IfOp>(nextIf->getPrevNode());2781    if (!prevIf)2782      return failure();2783 2784    // Determine the logical then/else blocks when prevIf's2785    // condition is used. Null means the block does not exist2786    // in that case (e.g. empty else). If neither of these2787    // are set, the two conditions cannot be compared.2788    Block *nextThen = nullptr;2789    Block *nextElse = nullptr;2790    if (nextIf.getCondition() == prevIf.getCondition()) {2791      nextThen = nextIf.thenBlock();2792      if (!nextIf.getElseRegion().empty())2793        nextElse = nextIf.elseBlock();2794    }2795    if (arith::XOrIOp notv =2796            nextIf.getCondition().getDefiningOp<arith::XOrIOp>()) {2797      if (notv.getLhs() == prevIf.getCondition() &&2798          matchPattern(notv.getRhs(), m_One())) {2799        nextElse = nextIf.thenBlock();2800        if (!nextIf.getElseRegion().empty())2801          nextThen = nextIf.elseBlock();2802      }2803    }2804    if (arith::XOrIOp notv =2805            prevIf.getCondition().getDefiningOp<arith::XOrIOp>()) {2806      if (notv.getLhs() == nextIf.getCondition() &&2807          matchPattern(notv.getRhs(), m_One())) {2808        nextElse = nextIf.thenBlock();2809        if (!nextIf.getElseRegion().empty())2810          nextThen = nextIf.elseBlock();2811      }2812    }2813 2814    if (!nextThen && !nextElse)2815      return failure();2816 2817    SmallVector<Value> prevElseYielded;2818    if (!prevIf.getElseRegion().empty())2819      prevElseYielded = prevIf.elseYield().getOperands();2820    // Replace all uses of return values of op within nextIf with the2821    // corresponding yields2822    for (auto it : llvm::zip(prevIf.getResults(),2823                             prevIf.thenYield().getOperands(), prevElseYielded))2824      for (OpOperand &use :2825           llvm::make_early_inc_range(std::get<0>(it).getUses())) {2826        if (nextThen && nextThen->getParent()->isAncestor(2827                            use.getOwner()->getParentRegion())) {2828          rewriter.startOpModification(use.getOwner());2829          use.set(std::get<1>(it));2830          rewriter.finalizeOpModification(use.getOwner());2831        } else if (nextElse && nextElse->getParent()->isAncestor(2832                                   use.getOwner()->getParentRegion())) {2833          rewriter.startOpModification(use.getOwner());2834          use.set(std::get<2>(it));2835          rewriter.finalizeOpModification(use.getOwner());2836        }2837      }2838 2839    SmallVector<Type> mergedTypes(prevIf.getResultTypes());2840    llvm::append_range(mergedTypes, nextIf.getResultTypes());2841 2842    IfOp combinedIf = IfOp::create(rewriter, nextIf.getLoc(), mergedTypes,2843                                   prevIf.getCondition(), /*hasElse=*/false);2844    rewriter.eraseBlock(&combinedIf.getThenRegion().back());2845 2846    rewriter.inlineRegionBefore(prevIf.getThenRegion(),2847                                combinedIf.getThenRegion(),2848                                combinedIf.getThenRegion().begin());2849 2850    if (nextThen) {2851      YieldOp thenYield = combinedIf.thenYield();2852      YieldOp thenYield2 = cast<YieldOp>(nextThen->getTerminator());2853      rewriter.mergeBlocks(nextThen, combinedIf.thenBlock());2854      rewriter.setInsertionPointToEnd(combinedIf.thenBlock());2855 2856      SmallVector<Value> mergedYields(thenYield.getOperands());2857      llvm::append_range(mergedYields, thenYield2.getOperands());2858      YieldOp::create(rewriter, thenYield2.getLoc(), mergedYields);2859      rewriter.eraseOp(thenYield);2860      rewriter.eraseOp(thenYield2);2861    }2862 2863    rewriter.inlineRegionBefore(prevIf.getElseRegion(),2864                                combinedIf.getElseRegion(),2865                                combinedIf.getElseRegion().begin());2866 2867    if (nextElse) {2868      if (combinedIf.getElseRegion().empty()) {2869        rewriter.inlineRegionBefore(*nextElse->getParent(),2870                                    combinedIf.getElseRegion(),2871                                    combinedIf.getElseRegion().begin());2872      } else {2873        YieldOp elseYield = combinedIf.elseYield();2874        YieldOp elseYield2 = cast<YieldOp>(nextElse->getTerminator());2875        rewriter.mergeBlocks(nextElse, combinedIf.elseBlock());2876 2877        rewriter.setInsertionPointToEnd(combinedIf.elseBlock());2878 2879        SmallVector<Value> mergedElseYields(elseYield.getOperands());2880        llvm::append_range(mergedElseYields, elseYield2.getOperands());2881 2882        YieldOp::create(rewriter, elseYield2.getLoc(), mergedElseYields);2883        rewriter.eraseOp(elseYield);2884        rewriter.eraseOp(elseYield2);2885      }2886    }2887 2888    SmallVector<Value> prevValues;2889    SmallVector<Value> nextValues;2890    for (const auto &pair : llvm::enumerate(combinedIf.getResults())) {2891      if (pair.index() < prevIf.getNumResults())2892        prevValues.push_back(pair.value());2893      else2894        nextValues.push_back(pair.value());2895    }2896    rewriter.replaceOp(prevIf, prevValues);2897    rewriter.replaceOp(nextIf, nextValues);2898    return success();2899  }2900};2901 2902/// Pattern to remove an empty else branch.2903struct RemoveEmptyElseBranch : public OpRewritePattern<IfOp> {2904  using OpRewritePattern<IfOp>::OpRewritePattern;2905 2906  LogicalResult matchAndRewrite(IfOp ifOp,2907                                PatternRewriter &rewriter) const override {2908    // Cannot remove else region when there are operation results.2909    if (ifOp.getNumResults())2910      return failure();2911    Block *elseBlock = ifOp.elseBlock();2912    if (!elseBlock || !llvm::hasSingleElement(*elseBlock))2913      return failure();2914    auto newIfOp = rewriter.cloneWithoutRegions(ifOp);2915    rewriter.inlineRegionBefore(ifOp.getThenRegion(), newIfOp.getThenRegion(),2916                                newIfOp.getThenRegion().begin());2917    rewriter.eraseOp(ifOp);2918    return success();2919  }2920};2921 2922/// Convert nested `if`s into `arith.andi` + single `if`.2923///2924///    scf.if %arg0 {2925///      scf.if %arg1 {2926///        ...2927///        scf.yield2928///      }2929///      scf.yield2930///    }2931///  becomes2932///2933///    %0 = arith.andi %arg0, %arg12934///    scf.if %0 {2935///      ...2936///      scf.yield2937///    }2938struct CombineNestedIfs : public OpRewritePattern<IfOp> {2939  using OpRewritePattern<IfOp>::OpRewritePattern;2940 2941  LogicalResult matchAndRewrite(IfOp op,2942                                PatternRewriter &rewriter) const override {2943    auto nestedOps = op.thenBlock()->without_terminator();2944    // Nested `if` must be the only op in block.2945    if (!llvm::hasSingleElement(nestedOps))2946      return failure();2947 2948    // If there is an else block, it can only yield2949    if (op.elseBlock() && !llvm::hasSingleElement(*op.elseBlock()))2950      return failure();2951 2952    auto nestedIf = dyn_cast<IfOp>(*nestedOps.begin());2953    if (!nestedIf)2954      return failure();2955 2956    if (nestedIf.elseBlock() && !llvm::hasSingleElement(*nestedIf.elseBlock()))2957      return failure();2958 2959    SmallVector<Value> thenYield(op.thenYield().getOperands());2960    SmallVector<Value> elseYield;2961    if (op.elseBlock())2962      llvm::append_range(elseYield, op.elseYield().getOperands());2963 2964    // A list of indices for which we should upgrade the value yielded2965    // in the else to a select.2966    SmallVector<unsigned> elseYieldsToUpgradeToSelect;2967 2968    // If the outer scf.if yields a value produced by the inner scf.if,2969    // only permit combining if the value yielded when the condition2970    // is false in the outer scf.if is the same value yielded when the2971    // inner scf.if condition is false.2972    // Note that the array access to elseYield will not go out of bounds2973    // since it must have the same length as thenYield, since they both2974    // come from the same scf.if.2975    for (const auto &tup : llvm::enumerate(thenYield)) {2976      if (tup.value().getDefiningOp() == nestedIf) {2977        auto nestedIdx = llvm::cast<OpResult>(tup.value()).getResultNumber();2978        if (nestedIf.elseYield().getOperand(nestedIdx) !=2979            elseYield[tup.index()]) {2980          return failure();2981        }2982        // If the correctness test passes, we will yield2983        // corresponding value from the inner scf.if2984        thenYield[tup.index()] = nestedIf.thenYield().getOperand(nestedIdx);2985        continue;2986      }2987 2988      // Otherwise, we need to ensure the else block of the combined2989      // condition still returns the same value when the outer condition is2990      // true and the inner condition is false. This can be accomplished if2991      // the then value is defined outside the outer scf.if and we replace the2992      // value with a select that considers just the outer condition. Since2993      // the else region contains just the yield, its yielded value is2994      // defined outside the scf.if, by definition.2995 2996      // If the then value is defined within the scf.if, bail.2997      if (tup.value().getParentRegion() == &op.getThenRegion()) {2998        return failure();2999      }3000      elseYieldsToUpgradeToSelect.push_back(tup.index());3001    }3002 3003    Location loc = op.getLoc();3004    Value newCondition = arith::AndIOp::create(rewriter, loc, op.getCondition(),3005                                               nestedIf.getCondition());3006    auto newIf = IfOp::create(rewriter, loc, op.getResultTypes(), newCondition);3007    Block *newIfBlock = rewriter.createBlock(&newIf.getThenRegion());3008 3009    SmallVector<Value> results;3010    llvm::append_range(results, newIf.getResults());3011    rewriter.setInsertionPoint(newIf);3012 3013    for (auto idx : elseYieldsToUpgradeToSelect)3014      results[idx] =3015          arith::SelectOp::create(rewriter, op.getLoc(), op.getCondition(),3016                                  thenYield[idx], elseYield[idx]);3017 3018    rewriter.mergeBlocks(nestedIf.thenBlock(), newIfBlock);3019    rewriter.setInsertionPointToEnd(newIf.thenBlock());3020    rewriter.replaceOpWithNewOp<YieldOp>(newIf.thenYield(), thenYield);3021    if (!elseYield.empty()) {3022      rewriter.createBlock(&newIf.getElseRegion());3023      rewriter.setInsertionPointToEnd(newIf.elseBlock());3024      YieldOp::create(rewriter, loc, elseYield);3025    }3026    rewriter.replaceOp(op, results);3027    return success();3028  }3029};3030 3031} // namespace3032 3033void IfOp::getCanonicalizationPatterns(RewritePatternSet &results,3034                                       MLIRContext *context) {3035  results.add<CombineIfs, CombineNestedIfs, ConditionPropagation,3036              ConvertTrivialIfToSelect, RemoveEmptyElseBranch,3037              RemoveStaticCondition, RemoveUnusedResults,3038              ReplaceIfYieldWithConditionOrValue>(context);3039}3040 3041Block *IfOp::thenBlock() { return &getThenRegion().back(); }3042YieldOp IfOp::thenYield() { return cast<YieldOp>(&thenBlock()->back()); }3043Block *IfOp::elseBlock() {3044  Region &r = getElseRegion();3045  if (r.empty())3046    return nullptr;3047  return &r.back();3048}3049YieldOp IfOp::elseYield() { return cast<YieldOp>(&elseBlock()->back()); }3050 3051//===----------------------------------------------------------------------===//3052// ParallelOp3053//===----------------------------------------------------------------------===//3054 3055void ParallelOp::build(3056    OpBuilder &builder, OperationState &result, ValueRange lowerBounds,3057    ValueRange upperBounds, ValueRange steps, ValueRange initVals,3058    function_ref<void(OpBuilder &, Location, ValueRange, ValueRange)>3059        bodyBuilderFn) {3060  result.addOperands(lowerBounds);3061  result.addOperands(upperBounds);3062  result.addOperands(steps);3063  result.addOperands(initVals);3064  result.addAttribute(3065      ParallelOp::getOperandSegmentSizeAttr(),3066      builder.getDenseI32ArrayAttr({static_cast<int32_t>(lowerBounds.size()),3067                                    static_cast<int32_t>(upperBounds.size()),3068                                    static_cast<int32_t>(steps.size()),3069                                    static_cast<int32_t>(initVals.size())}));3070  result.addTypes(initVals.getTypes());3071 3072  OpBuilder::InsertionGuard guard(builder);3073  unsigned numIVs = steps.size();3074  SmallVector<Type, 8> argTypes(numIVs, builder.getIndexType());3075  SmallVector<Location, 8> argLocs(numIVs, result.location);3076  Region *bodyRegion = result.addRegion();3077  Block *bodyBlock = builder.createBlock(bodyRegion, {}, argTypes, argLocs);3078 3079  if (bodyBuilderFn) {3080    builder.setInsertionPointToStart(bodyBlock);3081    bodyBuilderFn(builder, result.location,3082                  bodyBlock->getArguments().take_front(numIVs),3083                  bodyBlock->getArguments().drop_front(numIVs));3084  }3085  // Add terminator only if there are no reductions.3086  if (initVals.empty())3087    ParallelOp::ensureTerminator(*bodyRegion, builder, result.location);3088}3089 3090void ParallelOp::build(3091    OpBuilder &builder, OperationState &result, ValueRange lowerBounds,3092    ValueRange upperBounds, ValueRange steps,3093    function_ref<void(OpBuilder &, Location, ValueRange)> bodyBuilderFn) {3094  // Only pass a non-null wrapper if bodyBuilderFn is non-null itself. Make sure3095  // we don't capture a reference to a temporary by constructing the lambda at3096  // function level.3097  auto wrappedBuilderFn = [&bodyBuilderFn](OpBuilder &nestedBuilder,3098                                           Location nestedLoc, ValueRange ivs,3099                                           ValueRange) {3100    bodyBuilderFn(nestedBuilder, nestedLoc, ivs);3101  };3102  function_ref<void(OpBuilder &, Location, ValueRange, ValueRange)> wrapper;3103  if (bodyBuilderFn)3104    wrapper = wrappedBuilderFn;3105 3106  build(builder, result, lowerBounds, upperBounds, steps, ValueRange(),3107        wrapper);3108}3109 3110LogicalResult ParallelOp::verify() {3111  // Check that there is at least one value in lowerBound, upperBound and step.3112  // It is sufficient to test only step, because it is ensured already that the3113  // number of elements in lowerBound, upperBound and step are the same.3114  Operation::operand_range stepValues = getStep();3115  if (stepValues.empty())3116    return emitOpError(3117        "needs at least one tuple element for lowerBound, upperBound and step");3118 3119  // Check whether all constant step values are positive.3120  for (Value stepValue : stepValues)3121    if (auto cst = getConstantIntValue(stepValue))3122      if (*cst <= 0)3123        return emitOpError("constant step operand must be positive");3124 3125  // Check that the body defines the same number of block arguments as the3126  // number of tuple elements in step.3127  Block *body = getBody();3128  if (body->getNumArguments() != stepValues.size())3129    return emitOpError() << "expects the same number of induction variables: "3130                         << body->getNumArguments()3131                         << " as bound and step values: " << stepValues.size();3132  for (auto arg : body->getArguments())3133    if (!arg.getType().isIndex())3134      return emitOpError(3135          "expects arguments for the induction variable to be of index type");3136 3137  // Check that the terminator is an scf.reduce op.3138  auto reduceOp = verifyAndGetTerminator<scf::ReduceOp>(3139      *this, getRegion(), "expects body to terminate with 'scf.reduce'");3140  if (!reduceOp)3141    return failure();3142 3143  // Check that the number of results is the same as the number of reductions.3144  auto resultsSize = getResults().size();3145  auto reductionsSize = reduceOp.getReductions().size();3146  auto initValsSize = getInitVals().size();3147  if (resultsSize != reductionsSize)3148    return emitOpError() << "expects number of results: " << resultsSize3149                         << " to be the same as number of reductions: "3150                         << reductionsSize;3151  if (resultsSize != initValsSize)3152    return emitOpError() << "expects number of results: " << resultsSize3153                         << " to be the same as number of initial values: "3154                         << initValsSize;3155 3156  // Check that the types of the results and reductions are the same.3157  for (int64_t i = 0; i < static_cast<int64_t>(reductionsSize); ++i) {3158    auto resultType = getOperation()->getResult(i).getType();3159    auto reductionOperandType = reduceOp.getOperands()[i].getType();3160    if (resultType != reductionOperandType)3161      return reduceOp.emitOpError()3162             << "expects type of " << i3163             << "-th reduction operand: " << reductionOperandType3164             << " to be the same as the " << i3165             << "-th result type: " << resultType;3166  }3167  return success();3168}3169 3170ParseResult ParallelOp::parse(OpAsmParser &parser, OperationState &result) {3171  auto &builder = parser.getBuilder();3172  // Parse an opening `(` followed by induction variables followed by `)`3173  SmallVector<OpAsmParser::Argument, 4> ivs;3174  if (parser.parseArgumentList(ivs, OpAsmParser::Delimiter::Paren))3175    return failure();3176 3177  // Parse loop bounds.3178  SmallVector<OpAsmParser::UnresolvedOperand, 4> lower;3179  if (parser.parseEqual() ||3180      parser.parseOperandList(lower, ivs.size(),3181                              OpAsmParser::Delimiter::Paren) ||3182      parser.resolveOperands(lower, builder.getIndexType(), result.operands))3183    return failure();3184 3185  SmallVector<OpAsmParser::UnresolvedOperand, 4> upper;3186  if (parser.parseKeyword("to") ||3187      parser.parseOperandList(upper, ivs.size(),3188                              OpAsmParser::Delimiter::Paren) ||3189      parser.resolveOperands(upper, builder.getIndexType(), result.operands))3190    return failure();3191 3192  // Parse step values.3193  SmallVector<OpAsmParser::UnresolvedOperand, 4> steps;3194  if (parser.parseKeyword("step") ||3195      parser.parseOperandList(steps, ivs.size(),3196                              OpAsmParser::Delimiter::Paren) ||3197      parser.resolveOperands(steps, builder.getIndexType(), result.operands))3198    return failure();3199 3200  // Parse init values.3201  SmallVector<OpAsmParser::UnresolvedOperand, 4> initVals;3202  if (succeeded(parser.parseOptionalKeyword("init"))) {3203    if (parser.parseOperandList(initVals, OpAsmParser::Delimiter::Paren))3204      return failure();3205  }3206 3207  // Parse optional results in case there is a reduce.3208  if (parser.parseOptionalArrowTypeList(result.types))3209    return failure();3210 3211  // Now parse the body.3212  Region *body = result.addRegion();3213  for (auto &iv : ivs)3214    iv.type = builder.getIndexType();3215  if (parser.parseRegion(*body, ivs))3216    return failure();3217 3218  // Set `operandSegmentSizes` attribute.3219  result.addAttribute(3220      ParallelOp::getOperandSegmentSizeAttr(),3221      builder.getDenseI32ArrayAttr({static_cast<int32_t>(lower.size()),3222                                    static_cast<int32_t>(upper.size()),3223                                    static_cast<int32_t>(steps.size()),3224                                    static_cast<int32_t>(initVals.size())}));3225 3226  // Parse attributes.3227  if (parser.parseOptionalAttrDict(result.attributes) ||3228      parser.resolveOperands(initVals, result.types, parser.getNameLoc(),3229                             result.operands))3230    return failure();3231 3232  // Add a terminator if none was parsed.3233  ParallelOp::ensureTerminator(*body, builder, result.location);3234  return success();3235}3236 3237void ParallelOp::print(OpAsmPrinter &p) {3238  p << " (" << getBody()->getArguments() << ") = (" << getLowerBound()3239    << ") to (" << getUpperBound() << ") step (" << getStep() << ")";3240  if (!getInitVals().empty())3241    p << " init (" << getInitVals() << ")";3242  p.printOptionalArrowTypeList(getResultTypes());3243  p << ' ';3244  p.printRegion(getRegion(), /*printEntryBlockArgs=*/false);3245  p.printOptionalAttrDict(3246      (*this)->getAttrs(),3247      /*elidedAttrs=*/ParallelOp::getOperandSegmentSizeAttr());3248}3249 3250SmallVector<Region *> ParallelOp::getLoopRegions() { return {&getRegion()}; }3251 3252std::optional<SmallVector<Value>> ParallelOp::getLoopInductionVars() {3253  return SmallVector<Value>{getBody()->getArguments()};3254}3255 3256std::optional<SmallVector<OpFoldResult>> ParallelOp::getLoopLowerBounds() {3257  return getLowerBound();3258}3259 3260std::optional<SmallVector<OpFoldResult>> ParallelOp::getLoopUpperBounds() {3261  return getUpperBound();3262}3263 3264std::optional<SmallVector<OpFoldResult>> ParallelOp::getLoopSteps() {3265  return getStep();3266}3267 3268ParallelOp mlir::scf::getParallelForInductionVarOwner(Value val) {3269  auto ivArg = llvm::dyn_cast<BlockArgument>(val);3270  if (!ivArg)3271    return ParallelOp();3272  assert(ivArg.getOwner() && "unlinked block argument");3273  auto *containingOp = ivArg.getOwner()->getParentOp();3274  return dyn_cast<ParallelOp>(containingOp);3275}3276 3277namespace {3278// Collapse loop dimensions that perform a single iteration.3279struct ParallelOpSingleOrZeroIterationDimsFolder3280    : public OpRewritePattern<ParallelOp> {3281  using OpRewritePattern<ParallelOp>::OpRewritePattern;3282 3283  LogicalResult matchAndRewrite(ParallelOp op,3284                                PatternRewriter &rewriter) const override {3285    Location loc = op.getLoc();3286 3287    // Compute new loop bounds that omit all single-iteration loop dimensions.3288    SmallVector<Value> newLowerBounds, newUpperBounds, newSteps;3289    IRMapping mapping;3290    for (auto [lb, ub, step, iv] :3291         llvm::zip(op.getLowerBound(), op.getUpperBound(), op.getStep(),3292                   op.getInductionVars())) {3293      auto numIterations =3294          constantTripCount(lb, ub, step, /*isSigned=*/true, computeUbMinusLb);3295      if (numIterations.has_value()) {3296        // Remove the loop if it performs zero iterations.3297        if (*numIterations == 0) {3298          rewriter.replaceOp(op, op.getInitVals());3299          return success();3300        }3301        // Replace the loop induction variable by the lower bound if the loop3302        // performs a single iteration. Otherwise, copy the loop bounds.3303        if (*numIterations == 1) {3304          mapping.map(iv, getValueOrCreateConstantIndexOp(rewriter, loc, lb));3305          continue;3306        }3307      }3308      newLowerBounds.push_back(lb);3309      newUpperBounds.push_back(ub);3310      newSteps.push_back(step);3311    }3312    // Exit if none of the loop dimensions perform a single iteration.3313    if (newLowerBounds.size() == op.getLowerBound().size())3314      return failure();3315 3316    if (newLowerBounds.empty()) {3317      // All of the loop dimensions perform a single iteration. Inline3318      // loop body and nested ReduceOp's3319      SmallVector<Value> results;3320      results.reserve(op.getInitVals().size());3321      for (auto &bodyOp : op.getBody()->without_terminator())3322        rewriter.clone(bodyOp, mapping);3323      auto reduceOp = cast<ReduceOp>(op.getBody()->getTerminator());3324      for (int64_t i = 0, e = reduceOp.getReductions().size(); i < e; ++i) {3325        Block &reduceBlock = reduceOp.getReductions()[i].front();3326        auto initValIndex = results.size();3327        mapping.map(reduceBlock.getArgument(0), op.getInitVals()[initValIndex]);3328        mapping.map(reduceBlock.getArgument(1),3329                    mapping.lookupOrDefault(reduceOp.getOperands()[i]));3330        for (auto &reduceBodyOp : reduceBlock.without_terminator())3331          rewriter.clone(reduceBodyOp, mapping);3332 3333        auto result = mapping.lookupOrDefault(3334            cast<ReduceReturnOp>(reduceBlock.getTerminator()).getResult());3335        results.push_back(result);3336      }3337 3338      rewriter.replaceOp(op, results);3339      return success();3340    }3341    // Replace the parallel loop by lower-dimensional parallel loop.3342    auto newOp =3343        ParallelOp::create(rewriter, op.getLoc(), newLowerBounds,3344                           newUpperBounds, newSteps, op.getInitVals(), nullptr);3345    // Erase the empty block that was inserted by the builder.3346    rewriter.eraseBlock(newOp.getBody());3347    // Clone the loop body and remap the block arguments of the collapsed loops3348    // (inlining does not support a cancellable block argument mapping).3349    rewriter.cloneRegionBefore(op.getRegion(), newOp.getRegion(),3350                               newOp.getRegion().begin(), mapping);3351    rewriter.replaceOp(op, newOp.getResults());3352    return success();3353  }3354};3355 3356struct MergeNestedParallelLoops : public OpRewritePattern<ParallelOp> {3357  using OpRewritePattern<ParallelOp>::OpRewritePattern;3358 3359  LogicalResult matchAndRewrite(ParallelOp op,3360                                PatternRewriter &rewriter) const override {3361    Block &outerBody = *op.getBody();3362    if (!llvm::hasSingleElement(outerBody.without_terminator()))3363      return failure();3364 3365    auto innerOp = dyn_cast<ParallelOp>(outerBody.front());3366    if (!innerOp)3367      return failure();3368 3369    for (auto val : outerBody.getArguments())3370      if (llvm::is_contained(innerOp.getLowerBound(), val) ||3371          llvm::is_contained(innerOp.getUpperBound(), val) ||3372          llvm::is_contained(innerOp.getStep(), val))3373        return failure();3374 3375    // Reductions are not supported yet.3376    if (!op.getInitVals().empty() || !innerOp.getInitVals().empty())3377      return failure();3378 3379    auto bodyBuilder = [&](OpBuilder &builder, Location /*loc*/,3380                           ValueRange iterVals, ValueRange) {3381      Block &innerBody = *innerOp.getBody();3382      assert(iterVals.size() ==3383             (outerBody.getNumArguments() + innerBody.getNumArguments()));3384      IRMapping mapping;3385      mapping.map(outerBody.getArguments(),3386                  iterVals.take_front(outerBody.getNumArguments()));3387      mapping.map(innerBody.getArguments(),3388                  iterVals.take_back(innerBody.getNumArguments()));3389      for (Operation &op : innerBody.without_terminator())3390        builder.clone(op, mapping);3391    };3392 3393    auto concatValues = [](const auto &first, const auto &second) {3394      SmallVector<Value> ret;3395      ret.reserve(first.size() + second.size());3396      ret.assign(first.begin(), first.end());3397      ret.append(second.begin(), second.end());3398      return ret;3399    };3400 3401    auto newLowerBounds =3402        concatValues(op.getLowerBound(), innerOp.getLowerBound());3403    auto newUpperBounds =3404        concatValues(op.getUpperBound(), innerOp.getUpperBound());3405    auto newSteps = concatValues(op.getStep(), innerOp.getStep());3406 3407    rewriter.replaceOpWithNewOp<ParallelOp>(op, newLowerBounds, newUpperBounds,3408                                            newSteps, ValueRange(),3409                                            bodyBuilder);3410    return success();3411  }3412};3413 3414} // namespace3415 3416void ParallelOp::getCanonicalizationPatterns(RewritePatternSet &results,3417                                             MLIRContext *context) {3418  results3419      .add<ParallelOpSingleOrZeroIterationDimsFolder, MergeNestedParallelLoops>(3420          context);3421}3422 3423/// Given the region at `index`, or the parent operation if `index` is None,3424/// return the successor regions. These are the regions that may be selected3425/// during the flow of control. `operands` is a set of optional attributes that3426/// correspond to a constant value for each operand, or null if that operand is3427/// not a constant.3428void ParallelOp::getSuccessorRegions(3429    RegionBranchPoint point, SmallVectorImpl<RegionSuccessor> &regions) {3430  // Both the operation itself and the region may be branching into the body or3431  // back into the operation itself. It is possible for loop not to enter the3432  // body.3433  regions.push_back(RegionSuccessor(&getRegion()));3434  regions.push_back(RegionSuccessor(3435      getOperation(), ResultRange{getResults().end(), getResults().end()}));3436}3437 3438//===----------------------------------------------------------------------===//3439// ReduceOp3440//===----------------------------------------------------------------------===//3441 3442void ReduceOp::build(OpBuilder &builder, OperationState &result) {}3443 3444void ReduceOp::build(OpBuilder &builder, OperationState &result,3445                     ValueRange operands) {3446  result.addOperands(operands);3447  for (Value v : operands) {3448    OpBuilder::InsertionGuard guard(builder);3449    Region *bodyRegion = result.addRegion();3450    builder.createBlock(bodyRegion, {},3451                        ArrayRef<Type>{v.getType(), v.getType()},3452                        {result.location, result.location});3453  }3454}3455 3456LogicalResult ReduceOp::verifyRegions() {3457  // The region of a ReduceOp has two arguments of the same type as its3458  // corresponding operand.3459  for (int64_t i = 0, e = getReductions().size(); i < e; ++i) {3460    auto type = getOperands()[i].getType();3461    Block &block = getReductions()[i].front();3462    if (block.empty())3463      return emitOpError() << i << "-th reduction has an empty body";3464    if (block.getNumArguments() != 2 ||3465        llvm::any_of(block.getArguments(), [&](const BlockArgument &arg) {3466          return arg.getType() != type;3467        }))3468      return emitOpError() << "expected two block arguments with type " << type3469                           << " in the " << i << "-th reduction region";3470 3471    // Check that the block is terminated by a ReduceReturnOp.3472    if (!isa<ReduceReturnOp>(block.getTerminator()))3473      return emitOpError("reduction bodies must be terminated with an "3474                         "'scf.reduce.return' op");3475  }3476 3477  return success();3478}3479 3480MutableOperandRange3481ReduceOp::getMutableSuccessorOperands(RegionSuccessor point) {3482  // No operands are forwarded to the next iteration.3483  return MutableOperandRange(getOperation(), /*start=*/0, /*length=*/0);3484}3485 3486//===----------------------------------------------------------------------===//3487// ReduceReturnOp3488//===----------------------------------------------------------------------===//3489 3490LogicalResult ReduceReturnOp::verify() {3491  // The type of the return value should be the same type as the types of the3492  // block arguments of the reduction body.3493  Block *reductionBody = getOperation()->getBlock();3494  // Should already be verified by an op trait.3495  assert(isa<ReduceOp>(reductionBody->getParentOp()) && "expected scf.reduce");3496  Type expectedResultType = reductionBody->getArgument(0).getType();3497  if (expectedResultType != getResult().getType())3498    return emitOpError() << "must have type " << expectedResultType3499                         << " (the type of the reduction inputs)";3500  return success();3501}3502 3503//===----------------------------------------------------------------------===//3504// WhileOp3505//===----------------------------------------------------------------------===//3506 3507void WhileOp::build(::mlir::OpBuilder &odsBuilder,3508                    ::mlir::OperationState &odsState, TypeRange resultTypes,3509                    ValueRange inits, BodyBuilderFn beforeBuilder,3510                    BodyBuilderFn afterBuilder) {3511  odsState.addOperands(inits);3512  odsState.addTypes(resultTypes);3513 3514  OpBuilder::InsertionGuard guard(odsBuilder);3515 3516  // Build before region.3517  SmallVector<Location, 4> beforeArgLocs;3518  beforeArgLocs.reserve(inits.size());3519  for (Value operand : inits) {3520    beforeArgLocs.push_back(operand.getLoc());3521  }3522 3523  Region *beforeRegion = odsState.addRegion();3524  Block *beforeBlock = odsBuilder.createBlock(beforeRegion, /*insertPt=*/{},3525                                              inits.getTypes(), beforeArgLocs);3526  if (beforeBuilder)3527    beforeBuilder(odsBuilder, odsState.location, beforeBlock->getArguments());3528 3529  // Build after region.3530  SmallVector<Location, 4> afterArgLocs(resultTypes.size(), odsState.location);3531 3532  Region *afterRegion = odsState.addRegion();3533  Block *afterBlock = odsBuilder.createBlock(afterRegion, /*insertPt=*/{},3534                                             resultTypes, afterArgLocs);3535 3536  if (afterBuilder)3537    afterBuilder(odsBuilder, odsState.location, afterBlock->getArguments());3538}3539 3540ConditionOp WhileOp::getConditionOp() {3541  return cast<ConditionOp>(getBeforeBody()->getTerminator());3542}3543 3544YieldOp WhileOp::getYieldOp() {3545  return cast<YieldOp>(getAfterBody()->getTerminator());3546}3547 3548std::optional<MutableArrayRef<OpOperand>> WhileOp::getYieldedValuesMutable() {3549  return getYieldOp().getResultsMutable();3550}3551 3552Block::BlockArgListType WhileOp::getBeforeArguments() {3553  return getBeforeBody()->getArguments();3554}3555 3556Block::BlockArgListType WhileOp::getAfterArguments() {3557  return getAfterBody()->getArguments();3558}3559 3560Block::BlockArgListType WhileOp::getRegionIterArgs() {3561  return getBeforeArguments();3562}3563 3564OperandRange WhileOp::getEntrySuccessorOperands(RegionSuccessor successor) {3565  assert(successor.getSuccessor() == &getBefore() &&3566         "WhileOp is expected to branch only to the first region");3567  return getInits();3568}3569 3570void WhileOp::getSuccessorRegions(RegionBranchPoint point,3571                                  SmallVectorImpl<RegionSuccessor> &regions) {3572  // The parent op always branches to the condition region.3573  if (point.isParent()) {3574    regions.emplace_back(&getBefore(), getBefore().getArguments());3575    return;3576  }3577 3578  assert(llvm::is_contained(3579             {&getAfter(), &getBefore()},3580             point.getTerminatorPredecessorOrNull()->getParentRegion()) &&3581         "there are only two regions in a WhileOp");3582  // The body region always branches back to the condition region.3583  if (point.getTerminatorPredecessorOrNull()->getParentRegion() ==3584      &getAfter()) {3585    regions.emplace_back(&getBefore(), getBefore().getArguments());3586    return;3587  }3588 3589  regions.emplace_back(getOperation(), getResults());3590  regions.emplace_back(&getAfter(), getAfter().getArguments());3591}3592 3593SmallVector<Region *> WhileOp::getLoopRegions() {3594  return {&getBefore(), &getAfter()};3595}3596 3597/// Parses a `while` op.3598///3599/// op ::= `scf.while` assignments `:` function-type region `do` region3600///         `attributes` attribute-dict3601/// initializer ::= /* empty */ | `(` assignment-list `)`3602/// assignment-list ::= assignment | assignment `,` assignment-list3603/// assignment ::= ssa-value `=` ssa-value3604ParseResult scf::WhileOp::parse(OpAsmParser &parser, OperationState &result) {3605  SmallVector<OpAsmParser::Argument, 4> regionArgs;3606  SmallVector<OpAsmParser::UnresolvedOperand, 4> operands;3607  Region *before = result.addRegion();3608  Region *after = result.addRegion();3609 3610  OptionalParseResult listResult =3611      parser.parseOptionalAssignmentList(regionArgs, operands);3612  if (listResult.has_value() && failed(listResult.value()))3613    return failure();3614 3615  FunctionType functionType;3616  SMLoc typeLoc = parser.getCurrentLocation();3617  if (failed(parser.parseColonType(functionType)))3618    return failure();3619 3620  result.addTypes(functionType.getResults());3621 3622  if (functionType.getNumInputs() != operands.size()) {3623    return parser.emitError(typeLoc)3624           << "expected as many input types as operands " << "(expected "3625           << operands.size() << " got " << functionType.getNumInputs() << ")";3626  }3627 3628  // Resolve input operands.3629  if (failed(parser.resolveOperands(operands, functionType.getInputs(),3630                                    parser.getCurrentLocation(),3631                                    result.operands)))3632    return failure();3633 3634  // Propagate the types into the region arguments.3635  for (size_t i = 0, e = regionArgs.size(); i != e; ++i)3636    regionArgs[i].type = functionType.getInput(i);3637 3638  return failure(parser.parseRegion(*before, regionArgs) ||3639                 parser.parseKeyword("do") || parser.parseRegion(*after) ||3640                 parser.parseOptionalAttrDictWithKeyword(result.attributes));3641}3642 3643/// Prints a `while` op.3644void scf::WhileOp::print(OpAsmPrinter &p) {3645  printInitializationList(p, getBeforeArguments(), getInits(), " ");3646  p << " : ";3647  p.printFunctionalType(getInits().getTypes(), getResults().getTypes());3648  p << ' ';3649  p.printRegion(getBefore(), /*printEntryBlockArgs=*/false);3650  p << " do ";3651  p.printRegion(getAfter());3652  p.printOptionalAttrDictWithKeyword((*this)->getAttrs());3653}3654 3655/// Verifies that two ranges of types match, i.e. have the same number of3656/// entries and that types are pairwise equals. Reports errors on the given3657/// operation in case of mismatch.3658template <typename OpTy>3659static LogicalResult verifyTypeRangesMatch(OpTy op, TypeRange left,3660                                           TypeRange right, StringRef message) {3661  if (left.size() != right.size())3662    return op.emitOpError("expects the same number of ") << message;3663 3664  for (unsigned i = 0, e = left.size(); i < e; ++i) {3665    if (left[i] != right[i]) {3666      InFlightDiagnostic diag = op.emitOpError("expects the same types for ")3667                                << message;3668      diag.attachNote() << "for argument " << i << ", found " << left[i]3669                        << " and " << right[i];3670      return diag;3671    }3672  }3673 3674  return success();3675}3676 3677LogicalResult scf::WhileOp::verify() {3678  auto beforeTerminator = verifyAndGetTerminator<scf::ConditionOp>(3679      *this, getBefore(),3680      "expects the 'before' region to terminate with 'scf.condition'");3681  if (!beforeTerminator)3682    return failure();3683 3684  auto afterTerminator = verifyAndGetTerminator<scf::YieldOp>(3685      *this, getAfter(),3686      "expects the 'after' region to terminate with 'scf.yield'");3687  return success(afterTerminator != nullptr);3688}3689 3690namespace {3691/// Move a scf.if op that is directly before the scf.condition op in the while3692/// before region, and whose condition matches the condition of the3693/// scf.condition op, down into the while after region.3694///3695/// scf.while (..) : (...) -> ... {3696///  %additional_used_values = ...3697///  %cond = ...3698///  ...3699///  %res = scf.if %cond -> (...) {3700///    use(%additional_used_values)3701///    ... // then block3702///    scf.yield %then_value3703///  } else {3704///    scf.yield %else_value3705///  }3706///  scf.condition(%cond) %res, ...3707/// } do {3708/// ^bb0(%res_arg, ...):3709///    use(%res_arg)3710///    ...3711///3712/// becomes3713/// scf.while (..) : (...) -> ... {3714///  %additional_used_values = ...3715///  %cond = ...3716///  ...3717///  scf.condition(%cond) %else_value, ..., %additional_used_values3718/// } do {3719/// ^bb0(%res_arg ..., %additional_args): :3720///    use(%additional_args)3721///    ... // if then block3722///    use(%then_value)3723///    ...3724struct WhileMoveIfDown : public OpRewritePattern<scf::WhileOp> {3725  using OpRewritePattern<scf::WhileOp>::OpRewritePattern;3726 3727  LogicalResult matchAndRewrite(scf::WhileOp op,3728                                PatternRewriter &rewriter) const override {3729    auto conditionOp = op.getConditionOp();3730 3731    // Only support ifOp right before the condition at the moment. Relaxing this3732    // would require to:3733    // - check that the body does not have side-effects conflicting with3734    //    operations between the if and the condition.3735    // - check that results of the if operation are only used as arguments to3736    //    the condition.3737    auto ifOp = dyn_cast_or_null<scf::IfOp>(conditionOp->getPrevNode());3738 3739    // Check that the ifOp is directly before the conditionOp and that it3740    // matches the condition of the conditionOp. Also ensure that the ifOp has3741    // no else block with content, as that would complicate the transformation.3742    // TODO: support else blocks with content.3743    if (!ifOp || ifOp.getCondition() != conditionOp.getCondition() ||3744        (ifOp.elseBlock() && !ifOp.elseBlock()->without_terminator().empty()))3745      return failure();3746 3747    assert(ifOp->use_empty() || (llvm::all_equal(ifOp->getUsers()) &&3748                                 *ifOp->user_begin() == conditionOp) &&3749                                    "ifOp has unexpected uses");3750 3751    Location loc = op.getLoc();3752 3753    // Replace uses of ifOp results in the conditionOp with the yielded values3754    // from the ifOp branches.3755    for (auto [idx, arg] : llvm::enumerate(conditionOp.getArgs())) {3756      auto it = llvm::find(ifOp->getResults(), arg);3757      if (it != ifOp->getResults().end()) {3758        size_t ifOpIdx = it.getIndex();3759        Value thenValue = ifOp.thenYield()->getOperand(ifOpIdx);3760        Value elseValue = ifOp.elseYield()->getOperand(ifOpIdx);3761 3762        rewriter.replaceAllUsesWith(ifOp->getResults()[ifOpIdx], elseValue);3763        rewriter.replaceAllUsesWith(op.getAfterArguments()[idx], thenValue);3764      }3765    }3766 3767    // Collect additional used values from before region.3768    SetVector<Value> additionalUsedValuesSet;3769    visitUsedValuesDefinedAbove(ifOp.getThenRegion(), [&](OpOperand *operand) {3770      if (&op.getBefore() == operand->get().getParentRegion())3771        additionalUsedValuesSet.insert(operand->get());3772    });3773 3774    // Create new whileOp with additional used values as results.3775    auto additionalUsedValues = additionalUsedValuesSet.getArrayRef();3776    auto additionalValueTypes = llvm::map_to_vector(3777        additionalUsedValues, [](Value val) { return val.getType(); });3778    size_t additionalValueSize = additionalUsedValues.size();3779    SmallVector<Type> newResultTypes(op.getResultTypes());3780    newResultTypes.append(additionalValueTypes);3781 3782    auto newWhileOp =3783        scf::WhileOp::create(rewriter, loc, newResultTypes, op.getInits());3784 3785    rewriter.modifyOpInPlace(newWhileOp, [&] {3786      newWhileOp.getBefore().takeBody(op.getBefore());3787      newWhileOp.getAfter().takeBody(op.getAfter());3788      newWhileOp.getAfter().addArguments(3789          additionalValueTypes,3790          SmallVector<Location>(additionalValueSize, loc));3791    });3792 3793    rewriter.modifyOpInPlace(conditionOp, [&] {3794      conditionOp.getArgsMutable().append(additionalUsedValues);3795    });3796 3797    // Replace uses of additional used values inside the ifOp then region with3798    // the whileOp after region arguments.3799    rewriter.replaceUsesWithIf(3800        additionalUsedValues,3801        newWhileOp.getAfterArguments().take_back(additionalValueSize),3802        [&](OpOperand &use) {3803          return ifOp.getThenRegion().isAncestor(3804              use.getOwner()->getParentRegion());3805        });3806 3807    // Inline ifOp then region into new whileOp after region.3808    rewriter.eraseOp(ifOp.thenYield());3809    rewriter.inlineBlockBefore(ifOp.thenBlock(), newWhileOp.getAfterBody(),3810                               newWhileOp.getAfterBody()->begin());3811    rewriter.eraseOp(ifOp);3812    rewriter.replaceOp(op,3813                       newWhileOp->getResults().drop_back(additionalValueSize));3814    return success();3815  }3816};3817 3818/// Replace uses of the condition within the do block with true, since otherwise3819/// the block would not be evaluated.3820///3821/// scf.while (..) : (i1, ...) -> ... {3822///  %condition = call @evaluate_condition() : () -> i13823///  scf.condition(%condition) %condition : i1, ...3824/// } do {3825/// ^bb0(%arg0: i1, ...):3826///    use(%arg0)3827///    ...3828///3829/// becomes3830/// scf.while (..) : (i1, ...) -> ... {3831///  %condition = call @evaluate_condition() : () -> i13832///  scf.condition(%condition) %condition : i1, ...3833/// } do {3834/// ^bb0(%arg0: i1, ...):3835///    use(%true)3836///    ...3837struct WhileConditionTruth : public OpRewritePattern<WhileOp> {3838  using OpRewritePattern<WhileOp>::OpRewritePattern;3839 3840  LogicalResult matchAndRewrite(WhileOp op,3841                                PatternRewriter &rewriter) const override {3842    auto term = op.getConditionOp();3843 3844    // These variables serve to prevent creating duplicate constants3845    // and hold constant true or false values.3846    Value constantTrue = nullptr;3847 3848    bool replaced = false;3849    for (auto yieldedAndBlockArgs :3850         llvm::zip(term.getArgs(), op.getAfterArguments())) {3851      if (std::get<0>(yieldedAndBlockArgs) == term.getCondition()) {3852        if (!std::get<1>(yieldedAndBlockArgs).use_empty()) {3853          if (!constantTrue)3854            constantTrue = arith::ConstantOp::create(3855                rewriter, op.getLoc(), term.getCondition().getType(),3856                rewriter.getBoolAttr(true));3857 3858          rewriter.replaceAllUsesWith(std::get<1>(yieldedAndBlockArgs),3859                                      constantTrue);3860          replaced = true;3861        }3862      }3863    }3864    return success(replaced);3865  }3866};3867 3868/// Remove loop invariant arguments from `before` block of scf.while.3869/// A before block argument is considered loop invariant if :-3870///   1. i-th yield operand is equal to the i-th while operand.3871///   2. i-th yield operand is k-th after block argument which is (k+1)-th3872///      condition operand AND this (k+1)-th condition operand is equal to i-th3873///      iter argument/while operand.3874/// For the arguments which are removed, their uses inside scf.while3875/// are replaced with their corresponding initial value.3876///3877/// Eg:3878///    INPUT :-3879///    %res = scf.while <...> iter_args(%arg0_before = %a, %arg1_before = %b,3880///                                     ..., %argN_before = %N)3881///           {3882///                ...3883///                scf.condition(%cond) %arg1_before, %arg0_before,3884///                                     %arg2_before, %arg0_before, ...3885///           } do {3886///             ^bb0(%arg1_after, %arg0_after_1, %arg2_after, %arg0_after_2,3887///                  ..., %argK_after):3888///                ...3889///                scf.yield %arg0_after_2, %b, %arg1_after, ..., %argN3890///           }3891///3892///    OUTPUT :-3893///    %res = scf.while <...> iter_args(%arg2_before = %c, ..., %argN_before =3894///                                     %N)3895///           {3896///                ...3897///                scf.condition(%cond) %b, %a, %arg2_before, %a, ...3898///           } do {3899///             ^bb0(%arg1_after, %arg0_after_1, %arg2_after, %arg0_after_2,3900///                  ..., %argK_after):3901///                ...3902///                scf.yield %arg1_after, ..., %argN3903///           }3904///3905///    EXPLANATION:3906///      We iterate over each yield operand.3907///        1. 0-th yield operand %arg0_after_2 is 4-th condition operand3908///           %arg0_before, which in turn is the 0-th iter argument. So we3909///           remove 0-th before block argument and yield operand, and replace3910///           all uses of the 0-th before block argument with its initial value3911///           %a.3912///        2. 1-th yield operand %b is equal to the 1-th iter arg's initial3913///           value. So we remove this operand and the corresponding before3914///           block argument and replace all uses of 1-th before block argument3915///           with %b.3916struct RemoveLoopInvariantArgsFromBeforeBlock3917    : public OpRewritePattern<WhileOp> {3918  using OpRewritePattern<WhileOp>::OpRewritePattern;3919 3920  LogicalResult matchAndRewrite(WhileOp op,3921                                PatternRewriter &rewriter) const override {3922    Block &afterBlock = *op.getAfterBody();3923    Block::BlockArgListType beforeBlockArgs = op.getBeforeArguments();3924    ConditionOp condOp = op.getConditionOp();3925    OperandRange condOpArgs = condOp.getArgs();3926    Operation *yieldOp = afterBlock.getTerminator();3927    ValueRange yieldOpArgs = yieldOp->getOperands();3928 3929    bool canSimplify = false;3930    for (const auto &it :3931         llvm::enumerate(llvm::zip(op.getOperands(), yieldOpArgs))) {3932      auto index = static_cast<unsigned>(it.index());3933      auto [initVal, yieldOpArg] = it.value();3934      // If i-th yield operand is equal to the i-th operand of the scf.while,3935      // the i-th before block argument is a loop invariant.3936      if (yieldOpArg == initVal) {3937        canSimplify = true;3938        break;3939      }3940      // If the i-th yield operand is k-th after block argument, then we check3941      // if the (k+1)-th condition op operand is equal to either the i-th before3942      // block argument or the initial value of i-th before block argument. If3943      // the comparison results `true`, i-th before block argument is a loop3944      // invariant.3945      auto yieldOpBlockArg = llvm::dyn_cast<BlockArgument>(yieldOpArg);3946      if (yieldOpBlockArg && yieldOpBlockArg.getOwner() == &afterBlock) {3947        Value condOpArg = condOpArgs[yieldOpBlockArg.getArgNumber()];3948        if (condOpArg == beforeBlockArgs[index] || condOpArg == initVal) {3949          canSimplify = true;3950          break;3951        }3952      }3953    }3954 3955    if (!canSimplify)3956      return failure();3957 3958    SmallVector<Value> newInitArgs, newYieldOpArgs;3959    DenseMap<unsigned, Value> beforeBlockInitValMap;3960    SmallVector<Location> newBeforeBlockArgLocs;3961    for (const auto &it :3962         llvm::enumerate(llvm::zip(op.getOperands(), yieldOpArgs))) {3963      auto index = static_cast<unsigned>(it.index());3964      auto [initVal, yieldOpArg] = it.value();3965 3966      // If i-th yield operand is equal to the i-th operand of the scf.while,3967      // the i-th before block argument is a loop invariant.3968      if (yieldOpArg == initVal) {3969        beforeBlockInitValMap.insert({index, initVal});3970        continue;3971      } else {3972        // If the i-th yield operand is k-th after block argument, then we check3973        // if the (k+1)-th condition op operand is equal to either the i-th3974        // before block argument or the initial value of i-th before block3975        // argument. If the comparison results `true`, i-th before block3976        // argument is a loop invariant.3977        auto yieldOpBlockArg = llvm::dyn_cast<BlockArgument>(yieldOpArg);3978        if (yieldOpBlockArg && yieldOpBlockArg.getOwner() == &afterBlock) {3979          Value condOpArg = condOpArgs[yieldOpBlockArg.getArgNumber()];3980          if (condOpArg == beforeBlockArgs[index] || condOpArg == initVal) {3981            beforeBlockInitValMap.insert({index, initVal});3982            continue;3983          }3984        }3985      }3986      newInitArgs.emplace_back(initVal);3987      newYieldOpArgs.emplace_back(yieldOpArg);3988      newBeforeBlockArgLocs.emplace_back(beforeBlockArgs[index].getLoc());3989    }3990 3991    {3992      OpBuilder::InsertionGuard g(rewriter);3993      rewriter.setInsertionPoint(yieldOp);3994      rewriter.replaceOpWithNewOp<YieldOp>(yieldOp, newYieldOpArgs);3995    }3996 3997    auto newWhile = WhileOp::create(rewriter, op.getLoc(), op.getResultTypes(),3998                                    newInitArgs);3999 4000    Block &newBeforeBlock = *rewriter.createBlock(4001        &newWhile.getBefore(), /*insertPt*/ {},4002        ValueRange(newYieldOpArgs).getTypes(), newBeforeBlockArgLocs);4003 4004    Block &beforeBlock = *op.getBeforeBody();4005    SmallVector<Value> newBeforeBlockArgs(beforeBlock.getNumArguments());4006    // For each i-th before block argument we find it's replacement value as :-4007    //   1. If i-th before block argument is a loop invariant, we fetch it's4008    //      initial value from `beforeBlockInitValMap` by querying for key `i`.4009    //   2. Else we fetch j-th new before block argument as the replacement4010    //      value of i-th before block argument.4011    for (unsigned i = 0, j = 0, n = beforeBlock.getNumArguments(); i < n; i++) {4012      // If the index 'i' argument was a loop invariant we fetch it's initial4013      // value from `beforeBlockInitValMap`.4014      if (beforeBlockInitValMap.count(i) != 0)4015        newBeforeBlockArgs[i] = beforeBlockInitValMap[i];4016      else4017        newBeforeBlockArgs[i] = newBeforeBlock.getArgument(j++);4018    }4019 4020    rewriter.mergeBlocks(&beforeBlock, &newBeforeBlock, newBeforeBlockArgs);4021    rewriter.inlineRegionBefore(op.getAfter(), newWhile.getAfter(),4022                                newWhile.getAfter().begin());4023 4024    rewriter.replaceOp(op, newWhile.getResults());4025    return success();4026  }4027};4028 4029/// Remove loop invariant value from result (condition op) of scf.while.4030/// A value is considered loop invariant if the final value yielded by4031/// scf.condition is defined outside of the `before` block. We remove the4032/// corresponding argument in `after` block and replace the use with the value.4033/// We also replace the use of the corresponding result of scf.while with the4034/// value.4035///4036/// Eg:4037///    INPUT :-4038///    %res_input:K = scf.while <...> iter_args(%arg0_before = , ...,4039///                                             %argN_before = %N) {4040///                ...4041///                scf.condition(%cond) %arg0_before, %a, %b, %arg1_before, ...4042///           } do {4043///             ^bb0(%arg0_after, %arg1_after, %arg2_after, ..., %argK_after):4044///                ...4045///                some_func(%arg1_after)4046///                ...4047///                scf.yield %arg0_after, %arg2_after, ..., %argN_after4048///           }4049///4050///    OUTPUT :-4051///    %res_output:M = scf.while <...> iter_args(%arg0 = , ..., %argN = %N) {4052///                ...4053///                scf.condition(%cond) %arg0, %arg1, ..., %argM4054///           } do {4055///             ^bb0(%arg0, %arg3, ..., %argM):4056///                ...4057///                some_func(%a)4058///                ...4059///                scf.yield %arg0, %b, ..., %argN4060///           }4061///4062///     EXPLANATION:4063///       1. The 1-th and 2-th operand of scf.condition are defined outside the4064///          before block of scf.while, so they get removed.4065///       2. %res_input#1's uses are replaced by %a and %res_input#2's uses are4066///          replaced by %b.4067///       3. The corresponding after block argument %arg1_after's uses are4068///          replaced by %a and %arg2_after's uses are replaced by %b.4069struct RemoveLoopInvariantValueYielded : public OpRewritePattern<WhileOp> {4070  using OpRewritePattern<WhileOp>::OpRewritePattern;4071 4072  LogicalResult matchAndRewrite(WhileOp op,4073                                PatternRewriter &rewriter) const override {4074    Block &beforeBlock = *op.getBeforeBody();4075    ConditionOp condOp = op.getConditionOp();4076    OperandRange condOpArgs = condOp.getArgs();4077 4078    bool canSimplify = false;4079    for (Value condOpArg : condOpArgs) {4080      // Those values not defined within `before` block will be considered as4081      // loop invariant values. We map the corresponding `index` with their4082      // value.4083      if (condOpArg.getParentBlock() != &beforeBlock) {4084        canSimplify = true;4085        break;4086      }4087    }4088 4089    if (!canSimplify)4090      return failure();4091 4092    Block::BlockArgListType afterBlockArgs = op.getAfterArguments();4093 4094    SmallVector<Value> newCondOpArgs;4095    SmallVector<Type> newAfterBlockType;4096    DenseMap<unsigned, Value> condOpInitValMap;4097    SmallVector<Location> newAfterBlockArgLocs;4098    for (const auto &it : llvm::enumerate(condOpArgs)) {4099      auto index = static_cast<unsigned>(it.index());4100      Value condOpArg = it.value();4101      // Those values not defined within `before` block will be considered as4102      // loop invariant values. We map the corresponding `index` with their4103      // value.4104      if (condOpArg.getParentBlock() != &beforeBlock) {4105        condOpInitValMap.insert({index, condOpArg});4106      } else {4107        newCondOpArgs.emplace_back(condOpArg);4108        newAfterBlockType.emplace_back(condOpArg.getType());4109        newAfterBlockArgLocs.emplace_back(afterBlockArgs[index].getLoc());4110      }4111    }4112 4113    {4114      OpBuilder::InsertionGuard g(rewriter);4115      rewriter.setInsertionPoint(condOp);4116      rewriter.replaceOpWithNewOp<ConditionOp>(condOp, condOp.getCondition(),4117                                               newCondOpArgs);4118    }4119 4120    auto newWhile = WhileOp::create(rewriter, op.getLoc(), newAfterBlockType,4121                                    op.getOperands());4122 4123    Block &newAfterBlock =4124        *rewriter.createBlock(&newWhile.getAfter(), /*insertPt*/ {},4125                              newAfterBlockType, newAfterBlockArgLocs);4126 4127    Block &afterBlock = *op.getAfterBody();4128    // Since a new scf.condition op was created, we need to fetch the new4129    // `after` block arguments which will be used while replacing operations of4130    // previous scf.while's `after` blocks. We'd also be fetching new result4131    // values too.4132    SmallVector<Value> newAfterBlockArgs(afterBlock.getNumArguments());4133    SmallVector<Value> newWhileResults(afterBlock.getNumArguments());4134    for (unsigned i = 0, j = 0, n = afterBlock.getNumArguments(); i < n; i++) {4135      Value afterBlockArg, result;4136      // If index 'i' argument was loop invariant we fetch it's value from the4137      // `condOpInitMap` map.4138      if (condOpInitValMap.count(i) != 0) {4139        afterBlockArg = condOpInitValMap[i];4140        result = afterBlockArg;4141      } else {4142        afterBlockArg = newAfterBlock.getArgument(j);4143        result = newWhile.getResult(j);4144        j++;4145      }4146      newAfterBlockArgs[i] = afterBlockArg;4147      newWhileResults[i] = result;4148    }4149 4150    rewriter.mergeBlocks(&afterBlock, &newAfterBlock, newAfterBlockArgs);4151    rewriter.inlineRegionBefore(op.getBefore(), newWhile.getBefore(),4152                                newWhile.getBefore().begin());4153 4154    rewriter.replaceOp(op, newWhileResults);4155    return success();4156  }4157};4158 4159/// Remove WhileOp results that are also unused in 'after' block.4160///4161///  %0:2 = scf.while () : () -> (i32, i64) {4162///    %condition = "test.condition"() : () -> i14163///    %v1 = "test.get_some_value"() : () -> i324164///    %v2 = "test.get_some_value"() : () -> i644165///    scf.condition(%condition) %v1, %v2 : i32, i644166///  } do {4167///  ^bb0(%arg0: i32, %arg1: i64):4168///    "test.use"(%arg0) : (i32) -> ()4169///    scf.yield4170///  }4171///  return %0#0 : i324172///4173/// becomes4174///  %0 = scf.while () : () -> (i32) {4175///    %condition = "test.condition"() : () -> i14176///    %v1 = "test.get_some_value"() : () -> i324177///    %v2 = "test.get_some_value"() : () -> i644178///    scf.condition(%condition) %v1 : i324179///  } do {4180///  ^bb0(%arg0: i32):4181///    "test.use"(%arg0) : (i32) -> ()4182///    scf.yield4183///  }4184///  return %0 : i324185struct WhileUnusedResult : public OpRewritePattern<WhileOp> {4186  using OpRewritePattern<WhileOp>::OpRewritePattern;4187 4188  LogicalResult matchAndRewrite(WhileOp op,4189                                PatternRewriter &rewriter) const override {4190    auto term = op.getConditionOp();4191    auto afterArgs = op.getAfterArguments();4192    auto termArgs = term.getArgs();4193 4194    // Collect results mapping, new terminator args and new result types.4195    SmallVector<unsigned> newResultsIndices;4196    SmallVector<Type> newResultTypes;4197    SmallVector<Value> newTermArgs;4198    SmallVector<Location> newArgLocs;4199    bool needUpdate = false;4200    for (const auto &it :4201         llvm::enumerate(llvm::zip(op.getResults(), afterArgs, termArgs))) {4202      auto i = static_cast<unsigned>(it.index());4203      Value result = std::get<0>(it.value());4204      Value afterArg = std::get<1>(it.value());4205      Value termArg = std::get<2>(it.value());4206      if (result.use_empty() && afterArg.use_empty()) {4207        needUpdate = true;4208      } else {4209        newResultsIndices.emplace_back(i);4210        newTermArgs.emplace_back(termArg);4211        newResultTypes.emplace_back(result.getType());4212        newArgLocs.emplace_back(result.getLoc());4213      }4214    }4215 4216    if (!needUpdate)4217      return failure();4218 4219    {4220      OpBuilder::InsertionGuard g(rewriter);4221      rewriter.setInsertionPoint(term);4222      rewriter.replaceOpWithNewOp<ConditionOp>(term, term.getCondition(),4223                                               newTermArgs);4224    }4225 4226    auto newWhile =4227        WhileOp::create(rewriter, op.getLoc(), newResultTypes, op.getInits());4228 4229    Block &newAfterBlock = *rewriter.createBlock(4230        &newWhile.getAfter(), /*insertPt*/ {}, newResultTypes, newArgLocs);4231 4232    // Build new results list and new after block args (unused entries will be4233    // null).4234    SmallVector<Value> newResults(op.getNumResults());4235    SmallVector<Value> newAfterBlockArgs(op.getNumResults());4236    for (const auto &it : llvm::enumerate(newResultsIndices)) {4237      newResults[it.value()] = newWhile.getResult(it.index());4238      newAfterBlockArgs[it.value()] = newAfterBlock.getArgument(it.index());4239    }4240 4241    rewriter.inlineRegionBefore(op.getBefore(), newWhile.getBefore(),4242                                newWhile.getBefore().begin());4243 4244    Block &afterBlock = *op.getAfterBody();4245    rewriter.mergeBlocks(&afterBlock, &newAfterBlock, newAfterBlockArgs);4246 4247    rewriter.replaceOp(op, newResults);4248    return success();4249  }4250};4251 4252/// Replace operations equivalent to the condition in the do block with true,4253/// since otherwise the block would not be evaluated.4254///4255/// scf.while (..) : (i32, ...) -> ... {4256///  %z = ... : i324257///  %condition = cmpi pred %z, %a4258///  scf.condition(%condition) %z : i32, ...4259/// } do {4260/// ^bb0(%arg0: i32, ...):4261///    %condition2 = cmpi pred %arg0, %a4262///    use(%condition2)4263///    ...4264///4265/// becomes4266/// scf.while (..) : (i32, ...) -> ... {4267///  %z = ... : i324268///  %condition = cmpi pred %z, %a4269///  scf.condition(%condition) %z : i32, ...4270/// } do {4271/// ^bb0(%arg0: i32, ...):4272///    use(%true)4273///    ...4274struct WhileCmpCond : public OpRewritePattern<scf::WhileOp> {4275  using OpRewritePattern<scf::WhileOp>::OpRewritePattern;4276 4277  LogicalResult matchAndRewrite(scf::WhileOp op,4278                                PatternRewriter &rewriter) const override {4279    using namespace scf;4280    auto cond = op.getConditionOp();4281    auto cmp = cond.getCondition().getDefiningOp<arith::CmpIOp>();4282    if (!cmp)4283      return failure();4284    bool changed = false;4285    for (auto tup : llvm::zip(cond.getArgs(), op.getAfterArguments())) {4286      for (size_t opIdx = 0; opIdx < 2; opIdx++) {4287        if (std::get<0>(tup) != cmp.getOperand(opIdx))4288          continue;4289        for (OpOperand &u :4290             llvm::make_early_inc_range(std::get<1>(tup).getUses())) {4291          auto cmp2 = dyn_cast<arith::CmpIOp>(u.getOwner());4292          if (!cmp2)4293            continue;4294          // For a binary operator 1-opIdx gets the other side.4295          if (cmp2.getOperand(1 - opIdx) != cmp.getOperand(1 - opIdx))4296            continue;4297          bool samePredicate;4298          if (cmp2.getPredicate() == cmp.getPredicate())4299            samePredicate = true;4300          else if (cmp2.getPredicate() ==4301                   arith::invertPredicate(cmp.getPredicate()))4302            samePredicate = false;4303          else4304            continue;4305 4306          rewriter.replaceOpWithNewOp<arith::ConstantIntOp>(cmp2, samePredicate,4307                                                            1);4308          changed = true;4309        }4310      }4311    }4312    return success(changed);4313  }4314};4315 4316/// Remove unused init/yield args.4317struct WhileRemoveUnusedArgs : public OpRewritePattern<WhileOp> {4318  using OpRewritePattern<WhileOp>::OpRewritePattern;4319 4320  LogicalResult matchAndRewrite(WhileOp op,4321                                PatternRewriter &rewriter) const override {4322 4323    if (!llvm::any_of(op.getBeforeArguments(),4324                      [](Value arg) { return arg.use_empty(); }))4325      return rewriter.notifyMatchFailure(op, "No args to remove");4326 4327    YieldOp yield = op.getYieldOp();4328 4329    // Collect results mapping, new terminator args and new result types.4330    SmallVector<Value> newYields;4331    SmallVector<Value> newInits;4332    llvm::BitVector argsToErase;4333 4334    size_t argsCount = op.getBeforeArguments().size();4335    newYields.reserve(argsCount);4336    newInits.reserve(argsCount);4337    argsToErase.reserve(argsCount);4338    for (auto &&[beforeArg, yieldValue, initValue] : llvm::zip(4339             op.getBeforeArguments(), yield.getOperands(), op.getInits())) {4340      if (beforeArg.use_empty()) {4341        argsToErase.push_back(true);4342      } else {4343        argsToErase.push_back(false);4344        newYields.emplace_back(yieldValue);4345        newInits.emplace_back(initValue);4346      }4347    }4348 4349    Block &beforeBlock = *op.getBeforeBody();4350    Block &afterBlock = *op.getAfterBody();4351 4352    beforeBlock.eraseArguments(argsToErase);4353 4354    Location loc = op.getLoc();4355    auto newWhileOp =4356        WhileOp::create(rewriter, loc, op.getResultTypes(), newInits,4357                        /*beforeBody*/ nullptr, /*afterBody*/ nullptr);4358    Block &newBeforeBlock = *newWhileOp.getBeforeBody();4359    Block &newAfterBlock = *newWhileOp.getAfterBody();4360 4361    OpBuilder::InsertionGuard g(rewriter);4362    rewriter.setInsertionPoint(yield);4363    rewriter.replaceOpWithNewOp<YieldOp>(yield, newYields);4364 4365    rewriter.mergeBlocks(&beforeBlock, &newBeforeBlock,4366                         newBeforeBlock.getArguments());4367    rewriter.mergeBlocks(&afterBlock, &newAfterBlock,4368                         newAfterBlock.getArguments());4369 4370    rewriter.replaceOp(op, newWhileOp.getResults());4371    return success();4372  }4373};4374 4375/// Remove duplicated ConditionOp args.4376struct WhileRemoveDuplicatedResults : public OpRewritePattern<WhileOp> {4377  using OpRewritePattern::OpRewritePattern;4378 4379  LogicalResult matchAndRewrite(WhileOp op,4380                                PatternRewriter &rewriter) const override {4381    ConditionOp condOp = op.getConditionOp();4382    ValueRange condOpArgs = condOp.getArgs();4383 4384    llvm::SmallPtrSet<Value, 8> argsSet(llvm::from_range, condOpArgs);4385 4386    if (argsSet.size() == condOpArgs.size())4387      return rewriter.notifyMatchFailure(op, "No results to remove");4388 4389    llvm::SmallDenseMap<Value, unsigned> argsMap;4390    SmallVector<Value> newArgs;4391    argsMap.reserve(condOpArgs.size());4392    newArgs.reserve(condOpArgs.size());4393    for (Value arg : condOpArgs) {4394      if (!argsMap.count(arg)) {4395        auto pos = static_cast<unsigned>(argsMap.size());4396        argsMap.insert({arg, pos});4397        newArgs.emplace_back(arg);4398      }4399    }4400 4401    ValueRange argsRange(newArgs);4402 4403    Location loc = op.getLoc();4404    auto newWhileOp =4405        scf::WhileOp::create(rewriter, loc, argsRange.getTypes(), op.getInits(),4406                             /*beforeBody*/ nullptr,4407                             /*afterBody*/ nullptr);4408    Block &newBeforeBlock = *newWhileOp.getBeforeBody();4409    Block &newAfterBlock = *newWhileOp.getAfterBody();4410 4411    SmallVector<Value> afterArgsMapping;4412    SmallVector<Value> resultsMapping;4413    for (auto &&[i, arg] : llvm::enumerate(condOpArgs)) {4414      auto it = argsMap.find(arg);4415      assert(it != argsMap.end());4416      auto pos = it->second;4417      afterArgsMapping.emplace_back(newAfterBlock.getArgument(pos));4418      resultsMapping.emplace_back(newWhileOp->getResult(pos));4419    }4420 4421    OpBuilder::InsertionGuard g(rewriter);4422    rewriter.setInsertionPoint(condOp);4423    rewriter.replaceOpWithNewOp<ConditionOp>(condOp, condOp.getCondition(),4424                                             argsRange);4425 4426    Block &beforeBlock = *op.getBeforeBody();4427    Block &afterBlock = *op.getAfterBody();4428 4429    rewriter.mergeBlocks(&beforeBlock, &newBeforeBlock,4430                         newBeforeBlock.getArguments());4431    rewriter.mergeBlocks(&afterBlock, &newAfterBlock, afterArgsMapping);4432    rewriter.replaceOp(op, resultsMapping);4433    return success();4434  }4435};4436 4437/// If both ranges contain same values return mappping indices from args2 to4438/// args1. Otherwise return std::nullopt.4439static std::optional<SmallVector<unsigned>> getArgsMapping(ValueRange args1,4440                                                           ValueRange args2) {4441  if (args1.size() != args2.size())4442    return std::nullopt;4443 4444  SmallVector<unsigned> ret(args1.size());4445  for (auto &&[i, arg1] : llvm::enumerate(args1)) {4446    auto it = llvm::find(args2, arg1);4447    if (it == args2.end())4448      return std::nullopt;4449 4450    ret[std::distance(args2.begin(), it)] = static_cast<unsigned>(i);4451  }4452 4453  return ret;4454}4455 4456static bool hasDuplicates(ValueRange args) {4457  llvm::SmallDenseSet<Value> set;4458  for (Value arg : args) {4459    if (!set.insert(arg).second)4460      return true;4461  }4462  return false;4463}4464 4465/// If `before` block args are directly forwarded to `scf.condition`, rearrange4466/// `scf.condition` args into same order as block args. Update `after` block4467/// args and op result values accordingly.4468/// Needed to simplify `scf.while` -> `scf.for` uplifting.4469struct WhileOpAlignBeforeArgs : public OpRewritePattern<WhileOp> {4470  using OpRewritePattern::OpRewritePattern;4471 4472  LogicalResult matchAndRewrite(WhileOp loop,4473                                PatternRewriter &rewriter) const override {4474    auto oldBefore = loop.getBeforeBody();4475    ConditionOp oldTerm = loop.getConditionOp();4476    ValueRange beforeArgs = oldBefore->getArguments();4477    ValueRange termArgs = oldTerm.getArgs();4478    if (beforeArgs == termArgs)4479      return failure();4480 4481    if (hasDuplicates(termArgs))4482      return failure();4483 4484    auto mapping = getArgsMapping(beforeArgs, termArgs);4485    if (!mapping)4486      return failure();4487 4488    {4489      OpBuilder::InsertionGuard g(rewriter);4490      rewriter.setInsertionPoint(oldTerm);4491      rewriter.replaceOpWithNewOp<ConditionOp>(oldTerm, oldTerm.getCondition(),4492                                               beforeArgs);4493    }4494 4495    auto oldAfter = loop.getAfterBody();4496 4497    SmallVector<Type> newResultTypes(beforeArgs.size());4498    for (auto &&[i, j] : llvm::enumerate(*mapping))4499      newResultTypes[j] = loop.getResult(i).getType();4500 4501    auto newLoop = WhileOp::create(4502        rewriter, loop.getLoc(), newResultTypes, loop.getInits(),4503        /*beforeBuilder=*/nullptr, /*afterBuilder=*/nullptr);4504    auto newBefore = newLoop.getBeforeBody();4505    auto newAfter = newLoop.getAfterBody();4506 4507    SmallVector<Value> newResults(beforeArgs.size());4508    SmallVector<Value> newAfterArgs(beforeArgs.size());4509    for (auto &&[i, j] : llvm::enumerate(*mapping)) {4510      newResults[i] = newLoop.getResult(j);4511      newAfterArgs[i] = newAfter->getArgument(j);4512    }4513 4514    rewriter.inlineBlockBefore(oldBefore, newBefore, newBefore->begin(),4515                               newBefore->getArguments());4516    rewriter.inlineBlockBefore(oldAfter, newAfter, newAfter->begin(),4517                               newAfterArgs);4518 4519    rewriter.replaceOp(loop, newResults);4520    return success();4521  }4522};4523} // namespace4524 4525void WhileOp::getCanonicalizationPatterns(RewritePatternSet &results,4526                                          MLIRContext *context) {4527  results.add<RemoveLoopInvariantArgsFromBeforeBlock,4528              RemoveLoopInvariantValueYielded, WhileConditionTruth,4529              WhileCmpCond, WhileUnusedResult, WhileRemoveDuplicatedResults,4530              WhileRemoveUnusedArgs, WhileOpAlignBeforeArgs, WhileMoveIfDown>(4531      context);4532}4533 4534//===----------------------------------------------------------------------===//4535// IndexSwitchOp4536//===----------------------------------------------------------------------===//4537 4538/// Parse the case regions and values.4539static ParseResult4540parseSwitchCases(OpAsmParser &p, DenseI64ArrayAttr &cases,4541                 SmallVectorImpl<std::unique_ptr<Region>> &caseRegions) {4542  SmallVector<int64_t> caseValues;4543  while (succeeded(p.parseOptionalKeyword("case"))) {4544    int64_t value;4545    Region &region = *caseRegions.emplace_back(std::make_unique<Region>());4546    if (p.parseInteger(value) || p.parseRegion(region, /*arguments=*/{}))4547      return failure();4548    caseValues.push_back(value);4549  }4550  cases = p.getBuilder().getDenseI64ArrayAttr(caseValues);4551  return success();4552}4553 4554/// Print the case regions and values.4555static void printSwitchCases(OpAsmPrinter &p, Operation *op,4556                             DenseI64ArrayAttr cases, RegionRange caseRegions) {4557  for (auto [value, region] : llvm::zip(cases.asArrayRef(), caseRegions)) {4558    p.printNewline();4559    p << "case " << value << ' ';4560    p.printRegion(*region, /*printEntryBlockArgs=*/false);4561  }4562}4563 4564LogicalResult scf::IndexSwitchOp::verify() {4565  if (getCases().size() != getCaseRegions().size()) {4566    return emitOpError("has ")4567           << getCaseRegions().size() << " case regions but "4568           << getCases().size() << " case values";4569  }4570 4571  DenseSet<int64_t> valueSet;4572  for (int64_t value : getCases())4573    if (!valueSet.insert(value).second)4574      return emitOpError("has duplicate case value: ") << value;4575  auto verifyRegion = [&](Region &region, const Twine &name) -> LogicalResult {4576    auto yield = dyn_cast<YieldOp>(region.front().back());4577    if (!yield)4578      return emitOpError("expected region to end with scf.yield, but got ")4579             << region.front().back().getName();4580 4581    if (yield.getNumOperands() != getNumResults()) {4582      return (emitOpError("expected each region to return ")4583              << getNumResults() << " values, but " << name << " returns "4584              << yield.getNumOperands())4585                 .attachNote(yield.getLoc())4586             << "see yield operation here";4587    }4588    for (auto [idx, result, operand] :4589         llvm::enumerate(getResultTypes(), yield.getOperands())) {4590      if (!operand)4591        return yield.emitOpError() << "operand " << idx << " is null\n";4592      if (result == operand.getType())4593        continue;4594      return (emitOpError("expected result #")4595              << idx << " of each region to be " << result)4596                 .attachNote(yield.getLoc())4597             << name << " returns " << operand.getType() << " here";4598    }4599    return success();4600  };4601 4602  if (failed(verifyRegion(getDefaultRegion(), "default region")))4603    return failure();4604  for (auto [idx, caseRegion] : llvm::enumerate(getCaseRegions()))4605    if (failed(verifyRegion(caseRegion, "case region #" + Twine(idx))))4606      return failure();4607 4608  return success();4609}4610 4611unsigned scf::IndexSwitchOp::getNumCases() { return getCases().size(); }4612 4613Block &scf::IndexSwitchOp::getDefaultBlock() {4614  return getDefaultRegion().front();4615}4616 4617Block &scf::IndexSwitchOp::getCaseBlock(unsigned idx) {4618  assert(idx < getNumCases() && "case index out-of-bounds");4619  return getCaseRegions()[idx].front();4620}4621 4622void IndexSwitchOp::getSuccessorRegions(4623    RegionBranchPoint point, SmallVectorImpl<RegionSuccessor> &successors) {4624  // All regions branch back to the parent op.4625  if (!point.isParent()) {4626    successors.emplace_back(getOperation(), getResults());4627    return;4628  }4629 4630  llvm::append_range(successors, getRegions());4631}4632 4633void IndexSwitchOp::getEntrySuccessorRegions(4634    ArrayRef<Attribute> operands,4635    SmallVectorImpl<RegionSuccessor> &successors) {4636  FoldAdaptor adaptor(operands, *this);4637 4638  // If a constant was not provided, all regions are possible successors.4639  auto arg = dyn_cast_or_null<IntegerAttr>(adaptor.getArg());4640  if (!arg) {4641    llvm::append_range(successors, getRegions());4642    return;4643  }4644 4645  // Otherwise, try to find a case with a matching value. If not, the4646  // default region is the only successor.4647  for (auto [caseValue, caseRegion] : llvm::zip(getCases(), getCaseRegions())) {4648    if (caseValue == arg.getInt()) {4649      successors.emplace_back(&caseRegion);4650      return;4651    }4652  }4653  successors.emplace_back(&getDefaultRegion());4654}4655 4656void IndexSwitchOp::getRegionInvocationBounds(4657    ArrayRef<Attribute> operands, SmallVectorImpl<InvocationBounds> &bounds) {4658  auto operandValue = llvm::dyn_cast_or_null<IntegerAttr>(operands.front());4659  if (!operandValue) {4660    // All regions are invoked at most once.4661    bounds.append(getNumRegions(), InvocationBounds(/*lb=*/0, /*ub=*/1));4662    return;4663  }4664 4665  unsigned liveIndex = getNumRegions() - 1;4666  const auto *it = llvm::find(getCases(), operandValue.getInt());4667  if (it != getCases().end())4668    liveIndex = std::distance(getCases().begin(), it);4669  for (unsigned i = 0, e = getNumRegions(); i < e; ++i)4670    bounds.emplace_back(/*lb=*/0, /*ub=*/i == liveIndex);4671}4672 4673struct FoldConstantCase : OpRewritePattern<scf::IndexSwitchOp> {4674  using OpRewritePattern<scf::IndexSwitchOp>::OpRewritePattern;4675 4676  LogicalResult matchAndRewrite(scf::IndexSwitchOp op,4677                                PatternRewriter &rewriter) const override {4678    // If `op.getArg()` is a constant, select the region that matches with4679    // the constant value. Use the default region if no matche is found.4680    std::optional<int64_t> maybeCst = getConstantIntValue(op.getArg());4681    if (!maybeCst.has_value())4682      return failure();4683    int64_t cst = *maybeCst;4684    int64_t caseIdx, e = op.getNumCases();4685    for (caseIdx = 0; caseIdx < e; ++caseIdx) {4686      if (cst == op.getCases()[caseIdx])4687        break;4688    }4689 4690    Region &r = (caseIdx < op.getNumCases()) ? op.getCaseRegions()[caseIdx]4691                                             : op.getDefaultRegion();4692    Block &source = r.front();4693    Operation *terminator = source.getTerminator();4694    SmallVector<Value> results = terminator->getOperands();4695 4696    rewriter.inlineBlockBefore(&source, op);4697    rewriter.eraseOp(terminator);4698    // Replace the operation with a potentially empty list of results.4699    // Fold mechanism doesn't support the case where the result list is empty.4700    rewriter.replaceOp(op, results);4701 4702    return success();4703  }4704};4705 4706void IndexSwitchOp::getCanonicalizationPatterns(RewritePatternSet &results,4707                                                MLIRContext *context) {4708  results.add<FoldConstantCase>(context);4709}4710 4711//===----------------------------------------------------------------------===//4712// TableGen'd op method definitions4713//===----------------------------------------------------------------------===//4714 4715#define GET_OP_CLASSES4716#include "mlir/Dialect/SCF/IR/SCFOps.cpp.inc"4717