ConvertParallelToGPU1 pass crashes
help wanted
In RSBench benchmark, gpu.launch is lowered by parallel-lower into:
```
gpu_wrapper(grid0,grid1,grid2, block0,block1,block2) {
scf.parallel (3D grid) {
scf.parallel (3D block) {
enzyme.autodiff_region...
}
}
}
```
The first canonicalize pass (between parallel-lower and affine-cfg) applies MLIR general MergeNestedParallelLoops
Merge condition: outerBody has single element (the inner parallel)
→ merges 3D grid + 3D block into single 6D scf.parallel.
This is cause crashes in `ConvertParallelToGPU1` because:
- `parallelCount` for child ops drops from 2 (grid+block ancestors) to 1 (single 6D ancestor)
- `InnerParallelSerialization` requires `parallelCount` >= 2 to serialize child parallels
- With `parallelCount=1`, child parallels cannot be serialized
Later in the pipeline:
- affine-cfg (pass AffineParallelizePattern) generates `affine.for` inside `enzyme.autodiff_region` to `affine.parallel`
- lower-affine converts it to scf.parallel (child parallel)
- convert-parallel-to-gpu1 hits 3-way deadlock:
1. InnerParallelSerialization: parallelCount=1 < 2 → return failure()
2. SplitParallelOp: pop->walk() finds child parallel → child=true → return failure()
3. ParallelToGPULaunch: getDirectlyNestedSingleParallel skips allocas, hits pointer2memref → blockPop=nullptr → return failure()
There are three categories of gpu_wrapper in the dumped mlir before pass ConvertParallelToGPU1:
Type A (1 wrapper): 3D+3D with enzyme.autodiff_region containing child parallel.
→ After canonicalize merge: 6D + child → failes
Type B (2 wrappers): 3D+3D without child parallel.
→ After canonicalize merge: 6D → failed -> SplitParallelOp exactMatch fallback → OK
Type C (24 wrappers): already 3D+3D -> OK
Causes applyPatternsGreedily to exceed its iteration limit → signalPassFailure().
The minimal reproducible example is as following:
```
// Run: POLYGEIST_GPU_KERNEL_BLOCK_SIZE=256 enzymexlamlir-opt --pass-pipeline="builtin.module(canonicalize,convert-parallel-to-gpu1)" %s
module {
func.func @mre(%d0: index, %d1: index, %d2: index,
%d3: index, %d4: index, %d5: index,
%g0: index, %g1: index, %g2: index,
%b0: index, %b1: index, %b2: index,
%mem: memref<?xf64>) {
%c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index
%c1_i32 = arith.constant 1 : i32
%c1_i64 = arith.constant 1 : i64
%c4 = arith.constant 4 : index
%cst = arith.constant 0.0 : f64
// TYPE A: broken
// 3D grid + 3D block with enzyme.autodiff_region containing child scf.parallel.
%a = "enzymexla.gpu_wrapper"(%d0, %d1, %d2, %d3, %d4, %d5) ({
scf.parallel (%bx, %by, %bz) = (%c0, %c0, %c0) to (%d0, %d1, %d2) step (%c1, %c1, %c1) {
scf.parallel (%tx, %ty, %tz) = (%c0, %c0, %c0) to (%d3, %d4, %d5) step (%c1, %c1, %c1) {
%alloca = memref.alloca() : memref<1xf64>
%ptr0 = llvm.alloca %c1_i32 x !llvm.array<4 x f64> {alignment = 8 : i64} : (i32) -> !llvm.ptr
%ptr1 = llvm.alloca %c1_i64 x !llvm.struct<"S", (i32)> {alignment = 4 : i64} : (i64) -> !llvm.ptr
%as_memref = "enzymexla.pointer2memref"(%ptr1) : (!llvm.ptr) -> memref<?xi32>
memref.store %cst, %alloca[%c0] : memref<1xf64>
enzyme.autodiff_region(%ptr0) {
^bb0(%arg: !llvm.ptr):
scf.parallel (%j) = (%c0) to (%c4) step (%c1) {
memref.store %cst, %mem[%j] : memref<?xf64>
scf.reduce
}
enzyme.yield
} attributes {activity = [#enzyme<activity enzyme_const>], ret_activity = []} : (!llvm.ptr) -> ()
scf.reduce
}
scf.reduce
}
"enzymexla.polygeist_yield"() : () -> ()
}) : (index, index, index, index, index, index) -> index
// TYPE B: 3D+3D, no child parallel
%b1_ = "enzymexla.gpu_wrapper"(%d0, %d1, %d2, %d3, %d4, %d5) ({
scf.parallel (%bx, %by, %bz) = (%c0, %c0, %c0) to (%d0, %d1, %d2) step (%c1, %c1, %c1) {
scf.parallel (%tx, %ty, %tz) = (%c0, %c0, %c0) to (%d3, %d4, %d5) step (%c1, %c1, %c1) {
memref.store %cst, %mem[%tx] : memref<?xf64>
scf.reduce
}
scf.reduce
}
"enzymexla.polygeist_yield"() : () -> ()
}) : (index, index, index, index, index, index) -> index
%b2_ = "enzymexla.gpu_wrapper"(%d0, %d1, %d2, %d3, %d4, %d5) ({
scf.parallel (%bx, %by, %bz) = (%c0, %c0, %c0) to (%d0, %d1, %d2) step (%c1, %c1, %c1) {
scf.parallel (%tx, %ty, %tz) = (%c0, %c0, %c0) to (%d3, %d4, %d5) step (%c1, %c1, %c1) {
memref.store %cst, %mem[%ty] : memref<?xf64>
scf.reduce
}
scf.reduce
}
"enzymexla.polygeist_yield"() : () -> ()
}) : (index, index, index, index, index, index) -> index
return
}
}
```
0 条评论