EnzymeAD / EnzymeAD/Enzyme-JAX
Tracking: MFEM CUDA raising to xla-gpu
- Dominant language
- MLIR
- Stars
- 131
- Forks
- 53
- Avg merge
- 1d 10h
- Merged PRs (30d)
- 193
Description
Tracking issue for making the Reactant/EnzymeMLIR raising path a drop-in clang replacement for MFEM CUDA targeting the `xla-gpu` backend: every CUDA translation unit either raises fully to stablehlo (compiled to XLA executables, launched through the `reactantXLA*` runtime) or fails loudly at compile time.
## Where things stand
**Both MFEM mass TUs raise completely** (`fem/integ/bilininteg_mass_pa.cpp`: 10/10 kernels; `fem/integ/bilininteg_mass_kernels.cpp`: 129/129 kernels, 255 whiles) in strict mode with the open raising stack applied. Both compile end-to-end through the clang plugin, `libmfem_xla.a` and the unit-test binary link, XLA initializes, and the allocator TU (`general/cuda.cpp`) fully intercepts the CUDA runtime (no raw `cudaMalloc` left).
**Merged so far** (infrastructure + first half of the raising features):
- #2930 pointer-constant splat gate; #2926, #2928, #2931, #2932, #2934, #2942 (import-pipeline fixes)
- #2943 `xla_wrapper` scalar specialization + unified 6-arg `reactantXLAExec` (runtime side: Reactant.jl#3220)
- #2944 dynamic-extent parallel axes peeled into `enzymexla.parallel`-tagged whiles (scoped to raised regions)
- #2954 sequential `affine.for` over parallel bodies as `stablehlo.while` (symbolic bounds, buffer carrying)
- #2955 `memref.alloca` scratch as zero-splat tensors, loop-carried
- #2957 `affine.if` integer sets with symbols
- #2958 alloca-under-conditionals tests; #2959 barrier no-op; #2962 uniform buffer select
- #2963 **compile-time guard**: `convert-polygeist-to-llvm{backend=xla-*}` errors on any `gpu.launch_func` that survived raising, instead of emitting a binary that crashes on the two-worlds memory problem
- Reactant#86 (kernel-less TU round-trip), Reactant#87 (`REACTANT_BACKEND` / `REACTANT_XLA_BACKEND` predefines), Enzyme#3166/#3171/#3157
## Open PRs
**Raising stack** (required for the mass TUs to fully raise; stacked, base → tip):
- #2960 broadcast-store mask or-reduce (independent, on main)
- #2945 viewed shared-memory scratch flattening (independent, on main)
- #2948 scatter-mask or-reduce (stacked on the union spine)
- #2949 scalar control flow: scf loops, per-lane trip counts, racy writes (largest; carries the release-mode UB hardening)
- #2950 do-while rotation + guarded racy writes — with this, both TUs raise
**Reduction-kernel set** (from probing `linalg/vector.cpp` under `xla-gpu`; all independent, on main). MFEM's `general/reducers.hpp` two-phase reduction (`mfem::reduce`, used by every `Vector::Min/Max/Sum/Dot/Norml2`) was the next blocker class:
- #2964 map `llvm.minimumnum`/`maximumnum` `call_intrinsic`s to `arith.minnumf/maxnumf` in libdevice raising
- #2965 `MoveWhileToFor`: accept the duplicated increment `enzyme-lift-cf-to-scf` emits when synthesizing the multi-exit → single-condition conjunction (no `cse` runs between lifting and loop canonicalization)
- #2966 fold the carried `i1` live flag (residual `i < N` early-exit) into a pure induction-variable predicate: `iv == lb || f(cnt0 + iv - lb - 1) < N`
- #2967 convert the rotated halving do-while (`for (i = bs>>1; i > 0; i >>= 1)` tree reduction) to a log-space `scf.for` via `ctlz`; the existing `WhileShiftToInduction` only matches the canonical body-in-after form
**Output-quality** (not required for raisability):
- #2941 + #2946 loop-bound folds: delete the dead grid/block-stride residual loops clang's -O3 peel/rotate leaves in raised kernels
## Library coverage
Of 273 objects in `libmfem_xla.a`, 113 TUs still contain native kernel launches — but nearly all of those are the plugin's own cuda-backend output (round-tripped forall kernels), i.e. they become raise-or-error worklist items simply by recompiling with `backend=xla-gpu` under the #2963 guard. The one genuinely different class is **cub**: `general/scan.hpp` instantiates `cub::DeviceScan/DeviceSelect` whose internally-launched kernels cannot raise. That is now handled at the source level via Reactant#87's `REACTANT_XLA_BACKEND` predefine — MFEM's `scan.hpp` gates the cub path and takes its existing serial fallbacks under xla backends.
## Runtime blockers (next phase)
The **two-worlds memory problem**: `reactantXLAMalloc` returns host pointers into the runtime's buffer bookkeeping (PJRT buffers are opaque objects), while unraised CUDA kernels need raw device addresses; each world crashes on the other's pointers. Plan of record:
- make raw device pointers the single currency (PJRT `OpaqueDeviceMemoryDataPointer` + external buffer views)
- `input_output_aliasing` so XLA executions write in place at stable addresses (MFEM's memory manager assumes imperative semantics)
- implement device-to-device `reactantXLAMemcpy` (currently `llvm_unreachable`; MFEM needs it)
Also pending: rebuild `libReactantExtra` at merged Reactant.jl (the local build predates the 6-arg `reactantXLAExec` rename) and recompile the TU objects with the current plugin.
## Next steps
1. Land the raising stack (#2960, #2945, #2948, #2949, #2950) and the reduction set (#2964–#2967); re-probe `vector.cpp` — expected to fully raise.
2. Sweep the remaining library TUs with `backend=xla-gpu`; triage the guard errors per kernel family.
3. Raw-device-pointer currency in the runtime; then the PA Mass runtime test end-to-end on `xla-gpu`.
Contributor guide
No contributing guide indexed for this repository
Research direction
Start by reading the open raising PRs #2960, #2945, #2948, #2949, and #2950, then probe the listed MFEM translation units with backend=xla-gpu. Review general/reducers.hpp, general/scan.hpp, and general/cuda.cpp alongside the reduction issues #2964–#2967. Done means remaining CUDA translation units either raise fully to stablehlo or fail at compile time, followed by the raw-pointer runtime work and PA Mass end-to-end test.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- cpp
- Domain
- ai-infra-agents, backend, compilers
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 25/100