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

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: None (fabrizio-indirli)

<details>
<summary>Changes</summary>

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

``````````

</details>


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