[mlir][linalg] Add extra pattern to AffineMinSCFCanonicalization

[ThomasRaoux]

Summary

Add an ad hoc pattern to canonicalize affine.min ops after doing tile and distribute.

[mravishankar]

It is very specific use case :) .
I am fine with this since it is probably a stop-gap solution. Ill let @nicolasvasilache decide if this lives in core.
Treat as LGTM from me, but just dont want to stamp it till I see what Nicolas has to say.

mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp

523

Confusing name. The affine_min itself is not divisible. Cant suggest a better name right now.

565

Maybe make it simpler and just check that numResults() == 2?

[ThomasRaoux]

In D99558#2657637, @mravishankar wrote:

It is very specific use case :) .
I am fine with this since it is probably a stop-gap solution. Ill let @nicolasvasilache decide if this lives in core.
Treat as LGTM from me, but just dont want to stamp it till I see what Nicolas has to say.

Thanks for looking at it. It handles a very narrow case indeed. I'll wait for @nicolasvasilache to give his opinion and I'll address the comments.

mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp

if (stage3Lambda) {
         return failure();
       LLVM_DEBUG(DBGS() << "After 3rd stage, iter : " << iteration << "\n"
                         << *op);
     }
   }
   return success();
 }
 
+static bool isDivisible(Value v, int64_t dividend);
+
+/// Return true if we can prove that affineMinOp result is positive and
+/// divisible by the given |dividend|. This is true if all the the results of
+/// the associated affine map are positive and divisible by |dividend|.
+/// This speciically look for the following pattern:
+/// ```
+/// scf.for %iv = %lb to %ub step %step
+///   %affine.min affine_map<(d0, d1) -> (N, d0 - d1)>(%ub, %iv)
+/// ```
+static bool affineMinOpDivisible(AffineMinOp minOp, int64_t dividend) {
+  if (!minOp.getSymbolOperands().empty())
+    return {};
+
+  scf::ForOp forOp;
+  Value iv;
+  // Check if any of the dimensions are a ForOp induction variable.
+  for (auto dim : minOp.getDimOperands()) {
+    auto ivArg = dim.dyn_cast<BlockArgument>();
+    if (!ivArg)
+      continue;
+    Operation *containingOp = ivArg.getOwner()->getParentOp();
+    forOp = dyn_cast_or_null<scf::ForOp>(containingOp);
+    if (forOp && forOp.getInductionVar() == dim) {
+      iv = dim;
+      break;
+    }
+  }
+  if (!iv)
+    return false;
+  // Calculate the affine map representing `%ub - %iv`.
+  Value ub = forOp.upperBound();
+  AffineExpr ivDim;
+  AffineExpr ubDim;
+  for (auto dim : llvm::enumerate(minOp.getDimOperands())) {
+    if (dim.value() == iv)
+      ivDim = getAffineDimExpr(dim.index(), minOp.getContext());
+    else if (dim.value() == ub)
+      ubDim = getAffineDimExpr(dim.index(), minOp.getContext());
+    else
+      return false;
+  }
+
+  if (!ubDim) {
+    if (auto cstUb = ub.getDefiningOp<ConstantIndexOp>())
+      ubDim = getAffineConstantExpr(cstUb.getValue(), minOp.getContext());
+    else
+      return false;
+  }
+  AffineExpr diffExp = ubDim - ivDim;
+  // Check that all the affine map results are either constant divisible by
+  // `dividend` or equal to `%ub - %iv`.
+  for (AffineExpr result : minOp.getAffineMap().getResults()) {
+    if (auto cst = result.dyn_cast<AffineConstantExpr>()) {
+      if (cst.getValue() <= 0 || cst.getValue() % dividend != 0)
+        return false;
+    } else {
+      if (diffExp != result)
+        return false;
+    }
+  }
+  // Now check that for every value of the induction variable `%ub - %iv` is
+  // divisible by `dividend`. It is true if the lower bounder, the upper bound
+  // and the step are all divisible by `dividend`.
+  std::array<Value, 3> forOperands = {forOp.lowerBound(), forOp.step(),
+                                      forOp.upperBound()};
+  return llvm::all_of(forOperands,
+                      [dividend](Value v) { return isDivisible(v, dividend); });
+}
+
+/// Return true if we can prove that the value |v| is always divisible by the
+/// constant |dividend|. Return false otherwise.
+static bool isDivisible(Value v, int64_t dividend) {
+  MLIRContext *ctx = v.getContext();
+  // Create an expression (d0) -> (d0 % n) and try to simplify it.
+  AffineExpr mod = getAffineDimExpr(0, ctx) % dividend;
+  AffineMap modMap = AffineMap::get(1, 0, {mod}, ctx);
+  SmallVector<Value> ops(1, v);
+  fullyComposeAffineMapAndOperands(&modMap, &ops);
+  canonicalizeMapAndOperands(&modMap, &ops);
+  modMap = simplifyAffineMap(modMap);
+  auto cst = modMap.getResult(0).dyn_cast<AffineConstantExpr>();
+  if (cst)
+    return (cst.getValue() == 0);
+  // If the map doesn't fold to 0 but simplifies to (d0 %n) with d0 an
+  // affine.min, check if all the results of the affine.min's map are divisible
+  // by `dividend`.
+  if (modMap.getResult(0) != mod)
+    return false;
+  assert(ops.size() == 1);
+  auto minOp = ops[0].getDefiningOp<AffineMinOp>();
+  return (minOp && affineMinOpDivisible(minOp, dividend));
+}
+
+/// Try to fold a affine.min op by matching the following form:
+/// ```
+/// scf.for %iv = %lb to %ub step %step
+///   %affine.min affine_map<(d0, d1) -> (N, d0 - d1)>(%ub, %iv)
+/// ```
+/// With N a compile time constant. This operations can be replace by
+/// `%cN = constant N : index` if we can prove that %lb, %step and %ub are
+/// divisible by N.
+static Optional<int64_t> foldAffineMin(AffineMinOp minOp) {
+  AffineMap map = minOp.getAffineMap();
+  int64_t constantResult = 0;
+  for (AffineExpr result : map.getResults()) {
+    if (auto cst = result.dyn_cast<AffineConstantExpr>())
+      constantResult = cst.getValue();
+  }
+  if (constantResult == 0)
+    return {};
+  // If afine.min map's results are all positive and divisible by
+  // `constantResult` then it can be replaced by `constantResult`.
+  if (affineMinOpDivisible(minOp, constantResult))
+    return constantResult;
+  return {};
+}
+
 /// Given the `lbVal`, `ubVal` and `stepVal` of a loop, append `lbVal` and
 /// `ubVal` to `dims` and `stepVal` to `symbols`.
 /// Create new AffineDimExpr (`%lb` and `%ub`) and AffineSymbolExpr (`%step`)
 /// with positions matching the newly appended values. Substitute occurrences of
 /// `dimExpr` by either the min expression (i.e. `%lb`) or the max expression
 /// (i.e. `%lb + %step * floordiv(%ub -1 - %lb, %step)`), depending on whether
 /// the induction variable is used with a positive or negative  coefficient.
 static AffineExpr substituteLoopInExpr(AffineExpr expr, AffineExpr dimExpr,
AffineMapAndOperands mlir::linalg::substituteMin(
   return res;
 }
 
 LogicalResult AffineMinSCFCanonicalizationPattern::matchAndRewrite(
     AffineMinOp minOp, PatternRewriter &rewriter) const {
   LLVM_DEBUG(DBGS() << "Canonicalize AffineMinSCF: " << *minOp.getOperation()
                     << "\n");
 
+  // First check some special pattern where affine.min can be replaced by a
+  // constant.
+  if (Optional<int64_t> cst = foldAffineMin(minOp)) {
+    rewriter.replaceOpWithNewOp<ConstantOp>(minOp, rewriter.getIndexAttr(*cst));
+    return success();
+  }
+
   auto affineMapAndOperands = substituteMin(minOp);
   AffineMap map = affineMapAndOperands.map;
 
   LLVM_DEBUG(DBGS() << "Resulting map: " << map << "\n");
 
   // Check whether any of the expressions, when subtracted from all other
   // expressions, produces only >= 0 constants. If so, it is the min.
   for (auto e : minOp.getAffineMap().getResults()) {

mlir/test/Dialect/Linalg/fold-affine-min-scf.mlir

func @scf_parallel(%A : memref<i64>, %step : index) {
   scf.parallel (%i) = (%c0) to (%ub2) step (%step) {
     %1 = affine.min affine_map<(d0, d1, d2) -> (d0, d2 - d1)> (%step, %i, %ub2)
     %2 = index_cast %1: index to i64
     memref.store %2, %A[]: memref<i64>
   }
 
   return
 }
+
+// CHECK-LABEL: scf_for_dynamic
+func @scf_for_dynamic(%A : memref<i64>, %id1 : index, %count1 : index,
+                      %id2 : index, %count2 : index) {
+  %c1020 = constant 1020 : index
+  %c1024 = constant 1024 : index
+  %0 = affine.apply affine_map<()[s0] -> (s0 * 32)>()[%id1]
+  %1 = affine.apply affine_map<()[s0] -> (s0 * 32)>()[%count1]
+
+  //      CHECK: scf.for
+  //      CHECK:   %[[C32:.*]] = constant 32 : index
+  //      CHECK:   scf.for %{{.*}} = %{{.*}} to %[[C32]]
+  // CHECK-NEXT:     %[[C4:.*]] = constant 4 : index
+  // CHECK-NEXT:     %[[C4I64:.*]] = index_cast %[[C4:.*]]
+  // CHECK-NEXT:     memref.store %[[C4I64]], %{{.*}}[] : memref<i64>
+  scf.for %arg0 = %0 to %c1024 step %1 {
+    %2 = affine.min affine_map<(d0) -> (32, -d0 + 1024)>(%arg0)
+    %3 = affine.apply affine_map<()[s0] -> (s0 * 4)>()[%id2]
+    %4 = affine.apply affine_map<()[s0] -> (s0 * 4)>()[%count2]
+    scf.for %arg1 = %3 to %2 step %4 {
+      %5 = affine.min affine_map<(d0, d1) -> (4, d0 - d1)>(%2, %arg1)
+      %6 = index_cast %5: index to i64
+      memref.store %6, %A[]: memref<i64>
+    }
+  }
+
+  // In this case the first affine.min cannot be removed as the 1020 is not
+  // divisible by 32 but we can still remove the second level of affine.min
+  // since we know that %2 is always divisible by 4.
+  //      CHECK: scf.for
+  //      CHECK:   %[[MIN:.*]] = affine.min
+  //      CHECK:   scf.for %{{.*}} = %{{.*}} to %[[MIN]]
+  // CHECK-NEXT:     %[[C4:.*]] = constant 4 : index
+  // CHECK-NEXT:     %[[C4I64:.*]] = index_cast %[[C4:.*]]
+  // CHECK-NEXT:     memref.store %[[C4I64]], %{{.*}}[] : memref<i64>
+  scf.for %arg0 = %0 to %c1020 step %1 {
+    %2 = affine.min affine_map<(d0) -> (32, -d0 + 1020)>(%arg0)
+    %3 = affine.apply affine_map<()[s0] -> (s0 * 4)>()[%id2]
+    %4 = affine.apply affine_map<()[s0] -> (s0 * 4)>()[%count2]
+    scf.for %arg1 = %3 to %2 step %4 {
+      %5 = affine.min affine_map<(d0, d1) -> (4, d0 - d1)>(%2, %arg1)
+      %6 = index_cast %5: index to i64
+      memref.store %6, %A[]: memref<i64>
+    }
+  }
+
+  return
+}