[Mlir-commits] [mlir] [mlir][scf] Add parallelLoopUnrollByFactors() (#163806) (PR #164958)

[llvmlistbot at llvm.org]

https://github.com/fabrizio-indirli created https://github.com/llvm/llvm-project/pull/164958

- In the SCF Utils, add the `parallelLoopUnrollByFactors()` function to unroll scf::ParallelOp loops according to the specified unroll factors
- Add a test pass "TestParallelLoopUnrolling" and the related LIT test
- Expose `mlir::parallelLoopUnrollByFactors()`, `mlir::generateUnrolledLoop()`, and `mlir::scf::computeUbMinusLb()` functions in the mlir/Dialect/SCF/Utils/Utils.h and /IR/SCF.h headers to make them available to other passes.
- In `mlir::generateUnrolledLoop()`, add also an optional `IRMapping *clonedToSrcOpsMap` argument to map the new cloned operations to their original ones. In the function body, change the default `AnnotateFn` type to `static const` to silence potential warnings about dangling references when a function_ref is assigned  to a variable with automatic storage.

>From 59f967ed0ee917678bdf68b77fd7df383710c1af Mon Sep 17 00:00:00 2001
From: Fabrizio Indirli <Fabrizio.Indirli at arm.com>
Date: Fri, 24 Oct 2025 12:25:25 +0100
Subject: [PATCH] [mlir][scf] Add parallelLoopUnrollByFactors() (#163806)

- In the SCF Utils, add the `parallelLoopUnrollByFactors()` function
  to unroll scf::ParallelOp loops according to the specified unroll factors
- Add a test pass "TestParallelLoopUnrolling" and the related LIT test
- Expose `mlir::parallelLoopUnrollByFactors()`, `mlir::generateUnrolledLoop()`,
  and `mlir::scf::computeUbMinusLb()` functions in the
  mlir/Dialect/SCF/Utils/Utils.h and /IR/SCF.h headers to make them
  available to other passes.
- In `mlir::generateUnrolledLoop()`, add also an optional
  `IRMapping *clonedToSrcOpsMap` argument to map the new cloned
  operations to their original ones.
  In the function body, change the default `AnnotateFn` type to
  `static const` to silence potential warnings about dangling references
   when a function_ref is assigned  to a variable with automatic storage.

Signed-off-by: Fabrizio Indirli <Fabrizio.Indirli at arm.com>
---
 mlir/include/mlir/Dialect/SCF/IR/SCF.h      |  4 ++++
 mlir/include/mlir/Dialect/SCF/Utils/Utils.h |  6 ------
 mlir/lib/Dialect/SCF/IR/SCF.cpp             | 16 ++++++++++++++++
 mlir/lib/Dialect/SCF/Utils/Utils.cpp        | 16 ----------------
 4 files changed, 20 insertions(+), 22 deletions(-)

diff --git a/mlir/include/mlir/Dialect/SCF/IR/SCF.h b/mlir/include/mlir/Dialect/SCF/IR/SCF.h
index ba648181daecb..e754a04b0903a 100644
--- a/mlir/include/mlir/Dialect/SCF/IR/SCF.h
+++ b/mlir/include/mlir/Dialect/SCF/IR/SCF.h
@@ -112,6 +112,10 @@ SmallVector<Value> replaceAndCastForOpIterArg(RewriterBase &rewriter,
                                               Value replacement,
                                               const ValueTypeCastFnTy &castFn);
 
+/// Helper function to compute the difference between two values. This is used
+/// by the loop implementations to compute the trip count.
+std::optional<llvm::APSInt> computeUbMinusLb(Value lb, Value ub, bool isSigned);
+
 } // namespace scf
 } // namespace mlir
 #endif // MLIR_DIALECT_SCF_SCF_H
diff --git a/mlir/include/mlir/Dialect/SCF/Utils/Utils.h b/mlir/include/mlir/Dialect/SCF/Utils/Utils.h
index cdc52f4f3668c..3475bb2151777 100644
--- a/mlir/include/mlir/Dialect/SCF/Utils/Utils.h
+++ b/mlir/include/mlir/Dialect/SCF/Utils/Utils.h
@@ -254,12 +254,6 @@ FailureOr<scf::ParallelOp> parallelLoopUnrollByFactors(
 llvm::SmallVector<int64_t>
 getConstLoopTripCounts(mlir::LoopLikeOpInterface loopOp);
 
-namespace scf {
-/// Helper function to compute the difference between two values. This is used
-/// by the loop implementations to compute the trip count.
-std::optional<llvm::APSInt> computeUbMinusLb(Value lb, Value ub, bool isSigned);
-} // namespace scf
-
 } // namespace mlir
 
 #endif // MLIR_DIALECT_SCF_UTILS_UTILS_H_
diff --git a/mlir/lib/Dialect/SCF/IR/SCF.cpp b/mlir/lib/Dialect/SCF/IR/SCF.cpp
index 395b52fe46d25..bd12f8f00c917 100644
--- a/mlir/lib/Dialect/SCF/IR/SCF.cpp
+++ b/mlir/lib/Dialect/SCF/IR/SCF.cpp
@@ -112,6 +112,22 @@ static TerminatorTy verifyAndGetTerminator(Operation *op, Region &region,
   return nullptr;
 }
 
+std::optional<llvm::APSInt> mlir::scf::computeUbMinusLb(Value lb, Value ub,
+                                                        bool isSigned) {
+  llvm::APSInt diff;
+  auto addOp = ub.getDefiningOp<arith::AddIOp>();
+  if (!addOp)
+    return std::nullopt;
+  if ((isSigned && !addOp.hasNoSignedWrap()) ||
+      (!isSigned && !addOp.hasNoUnsignedWrap()))
+    return std::nullopt;
+
+  if (addOp.getLhs() != lb ||
+      !matchPattern(addOp.getRhs(), m_ConstantInt(&diff)))
+    return std::nullopt;
+  return diff;
+}
+
 //===----------------------------------------------------------------------===//
 // ExecuteRegionOp
 //===----------------------------------------------------------------------===//
diff --git a/mlir/lib/Dialect/SCF/Utils/Utils.cpp b/mlir/lib/Dialect/SCF/Utils/Utils.cpp
index 2d989d50bb8ac..888dd448b66f9 100644
--- a/mlir/lib/Dialect/SCF/Utils/Utils.cpp
+++ b/mlir/lib/Dialect/SCF/Utils/Utils.cpp
@@ -1559,22 +1559,6 @@ bool mlir::isPerfectlyNestedForLoops(
   return true;
 }
 
-std::optional<llvm::APSInt> mlir::scf::computeUbMinusLb(Value lb, Value ub,
-                                                        bool isSigned) {
-  llvm::APSInt diff;
-  auto addOp = ub.getDefiningOp<arith::AddIOp>();
-  if (!addOp)
-    return std::nullopt;
-  if ((isSigned && !addOp.hasNoSignedWrap()) ||
-      (!isSigned && !addOp.hasNoUnsignedWrap()))
-    return std::nullopt;
-
-  if (addOp.getLhs() != lb ||
-      !matchPattern(addOp.getRhs(), m_ConstantInt(&diff)))
-    return std::nullopt;
-  return diff;
-}
-
 llvm::SmallVector<int64_t>
 mlir::getConstLoopTripCounts(mlir::LoopLikeOpInterface loopOp) {
   std::optional<SmallVector<OpFoldResult>> loBnds = loopOp.getLoopLowerBounds();