[mlir][NFC] Move helper substWithMin into Affine utils

[ThomasRaoux]

Summary

This allow using this helper outside of the linalg canonicalization.

mlir/include/mlir/Dialect/Affine/Utils.h

 //===- Utils.h - Affine dialect utilities -----------------------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This header file declares a set of utilities for the affine dialect ops.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_AFFINE_UTILS_H
 #define MLIR_DIALECT_AFFINE_UTILS_H
 
+#include "mlir/IR/AffineExpr.h"
 #include "mlir/Support/LLVM.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/SmallVector.h"
 
 namespace mlir {
 
 class AffineForOp;
 class AffineIfOp;

 vectorizeAffineLoopNest(std::vector<SmallVector<AffineForOp, 2>> &loops,
                         const VectorizationStrategy &strategy);
 
 /// Normalize a affine.parallel op so that lower bounds are 0 and steps are 1.
 /// As currently implemented, this transformation cannot fail and will return
 /// early if the op is already in a normalized form.
 void normalizeAffineParallel(AffineParallelOp op);
 
+/// 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.
+AffineExpr substWithMin(AffineExpr e, AffineExpr dim, AffineExpr min,
+                        AffineExpr max, bool positivePath = true);
+
 } // namespace mlir
 
 #endif // MLIR_DIALECT_AFFINE_UTILS_H

mlir/lib/Dialect/Affine/Utils/Utils.cpp

LogicalResult mlir::hoistAffineIfOp(AffineIfOp ifOp, bool *folded) {
   // Canonicalize to remove dead else blocks (happens whenever an 'if' moves up
   // a sequence of affine.fors that are all perfectly nested).
   applyPatternsAndFoldGreedily(
       hoistedIfOp->getParentWithTrait<OpTrait::IsIsolatedFromAbove>(),
       frozenPatterns);
 
   return success();
 }
+
+// Return the min expr after replacing the given dim.
+AffineExpr mlir::substWithMin(AffineExpr e, AffineExpr dim, AffineExpr min,
+                              AffineExpr max, bool positivePath) {
+  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;
+}

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

 //===- LinalgTransforms.cpp - Linalg transformations as patterns ----------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements logic and helpers to expose Linalg transforms as rewrite
 // patterns.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Linalg/Transforms/Transforms.h"
+#include "mlir/Dialect/Affine/Utils.h"
 #include "mlir/Dialect/Linalg/Analysis/DependenceAnalysis.h"
 #include "mlir/Dialect/Linalg/IR/LinalgOps.h"
 #include "mlir/Dialect/Linalg/Utils/Utils.h"
 #include "mlir/Dialect/StandardOps/EDSC/Intrinsics.h"
 #include "mlir/Dialect/Utils/StructuredOpsUtils.h"
 #include "mlir/Dialect/Vector/EDSC/Intrinsics.h"
 #include "mlir/Dialect/Vector/VectorOps.h"
 #include "mlir/IR/AffineExpr.h"
if (stage3Lambda) {
         return failure();
       LLVM_DEBUG(DBGS() << "After 3rd stage, iter : " << iteration << "\n"
                         << *op);
     }
   }
   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,