[mlir][linalg] Replace AffineMinSCFCanonicalizationPattern with SCF reimplementation

Use the new canonicalization pattern in the SCF dialect.

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

Author: matthias-springer

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

@@ -982,69 +982,6 @@ struct LinalgCopyVTWForwardingPattern
                                 PatternRewriter &rewriter) const override;
 };
 
-using GetMinMaxExprFn =
-    std::function<Optional<std::pair<AffineExpr, AffineExpr>>(
-        Value value, SmallVectorImpl<Value> &dims,
-        SmallVectorImpl<Value> &symbols)>;
-
-/// Canonicalize AffineMinOp operations in the context of ops with a known range
-/// by:
-///   1. building an affine map where uses of the known ops are replaced by
-///   their min annd max expressions returned by the lambda `getMinMaxFn`.
-///   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).
-struct AffineMinRangeCanonicalizationPattern
-    : public OpRewritePattern<AffineMinOp> {
-  AffineMinRangeCanonicalizationPattern(MLIRContext *context,
-                                        GetMinMaxExprFn getMinMaxFn)
-      : OpRewritePattern<AffineMinOp>(context), getMinMaxFn(getMinMaxFn) {}
-  LogicalResult matchAndRewrite(AffineMinOp minOp,
-                                PatternRewriter &rewriter) const override;
-
-protected:
-  GetMinMaxExprFn getMinMaxFn;
-};
-
-/// Specialized version of `AffineMinRangeCanonicalizationPattern` pattern
-/// using `getSCFMinMaxExpr` to know the min and max expression of induction
-/// variables from scf loops.
-// 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 AffineMinRangeCanonicalizationPattern {
-  static Optional<std::pair<AffineExpr, AffineExpr>>
-  getMinMax(Value value, SmallVectorImpl<Value> &dims,
-            SmallVectorImpl<Value> &symbols) {
-    return getSCFMinMaxExpr(value, dims, symbols);
-  }
-  AffineMinSCFCanonicalizationPattern(MLIRContext *context)
-      : AffineMinRangeCanonicalizationPattern(context, getMinMax) {}
-};
-
-/// Helper struct to return the results of `substituteMin`.
-struct AffineMapAndOperands {
-  AffineMap map;
-  SmallVector<Value> dims;
-  SmallVector<Value> symbols;
-};
-
-/// Traverse the dims of the AffineMap of `affineMinOp` and substitute
-/// dimensions with known range by new expressions involving the min or max
-/// expression:
-///   - If the AffineDimExpr mapped to a known value has a positive sign, it
-///     is replaced by the min expression.
-///   - If the AffineDimExpr mapped to a known value has a negative sign, it is
-///     replaced by the max expression.
-/// All known values are iteratively replaced.
-/// This is used as an intermediate step in computing bounding boxes and
-/// canonicalize AffineMinOps. All dim and symbol operands are assumed to have
-/// positive values (positive orthant assumptions).
-/// Return a new AffineMap, dims and symbols that have been canonicalized and
-/// simplified.
-AffineMapAndOperands substituteMin(AffineMinOp affineMinOp,
-                                   GetMinMaxExprFn getMinMaxExpr);
-
 /// Converts Convolution op into vector contraction.
 ///
 /// Conversion expects ConvOp to have dimensions marked in the *mask* as

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

@@ -14,6 +14,7 @@
 #include "mlir/Dialect/Linalg/Transforms/CodegenStrategy.h"
 
 #include "mlir/Dialect/Linalg/Transforms/Hoisting.h"
+#include "mlir/Dialect/SCF/Transforms.h"
 #include "mlir/Dialect/Vector/VectorOps.h"
 #include "mlir/Dialect/Vector/VectorTransforms.h"
 #include "mlir/Pass/PassManager.h"
