EnzymeAD/Enzyme-JAX

ConvertParallelToGPU1 pass crashes

Offen

#2.358 geöffnet am 30.03.2026

 (1 Kommentar) (0 Reaktionen) (3 zugewiesene Personen)MLIR (45 Forks)auto 404
help wanted

Repository-Metriken

Stars
 (126 Sterne)
PR-Merge-Metriken
 (PR-Metriken ausstehend)

Beschreibung

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
  }
}

Contributor Guide