[mlir][Interfaces] Add `promoteIfSingleIteration` to `LoopLikeOpInterface`

[springerm]

Summary

There are existing implementations for scf.for, scf.forall and affine.for. This revision adds an interface method to the LoopLikeOpInterface.

• scf.forall now implements the LoopLikeOpInterface.
• The implementations of scf.for and scf.forall become interface method implementations. affine.for remains as is for the moment. (The implementation of promoteIfSingleIteration depepends on helper functions from MLIRAffineAnalysis, which cannot be used from MLIRAffineDialect, where the interface is currently implemented.)
• More efficient implementations of promoteIfSingleIteration. In particular, the scf.forall operation now inlines operations instead of cloning them. This also preserves handles when used from the transform dialect.

[nicolasvasilache]

nice !

mlir/include/mlir/Dialect/SCF/IR/SCF.h

 /// If the provided value is not a thread index variable, then return nullptr.
 ForallOp getForallOpThreadIndexOwner(Value val);
 
 /// Return true if ops a and b (or their ancestors) are in mutually exclusive
 /// regions/blocks of an IfOp.
 // TODO: Consider moving this functionality to RegionBranchOpInterface.
 bool insideMutuallyExclusiveBranches(Operation *a, Operation *b);
 
-/// Promotes the loop body of a scf::ForallOp to its containing block if the
-/// loop was known to have a single iteration.
-LogicalResult promoteIfSingleIteration(PatternRewriter &rewriter,
-                                       scf::ForallOp forallOp);
-
 /// Promotes the loop body of a scf::ForallOp to its containing block.
