EnzymeAD / EnzymeAD/Enzyme-JAX

Tracking: MFEM CUDA raising to xla-gpu

Open
#2,968 55 comments 0 reactions 0 assignees View on GitHub
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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.