[mlir][scf] Add parallelLoopUnrollByFactors() #163806
Merged
+443
−35
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
|
@llvm/pr-subscribers-mlir-core @llvm/pr-subscribers-mlir Author: None (fabrizio-indirli) ChangesDefine Full diff: https://github.com/llvm/llvm-project/pull/163806.diff 2 Files Affected:
diff --git a/mlir/include/mlir/Dialect/SCF/Utils/Utils.h b/mlir/include/mlir/Dialect/SCF/Utils/Utils.h
index ecd829ed14add..5c0908bb71fcf 100644
--- a/mlir/include/mlir/Dialect/SCF/Utils/Utils.h
+++ b/mlir/include/mlir/Dialect/SCF/Utils/Utils.h
@@ -221,6 +221,20 @@ FailureOr<scf::ForallOp> normalizeForallOp(RewriterBase &rewriter,
/// 4. Each region iter arg and result has exactly one use
bool isPerfectlyNestedForLoops(MutableArrayRef<LoopLikeOpInterface> loops);
+/// Generate unrolled copies of an scf loop's 'loopBodyBlock', with 'iterArgs'
+/// and 'yieldedValues' as the block arguments and yielded values of the loop.
+/// The content of the loop body is replicated 'unrollFactor' times, calling 'ivRemapFn' to remap
+/// 'iv' for each unrolled body. If specified, annotates the Ops in each
+/// unrolled iteration using annotateFn.
+/// If provided, 'clonedToSrcOpsMap' is populated with the mappings from the cloned
+/// ops to the original op.
+void generateUnrolledLoop(
+ Block *loopBodyBlock, Value iv, uint64_t unrollFactor,
+ function_ref<Value(unsigned, Value, OpBuilder)> ivRemapFn,
+ function_ref<void(unsigned, Operation *, OpBuilder)> annotateFn,
+ ValueRange iterArgs, ValueRange yieldedValues,
+ IRMapping* clonedToSrcOpsMap = nullptr);
+
} // namespace mlir
#endif // MLIR_DIALECT_SCF_UTILS_UTILS_H_
diff --git a/mlir/lib/Dialect/SCF/Utils/Utils.cpp b/mlir/lib/Dialect/SCF/Utils/Utils.cpp
index 10eae8906ce31..5220ba6c73252 100644
--- a/mlir/lib/Dialect/SCF/Utils/Utils.cpp
+++ b/mlir/lib/Dialect/SCF/Utils/Utils.cpp
@@ -291,47 +291,61 @@ static Value ceilDivPositive(OpBuilder &builder, Location loc, Value dividend,
return arith::DivUIOp::create(builder, loc, sum, divisor);
}
-/// Generates unrolled copies of scf::ForOp 'loopBodyBlock', with
-/// associated 'forOpIV' by 'unrollFactor', calling 'ivRemapFn' to remap
-/// 'forOpIV' for each unrolled body. If specified, annotates the Ops in each
-/// unrolled iteration using annotateFn.
-static void generateUnrolledLoop(
- Block *loopBodyBlock, Value forOpIV, uint64_t unrollFactor,
+void mlir::generateUnrolledLoop(
+ Block *loopBodyBlock, Value iv, uint64_t unrollFactor,
function_ref<Value(unsigned, Value, OpBuilder)> ivRemapFn,
function_ref<void(unsigned, Operation *, OpBuilder)> annotateFn,
- ValueRange iterArgs, ValueRange yieldedValues) {
+ ValueRange iterArgs, ValueRange yieldedValues,
+ IRMapping* clonedToSrcOpsMap) {
+
+ // check if the op was cloned from another source op, and return it if found
+ // (or the same op if not found)
+ auto findOriginalSrcOp =
+ [](Operation *op, const IRMapping &clonedToSrcOpsMap) -> Operation * {
+ Operation *srcOp = op;
+ // if the source op derives from another op: traverse the chain to find the
+ // original source op
+ while (srcOp && clonedToSrcOpsMap.contains(srcOp))
+ srcOp = clonedToSrcOpsMap.lookup(srcOp);
+ return srcOp;
+ };
+
// Builder to insert unrolled bodies just before the terminator of the body of
- // 'forOp'.
+ // the loop.
auto builder = OpBuilder::atBlockTerminator(loopBodyBlock);
- constexpr auto defaultAnnotateFn = [](unsigned, Operation *, OpBuilder) {};
+ static const auto noopAnnotateFn = [](unsigned, Operation *, OpBuilder) {};
if (!annotateFn)
- annotateFn = defaultAnnotateFn;
+ annotateFn = noopAnnotateFn;
// Keep a pointer to the last non-terminator operation in the original block
// so that we know what to clone (since we are doing this in-place).
Block::iterator srcBlockEnd = std::prev(loopBodyBlock->end(), 2);
- // Unroll the contents of 'forOp' (append unrollFactor - 1 additional copies).
+ // Unroll the contents of the loop body (append unrollFactor - 1 additional
+ // copies).
SmallVector<Value, 4> lastYielded(yieldedValues);
for (unsigned i = 1; i < unrollFactor; i++) {
- IRMapping operandMap;
-
// Prepare operand map.
+ IRMapping operandMap;
operandMap.map(iterArgs, lastYielded);
// If the induction variable is used, create a remapping to the value for
// this unrolled instance.
- if (!forOpIV.use_empty()) {
- Value ivUnroll = ivRemapFn(i, forOpIV, builder);
- operandMap.map(forOpIV, ivUnroll);
+ if (!iv.use_empty()) {
+ Value ivUnroll = ivRemapFn(i, iv, builder);
+ operandMap.map(iv, ivUnroll);
}
// Clone the original body of 'forOp'.
for (auto it = loopBodyBlock->begin(); it != std::next(srcBlockEnd); it++) {
- Operation *clonedOp = builder.clone(*it, operandMap);
+ Operation *srcOp = &(*it);
+ Operation *clonedOp = builder.clone(*srcOp, operandMap);
annotateFn(i, clonedOp, builder);
+ if(clonedToSrcOpsMap)
+ clonedToSrcOpsMap->map(clonedOp,
+ findOriginalSrcOp(srcOp, *clonedToSrcOpsMap));
}
// Update yielded values.
|
|
@llvm/pr-subscribers-mlir-scf Author: None (fabrizio-indirli) ChangesDefine Full diff: https://github.com/llvm/llvm-project/pull/163806.diff 2 Files Affected:
diff --git a/mlir/include/mlir/Dialect/SCF/Utils/Utils.h b/mlir/include/mlir/Dialect/SCF/Utils/Utils.h
index ecd829ed14add..5c0908bb71fcf 100644
--- a/mlir/include/mlir/Dialect/SCF/Utils/Utils.h
+++ b/mlir/include/mlir/Dialect/SCF/Utils/Utils.h
@@ -221,6 +221,20 @@ FailureOr<scf::ForallOp> normalizeForallOp(RewriterBase &rewriter,
/// 4. Each region iter arg and result has exactly one use
bool isPerfectlyNestedForLoops(MutableArrayRef<LoopLikeOpInterface> loops);
+/// Generate unrolled copies of an scf loop's 'loopBodyBlock', with 'iterArgs'
+/// and 'yieldedValues' as the block arguments and yielded values of the loop.
+/// The content of the loop body is replicated 'unrollFactor' times, calling 'ivRemapFn' to remap
+/// 'iv' for each unrolled body. If specified, annotates the Ops in each
+/// unrolled iteration using annotateFn.
+/// If provided, 'clonedToSrcOpsMap' is populated with the mappings from the cloned
+/// ops to the original op.
+void generateUnrolledLoop(
+ Block *loopBodyBlock, Value iv, uint64_t unrollFactor,
+ function_ref<Value(unsigned, Value, OpBuilder)> ivRemapFn,
+ function_ref<void(unsigned, Operation *, OpBuilder)> annotateFn,
+ ValueRange iterArgs, ValueRange yieldedValues,
+ IRMapping* clonedToSrcOpsMap = nullptr);
+
} // namespace mlir
#endif // MLIR_DIALECT_SCF_UTILS_UTILS_H_
diff --git a/mlir/lib/Dialect/SCF/Utils/Utils.cpp b/mlir/lib/Dialect/SCF/Utils/Utils.cpp
index 10eae8906ce31..5220ba6c73252 100644
--- a/mlir/lib/Dialect/SCF/Utils/Utils.cpp
+++ b/mlir/lib/Dialect/SCF/Utils/Utils.cpp
@@ -291,47 +291,61 @@ static Value ceilDivPositive(OpBuilder &builder, Location loc, Value dividend,
return arith::DivUIOp::create(builder, loc, sum, divisor);
}
-/// Generates unrolled copies of scf::ForOp 'loopBodyBlock', with
-/// associated 'forOpIV' by 'unrollFactor', calling 'ivRemapFn' to remap
-/// 'forOpIV' for each unrolled body. If specified, annotates the Ops in each
-/// unrolled iteration using annotateFn.
-static void generateUnrolledLoop(
- Block *loopBodyBlock, Value forOpIV, uint64_t unrollFactor,
+void mlir::generateUnrolledLoop(
+ Block *loopBodyBlock, Value iv, uint64_t unrollFactor,
function_ref<Value(unsigned, Value, OpBuilder)> ivRemapFn,
function_ref<void(unsigned, Operation *, OpBuilder)> annotateFn,
- ValueRange iterArgs, ValueRange yieldedValues) {
+ ValueRange iterArgs, ValueRange yieldedValues,
+ IRMapping* clonedToSrcOpsMap) {
+
+ // check if the op was cloned from another source op, and return it if found
+ // (or the same op if not found)
+ auto findOriginalSrcOp =
+ [](Operation *op, const IRMapping &clonedToSrcOpsMap) -> Operation * {
+ Operation *srcOp = op;
+ // if the source op derives from another op: traverse the chain to find the
+ // original source op
+ while (srcOp && clonedToSrcOpsMap.contains(srcOp))
+ srcOp = clonedToSrcOpsMap.lookup(srcOp);
+ return srcOp;
+ };
+
// Builder to insert unrolled bodies just before the terminator of the body of
- // 'forOp'.
+ // the loop.
auto builder = OpBuilder::atBlockTerminator(loopBodyBlock);
- constexpr auto defaultAnnotateFn = [](unsigned, Operation *, OpBuilder) {};
+ static const auto noopAnnotateFn = [](unsigned, Operation *, OpBuilder) {};
if (!annotateFn)
- annotateFn = defaultAnnotateFn;
+ annotateFn = noopAnnotateFn;
// Keep a pointer to the last non-terminator operation in the original block
// so that we know what to clone (since we are doing this in-place).
Block::iterator srcBlockEnd = std::prev(loopBodyBlock->end(), 2);
- // Unroll the contents of 'forOp' (append unrollFactor - 1 additional copies).
+ // Unroll the contents of the loop body (append unrollFactor - 1 additional
+ // copies).
SmallVector<Value, 4> lastYielded(yieldedValues);
for (unsigned i = 1; i < unrollFactor; i++) {
- IRMapping operandMap;
-
// Prepare operand map.
+ IRMapping operandMap;
operandMap.map(iterArgs, lastYielded);
// If the induction variable is used, create a remapping to the value for
// this unrolled instance.
- if (!forOpIV.use_empty()) {
- Value ivUnroll = ivRemapFn(i, forOpIV, builder);
- operandMap.map(forOpIV, ivUnroll);
+ if (!iv.use_empty()) {
+ Value ivUnroll = ivRemapFn(i, iv, builder);
+ operandMap.map(iv, ivUnroll);
}
// Clone the original body of 'forOp'.
for (auto it = loopBodyBlock->begin(); it != std::next(srcBlockEnd); it++) {
- Operation *clonedOp = builder.clone(*it, operandMap);
+ Operation *srcOp = &(*it);
+ Operation *clonedOp = builder.clone(*srcOp, operandMap);
annotateFn(i, clonedOp, builder);
+ if(clonedToSrcOpsMap)
+ clonedToSrcOpsMap->map(clonedOp,
+ findOriginalSrcOp(srcOp, *clonedToSrcOpsMap));
}
// Update yielded values.
|
|
✅ With the latest revision this PR passed the C/C++ code formatter. |
fabrizio-indirli
force-pushed
the
04892_unrPar
branch
from
8275239 to
6beb038
Compare
|
Looks good but can you call out the new mapping argument and the constexpr change in the commit message? |
fabrizio-indirli
force-pushed
the
04892_unrPar
branch
2 times, most recently
from
0f7cf97 to
95c44d9
Compare
|
Done, thanks for having a look! :) |
|
Not sure who should I tag to review this, please? |
fabrizio-indirli
force-pushed
the
04892_unrPar
branch
from
20d99ba to
d139d6f
Compare
fabrizio-indirli
force-pushed
the
04892_unrPar
branch
from
d139d6f to
c72da0e
Compare
fabrizio-indirli
changed the title
|
I added a test, and for that I had to add also a new function |
| function_ref<void(unsigned, Operation *, OpBuilder)> annotateFn = nullptr, | ||
| IRMapping *clonedToSrcOpsMap = nullptr); | ||
|
|
||
| namespace scf { |
There was a problem hiding this comment.
Why do we need this namespace?
There was a problem hiding this comment.
I'm not so sure of this either: I wrapped it around the namespace because this is was an internal helper function used only by other SCF functions; I exposed it in the header so that it can be used across different files (of the SCF dialect), but I wanted to make it clear that it's specific to SCF. I don't have any strong opinion on this, though
fabrizio-indirli
force-pushed
the
04892_unrPar
branch
2 times, most recently
from
09de4dc to
02ddf25
Compare
| %c16 = arith.constant 16 : index | ||
| %c0 = arith.constant 0 : index | ||
| %c1 = arith.constant 1 : index | ||
| scf.parallel (%arg2, %arg3, %arg4) = (%c0, %c0, %c0) to (%c1, %c16, %c12) step (%c1, %c1, %c1) { |
There was a problem hiding this comment.
[nit]
Suggested change
| scf.parallel (%arg2, %arg3, %arg4) = (%c0, %c0, %c0) to (%c1, %c16, %c12) step (%c1, %c1, %c1) { | |
| scf.parallel (%iv0, %iv1, %iv2) = (%c0, %c0, %c0) to (%c1, %c16, %c12) step (%c1, %c1, %c1) { |
Same comment elsewhere.
There was a problem hiding this comment.
I'm not a big fan of this tbh because when you run the pass manually mlir-opt will rename al the variable names to %arg0, %arg1, etc. and it becomes harder to map them to the original names in the test
There was a problem hiding this comment.
mlir-opt will rename al the variable names to %arg0, %arg1
Indeed, we are not in control of what mlir-opt does with names. So, even %arg2, %arg3, can (and probably will) be replaced with %arg0, %arg2. For this reason, we should optimise for what makes parsing for humans easier.
fabrizio-indirli
force-pushed
the
04892_unrPar
branch
4 times, most recently
from
888e728 to
35cb6c7
Compare
banach-space
approved these changes
Oct 24, 2025
There was a problem hiding this comment.
LGTM, thank you for your contribution and for addressing my comments!
My outstanding comments are minor (nits), no need to wait for me take another look. Also, I am travelling today and next week, so won't be able to re-visit this until some later time :) (i.e. feel free to land this)
| %c16 = arith.constant 16 : index | ||
| %c0 = arith.constant 0 : index | ||
| %c1 = arith.constant 1 : index | ||
| scf.parallel (%arg2, %arg3, %arg4) = (%c0, %c0, %c0) to (%c1, %c16, %c12) step (%c1, %c1, %c1) { |
There was a problem hiding this comment.
mlir-opt will rename al the variable names to %arg0, %arg1
Indeed, we are not in control of what mlir-opt does with names. So, even %arg2, %arg3, can (and probably will) be replaced with %arg0, %arg2. For this reason, we should optimise for what makes parsing for humans easier.
- 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 header to make thme 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 <[email protected]>
fabrizio-indirli
force-pushed
the
04892_unrPar
branch
from
0eeeb38 to
0935f5d
Compare
|
LLVM Buildbot has detected a new failure on builder Full details are available at: https://lab.llvm.org/buildbot/#/builders/117/builds/14487 Here is the relevant piece of the build log for the reference |
|
LLVM Buildbot has detected a new failure on builder Full details are available at: https://lab.llvm.org/buildbot/#/builders/140/builds/32977 Here is the relevant piece of the build log for the reference |
|
LLVM Buildbot has detected a new failure on builder Full details are available at: https://lab.llvm.org/buildbot/#/builders/138/builds/20873 Here is the relevant piece of the build log for the reference |
|
LLVM Buildbot has detected a new failure on builder Full details are available at: https://lab.llvm.org/buildbot/#/builders/205/builds/25988 Here is the relevant piece of the build log for the reference |
|
LLVM Buildbot has detected a new failure on builder Full details are available at: https://lab.llvm.org/buildbot/#/builders/53/builds/21538 Here is the relevant piece of the build log for the reference |
|
LLVM Buildbot has detected a new failure on builder Full details are available at: https://lab.llvm.org/buildbot/#/builders/203/builds/27199 Here is the relevant piece of the build log for the reference |
|
LLVM Buildbot has detected a new failure on builder Full details are available at: https://lab.llvm.org/buildbot/#/builders/204/builds/26011 Here is the relevant piece of the build log for the reference |
fabrizio-indirli
added a commit
that referenced
this pull request
Oct 24, 2025
This reverts commit 86a2073.
fabrizio-indirli
added a commit
that referenced
this pull request
Oct 24, 2025
Reverts #163806 due to linking errors on the function `mlir::scf::computeUbMinusLb`
|
LLVM Buildbot has detected a new failure on builder Full details are available at: https://lab.llvm.org/buildbot/#/builders/172/builds/16135 Here is the relevant piece of the build log for the reference |
|
LLVM Buildbot has detected a new failure on builder Full details are available at: https://lab.llvm.org/buildbot/#/builders/29/builds/17517 Here is the relevant piece of the build log for the reference |
|
LLVM Buildbot has detected a new failure on builder Full details are available at: https://lab.llvm.org/buildbot/#/builders/80/builds/16967 Here is the relevant piece of the build log for the reference |
|
LLVM Buildbot has detected a new failure on builder Full details are available at: https://lab.llvm.org/buildbot/#/builders/152/builds/5863 Here is the relevant piece of the build log for the reference |
llvm-sync bot
pushed a commit
to arm/arm-toolchain
that referenced
this pull request
Oct 24, 2025
…64949) Reverts llvm/llvm-project#163806 due to linking errors on the function `mlir::scf::computeUbMinusLb`
fabrizio-indirli
added a commit
to fabrizio-indirli/llvm-project
that referenced
this pull request
Oct 24, 2025
- 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 <[email protected]>
|
LLVM Buildbot has detected a new failure on builder Full details are available at: https://lab.llvm.org/buildbot/#/builders/130/builds/15825 Here is the relevant piece of the build log for the reference |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
5 participants
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
parallelLoopUnrollByFactors()functionto unroll scf::ParallelOp loops according to the specified unroll factors
mlir::parallelLoopUnrollByFactors(),mlir::generateUnrolledLoop(),and
mlir::scf::computeUbMinusLb()functions in themlir/Dialect/SCF/Utils/Utils.h header to make them available to other passes.
mlir::generateUnrolledLoop(), add also an optionalIRMapping *clonedToSrcOpsMapargument to map the new cloned operations to their original ones.
In the function body, change the default
AnnotateFntype tostatic constto silence potential warnings about dangling references when a function_ref
is assigned to a variable with automatic storage.