[mlir] Add multi-level AffineMinSCF canonicalization support

[nicolasvasilache]

Summary

[mlir] Add multi-level AffineMinSCF floding support

This revision allows chains of AffineMinOp that depend on SCF induction variables to fold. This allows the foldings of IR such as:

#map1 = affine_map<(d0) -> (128, -d0 + 192)>
#map2 = affine_map<(d0, d1) -> (16, d0 - d1)>

scf.parallel (%i) = (%c0) to (%c192) step (%c128) {
  %1 = affine.min #map1(%i)
  scf.for %j = %c0 to %1 step %c16 {
    %2 = affine.min #map2(%1, %j)
    %3 = index_cast %2: index to i64
    store %3, %A[]: memref<i64>

into:

#map1 = affine_map<(d0) -> (128, -d0 + 192)>

scf.parallel (%arg1) = (%c0) to (%c192) step (%c128) {
  %0 = affine.min #map1(%arg1)
  scf.for %arg2 = %c0 to %0 step %c16 {
    store %c16_i64, %arg0[] : memref<i64>

Note that handling chains is required because 192 % 128 != 0 but 192 % 16 == 0.

[bondhugula]

This revision allows chains of AffineMinOp that depend on SCF induction variables to canonicalize.

to canonicalize -> to fold.

What's map1 and map3 in the commit summary example?

[nicolasvasilache]

Reword.

[ftynse]

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

310

Nit: can we invert the condition and remove one indentaion level?

311

Nit: llvm::to_vector

339

Nit: consider adding a debug print here, we may want to know when we hit pathological cases

[nicolasvasilache]

Address comments.

[bondhugula]

mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h

509

outerloops -> outer loops

Again, canonicalize -> fold?

519

It's not %1 that canonicalizes, but it's the affine.min that folds.

mlir/include/mlir/IR/AffineExpr.h

137

Should this just be shiftSymbols or replaceByShiftedSymbols?

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

328

Comment please here.

mlir/lib/IR/AffineExpr.cpp

129

Missing assertion on valid numDims values.

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

120

Please add a one line comment on what it's testing.

124

Missing CHECK-LABEL.

133

Nit: indent by two for readability.

145

fold

[limo1996]

mlir/lib/IR/AffineExpr.cpp

143

This should be getAffineSymbolExpr

mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h

 
 /// Canonicalize AffineMinOp operations in the context of enclosing scf.for and
 /// scf.parallel by:
 ///   1. building an affine map where uses of the induction variable of a loop
 ///   are replaced by a `%lb + %step * floorDiv(%iv - %lb, %step)` expression.
 ///   2. checking whether any of the results of this affine map divides all the
 ///   other results (in which case it is also guaranteed to be the min).
 ///   3. replacing the AffineMinOp by the result of (2).
+///
+/// This composes chains of AffineMinOp recursively which is necessary because,
+/// in a multi-loop context, affine.min fed by outerloops may not canonicalize
+/// transitively.
+///
+/// Example:
+/// ```
+///    scf.parallel (%i) = (%c0) to (%c192) step (%c128) {
+///      %0 = affine.min affine_map<(d0) -> (128, 192 - d0)> (%i)
+///      scf.for %j = %c0 to %0 step %c16 {
+///        %1 = affine.min affine_map<(d0, d1) -> (16, d0 - d1)> (%0, %j)
+/// ```
+/// %1 can canonicalize despite the fact that %0 cannot:
+/// ```
+///    scf.parallel (%i) = (%c0) to (%c192) step (%c128) {
+///      %0 = affine.min affine_map<(d0) -> (128, 192 - d0)> (%i)
+///      scf.for %j = %c0 to %0 step %c16 {
+///        ... %c16
+/// ```
 // TODO: move to a more appropriate place when it is determined. For now Linalg
 // depends both on Affine and SCF but they do not depend on each other.
 struct AffineMinSCFCanonicalizationPattern
     : public OpRewritePattern<AffineMinOp> {
   using OpRewritePattern<AffineMinOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(AffineMinOp minOp,
                                 PatternRewriter &rewriter) const override;

mlir/include/mlir/IR/AffineExpr.h

public:
   /// modified expression tree.
   AffineExpr replace(AffineExpr expr, AffineExpr replacement) const;
 
   /// Sparse replace method. If `*this` appears in `map` replaces it by
   /// `map[*this]` and return the modified expression tree. Otherwise traverse
   /// `*this` and apply replace with `map` on its subexpressions.
   AffineExpr replace(const DenseMap<AffineExpr, AffineExpr> &map) const;
 
+  /// Replace dims[0 .. numDims - 1] by dims[shift .. shift + numDims - 1].
+  AffineExpr replaceDimsByShifting(unsigned numDims, unsigned shift) const;
+
+  /// Replace symbols[0 .. numDims - 1] by
+  /// symbols[shift .. shift + numDims - 1].
+  AffineExpr replaceSymbolsByShifting(unsigned numSymbols,
+                                      unsigned shift) const;
+
   AffineExpr operator+(int64_t v) const;
   AffineExpr operator+(AffineExpr other) const;
   AffineExpr operator-() const;
   AffineExpr operator-(int64_t v) const;
   AffineExpr operator-(AffineExpr other) const;
   AffineExpr operator*(int64_t v) const;
   AffineExpr operator*(AffineExpr other) const;
   AffineExpr floorDiv(uint64_t v) const;

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

static void substituteLoop(unsigned dimIdx, Value lbVal, Value ubVal,
   // 3. Rewrite `exprs` in place by replacing `dim[dimIdx]` by `lb + step * iv`.
   AffineExpr iv = getAffineDimExpr(dimIdx, ctx);
   for (auto &e : exprs)
     e = e.replace(iv, lb + step * (iv - lb).floorDiv(step));
 }
 
 /// Traverse the `dims` and substitute linear expressions in place of induction
 /// variables in `exprs`.
-static void substitute(SmallVectorImpl<AffineExpr> &exprs,
+static LogicalResult substitute(SmallVectorImpl<AffineExpr> &exprs,
                                 SmallVectorImpl<Value> &dims,
                                 SmallVectorImpl<Value> &symbols) {
+  constexpr unsigned kLimit = 1024;
+
   assert(!exprs.empty() && "Unexpected empty exprs");
   LLVM_DEBUG(llvm::interleaveComma(dims, DBGS() << "Start subst with dims: "));
   LLVM_DEBUG(llvm::dbgs() << "\n");
 
+  MLIRContext *ctx = exprs[0].getContext();
+
   // Note: `dims` and `symbols` grow as we iterate, upper bound is dynamic.
   for (unsigned dimIdx = 0; dimIdx < dims.size(); ++dimIdx) {
     Value dim = dims[dimIdx];
     LLVM_DEBUG(DBGS() << "Subst: " << dim << "\n");
 
     // Replace dim @ pos[dimIdx] by `%lb + %step * new_dim`
     // Where new dim / symbols are added depending on whether the values are
     // static or not.
     if (auto forOp = scf::getForInductionVarOwner(dim)) {
       substituteLoop(dimIdx, forOp.lowerBound(), forOp.upperBound(),
                      forOp.step(), exprs, dims, symbols);
       continue;
     }
     if (auto parallelForOp = scf::getParallelForInductionVarOwner(dim)) {
       for (unsigned idx = 0, e = parallelForOp.getNumLoops(); idx < e; ++idx)
         substituteLoop(dimIdx, parallelForOp.lowerBound()[idx],
                        parallelForOp.upperBound()[idx],
                        parallelForOp.step()[idx], exprs, dims, symbols);
       continue;
     }
+
+    auto parentMinOp = dim.getDefiningOp<AffineMinOp>();
+    if (!parentMinOp)
+      continue;
+
+    // Handle AffineMinOp recursively.
+    auto parentMinOpExprs =
+        llvm::to_vector<4>(parentMinOp.getAffineMap().getResults());
+    SmallVector<Value, 4> parentMinOpDims(parentMinOp.getDimOperands()),
+        parentMinOpSymbols(parentMinOp.getSymbolOperands());
+
+    // Go recursive.
+    substitute(parentMinOpExprs, parentMinOpDims, parentMinOpSymbols);
+
+    // In-place replacement of dims and symbols of exprs by shifting.
+    for (auto &e : parentMinOpExprs)
+      e = e.replaceDimsByShifting(parentMinOpDims.size(), dims.size())
+              .replaceSymbolsByShifting(parentMinOpSymbols.size(),
+                                        symbols.size());
+
+    SmallVector<AffineExpr, 4> extraExprs;
+    extraExprs.reserve(exprs.size());
+    AffineExpr dimToReplace = getAffineDimExpr(dimIdx, ctx);
+
+    // Potential N^2 blowup in pathological cases: bail after kLimit exprs.
+    for (auto &expr : exprs) {
+      AffineExpr lastSubstitutedExpr;
+      for (auto parentExpr : parentMinOpExprs) {
+        auto substitutedExpr = expr.replace(dimToReplace, parentExpr);
+        if (substitutedExpr == expr)
+          continue;
+        lastSubstitutedExpr = substitutedExpr;
+        extraExprs.push_back(lastSubstitutedExpr);
+        if (exprs.size() + extraExprs.size() >= kLimit) {
+          LLVM_DEBUG(DBGS() << "Bailing on expression size limit: " << kLimit);
+          return failure();
+        }
+      }
+      // In-place update `expr`.
+      if (lastSubstitutedExpr)
+        expr = extraExprs.pop_back_val();
+    }
+    exprs.append(extraExprs.begin(), extraExprs.end());
+    dims.append(parentMinOpDims.begin(), parentMinOpDims.end());
+    symbols.append(parentMinOpSymbols.begin(), parentMinOpSymbols.end());
+    continue;
   }
 
   // Cleanup and simplify the results.
   SmallVector<Value, 4> operands(dims.begin(), dims.end());
   operands.append(symbols.begin(), symbols.end());
   auto map = AffineMap::get(dims.size(), symbols.size(), exprs,
                             exprs.front().getContext());
   // Pull in affine.apply operations and compose them fully into the result.
   fullyComposeAffineMapAndOperands(&map, &operands);
   canonicalizeMapAndOperands(&map, &operands);
   map = simplifyAffineMap(map);
   // Assign the results.
   exprs.assign(map.getResults().begin(), map.getResults().end());
   dims.assign(operands.begin(), operands.begin() + map.getNumDims());
   symbols.assign(operands.begin() + map.getNumDims(), operands.end());
+
+  return success();
 }
 
 LogicalResult AffineMinSCFCanonicalizationPattern::matchAndRewrite(
     AffineMinOp minOp, PatternRewriter &rewriter) const {
   // At least one loop is needed to canonicalize affine.min + SCF.
   auto isLoopLike = [](Value v) {
     return scf::getParallelForInductionVarOwner(v) ||
            scf::getForInductionVarOwner(v);
   };
   if (llvm::none_of(minOp.getDimOperands(), isLoopLike))
     return failure();
 
   LLVM_DEBUG(DBGS() << "Canonicalize AffineMinSCF: " << *minOp.getOperation()
                     << "\n");
 
   auto exprs = llvm::to_vector<4>(minOp.getAffineMap().getResults());
   SmallVector<Value, 4> dims(minOp.getDimOperands()),
       symbols(minOp.getSymbolOperands());
-  substitute(exprs, dims, symbols);
+  if (failed(substitute(exprs, dims, symbols)))
+    return failure();
 
   MLIRContext *ctx = minOp.getContext();
   auto map = AffineMap::get(dims.size(), symbols.size(), exprs, ctx);
   LLVM_DEBUG(DBGS() << "Resulting map: " << map << "\n");
+  LLVM_DEBUG(llvm::interleaveComma(dims, DBGS() << "Resulting dims: "));
+  LLVM_DEBUG(llvm::interleaveComma(symbols, DBGS() << "Resulting symbols: "));
 
   // Check whether any of the expressions divides all expressions. In which case
   // it is guaranteed to be the min.
   for (auto e : map.getResults()) {
     LLVM_DEBUG(DBGS() << "Candidate mod: " << e << "\n");
     if (!e.isSymbolicOrConstant())
       continue;
 

mlir/lib/IR/AffineExpr.cpp

 
 /// Sparse replace method. Return the modified expression tree.
 AffineExpr AffineExpr::replace(AffineExpr expr, AffineExpr replacement) const {
   DenseMap<AffineExpr, AffineExpr> map;
   map.insert(std::make_pair(expr, replacement));
   return replace(map);
 }
 
+/// Replace dims[0 .. numDims - 1] by dims[shift .. shift + numDims - 1].
+AffineExpr AffineExpr::replaceDimsByShifting(unsigned numDims,
+                                             unsigned shift) const {
+  DenseMap<AffineExpr, AffineExpr> map;
+  for (unsigned idx = 0; idx < numDims; ++idx)
+    map.insert({getAffineDimExpr(idx, getContext()),
+                getAffineDimExpr(idx + shift, getContext())});
+  return replace(map);
+}
+
+/// Replace symbols[0 .. numDims - 1] by symbols[shift .. shift + numDims - 1].
+AffineExpr AffineExpr::replaceSymbolsByShifting(unsigned numSymbols,
+                                                unsigned shift) const {
+  DenseMap<AffineExpr, AffineExpr> map;
+  for (unsigned idx = 0; idx < numSymbols; ++idx)
+    map.insert({getAffineDimExpr(idx, getContext()),
+                getAffineDimExpr(idx + shift, getContext())});
+  return replace(map);
+}
+
 /// Returns true if this expression is made out of only symbols and
 /// constants (no dimensional identifiers).
 bool AffineExpr::isSymbolicOrConstant() const {
   switch (getKind()) {
   case AffineExprKind::Constant:
     return true;
   case AffineExprKind::DimId:
     return false;

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
     store %2, %A[]: memref<i64>
   }
 
   return
 }
+
+#map1 = affine_map<(d0) -> (128, -d0 + 192)>
+#map2 = affine_map<(d0, d1) -> (16, d0 - d1)>
+
+func @mixed(%A: memref<i64>) {
+  %c128 = constant 128 : index
+  %c192 = constant 192 : index
+  %c32 = constant 32 : index
+  %c0 = constant 0 : index
+  %c16 = constant 16 : index
+
+  //     CHECK: scf.parallel
+  //     CHECK:   scf.for
+  // CHECK-NEXT:   %[[C16:.*]] = constant 16 : index
+  // CHECK-NEXT:   %[[C16I64:.*]] = index_cast %[[C16:.*]]
+  // CHECK-NEXT:   store %[[C16I64]], %{{.*}}[] : memref<i64>
+  scf.parallel (%i) = (%c0) to (%c192) step (%c128) {
+    %1 = affine.min #map1(%i)
+    scf.for %j = %c0 to %1 step %c16 {
+      %2 = affine.min #map2(%1, %j)
+      %3 = index_cast %2: index to i64
+      store %3, %A[]: memref<i64>
+    }
+  }
+
+  // This should not canonicalize because:
+  //   `129 * floordiv(%iv - 16, 129) % 16 != 0`.
+  //     CHECK: scf.parallel
+  //     CHECK:   scf.for
+  //     CHECK:     index_cast
+  %c129 = constant 129 : index
+  scf.parallel (%i) = (%c0) to (%c192) step (%c129) {
+    %1 = affine.min #map1(%i)
+    scf.for %j = %c0 to %1 step %c16 {
+      %2 = affine.min #map2(%1, %j)
+      %3 = index_cast %2: index to i64
+      store %3, %A[]: memref<i64>
+    }
+  }
+  return
+}