[mlir][scf] Add parallelLoopUnrollByFactors()

[fabrizio-indirli]

• 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.

[llvmbot]

@llvm/pr-subscribers-mlir-core
@llvm/pr-subscribers-mlir-scf

@llvm/pr-subscribers-mlir

Author: None (fabrizio-indirli)

Changes

• 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.

---

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

4 Files Affected:

• (modified) mlir/include/mlir/Dialect/SCF/IR/SCF.h (+4)
• (modified) mlir/include/mlir/Dialect/SCF/Utils/Utils.h (-6)
• (modified) mlir/lib/Dialect/SCF/IR/SCF.cpp (+16)
• (modified) mlir/lib/Dialect/SCF/Utils/Utils.cpp (-16)

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();

[fabrizio-indirli]

This is the same PR as #163806 , which was reverted in #164949 due to a linking error: computeUbMinusLb had been erroneously moved to a different file (Utils.cpp instead of SCF.cpp). Here the issue was fixed ; everything else is the same