@@ -47,7 +48,7 @@ void mlir::linalg::CodegenStrategy::transform(FuncOp func) const {
 
   RewritePatternSet stage2Patterns =
       linalg::getLinalgTilingCanonicalizationPatterns(context);
-  stage2Patterns.add<AffineMinSCFCanonicalizationPattern>(context);
+  scf::populateSCFLoopBodyCanonicalizationPatterns(stage2Patterns);
 
   auto stage3Transforms = [&](Operation *op) {
     // Some of these may be too aggressive as a stage 3 that is applied on each

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

@@ -16,6 +16,7 @@
 #include "mlir/Dialect/Linalg/Transforms/Transforms.h"
 #include "mlir/Dialect/Linalg/Utils/Utils.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Dialect/SCF/Transforms.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/IR/AffineExpr.h"
 #include "mlir/IR/AffineMap.h"
@@ -536,7 +537,7 @@ applyTilingToLoopPatterns(LinalgTilingLoopType loopType, FuncOp funcOp,
   MLIRContext *ctx = funcOp.getContext();
   RewritePatternSet patterns(ctx);
   insertTilingPatterns(patterns, options);
-  patterns.add<AffineMinSCFCanonicalizationPattern>(patterns.getContext());
+  scf::populateSCFLoopBodyCanonicalizationPatterns(patterns);
   (void)applyPatternsAndFoldGreedily(funcOp, std::move(patterns));
   (void)applyPatternsAndFoldGreedily(
       funcOp, getLinalgTilingCanonicalizationPatterns(ctx));

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

@@ -494,145 +494,6 @@ LogicalResult mlir::linalg::applyStagedPatterns(
   return success();
 }
 
-/// Traverse the `dims` and substitute known min or max expressions returned by
-/// the lambda |getMinMaxExpr|.
-static AffineMap substitute(AffineMap map, SmallVectorImpl<Value> &dims,
-                            SmallVectorImpl<Value> &symbols,
-                            GetMinMaxExprFn getMinMaxExpr) {
-  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];
-        auto minMax = getMinMaxExpr(dim, dims, symbols);
-        if (!minMax)
-          continue;
-        AffineExpr dimExpr = getAffineDimExpr(dimIdx, expr.getContext());
-        LLVM_DEBUG(DBGS() << "Subst: " << dim << " @ " << dimExpr << "\n");
-        LLVM_DEBUG(DBGS() << "Before: " << expr << "\n");
-        // Substitute occurrences of `dimExpr` by either the min expression or
-        // the max expression depending on whether the value is used with a
-        // positive or negative  coefficient.
-        AffineExpr substitutedExpr =
-            substWithMin(expr, dimExpr, minMax->first, minMax->second);
-        LLVM_DEBUG(DBGS() << "After: " << substitutedExpr << "\n");
-        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";
-      DBGS() << "Operands:\n";
-      for (Value v : operands)
-        DBGS() << v << "\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());
-}
-
-/// Traverse the dims of the AffineMap of `affineMinOp` and substitute
-/// dimensions with known range by new expressions involving the min or max
-/// expression:
-///   - If the AffineDimExpr mapped to a known value has a positive sign, it
-///     is replaced by the min expression.
-///   - If the AffineDimExpr mapped to a known value has a negative sign, it is
-///     replaced by the max expression.
-/// All known values are iteratively replaced.
-/// This is used as an intermediate step in computing bounding boxes and
-/// canonicalize AffineMinOps. All dim and symbol operands are assumed to have
-/// positive values (positive orthant assumptions).
-/// Return a new AffineMap, dims and symbols that have been canonicalized and
-/// simplified.
-AffineMapAndOperands
-mlir::linalg::substituteMin(AffineMinOp affineMinOp,
-                            GetMinMaxExprFn getMinMaxExpr) {
-  AffineMapAndOperands res{affineMinOp.getAffineMap(),
-                           SmallVector<Value>(affineMinOp.getDimOperands()),
-                           SmallVector<Value>(affineMinOp.getSymbolOperands())};
-  res.map = substitute(affineMinOp.getAffineMap(), res.dims, res.symbols,
-                       getMinMaxExpr);
-  return res;
-}
-
-LogicalResult AffineMinRangeCanonicalizationPattern::matchAndRewrite(
-    AffineMinOp minOp, PatternRewriter &rewriter) const {
-  LLVM_DEBUG(DBGS() << "Canonicalize AffineMinSCF: " << *minOp.getOperation()
-                    << "\n");
-
-  auto affineMapAndOperands = substituteMin(minOp, getMinMaxFn);
-  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()) {
-    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> resultOperands = affineMapAndOperands.dims;
-      llvm::append_range(resultOperands, affineMapAndOperands.symbols);
-      canonicalizeMapAndOperands(&resultMap, &resultOperands);
-      resultMap = simplifyAffineMap(resultMap);
-      rewriter.replaceOpWithNewOp<AffineApplyOp>(minOp, resultMap,
-                                                 resultOperands);
-    }
-    return success();
-  }
-
-  return failure();
-}
-
 static SmallVector<StringRef> getNParallelLoopsAttrs(unsigned nParallelLoops) {
   return SmallVector<StringRef>(nParallelLoops, getParallelIteratorTypeName());
 }

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

@@ -145,43 +145,3 @@ bool mlir::getInnermostParallelLoops(Operation *rootOp,
   }
   return rootEnclosesPloops;
 }
-
-/// 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. Then create a min
-/// expression (i.e. `%lb`) and a max expression
-/// (i.e. `%lb + %step * floordiv(%ub -1 - %lb, %step)`.
-static std::pair<AffineExpr, AffineExpr>
-getMinMaxLoopIndVar(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);
-  return std::make_pair(lb, lb + step * ((ub - 1) - lb).floorDiv(step));
-}
-
-/// Return the min/max expressions for `value` if it is an induction variable
-/// from scf.for or scf.parallel loop.
-/// if `loopFilter` is passed, the filter determines which loop to consider.
-/// Other induction variables are ignored.
-Optional<std::pair<AffineExpr, AffineExpr>> mlir::getSCFMinMaxExpr(
-    Value value, SmallVectorImpl<Value> &dims, SmallVectorImpl<Value> &symbols,
-    llvm::function_ref<bool(Operation *)> substituteOperation) {
-  if (auto forOp = scf::getForInductionVarOwner(value))
-    return getMinMaxLoopIndVar(forOp.lowerBound(), forOp.upperBound(),
-                               forOp.step(), dims, symbols);
-
-  if (auto parallelForOp = scf::getParallelForInductionVarOwner(value))
-    for (unsigned idx = 0, e = parallelForOp.getNumLoops(); idx < e; ++idx)
-      if (parallelForOp.getInductionVars()[idx] == value)
-        return getMinMaxLoopIndVar(parallelForOp.lowerBound()[idx],
-                                   parallelForOp.upperBound()[idx],
-                                   parallelForOp.step()[idx], dims, symbols);
-  return {};
-}