EnzymeAD / EnzymeAD/Enzyme-JAX
Raising miscompiles mfem's shared-memory simplex diffusion apply (thread-0-guarded sections yield zeros)
- Dominant language
- MLIR
- Stars
- 131
- Forks
- 53
- Avg merge
- 1d 10h
- Merged PRs (30d)
- 193
Description
The raised `mfem::internal::SmemPADiffusionApplyTetrahedron<2,2>` (mfem `fem/integ/bilininteg_diffusion_pa_simplices.hpp`) compiles strict but computes an **all-zero output** at runtime (beam-tet order 1, Q1D=2; the <2,1> single-point instantiation is exact, as are the 2D triangle kernels).
Post-pipeline module: https://gist.github.com/wsmoses/e438453754885ba06b6be3d81815477a
The suspicious shape in the raised IR:
```mlir
%114 = stablehlo.reduce(%93 ...) // tid.x + tid.y
%115 = stablehlo.compare LT, %114, %c_56 // thread-0 guard: tid sum < 1
...
%126 = stablehlo.select %115, %125#7, %cst_20 // carried 2x4xf64 accumulators ...
%127 = stablehlo.select %115, %125#6, %cst_20 // ... become ZERO on non-zero lanes
```
The kernel is a hard case: `MFEM_FOREACH_THREAD_DIRECT` sections, `MFEM_SHARED` scratch **unioned under several shapes** (`sm0`/`sm1` reinterpreted as X/DDQ/DQQ/QQQ/QQD/QDD views across phases), and thread-0-only writes. The masked raise appears to zero the loop-carried accumulators for lanes outside the guard rather than preserving them.
Reproduction: mfem `PA Simplices` GPU unit test through the raising path (#2968); a minimal driver is
```c++
// beam-tet.mesh, H1 order 1 BasisType::Positive, DiffusionIntegrator with
// StroudIntRules order 2p+OrderW (nq=8): AddMultPA returns y == 0.
```
Worked around in the mfem harness by dispatching to the per-element fallback (`PADiffusionApplyTetrahedron`) under the XLA backend, which raises and runs bit-correct — so this is not blocking, but the shared-memory shape-punned scratch + direct-thread-section pattern will recur (mfem uses it widely).
🤖 Generated with [Claude Code](https://claude.com/claude-code)
https://claude.ai/code/session_016zErYp7upmqr4NHfhod9UD
Contributor guide
No contributing guide indexed for this repository
Research direction
Start with mfem/fem/integ/bilininteg_diffusion_pa_simplices.hpp and reproduce the beam-tet order-1 case through the PA Simplices GPU unit test and raising path from #2968. Inspect the raised IR around the thread-0 guard, loop-carried accumulators, and shared-memory views, then compare against the per-element PADiffusionApplyTetrahedron fallback. Done means the shared-memory tetrahedron instantiation produces the expected nonzero output through XLA.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- cpp
- Domain
- compilers
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 48/100