Fix AMD wave barrier in HIP kernels; add NequIP-OAM-L test problems

[asglover]

Fixes #211 (see also mir-group/nequip#610).

Problem

On AMD (MI250X / gfx90a via ROCm), the weight gradient of NequIP-OAM-L's 23-instruction main-layer TensorProductConv was wrong in float32 atomic mode (relative error ~0.5 on one instruction, 32x2e x 1x3o -> 32x3o uvu). Energies were unaffected, so the symptom was non-conservative forces in MD. CUDA, the other NequIP-OAM sizes, float64, and deterministic mode were all fine.

Root cause (found by @kavanase on Frontier): on HIP, __syncwarp() was rewritten to __threadfence_block(), which is a memory fence, not an execution barrier. The compiler was free to reorder shared-memory accesses across it inside one unrolled backward segment.

Fix

jinja_utils.py: on HIP, syncwarp now expands to a release fence, __builtin_amdgcn_wave_barrier(), and an acquire fence, keeping the memory-ordering semantics and adding a real wavefront barrier.

Test

The failing tensor product was not in the test suite. This PR adds:

• NequIPTPP in benchmark/problems.py, taken from nequip's InteractionBlock. It emits the same instruction set as ChannelwiseTPP but in nequip's original order. ComputationSchedule never sorts instructions and greedily cuts segments in the given order, so the order determines segment boundaries and the generated kernel. Using nequip's order reproduces the exact kernel that miscompiled.
• nequip_oam_problems(): the three distinct NequIP-OAM-L problems (first layer, main layers, last layer), built from the L preset in nequip/model/nequip_models.py. Instruction 16 of the main-layer problem is 32x2e x 1x3o -> 32x3o uvu, matching the reporter's breakdown.
• Both are added to TestProductionModels in batch_test.py and conv_test.py. The conv suite (atomic mode, F32) is the one that reproduces the report.

Status

• Not yet verified on AMD hardware. The new test cases are expected to fail on gfx90a before the jinja_utils.py change and pass after; that has not been run here.
• Verify no regression on CUDA.

🤖 Generated with Claude Code