EnzymeAD / EnzymeAD/Enzyme-JAX

Raising miscompiles mfem's shared-memory simplex diffusion apply (thread-0-guarded sections yield zeros)

Open
#3,003 1 comment 0 reactions 0 assignees View on GitHub
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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.