[mlir] Add a pattern to fold single- and zero-iteration scf.forall ops.

[nicolasvasilache]

let's give some time to review

[pifon2a]

update the test.

[nicolasvasilache]

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

1477

How about generalizing promoteIfSingleIteration to also have an implementation for scf.forall ?
This is a common API already used for both scf.for and affine.for (with different implementations).

We can certainly wrap a pattern around and expose with populate functions but by now it is pretty well established the key functionality is better exposed via single functional-style functions.

Additionally, the promoteIfSingleIteration are in need of a refresh while we are at it: they predate the proper use of rewriter APIs and perform unsafe RAUW, this is a good opportunity to tighten things up.

1484

Please add a helper in StaticValueUtils that returns the optional number of steps:

optional<int64_t> constantTripCount(OpFoldResult lb, OpFoldResult ub, OpFoldResult step);

We could go as far as to implement that with OpFoldResult applied = makeComposedFoldedAffineApply which would capture a bunch of nontrivial cases (but would potentially build an op unless evolved to guarantee it remains in type-land).

Then your loop goes away and the code is significantly more idiomatic:

if (llvm::any_of(... [](){ return ... == 0; }))
  rewriter.replaceOp(op, op.getOutputs());
      
llvm::make_filter_range(... [](){ if (... != 1) ...; })

etc

1511

rewriter.notifyMatchFailure

1515

This seems like the part that performs the same as promoteIfSingleIteration for scf.forall.
This seems can be extracted in a well-defined API, brought closer to the one for scf.for and reused here, instead of inlining implementations.

1553

This feels like a generally useful utility on the op itself to provide a helper that "projects out" one of the dimensions and returning an (N-1)-D version of the op.
This can then be also reused to implement other functionality (e.g. mapping to a physical thread id that currently may jump the abstraction gap a bit too quickly).

1567

This may need to be wrapped in a rewriter.startRootUpdate

1581

overall the pattern itself should be ~5 lines and we gain a few new reusable APIs

[pifon2a]

Address the comments.

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

1477

added promoteIfSingleIteration to SCF.h instead of scf utils.h, because otherwise there is a circular dependency.

1484

I introduced the helper, but I am not sure about idiomatic part. I tried it and it was similar amount of code, but less efficient.

1553

Shall we refactor it later, when we know the exact use case?

[herhut]

Looks good conditional on small nits.

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

529

nit: please format.

1553

This seems reasonable to me.

2738

nit: Give this a better name. ParallelOpSingleOrZeroIterationDimsFolder maybe? That would be consistent with the other one.

mlir/lib/Dialect/Utils/StaticValueUtils.cpp

9 ↗

(On Diff #504449)

nit: is this needed?

[tschuett]

mlir/lib/Dialect/Utils/StaticValueUtils.cpp

202 ↗

(On Diff #504449)

You use .has_value() in several places. They are all redundant.

[pifon2a]

Address the comments.

[nicolasvasilache]

thanks for picking up the slack @herhut , won't have time to look deeper atm so please go with this

[pifon2a]

address comments

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

void ForOp::getSuccessorRegions(std::optional<unsigned> index,
   assert(*index == 0 && "expected loop region");
   regions.push_back(RegionSuccessor(&getLoopBody(), getRegionIterArgs()));
   regions.push_back(RegionSuccessor(getResults()));
 }
 
 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() &&
          "expected the same number of lower and upper bounds");
   assert(lbs.size() == steps.size() &&
          "expected the same number of lower bounds and steps");
 
   // If there are no bounds, call the body-building function and return early.
   if (lbs.empty()) {
rewriter.updateRootInPlace(op, [&]() {
       dispatchIndexOpFoldResults(mixedStep, dynamicStep, staticStep);
       op.getDynamicStepMutable().assign(dynamicStep);
       op.setStaticStep(staticStep);
     });
     return success();
   }
 };
 
