[mlir] Refactor AffineMinSCFCanonicalizationPattern to make it more generic

[ThomasRaoux]

Summary

Replace the pattern with a more generic one allowing to inject replacement code for any known ops.

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

 struct LinalgCopyVTWForwardingPattern
     : public OpRewritePattern<vector::TransferWriteOp> {
   using OpRewritePattern<vector::TransferWriteOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(vector::TransferWriteOp xferOp,
                                 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.
+/// 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 value based on the lambda `substituteFn`.
 ///   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
+struct AffineMinRangeCanonicalizationPattern
     : public OpRewritePattern<AffineMinOp> {
-  using OpRewritePattern<AffineMinOp>::OpRewritePattern;
-
+  using SubstituteFn = std::function<AffineExpr(Value, AffineExpr, AffineExpr,
+                                                SmallVectorImpl<Value> &,
+                                                SmallVectorImpl<Value> &)>;
+  AffineMinRangeCanonicalizationPattern(MLIRContext *context,
+                                        SubstituteFn substituteFn)
+      : OpRewritePattern<AffineMinOp>(context), substituteFn(substituteFn) {}
   LogicalResult matchAndRewrite(AffineMinOp minOp,
                                 PatternRewriter &rewriter) const override;
+
+protected:
+  // Lambda to substitute a dimension based on the semantic of its definition.
+  SubstituteFn substituteFn;
+};
+
+/// Replace the induction variable of a scf loop by the min (i.e. `%lb`) or 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.
+// 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.
+AffineExpr substituteSCFOps(Value dim, AffineExpr expr, AffineExpr dimExpr,
+                            SmallVectorImpl<Value> &dims,
+                            SmallVectorImpl<Value> &symbols);
+
+/// Specialized version of `AffineMinRangeCanonicalizationPattern` pattern
+/// handling scf ops.
+struct AffineMinSCFCanonicalizationPattern
+    : public AffineMinRangeCanonicalizationPattern {
+  AffineMinSCFCanonicalizationPattern(MLIRContext *context)
+      : AffineMinRangeCanonicalizationPattern(context, substituteSCFOps) {}
 };
 
 /// Converts Convolution op into vector contraction.
 ///
 /// Conversion expects ConvOp to have dimensions marked in the *mask* as
 /// false of size 1. This ensures that the ConvOp can be lowered to vector
 /// contraction of dimensions marked in the *mask* as true.
 ///

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

static AffineExpr substituteLoopInExpr(AffineExpr expr, AffineExpr dimExpr,
   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,
+/// substitute known min or max expressions in place of induction variables in
+/// `exprs`.
+AffineExpr mlir::linalg::substituteSCFOps(Value dim, AffineExpr expr,
+                                          AffineExpr dimExpr,
+                                          SmallVectorImpl<Value> &dims,
                                           SmallVectorImpl<Value> &symbols) {
+  AffineExpr substitutedExpr;
+  if (auto forOp = scf::getForInductionVarOwner(dim)) {
+    substitutedExpr =
+        substituteLoopInExpr(expr, dimExpr, forOp.lowerBound(),
+                             forOp.upperBound(), forOp.step(), dims, symbols);
+  } else 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);
+  }
+  return substitutedExpr;
+}
+
+/// Traverse the `dims` and substitute known min or max expressions in `exprs`.
+static AffineMap
+substitute(AffineMap map, SmallVectorImpl<Value> &dims,
+           SmallVectorImpl<Value> &symbols,
+           AffineMinRangeCanonicalizationPattern::SubstituteFn substituteFn) {
   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);
+        AffineExpr substitutedExpr =
+            substituteFn(dim, expr, dimExpr, dims, symbols);
 
         if (!substitutedExpr)
           continue;
 
         substituted = (substitutedExpr != expr);
         expr = substitutedExpr;
       }
     }
for (AffineExpr &expr : exprs) {
 
     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(
+LogicalResult AffineMinRangeCanonicalizationPattern::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);
+  AffineMap map = substitute(minOp.getAffineMap(), dims, symbols, substituteFn);
 
   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())