[mlir][scf]: Expose emitNormalizedLoopBounds/denormalizeInductionVariable util functions (NFC)

[AviadCo]

No description provided.

[llvmbot]

@llvm/pr-subscribers-mlir-arith
@llvm/pr-subscribers-mlir-affine
@llvm/pr-subscribers-mlir-scf

@llvm/pr-subscribers-mlir

Author: Aviad Cohen (AviadCo)

Changes

---

Full diff: https://github.com/llvm/llvm-project/pull/94429.diff

2 Files Affected:

• (modified) mlir/include/mlir/Dialect/SCF/Utils/Utils.h (+25)
• (modified) mlir/lib/Dialect/SCF/Utils/Utils.cpp (+5-25)

diff --git a/mlir/include/mlir/Dialect/SCF/Utils/Utils.h b/mlir/include/mlir/Dialect/SCF/Utils/Utils.h
index bc09cc7f7fa5e..dde2daddb81e0 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 {
+  Value lowerBound;
+  Value upperBound;
+  Value 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..dd8a650cf1089 100644
--- a/mlir/lib/Dialect/SCF/Utils/Utils.cpp
+++ b/mlir/lib/Dialect/SCF/Utils/Utils.cpp
@@ -29,16 +29,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 +463,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 +498,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);

[AviadCo]

Thanks for the code review!

[MaheshRavishankar]

One more comment, but looks good to me!

[AviadCo]

@MaheshRavishankar can please take a look at the last revision? thanks!

[MaheshRavishankar]

If changing the function to take OpFoldResult arguments works without issues, this looks good to me.

[AviadCo]

If changing the function to take OpFoldResult arguments works without issues, this looks good to me.

I refactored the func args, it added a little more wrapping functions but works the same.

[AviadCo]

@MaheshRavishankar can u please take another look?

[MaheshRavishankar]

Looks fine to me! I already approved. Are you wait for me to land this?

[MaheshRavishankar]

Actually I just realized that LoopParams here is same as Range here . I should have pointed it out earlier. Fine to do as a follow up, but would appreciate just using Range instead of LoopParams. If not, I can do as follow up.

[AviadCo]

Looks fine to me! I already approved. Are you wait for me to land this?

I thought I need reapproval after new changes, will land this one

[AviadCo]

Actually I just realized that LoopParams here is same as Range here . I should have pointed it out earlier. Fine to do as a follow up, but would appreciate just using Range instead of LoopParams. If not, I can do as follow up.

I will do the follow up PR it is o.k.

[AviadCo]

@MaheshRavishankar Please review the follow up PR: #95501
@kuhar thanks for the review here also :)