[Mlir-commits] [mlir] [mlir][scf] Expose isPerfectlyNestedForLoops (PR #152115)
Shay Kleiman
llvmlistbot at llvm.org
Mon Aug 25 06:48:47 PDT 2025
https://github.com/shay-kl updated https://github.com/llvm/llvm-project/pull/152115
>From 9e9093fb6a2396c4c682aa880d6c4d41a65cd086 Mon Sep 17 00:00:00 2001
From: Shay Kleiman <shay.kleiman at mobileye.com>
Date: Tue, 29 Jul 2025 18:02:59 +0300
Subject: [PATCH 1/2] [mlir][scf] Expose isPerfectlyNestedForLoops
---
.../mlir/Dialect/SCF/Transforms/TileUsingInterface.h | 8 ++++++++
mlir/lib/Dialect/SCF/Transforms/TileUsingInterface.cpp | 6 +++---
2 files changed, 11 insertions(+), 3 deletions(-)
diff --git a/mlir/include/mlir/Dialect/SCF/Transforms/TileUsingInterface.h b/mlir/include/mlir/Dialect/SCF/Transforms/TileUsingInterface.h
index 3205da6e448fc..1177aee560f6a 100644
--- a/mlir/include/mlir/Dialect/SCF/Transforms/TileUsingInterface.h
+++ b/mlir/include/mlir/Dialect/SCF/Transforms/TileUsingInterface.h
@@ -364,6 +364,14 @@ FailureOr<scf::SCFTilingResult>
tileReductionUsingScf(RewriterBase &b, PartialReductionOpInterface op,
ArrayRef<OpFoldResult> tileSizes);
+/// Check if the provided loops are perfectly nested for-loops. Perfect nesting
+/// means:
+/// 1. All loops are scf.for operations
+/// 2. Each outer loop's region iter args match the inner loop's init args
+/// 3. Each outer loop's yields match the inner loop's results
+/// 4. Each region iter arg and result has exactly one use
+bool isPerfectlyNestedForLoops(MutableArrayRef<LoopLikeOpInterface> loops);
+
} // namespace scf
} // namespace mlir
diff --git a/mlir/lib/Dialect/SCF/Transforms/TileUsingInterface.cpp b/mlir/lib/Dialect/SCF/Transforms/TileUsingInterface.cpp
index c0e47ee1e74fc..00f1b25681a3d 100644
--- a/mlir/lib/Dialect/SCF/Transforms/TileUsingInterface.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/TileUsingInterface.cpp
@@ -1927,8 +1927,8 @@ static FailureOr<OpOperand *> getConsumerFromLoopUses(RewriterBase &rewriter,
/// yield %1
/// ```
/// Here loops should be [%0, %1].
-static bool
-isPerfectlyNestedForLoops(MutableArrayRef<LoopLikeOpInterface> loops) {
+bool mlir::scf::isPerfectlyNestedForLoops(
+ MutableArrayRef<LoopLikeOpInterface> loops) {
assert(!loops.empty() && "unexpected empty loop nest");
if (loops.size() == 1) {
return isa_and_nonnull<scf::ForOp>(loops.front().getOperation());
@@ -1991,7 +1991,7 @@ getUntiledConsumerFromSlice(RewriterBase &rewriter,
}
// 2. Check that the loop is perfectly nested.
- if (!isPerfectlyNestedForLoops(loops)) {
+ if (!mlir::scf::isPerfectlyNestedForLoops(loops)) {
return rewriter.notifyMatchFailure(
candidateSliceOp, "expected passed loops to be perfectly nested.");
}
>From df9ccb7e18c7911bb517ba05748d5e763bd66225 Mon Sep 17 00:00:00 2001
From: Shay Kleiman <shay.kleiman at mobileye.com>
Date: Mon, 25 Aug 2025 16:48:25 +0300
Subject: [PATCH 2/2] moved
---
mlir/include/mlir/Dialect/SCF/Utils/Utils.h | 8 +++
.../SCF/Transforms/TileUsingInterface.cpp | 59 +------------------
mlir/lib/Dialect/SCF/Utils/Utils.cpp | 44 ++++++++++++++
3 files changed, 53 insertions(+), 58 deletions(-)
diff --git a/mlir/include/mlir/Dialect/SCF/Utils/Utils.h b/mlir/include/mlir/Dialect/SCF/Utils/Utils.h
index e620067c15be9..ecd829ed14add 100644
--- a/mlir/include/mlir/Dialect/SCF/Utils/Utils.h
+++ b/mlir/include/mlir/Dialect/SCF/Utils/Utils.h
@@ -213,6 +213,14 @@ scf::ForOp fuseIndependentSiblingForLoops(scf::ForOp target, scf::ForOp source,
FailureOr<scf::ForallOp> normalizeForallOp(RewriterBase &rewriter,
scf::ForallOp forallOp);
+/// Check if the provided loops are perfectly nested for-loops. Perfect nesting
+/// means:
+/// 1. All loops are scf.for operations
+/// 2. Each outer loop's region iter args match the inner loop's init args
+/// 3. Each outer loop's yields match the inner loop's results
+/// 4. Each region iter arg and result has exactly one use
+bool isPerfectlyNestedForLoops(MutableArrayRef<LoopLikeOpInterface> loops);
+
} // namespace mlir
#endif // MLIR_DIALECT_SCF_UTILS_UTILS_H_
diff --git a/mlir/lib/Dialect/SCF/Transforms/TileUsingInterface.cpp b/mlir/lib/Dialect/SCF/Transforms/TileUsingInterface.cpp
index 00f1b25681a3d..c7726640fb579 100644
--- a/mlir/lib/Dialect/SCF/Transforms/TileUsingInterface.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/TileUsingInterface.cpp
@@ -1914,63 +1914,6 @@ static FailureOr<OpOperand *> getConsumerFromLoopUses(RewriterBase &rewriter,
return failure();
}
-/// Check that the loop is perfectly nested.
-/// The loops are expected to be ordered from outer most to inner most.
-/// For example:
-/// ```
-/// %0 = scf.for()
-/// %1 = scf.for()
-/// %2 = scf.for()
-/// %3 = ...
-/// yield %3
-/// yield %2
-/// yield %1
-/// ```
-/// Here loops should be [%0, %1].
-bool mlir::scf::isPerfectlyNestedForLoops(
- MutableArrayRef<LoopLikeOpInterface> loops) {
- assert(!loops.empty() && "unexpected empty loop nest");
- if (loops.size() == 1) {
- return isa_and_nonnull<scf::ForOp>(loops.front().getOperation());
- }
- for (auto [outerLoop, innerLoop] :
- llvm::zip_equal(loops.drop_back(), loops.drop_front())) {
- auto outerFor = dyn_cast_or_null<scf::ForOp>(outerLoop.getOperation());
- auto innerFor = dyn_cast_or_null<scf::ForOp>(innerLoop.getOperation());
- if (!outerFor || !innerFor) {
- return false;
- }
- auto outerBBArgs = outerFor.getRegionIterArgs();
- auto innerIterArgs = innerFor.getInitArgs();
- if (outerBBArgs.size() != innerIterArgs.size()) {
- return false;
- }
-
- for (auto [outerBBArg, innerIterArg] :
- llvm::zip_equal(outerBBArgs, innerIterArgs)) {
- if (!llvm::hasSingleElement(outerBBArg.getUses()) ||
- innerIterArg != outerBBArg) {
- return false;
- }
- }
-
- ValueRange outerYields =
- cast<scf::YieldOp>(outerFor.getBody()->getTerminator())->getOperands();
- ValueRange innerResults = innerFor.getResults();
- if (outerYields.size() != innerResults.size()) {
- return false;
- }
- for (auto [outerYield, innerResult] :
- llvm::zip_equal(outerYields, innerResults)) {
- if (!llvm::hasSingleElement(innerResult.getUses()) ||
- outerYield != innerResult) {
- return false;
- }
- }
- }
- return true;
-}
-
/// Fetch the untiled consumer of the outermost scf.for's result which is
/// yielded by a tensor.insert_slice from the innermost scf.for. This function
/// makes the following assumptions :
@@ -1991,7 +1934,7 @@ getUntiledConsumerFromSlice(RewriterBase &rewriter,
}
// 2. Check that the loop is perfectly nested.
- if (!mlir::scf::isPerfectlyNestedForLoops(loops)) {
+ if (!mlir::isPerfectlyNestedForLoops(loops)) {
return rewriter.notifyMatchFailure(
candidateSliceOp, "expected passed loops to be perfectly nested.");
}
diff --git a/mlir/lib/Dialect/SCF/Utils/Utils.cpp b/mlir/lib/Dialect/SCF/Utils/Utils.cpp
index 57317951d609c..dfdf019354d5f 100644
--- a/mlir/lib/Dialect/SCF/Utils/Utils.cpp
+++ b/mlir/lib/Dialect/SCF/Utils/Utils.cpp
@@ -1506,3 +1506,47 @@ FailureOr<scf::ForallOp> mlir::normalizeForallOp(RewriterBase &rewriter,
rewriter.replaceOp(forallOp, normalizedForallOp);
return normalizedForallOp;
}
+
+bool mlir::isPerfectlyNestedForLoops(
+ MutableArrayRef<LoopLikeOpInterface> loops) {
+ assert(!loops.empty() && "unexpected empty loop nest");
+ if (loops.size() == 1) {
+ return isa_and_nonnull<scf::ForOp>(loops.front().getOperation());
+ }
+ for (auto [outerLoop, innerLoop] :
+ llvm::zip_equal(loops.drop_back(), loops.drop_front())) {
+ auto outerFor = dyn_cast_or_null<scf::ForOp>(outerLoop.getOperation());
+ auto innerFor = dyn_cast_or_null<scf::ForOp>(innerLoop.getOperation());
+ if (!outerFor || !innerFor) {
+ return false;
+ }
+ auto outerBBArgs = outerFor.getRegionIterArgs();
+ auto innerIterArgs = innerFor.getInitArgs();
+ if (outerBBArgs.size() != innerIterArgs.size()) {
+ return false;
+ }
+
+ for (auto [outerBBArg, innerIterArg] :
+ llvm::zip_equal(outerBBArgs, innerIterArgs)) {
+ if (!llvm::hasSingleElement(outerBBArg.getUses()) ||
+ innerIterArg != outerBBArg) {
+ return false;
+ }
+ }
+
+ ValueRange outerYields =
+ cast<scf::YieldOp>(outerFor.getBody()->getTerminator())->getOperands();
+ ValueRange innerResults = innerFor.getResults();
+ if (outerYields.size() != innerResults.size()) {
+ return false;
+ }
+ for (auto [outerYield, innerResult] :
+ llvm::zip_equal(outerYields, innerResults)) {
+ if (!llvm::hasSingleElement(innerResult.getUses()) ||
+ outerYield != innerResult) {
+ return false;
+ }
+ }
+ }
+ return true;
+}
\ No newline at end of file
More information about the Mlir-commits
mailing list