[Mlir-commits] [mlir] Codex/affine parallel scope depth (PR #228919)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Sun Oct 4 10:01:27 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir
Author: hanbeom (ParkHanbum)
<details>
<summary>Changes</summary>
[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
Assisted-by: GPT 6
---
Full diff: https://github.com/llvm/llvm-project/pull/228919.diff
2 Files Affected:
- (modified) mlir/lib/Dialect/Affine/Analysis/AffineAnalysis.cpp (+11-2)
- (added) mlir/test/Dialect/Affine/parallelize-affine-scope.mlir (+132)
``````````diff
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
new file mode 100644
index 0000000000000..8e31540ae175b
--- /dev/null
+++ b/mlir/test/Dialect/Affine/parallelize-affine-scope.mlir
@@ -0,0 +1,132 @@
+// 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
+// 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.
+// SCOPE-LABEL: func.func @scope_depth
+// SCOPE: affine.for {{.*}} = 0 to 1 {
+// SCOPE-NEXT: gpu.launch
+// 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
+ 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.
+// OUTPUT-LABEL: func.func @scope_depth_output_dependence
+// OUTPUT: affine.for {{.*}} = 0 to 1 {
+// OUTPUT-NEXT: gpu.launch
+// 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
+ 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
+}
+
+// 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
+}
``````````
</details>
https://github.com/llvm/llvm-project/pull/228919
More information about the Mlir-commits
mailing list