Domain stream: the EBM Langevin loop through the MLIR/LLVM backbone — CPU JIT and one cooperative kernel on gfx1151/gfx1201/sm_120 (W4-PRODUCT-1 / AD-SOLVER-IFT-1)

[gstoner]

Summary

Domain stream, fourth and fifth slices, owner W4-PRODUCT-1 / AD-SOLVER-IFT-1. Sync keys EBM-NATIVE-QUADRATIC-2026-09-16 (CPU lane) and EBM-NATIVE-GPU-2026-09-16 (device lane). The GA/EBM review's first acceptance clause for "an energy is a typed program" is met on the CPU lane and on gfx1151, gfx1201 and sm_120: the energy is a Graph IR function, its gradient is the compiler's, and the K-step Langevin loop with on-device Philox noise runs as one native call on the CPU and as one cooperative kernel on each GPU, bit-exact with the declared policy everywhere.

CPU lane (fourth slice)

• The energy is a program the compiler differentiates. E(y, x) = 0.5·Σ(x − y)² is a Graph IR function marked tessera.autodiff = "reverse"; the paired autodiff pass derives @E__bwd; nothing hand-writes the gradient.
• The EBM dialect's first lowering. tessera-ebm-lower-langevin turns energy / inner_step / langevin_step (euclidean) into arith + linalg over that gradient with Philox-4x32-10 / Box-Muller noise in a linalg.generic, so a K-step scf.for compiles as one function. langevin_step gains a variadic captures operand. Non-euclidean manifolds, a missing @E__bwd and mistyped operands fail closed. The RNG policy is stated in the pass header and mirrored bit-for-bit by native_langevin.reference_langevin_loop.
• Consumers. native_langevin.{native_langevin_loop, native_quadratic_energy, package_ebm_langevin_cpu}, row cpu / cpu_ebm_langevin_llvm_jit.
• Three shared-compiler gaps closed on the way: tessera.sub had no adjoint; tessera.unsqueeze / tessera.broadcast (the sum-reduce adjoint's output) had no linalg lowering; the JIT's DPS out-param rewrite did not follow intra-module call sites (and must not re-pair already-paired modules).

Device lane (fifth slice) — the tensor-level gradient lowered inside a kernel

• One driver runs the chain (T1). tessera-opt registers the EBM and Clifford dialects and passes when built with them (TESSERA_HAVE_EBM / TESSERA_HAVE_CLIFFORD) plus convert-elementwise-to-linalg; lit gains a tessera-ebm feature. Chain: --tessera-autodiff-paired --tessera-ebm-canonicalize --tessera-ebm-lower-langevin --tessera-to-linalg --inline --convert-elementwise-to-linalg --canonicalize --cse --tessera-row-program-to-gpu.
• The row-program emitter (G2), src/transforms/lib/RowProgramToGPUPass.cpp. A new core pass over upstream dialects: a [rows, features] row program (parallel linalg.generic bodies, feature-axis linalg.reduce, small uniform integer vectors, scf.for) becomes one gpu.func — one block per row, one lane per feature (≤ 1024), the K-step loop carried in registers (state, gradient and noise never leave them), Philox inside the loop, ordered shared-memory reductions (sequential over the feature index, so reduced quantities are bit-exact with a host f32 fold), guarded !llvm.ptr<1> loads/stores, the arena marker, and tessera.row_program.* provenance. Calls are refused with "inline every call before lowering"; anything outside the envelope names the op. The scoped schedule.ebm_langevin / tile.langevin_kernel contract was not needed: after lowering, the loop function already is the scalar body, and one pass maps it (Decision Apple GPU: Gumbel-max inference sampler (#18-safe, reproducible) #31 — no second scalar emitter). G1 serial residency stays scoped only as the fallback outside the envelope.
• Packaging seam. compiler/native_row_program.py (compile → tensor contract → build_native_gpu_storage with replay check → cached HostArrayProgram) is what any [rows, features] domain program packages through; native_host_program.py factors the host-array runner out of the Clifford route (CliffordDeviceProgram is now a subclass); the contract producer merges into a module that already carries compiler provenance. Rows rocm / rocm_ebm_langevin_native_compiled and nvidia_sm120 / nvidia_ebm_langevin_native_compiled; package_ebm_langevin_native.
• Reduction proof. The quadratic gradient is elementwise, so the Langevin kernel carries no reduction; a row-normalization program (x / sqrt(Σ_f x²), F up to 1024) proves the emitter's ordered fold bit-exact with the sequential f32 fold on all three devices, host-free in phase_f5/row_program_to_gpu_reduce.mlir (with entry-missing and too-wide refusals and both backends' sqrt lowering).

Four defects only the assertions driver and sm_120 could show

1. --inline builds its interface collection over every loaded dialect; the LLVM dialect promises a DialectInlinerInterface that registerAllExtensions does not provide → tessera-opt now calls LLVM::registerInlinerInterface.
2. The pass parsed a skeleton naming tile.alloc_shared without declaring the tile dialect ("Loading a dialect (tile) while in a multi-threaded execution context").
3. Its tessera.* provenance attributes load the Tessera dialect, which an op-free input (the reduction fixture) never loaded.
4. On the NVVM arena route math.sqrt lowers to libdevice's __nv_sqrtf, whose precise branch is gated on __CUDA_PREC_SQRT, which MLIR never sets (LLVM 23 exposes only nvvm-reflect-ftz): the kernel ran MUFU.SQRT and one row of the reduction proof was 1 ulp off on sm_120 while div was div.rn and both RDNA parts were exact. convert-gpu-to-nvvm outlaws llvm.intr.sqrt, so the emitter calls libdevice's rounding-explicit __nv_fsqrt_rn on NVIDIA and pins the intrinsic on ROCm. Every other libdevice f32 path on that route is recorded as unmeasured (docs/audit/backend/nvidia/todo.md).

Items 1–3 ran green on every NDEBUG driver in the fleet (Decision #19's standing lesson, third instance, recorded in CLAUDE.md).

Evidence

• Packets benchmarks/baselines/ebm_langevin_native_gpu_20260916/ (README, source fingerprints): six (shape, K, T) rows per device — K ∈ {1,4,5,8,12}, F ∈ {5,8,33,100,1024}, T ∈ {0,0.3,0.5,0.7} — worst abs error 0 on gfx1151 (Princess-Luna), gfx1201 (Tajasarus, TESSERA_ROCM_CHIP=gfx1201, assertions-ON driver) and sm_120 (Super-Bear), the device f64 log/cos included; kernel structure read from the arena IR (one gpu.func, one scf.for carrying (f32, i64), 18 Philox mului_extended, no linalg/tensor op).
• Device tests tests/unit/test_ebm_native_langevin_gpu.py, all unsandboxed: 16/16 on Princess-Luna, 16/16 on Tajasarus under the assertions compiler, 16/16 on Super-Bear (both build/ and the reconfigured build-nvidia-cuda/, which had been configured without the EBM/Clifford backends). Host-free half on the Mac (chain shape, arena replay for both backends, envelope refusals, both fixtures).
• Lit: phase2_autodiff/row_program_to_gpu_langevin.mlir + phase_f5/row_program_to_gpu_reduce.mlir pass on the Mac, both ROCm boxes and both Super-Bear trees; phase2_autodiff + phase8 + phase_f5 under Tajasarus's assertions driver 61 pass / 55 unsupported / 0 fail.
• Mac gates: drift/registry/matrix/audit suites 543 passed, plan gate, generated docs, audit docs, domain routing, mypy, ruff green. CPU-lane evidence as before (Mac 11/11; Princess-Luna, Super-Bear, Tajasarus EBM lit 14/14; Super-Bear bf16 JIT failures pre-existing, bisected against main).

Records

EBM_NATIVE_LOOP_ARCHITECTURE.md (G2 as shipped and why the Tile contract was not needed, T1, measured risks), log entry, plan gate + Latest, GA/EBM review, domain audit (EBM ladder row gains rocm and nvidia_sm120 from the matrix), four backend todos under EBM-NATIVE-GPU-2026-09-16, CLAUDE.md.

Not claimed / next

No performance claim — per-call host transfers; the Python-emitted rocm_ebm_langevin_compiled / x86_ebm_langevin_compiled / Apple kernels remain the lanes until the dispatch/allocation/traffic/kernel-time comparison is recorded with route and latency_source. Rows wider than 1024 features, data-dependent control flow and cross-row-coupled energies fail closed. Nonlinear (N1) and manifold (M1 sphere, M2 bivector) energies are the next slices and are row programs the same emitter maps; opaque callbacks keep the reference path; no Apple package route (the Apple arena lane is MSL).

🤖 Generated with Claude Code

[chatgpt-codex-connector]

You have reached your Codex usage limits for code reviews. You can see your limits in the Codex usage dashboard.
To continue using code reviews, you can upgrade your account or add credits to your account and enable them for code reviews in your settings.