-void promote(PatternRewriter &rewriter, scf::ForallOp forallOp);
+void promote(RewriterBase &rewriter, scf::ForallOp forallOp);
 
 /// An owning vector of values, handy to return from functions.
 using ValueVector = SmallVector<Value>;
 using LoopVector = SmallVector<scf::ForOp>;
 struct LoopNest {
   LoopVector loops;
   ValueVector results;
 };

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

 
 //===----------------------------------------------------------------------===//
 // ForOp
 //===----------------------------------------------------------------------===//
 
 def ForOp : SCF_Op<"for",
       [AutomaticAllocationScope, DeclareOpInterfaceMethods<LoopLikeOpInterface,
        ["getSingleInductionVar", "getSingleLowerBound", "getSingleStep",
-        "getSingleUpperBound"]>,
+        "getSingleUpperBound", "promoteIfSingleIteration"]>,
        AllTypesMatch<["lowerBound", "upperBound", "step"]>,
        ConditionallySpeculatable,
        DeclareOpInterfaceMethods<RegionBranchOpInterface>,
        SingleBlockImplicitTerminator<"scf::YieldOp">,
        RecursiveMemoryEffects]> {
   let summary = "for operation";
   let description = [{
     The `scf.for` operation represents a loop taking 3 SSA value as operands

 
 //===----------------------------------------------------------------------===//
 // ForallOp
 //===----------------------------------------------------------------------===//
 
 def ForallOp : SCF_Op<"forall", [
        AttrSizedOperandSegments,
        AutomaticAllocationScope,
+       DeclareOpInterfaceMethods<LoopLikeOpInterface,
+          ["promoteIfSingleIteration"]>,
        RecursiveMemoryEffects,
        SingleBlockImplicitTerminator<"scf::InParallelOp">,
        DeclareOpInterfaceMethods<RegionBranchOpInterface>,
      ]> {
   let summary = "evaluate a block multiple times in parallel";
   let description = [{
     `scf.forall` is a target-independent multi-dimensional parallel
     region application operation. It has exactly one block that represents the

mlir/include/mlir/Dialect/SCF/Utils/Utils.h

 LogicalResult coalesceLoops(MutableArrayRef<scf::ForOp> loops);
 
 /// Take the ParallelLoop and for each set of dimension indices, combine them
 /// into a single dimension. combinedDimensions must contain each index into
 /// loops exactly once.
 void collapseParallelLoops(scf::ParallelOp loops,
                            ArrayRef<std::vector<unsigned>> combinedDimensions);
 
-/// Promotes the loop body of a scf::ForOp to its containing block if the loop
-/// was known to have a single iteration.
-LogicalResult promoteIfSingleIteration(scf::ForOp forOp);
-
 /// Unrolls this for operation by the specified unroll factor. Returns failure
 /// if the loop cannot be unrolled either due to restrictions or due to invalid
 /// unroll factors. Requires positive loop bounds and step. If specified,
 /// annotates the Ops in each unrolled iteration by applying `annotateFn`.
 LogicalResult loopUnrollByFactor(
     scf::ForOp forOp, uint64_t unrollFactor,
     function_ref<void(unsigned, Operation *, OpBuilder)> annotateFn = nullptr);
 

mlir/include/mlir/Interfaces/LoopLikeInterface.h

 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_INTERFACES_LOOPLIKEINTERFACE_H_
 #define MLIR_INTERFACES_LOOPLIKEINTERFACE_H_
 
 #include "mlir/IR/OpDefinition.h"
 
+namespace mlir {
+class RewriterBase;
+} // namespace mlir
+
 /// Include the generated interface declarations.
 #include "mlir/Interfaces/LoopLikeInterface.h.inc"
 
 #endif // MLIR_INTERFACES_LOOPLIKEINTERFACE_H_

mlir/include/mlir/Interfaces/LoopLikeInterface.td

InterfaceMethod<[{
         Moves the given loop-invariant operation out of the loop.
       }],
       "void", "moveOutOfLoop",
       (ins "::mlir::Operation *":$op), [{}], [{
         op->moveBefore($_op);
       }]
     >,
     InterfaceMethod<[{
+        Promotes the loop body to its containing block if the loop is known to
+        have a single iteration. Returns "success" if the promotion was
+        successful.
+      }],
+      /*retTy=*/"::mlir::LogicalResult",
+      /*methodName=*/"promoteIfSingleIteration",
+      /*args=*/(ins "::mlir::RewriterBase &":$rewriter),
+      /*methodBody=*/"",
+      /*defaultImplementation=*/[{
+        return ::mlir::failure();
+      }]
+    >,
+    InterfaceMethod<[{
         If there is a single induction variable return it, otherwise return
         std::nullopt.
       }],
       /*retTy=*/"::std::optional<::mlir::Value>",
       /*methodName=*/"getSingleInductionVar",
       /*args=*/(ins),
       /*methodBody=*/"",
       /*defaultImplementation=*/[{

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

 std::optional<OpFoldResult> ForOp::getSingleStep() {
   return OpFoldResult(getStep());
 }
 
 std::optional<OpFoldResult> ForOp::getSingleUpperBound() {
   return OpFoldResult(getUpperBound());
 }
 
+/// Promotes the loop body of a forOp to its containing block if the forOp
+/// it can be determined that the loop has a single iteration.
+LogicalResult ForOp::promoteIfSingleIteration(RewriterBase &rewriter) {
+  std::optional<int64_t> tripCount =
+      constantTripCount(getLowerBound(), getUpperBound(), getStep());
+  if (!tripCount.has_value() || tripCount != 1)
+    return failure();
+
+  // Replace all results with the yielded values.
+  auto yieldOp = cast<scf::YieldOp>(getBody()->getTerminator());
+  rewriter.replaceAllUsesWith(getResults(), yieldOp.getOperands());
+
+  // Replace block arguments with lower bound (replacement for IV) and
+  // iter_args.
+  SmallVector<Value> bbArgReplacements;
+  bbArgReplacements.push_back(getLowerBound());
+  bbArgReplacements.append(getIterOperands().begin(), getIterOperands().end());
+
+  // Move the loop body operations to the loop's containing block.
+  rewriter.inlineBlockBefore(getBody(), getOperation()->getBlock(),
+                             getOperation()->getIterator(), bbArgReplacements);
+
+  // Erase the old terminator and the loop.
+  rewriter.eraseOp(yieldOp);
+  rewriter.eraseOp(*this);
+
+  return success();
+}
+
 /// Prints the initialization list in the form of
 ///   <prefix>(%inner = %outer, %inner2 = %outer2, <...>)
 /// where 'inner' values are assumed to be region arguments and 'outer' values
 /// are regular SSA values.
 static void printInitializationList(OpAsmPrinter &p,
                                     Block::BlockArgListType blocksArgs,
                                     ValueRange initializers,
                                     StringRef prefix = "") {
void ForOp::getSuccessorRegions(std::optional<unsigned> index,
   }
 
   // Otherwise, the loop may branch back to itself or the parent operation.
   assert(*index == 0 && "expected loop region");
   regions.push_back(RegionSuccessor(&getLoopBody(), getRegionIterArgs()));
   regions.push_back(RegionSuccessor(getResults()));
 }
 
+Region &ForallOp::getLoopBody() { return getRegion(); }
+
 /// Promotes the loop body of a forallOp to its containing block if it can be
 /// determined that the loop has a single iteration.
-LogicalResult mlir::scf::promoteIfSingleIteration(PatternRewriter &rewriter,
-                                                  scf::ForallOp forallOp) {
+LogicalResult scf::ForallOp::promoteIfSingleIteration(RewriterBase &rewriter) {
   for (auto [lb, ub, step] :
-       llvm::zip(forallOp.getMixedLowerBound(), forallOp.getMixedUpperBound(),
-                 forallOp.getMixedStep())) {
+       llvm::zip(getMixedLowerBound(), getMixedUpperBound(), getMixedStep())) {
     auto tripCount = constantTripCount(lb, ub, step);
     if (!tripCount.has_value() || *tripCount != 1)
       return failure();
   }
 
-  promote(rewriter, forallOp);
+  promote(rewriter, *this);
   return success();
 }
 
 /// Promotes the loop body of a scf::ForallOp to its containing block.
-void mlir::scf::promote(PatternRewriter &rewriter, scf::ForallOp forallOp) {
-  IRMapping mapping;
-  mapping.map(forallOp.getInductionVars(), forallOp.getLowerBound(rewriter));
-  mapping.map(forallOp.getOutputBlockArguments(), forallOp.getOutputs());
-  for (auto &bodyOp : forallOp.getBody()->without_terminator())
-    rewriter.clone(bodyOp, mapping);
+void mlir::scf::promote(RewriterBase &rewriter, scf::ForallOp forallOp) {
+  OpBuilder::InsertionGuard g(rewriter);
+  scf::InParallelOp terminator = forallOp.getTerminator();
+
+  // Replace block arguments with lower bounds (replacements for IVs) and
+  // outputs.
+  SmallVector<Value> bbArgReplacements = forallOp.getLowerBound(rewriter);
+  bbArgReplacements.append(forallOp.getOutputs().begin(),
+                           forallOp.getOutputs().end());
+
+  // Move the loop body operations to the loop's containing block.
+  rewriter.inlineBlockBefore(forallOp.getBody(), forallOp->getBlock(),
+                             forallOp->getIterator(), bbArgReplacements);
 
+  // Replace the terminator with tensor.insert_slice ops.
+  rewriter.setInsertionPointAfter(forallOp);
   SmallVector<Value> results;
   results.reserve(forallOp.getResults().size());
-  scf::InParallelOp terminator = forallOp.getTerminator();
   for (auto &yieldingOp : terminator.getYieldingOps()) {
     auto parallelInsertSliceOp =
         cast<tensor::ParallelInsertSliceOp>(yieldingOp);
 
     Value dst = parallelInsertSliceOp.getDest();
     Value src = parallelInsertSliceOp.getSource();
-
-    auto getMappedValues = [&](ValueRange values) {
-      return llvm::to_vector(llvm::map_range(
-          values, [&](Value value) { return mapping.lookupOrDefault(value); }));
-    };
-
-    Value srcVal = mapping.lookupOrDefault(src);
-    if (llvm::isa<TensorType>(srcVal.getType())) {
+    if (llvm::isa<TensorType>(src.getType())) {
       results.push_back(rewriter.create<tensor::InsertSliceOp>(
-          forallOp.getLoc(), dst.getType(), srcVal,
-          mapping.lookupOrDefault(dst),
-          getMappedValues(parallelInsertSliceOp.getOffsets()),
-          getMappedValues(parallelInsertSliceOp.getSizes()),
-          getMappedValues(parallelInsertSliceOp.getStrides()),
+          forallOp.getLoc(), dst.getType(), src, dst,
+          parallelInsertSliceOp.getOffsets(), parallelInsertSliceOp.getSizes(),
+          parallelInsertSliceOp.getStrides(),
           parallelInsertSliceOp.getStaticOffsets(),
           parallelInsertSliceOp.getStaticSizes(),
           parallelInsertSliceOp.getStaticStrides()));
+    } else {
+      llvm_unreachable("unsupported terminator");
     }
   }
-  rewriter.replaceOp(forallOp, results);
+  rewriter.replaceAllUsesWith(forallOp.getResults(), results);
+
+  // Erase the old terminator and the loop.
+  rewriter.eraseOp(terminator);
+  rewriter.eraseOp(forallOp);
 }
 
 LoopNest mlir::scf::buildLoopNest(
     OpBuilder &builder, Location loc, ValueRange lbs, ValueRange ubs,
     ValueRange steps, ValueRange iterArgs,
     function_ref<ValueVector(OpBuilder &, Location, ValueRange, ValueRange)>
         bodyBuilder) {
   assert(lbs.size() == ubs.size() &&

mlir/lib/Dialect/SCF/Utils/Utils.cpp

static Value ceilDivPositive(OpBuilder &builder, Location loc, Value dividend,
   assert(dividend.getType().isIndex() && "expected index-typed value");
 
   Value cstOne = builder.create<arith::ConstantIndexOp>(loc, 1);
   Value divisorMinusOne = builder.create<arith::SubIOp>(loc, divisor, cstOne);
   Value sum = builder.create<arith::AddIOp>(loc, dividend, divisorMinusOne);
   return builder.create<arith::DivUIOp>(loc, sum, divisor);
 }
 
-/// Helper to replace uses of loop carried values (iter_args) and loop
-/// yield values while promoting single iteration scf.for ops.
-static void replaceIterArgsAndYieldResults(scf::ForOp forOp) {
-  // Replace uses of iter arguments with iter operands (initial values).
-  auto iterOperands = forOp.getIterOperands();
-  auto iterArgs = forOp.getRegionIterArgs();
-  for (auto e : llvm::zip(iterOperands, iterArgs))
-    std::get<1>(e).replaceAllUsesWith(std::get<0>(e));
-
-  // Replace uses of loop results with the values yielded by the loop.
-  auto outerResults = forOp.getResults();
-  auto innerResults = forOp.getBody()->getTerminator()->getOperands();
-  for (auto e : llvm::zip(outerResults, innerResults))
-    std::get<0>(e).replaceAllUsesWith(std::get<1>(e));
-}
-
-/// Promotes the loop body of a forOp to its containing block if the forOp
-/// it can be determined that the loop has a single iteration.
-LogicalResult mlir::promoteIfSingleIteration(scf::ForOp forOp) {
-  std::optional<int64_t> tripCount = constantTripCount(
-      forOp.getLowerBound(), forOp.getUpperBound(), forOp.getStep());
-  if (!tripCount.has_value() || tripCount != 1)
-    return failure();
-  auto iv = forOp.getInductionVar();
-  iv.replaceAllUsesWith(forOp.getLowerBound());
-
-  replaceIterArgsAndYieldResults(forOp);
-
-  // Move the loop body operations, except for its terminator, to the loop's
-  // containing block.
-  auto *parentBlock = forOp->getBlock();
-  forOp.getBody()->getTerminator()->erase();
-  parentBlock->getOperations().splice(Block::iterator(forOp),
-                                      forOp.getBody()->getOperations());
-  forOp.erase();
-  return success();
-}
-
 /// Generates unrolled copies of scf::ForOp 'loopBodyBlock', with
 /// associated 'forOpIV' by 'unrollFactor', calling 'ivRemapFn' to remap
 /// 'forOpIV' for each unrolled body. If specified, annotates the Ops in each
 /// unrolled iteration using annotateFn.
 static void generateUnrolledLoop(
     Block *loopBodyBlock, Value forOpIV, uint64_t unrollFactor,
     function_ref<Value(unsigned, Value, OpBuilder)> ivRemapFn,
     function_ref<void(unsigned, Operation *, OpBuilder)> annotateFn,
LogicalResult mlir::loopUnrollByFactor(
 
   // Return if the loop body is empty.
   if (llvm::hasSingleElement(forOp.getBody()->getOperations()))
     return success();
 
   // Compute tripCount = ceilDiv((upperBound - lowerBound), step) and populate
   // 'upperBoundUnrolled' and 'stepUnrolled' for static and dynamic cases.
   OpBuilder boundsBuilder(forOp);
+  IRRewriter rewriter(forOp.getContext());
   auto loc = forOp.getLoc();
   Value step = forOp.getStep();
   Value upperBoundUnrolled;
   Value stepUnrolled;
   bool generateEpilogueLoop = true;
 
   auto lbCstOp = forOp.getLowerBound().getDefiningOp<arith::ConstantIndexOp>();
   auto ubCstOp = forOp.getUpperBound().getDefiningOp<arith::ConstantIndexOp>();
   auto stepCstOp = forOp.getStep().getDefiningOp<arith::ConstantIndexOp>();
   if (lbCstOp && ubCstOp && stepCstOp) {
     // Constant loop bounds computation.
     int64_t lbCst = lbCstOp.value();
     int64_t ubCst = ubCstOp.value();
     int64_t stepCst = stepCstOp.value();
     assert(lbCst >= 0 && ubCst >= 0 && stepCst >= 0 &&
            "expected positive loop bounds and step");
     int64_t tripCount = mlir::ceilDiv(ubCst - lbCst, stepCst);
 
     if (unrollFactor == 1) {
-      if (tripCount == 1 && failed(promoteIfSingleIteration(forOp)))
+      if (tripCount == 1 && failed(forOp.promoteIfSingleIteration(rewriter)))
         return failure();
       return success();
     }
 
     int64_t tripCountEvenMultiple = tripCount - (tripCount % unrollFactor);
     int64_t upperBoundUnrolledCst = lbCst + tripCountEvenMultiple * stepCst;
     assert(upperBoundUnrolledCst <= ubCst);
     int64_t stepUnrolledCst = stepCst * unrollFactor;
if (generateEpilogueLoop) {
     auto results = forOp.getResults();
     auto epilogueResults = epilogueForOp.getResults();
 
     for (auto e : llvm::zip(results, epilogueResults)) {
       std::get<0>(e).replaceAllUsesWith(std::get<1>(e));
     }
     epilogueForOp->setOperands(epilogueForOp.getNumControlOperands(),
                                epilogueForOp.getNumIterOperands(), results);
-    (void)promoteIfSingleIteration(epilogueForOp);
+    (void)epilogueForOp.promoteIfSingleIteration(rewriter);
   }
 
   // Create unrolled loop.
   forOp.setUpperBound(upperBoundUnrolled);
   forOp.setStep(stepUnrolled);
 
   auto iterArgs = ValueRange(forOp.getRegionIterArgs());
   auto yieldedValues = forOp.getBody()->getTerminator()->getOperands();
 
   generateUnrolledLoop(
       forOp.getBody(), forOp.getInductionVar(), unrollFactor,
       [&](unsigned i, Value iv, OpBuilder b) {
         // iv' = iv + step * i;
         auto stride = b.create<arith::MulIOp>(
             loc, step, b.create<arith::ConstantIndexOp>(loc, i));
         return b.create<arith::AddIOp>(loc, iv, stride);
       },
       annotateFn, iterArgs, yieldedValues);
   // Promote the loop body up if this has turned into a single iteration loop.
-  (void)promoteIfSingleIteration(forOp);
+  (void)forOp.promoteIfSingleIteration(rewriter);
   return success();
 }
 
 /// Return the new lower bound, upper bound, and step in that order. Insert any
 /// additional bounds calculations before the given builder and any additional
 /// conversion back to the original loop induction value inside the given Block.
 static LoopParams normalizeLoop(OpBuilder &boundsBuilder,
                                 OpBuilder &insideLoopBuilder, Location loc,

mlir/test/Dialect/SCF/one-shot-bufferize.mlir

func.func @parallel_insert_slice_no_conflict(
     %idx: index,
     %idx2: index,
     %arg1: tensor<?xf32> {bufferization.writable = true},
     %arg2: tensor<?xf32> {bufferization.writable = true}) -> (tensor<?xf32>, f32) {
   %cst = arith.constant 4.200000e+01 : f32
   %c0 = arith.constant 0 : index
   %c1 = arith.constant 1 : index
 
+  // CHECK: %[[subview:.*]] = memref.subview %[[arg2]][5] [%[[idx]]] [1]
   // CHECK: scf.forall (%[[tidx:.*]]) in (%[[idx2]])
   %2 = scf.forall (%arg3) in (%idx2) shared_outs(%o = %arg2) -> (tensor<?xf32>) {
-      // CHECK: %[[subview:.*]] = memref.subview %[[arg2]][5] [%[[idx]]] [1]
       %6 = tensor.extract_slice %o[5] [%idx] [%c1] : tensor<?xf32> to tensor<?xf32>
       // CHECK: linalg.fill ins(%{{.*}}) outs(%[[subview]] : memref<?xf32
       %8 = linalg.fill ins(%cst : f32) outs(%6 : tensor<?xf32>) -> tensor<?xf32>
       // Self-copy will DCE away later.
       // CHECK: memref.copy %[[subview]], %[[subview]]
 
       // Empty terminator is elided from pretty-printing.
       // CHECK-NOT: scf.forall.in_parallel
func.func @parallel_insert_slice_with_conflict(
   %c0 = arith.constant 0 : index
   %c1 = arith.constant 1 : index
 
   // The parallel_insert_slice_op bufferizes out-of-place due to a RAW conflict
   // on %arg2, so we need an allocation.
   // CHECK: %[[alloc1:.*]] = memref.alloc
   // CHECK: memref.copy %[[arg2]], %[[alloc1]]
 
+  // CHECK: %[[subview1:.*]] = memref.subview %[[alloc1]][5] [%[[idx]]] [1]
   // CHECK: scf.forall (%[[tidx:.*]]) in (%[[idx2]])
   %2 = scf.forall (%arg3) in (%idx2) shared_outs(%o = %arg2) -> (tensor<?xf32>) {
-      // CHECK: %[[subview1:.*]] = memref.subview %[[alloc1]][5] [%[[idx]]] [1]
       %6 = tensor.extract_slice %o[5] [%idx] [%c1] : tensor<?xf32> to tensor<?xf32>
 
       // CHECK: linalg.fill ins(%{{.*}}) outs(%[[subview1]] : memref<?xf32
       %8 = linalg.fill ins(%cst : f32) outs(%6 : tensor<?xf32>) -> tensor<?xf32>
 
       // Now the copy of the actual insert_slice. (It will fold away.)
       // CHECK: memref.copy %[[subview1]], %[[subview1]]