[Mlir-commits] [mlir] Codex/affine parallel scope depth (PR #228919)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Sun Oct 4 10:00:47 PDT 2026
https://github.com/ParkHanbum created https://github.com/llvm/llvm-project/pull/228919
[mlir][affine] Fix dependence-check depth in isLoopMemoryParallel
isLoopMemoryParallel computed the dependence-check depth with
getNestingDepth(), which also counts loops outside the nearest
AffineScope (e.g. an affine.for around a gpu.launch). The access
domains built by getOpIndexSet() stop at the scope boundary, so the
depth could disagree with the number of domain dimensions.
With a single out-of-scope loop, only same-iteration dependences were
checked, so a loop carrying a dependence was wrongly converted to
affine.parallel. With more out-of-scope loops, the depth exceeded the
number of common loops and hit the "Invalid depth" assertion.
Compute the depth from the ops returned by getEnclosingAffineOps(),
which stops at the scope boundary: one dimension per affine.for and
getNumDims() per affine.parallel; affine.if adds none.
Fixes https://github.com/llvm/llvm-project/issues/228862
>From 665ac9e5189c7f0c9f05fc8621f0b6a56ce39d1d Mon Sep 17 00:00:00 2001
From: hanbeom <kese111 at gmail.com>
Date: Sun, 4 Oct 2026 19:32:21 +0900
Subject: [PATCH 1/2] [mlir][affine] Precommit scoped parallelization
regressions
Record the existing incorrect parallelization across an affine scope: the
dependent %i loops in scope_depth and scope_depth_output_dependence are
currently converted to affine.parallel. Keep baseline check lines and mark
this behavior with FIXME comments for the follow-up fix.
Also include an in-scope parallel-dimension control with unchanged, correct
behavior.
---
.../Affine/parallelize-affine-scope.mlir | 69 +++++++++++++++++++
1 file changed, 69 insertions(+)
create mode 100644 mlir/test/Dialect/Affine/parallelize-affine-scope.mlir
diff --git a/mlir/test/Dialect/Affine/parallelize-affine-scope.mlir b/mlir/test/Dialect/Affine/parallelize-affine-scope.mlir
new file mode 100644
index 0000000000000..426392369a723
--- /dev/null
+++ b/mlir/test/Dialect/Affine/parallelize-affine-scope.mlir
@@ -0,0 +1,69 @@
+// RUN: split-file %s %t
+// RUN: mlir-opt %t/scope-depth.mlir -affine-parallelize | FileCheck %s --check-prefix=SCOPE
+// RUN: mlir-opt %t/parallel-dimensions.mlir -affine-parallelize | FileCheck %s --check-prefix=PARALLEL
+// RUN: mlir-opt %t/output-dependence.mlir -affine-parallelize | FileCheck %s --check-prefix=OUTPUT
+
+//--- scope-depth.mlir
+// Loops outside gpu.launch must not shift the dependence-check dimension.
+// FIXME: %i is incorrectly parallelized; the follow-up fix must keep it sequential.
+// SCOPE-LABEL: func.func @scope_depth
+// SCOPE: affine.for {{.*}} = 0 to 1 {
+// SCOPE-NEXT: gpu.launch
+// SCOPE-NEXT: affine.parallel ({{.*}}) = (1) to (5) {
+// SCOPE-NEXT: affine.parallel ({{.*}}) = (0) to (3) {
+func.func @scope_depth(%a: memref<6x4xf32>) {
+ %c1 = arith.constant 1 : index
+ affine.for %q = 0 to 1 {
+ gpu.launch blocks(%bx, %by, %bz) in (%gx = %c1, %gy = %c1, %gz = %c1)
+ threads(%tx, %ty, %tz) in (%sx = %c1, %sy = %c1, %sz = %c1) {
+ affine.for %i = 1 to 5 {
+ affine.for %j = 0 to 3 {
+ %v = affine.load %a[%i - 1, %j + 1] : memref<6x4xf32>
+ affine.store %v, %a[%i, %j] : memref<6x4xf32>
+ }
+ }
+ gpu.terminator
+ }
+ }
+ return
+}
+
+//--- parallel-dimensions.mlir
+// Each IV of an enclosing affine.parallel contributes an iteration dimension.
+// PARALLEL-LABEL: func.func @parallel_dimensions
+// PARALLEL: affine.parallel ({{.*}}, {{.*}}) = (0, 0) to (2, 2) {
+// PARALLEL-NEXT: affine.for {{.*}} = 1 to 5 {
+// PARALLEL-NEXT: affine.parallel ({{.*}}) = (0) to (3) {
+func.func @parallel_dimensions(%a: memref<2x2x6x4xf32>) {
+ affine.parallel (%p, %q) = (0, 0) to (2, 2) {
+ affine.for %i = 1 to 5 {
+ affine.for %j = 0 to 3 {
+ %v = affine.load %a[%p, %q, %i - 1, %j + 1] : memref<2x2x6x4xf32>
+ affine.store %v, %a[%p, %q, %i, %j] : memref<2x2x6x4xf32>
+ }
+ }
+ }
+ return
+}
+
+//--- output-dependence.mlir
+// A store to a fixed address has an output dependence across loop iterations.
+// FIXME: %i is incorrectly parallelized; the follow-up fix must keep it sequential.
+// OUTPUT-LABEL: func.func @scope_depth_output_dependence
+// OUTPUT: affine.for {{.*}} = 0 to 1 {
+// OUTPUT-NEXT: gpu.launch
+// OUTPUT-NEXT: affine.parallel ({{.*}}) = (0) to (4) {
+// OUTPUT-NEXT: affine.store
+func.func @scope_depth_output_dependence(%a: memref<1xindex>) {
+ %c1 = arith.constant 1 : index
+ affine.for %q = 0 to 1 {
+ gpu.launch blocks(%bx, %by, %bz) in (%gx = %c1, %gy = %c1, %gz = %c1)
+ threads(%tx, %ty, %tz) in (%sx = %c1, %sy = %c1, %sz = %c1) {
+ affine.for %i = 0 to 4 {
+ affine.store %i, %a[0] : memref<1xindex>
+ }
+ gpu.terminator
+ }
+ }
+ return
+}
>From 467429514ce5578c4377987de1993ddae2233faf Mon Sep 17 00:00:00 2001
From: hanbeom <kese111 at gmail.com>
Date: Sun, 4 Oct 2026 19:34:20 +0900
Subject: [PATCH 2/2] [mlir][affine] Fix dependence-check depth in
isLoopMemoryParallel
isLoopMemoryParallel computed the dependence-check depth with
getNestingDepth(), which also counts loops outside the nearest
AffineScope (e.g. an affine.for around a gpu.launch). The access
domains built by getOpIndexSet() stop at the scope boundary, so the
depth could disagree with the number of domain dimensions.
With a single out-of-scope loop, only same-iteration dependences were
checked, so a loop carrying a dependence was wrongly converted to
affine.parallel. With more out-of-scope loops, the depth exceeded the
number of common loops and hit the "Invalid depth" assertion.
Compute the depth from the ops returned by getEnclosingAffineOps(),
which stops at the scope boundary: one dimension per affine.for and
getNumDims() per affine.parallel; affine.if adds none.
Fixes #228862
---
.../Affine/Analysis/AffineAnalysis.cpp | 13 +++-
.../Affine/parallelize-affine-scope.mlir | 71 +++++++++++++++++--
2 files changed, 78 insertions(+), 6 deletions(-)
diff --git a/mlir/lib/Dialect/Affine/Analysis/AffineAnalysis.cpp b/mlir/lib/Dialect/Affine/Analysis/AffineAnalysis.cpp
index 27a8c7b8ea863..a5e1acebcb7bd 100644
--- a/mlir/lib/Dialect/Affine/Analysis/AffineAnalysis.cpp
+++ b/mlir/lib/Dialect/Affine/Analysis/AffineAnalysis.cpp
@@ -166,8 +166,17 @@ bool mlir::affine::isLoopMemoryParallel(AffineForOp forOp) {
if (walkResult.wasInterrupted())
return false;
- // Dep check depth would be number of enclosing loops + 1.
- unsigned depth = getNestingDepth(forOp) + 1;
+ // Match the dimensions collected by getOpIndexSet: loops outside the
+ // surrounding affine scope do not contribute to the access domain.
+ SmallVector<Operation *> enclosingOps;
+ getEnclosingAffineOps(*forOp, &enclosingOps);
+ unsigned depth = 1;
+ for (Operation *op : enclosingOps) {
+ if (isa<AffineForOp>(op))
+ ++depth;
+ else if (auto parallelOp = dyn_cast<AffineParallelOp>(op))
+ depth += parallelOp.getNumDims();
+ }
// Check dependences between all pairs of ops in 'loadAndStoreOps'.
for (auto *srcOp : loadAndStoreOps) {
diff --git a/mlir/test/Dialect/Affine/parallelize-affine-scope.mlir b/mlir/test/Dialect/Affine/parallelize-affine-scope.mlir
index 426392369a723..8e31540ae175b 100644
--- a/mlir/test/Dialect/Affine/parallelize-affine-scope.mlir
+++ b/mlir/test/Dialect/Affine/parallelize-affine-scope.mlir
@@ -2,14 +2,15 @@
// RUN: mlir-opt %t/scope-depth.mlir -affine-parallelize | FileCheck %s --check-prefix=SCOPE
// RUN: mlir-opt %t/parallel-dimensions.mlir -affine-parallelize | FileCheck %s --check-prefix=PARALLEL
// RUN: mlir-opt %t/output-dependence.mlir -affine-parallelize | FileCheck %s --check-prefix=OUTPUT
+// RUN: mlir-opt %t/multiple-outer-loops.mlir -affine-parallelize | FileCheck %s --check-prefix=MULTIPLE
+// RUN: mlir-opt %t/scoped-parallel-dimensions.mlir -affine-parallelize | FileCheck %s --check-prefix=SCOPED
//--- scope-depth.mlir
// Loops outside gpu.launch must not shift the dependence-check dimension.
-// FIXME: %i is incorrectly parallelized; the follow-up fix must keep it sequential.
// SCOPE-LABEL: func.func @scope_depth
// SCOPE: affine.for {{.*}} = 0 to 1 {
// SCOPE-NEXT: gpu.launch
-// SCOPE-NEXT: affine.parallel ({{.*}}) = (1) to (5) {
+// SCOPE-NEXT: affine.for {{.*}} = 1 to 5 {
// SCOPE-NEXT: affine.parallel ({{.*}}) = (0) to (3) {
func.func @scope_depth(%a: memref<6x4xf32>) {
%c1 = arith.constant 1 : index
@@ -48,11 +49,10 @@ func.func @parallel_dimensions(%a: memref<2x2x6x4xf32>) {
//--- output-dependence.mlir
// A store to a fixed address has an output dependence across loop iterations.
-// FIXME: %i is incorrectly parallelized; the follow-up fix must keep it sequential.
// OUTPUT-LABEL: func.func @scope_depth_output_dependence
// OUTPUT: affine.for {{.*}} = 0 to 1 {
// OUTPUT-NEXT: gpu.launch
-// OUTPUT-NEXT: affine.parallel ({{.*}}) = (0) to (4) {
+// OUTPUT-NEXT: affine.for {{.*}} = 0 to 4 {
// OUTPUT-NEXT: affine.store
func.func @scope_depth_output_dependence(%a: memref<1xindex>) {
%c1 = arith.constant 1 : index
@@ -67,3 +67,66 @@ func.func @scope_depth_output_dependence(%a: memref<1xindex>) {
}
return
}
+
+// These inputs used to abort with "Invalid depth" because loops outside the
+// affine scope inflated the dependence-check depth. Check normal completion
+// and the parallelization decisions after fixing the scoped depth.
+
+//--- multiple-outer-loops.mlir
+// Multiple out-of-scope loops must not exceed the scoped access-domain depth.
+// MULTIPLE-LABEL: func.func @scope_depth_multiple_outer_loops
+// MULTIPLE: affine.for {{.*}} = 0 to 1 {
+// MULTIPLE-NEXT: affine.for {{.*}} = 0 to 1 {
+// MULTIPLE-NEXT: gpu.launch
+// MULTIPLE-NEXT: affine.for {{.*}} = 1 to 5 {
+// MULTIPLE: affine.parallel ({{.*}}) = (0) to (4) {
+func.func @scope_depth_multiple_outer_loops(%a: memref<6xf32>, %b: memref<4xf32>) {
+ %c1 = arith.constant 1 : index
+ affine.for %q = 0 to 1 {
+ affine.for %r = 0 to 1 {
+ gpu.launch blocks(%bx, %by, %bz) in (%gx = %c1, %gy = %c1, %gz = %c1)
+ threads(%tx, %ty, %tz) in (%sx = %c1, %sy = %c1, %sz = %c1) {
+ affine.for %i = 1 to 5 {
+ %v = affine.load %a[%i - 1] : memref<6xf32>
+ affine.store %v, %a[%i] : memref<6xf32>
+ }
+ affine.for %j = 0 to 4 {
+ %v = affine.load %b[%j] : memref<4xf32>
+ affine.store %v, %b[%j] : memref<4xf32>
+ }
+ gpu.terminator
+ }
+ }
+ }
+ return
+}
+
+//--- scoped-parallel-dimensions.mlir
+// Count in-scope parallel dimensions, exclude out-of-scope ones, and ignore affine.if.
+// SCOPED-LABEL: func.func @scope_depth_parallel_dimensions
+// SCOPED: affine.parallel ({{.*}}, {{.*}}) = (0, 0) to (1, 1) {
+// SCOPED-NEXT: gpu.launch
+// SCOPED-NEXT: affine.parallel ({{.*}}, {{.*}}) = (0, 0) to (2, 2) {
+// SCOPED-NEXT: affine.for {{.*}} = 1 to 5 {
+// SCOPED-NEXT: affine.if
+// SCOPED-NEXT: affine.parallel ({{.*}}) = (0) to (3) {
+func.func @scope_depth_parallel_dimensions(%a: memref<2x2x6x4xf32>) {
+ %c1 = arith.constant 1 : index
+ affine.parallel (%r, %s) = (0, 0) to (1, 1) {
+ gpu.launch blocks(%bx, %by, %bz) in (%gx = %c1, %gy = %c1, %gz = %c1)
+ threads(%tx, %ty, %tz) in (%sx = %c1, %sy = %c1, %sz = %c1) {
+ affine.parallel (%p, %q) = (0, 0) to (2, 2) {
+ affine.for %i = 1 to 5 {
+ affine.if affine_set<(d0) : (d0 - 1 >= 0)>(%i) {
+ affine.for %j = 0 to 3 {
+ %v = affine.load %a[%p, %q, %i - 1, %j + 1] : memref<2x2x6x4xf32>
+ affine.store %v, %a[%p, %q, %i, %j] : memref<2x2x6x4xf32>
+ }
+ }
+ }
+ }
+ gpu.terminator
+ }
+ }
+ return
+}
More information about the Mlir-commits
mailing list