[mlir] Introduce AffineMinSCF folding as a pattern

This revision adds a folding pattern to replace affine.min ops by the actual min value, when it can be determined statically from the strides and bounds of enclosing scf loop .

This matches the type of expressions that Linalg produces during tiling and simplifies boundary checks. For now Linalg depends both on Affine and SCF but they do not depend on each other, so the pattern is added there.
In the future this will move to a more appropriate place when it is determined.

The canonicalization of AffineMinOp operations in the context of enclosing scf.for and scf.parallel proceeds by:
  1. building an affine map where uses of the induction variable of a loop
  are replaced by `%lb + %step * floordiv(%iv - %lb, %step)` expressions.
  2. checking if 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).

The algorithm is functional in simple parametric tiling cases by using semi-affine maps. However simplifications of such semi-affine maps are not yet available and the canonicalization does not succeed yet.

Differential Revision: https://reviews.llvm.org/D82009

Author: Nicolas Vasilache

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

@@ -502,6 +502,26 @@ struct LinalgCopyVTWForwardingPattern
                                 PatternRewriter &rewriter) const override;
 };
 
+/// 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 either the min (i.e. `%lb`) of the max
+///   (i.e. `%lb + %step * floordiv(%ub -1 - %lb, %step)`) expression, depending
+///   on whether the induction variable is used with a positive or negative
+///   coefficient.
+///   2. checking whether any of the results of this affine map is known to be
+///   greater than all other results.
+///   3. replacing the AffineMinOp by the result of (2).
+// 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;
+};
+
 //===----------------------------------------------------------------------===//
 // Support for staged pattern application.
 //===----------------------------------------------------------------------===//
@@ -519,6 +539,7 @@ LogicalResult applyStagedPatterns(
     Operation *op, ArrayRef<OwningRewritePatternList> stage1Patterns,
     const OwningRewritePatternList &stage2Patterns,
     function_ref<LogicalResult(Operation *)> stage3Lambda = nullptr);
+
 } // namespace linalg
 } // namespace mlir
 

mlir/include/mlir/IR/AffineExpr.h

@@ -115,9 +115,20 @@ class AffineExpr {
 
   /// This method substitutes any uses of dimensions and symbols (e.g.
   /// dim#0 with dimReplacements[0]) and returns the modified expression tree.
+  /// This is a dense replacement method: a replacement must be specified for
+  /// every single dim and symbol.
   AffineExpr replaceDimsAndSymbols(ArrayRef<AffineExpr> dimReplacements,
                                    ArrayRef<AffineExpr> symReplacements) const;
 
+  /// Sparse replace method. Replace `expr` by `replacement` and return the
+  /// 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 symbols[0 .. numDims - 1] by
   /// symbols[shift .. shift + numDims - 1].
   AffineExpr shiftSymbols(unsigned numSymbols, unsigned shift) const;

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

@@ -36,6 +36,7 @@ using namespace mlir::edsc::intrinsics;
 using namespace mlir::linalg;
 
 #define DBGS() (llvm::dbgs() << "[" DEBUG_TYPE << "]: ")
+
 //===----------------------------------------------------------------------===//
 // Transformations exposed as rewrite patterns.
 //===----------------------------------------------------------------------===//
@@ -235,3 +236,177 @@ LogicalResult mlir::linalg::applyStagedPatterns(
   }
   return success();
 }
