EnzymeAD/Enzyme-JAX

ConvertParallelToGPU1 pass crashes

Aberta

#2.358 aberto em 30 de mar. de 2026

 (1 comentário) (0 reação) (3 responsáveis)MLIR (45 forks)auto 404
help wanted

Métricas do repositório

Stars
 (126 estrelas)
Métricas de merge de PR
 (Métricas PR pendentes)

Description

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

Guia do colaborador