[Mlir-commits] [mlir] [mlir][scf]: Expose emitNormalizedLoopBounds/denormalizeInductionVariable util functions (NFC) (PR #94429)

[Aviad Cohen]

https://github.com/AviadCo updated https://github.com/llvm/llvm-project/pull/94429

>From 5bc850b6e25f20aa489c3e758368d05b25ac4744 Mon Sep 17 00:00:00 2001
From: Aviad Cohen <aviad.cohen2 at mobileye.com>
Date: Wed, 5 Jun 2024 08:07:16 +0300
Subject: [PATCH] [mlir][scf]: Expose
 emitNormalizedLoopBounds/denormalizeInductionVariable util functions (NFC)

---
 mlir/include/mlir/Dialect/SCF/Utils/Utils.h | 25 ++++++++++++
 mlir/lib/Dialect/SCF/Utils/Utils.cpp        | 44 ++++++---------------
 2 files changed, 36 insertions(+), 33 deletions(-)

diff --git a/mlir/include/mlir/Dialect/SCF/Utils/Utils.h b/mlir/include/mlir/Dialect/SCF/Utils/Utils.h
index bc09cc7f7fa5e..8eb4bd6bceff1 100644
--- a/mlir/include/mlir/Dialect/SCF/Utils/Utils.h
+++ b/mlir/include/mlir/Dialect/SCF/Utils/Utils.h
@@ -120,6 +120,31 @@ LogicalResult loopUnrollByFactor(
     scf::ForOp forOp, uint64_t unrollFactor,
     function_ref<void(unsigned, Operation *, OpBuilder)> annotateFn = nullptr);
 
+/// This structure is to pass and return sets of loop parameters without
+/// confusing the order.
+struct LoopParams {
+  OpFoldResult lowerBound;
+  OpFoldResult upperBound;
+  OpFoldResult step;
+};
+
+/// Transform a loop with a strictly positive step
+///   for %i = %lb to %ub step %s
+/// into a 0-based loop with step 1
+///   for %ii = 0 to ceildiv(%ub - %lb, %s) step 1 {
+///     %i = %ii * %s + %lb
+/// Insert the induction variable remapping in the body of `inner`, which is
+/// expected to be either `loop` or another loop perfectly nested under `loop`.
+/// Insert the definition of new bounds immediate before `outer`, which is
+/// expected to be either `loop` or its parent in the loop nest.
+LoopParams emitNormalizedLoopBounds(RewriterBase &rewriter, Location loc,
+                                    Value lb, Value ub, Value step);
+
+/// Get back the original induction variable values after loop normalization.
+void denormalizeInductionVariable(RewriterBase &rewriter, Location loc,
+                                  Value normalizedIv, Value origLb,
+                                  Value origStep);
+
 /// Tile a nest of standard for loops rooted at `rootForOp` by finding such
 /// parametric tile sizes that the outer loops have a fixed number of iterations
 /// as defined in `sizes`.
diff --git a/mlir/lib/Dialect/SCF/Utils/Utils.cpp b/mlir/lib/Dialect/SCF/Utils/Utils.cpp
index 6658cca03eba7..94e8c066d5414 100644
--- a/mlir/lib/Dialect/SCF/Utils/Utils.cpp
+++ b/mlir/lib/Dialect/SCF/Utils/Utils.cpp
@@ -18,6 +18,7 @@
 #include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/IRMapping.h"
+#include "mlir/IR/OpDefinition.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 #include "mlir/Support/MathExtras.h"
@@ -29,16 +30,6 @@
 
 using namespace mlir;
 
-namespace {
-// This structure is to pass and return sets of loop parameters without
-// confusing the order.
-struct LoopParams {
-  Value lowerBound;
-  Value upperBound;
-  Value step;
-};
-} // namespace
-
 SmallVector<scf::ForOp> mlir::replaceLoopNestWithNewYields(
     RewriterBase &rewriter, MutableArrayRef<scf::ForOp> loopNest,
     ValueRange newIterOperands, const NewYieldValuesFn &newYieldValuesFn,
@@ -473,17 +464,8 @@ LogicalResult mlir::loopUnrollByFactor(
   return success();
 }
 
-/// Transform a loop with a strictly positive step
-///   for %i = %lb to %ub step %s
-/// into a 0-based loop with step 1
-///   for %ii = 0 to ceildiv(%ub - %lb, %s) step 1 {
-///     %i = %ii * %s + %lb
-/// Insert the induction variable remapping in the body of `inner`, which is
-/// expected to be either `loop` or another loop perfectly nested under `loop`.
-/// Insert the definition of new bounds immediate before `outer`, which is
-/// expected to be either `loop` or its parent in the loop nest.
-static LoopParams emitNormalizedLoopBounds(RewriterBase &rewriter, Location loc,
-                                           Value lb, Value ub, Value step) {
+LoopParams mlir::emitNormalizedLoopBounds(RewriterBase &rewriter, Location loc,
+                                          Value lb, Value ub, Value step) {
   // For non-index types, generate `arith` instructions
   // Check if the loop is already known to have a constant zero lower bound or
   // a constant one step.
@@ -517,10 +499,9 @@ static LoopParams emitNormalizedLoopBounds(RewriterBase &rewriter, Location loc,
   return {newLowerBound, newUpperBound, newStep};
 }
 
-/// Get back the original induction variable values after loop normalization
-static void denormalizeInductionVariable(RewriterBase &rewriter, Location loc,
-                                         Value normalizedIv, Value origLb,
-                                         Value origStep) {
+void mlir::denormalizeInductionVariable(RewriterBase &rewriter, Location loc,
+                                        Value normalizedIv, Value origLb,
+                                        Value origStep) {
   Value denormalizedIv;
   SmallPtrSet<Operation *, 2> preserve;
   bool isStepOne = isConstantIntValue(origStep, 1);
@@ -638,9 +619,9 @@ LogicalResult mlir::coalesceLoops(RewriterBase &rewriter,
         emitNormalizedLoopBounds(rewriter, loop.getLoc(), lb, ub, step);
 
     rewriter.modifyOpInPlace(loop, [&]() {
-      loop.setLowerBound(newLoopParams.lowerBound);
-      loop.setUpperBound(newLoopParams.upperBound);
-      loop.setStep(newLoopParams.step);
+      loop.setLowerBound(cast<Value>(newLoopParams.lowerBound));
+      loop.setUpperBound(cast<Value>(newLoopParams.upperBound));
+      loop.setStep(cast<Value>(newLoopParams.step));
     });
 
     rewriter.setInsertionPointToStart(innermost.getBody());
@@ -778,8 +759,7 @@ void mlir::collapseParallelLoops(
     llvm::sort(dims);
 
   // Normalize ParallelOp's iteration pattern.
-  SmallVector<Value, 3> normalizedLowerBounds, normalizedSteps,
-      normalizedUpperBounds;
+  SmallVector<Value, 3> normalizedUpperBounds;
   for (unsigned i = 0, e = loops.getNumLoops(); i < e; ++i) {
     OpBuilder::InsertionGuard g2(rewriter);
     rewriter.setInsertionPoint(loops);
@@ -787,9 +767,7 @@ void mlir::collapseParallelLoops(
     Value ub = loops.getUpperBound()[i];
     Value step = loops.getStep()[i];
     auto newLoopParams = emitNormalizedLoopBounds(rewriter, loc, lb, ub, step);
-    normalizedLowerBounds.push_back(newLoopParams.lowerBound);
-    normalizedUpperBounds.push_back(newLoopParams.upperBound);
-    normalizedSteps.push_back(newLoopParams.step);
+    normalizedUpperBounds.push_back(cast<Value>(newLoopParams.upperBound));
 
     rewriter.setInsertionPointToStart(loops.getBody());
     denormalizeInductionVariable(rewriter, loc, loops.getInductionVars()[i], lb,