+
+/// Traverse `e` and return an AffineExpr where all occurrences of `dim` have
+/// been replaced by either:
+///  - `min` if `positivePath` is true when we reach an occurrence of `dim`
+///  - `max` if `positivePath` is true when we reach an occurrence of `dim`
+/// `positivePath` is negated each time we hit a multiplicative or divisive
+/// binary op with a constant negative coefficient.
+static AffineExpr substWithMin(AffineExpr e, AffineExpr dim, AffineExpr min,
+                               AffineExpr max, bool positivePath = true) {
+  if (e == dim)
+    return positivePath ? min : max;
+  if (auto bin = e.dyn_cast<AffineBinaryOpExpr>()) {
+    AffineExpr lhs = bin.getLHS();
+    AffineExpr rhs = bin.getRHS();
+    if (bin.getKind() == mlir::AffineExprKind::Add)
+      return substWithMin(lhs, dim, min, max, positivePath) +
+             substWithMin(rhs, dim, min, max, positivePath);
+
+    auto c1 = bin.getLHS().dyn_cast<AffineConstantExpr>();
+    auto c2 = bin.getRHS().dyn_cast<AffineConstantExpr>();
+    if (c1 && c1.getValue() < 0)
+      return getAffineBinaryOpExpr(
+          bin.getKind(), c1, substWithMin(rhs, dim, min, max, !positivePath));
+    if (c2 && c2.getValue() < 0)
+      return getAffineBinaryOpExpr(
+          bin.getKind(), substWithMin(lhs, dim, min, max, !positivePath), c2);
+    return getAffineBinaryOpExpr(
+        bin.getKind(), substWithMin(lhs, dim, min, max, positivePath),
+        substWithMin(rhs, dim, min, max, positivePath));
+  }
+  return e;
+}
+
+/// 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,
+                                       Value lbVal, Value ubVal, Value stepVal,
+                                       SmallVectorImpl<Value> &dims,
+                                       SmallVectorImpl<Value> &symbols) {
+  MLIRContext *ctx = lbVal.getContext();
+  AffineExpr lb = getAffineDimExpr(dims.size(), ctx);
+  dims.push_back(lbVal);
+  AffineExpr ub = getAffineDimExpr(dims.size(), ctx);
+  dims.push_back(ubVal);
+  AffineExpr step = getAffineSymbolExpr(symbols.size(), ctx);
+  symbols.push_back(stepVal);
+  LLVM_DEBUG(DBGS() << "Before: " << expr << "\n");
+  AffineExpr ee = substWithMin(expr, dimExpr, lb,
+                               lb + step * ((ub - 1) - lb).floorDiv(step));
+  LLVM_DEBUG(DBGS() << "After: " << expr << "\n");
+  return ee;
+}
+
+/// Traverse the `dims` and substitute known min or max expressions in place of
+/// induction variables in `exprs`.
+static AffineMap substitute(AffineMap map, SmallVectorImpl<Value> &dims,
+                            SmallVectorImpl<Value> &symbols) {
+  auto exprs = llvm::to_vector<4>(map.getResults());
+  for (AffineExpr &expr : exprs) {
+    bool substituted = true;
+    while (substituted) {
+      substituted = false;
+      for (unsigned dimIdx = 0; dimIdx < dims.size(); ++dimIdx) {
+        Value dim = dims[dimIdx];
+        AffineExpr dimExpr = getAffineDimExpr(dimIdx, expr.getContext());
+        LLVM_DEBUG(DBGS() << "Subst: " << dim << " @ " << dimExpr << "\n");
+        AffineExpr substitutedExpr;
+        if (auto forOp = scf::getForInductionVarOwner(dim))
+          substitutedExpr = substituteLoopInExpr(
+              expr, dimExpr, forOp.lowerBound(), forOp.upperBound(),
+              forOp.step(), dims, symbols);
+
+        if (auto parallelForOp = scf::getParallelForInductionVarOwner(dim))
+          for (unsigned idx = 0, e = parallelForOp.getNumLoops(); idx < e;
+               ++idx)
+            substitutedExpr = substituteLoopInExpr(
+                expr, dimExpr, parallelForOp.lowerBound()[idx],
+                parallelForOp.upperBound()[idx], parallelForOp.step()[idx],
+                dims, symbols);
+
+        if (!substitutedExpr)
+          continue;
+
+        substituted = (substitutedExpr != expr);
+        expr = substitutedExpr;
+      }
+    }
+
+    // Cleanup and simplify the results.
+    // This needs to happen outside of the loop iterating on dims.size() since
+    // it modifies dims.
+    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());
+
+    LLVM_DEBUG(DBGS() << "Map to simplify: " << map << "\n");
+
+    // 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());
+
+    LLVM_DEBUG(DBGS() << "Map simplified: " << map << "\n");
+  }
+
+  assert(!exprs.empty() && "Unexpected empty exprs");
+  return AffineMap::get(dims.size(), symbols.size(), exprs, map.getContext());
+}
+
+LogicalResult AffineMinSCFCanonicalizationPattern::matchAndRewrite(
+    AffineMinOp minOp, PatternRewriter &rewriter) const {
+  LLVM_DEBUG(DBGS() << "Canonicalize AffineMinSCF: " << *minOp.getOperation()
+                    << "\n");
+
+  SmallVector<Value, 4> dims(minOp.getDimOperands()),
+      symbols(minOp.getSymbolOperands());
+  AffineMap map = substitute(minOp.getAffineMap(), dims, symbols);
+
+  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()) {
+    LLVM_DEBUG(DBGS() << "Candidate min: " << e << "\n");
+    if (!e.isSymbolicOrConstant())
+      continue;
+
+    auto isNonPositive = [](AffineExpr e) {
+      if (auto cst = e.dyn_cast<AffineConstantExpr>())
+        return cst.getValue() < 0;
+      return true;
+    };
+
+    // Build the subMap and check everything is statically known to be
+    // positive.
+    SmallVector<AffineExpr, 4> subExprs;
+    subExprs.reserve(map.getNumResults());
+    for (auto ee : map.getResults())
+      subExprs.push_back(ee - e);
+    MLIRContext *ctx = minOp.getContext();
+    AffineMap subMap = simplifyAffineMap(
+        AffineMap::get(map.getNumDims(), map.getNumSymbols(), subExprs, ctx));
+    LLVM_DEBUG(DBGS() << "simplified subMap: " << subMap << "\n");
+    if (llvm::any_of(subMap.getResults(), isNonPositive))
+      continue;
+
+    // Static min found.
+    if (auto cst = e.dyn_cast<AffineConstantExpr>()) {
+      rewriter.replaceOpWithNewOp<ConstantIndexOp>(minOp, cst.getValue());
+    } else {
+      auto resultMap = AffineMap::get(0, map.getNumSymbols(), {e}, ctx);
+      SmallVector<Value, 4> resultOperands = dims;
+      resultOperands.append(symbols.begin(), symbols.end());
+      canonicalizeMapAndOperands(&resultMap, &resultOperands);
+      resultMap = simplifyAffineMap(resultMap);
+      rewriter.replaceOpWithNewOp<AffineApplyOp>(minOp, resultMap,
+                                                 resultOperands);
+    }
+    return success();
+  }
+
+  return failure();
+}