+struct ForallOpSingleOrZeroIterationDimsFolder
+    : public OpRewritePattern<ForallOp> {
+  using OpRewritePattern<ForallOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(ForallOp op,
+                                PatternRewriter &rewriter) const override {
+    // Do not fold dimensions if they are mapped to processing units.
+    if (op.getMapping().has_value())
+      return failure();
+    Location loc = op.getLoc();
+
+    // Compute new loop bounds that omit all single-iteration loop dimensions.
+    SmallVector<OpFoldResult> newMixedLowerBounds, newMixedUpperBounds,
+        newMixedSteps;
+    IRMapping mapping;
+    for (auto [lowerBound, upperBound, step, iv] :
+         llvm::zip(op.getMixedLowerBound(), op.getMixedUpperBound(),
+                   op.getMixedStep(), op.getInductionVars())) {
+      // Collect the statically known loop bounds.
+      std::optional<int64_t> lowerBoundConstant =
+          getConstantIntValue(lowerBound);
+      std::optional<int64_t> upperBoundConstant =
+          getConstantIntValue(upperBound);
+      std::optional<int64_t> stepConstant = getConstantIntValue(step);
+      // Remove the loop if it performs zero iterations.
+      if (lowerBoundConstant && upperBoundConstant &&
+          *lowerBoundConstant == *upperBoundConstant) {
+        rewriter.replaceOp(op, op.getOutputs());
+        return success();
+      }
+      // Replace the loop induction variable by the lower bound if the loop
+      // performs a single iteration. Otherwise, copy the loop bounds.
+      if (lowerBoundConstant && upperBoundConstant && stepConstant &&
+          (*upperBoundConstant - *lowerBoundConstant) > 0 &&
+          (*upperBoundConstant - *lowerBoundConstant) <= *stepConstant) {
+        mapping.map(iv,
+                    getValueOrCreateConstantIndexOp(rewriter, loc, lowerBound));
+      } else {
+        newMixedLowerBounds.push_back(lowerBound);
+        newMixedUpperBounds.push_back(upperBound);
+        newMixedSteps.push_back(step);
+      }
+    }
+    // Exit if none of the loop dimensions perform a single iteration.
+    if (newMixedLowerBounds.size() == static_cast<unsigned>(op.getRank())) {
+      return failure();
+    }
+
+    // All of the loop dimensions perform a single iteration. Inline loop body.
+    if (newMixedLowerBounds.empty()) {
+      mapping.map(op.getOutputBlockArguments(), op.getOutputs());
+      for (auto &bodyOp : op.getBody()->without_terminator())
+        rewriter.clone(bodyOp, mapping);
+      SmallVector<Value> results;
+      results.reserve(op.getResults().size());
+      scf::InParallelOp terminator = op.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 (srcVal.getType().isa<TensorType>()) {
+          results.push_back(rewriter.create<tensor::InsertSliceOp>(
+              op.getLoc(), dst.getType(), srcVal, mapping.lookupOrDefault(dst),
+              getMappedValues(parallelInsertSliceOp.getOffsets()),
+              getMappedValues(parallelInsertSliceOp.getSizes()),
+              getMappedValues(parallelInsertSliceOp.getStrides()),
+              parallelInsertSliceOp.getStaticOffsets(),
+              parallelInsertSliceOp.getStaticSizes(),
+              parallelInsertSliceOp.getStaticStrides()));
+        }
+      }
+      rewriter.replaceOp(op, results);
+      return success();
+    }
+
+    // Replace the loop by a lower-dimensional loop.
+    ForallOp newOp;
+    newOp = rewriter.create<ForallOp>(loc, newMixedLowerBounds,
+                                      newMixedUpperBounds, newMixedSteps,
+                                      op.getOutputs(), std::nullopt, nullptr);
+    newOp.getBodyRegion().getBlocks().clear();
+    // The new loop needs to keep all attributes from the old one, except for
+    // "operand_segment_sizes" and static loop bound attributes which capture
+    // the outdated information of the old iteration domain.
+    SmallVector<StringAttr> elidedAttrs{newOp.getOperandSegmentSizesAttrName(),
+                                        newOp.getStaticLowerBoundAttrName(),
+                                        newOp.getStaticUpperBoundAttrName(),
+                                        newOp.getStaticStepAttrName()};
+    for (const auto &namedAttr : op->getAttrs()) {
+      if (llvm::is_contained(elidedAttrs, namedAttr.getName()))
+        continue;
+      newOp->setAttr(namedAttr.getName(), namedAttr.getValue());
+    }
+    rewriter.cloneRegionBefore(op.getRegion(), newOp.getRegion(),
+                               newOp.getRegion().begin(), mapping);
+    rewriter.replaceOp(op, newOp.getResults());
+    return success();
+  }
+};
+
 } // namespace
 
 void ForallOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                            MLIRContext *context) {
-  results.add<DimOfForallOp, ForallOpControlOperandsFolder>(context);
+  results.add<DimOfForallOp, ForallOpControlOperandsFolder,
+              ForallOpSingleOrZeroIterationDimsFolder>(context);
 }
 
 //===----------------------------------------------------------------------===//
 // InParallelOp
 //===----------------------------------------------------------------------===//
 
 // Build a InParallelOp with mixed static and dynamic entries.
 void InParallelOp::build(OpBuilder &b, OperationState &result) {
if (!ivArg)
     return ParallelOp();
   assert(ivArg.getOwner() && "unlinked block argument");
   auto *containingOp = ivArg.getOwner()->getParentOp();
   return dyn_cast<ParallelOp>(containingOp);
 }
 
 namespace {
 // Collapse loop dimensions that perform a single iteration.
 struct CollapseSingleIterationLoops : public OpRewritePattern<ParallelOp> {
   using OpRewritePattern<ParallelOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(ParallelOp op,
                                 PatternRewriter &rewriter) const override {
     IRMapping mapping;
     // Compute new loop bounds that omit all single-iteration loop dimensions.
     SmallVector<Value, 2> newLowerBounds;
     SmallVector<Value, 2> newUpperBounds;

mlir/test/Dialect/SCF/canonicalize.mlir

scf.forall.in_parallel {
       tensor.parallel_insert_slice %slice into %o[%i, %j] [1, 1] [1, 1]
         : tensor<1x1xf32> into tensor<?x10xf32>
     }
   }
 
   return %result : tensor<?x10xf32>
 }
 // CHECK: forall (%{{.*}}, %{{.*}}) in (%{{.*}}, 10)
+
+// -----
+
+func.func @inline_forall_loop(%in: tensor<8x8xf32>) -> tensor<8x8xf32> {
+  %c8 = arith.constant 8 : index
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %cst = arith.constant 0.000000e+00 : f32
+  %0 = tensor.empty() : tensor<8x8xf32>
+  %1 = scf.forall (%i, %j) = (%c0, %c0) to (%c1, %c1)
+        step (%c8, %c8) shared_outs (%out_ = %0) -> (tensor<8x8xf32>) {
+    %slice = tensor.extract_slice %out_[%i, %j] [2, 3] [1, 1]
+      : tensor<8x8xf32> to tensor<2x3xf32>
+    %fill = linalg.fill ins(%cst : f32) outs(%slice : tensor<2x3xf32>)
+          -> tensor<2x3xf32>
+    scf.forall.in_parallel {
+      tensor.parallel_insert_slice %fill into %out_[%i, %j] [2, 3] [1, 1]
+        : tensor<2x3xf32> into tensor<8x8xf32>
+    }
+  }
+  return %1 : tensor<8x8xf32>
+}
+// CHECK-LABEL: @inline_forall_loop
+// CHECK-NOT:     scf.forall
+// CHECK:         %[[OUT:.*]] = tensor.empty
+
+// CHECK-NEXT:    %[[SLICE:.*]] = tensor.extract_slice %[[OUT]]
+// CHECK-SAME:      : tensor<8x8xf32> to tensor<2x3xf32>
+
+// CHECK-NEXT:    %[[FILL:.*]] = linalg.fill
+// CHECK-SAME:      outs(%[[SLICE]]
+
+// CHECK-NEXT:    tensor.insert_slice %[[FILL]]
+// CHECK-SAME:      : tensor<2x3xf32> into tensor<8x8xf32>
+
+// -----
+
+func.func @do_not_inline_distributed_forall_loop(
+    %in: tensor<8x8xf32>) -> tensor<8x8xf32> {
+  %cst = arith.constant 0.000000e+00 : f32
+  %0 = tensor.empty() : tensor<8x8xf32>
+  %1 = scf.forall (%i, %j) = (0, 0) to (1, 1) step (8, 8)
+      shared_outs (%out_ = %0) -> (tensor<8x8xf32>) {
+    %slice = tensor.extract_slice %out_[%i, %j] [2, 3] [1, 1]
+      : tensor<8x8xf32> to tensor<2x3xf32>
+    %fill = linalg.fill ins(%cst : f32) outs(%slice : tensor<2x3xf32>)
+          -> tensor<2x3xf32>
+    scf.forall.in_parallel {
+      tensor.parallel_insert_slice %fill into %out_[%i, %j] [2, 3] [1, 1]
+        : tensor<2x3xf32> into tensor<8x8xf32>
+    }
+  }{ mapping = [#gpu.thread<y>, #gpu.thread<x>] }
+  return %1 : tensor<8x8xf32>
+}
+// CHECK-LABEL: @do_not_inline_distributed_forall_loop
+// CHECK: scf.forall
+
+// -----
+
+func.func @collapse_one_dim_parallel(%in: tensor<8x8xf32>) -> tensor<8x8xf32> {
+  %c8 = arith.constant 8 : index
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %c16 = arith.constant 16 : index
+  %cst = arith.constant 0.000000e+00 : f32
+  %0 = tensor.empty() : tensor<8x8xf32>
+  %1 = scf.forall (%i, %j) = (0, %c0) to (1, %c16)
+        step (8, %c8) shared_outs (%out_ = %0) -> (tensor<8x8xf32>) {
+    %fill = linalg.fill ins(%cst : f32) outs(%out_ : tensor<8x8xf32>)
+          -> tensor<8x8xf32>
+    scf.forall.in_parallel {
+      tensor.parallel_insert_slice %fill into %out_[%i, %j] [8, 8] [1, 1]
+        : tensor<8x8xf32> into tensor<8x8xf32>
+    }
+  }
+  return %1 : tensor<8x8xf32>
+}
+// CHECK-LABEL: @collapse_one_dim_parallel
+// CHECK:         scf.forall (%[[ARG:.*]]) = (0) to (16) step (8)
+// CHECK:           linalg.fill
+// CHECK:           tensor.parallel_insert_slice
+
+// -----
+
+func.func @remove_empty_forall(%in: tensor<8x8xf32>) -> tensor<8x8xf32> {
+  %c8 = arith.constant 8 : index
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %c16 = arith.constant 16 : index
+  %cst = arith.constant 0.000000e+00 : f32
+  %0 = tensor.empty() : tensor<8x8xf32>
+  %1 = scf.forall (%i, %j) = (%c0, %c16) to (%c1, %c16)
+        step (%c8, %c8) shared_outs (%out_ = %0) -> (tensor<8x8xf32>) {
+    %fill = linalg.fill ins(%cst : f32) outs(%out_ : tensor<8x8xf32>)
+          -> tensor<8x8xf32>
+    scf.forall.in_parallel {
+      tensor.parallel_insert_slice %fill into %out_[%i, %j] [8, 8] [1, 1]
+        : tensor<8x8xf32> into tensor<8x8xf32>
+    }
+  }
+  return %1 : tensor<8x8xf32>
+}
+// CHECK-LABEL: @remove_empty_forall
+// CHECK-NOT:   scf.forall
+// CHECK:       %[[EMPTY:.*]] = tensor.empty
+// CHECK:       return %[[EMPTY]]
+

mlir/test/Dialect/Tensor/fold-consecutive-insert-extract-slice.mlir

 
 //   CHECK-LABEL: func.func @insert_slice_rank_reducing_dynamic_shape
 // CHECK-COUNT-2:   tensor.insert_slice
 
 // -----
 
 // CHECK-LABEL: func.func @parallel_insert_slice
 //   CHECK-NOT:   tensor.insert_slice
-//       CHECK:   tensor.parallel_insert_slice %{{.*}} into %{{.*}}[%{{.*}}, %{{.*}}] [1, 1] [1, 1] : tensor<f32> into tensor<1x2xf32>
+//       CHECK:   tensor.parallel_insert_slice %{{.*}} into %{{.*}}[0, %{{.*}}] [1, 1] [1, 1] : tensor<f32> into tensor<1x2xf32>
 func.func @parallel_insert_slice(%t0: tensor<1x2xf32>, %t1: tensor<f32>, %t2: tensor<1x1xf32>) -> tensor<1x2xf32> {
   %c1 = arith.constant 1 : index
   %c2 = arith.constant 2 : index
   %r = scf.forall (%arg2, %arg3) in (%c1, %c2) shared_outs(%arg4 = %t0) -> (tensor<1x2xf32>) {
     %inserted_slice = tensor.insert_slice %t1 into %t2[0, 0] [1, 1] [1, 1] : tensor<f32> into tensor<1x1xf32>
     scf.forall.in_parallel {
       tensor.parallel_insert_slice %inserted_slice into %arg4[%arg2, %arg3] [1, 1] [1, 1] : tensor<1x1xf32> into tensor<1x2xf32>
     }
   }
   return %r : tensor<1x2xf32>
 }