mlir/lib/IR/AffineExpr.cpp

@@ -101,6 +101,37 @@ AffineExpr AffineExpr::shiftSymbols(unsigned numSymbols, unsigned shift) const {
   return replaceDimsAndSymbols({}, symbols);
 }
 
+/// Sparse replace method. Return the modified expression tree.
+AffineExpr
+AffineExpr::replace(const DenseMap<AffineExpr, AffineExpr> &map) const {
+  auto it = map.find(*this);
+  if (it != map.end())
+    return it->second;
+  switch (getKind()) {
+  default:
+    return *this;
+  case AffineExprKind::Add:
+  case AffineExprKind::Mul:
+  case AffineExprKind::FloorDiv:
+  case AffineExprKind::CeilDiv:
+  case AffineExprKind::Mod:
+    auto binOp = cast<AffineBinaryOpExpr>();
+    auto lhs = binOp.getLHS(), rhs = binOp.getRHS();
+    auto newLHS = lhs.replace(map);
+    auto newRHS = rhs.replace(map);
+    if (newLHS == lhs && newRHS == rhs)
+      return *this;
+    return getAffineBinaryOpExpr(getKind(), newLHS, newRHS);
+  }
+  llvm_unreachable("Unknown AffineExpr");
+}
+
+/// 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);
+}
 /// Returns true if this expression is made out of only symbols and
 /// constants (no dimensional identifiers).
 bool AffineExpr::isSymbolicOrConstant() const {