diff --git a/CLAUDE.md b/CLAUDE.md index aab380ff9..0294362f4 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -56,7 +56,7 @@ x86 AMX/AVX512, Apple M-series CPU/GPU. **Execution reality (updated 2026-09-15):** the **x86 AVX-512** backend and **Apple CPU (Accelerate) + GPU (MPS/MSL/MPSGraph + Metal 4 MPP `matmul2d`)** -backends execute natively. **ROCm** executes on two RDNA parts: **gfx1151** +backends execute natively. Since 2026-09-16 the EBM Langevin loop's compiler-derived gradient runs *inside* one cooperative kernel on gfx1151, gfx1201 and sm_120 through the row-program emitter (`tessera-row-program-to-gpu`), bit-exact with the CPU JIT lane's policy — a correctness lane, not a promotion. **ROCm** executes on two RDNA parts: **gfx1151** (Strix Halo, RDNA 3.5, `Princess-Luna`) has broad native execution across the attention family, norms/activations, matmul compositions, MoE transport, SSM, EBM and warp-shuffle lanes; **gfx1201** (RX 9070 XT, RDNA4, `Tajasarus`, added @@ -434,6 +434,23 @@ Per-phase deliverables and the open-work priority queue live in falsifier" clause above is history. Route MLIR promise/contract claims to Tajasarus (or any assertions build) before recording "does not reproduce". + **Third instance (2026-09-16): the EBM row-program chain ran green on every + NDEBUG driver in the fleet and aborted twice on Tajasarus's assertions-ON + driver** — `--inline` builds its interface collection over every loaded + dialect and the LLVM dialect *promises* a `DialectInlinerInterface` that + `registerAllExtensions` does not provide (`tessera-opt` now calls + `LLVM::registerInlinerInterface`), and a pass that parses a textual + skeleton naming `tile.alloc_shared` had not declared the tile dialect in + `getDependentDialects` ("Loading a dialect (tile) while in a + multi-threaded execution context"), nor the Tessera dialect its + `tessera.*` provenance attributes load. All one-line fixes; all invisible + without that box. (Separately, sm_120 showed the NVVM arena route's + `math.sqrt` reaches libdevice's *approximate* `__nv_sqrtf` because MLIR + never sets the precise-sqrt reflect flag, and `convert-gpu-to-nvvm` + outlaws `llvm.intr.sqrt` — call libdevice's rounding-explicit + `__nv_fsqrt_rn` when a result must be IEEE; `docs/audit/backend/nvidia/todo.md`.) Run every new pass through Tajasarus before recording + "passes" (sync `EBM-NATIVE-GPU-2026-09-16`). + **Separately, still open: `TileToX86Pass` loads `tessera_x86` from inside `runOnOperation()`** (`src/transforms/lib/TileToX86Pass.cpp:1045`, a by-name `getOrLoadDialect` used to avoid linking the optional backend). MLIR forbids @@ -932,9 +949,13 @@ ninja -C build tessera-opt # 32 threads; ~1-2 min cold # Since 2026-09-16 the Clifford backend is also what puts the geometric-algebra # lane into libtessera_jit (tessera_jit_has_clifford()) and builds # ts-clifford-opt, which the native GPU route (native_clifford_gpu.py) shells -# out to; every WSL box's `build/` (Princess-Luna, Super-Bear, Tajasarus) and -# the Mac configure it ON, so domain fixtures and the GA lanes are -# fleet-covered, not Mac-only. +# out to, and the EBM backend puts the energy lane in (tessera_jit_has_ebm(); +# paired autodiff + tessera-ebm-lower-langevin). Every WSL box's `build/` +# (Princess-Luna, Super-Bear, Tajasarus, plus Tajasarus's build-assertions) +# and the Mac configure both ON, so domain fixtures and the GA/EBM lanes are +# fleet-covered, not Mac-only. bf16 JIT tests fail on Super-Bear (Zen 2 has +# no AVX512-BF16; unresolved `_mlir_ciface_*` symbols) independent of any +# branch — bf16 JIT proof belongs on the Zen 5 hosts. cmake -S . -B build -G Ninja -DTESSERA_BUILD_EBM_BACKEND=ON -DTESSERA_BUILD_CLIFFORD_BACKEND=ON # Re-verify a C++ pass change end-to-end: rebuild → lit fixture + FileCheck → drift test diff --git a/benchmarks/baselines/ebm_langevin_native_gpu_20260916/README.md b/benchmarks/baselines/ebm_langevin_native_gpu_20260916/README.md new file mode 100644 index 000000000..13f0c63b4 --- /dev/null +++ b/benchmarks/baselines/ebm_langevin_native_gpu_20260916/README.md @@ -0,0 +1,44 @@ +# The EBM Langevin loop as one cooperative kernel, on the owning devices + +Owner: W4-PRODUCT-1 / AD-SOLVER-IFT-1 (sync `EBM-NATIVE-GPU-2026-09-16`). +`record_ebm_langevin_native_gpu.py` executed independently on three owning +devices; `source-hashes.json` binds the implementation. Packets are +correctness evidence; no performance is measured or claimed, and gfx1151 / +gfx1201 / sm_120 proofs never transfer. + +| Packet | Host | Device | Rows | Worst abs error | +|---|---|---|---|---| +| `rocm_gfx1151.json` | Princess-Luna (WSL2, ROCm 10.0) | gfx1151 | 6 | 0 | +| `rocm_gfx1201.json` | Tajasarus (WSL2, ROCm 10.0, `TESSERA_ROCM_CHIP=gfx1201`, assertions-ON driver) | gfx1201 | 6 | 0 | +| `nvidia_sm120.json` | The-Super-Bear (WSL2, CUDA 13.4 / driver 610.88) | sm_120 (RTX 5070) | 6 | 0 | + +Each row is one K-step loop over the quadratic energy `0.5·Σ(x − y)²` at one +`(shape, K, T)` — `(4×8, 1, 0.7)`, `(6×8, 5, 0.7)`, `(3×5, 12, 0.7)`, +`(9×100, 4, 0.3)`, `(16×33, 8, 0)`, `(2×1024, 3, 0.5)` — compared with +`native_langevin.reference_langevin_loop`, the numpy statement of the declared +Philox-4x32-10 / Box–Muller policy, and the returned key. `kernel_structure` +is read from the arena IR of the packaged kernel: one `gpu.func`, one +`scf.for` carrying `(state f32, key i64)` in registers (the loop-invariant key +word is hoisted by canonicalization), 18 `arith.mului_extended` Philox rounds +inside the loop, no barrier (the quadratic gradient is elementwise), and no +`linalg`/`tensor` op left. + +How the kernel is produced (no Python-emitted arithmetic): one `tessera-opt` +invocation runs `--tessera-autodiff-paired` (the gradient), +`--tessera-ebm-canonicalize --tessera-ebm-lower-langevin` (the step with the +noise), `--tessera-to-linalg --inline --convert-elementwise-to-linalg +--canonicalize --cse` (a `[rows, features]` row program in linalg) and +`--tessera-row-program-to-gpu` (one block per row, one lane per feature, +the K-step loop in registers, ordered shared-memory reductions); the tensor +contract is attached, `build_native_gpu_storage` packages it and the native +storage binding launches it. The emitter's reduction path is exercised by the +row-normalization device test in `tests/unit/test_ebm_native_langevin_gpu.py` +(bit-exact with the sequential f32 fold on the same three devices), not by +these rows. + +Reproduce on an owning host (EBM backend configured ON, toolkit env sourced): + +```bash +PYTHONPATH=python:. python benchmarks/record_ebm_langevin_native_gpu.py --backend rocm --chip gfx1151 \ + --compiler build/tools/tessera-opt/tessera-opt --output /rocm_gfx1151.json +``` diff --git a/benchmarks/baselines/ebm_langevin_native_gpu_20260916/nvidia_sm120.json b/benchmarks/baselines/ebm_langevin_native_gpu_20260916/nvidia_sm120.json new file mode 100644 index 000000000..bf16fa486 --- /dev/null +++ b/benchmarks/baselines/ebm_langevin_native_gpu_20260916/nvidia_sm120.json @@ -0,0 +1,114 @@ +{ + "schema": 1, + "backend": "nvidia", + "chip": "sm_120", + "host": "The-Super-Bear", + "compiler_sha256": "7715c2cae3b62acda99052651b84eaeaa73b3ad317deaa0a2586af07ef6543f9", + "llvm_bin": "/usr/lib/llvm-23/bin", + "recorder_sha256": "8d1a24427a4d9a8395aaffa4b6a4c2c69c0c983008dd0a6d54b851142c508ced", + "rows": [ + { + "shape": [ + 4, + 8 + ], + "steps": 1, + "temperature": 0.7, + "max_abs_error": 0.0, + "next_key": [ + 20017429977207, + 43 + ], + "binding": "ff3f26a07d47eb4b4791c6db506488898dfb84f45f9c3b317910b7ce79858157" + }, + { + "shape": [ + 6, + 8 + ], + "steps": 5, + "temperature": 0.7, + "max_abs_error": 0.0, + "next_key": [ + 20017429977211, + 47 + ], + "binding": "7f8adc6b089de3ff02502af88e2c6c87e062e2291e3acd23826e37877e9f6e5c" + }, + { + "shape": [ + 3, + 5 + ], + "steps": 12, + "temperature": 0.7, + "max_abs_error": 0.0, + "next_key": [ + 20017429977218, + 54 + ], + "binding": "c4d03d8d12909999d1a181b0cd6602bfb228f2e54382fac0e7ccfc15bc2ec15c" + }, + { + "shape": [ + 9, + 100 + ], + "steps": 4, + "temperature": 0.3, + "max_abs_error": 0.0, + "next_key": [ + 20017429977210, + 46 + ], + "binding": "82184df27e9b8e926c06630d60c2692044cabd96681cb71f7847e68340705d36" + }, + { + "shape": [ + 16, + 33 + ], + "steps": 8, + "temperature": 0.0, + "max_abs_error": 0.0, + "next_key": [ + 20017429977214, + 50 + ], + "binding": "dee9be1aac695b75ad94620852cf7b1e9ef85b0a6ae417b3f26b67e415a85d6f" + }, + { + "shape": [ + 2, + 1024 + ], + "steps": 3, + "temperature": 0.5, + "max_abs_error": 0.0, + "next_key": [ + 20017429977209, + 45 + ], + "binding": "85a7eb9f04b41abc4e2a1cd4a608e36b25bd498d2630e393659075e25374df3b" + } + ], + "kernel_structure": { + "gpu_funcs": 1, + "loops": 1, + "loop_carried": [ + "f32, i64" + ], + "philox_mul": 18, + "barriers": 0, + "linalg_or_tensor_ops": 0 + }, + "proofs": [ + "the K-step loop is one device launch, bit-exact with the declared Philox/Box-Muller policy for every case (K in 1..12, F in 5..1024, T in {0, 0.3, 0.5, 0.7})", + "the gradient is the paired autodiff pass's own adjoint, lowered inside the kernel (no host gradient, no host noise, no Python-emitted kernel)", + "one gpu.func, one scf.for carrying (state, key, key) in registers, Philox inside the loop" + ], + "promotion_eligible": false, + "measured_performance": false, + "envelope": "f32, quadratic energy, [rows, features] with features <= 1024; per-call host transfers; correctness only", + "rocm_chip_env": null +} diff --git a/benchmarks/baselines/ebm_langevin_native_gpu_20260916/rocm_gfx1151.json b/benchmarks/baselines/ebm_langevin_native_gpu_20260916/rocm_gfx1151.json new file mode 100644 index 000000000..51de9aa72 --- /dev/null +++ b/benchmarks/baselines/ebm_langevin_native_gpu_20260916/rocm_gfx1151.json @@ -0,0 +1,114 @@ +{ + "schema": 1, + "backend": "rocm", + "chip": "gfx1151", + "host": "Princess-Luna", + "compiler_sha256": "a68c821632fa74eca218d2519a58a6749c4dba7580335d52cbe7d29dcd0c1949", + "llvm_bin": "/usr/lib/llvm-23/bin", + "recorder_sha256": "8d1a24427a4d9a8395aaffa4b6a4c2c69c0c983008dd0a6d54b851142c508ced", + "rows": [ + { + "shape": [ + 4, + 8 + ], + "steps": 1, + "temperature": 0.7, + "max_abs_error": 0.0, + "next_key": [ + 20017429977207, + 43 + ], + "binding": "abc46be93febf94155207757e2bc85fc641b8e02723d6e84c98ede2900c579d2" + }, + { + "shape": [ + 6, + 8 + ], + "steps": 5, + "temperature": 0.7, + "max_abs_error": 0.0, + "next_key": [ + 20017429977211, + 47 + ], + "binding": "80b884cb338c99a51930415b2f9553e5e5eff7a7d9b08eb88ef7c5a9909aeff3" + }, + { + "shape": [ + 3, + 5 + ], + "steps": 12, + "temperature": 0.7, + "max_abs_error": 0.0, + "next_key": [ + 20017429977218, + 54 + ], + "binding": "94dd26cd5e765ce817148e7f4957ef8e53da876710d6dd9e3fd9db08a438cfba" + }, + { + "shape": [ + 9, + 100 + ], + "steps": 4, + "temperature": 0.3, + "max_abs_error": 0.0, + "next_key": [ + 20017429977210, + 46 + ], + "binding": "d41479668d50b16cb233347b6d692e5c97cedfd389c96ecea4a6920b0b49f81d" + }, + { + "shape": [ + 16, + 33 + ], + "steps": 8, + "temperature": 0.0, + "max_abs_error": 0.0, + "next_key": [ + 20017429977214, + 50 + ], + "binding": "8881cc2b4a96b206b5d6ec9875603fd9f436f46f963bc2c53469d6805344493e" + }, + { + "shape": [ + 2, + 1024 + ], + "steps": 3, + "temperature": 0.5, + "max_abs_error": 0.0, + "next_key": [ + 20017429977209, + 45 + ], + "binding": "be4fe83e60ca9158b5f817d68f5a555afeeb949b79d9f57598393f9b21b62d0b" + } + ], + "kernel_structure": { + "gpu_funcs": 1, + "loops": 1, + "loop_carried": [ + "f32, i64" + ], + "philox_mul": 18, + "barriers": 0, + "linalg_or_tensor_ops": 0 + }, + "proofs": [ + "the K-step loop is one device launch, bit-exact with the declared Philox/Box-Muller policy for every case (K in 1..12, F in 5..1024, T in {0, 0.3, 0.5, 0.7})", + "the gradient is the paired autodiff pass's own adjoint, lowered inside the kernel (no host gradient, no host noise, no Python-emitted kernel)", + "one gpu.func, one scf.for carrying (state, key, key) in registers, Philox inside the loop" + ], + "promotion_eligible": false, + "measured_performance": false, + "envelope": "f32, quadratic energy, [rows, features] with features <= 1024; per-call host transfers; correctness only", + "rocm_chip_env": null +} diff --git a/benchmarks/baselines/ebm_langevin_native_gpu_20260916/rocm_gfx1201.json b/benchmarks/baselines/ebm_langevin_native_gpu_20260916/rocm_gfx1201.json new file mode 100644 index 000000000..8a8d63b53 --- /dev/null +++ b/benchmarks/baselines/ebm_langevin_native_gpu_20260916/rocm_gfx1201.json @@ -0,0 +1,114 @@ +{ + "schema": 1, + "backend": "rocm", + "chip": "gfx1201", + "host": "tajasarus", + "compiler_sha256": "f38e5784277e079898a3d32c71c373790eda6e6a6bc1501bac54b161630df0c8", + "llvm_bin": "/home/angstorms/.local/share/tessera-toolchains/llvm-23.1.1-assertions/bin", + "recorder_sha256": "8d1a24427a4d9a8395aaffa4b6a4c2c69c0c983008dd0a6d54b851142c508ced", + "rows": [ + { + "shape": [ + 4, + 8 + ], + "steps": 1, + "temperature": 0.7, + "max_abs_error": 0.0, + "next_key": [ + 20017429977207, + 43 + ], + "binding": "666c29d8309fda172c30a8b2dcbd3a1acaa097e094e8f841c635916c8be24b36" + }, + { + "shape": [ + 6, + 8 + ], + "steps": 5, + "temperature": 0.7, + "max_abs_error": 0.0, + "next_key": [ + 20017429977211, + 47 + ], + "binding": "139d7d1a8e3c5ef9ac99fba79cb976e85636025bb711c95b3ebdb3a1dbe803c2" + }, + { + "shape": [ + 3, + 5 + ], + "steps": 12, + "temperature": 0.7, + "max_abs_error": 0.0, + "next_key": [ + 20017429977218, + 54 + ], + "binding": "fff68289171083e8048ea3eb21f1ffd982e1868603e19ee1e213fd773fbd53c6" + }, + { + "shape": [ + 9, + 100 + ], + "steps": 4, + "temperature": 0.3, + "max_abs_error": 0.0, + "next_key": [ + 20017429977210, + 46 + ], + "binding": "1a75aa6c44c19e70808f17dea2e8847c04193f4685ee4d5acd9b87f6c8a1d89f" + }, + { + "shape": [ + 16, + 33 + ], + "steps": 8, + "temperature": 0.0, + "max_abs_error": 0.0, + "next_key": [ + 20017429977214, + 50 + ], + "binding": "b33b19daed814c4edc38a7a8d338f13e2e0f018db56faa0030bd2f58fa79e317" + }, + { + "shape": [ + 2, + 1024 + ], + "steps": 3, + "temperature": 0.5, + "max_abs_error": 0.0, + "next_key": [ + 20017429977209, + 45 + ], + "binding": "30af8b8c141e739e90ef17cf4d1b2744006431ad3e85f3eab2467ce92e425788" + } + ], + "kernel_structure": { + "gpu_funcs": 1, + "loops": 1, + "loop_carried": [ + "f32, i64" + ], + "philox_mul": 18, + "barriers": 0, + "linalg_or_tensor_ops": 0 + }, + "proofs": [ + "the K-step loop is one device launch, bit-exact with the declared Philox/Box-Muller policy for every case (K in 1..12, F in 5..1024, T in {0, 0.3, 0.5, 0.7})", + "the gradient is the paired autodiff pass's own adjoint, lowered inside the kernel (no host gradient, no host noise, no Python-emitted kernel)", + "one gpu.func, one scf.for carrying (state, key, key) in registers, Philox inside the loop" + ], + "promotion_eligible": false, + "measured_performance": false, + "envelope": "f32, quadratic energy, [rows, features] with features <= 1024; per-call host transfers; correctness only", + "rocm_chip_env": "gfx1201" +} diff --git a/benchmarks/baselines/ebm_langevin_native_gpu_20260916/source-hashes.json b/benchmarks/baselines/ebm_langevin_native_gpu_20260916/source-hashes.json new file mode 100644 index 000000000..0ce5673d1 --- /dev/null +++ b/benchmarks/baselines/ebm_langevin_native_gpu_20260916/source-hashes.json @@ -0,0 +1,12 @@ +{ + "python/tessera/compiler/native_row_program.py": "b50709a6cf41bb64d9e755e8a5660a148fdfa9df79444f9df8c01c028b96664c", + "python/tessera/compiler/native_host_program.py": "d003efc0ebe03686f3ee62147733b403c2c1aaf2120d53240be93c6cfeda7aea", + "python/tessera/compiler/native_gpu_storage.py": "d862ef1d36bca98d74d3afc20eba1f13ef89053fd773529bb64d961bdd9ad5f8", + "python/tessera/compiler/native_storage_contract.py": "75fcc10ce8c864164207a345fe9bd9b7d5d096f0f10af1a5f6c39667c60f8283", + "python/tessera/ebm/native_langevin.py": "91f78a01dd8bb819b54e5bcac61e059f6305275d625084aceffa9b48b3b73bdc", + "src/transforms/lib/RowProgramToGPUPass.cpp": "cfc732ae3d0261726d7ec0f990b1fea381f87252c6dddafb934da7d3cc78963d", + "src/solvers/ebm/lib/Passes/LowerLangevin.cpp": "0646284d848bc5a46b321976747feb24b1e620cd858dfb8fe21c6ce443135883", + "tools/tessera-opt/tessera-opt.cpp": "b1b60bb9ba0bce43a5f7b6da8aabb800a879070ca2fa82b0add9499040aa73f8", + "benchmarks/record_ebm_langevin_native_gpu.py": "8d1a24427a4d9a8395aaffa4b6a4c2c69c0c983008dd0a6d54b851142c508ced", + "tests/unit/test_ebm_native_langevin_gpu.py": "0f99d1182e9117b579227f57a4e6197ad3fe297eacdaa4defddcde25eb265c26" +} diff --git a/benchmarks/record_ebm_langevin_native_gpu.py b/benchmarks/record_ebm_langevin_native_gpu.py new file mode 100644 index 000000000..68cb016bf --- /dev/null +++ b/benchmarks/record_ebm_langevin_native_gpu.py @@ -0,0 +1,98 @@ +#!/usr/bin/env python3 +"""The EBM Langevin loop as one cooperative device kernel, on the owning device. + +Records, per (shape, steps, temperature), the max absolute deviation of the +device loop from the declared numpy policy (`reference_langevin_loop`), the +returned Philox key, and the structure of the emitted kernel taken from the +arena IR (one gpu.func, one scf.for carrying the state in registers, Philox +inside the loop, no linalg/tensor op left). Correctness evidence only: the +per-call host transfers are not a performance path and no promotion is +claimed. gfx1151, gfx1201 and sm_120 packets are separate proofs. +""" +import argparse +import hashlib +import json +import os +import platform +import re +import sys +from pathlib import Path + +import numpy as np + +ROOT = Path(__file__).resolve().parents[1] +sys.path[:0] = [str(ROOT), str(ROOT / "python")] +from tessera.compiler.llvm_tools import llvm_bin_dir # noqa: E402 +from tessera.compiler.native_gpu_storage import replay_arena_ir # noqa: E402 +from tessera.ebm import native_langevin as nl # noqa: E402 + +CASES = [((4, 8), 1, 0.7), ((6, 8), 5, 0.7), ((3, 5), 12, 0.7), ((9, 100), 4, 0.3), + ((16, 33), 8, 0.0), ((2, 1024), 3, 0.5)] + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--backend", choices=["nvidia", "rocm"], required=True) + parser.add_argument("--chip", required=True) + parser.add_argument("--compiler", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + args = parser.parse_args() + llvm = llvm_bin_dir() + if llvm is None: + raise SystemExit("matched LLVM tools are required") + rng = np.random.default_rng(20260916) + rows = [] + for shape, steps, temperature in CASES: + y0, x = (rng.standard_normal(shape).astype(np.float32) for _ in range(2)) + key = [0x1234ABCD9876 + steps, 42] + out, next_key = nl.native_langevin_loop_device(y0, x, key, eta=0.1, temperature=temperature, steps=steps, + backend=args.backend, chip=args.chip, + compiler=args.compiler, llvm_bin=llvm) + expect, expect_key = nl.reference_langevin_loop(y0, x, key, eta=0.1, temperature=temperature, steps=steps) + error = float(np.max(np.abs(np.asarray(out) - expect))) + if list(np.asarray(next_key)) != list(expect_key): + raise SystemExit(f"{shape} K={steps}: returned key {list(next_key)} != {list(expect_key)}") + if not np.allclose(out, expect, rtol=1e-5, atol=1e-5): + raise SystemExit(f"{shape} K={steps} T={temperature} disagrees with the declared policy (max abs error {error})") + if temperature == 0.0 and not np.allclose(out, nl.reference_langevin_loop(y0, x, key, eta=0.1, temperature=0.0, + steps=steps)[0]): + raise SystemExit("zero-temperature loop is not the plain descent") + program = nl.ebm_langevin_program(shape, eta=0.1, temperature=temperature, steps=steps, backend=args.backend, + chip=args.chip, compiler=args.compiler, llvm_bin=llvm) + rows.append(dict(shape=list(shape), steps=steps, temperature=temperature, max_abs_error=error, + next_key=[int(v) for v in np.asarray(next_key)], binding=program.package.binding_digest)) + # The kernel structure is in the device code: replay the packaged source. + source, _ = nl.langevin_device_source((4, 8), eta=0.1, temperature=0.7, steps=3, backend=args.backend, + compiler=args.compiler) + arena = replay_arena_ir(args.compiler, source) + body = arena.split("gpu.func @row_program(", 1)[1] + structure = dict(gpu_funcs=arena.count("gpu.func "), loops=len(re.findall(r"scf\.for ", body)), + loop_carried=re.findall(r"iter_args\([^)]*\) -> \(([^)]*)\)", body), + philox_mul=body.count("arith.mului_extended"), barriers=body.count("gpu.barrier"), + linalg_or_tensor_ops=sum(("linalg." in line or "tensor." in line) for line in body.splitlines())) + # The arena pipeline canonicalizes the packaged kernel: the loop-invariant + # key word (key[0]) is hoisted out of the loop, so the carried tuple is + # (state f32, key[1] i64); before the arena it is (f32, i64, i64). + carried = structure["loop_carried"] + if structure["gpu_funcs"] != 1 or structure["loops"] != 1 or structure["linalg_or_tensor_ops"] != 0 \ + or len(carried) != 1 or not carried[0].startswith("f32") or "i64" not in carried[0]: + raise SystemExit(f"unexpected kernel structure {structure}") + packet = dict(schema=1, backend=args.backend, chip=args.chip, host=platform.node(), + compiler_sha256=hashlib.sha256(args.compiler.read_bytes()).hexdigest(), + llvm_bin=str(llvm), recorder_sha256=hashlib.sha256(Path(__file__).read_bytes()).hexdigest(), + rows=rows, kernel_structure=structure, + proofs=["the K-step loop is one device launch, bit-exact with the declared Philox/Box-Muller policy " + "for every case (K in 1..12, F in 5..1024, T in {0, 0.3, 0.5, 0.7})", + "the gradient is the paired autodiff pass's own adjoint, lowered inside the kernel " + "(no host gradient, no host noise, no Python-emitted kernel)", + "one gpu.func, one scf.for carrying (state, key, key) in registers, Philox inside the loop"], + promotion_eligible=False, measured_performance=False, + envelope="f32, quadratic energy, [rows, features] with features <= 1024; per-call host transfers; " + "correctness only", rocm_chip_env=os.environ.get("TESSERA_ROCM_CHIP")) + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(json.dumps(packet, indent=2) + "\n") + print(json.dumps(dict(output=str(args.output), rows=len(rows), kernel_structure=structure))) + + +if __name__ == "__main__": + main() diff --git a/docs/audit/backend/apple/todo.md b/docs/audit/backend/apple/todo.md index 37d88b1f4..b8391d6f2 100644 --- a/docs/audit/backend/apple/todo.md +++ b/docs/audit/backend/apple/todo.md @@ -9070,3 +9070,19 @@ Sync `GA-NATIVE-GPU-2026-09-16`; owner W6.4. Follow-up required: the arena pipeline route is NVVM/ROCDL only; the Apple arena lane emits MSL (`apple_native_arena.py`) and the hand-written Apple GPU Clifford kernels stay the `apple_gpu` GA row. An Apple package route would expand the same skeleton through the MSL arena emitter. See the [plan log entry](../../compiler/INTEGRATED_COMPILER_LOG.md#2026-09-16--the-clifford-family-reaches-rocm-and-sm120-through-the-arena-pipeline) and the [device packets](../../../../benchmarks/baselines/clifford_native_gpu_20260916/README.md). + +## EBM Langevin loop as one cooperative kernel — 2026-09-16 + +Sync `EBM-NATIVE-GPU-2026-09-16`; owner W4-PRODUCT-1 / AD-SOLVER-IFT-1. + +Not applicable on the device: the row-program emitter targets the NVVM/ROCDL arena pipeline; the Apple GPU EBM rows remain the hand-written MSL kernels and the Apple arena lane (MSL) is a separate emitter (`EBM_NATIVE_LOOP_ARCHITECTURE.md` §3.4). Host-free half validated on the M1 Max: the single-driver chain, the arena replay for both backends, the envelope refusals and both lit fixtures (`row_program_to_gpu_langevin.mlir`, `row_program_to_gpu_reduce.mlir`); device lanes skip honestly. + +See the [plan log entry](../../compiler/INTEGRATED_COMPILER_LOG.md#2026-09-16--the-ebm-langevin-loop-runs-as-one-cooperative-kernel-on-gfx1151-gfx1201-and-sm120). + +## EBM quadratic energy loop through the backbone — 2026-09-16 + +Sync `EBM-NATIVE-QUADRATIC-2026-09-16`; owner W4-PRODUCT-1 / AD-SOLVER-IFT-1. + +Parity validated on the M1 Max CPU lane (11/11). Follow-up unchanged: the Apple GPU EBM kernels (`apple_gpu` EBM rows) are hand-written MSL; no native package route for the loop. + +See the [plan log entry](../../compiler/INTEGRATED_COMPILER_LOG.md#2026-09-16--the-ebm-quadratic-energy-loop-executes-through-the-mlirllvm-backbone). diff --git a/docs/audit/backend/nvidia/todo.md b/docs/audit/backend/nvidia/todo.md index ec9b09ada..98852b028 100644 --- a/docs/audit/backend/nvidia/todo.md +++ b/docs/audit/backend/nvidia/todo.md @@ -7647,3 +7647,19 @@ Sync `GA-NATIVE-GPU-2026-09-16`; owner W6.4. Parity validated on owning sm_120 (RTX 5070, CUDA 13.4 / driver 610.88): ten Clifford ops × three shapes through the native storage route match the GA reference (worst abs error 2.4e-7); grade-2 pruning emits 24 of 64 products in the NVVM kernel; `runtime.launch` row `nvidia_sm120` / `nvidia_clifford_native_compiled` reports native_gpu. This is the first sm_120 GA row in the proof ladder. No performance measured; the route pays per-call host transfers and is correctness evidence only. See the [plan log entry](../../compiler/INTEGRATED_COMPILER_LOG.md#2026-09-16--the-clifford-family-reaches-rocm-and-sm120-through-the-arena-pipeline) and the [device packets](../../../../benchmarks/baselines/clifford_native_gpu_20260916/README.md). + +## EBM Langevin loop as one cooperative kernel — 2026-09-16 + +Sync `EBM-NATIVE-GPU-2026-09-16`; owner W4-PRODUCT-1 / AD-SOLVER-IFT-1. + +Parity validated on sm_120 (Super-Bear): the loop is one launch through the row-program emitter + native storage package, bit-exact with the declared policy in every packet row (worst abs error 0); the reduction program bit-exact with the sequential fold; row `nvidia_sm120` / `nvidia_ebm_langevin_native_compiled`, 11/11 device tests. No performance claim (per-call host transfers; WSL). `build-nvidia-cuda/` — the tree whose driver the runtime resolves — was configured without the EBM/Clifford backends and is now reconfigured with `TESSERA_BUILD_EBM_BACKEND=ON -DTESSERA_BUILD_CLIFFORD_BACKEND=ON`; the fixture passes under both trees there. **Finding, NVIDIA-specific:** on the NVVM arena route `math.sqrt` lowers to libdevice `__nv_sqrtf`, whose precise branch is gated on `__CUDA_PREC_SQRT`, which MLIR's pipeline 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; `div` was `div.rn`. The row-program emitter now calls libdevice's rounding-explicit `__nv_fsqrt_rn` (`convert-gpu-to-nvvm` marks the LLVM math intrinsics illegal, so `llvm.intr.sqrt` is not an option on this route). Follow-up required: every other libdevice f32 path any NVIDIA arena package takes (`__nv_rsqrtf`, `__nv_expf`, `__nv_tanhf`…) is subject to the same default; measure before calling any such result exact, or set the reflect defaults at the packager once LLVM exposes them. + +See the [plan log entry](../../compiler/INTEGRATED_COMPILER_LOG.md#2026-09-16--the-ebm-langevin-loop-runs-as-one-cooperative-kernel-on-gfx1151-gfx1201-and-sm120) and the [device packets](../../../../benchmarks/baselines/ebm_langevin_native_gpu_20260916/README.md). + +## EBM quadratic energy loop through the backbone — 2026-09-16 + +Sync `EBM-NATIVE-QUADRATIC-2026-09-16`; owner W4-PRODUCT-1 / AD-SOLVER-IFT-1. + +Follow-up required: CPU-lane only; no sm_120 EBM lane exists and none is claimed. EBM lit 14/14 on Super-Bear. + +See the [plan log entry](../../compiler/INTEGRATED_COMPILER_LOG.md#2026-09-16--the-ebm-quadratic-energy-loop-executes-through-the-mlirllvm-backbone). diff --git a/docs/audit/backend/rocm/todo.md b/docs/audit/backend/rocm/todo.md index 3cb4d335f..41011be5f 100644 --- a/docs/audit/backend/rocm/todo.md +++ b/docs/audit/backend/rocm/todo.md @@ -7988,3 +7988,19 @@ Sync `GA-NATIVE-GPU-2026-09-16`; owner W6.4. Parity validated on both owning devices: gfx1151 (Princess-Luna) and gfx1201 (Tajasarus, `TESSERA_ROCM_CHIP=gfx1201`, Clifford backend now configured ON in its `build/`), ten ops × three shapes, worst abs error 0; row `rocm` / `rocm_clifford_native_compiled`. The Python-emitted `rocm_clifford_compiled` gfx1151 kernel is untouched and remains the ROCm GA lane until the native route is measured against it (dispatch/allocation/kernel time separately; no `/dev/kfd` here, so no counters). See the [plan log entry](../../compiler/INTEGRATED_COMPILER_LOG.md#2026-09-16--the-clifford-family-reaches-rocm-and-sm120-through-the-arena-pipeline) and the [device packets](../../../../benchmarks/baselines/clifford_native_gpu_20260916/README.md). + +## EBM Langevin loop as one cooperative kernel — 2026-09-16 + +Sync `EBM-NATIVE-GPU-2026-09-16`; owner W4-PRODUCT-1 / AD-SOLVER-IFT-1. + +Parity validated on both owning devices: gfx1151 (Princess-Luna) and gfx1201 (Tajasarus, `TESSERA_ROCM_CHIP=gfx1201`, driver built against the assertions-ON LLVM): the K-step quadratic-energy Langevin loop is one launch through the row-program emitter + native storage package, bit-exact with the declared policy in every packet row (worst abs error 0, K up to 12, F up to 1024, T ∈ {0,…,0.7}); the row-normalization reduction program is bit-exact with the sequential f32 fold; row `rocm` / `rocm_ebm_langevin_native_compiled`, 11/11 device tests on each box. The Python-emitted `rocm_ebm_langevin_compiled` (gfx1151) stays the ROCm lane until the dispatch/allocation/traffic/kernel-time comparison is recorded — no promotion. `arith.mului_extended` and the f64 `math.log`/`cos` lowered and matched on both RDNA parts. Tajasarus's assertions driver caught the two dialect-promise defects recorded in the log entry. + +See the [plan log entry](../../compiler/INTEGRATED_COMPILER_LOG.md#2026-09-16--the-ebm-langevin-loop-runs-as-one-cooperative-kernel-on-gfx1151-gfx1201-and-sm120) and the [device packets](../../../../benchmarks/baselines/ebm_langevin_native_gpu_20260916/README.md). + +## EBM quadratic energy loop through the backbone — 2026-09-16 + +Sync `EBM-NATIVE-QUADRATIC-2026-09-16`; owner W4-PRODUCT-1 / AD-SOLVER-IFT-1. + +Follow-up required: CPU-lane only; `rocm_ebm_langevin_compiled` (Python-emitted gfx1151 kernel) stays the ROCm lane. A device package for the loop needs the tensor-level gradient lowered inside a kernel (tile pipeline), which the arena skeleton route used for Clifford does not provide. EBM lit 14/14 on Princess-Luna and, under the assertions LLVM, on Tajasarus. + +See the [plan log entry](../../compiler/INTEGRATED_COMPILER_LOG.md#2026-09-16--the-ebm-quadratic-energy-loop-executes-through-the-mlirllvm-backbone). diff --git a/docs/audit/backend/x86/todo.md b/docs/audit/backend/x86/todo.md index be6a43259..73b209aad 100644 --- a/docs/audit/backend/x86/todo.md +++ b/docs/audit/backend/x86/todo.md @@ -5314,3 +5314,19 @@ Sync `GA-NATIVE-GPU-2026-09-16`; owner W6.4. Not applicable with a reason: the GPU storage route has no CPU package; x86's native GA lane is the MLIR/LLVM JIT recorded under `GA-NATIVE-FAMILY-2026-09-16`. Host-free halves of the new tests (skeleton expansion, pruning counts) run on any host with the tools. See the [plan log entry](../../compiler/INTEGRATED_COMPILER_LOG.md#2026-09-16--the-clifford-family-reaches-rocm-and-sm120-through-the-arena-pipeline) and the [device packets](../../../../benchmarks/baselines/clifford_native_gpu_20260916/README.md). + +## EBM Langevin loop as one cooperative kernel — 2026-09-16 + +Sync `EBM-NATIVE-GPU-2026-09-16`; owner W4-PRODUCT-1 / AD-SOLVER-IFT-1. + +Not applicable: a GPU-only slice. The CPU JIT lane (`cpu_ebm_langevin_llvm_jit`, `EBM-NATIVE-QUADRATIC-2026-09-16`) is unchanged and its numpy policy is the oracle the device packets are bit-exact against. The x86-hosted drivers on Princess-Luna and Tajasarus ran the host-free chain and fixtures (the Zen 5 hosts), Super-Bear's Zen 2 driver too. + +See the [plan log entry](../../compiler/INTEGRATED_COMPILER_LOG.md#2026-09-16--the-ebm-langevin-loop-runs-as-one-cooperative-kernel-on-gfx1151-gfx1201-and-sm120). + +## EBM quadratic energy loop through the backbone — 2026-09-16 + +Sync `EBM-NATIVE-QUADRATIC-2026-09-16`; owner W4-PRODUCT-1 / AD-SOLVER-IFT-1. + +Parity validated on Princess-Luna (Zen 5): the quadratic-energy Langevin loop compiles as one function through libtessera_jit with the compiler-derived gradient and on-device Philox noise; forward/gradient match the independent formulas and fixed-key samples are bit-exact with the declared policy for 1/5/12 steps; JIT + autodiff suites 174 passed; EBM lit 14/14; autodiff fixtures 21/21. Super-Bear (Zen 2) passes the same EBM and autodiff suites; its bf16 JIT tests fail independent of this branch (bisected against main's JIT sources; Zen 2 has no AVX512-BF16). The Python-emitted `x86_ebm_langevin_compiled` AVX-512 kernel is unchanged; no performance claim. + +See the [plan log entry](../../compiler/INTEGRATED_COMPILER_LOG.md#2026-09-16--the-ebm-quadratic-energy-loop-executes-through-the-mlirllvm-backbone). diff --git a/docs/audit/compiler/INTEGRATED_COMPILER_LOG.md b/docs/audit/compiler/INTEGRATED_COMPILER_LOG.md index 0066f92de..e5771208c 100644 --- a/docs/audit/compiler/INTEGRATED_COMPILER_LOG.md +++ b/docs/audit/compiler/INTEGRATED_COMPILER_LOG.md @@ -4157,3 +4157,31 @@ Remaining: the Python-emitted `rocm_clifford_compiled` / `x86_clifford_compiled` Evidence: [three device packets](../../../benchmarks/baselines/clifford_native_gpu_20260916/README.md) with source fingerprints, `tests/unit/test_clifford_native_gpu.py` (host-free expansion/pruning half on the Mac; device half on the three boxes), `benchmarks/record_clifford_native_gpu.py`, `docs/audit/generated/domain_proof_ladder.md`. + +### 2026-09-16 — the EBM quadratic energy loop executes through the MLIR/LLVM backbone + +Owner: [W4-PRODUCT-1](INTEGRATED_COMPILER_PLAN.md#w4-product-1) + +PRs: domain-support stream, fourth slice (sync `EBM-NATIVE-QUADRATIC-2026-09-16`); co-owner [AD-SOLVER-IFT-1](INTEGRATED_COMPILER_PLAN.md#ad-solver-ift-1). + +Outcome: The GA/EBM review's first acceptance clause for "an energy is a typed program" is met on the CPU lane. `E(y, x) = 0.5·Σ(x − y)²` is a Graph IR function marked for reverse-mode; inside `libtessera_jit` the paired autodiff pass derives `@E__bwd`, the EBM dialect's first lowering pass (`tessera-ebm-lower-langevin`) turns `energy` / `inner_step` / `langevin_step` into arith + linalg over that gradient with Philox-4x32-10 / Box-Muller noise generated in a `linalg.generic`, and the K-step `scf.for` compiles as one function — no per-step host gradient or noise transfers. `langevin_step` gained a variadic `captures` operand for the energy's context. The declared RNG policy is stated in the pass and mirrored bit-for-bit by `native_langevin.reference_langevin_loop`; forward energy and the T = 0 step match the independent formulas; fixed-key samples match for 1, 5 and 12 steps; `runtime.launch` row `cpu` / `cpu_ebm_langevin_llvm_jit`. Verified on the M1 Max, Princess-Luna (Zen 5) and Super-Bear (Zen 2); EBM lit 14/14 on the Mac, both Zen hosts and Tajasarus under its assertions LLVM. Three compiler gaps found and closed on the way: `tessera.sub` had no adjoint (reverse-mode stopped on any subtraction); `tessera.unsqueeze` / `tessera.broadcast`, which the sum-reduce adjoint emits, had no linalg lowering (no reduce-sum gradient could reach the JIT); and the JIT's DPS out-param rewrite did not follow intra-module call sites. + +Remaining: a GPU package for the loop (the tensor-level gradient needs the tile pipeline, not the arena skeleton); nonlinear and manifold energies (sphere / bivector integrators fail closed); opaque-callback energies keep the reported reference path; no performance measurement. The engineering plan for all of these — with the measured stop of each existing device route on this loop (the serial native-tape route rejects it for exactly two reasons; the cooperative route has no EBM Schedule op) — is [EBM_NATIVE_LOOP_ARCHITECTURE.md](../domain/EBM_NATIVE_LOOP_ARCHITECTURE.md). Pre-existing and unrelated: every bf16 JIT test fails on Super-Bear (Zen 2, no AVX512-BF16) with unresolved `_mlir_ciface_*` symbols — bisected by building the JIT from main's sources on that box; bf16 JIT proof belongs on the Zen 5 hosts. + +Evidence: `tests/unit/test_ebm_native_langevin.py` (three CPU hosts), `src/solvers/ebm/test/ir/passes/lower_langevin_quadratic.mlir` + `lower_langevin_rejects.mlir`, `tests/tessera-ir/phase2_autodiff/autodiff_paired_sub.mlir`, `docs/audit/generated/domain_proof_ladder.md` (`cpu=1` under EBM). + + + +### 2026-09-16 — the EBM Langevin loop runs as one cooperative kernel on gfx1151, gfx1201 and sm_120 + +Owner: [W4-PRODUCT-1](INTEGRATED_COMPILER_PLAN.md#w4-product-1) + +PRs: domain-support stream, fifth slice, on the EBM branch (sync `EBM-NATIVE-GPU-2026-09-16`); co-owner [AD-SOLVER-IFT-1](INTEGRATED_COMPILER_PLAN.md#ad-solver-ift-1). + +Outcome: The tensor-level gradient is lowered *inside* a device kernel, through the compiler alone. A new core pass, `tessera-row-program-to-gpu` (`RowProgramToGPUPass.cpp`), consumes the loop after paired autodiff → EBM lowering → `tessera-to-linalg` → inlining and emits one `gpu.func`: one block per row, one lane per feature (≤ 1024), the K-step `scf.for` carried in registers, Philox-4x32-10 / Box–Muller inside the loop, ordered shared-memory reductions (sequential over the feature index, so reduced quantities are bit-exact with a host f32 fold, not tolerance-based), guarded `!llvm.ptr<1>` loads/stores and the arena marker — then `build_native_gpu_storage` packages it and the native tensor binding launches it. `tessera-opt` now registers the EBM and Clifford dialects and passes when built with them, so the whole chain is one driver invocation (T1). Rows `rocm` / `rocm_ebm_langevin_native_compiled` and `nvidia_sm120` / `nvidia_ebm_langevin_native_compiled`; `package_ebm_langevin_native`; `compiler/native_row_program.py` is the packaging seam any `[rows, features]` domain program reuses. Measured on gfx1151 (Princess-Luna), gfx1201 (Tajasarus, assertions-ON driver) and sm_120 (Super-Bear): one launch per loop, bit-exact with the declared policy for K ∈ {1,4,5,8,12}, F ∈ {5,…,1024}, T ∈ {0,…,0.7} — worst abs error 0 in every packet row, the device f64 `log`/`cos` included; the row-normalization reduction program bit-exact with the sequential fold. The scoped G2 Tile contract was not needed: the lowered linalg loop already *is* the scalar body, and one pass over upstream dialects maps it (Decision #31 — no second scalar emitter); G1 serial residency stays scoped only as the fallback outside the envelope. The assertions-enabled driver falsified three promises every NDEBUG driver ran green through — `--inline` needs the LLVM dialect's promised inliner interface (now registered in `tessera-opt`), the pass emitted `tile.alloc_shared` without declaring the tile dialect, and its `tessera.*` provenance attributes load the Tessera dialect an op-free input never loaded (both now declared) — Decision #19's standing lesson, third instance. sm_120 exposed a fourth defect: `math.sqrt` on the NVVM route lowers to libdevice's `__nv_sqrtf`, whose precise branch is gated on a reflect flag MLIR never sets, so the kernel ran `MUFU.SQRT` and one row of the reduction proof was 1 ulp off; the emitter now calls libdevice's rounding-explicit `__nv_fsqrt_rn` on NVIDIA (`convert-gpu-to-nvvm` outlaws the LLVM math intrinsics) and pins `llvm.intr.sqrt` on ROCm and every other libdevice f32 path on that route is recorded as unmeasured. + +Remaining: no performance claim — 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; no Apple package route (the Apple arena lane is MSL). Super-Bear's `build-nvidia-cuda/` tree was configured without the EBM/Clifford backends and is now reconfigured with both ON (the runtime resolves that tree's driver). + +Evidence: [three device packets](../../../benchmarks/baselines/ebm_langevin_native_gpu_20260916/README.md) with source fingerprints, `tests/unit/test_ebm_native_langevin_gpu.py` (host-free chain/replay/envelope half on the Mac; device half on the three boxes), `benchmarks/record_ebm_langevin_native_gpu.py`, `tests/tessera-ir/phase2_autodiff/row_program_to_gpu_langevin.mlir`, `tests/tessera-ir/phase_f5/row_program_to_gpu_reduce.mlir`, `docs/audit/generated/domain_proof_ladder.md` (EBM row gains the two device columns from the matrix). + + diff --git a/docs/audit/compiler/INTEGRATED_COMPILER_PLAN.md b/docs/audit/compiler/INTEGRATED_COMPILER_PLAN.md index 310c5b360..1d5bb2b51 100644 --- a/docs/audit/compiler/INTEGRATED_COMPILER_PLAN.md +++ b/docs/audit/compiler/INTEGRATED_COMPILER_PLAN.md @@ -212,8 +212,8 @@ the next action's host requirement; it is not a live fleet-availability claim. - Next architecture gate: [gated owner and isolated recovery](HEAP_BARRIER_ARCHITECTURE_REVIEW.md#gated-metadata-owner-and-isolated-recovery-2026-09-11): every admitted metadata operation in the opt-in gated owner has replay/device proof; live metadata and legacy import/snapshot paths refuse. Process-owned heap recovery requires confirmed death. Since 2026-09-15 ([probed admission and replacement](HEAP_BARRIER_ARCHITECTURE_REVIEW.md#probed-admission-and-health-checked-replacement-2026-09-15)) a worker is admitted only after its in-process device probe verifies the admitted producers, and a replacement is admitted only after confirmed predecessor death plus its own probe (RTX 5070 and gfx1151). Next: migrate legacy snapshot/import callers to a gated producer and validate actual driver-failure behavior (the recorded fault is an injected stall). Epoch ordering remains; no measured overlap or promotion. - Depends on: — - Start: host-free -- AD gate: [native composed HVP execution](AUTODIFF_EXECUTION_PLAN.md#native-composed-hvp-execution-2026-09-14) adds static CPU saved-product tangents and bounded gfx1201 HVP execution; arbitrary CFG/effects, dynamic results, higher orders and general GPU binding remain open. -- Latest: [probed admission and replacement](INTEGRATED_COMPILER_LOG.md#2026-09-15--probed-admission-and-health-checked-replacement-of-isolated-heap-workers) +- AD gate: [native composed HVP execution](AUTODIFF_EXECUTION_PLAN.md#native-composed-hvp-execution-2026-09-14) adds static CPU saved-product tangents and bounded gfx1201 HVP execution; arbitrary CFG/effects, dynamic results, higher orders and general GPU binding remain open. The energy-as-typed-program acceptance's first clause is met on the CPU lane (2026-09-16, `EBM-NATIVE-QUADRATIC-2026-09-16`): the quadratic energy is a Graph IR function, its gradient is the compiler's, the K-step Langevin loop with on-device Philox noise compiles as one function and is bit-exact with the declared policy on the M1 Max, Zen 5 and Zen 2; the device clause followed the same day (`EBM-NATIVE-GPU-2026-09-16`): the row-program emitter lowers the compiler-derived gradient inside one cooperative kernel (one block per row, lanes per feature, the K-step loop and Philox in registers, ordered reductions), bit-exact on gfx1151, gfx1201 and sm_120 with one launch per loop and one `tessera-opt` invocation for the whole chain; the nonlinear/manifold energies (N1, M1, M2) and the overhead measurement against the Python-emitted lanes remain open — [the EBM native loop architecture](../domain/EBM_NATIVE_LOOP_ARCHITECTURE.md) records what shipped and why the scoped Tile contract was not needed. +- Latest: [the Langevin loop as one cooperative kernel on three GPUs](INTEGRATED_COMPILER_LOG.md#2026-09-16--the-ebm-langevin-loop-runs-as-one-cooperative-kernel-on-gfx1151-gfx1201-and-sm120) ### AD-RESIDUAL-EVAL-1 @@ -264,6 +264,7 @@ the next action's host requirement; it is not a live fleet-availability claim. - Owner: [AUTODIFF_EXECUTION_PLAN.md](AUTODIFF_EXECUTION_PLAN.md) - Gate: Extend residual/predicate/solver envelopes and Apple/NVIDIA consumers through compiler-owned children; require convergence/conditioning certificates and owning-device proof. +- Current increment: The first EBM energy differentiates through the compiler (2026-09-16, `EBM-NATIVE-QUADRATIC-2026-09-16`): `tessera.sub` gained its adjoint and the sum-reduce adjoint's `unsqueeze`/`broadcast` gained linalg lowerings, so the paired autodiff pass now carries a quadratic energy to a native gradient the EBM Langevin lowering consumes; implicit differentiation and OT primitives are untouched. - Depends on: [E2E-REAL-6](#e2e-real-6): canonical artifact boundary for this workload, not every family migration. - Start: host-free - Latest: [recorded increment](INTEGRATED_COMPILER_LOG.md#2026-09-06--functional-analysis-contracts--consolidated-ownership) diff --git a/docs/audit/domain/DOMAIN_AUDIT.md b/docs/audit/domain/DOMAIN_AUDIT.md index aa83fcb76..13006e6f3 100644 --- a/docs/audit/domain/DOMAIN_AUDIT.md +++ b/docs/audit/domain/DOMAIN_AUDIT.md @@ -49,7 +49,7 @@ owners, as the plan's routing index states them: | Domain | Existing capability | Remaining architectural boundary | |---|---|---| | Geometric algebra / Clifford | Signature/product-table references, canonical `clifford_*` operations, differentiated tensor shims, specialized kernels and native grade-pruning passes; since 2026-09-16 the batched geometric product lowers natively (`ExpandProductTable` over any static rank) and executes through MLIR/LLVM on the CPU lane (`cpu` row in the ladder). | The linear/bilinear family (products, contractions, inner/norm, involutions, Hodge star, grade projection, rotor sandwich) lowers natively and executes on the CPU lane and, as native storage packages, on gfx1151/gfx1201/sm_120 (`GA-NATIVE-GPU-2026-09-16`). Still open: exp/log and the field ops, an Apple package route, and measuring the native route against the Python-emitted x86/ROCm/Apple kernels, which remain a second implementation until displaced. General signatures, packed grades and physical derivatives require their own proof. | -| Energy-based models | Reference energies/samplers/losses and specialized update/loss kernels. `geo_sampling.py` uses tape gradients for traceable energies and finite differences otherwise. | The generic `energy.py` Langevin route still uses numerical gradients when no `grad_fn` is supplied. Host gradient evaluation is not a fully resident sampler. Trace an energy body into the native shared AD/loop/ownership path. | +| Energy-based models | Reference energies/samplers/losses and specialized update/loss kernels. `geo_sampling.py` uses tape gradients for traceable energies and finite differences otherwise. Since 2026-09-16 the quadratic energy loop differentiates and executes through the MLIR/LLVM CPU lane (`cpu` row in the ladder): compiler-derived gradient, on-device Philox noise, one native call per K-step loop. | The generic `energy.py` Langevin route still uses numerical gradients when no `grad_fn` is supplied. Host gradient evaluation is not a fully resident sampler. Trace an energy body into the native shared AD/loop/ownership path. | | Attention / persistent state | Canonical families, scheduled packages and bounded native AD; resident O/LSE and isolated CUDA Q/K JVP have explicit packets. | Composed AD, variant breadth, general state lifetimes and cross-target evidence remain separate. KV tiering or a prefetch annotation does not prove overlap. | | Matrix/field calculus, PDE and spectral | Reference/domain contracts and shared transform/solver surfaces exist. | Coordinate/boundary-condition semantics must reach native operators; use the existing layout, numerical-policy, solver and AD owners. Follow [MSW](../compiler/MATH_SOURCE_WORKSTREAM.md) and the [PDE plan](../compiler/PDE_STENCIL_CAPABILITY_PLAN.md). | | Game theory / structured contractions | Reference butterfly/coalition operations and derivative laws provide consumers for shared transforms. | Native shared butterfly lowering, batching, numerical-policy transport and measured execution remain workload-specific gates. | @@ -156,3 +156,40 @@ device kernels (dispatch/allocation/kernel time separately) before any lane change, an Apple package route (the ladder's `rocm`/`apple_gpu` GA rows are still the Python-emitted kernels), then the traceable quadratic energy loop (EBM). +## EBM Langevin loop as one cooperative GPU kernel — 2026-09-16 + +Fifth slice, sync `EBM-NATIVE-GPU-2026-09-16` (W4-PRODUCT-1 / AD-SOLVER-IFT-1). The +tensor-level gradient now lowers inside a device kernel through the compiler +alone: the new row-program emitter turns the lowered `[rows, features]` loop +into one cooperative `gpu.func` (block per row, lane per feature, K steps and +Philox in registers, ordered reductions), the native storage route packages +it, and the ladder's EBM row gains `rocm` and `nvidia_sm120` from the matrix +rows `rocm_ebm_langevin_native_compiled` / `nvidia_ebm_langevin_native_compiled`. +Bit-exact on gfx1151, gfx1201 and sm_120; one launch per loop. Driving it +through the compiler found what no EBM kernel would have: `tessera-opt` had +no way to run a domain dialect's passes in one invocation (now it does, with +EBM and Clifford registered when built), and the assertions-ON driver on +Tajasarus falsified two dialect promises the NDEBUG fleet ran green through. +The Python-emitted device kernels remain the lanes until measured against +this route. Next: the sphere and bivector integrators and the nonlinear +energies — all row programs the same emitter maps — per +[EBM_NATIVE_LOOP_ARCHITECTURE.md](EBM_NATIVE_LOOP_ARCHITECTURE.md). + +## EBM quadratic energy through the backbone — 2026-09-16 + +Fourth slice, sync `EBM-NATIVE-QUADRATIC-2026-09-16` (W4-PRODUCT-1 / AD-SOLVER-IFT-1). Driving the +energy through the compiler rather than through `energy_jit`'s Python plan +found three shared-compiler gaps that no EBM-specific kernel would have +exposed: the Graph IR subtraction had no reverse-mode adjoint, the +sum-reduce adjoint's `unsqueeze`/`broadcast` had no linalg lowering, and the +JIT's DPS rewrite did not follow intra-module calls. All three are closed on +the branch and are what let the paired autodiff pass carry `0.5·Σ(x−y)²` to a +native gradient. The loop is a CPU-lane result on three hosts; the ladder's +EBM row gains `cpu=1` from the matrix. The Python-emitted +`x86_ebm_langevin_compiled` / `rocm_ebm_langevin_compiled` kernels are +unchanged and remain the device lanes. Next in this stream, planned in +[EBM_NATIVE_LOOP_ARCHITECTURE.md](EBM_NATIVE_LOOP_ARCHITECTURE.md): G1 serial +device residency through the native-tape route (two precise admission changes), +then the sphere (M1) and bivector (M2) integrators, nonlinear energies (N1), +the cooperative Tile contract (G2) and one compiler driver (T1). + diff --git a/docs/audit/domain/EBM_NATIVE_LOOP_ARCHITECTURE.md b/docs/audit/domain/EBM_NATIVE_LOOP_ARCHITECTURE.md new file mode 100644 index 000000000..4e939c7ff --- /dev/null +++ b/docs/audit/domain/EBM_NATIVE_LOOP_ARCHITECTURE.md @@ -0,0 +1,449 @@ +--- +last_updated: 2026-09-16 +audit_role: reference +scope: EBM sampling loops as native GPU packages; nonlinear and manifold energies +--- + +# The EBM sampling loop as a native GPU package + +Owner: [W4-PRODUCT-1](../compiler/INTEGRATED_COMPILER_PLAN.md#w4-product-1) +with [AD-SOLVER-IFT-1](../compiler/INTEGRATED_COMPILER_PLAN.md#ad-solver-ift-1); +acceptance from the [GA/EBM review](GA_EBM_ARCHITECTURE_REVIEW.md) §"an +energy is a typed program". Sync keys `EBM-NATIVE-QUADRATIC-2026-09-16` +(CPU lane) and `EBM-NATIVE-GPU-2026-09-16` (device lane). +Written after the CPU-lane slice landed (`tests/unit/test_ebm_native_langevin.py`) +and after driving that same loop as far as the existing device routes take it; +every "today" claim below is a measured stop, not a reading of prose. + +> **Updated 2026-09-16, same day — G2 and T1 landed, and G2 landed by a +> different mechanism than §3.3 scoped.** The cooperative kernel is produced +> by a new pass over the *already lowered* loop +> (`tessera-row-program-to-gpu`, §3.3), not by a `schedule.ebm_langevin` / +> `tile.langevin_kernel` contract: once the paired autodiff pass, the EBM +> lowering, `tessera-to-linalg` and inlining have run, the loop function *is* +> a `[rows, features]` row program in linalg, and that program is the +> scalar body the Tile route would otherwise have carried as a region. The +> row-program emitter maps it directly (one block per row, one lane per +> feature, the K-step `scf.for` and the Philox draw inside the kernel, state +> in registers, ordered shared-memory reductions). Measured on gfx1151, +> gfx1201 and sm_120: one launch per loop, bit-exact with the declared +> policy for every case in the packets. §3.3 below is rewritten to what +> shipped; §3.2's G1 (serial residency through the native-tape route) was +> **not** built — the cooperative route made it unnecessary for this loop +> and it stays scoped only as the fallback for programs outside the +> row-program envelope. `tessera-opt` now registers the EBM and Clifford +> dialects and passes when built with them (T1), so the whole chain is one +> driver invocation, and the assertions-enabled driver on Tajasarus +> falsified three dialect promises the NDEBUG fleet had run green through, +> and sm_120 exposed libdevice's approximate `sqrtf` default (see §6). + +## 1. What the CPU slice established, and what it did not + +The loop is now a program the compiler owns end to end on the CPU lane: + +``` +@quadratic_energy(y, x) tessera.autodiff = "reverse" (Graph IR) + └─ tessera-autodiff-paired ──▶ @quadratic_energy__bwd(y, x, cot) -> (dy, dx) +scf.for K { tessera_ebm.langevin_step(y, key, x) } + └─ tessera-ebm-lower-langevin ──▶ call @E__bwd; y - η·g + √(2ηT)·z; key+[0,1] + z from Philox-4x32-10 / Box-Muller in a linalg.generic + └─ tessera-to-linalg, one-shot bufferize, DPS, LLVM (libtessera_jit) +``` + +One native invocation executes all K steps; noise is generated inside the +compiled loop; the gradient is derived, not written. Three shared-compiler +gaps were closed to get there (`sub` adjoint; `unsqueeze`/`broadcast` linalg +lowering; the JIT's DPS rewrite now follows intra-module calls). + +What it did **not** establish, and what this document plans: + +1. **A device package for the loop.** The Clifford GPU route + (`native_clifford_gpu.py`) works because a rank-1 Clifford op expands to + *scalar* arithmetic that a hand-written per-thread skeleton can host. The + Langevin body is not scalar: `@E__bwd` is tensor-level linalg with a + row reduction (`0.5·Σ(x−y)²` reduces over features), and the loop carries + a whole state tensor across steps. The skeleton route cannot host that. +2. **Nonlinear and manifold energies.** `manifold = "sphere"` and + `"bivector"` fail closed in `tessera-ebm-lower-langevin`; only the + Euclidean integrator lowers. The reference integrators exist in + `python/tessera/ebm/geo_sampling.py` and the GA/EBM review asks for + "explicit singularity/branch handling" before promotion. + +## 2. Measured position of the existing device routes (2026-09-16) + +Three routes were tried on the Mac with the EBM loop module produced by +`native_langevin.langevin_loop_module((4, 8), eta=0.1, temperature=0.7, steps=3)`. + +| Route | Mechanism | Result on the EBM loop | Why | +|---|---|---|---| +| Arena skeleton (`native_clifford_gpu.py`) | Python writes a per-thread `gpu.func`; a dialect driver expands one op; the arena pipeline folds tensors to scalars | **Not applicable** | The body is tensor-level linalg with a reduction and an `scf.for` carrying a tensor; there is no single scalar op to expand | +| Serial native tape (`tessera-native-tape-to-gpu`, the SSD serial lane in `native_ssd._prepare`) | mlir-opt bufferizes the loop function (`one-shot-bufferize` + `buffer-results-to-out-params` + `convert-linalg-to-loops`); the Tessera pass wraps the single host function into a `gpu.module` kernel with device-pointer args at grid (1,1,1) block (1,1,1); `build_native_gpu_storage` packages it | **Rejected at the pass's admission check** — `native tape GPU requires an isolated static f32/f64 buffer product with bounded for/if control …` | Two precise reasons, both in `NativeTapeToGPUPass.cpp:46-51`: (a) the module must hold **exactly one** function, and the loop module holds four (`@quadratic_energy`, its `__bwd`, the energy entry, the loop); (b) the module must carry exactly one owner marker (`tessera.ssd.source`, `tessera.ann.source`, the AD product pair, `tessera.native_result_program`, or `tessera.source_state`), and no EBM marker exists | +| Cooperative Tile lowering (`--tessera-schedule-to-tile=ssd-gpu=`, `NativeSSD.h`) | The Schedule op carries the loop contract; the Tile lowering emits a block-per-(row, column) kernel with lanes owning state elements, LDS + barriers for the reduction | **No EBM Schedule op exists** | `tessera_ebm.langevin_step` is an EBM-dialect op with a symbolic `energy_fn`; nothing projects it into `schedule.*`, so `PMPasses` has nothing to lower | + +Two tooling facts shaped the experiment and are gaps in their own right: + +* The front of the chain needs **three drivers** (`tessera-opt` for the + paired autodiff and `tessera-to-linalg`, `ts-ebm-opt` for the EBM lowering, + `mlir-opt` for bufferization), and `tessera-opt` prints Tessera ops in + custom syntax that `ts-ebm-opt` cannot parse — only `--mlir-print-op-generic` + round-trips. `libtessera_jit` avoids this because it links every dialect + in-process. `tessera-opt` registers neither the EBM nor the Clifford + dialect even when they are built. +* `tessera-opt` already registers the upstream **linalg → scf.parallel → + gpu.launch → outlining** spine (`kEmitSpinePrefix`, `tessera-opt.cpp:249`), + labelled "EMISSION ONLY … launch stays hardware-gated" in June. Launch is no + longer gated: `build_native_gpu_storage` + `NativeTensorCall` launch any + `gpu.module` with the storage ABI on gfx1151, gfx1201 and sm_120. That + spine is therefore a third candidate, not just an inspection aid. + +## 3. Architecture + +### 3.1 Principle: one lowering authority, three physical routes + +Decision #31 (one production lowering per boundary) and #28 (measured +selection among implementations) settle the shape. The EBM lowering +(`tessera-ebm-lower-langevin`) stays the **only** authority for what a +Langevin step means — gradient call, integrator, noise policy, key advance — +and produces tensor-level IR. What differs per target is only how that IR +becomes a kernel: + +``` + tessera_ebm.langevin_step (+ energy_fn, captures) + │ tessera-ebm-lower-langevin (one authority) + ▼ + scf.for over tensors: call @E__bwd, arith, linalg.generic(noise) + ┌──────────────────────┬──────────────────────┬────────────────────────┐ + ▼ ▼ ▼ ▼ + CPU JIT (landed) G1 serial device kernel G2 cooperative kernel Apple arena (MSL) + libtessera_jit native-tape-to-gpu linalg→parallel→gpu apple_native_arena + one call, K steps grid 1, one call, K steps rows→blocks, feats→lanes (later; same IR) +``` + +Every route validates by **replay** (`replay_arena_ir` byte-compares the +packaged IR with a fresh lowering — the same discipline the heap, SSD and +Clifford packages use) and reports the route in its packet (Decision #12). + +### 3.2 G1 — serial device residency (the acceptance's "no per-step host transfers" on a device) + +Purpose: prove the loop is device-resident and bit-exact with the declared +policy on gfx1151/gfx1201/sm_120 with the smallest possible change, before +any parallel mapping. Serial means one thread executes the whole loop; it is +a correctness route (as the SSD serial lane is), never a performance claim. + +Changes, each small and testable: + +1. **Isolate the loop function.** After `tessera-ebm-lower-langevin`, inline + `@E__bwd` (and `@E` if the energy entry is exported) into the loop function + and drop the dead originals — `mlir::inlineCall` on the call sites the + lowering just emitted, gated on a pass option `isolate-loop=true`. The + native tape pass then sees exactly one function. The JIT keeps calls + (its DPS rewrite now follows them); only the device route inlines. +2. **An EBM owner marker.** The lowering stamps + `tessera.ebm.source = ""` on the module + when isolating, and `NativeTapeToGPUPass` admits it as a sixth owner kind + (the `ssd` branch is the template: one function, static buffers, bounded + `for/if`). The marker is what lets `validate()` re-derive the kernel from + the source and byte-compare it — the replay contract. +3. **Admission audit of the body.** The pass's whitelist ("bounded for/if + control, static f32/f64 buffers, ≤ 4096 temporary bytes") must be checked + against what the loop actually contains after `convert-linalg-to-loops`: + integer `arith` on i32/i64 (Philox), `arith.mului_extended`, `math.log`, + `math.cos`, `math.sqrt`, `index_cast`, and the key tensor `memref<2xi64>`. + Expected finding: the integer/`mului_extended` ops are new to that pass + and its op whitelist needs them; the NVVM/ROCDL lowerings handle them + (the Clifford packets already execute `mului_extended`-free scalar code; + `math.*` lowers through `convert-math-to-llvm`, as the arena pipeline does). +4. **Python:** `native_langevin.package_ebm_langevin_native(shape, …, target)` + mirroring `package_clifford_native`, a `run_*` through `NativeTensorCall` + (host arrays in/out via the driver handle, as `CliffordDeviceProgram` + does), execution-matrix rows `rocm` / `rocm_ebm_langevin_native_compiled` + and `nvidia_sm120` / `nvidia_ebm_langevin_native_compiled`, and a recorder + writing per-device packets. + +Acceptance G1: fixed-key samples bit-exact with `reference_langevin_loop` for +K ∈ {1, 5, 12} on all three devices; the packet records that the loop ran as +one launch (the binding's launch count) and that no host buffer was touched +between steps (the kernel's only host contact is the argument copy-in and +result copy-out). Not a performance route; no promotion. + +### 3.3 G2 — cooperative kernel (rows to blocks, features to lanes) — landed 2026-09-16 + +Purpose: the loop at real batch sizes, as one launch. What shipped is the +**row-program emitter**, `src/transforms/lib/RowProgramToGPUPass.cpp` +(`--tessera-row-program-to-gpu=backend={nvidia,rocm} entry=`), consuming +the lowered loop after +`--tessera-autodiff-paired --tessera-ebm-canonicalize --tessera-ebm-lower-langevin +--tessera-to-linalg --inline --convert-elementwise-to-linalg --canonicalize --cse` +in one `tessera-opt` invocation: + +* **Input envelope (fails closed outside it).** One entry `func.func` whose + operands and results are `[R, F]` f32 tensors (lane values), `[R, 1]` / + `[R]` f32 tensors (row values) and small integer vectors of ≤ 8 elements + (uniform values, e.g. the Philox key); a body of parallel + `linalg.generic` (identity or row-broadcast maps), feature-axis + `linalg.reduce`, `tensor.empty`/`linalg.fill`, `expand_shape`/`collapse_shape` + between row shapes, `tensor.extract` on a uniform vector with a constant + index, splat constants, and `scf.for` over any mix of those. `F ≤ 1024` + (one lane each; lanes = next power of two ≥ F, spare lanes masked). Any + call is refused with "inline every call before lowering"; anything else + names the op it cannot classify. +* **Mapping.** One `gpu.func @row_program` in `gpu.module @native_row`, + `known_block_size = [lanes, 1, 1]`, one block per row. Lane values live in + registers across the whole program including every `scf.for` iteration + (the loop is carried as `iter_args` of scalars, so a K-step Langevin loop + is *one* kernel with the state, the gradient and the noise never leaving + registers). Arguments are `!llvm.ptr<1>` per tensor plus the index scratch + the arena pipeline expects; loads and stores are guarded by the lane mask + (lane values) or the block leader (row / uniform values). The + `tile.alloc_shared` marker is emitted so `build_native_gpu_storage`'s + shared-memory sizer sees the kernel it already knows. +* **Reductions.** A feature-axis `linalg.reduce` becomes an **ordered** + fold: every active lane stores its element into a + `@row_reduction` array in address space 3, `gpu.barrier`, the leader folds + the combiner over lanes **in index order** into a scalar and writes it + back, `gpu.barrier`, every lane reads the broadcast, `gpu.barrier`. The + declared reduction order is therefore "sequential over the feature index", + the same order a host `for` loop in f32 produces — which is why the row + normalization proof below is bit-exact, not tolerance-based. +* **Noise.** The EBM lowering's Philox-4x32-10 / Box–Muller `linalg.generic` + is an ordinary lane body to the emitter: `linalg.index 0/1` become + `(block_id, thread_id)`, the ten `arith.mului_extended` rounds and the + f64 `math.log`/`math.cos` stay inside the loop, and the key words are + uniform values carried through the loop (canonicalization hoists the + invariant `key[0]`, so the packaged loop carries `(state f32, key[1] i64)`). +* **Provenance.** The module carries `tessera.row_program.{source, entry, + rows, features, backend}`; the tensor contract producer merges into that + attribute dictionary instead of requiring a bare module, so the package + records the lowered program it was built from. + +Python: `compiler/native_row_program.py` (`row_program_kernel` → +`row_program_device_source` → `bind_row_program` → `row_program_device`, a +cached `HostArrayProgram`) is the seam every row-shaped domain loop packages +through; `ebm/native_langevin.py` builds the Langevin module and its specs +on top of it (`langevin_device_source`, `native_langevin_loop_device`, +`package_ebm_langevin_native`); `runtime.launch` rows +`rocm` / `rocm_ebm_langevin_native_compiled` and +`nvidia_sm120` / `nvidia_ebm_langevin_native_compiled`. + +Measured (packets in `benchmarks/baselines/ebm_langevin_native_gpu_20260916/`, +`tests/unit/test_ebm_native_langevin_gpu.py`): on gfx1151 (Princess-Luna), +gfx1201 (Tajasarus, `TESSERA_ROCM_CHIP=gfx1201`, assertions-ON driver) and +sm_120 (Super-Bear) the loop is one launch and **bit-exact** with +`reference_langevin_loop` for K ∈ {1, 4, 5, 8, 12}, F ∈ {5, 8, 33, 100, 1024}, +T ∈ {0, 0.3, 0.5, 0.7} — worst absolute error 0 in every row, including the +f64 `log`/`cos` of the noise (the libm risk in §6 did not materialize on any +of the three devices). The T = 0 loop is the plain descent. The quadratic +gradient is elementwise (the sum-reduce adjoint is a broadcast), so the +Langevin kernel carries no reduction; the reduction path is proven +separately by a row-normalization program (`x / sqrt(Σ_f x²)`, F up to +1024) that matches the sequential-order f32 fold bit-for-bit on the same +three devices. + +What this is not: a performance claim. The route's per-call host transfers +are correctness-only; the Python-emitted `rocm_ebm_langevin_compiled` / +`x86_ebm_langevin_compiled` / Apple kernels remain their lanes until the +dispatch/allocation/traffic/kernel-time comparison the GA/EBM review asks for +is recorded with `route` and `latency_source`. The envelope is also narrow +on purpose: rows wider than 1024 features (a second reduction level), row +programs with data-dependent control flow, and energies whose gradient is +not a row program (cross-row coupling) fail closed and would go through +§3.2's serial route or a real Tile contract. + +Why not the Tile contract scoped before: `schedule.ebm_langevin` with a body +region would have carried exactly the program the emitter now reads from +linalg, and would have needed its own body lowering — a second scalar +emitter beside the one the Tile kernels use (Decision #31). The row-program +pass is the smaller authority: it adds one pass over upstream dialects, no +new ops, and the same pass serves any `[rows, features]` domain program (the +GA family's batched products qualify once their ExpandProductTable output is +inlined). If a future energy needs multi-level reductions or cross-row +coupling, the Tile contract is the next step and this pass is its +measurement baseline. + +Acceptance G2, as met: bit-exactness (samples *and* the ordered reductions); +one launch per loop; the packets. As not met: the separate overhead +measurement against the Python-emitted lanes — promotion waits for it. + +### 3.4 Apple + +The Apple GPU route is the MSL arena (`apple_native_arena.py`, +`--tessera-tile-buffer-arena=emit-apple-msl=true`), not the NVVM/ROCDL +pipeline. G1's isolated loop function is the input it needs; the open work +is that the arena emitter takes a Tile-level buffer program, not a +bufferized scf loop. Route through G2b's `tile.langevin_kernel` once it +exists; do not write a third emitter (`CLAUDE.md` §"the real Apple gap"). + +## 4. Nonlinear and manifold energies + +### 4.1 Nonlinear energies: nothing manifold-specific is needed + +The Euclidean integrator is energy-agnostic: the lowering calls `@E__bwd` +whatever `E` is. A nonlinear energy is therefore a **differentiation +coverage** question, not a lowering one. The quadratic case already exposed +that coverage gap twice (`sub`; `unsqueeze`/`broadcast`). The next energies +to admit, chosen because each exercises one more adjoint on the Graph IR +path and each has an Apple or ROCm reference kernel to compare against: + +| Energy | Graph ops on the gradient path | Adjoint status today | Branch / singularity handling | +|---|---|---|---| +| Huber / smooth-L1 on `x − y` | `tessera.loss.huber` | adjoint exists (`HuberLossOp::buildAdjoint`) | the kink at `|r| = δ` is a `compare_scalar` + `masked_fill`, already the pattern ReLU uses | +| Softplus / log-cosh energy | `softplus`, `log`, `cosh` | `tanh`/`sigmoid`/`gelu` have adjoints; `softplus`/`cosh` need checking | overflow guard: the adjoint must use the stable form `sigmoid(r)`, never `exp(r)/(1+exp(r))` | +| Softmax-normalized energy (contrastive rows) | `tessera.softmax`, `reduce` | adjoints exist | the row maximum subtraction is inside the softmax adjoint already | +| Energy with a learned `W`: `0.5·‖x − Wy‖²` | `tessera.matmul` | adjoint exists; this is the paired-matmul row | none | + +Rule for admission (fails closed, as today): the paired autodiff pass must +produce `@E__bwd` with **no** `tessera.custom_adjoint_call` in it — a +placeholder means the gradient would round-trip through the Python VJP +registry, which is a host transfer per step. The EBM lowering checks for the +placeholder and refuses. + +### 4.2 Sphere: tangent projection + retraction + +Reference (`geo_sampling.sphere_langevin_step`): with `x` unit-norm per row, +`ξ ~ N(0, I)`, + +``` +g_t = g − ⟨g, x⟩ x (project ∇E to the tangent plane) +ξ_t = ξ − ⟨ξ, x⟩ x (project the noise likewise) +x' = (x − η g_t + √(2ηT) ξ_t) / ‖x − η g_t + √(2ηT) ξ_t‖ (retraction = normalize) +``` + +Lowering: three per-row dot products and one norm, i.e. **row reductions +over the feature axis**, then elementwise work. That is precisely what G2b's +kernel provides and G1's serial kernel trivially provides; on the CPU lane +it is `linalg.reduce` + broadcast — every piece already lowers after this +slice's `unsqueeze`/`broadcast` fix. + +Singularities and their contract (fail closed, Decision #21a): + +* **Entry precondition** `|‖x‖ − 1| ≤ 1e-3` per row. The reference asserts + it on the host. Natively: the kernel writes a per-row status word (the + pass already has `status-buffer` machinery for structured assertions), and + the package refuses the result when any row violated it; no silent + renormalization on entry. +* **Retraction at the origin.** If `‖x − η g_t + …‖` underflows (a step that + cancels `x` exactly, or a huge η), normalization divides by ~0. Contract: + `‖·‖ < ε` (ε = 1e-12 in f32 squared-norm terms) sets the row status and + keeps the previous `x`; the reference gets the same guard so both agree. +* **Projection when `⟨g, x⟩ x` dominates.** Numerically benign (a + subtraction of two O(‖g‖) terms), but the *order* of the dot-product + reduction must be declared: sequential over features in G1/CPU, the + block-reduction order in G2b, recorded in `numeric_policy` and matched by + the reference's `np.float32` accumulate; the parity tolerance is 1e-5, not + bit-exactness, for the projected quantities (the noise itself stays + bit-exact). +* **Temperature 0** must reduce to projected gradient descent with + retraction; a T = 0 fixture checks `‖x'‖ = 1` and monotone energy decrease + for the quadratic energy on the sphere. + +`manifold = "sphere"` therefore lowers as: G1/CPU immediately (reductions +already lower), G2b when the Tile contract lands. The per-row status word is +the new piece. + +### 4.3 Bivector (grade-restricted Lie-algebra sampling) + +Reference (`geo_sampling.bivector_langevin_step`): state is a multivector +constrained to grade `k` (default 2, so(3) in Cl(3,0)); both the gradient +and the noise are grade-projected, then the Euclidean affine step applies: + +``` +g_k = grade_k(∇E) ξ_k = grade_k(ξ) s' = s − η g_k + √(2ηT) ξ_k +``` + +Lowering: the grade projection is exactly the `tessera_clifford.grade` op +the Clifford slice lowered (a compile-time keep-mask: `UnaryMapPattern`, +`Unary::Grade`). The EBM lowering for `manifold = "bivector"` emits +`clifford.grade(grade=k)` on the gradient and on the noise before the affine +update — the EBM dialect depends on the Clifford dialect for this integrator +only, which is a build-time dependency `TesseraEBM → TesseraClifford` gated +on `TESSERA_BUILD_CLIFFORD_BACKEND` (both are ON fleet-wide since this +stream). No new arithmetic, no singularity: the projection is linear and the +step stays in the subspace by construction (the reference asserts the input +grade; natively the input-grade check is a status word like the sphere's +norm check). + +Two facts to keep honest: this integrator samples on the **Lie algebra** +(grade-k multivectors), not on the group — mapping to rotors is `exp`, which +has no native lowering yet (`GA-NATIVE-FAMILY-2026-09-16` listed `exp`/`log` +as open) — and the state layout is `[..., 2^n]` coefficients, so a batch of +rotors is the batched Clifford tensor the W6.4 lowering already handles. + +### 4.4 What stays closed + +* `manifold` values outside {euclidean, sphere, bivector} — the enum is + closed by `EBM_ManifoldAttr`. +* Energies whose `__bwd` needs a `custom_adjoint_call` (host VJP). +* Dynamic shapes (the loop nest and Philox counter need static extents; the + reference keys the counter on the flat index). +* `exp`/`log` of multivectors (rotor sampling on the group). + +## 5. Slices and acceptance, in order + +| Slice | Deliverable | Acceptance | Fleet | +|---|---|---|---| +| G1 | isolate-loop inlining + `tessera.ebm.source` owner + admission audit; `package_ebm_langevin_native`; rows for `rocm` and `nvidia_sm120`; recorder | bit-exact K ∈ {1,5,12} on gfx1151, gfx1201, sm_120; one launch per loop | Princess-Luna, Tajasarus, Super-Bear | +| M1 | `manifold = "sphere"` on CPU + G1: tangent projection, retraction, per-row status word; reference guard mirrored | ‖x'‖ = 1 per row, T = 0 monotone descent, 1e-5 parity on projected quantities, bit-exact noise | Mac + x86 + G1 devices | +| M2 | `manifold = "bivector"` via `clifford.grade` on gradient and noise; EBM → Clifford build dependency | parity with `bivector_langevin_step` on Cl(3,0) grade 2; state stays grade-2 over 100 steps | same | +| N1 | Huber and softplus energies through the paired pass with no placeholder; the `softplus`/`cosh` adjoints if missing | gradient matches finite differences and the Apple/ROCm reference kernels | Mac + x86 | +| G2 | **landed 2026-09-16** as the row-program emitter over the lowered loop (§3.3), not a Schedule/Tile op; rows `rocm_ebm_langevin_native_compiled` / `nvidia_ebm_langevin_native_compiled`; recorder | bit-exact on gfx1151, gfx1201, sm_120 (packets); one launch per loop; **open:** the dispatch/alloc/traffic/kernel-time packets vs the Python-emitted lanes | all three GPUs | +| T1 | **landed 2026-09-16**: `tessera-opt` registers EBM + Clifford dialects and passes when built with them (`TESSERA_HAVE_EBM` / `TESSERA_HAVE_CLIFFORD`) plus `convert-elementwise-to-linalg`; the whole chain is one invocation, and lit gains a `tessera-ebm` feature | one `tessera-opt` invocation runs the G2 chain; fixtures `phase2_autodiff/row_program_to_gpu_langevin.mlir`, `phase_f5/row_program_to_gpu_reduce.mlir` | host-free | + +G2 and T1 landed first, and G1 was skipped: the cooperative route covers the +quadratic loop outright, so serial residency is only the fallback for +programs outside the row-program envelope. M1/M2 now build on G2 (the +tangent projection and retraction are row programs: a feature-axis +reduction followed by lane arithmetic, exactly what the emitter maps); N1 is +independent and can run in parallel on the CPU lane. + +## 6. Risks named now + +* **Reduction order across lanes.** *Resolved by construction (2026-09-16):* + the emitter's fold is sequential over the feature index, so it reproduces + the host's sequential f32 sum bit-for-bit (row normalization proof, three + devices). A tree or warp-shuffle reduction, if one is ever adopted for + speed, changes the declared order and must re-derive the reference — never + hide it behind a looser tolerance. +* **Philox on ROCDL.** `arith.mului_extended` lowers to `llvm.umul_with_overflow` + or a 64-bit multiply; both targets have it, but the first G1 run on each + device is the proof, not the ISA table. The `philox_msl_source` template + already exists for Apple. +* **`math.log`/`math.cos` precision.** *Measured 2026-09-16:* the f64 + `log`/`cos` of the Box–Muller draw matched the host bit-for-bit on + gfx1151, gfx1201 and sm_120 in every packet row (worst abs error 0 after + the single f32 rounding). The risk stays named: a toolkit upgrade can + change a libm ulp, and the recorder will show it as a non-zero + `max_abs_error` on a noise row — amend the policy to a tolerance for `z` + then, never silently. +* **libdevice's approximate f32 paths are the NVVM route's default.** + *Measured 2026-09-16:* `math.sqrt` on the NVVM route lowers to + `__nv_sqrtf`, whose precise branch is gated on the `__CUDA_PREC_SQRT` + reflect value that MLIR's pipeline never sets (LLVM 23's NVVMReflect + exposes only `nvvm-reflect-ftz`), so the kernel ran `MUFU.SQRT` and one + row of the row-normalization proof was 1 ulp off on sm_120 — the + division was `div.rn` and both RDNA parts were exact. `convert-gpu-to-nvvm` + marks the LLVM math intrinsics illegal, so the emitter now calls + libdevice's rounding-explicit `__nv_fsqrt_rn` on NVIDIA (correctly rounded + whatever the reflect flags say) and pins `llvm.intr.sqrt` on ROCm (AMDGPU's + correctly rounded expansion). Every other `math.*` f32 op that reaches libdevice + on this route (`rsqrt`, `exp`, `log`, `tanh`…) is still subject to the + same default and must be measured before its result is called exact; the + f64 `log`/`cos` of the noise are bit-exact because libdevice's f64 paths + carry no approximate branch. +* **NDEBUG hides dialect-promise defects (Decision #19, third instance).** + The chain ran green on every NDEBUG driver in the fleet and aborted twice + on Tajasarus's assertions-ON driver: `--inline` needs the LLVM dialect's + promised `DialectInlinerInterface`, which `registerAllExtensions` does not + provide (`tessera-opt` now registers it), and the row-program pass emits + `tile.alloc_shared` without declaring the tile dialect as a dependency + (now declared), and parsing its `tessera.*`-prefixed provenance + attributes loads the Tessera dialect, which an input without Tessera ops + never loaded (now declared). All one-line fixes; none visible without that + box. Route every new pass through it before recording "passes". +* **Register pressure in G2b.** Keeping a row's state and gradient in + registers across K steps bounds the feature width per lane; wider rows + need the second reduction level. Declare the admitted width in the + contract and fail closed above it. +* **Owner-marker sprawl in `NativeTapeToGPUPass`.** Adding a sixth owner + kind is the fifth time that pass grew a special case. Record it in the + compiler audit as technical debt to fold into a generic "isolated buffer + program" admission with a required `tessera.source_kind` attribute. diff --git a/docs/audit/domain/GA_EBM_ARCHITECTURE_REVIEW.md b/docs/audit/domain/GA_EBM_ARCHITECTURE_REVIEW.md index c56e81ad5..44e06147a 100644 --- a/docs/audit/domain/GA_EBM_ARCHITECTURE_REVIEW.md +++ b/docs/audit/domain/GA_EBM_ARCHITECTURE_REVIEW.md @@ -145,3 +145,41 @@ claim is made, and the Python `x86_clifford_compiled` / `rocm_clifford_compiled` kernels are unchanged and remain the device lanes until displaced by measured evidence. +## The Langevin loop as one cooperative kernel — 2026-09-16 + +Sync `EBM-NATIVE-GPU-2026-09-16` (W4-PRODUCT-1 / AD-SOLVER-IFT-1). The device +half of the first acceptance clause: the compiler-derived gradient is +lowered *inside* a kernel by the new row-program emitter +(`tessera-row-program-to-gpu`) over the lowered loop — one block per row, +one lane per feature, the K-step loop with its Philox draw carried in +registers, ordered shared-memory reductions — packaged by the native GPU +storage route and launched as one call. Bit-exact with the declared policy +on gfx1151, gfx1201 and sm_120 (packets, worst abs error 0), reductions +bit-exact with the sequential f32 fold, one launch per loop, one driver +invocation for the chain. Still not met: the overhead measurement that +would let this lane displace the Python-emitted `*_ebm_langevin_compiled` +kernels (no promotion is claimed); nonlinear and manifold energies (the +sphere and bivector integrators are row programs the same emitter maps and +are the next slices); opaque callbacks; Apple. Details and the reason the +scoped Tile contract was not needed: [EBM_NATIVE_LOOP_ARCHITECTURE.md](EBM_NATIVE_LOOP_ARCHITECTURE.md) §3.3. + +## Quadratic energy through the backbone — 2026-09-16 + +Sync `EBM-NATIVE-QUADRATIC-2026-09-16` (W4-PRODUCT-1 / AD-SOLVER-IFT-1). The first acceptance clause +above is met on the CPU lane: the quadratic energy is a Graph IR function, +its gradient is the compiler's (paired reverse-mode), fixed-key samples agree +bit-for-bit with the declared Philox / Box-Muller policy, and a K-step loop +executes as one native call without per-step host gradient transfers, on the +M1 Max, Princess-Luna and Super-Bear. The EBM dialect gained its first +lowering pass and a `captures` operand on `langevin_step`; the shared +compiler gained the `sub` adjoint and `unsqueeze`/`broadcast` lowerings the +energy's gradient needed. Not yet met: the nonlinear and manifold cases +(sphere / bivector integrators fail closed), opaque callbacks (still the +reported reference path), a GPU package for the loop, and the "no per-step +transfers" claim on a device — this is the CPU lane. `energy.py::langevin_step` +still takes finite differences without `grad_fn`; the native lane is opt-in. +The device package, the sphere and bivector integrators and the nonlinear +energies are scoped as slices G1/M1/M2/N1/G2/T1 in +[EBM_NATIVE_LOOP_ARCHITECTURE.md](EBM_NATIVE_LOOP_ARCHITECTURE.md), which +records where each existing device route stops on this loop today. + diff --git a/docs/audit/generated/autodiff_connection_ledger.csv b/docs/audit/generated/autodiff_connection_ledger.csv index df1cfa1b1..7611685d6 100644 --- a/docs/audit/generated/autodiff_connection_ledger.csv +++ b/docs/audit/generated/autodiff_connection_ledger.csv @@ -292,7 +292,7 @@ stack,tensor_algebra,yes,none,none,,,,,,,,python_reference=python-unit-registry, std,reduction,yes,none,none,,,,,,,,python_reference=python-unit-registry,,, stft,spectral,yes,native,native,,,,x86_avx512,x86_avx512,,x86_avx512,python_reference=python-unit-registry;ir_adjoint=llvm23-core;ir_tangent=llvm23-core;device[x86_avx512=LLVM/MLIR 23; Ryzen AI MAX+ 395 AVX-512],native compiler adjoint; native backward executes on x86_avx512 (Phase 4),x86_avx512=save_inputs,x86_avx512=dedicated stop_gradient,layout_transform,yes,native,native,,,,,,,,python_reference=python-unit-registry;ir_adjoint=llvm23-core;ir_tangent=llvm23-core,native compiler adjoint,, -sub,elementwise,yes,none,native,,,,,,,,python_reference=python-unit-registry;ir_tangent=llvm23-core,,, +sub,elementwise,yes,native,native,,,,,,,,python_reference=python-unit-registry;ir_adjoint=llvm23-core;ir_tangent=llvm23-core,native compiler adjoint,, sum,stable_reduction,yes,native,none,,cpu,,,,,,python_reference=python-unit-registry;ir_adjoint=llvm23-core;bwd_cpu_ir_oracle=llvm23-core,native compiler adjoint (static extent required for mean),, svd,linalg_decomposition,yes,none,none,,,,,,,,python_reference=python-unit-registry,,, take,indexing,yes,none,none,,,,,,,,python_reference=python-unit-registry,,, diff --git a/docs/audit/generated/autodiff_connection_ledger.md b/docs/audit/generated/autodiff_connection_ledger.md index b895df7dd..acdd7397e 100644 --- a/docs/audit/generated/autodiff_connection_ledger.md +++ b/docs/audit/generated/autodiff_connection_ledger.md @@ -19,7 +19,7 @@ One row per differentiable **op family**, over the independent proof axes of [`A - Differentiable families tracked: **314** - `python_reference` (Python VJP/JVP): **314** -- `ir_adjoint = native`: **54** (adam, adam_w, add, all_gather, all_reduce, all_to_all, binary_cross_entropy_loss, broadcast, cross_entropy_loss, dct, depth_attn, dropout, expand, fft, flash_attn, flatten, gelu, huber_loss, ifft, irfft, istft, j_s_divergence_loss, k_l_divergence_loss, layer_norm, mae_loss, matmul, momentum, mse_loss, mul, nesterov, permute, reduce, reduce_scatter, relu, reshape, rfft, rmsnorm, sgd, sigmoid, silu, smooth_l1_loss, softmax, spectral_conv, spectral_filter, squeeze, stft, stop_gradient, tanh, transpose, unsqueeze, view) +- `ir_adjoint = native`: **55** (adam, adam_w, add, all_gather, all_reduce, all_to_all, binary_cross_entropy_loss, broadcast, cross_entropy_loss, dct, depth_attn, dropout, expand, fft, flash_attn, flatten, gelu, huber_loss, ifft, irfft, istft, j_s_divergence_loss, k_l_divergence_loss, layer_norm, mae_loss, matmul, momentum, mse_loss, mul, nesterov, permute, reduce, reduce_scatter, relu, reshape, rfft, rmsnorm, sgd, sigmoid, silu, smooth_l1_loss, softmax, spectral_conv, spectral_filter, squeeze, stft, stop_gradient, sub, tanh, transpose, unsqueeze, view) - `ir_adjoint = placeholder` (Python round-trip, not native): **3** (log_softmax, sin, softplus) - `ir_adjoint = mixed` (kind-aware native + placeholder): **0** - `ir_tangent = native`: **37** (add, all_gather, all_reduce, all_to_all, broadcast, dct, depth_attn, dropout, es_low_rank_correction, expand, fft, flash_attn, flatten, ifft, irfft, istft, layer_norm, matmul, mul, permute, reduce, reduce_scatter, reshape, rfft, rmsnorm, sigmoid, softmax, spectral_conv, spectral_filter, squeeze, stft, stop_gradient, sub, tanh, transpose, unsqueeze, view) @@ -367,7 +367,7 @@ One row per differentiable **op family**, over the independent proof axes of [`A | `std` | reduction | yes | none | none | — | — | — | — | — | — | — | — | — | python_reference=python-unit-registry | | | `stft` | spectral | yes | native | native | — | — | — | x86_avx512 | x86_avx512 | — | x86_avx512 | x86_avx512=save_inputs | x86_avx512=dedicated | python_reference=python-unit-registry; ir_adjoint=llvm23-core; ir_tangent=llvm23-core; device[x86_avx512=LLVM/MLIR 23; Ryzen AI MAX+ 395 AVX-512] | native compiler adjoint; native backward executes on x86_avx512 (Phase 4) | | `stop_gradient` | layout_transform | yes | native | native | — | — | — | — | — | — | — | — | — | python_reference=python-unit-registry; ir_adjoint=llvm23-core; ir_tangent=llvm23-core | native compiler adjoint | -| `sub` | elementwise | yes | none | native | — | — | — | — | — | — | — | — | — | python_reference=python-unit-registry; ir_tangent=llvm23-core | | +| `sub` | elementwise | yes | native | native | — | — | — | — | — | — | — | — | — | python_reference=python-unit-registry; ir_adjoint=llvm23-core; ir_tangent=llvm23-core | native compiler adjoint | | `sum` | stable_reduction | yes | native | none | — | cpu | — | — | — | — | — | — | — | python_reference=python-unit-registry; ir_adjoint=llvm23-core; bwd_cpu_ir_oracle=llvm23-core | native compiler adjoint (static extent required for mean) | | `svd` | linalg_decomposition | yes | none | none | — | — | — | — | — | — | — | — | — | python_reference=python-unit-registry | | | `take` | indexing | yes | none | none | — | — | — | — | — | — | — | — | — | python_reference=python-unit-registry | | diff --git a/docs/audit/generated/compiler_progress.csv b/docs/audit/generated/compiler_progress.csv index 0a18e583d..996e5e5d9 100644 --- a/docs/audit/generated/compiler_progress.csv +++ b/docs/audit/generated/compiler_progress.csv @@ -15,18 +15,18 @@ primitive_axis,lowering_rule,closed,520,520,0,"primitive contract axis; open mea primitive_axis,backend_kernel,mixed,100,520,420,"primitive contract axis; open means partial or planned, not necessarily missing API support",docs/audit/generated/s_series_status.md,Promote by backend/pathway; do not treat every target as an all-up compiler veto. integration,Verifier coverage,closed,244,244,0,real=244,docs/audit/generated/verifier_coverage.csv,No action unless this row reopens. integration,Direct test evidence,mixed,408,520,112,"covered_by_family=26, directly_tested=415, hardware_gated=4, structural_only=75",tessera.compiler.generated_docs:test_coverage (live renderer),Convert structural_only and needs_direct_test rows into direct compare fixtures; keep hardware_gated tied to backend proof. -integration,Runtime execution matrix,closed,265,265,0,"apple_cpu=2, apple_gpu=24, cpu=6, nvidia_sm120=53, rocm=93, x86=87",docs/audit/generated/runtime_execution_matrix.csv,Add rows only when a launch path actually executes. +integration,Runtime execution matrix,closed,268,268,0,"apple_cpu=2, apple_gpu=24, cpu=7, nvidia_sm120=54, rocm=94, x86=87",docs/audit/generated/runtime_execution_matrix.csv,Add rows only when a launch path actually executes. integration,Runtime ABI symbols,mixed,571,900,329,"apple=719, nvidia=8, rocm=13, x86=160",docs/audit/generated/runtime_abi.csv,Reduce stub-only ABI rows where a backend claims native execution. integration,Audited repo surfaces,mixed,31,64,33,"archived=4, compile_only=18, runnable=31, runnable_optional=1, scaffold=10",docs/audit/generated/surface_status.csv,Graduate compile_only/scaffold entries that exercise compiler pathways; archive dead surfaces. -codegen_pathway,Portable CPU reference,closed,370,370,0,"manifest: exact_verified=0, implementation_present=0, reference=370, artifact_or_planned=0, other=0, missing_target_row=150; runtime_paths: executable=6/6 (cpu=6)",docs/audit/generated/s_series_status.md,Keep portable reference execution distinct from native x86 proof. +codegen_pathway,Portable CPU reference,closed,370,370,0,"manifest: exact_verified=0, implementation_present=0, reference=370, artifact_or_planned=0, other=0, missing_target_row=150; runtime_paths: executable=7/7 (cpu=7)",docs/audit/generated/s_series_status.md,Keep portable reference execution distinct from native x86 proof. codegen_pathway,x86 / AVX-512,mixed,294,349,55,"manifest: exact_verified=294, implementation_present=49, reference=6, artifact_or_planned=0, other=0, missing_target_row=171; runtime_paths: executable=87/87 (x86=87)",docs/audit/generated/s_series_status.md,Promote implementation-only rows with exact Zen 5 execute-and-compare; keep AMX separately gated. codegen_pathway,Apple CPU,mixed,398,402,4,"manifest: exact_verified=0, implementation_present=4, reference=398, artifact_or_planned=0, other=0, missing_target_row=118; runtime_paths: executable=2/2 (apple_cpu=2)",docs/audit/generated/s_series_status.md,Exact-device verify implementation-only rows or retain them explicitly as reference execution. codegen_pathway,Apple GPU,mixed,148,207,59,"manifest: exact_verified=148, implementation_present=51, reference=4, artifact_or_planned=4, other=0, missing_target_row=313; runtime_paths: executable=24/24 (apple_gpu=24)",docs/audit/generated/s_series_status.md,Promote implementation/artifact rows only with exact-device execute-and-compare. -codegen_pathway,ROCm / gfx1151,mixed,346,352,6,"manifest: exact_verified=346, implementation_present=0, reference=2, artifact_or_planned=4, other=0, missing_target_row=168; runtime_paths: executable=93/93 (rocm=93)",docs/audit/generated/s_series_status.md,Promote remaining reference rows only with exact gfx1151 execute-and-compare. +codegen_pathway,ROCm / gfx1151,mixed,346,352,6,"manifest: exact_verified=346, implementation_present=0, reference=2, artifact_or_planned=4, other=0, missing_target_row=168; runtime_paths: executable=94/94 (rocm=94)",docs/audit/generated/s_series_status.md,Promote remaining reference rows only with exact gfx1151 execute-and-compare. codegen_pathway,NVIDIA SM80,open,0,72,72,"manifest: exact_verified=0, implementation_present=0, reference=0, artifact_or_planned=72, other=0, missing_target_row=448; runtime_paths: executable=0/0 (no runtime rows)",docs/audit/generated/s_series_status.md,Retain as declared/open until architecture-owned execution evidence exists. codegen_pathway,NVIDIA SM90,open,0,104,104,"manifest: exact_verified=0, implementation_present=0, reference=2, artifact_or_planned=102, other=0, missing_target_row=416; runtime_paths: executable=0/0 (no runtime rows)",docs/audit/generated/s_series_status.md,Keep compile/artifact evidence separate from SM120 exact-device execution. codegen_pathway,NVIDIA SM100,open,0,72,72,"manifest: exact_verified=0, implementation_present=0, reference=0, artifact_or_planned=72, other=0, missing_target_row=448; runtime_paths: executable=0/0 (no runtime rows)",docs/audit/generated/s_series_status.md,Retain as declared/open until architecture-owned execution evidence exists. -codegen_pathway,NVIDIA SM120,mixed,69,134,65,"manifest: exact_verified=69, implementation_present=0, reference=0, artifact_or_planned=65, other=0, missing_target_row=386; runtime_paths: executable=53/53 (nvidia_sm120=53)",docs/audit/generated/s_series_status.md,Promote artifact rows with SM120 execute-and-compare evidence. +codegen_pathway,NVIDIA SM120,mixed,69,134,65,"manifest: exact_verified=69, implementation_present=0, reference=0, artifact_or_planned=65, other=0, missing_target_row=386; runtime_paths: executable=54/54 (nvidia_sm120=54)",docs/audit/generated/s_series_status.md,Promote artifact rows with SM120 execute-and-compare evidence. open_work,Target IR native/fused codegen,mixed,329,353,24,"device_verified_abi=27, device_verified_jit=196, fused=72, no_kernel_required=8, not_applicable=26, reference=24",docs/audit/generated/support_table.csv,Promote high-use reference rows into native/fused Target IR or mark intentional reference-only lanes. open_work,backend_kernel,mixed,100,520,420,"primitive contract axis; open means partial or planned, not necessarily missing API support",docs/audit/generated/s_series_status.md,Promote by backend/pathway; do not treat every target as an all-up compiler veto. open_work,Direct test evidence,mixed,408,520,112,"covered_by_family=26, directly_tested=415, hardware_gated=4, structural_only=75",tessera.compiler.generated_docs:test_coverage (live renderer),Convert structural_only and needs_direct_test rows into direct compare fixtures; keep hardware_gated tied to backend proof. diff --git a/docs/audit/generated/compiler_progress.md b/docs/audit/generated/compiler_progress.md index 188da0287..a82f4d611 100644 --- a/docs/audit/generated/compiler_progress.md +++ b/docs/audit/generated/compiler_progress.md @@ -43,7 +43,7 @@ A row is not marked incomplete merely because Apple, x86, ROCm, and CUDA are not |---|---|---:|---:|---:|---|---| | `Verifier coverage` | closed | 244 | 244 | 0 | real=244 | No action unless this row reopens. | | `Direct test evidence` | mixed | 408 | 520 | 112 | covered_by_family=26, directly_tested=415, hardware_gated=4, structural_only=75 | Convert structural_only and needs_direct_test rows into direct compare fixtures; keep hardware_gated tied to backend proof. | -| `Runtime execution matrix` | closed | 265 | 265 | 0 | apple_cpu=2, apple_gpu=24, cpu=6, nvidia_sm120=53, rocm=93, x86=87 | Add rows only when a launch path actually executes. | +| `Runtime execution matrix` | closed | 268 | 268 | 0 | apple_cpu=2, apple_gpu=24, cpu=7, nvidia_sm120=54, rocm=94, x86=87 | Add rows only when a launch path actually executes. | | `Runtime ABI symbols` | mixed | 571 | 900 | 329 | apple=719, nvidia=8, rocm=13, x86=160 | Reduce stub-only ABI rows where a backend claims native execution. | | `Audited repo surfaces` | mixed | 31 | 64 | 33 | archived=4, compile_only=18, runnable=31, runnable_optional=1, scaffold=10 | Graduate compile_only/scaffold entries that exercise compiler pathways; archive dead surfaces. | @@ -51,15 +51,15 @@ A row is not marked incomplete merely because Apple, x86, ROCm, and CUDA are not | Item | Status | Ready | Total | Open | Detail | Next | |---|---|---:|---:|---:|---|---| -| `Portable CPU reference` | closed | 370 | 370 | 0 | manifest: exact_verified=0, implementation_present=0, reference=370, artifact_or_planned=0, other=0, missing_target_row=150; runtime_paths: executable=6/6 (cpu=6) | Keep portable reference execution distinct from native x86 proof. | +| `Portable CPU reference` | closed | 370 | 370 | 0 | manifest: exact_verified=0, implementation_present=0, reference=370, artifact_or_planned=0, other=0, missing_target_row=150; runtime_paths: executable=7/7 (cpu=7) | Keep portable reference execution distinct from native x86 proof. | | `x86 / AVX-512` | mixed | 294 | 349 | 55 | manifest: exact_verified=294, implementation_present=49, reference=6, artifact_or_planned=0, other=0, missing_target_row=171; runtime_paths: executable=87/87 (x86=87) | Promote implementation-only rows with exact Zen 5 execute-and-compare; keep AMX separately gated. | | `Apple CPU` | mixed | 398 | 402 | 4 | manifest: exact_verified=0, implementation_present=4, reference=398, artifact_or_planned=0, other=0, missing_target_row=118; runtime_paths: executable=2/2 (apple_cpu=2) | Exact-device verify implementation-only rows or retain them explicitly as reference execution. | | `Apple GPU` | mixed | 148 | 207 | 59 | manifest: exact_verified=148, implementation_present=51, reference=4, artifact_or_planned=4, other=0, missing_target_row=313; runtime_paths: executable=24/24 (apple_gpu=24) | Promote implementation/artifact rows only with exact-device execute-and-compare. | -| `ROCm / gfx1151` | mixed | 346 | 352 | 6 | manifest: exact_verified=346, implementation_present=0, reference=2, artifact_or_planned=4, other=0, missing_target_row=168; runtime_paths: executable=93/93 (rocm=93) | Promote remaining reference rows only with exact gfx1151 execute-and-compare. | +| `ROCm / gfx1151` | mixed | 346 | 352 | 6 | manifest: exact_verified=346, implementation_present=0, reference=2, artifact_or_planned=4, other=0, missing_target_row=168; runtime_paths: executable=94/94 (rocm=94) | Promote remaining reference rows only with exact gfx1151 execute-and-compare. | | `NVIDIA SM80` | open | 0 | 72 | 72 | manifest: exact_verified=0, implementation_present=0, reference=0, artifact_or_planned=72, other=0, missing_target_row=448; runtime_paths: executable=0/0 (no runtime rows) | Retain as declared/open until architecture-owned execution evidence exists. | | `NVIDIA SM90` | open | 0 | 104 | 104 | manifest: exact_verified=0, implementation_present=0, reference=2, artifact_or_planned=102, other=0, missing_target_row=416; runtime_paths: executable=0/0 (no runtime rows) | Keep compile/artifact evidence separate from SM120 exact-device execution. | | `NVIDIA SM100` | open | 0 | 72 | 72 | manifest: exact_verified=0, implementation_present=0, reference=0, artifact_or_planned=72, other=0, missing_target_row=448; runtime_paths: executable=0/0 (no runtime rows) | Retain as declared/open until architecture-owned execution evidence exists. | -| `NVIDIA SM120` | mixed | 69 | 134 | 65 | manifest: exact_verified=69, implementation_present=0, reference=0, artifact_or_planned=65, other=0, missing_target_row=386; runtime_paths: executable=53/53 (nvidia_sm120=53) | Promote artifact rows with SM120 execute-and-compare evidence. | +| `NVIDIA SM120` | mixed | 69 | 134 | 65 | manifest: exact_verified=69, implementation_present=0, reference=0, artifact_or_planned=65, other=0, missing_target_row=386; runtime_paths: executable=54/54 (nvidia_sm120=54) | Promote artifact rows with SM120 execute-and-compare evidence. | ## Open Work Summary diff --git a/docs/audit/generated/docs_freshness.md b/docs/audit/generated/docs_freshness.md index a368eb481..a3f0ba72f 100644 --- a/docs/audit/generated/docs_freshness.md +++ b/docs/audit/generated/docs_freshness.md @@ -6,9 +6,9 @@ Reference date for staleness: **2026-09-16**. ## Headline -- **157** docs catalogued across the canonical doc tree. -- **156** carry a `last_updated:` marker; **1** are undated (invisible to the freshness audit until tagged). -- **57** updated within the last 30 days. +- **158** docs catalogued across the canonical doc tree. +- **157** carry a `last_updated:` marker; **1** are undated (invisible to the freshness audit until tagged). +- **58** updated within the last 30 days. - **38** older than 90 days; **0** older than 180 days. ## Undated docs (no parseable `last_updated`) @@ -206,6 +206,7 @@ These docs need either YAML frontmatter (`last_updated: YYYY-MM-DD`) or a body-f | `compiler/compiler_enhancement.md` | - | 2026-09-08 | 8 | ✓ | | `coverage/COVERAGE_AUDIT.md` | - | 2026-09-04 | 12 | ✓ | | `domain/DOMAIN_AUDIT.md` | - | 2026-09-15 | 1 | ✓ | +| `domain/EBM_NATIVE_LOOP_ARCHITECTURE.md` | - | 2026-09-16 | 0 | ✓ | | `domain/GA_EBM_ARCHITECTURE_REVIEW.md` | - | 2026-09-15 | 1 | ✓ | | `roadmap/CF_CROSS_ELEMENT_PLAN.md` | - | 2026-06-30 | 78 | ✓ | | `roadmap/MODEL_CLASS_ROADMAP.md` | - | 2026-08-12 | 35 | ✓ | diff --git a/docs/audit/generated/domain_proof_ladder.csv b/docs/audit/generated/domain_proof_ladder.csv index 70d93e524..6b58e7e17 100644 --- a/docs/audit/generated/domain_proof_ladder.csv +++ b/docs/audit/generated/domain_proof_ladder.csv @@ -1,6 +1,6 @@ domain,label,registry_total,registry_planned,ledger_total,ledger_adjoint,ledger_device_verified,native_targets,owners geometric_algebra,Geometric algebra / Clifford,18,17,17,0,0,apple_gpu=1;cpu=1;nvidia_sm120=1;rocm=2;x86=1,W6.4(owner);AD-HIGHER-1(owner) -energy_based_models,Energy-based models,16,5,4,0,0,apple_gpu=2;rocm=3;x86=3,W4-PRODUCT-1(owner);AD-SOLVER-IFT-1(owner);AD-RESIDUAL-EVAL-1(owner) +energy_based_models,Energy-based models,16,5,4,0,0,apple_gpu=2;cpu=1;nvidia_sm120=1;rocm=4;x86=3,W4-PRODUCT-1(owner);AD-SOLVER-IFT-1(owner);AD-RESIDUAL-EVAL-1(owner) attention_persistent_state,Attention / persistent state,25,0,20,1,7,apple_gpu=1;nvidia_sm120=8;rocm=8;x86=11,W5.2(owner);AD-RESIDUAL-EVAL-1(owner);W2.4a(owner) field_pde_spectral,"Matrix/field calculus, PDE and spectral",13,0,7,6,2,nvidia_sm120=5;rocm=4;x86=4,MSW-9(owner);TSOL-POLICY-PHYS-1(owner);TSOL-PHYS-TAIL-1(owner);PDE-STENCIL-FOUNDATION-1(owner) game_theory_structured,Game theory / structured contractions,9,0,8,0,0,rocm=1,TSOL-PHYS-TAIL-1(owner) diff --git a/docs/audit/generated/domain_proof_ladder.md b/docs/audit/generated/domain_proof_ladder.md index 9b6b1ab6e..dedd91ee1 100644 --- a/docs/audit/generated/domain_proof_ladder.md +++ b/docs/audit/generated/domain_proof_ladder.md @@ -18,7 +18,7 @@ is deliberately not summarized here. | Domain | Registry rows (of which still `planned`) | AD ledger rows (adjoint / device) | Native executable rows per target | Plan owners | |---|---|---|---|---| | Geometric algebra / Clifford | 18 (17) | 17 (0 / 0) | apple_gpu=1, cpu=1, nvidia_sm120=1, rocm=2, x86=1 | `W6.4(owner)`, `AD-HIGHER-1(owner)` | -| Energy-based models | 16 (5) | 4 (0 / 0) | apple_gpu=2, rocm=3, x86=3 | `W4-PRODUCT-1(owner)`, `AD-SOLVER-IFT-1(owner)`, `AD-RESIDUAL-EVAL-1(owner)` | +| Energy-based models | 16 (5) | 4 (0 / 0) | apple_gpu=2, cpu=1, nvidia_sm120=1, rocm=4, x86=3 | `W4-PRODUCT-1(owner)`, `AD-SOLVER-IFT-1(owner)`, `AD-RESIDUAL-EVAL-1(owner)` | | Attention / persistent state | 25 (0) | 20 (1 / 7) | apple_gpu=1, nvidia_sm120=8, rocm=8, x86=11 | `W5.2(owner)`, `AD-RESIDUAL-EVAL-1(owner)`, `W2.4a(owner)` | | Matrix/field calculus, PDE and spectral | 13 (0) | 7 (6 / 2) | nvidia_sm120=5, rocm=4, x86=4 | `MSW-9(owner)`, `TSOL-POLICY-PHYS-1(owner)`, `TSOL-PHYS-TAIL-1(owner)`, `PDE-STENCIL-FOUNDATION-1(owner)` | | Game theory / structured contractions | 9 (0) | 8 (0 / 0) | rocm=1 | `TSOL-PHYS-TAIL-1(owner)` | diff --git a/docs/audit/generated/runtime_execution_matrix.csv b/docs/audit/generated/runtime_execution_matrix.csv index 8c2907c01..d996df658 100644 --- a/docs/audit/generated/runtime_execution_matrix.csv +++ b/docs/audit/generated/runtime_execution_matrix.csv @@ -29,6 +29,7 @@ cpu,cpu_autodiff_matmul_llvm_jit,native_cpu,1,cpu_autodiff_paired_llvm_jit,succe cpu,cpu_autodiff_sigmoid_llvm_jit,native_cpu,1,cpu_autodiff_paired_llvm_jit,success,mlir_llvm_jit,Compiler-generated sigmoid adjoint in the paired matmul activation slice launches through MLIR/LLVM JIT and matches NumPy. cpu,cpu_autodiff_tanh_llvm_jit,native_cpu,1,cpu_autodiff_paired_llvm_jit,success,mlir_llvm_jit,Compiler-generated tanh adjoint in the paired matmul activation slice launches through MLIR/LLVM JIT and matches NumPy. cpu,cpu_clifford_llvm_jit,native_cpu,1,cpu_clifford_llvm_jit,success,mlir_llvm_jit,"W6.4 batched native GA: the tessera_clifford product family (geo_product with optional output-grade restriction, wedge, left_contract, inner, norm, reverse, grade_involute, conjugate, hodge_star, grade, rotor_sandwich) on [..., 2**n] f32 tensors of Cl(p,q,r) lowers through GradeFusion + ExpandProductTable to an scf.for nest over the compile-time Cayley table inside libtessera_jit and executes through MLIR/LLVM on the host CPU (M1 Max, Zen 5, Zen 2 parity with the standalone GA reference, 2026-09-16). No numpy fallback; out-of-envelope requests raise." +cpu,cpu_ebm_langevin_llvm_jit,native_cpu,1,cpu_ebm_langevin_llvm_jit,success,mlir_llvm_jit,"W4-PRODUCT-1 quadratic energy loop (2026-09-16): E(y, x) = 0.5*||x - y||^2 is a Graph IR function marked for reverse-mode; inside libtessera_jit the paired autodiff pass derives @E__bwd, the EBM lowering turns each tessera_ebm.langevin_step into y - eta*grad + sqrt(2 eta T)*z with Philox-4x32-10 / Box-Muller noise generated in a linalg.generic, and the K-step scf.for compiles as one function: no per-step host gradient or noise transfers. Bit-exact with the declared numpy policy on the M1 Max and the x86 hosts. No fallback." cpu,jit_cpu_numpy,reference_cpu,1,jit_cpu_numpy,success,,CPU JIT artifact runs through the numpy reference path. cpu,native_cpu,native_cpu,1,native_cpu,success,,CPU artifact runs through the x86 AMX / native CPU runtime. nvidia_sm120,nvidia_adafactor_bwd_compiled,native_gpu,1,nvidia_adafactor_bwd_compiled,success,cuda_driver,SM120 PTX executes full or factored Adafactor VJP; factored reductions have one deterministic owner. @@ -44,6 +45,7 @@ nvidia_sm120,nvidia_deltanet_bwd_compiled,native_gpu,1,nvidia_deltanet_bwd_compi nvidia_sm120,nvidia_deltanet_compiled,native_gpu,1,nvidia_deltanet_compiled,success,cuda_runtime,"NVIDIA sm_120 causal DeltaNet recurrence covering plain, gated, beta/decay, erase, and modified updates in one CUDA launch." nvidia_sm120,nvidia_dense_krylov_compiled,native_gpu,1,nvidia_dense_krylov_compiled,success,cuda_runtime,"A cooperative CUDA grid executes dense non-diagonal CG or twice-MGS restarted GMRES. Krylov/Arnoldi state remains on device, dot/norm reductions span multiple CTAs in a fixed order, and only an fp32 true residual can establish convergence." nvidia_sm120,nvidia_dequant_gemm_compiled,native_gpu,1,nvidia_dequant_gemm_compiled,success,cuda_runtime,"NVIDIA sm_120 fused int4/int8 packed-weight dequantization inside single and grouped GEMM, with per-group scales and device offsets." +nvidia_sm120,nvidia_ebm_langevin_native_compiled,native_gpu,1,nvidia_ebm_langevin_native_compiled,success,cuda_runtime,"EBM quadratic-energy Langevin loop on sm_120 (2026-09-16): the K-step scf.for over the compiler-derived gradient becomes one cooperative kernel through the row-program emitter (one block per row, one lane per feature, Philox inside the kernel), launched through the native storage package. One launch per loop; bit-exact with the declared numpy policy. No Python-emitted kernel; no fallback." nvidia_sm120,nvidia_fft_compiled,native_gpu,1,nvidia_fft_compiled,success,cuda_runtime,Canonical CUDA FFT/workspace v2 package executes C2C/R2C/C2R FFT/IFFT/RFFT/IRFFT through reusable cuFFT plans with auto-allocation disabled and explicit caller-owned device workspace. nvidia_sm120,nvidia_flash_attn_bwd_compiled,native_gpu,1,nvidia_flash_attn_bwd_compiled,success,cuda_runtime,"NVIDIA sm_120 compiler-emitted Flash Attention backward: f32/fp16/bf16 storage, f32 dQ/dK/dV accumulation recomputes online-softmax statistics, atomically accumulates shared GQA/MQA dK/dV, and differentiates causal/window masks, dense additive bias, and logit soft-cap. The P0 f32 saved-LSE form consumes an explicit row_lse pointer and is measured separately from recompute." nvidia_sm120,nvidia_flash_attn_compiled,native_gpu,1,nvidia_flash_attn_compiled,success,cuda_runtime,"NVIDIA sm_120 compiler-emitted Flash Attention forward: f32/fp16/bf16 storage, f32 accumulation, MHA/GQA/MQA KV-head mapping, online softmax, causal/window masks, dense additive bias and logit soft-cap; one CUDA launch over all B*Hq*Sq query rows. The P0 f32 checkpoint form has a distinct O,row_lse physical ABI; the default remains no-save/recompute." @@ -113,6 +115,7 @@ rocm,rocm_distribution_loss_bwd_compiled,native_gpu,1,rocm_distribution_loss_bwd rocm,rocm_dspark_draft_block_compiled,native_gpu,1,rocm_dspark_draft_block_compiled,success,hip_runtime,"ROCm DSpark draft-block artifact launches the compiler-generated fused HIP/ROCDL draft-block kernel (generate-rocm-dspark-draft-block-kernel) for logits, confidence logits, greedy tokens, and hidden states. The executor falls back to the DS1 oracle only when ROCm hardware or tessera-opt is unavailable." rocm,rocm_ebm_compute_compiled,native_gpu,1,rocm_ebm_compute_compiled,success,hip_runtime,"ROCm EBM compute artifact runs energy_quadratic / inner_step / refinement / self_verify: the diff/square and reductions run on the gfx1151 binary + reduce kernels, the structure on the host. ROCm mirror of x86_ebm_compute. f32." rocm,rocm_ebm_langevin_compiled,native_gpu,1,rocm_ebm_langevin_compiled,success,hip_runtime,"ROCm EBM Langevin sampling artifact runs y − η·grad + noise_scale·z where z is Box-Muller Gaussian noise drawn ON-DEVICE from counter-based Philox-4x32-10 via the COMPILER-GENERATED gfx1151 kernel (generate-rocm-ebm-langevin-kernel). f32, matches tessera.ebm.langevin_step_philox." +rocm,rocm_ebm_langevin_native_compiled,native_gpu,1,rocm_ebm_langevin_native_compiled,success,hip_runtime,"EBM quadratic-energy Langevin loop on gfx1151/gfx1201 (2026-09-16): the paired autodiff pass derives the gradient, the EBM lowering emits the step with Philox noise, tessera-to-linalg + the row-program emitter turn the K-step scf.for into ONE cooperative kernel (one block per row, one lane per feature, state in registers across steps, ordered shared-memory reductions), and the native storage package launches it. One launch per loop; no per-step host gradient or noise transfer; bit-exact with the declared numpy policy. No Python-emitted kernel; no fallback." rocm,rocm_ebm_loss_compiled,native_gpu,1,rocm_ebm_loss_compiled,success,hip_runtime,"ROCm EBM/diffusion loss artifact runs score_matching / denoising / implicit / contrastive_divergence / persistent_cd / ddpm_noise_pred / vlb / load_balance: the diff/square and reductions run on the gfx1151 binary + reduce kernels, the structure on the host. ROCm mirror of x86_ebm_loss. f32." rocm,rocm_exotic_attn_compiled,native_gpu,1,rocm_exotic_attn_compiled,success,hip_runtime,"ROCm exotic-attention artifact composes the COMPILER-GENERATED WMMA flash_attn kernel with the WMMA GEMM kernel (MLA latent projections) + an elementwise gate — gated_attention, mla_decode, mla_decode_fused — and routes mla_decode_step through the DK1 absorbed-latent decode kernel against stdlib.attention.mla_decode_step. f16/bf16, f32 softmax+accumulate." rocm,rocm_fft_compiled,native_gpu,1,rocm_fft_compiled,success,hip_runtime,"ROCm FFT artifact runs fft / ifft / rfft / irfft over any axis length on the shipping mixed-radix Stockham/Bluestein package on gfx1151, with batched rows, r2c/c2r pack-unpack, and plan-owned normalization. The O(N^2) compiler DFT is retained only as an explicitly named diagnostic/oracle. complex64/f32." diff --git a/docs/audit/generated/runtime_execution_matrix.md b/docs/audit/generated/runtime_execution_matrix.md index 3682de0d7..0e9e7b303 100644 --- a/docs/audit/generated/runtime_execution_matrix.md +++ b/docs/audit/generated/runtime_execution_matrix.md @@ -43,6 +43,7 @@ Single source of truth for what `runtime.launch()` does with each `(target, comp | `cpu` | `cpu_autodiff_sigmoid_llvm_jit` | `cpu_autodiff_paired_llvm_jit` | `native_cpu` | `mlir_llvm_jit` | Compiler-generated sigmoid adjoint in the paired matmul activation slice launches through MLIR/LLVM JIT and matches NumPy. | | `cpu` | `cpu_autodiff_tanh_llvm_jit` | `cpu_autodiff_paired_llvm_jit` | `native_cpu` | `mlir_llvm_jit` | Compiler-generated tanh adjoint in the paired matmul activation slice launches through MLIR/LLVM JIT and matches NumPy. | | `cpu` | `cpu_clifford_llvm_jit` | `cpu_clifford_llvm_jit` | `native_cpu` | `mlir_llvm_jit` | W6.4 batched native GA: the tessera_clifford product family (geo_product with optional output-grade restriction, wedge, left_contract, inner, norm, reverse, grade_involute, conjugate, hodge_star, grade, rotor_sandwich) on [..., 2**n] f32 tensors of Cl(p,q,r) lowers through GradeFusion + ExpandProductTable to an scf.for nest over the compile-time Cayley table inside libtessera_jit and executes through MLIR/LLVM on the host CPU (M1 Max, Zen 5, Zen 2 parity with the standalone GA reference, 2026-09-16). No numpy fallback; out-of-envelope requests raise. | +| `cpu` | `cpu_ebm_langevin_llvm_jit` | `cpu_ebm_langevin_llvm_jit` | `native_cpu` | `mlir_llvm_jit` | W4-PRODUCT-1 quadratic energy loop (2026-09-16): E(y, x) = 0.5*||x - y||^2 is a Graph IR function marked for reverse-mode; inside libtessera_jit the paired autodiff pass derives @E__bwd, the EBM lowering turns each tessera_ebm.langevin_step into y - eta*grad + sqrt(2 eta T)*z with Philox-4x32-10 / Box-Muller noise generated in a linalg.generic, and the K-step scf.for compiles as one function: no per-step host gradient or noise transfers. Bit-exact with the declared numpy policy on the M1 Max and the x86 hosts. No fallback. | | `cpu` | `jit_cpu_numpy` | `jit_cpu_numpy` | `reference_cpu` | - | CPU JIT artifact runs through the numpy reference path. | | `cpu` | `native_cpu` | `native_cpu` | `native_cpu` | - | CPU artifact runs through the x86 AMX / native CPU runtime. | | `nvidia_sm120` | `nvidia_adafactor_bwd_compiled` | `nvidia_adafactor_bwd_compiled` | `native_gpu` | `cuda_driver` | SM120 PTX executes full or factored Adafactor VJP; factored reductions have one deterministic owner. | @@ -58,6 +59,7 @@ Single source of truth for what `runtime.launch()` does with each `(target, comp | `nvidia_sm120` | `nvidia_deltanet_compiled` | `nvidia_deltanet_compiled` | `native_gpu` | `cuda_runtime` | NVIDIA sm_120 causal DeltaNet recurrence covering plain, gated, beta/decay, erase, and modified updates in one CUDA launch. | | `nvidia_sm120` | `nvidia_dense_krylov_compiled` | `nvidia_dense_krylov_compiled` | `native_gpu` | `cuda_runtime` | A cooperative CUDA grid executes dense non-diagonal CG or twice-MGS restarted GMRES. Krylov/Arnoldi state remains on device, dot/norm reductions span multiple CTAs in a fixed order, and only an fp32 true residual can establish convergence. | | `nvidia_sm120` | `nvidia_dequant_gemm_compiled` | `nvidia_dequant_gemm_compiled` | `native_gpu` | `cuda_runtime` | NVIDIA sm_120 fused int4/int8 packed-weight dequantization inside single and grouped GEMM, with per-group scales and device offsets. | +| `nvidia_sm120` | `nvidia_ebm_langevin_native_compiled` | `nvidia_ebm_langevin_native_compiled` | `native_gpu` | `cuda_runtime` | EBM quadratic-energy Langevin loop on sm_120 (2026-09-16): the K-step scf.for over the compiler-derived gradient becomes one cooperative kernel through the row-program emitter (one block per row, one lane per feature, Philox inside the kernel), launched through the native storage package. One launch per loop; bit-exact with the declared numpy policy. No Python-emitted kernel; no fallback. | | `nvidia_sm120` | `nvidia_fft_compiled` | `nvidia_fft_compiled` | `native_gpu` | `cuda_runtime` | Canonical CUDA FFT/workspace v2 package executes C2C/R2C/C2R FFT/IFFT/RFFT/IRFFT through reusable cuFFT plans with auto-allocation disabled and explicit caller-owned device workspace. | | `nvidia_sm120` | `nvidia_flash_attn_bwd_compiled` | `nvidia_flash_attn_bwd_compiled` | `native_gpu` | `cuda_runtime` | NVIDIA sm_120 compiler-emitted Flash Attention backward: f32/fp16/bf16 storage, f32 dQ/dK/dV accumulation recomputes online-softmax statistics, atomically accumulates shared GQA/MQA dK/dV, and differentiates causal/window masks, dense additive bias, and logit soft-cap. The P0 f32 saved-LSE form consumes an explicit row_lse pointer and is measured separately from recompute. | | `nvidia_sm120` | `nvidia_flash_attn_compiled` | `nvidia_flash_attn_compiled` | `native_gpu` | `cuda_runtime` | NVIDIA sm_120 compiler-emitted Flash Attention forward: f32/fp16/bf16 storage, f32 accumulation, MHA/GQA/MQA KV-head mapping, online softmax, causal/window masks, dense additive bias and logit soft-cap; one CUDA launch over all B*Hq*Sq query rows. The P0 f32 checkpoint form has a distinct O,row_lse physical ABI; the default remains no-save/recompute. | @@ -127,6 +129,7 @@ Single source of truth for what `runtime.launch()` does with each `(target, comp | `rocm` | `rocm_dspark_draft_block_compiled` | `rocm_dspark_draft_block_compiled` | `native_gpu` | `hip_runtime` | ROCm DSpark draft-block artifact launches the compiler-generated fused HIP/ROCDL draft-block kernel (generate-rocm-dspark-draft-block-kernel) for logits, confidence logits, greedy tokens, and hidden states. The executor falls back to the DS1 oracle only when ROCm hardware or tessera-opt is unavailable. | | `rocm` | `rocm_ebm_compute_compiled` | `rocm_ebm_compute_compiled` | `native_gpu` | `hip_runtime` | ROCm EBM compute artifact runs energy_quadratic / inner_step / refinement / self_verify: the diff/square and reductions run on the gfx1151 binary + reduce kernels, the structure on the host. ROCm mirror of x86_ebm_compute. f32. | | `rocm` | `rocm_ebm_langevin_compiled` | `rocm_ebm_langevin_compiled` | `native_gpu` | `hip_runtime` | ROCm EBM Langevin sampling artifact runs y − η·grad + noise_scale·z where z is Box-Muller Gaussian noise drawn ON-DEVICE from counter-based Philox-4x32-10 via the COMPILER-GENERATED gfx1151 kernel (generate-rocm-ebm-langevin-kernel). f32, matches tessera.ebm.langevin_step_philox. | +| `rocm` | `rocm_ebm_langevin_native_compiled` | `rocm_ebm_langevin_native_compiled` | `native_gpu` | `hip_runtime` | EBM quadratic-energy Langevin loop on gfx1151/gfx1201 (2026-09-16): the paired autodiff pass derives the gradient, the EBM lowering emits the step with Philox noise, tessera-to-linalg + the row-program emitter turn the K-step scf.for into ONE cooperative kernel (one block per row, one lane per feature, state in registers across steps, ordered shared-memory reductions), and the native storage package launches it. One launch per loop; no per-step host gradient or noise transfer; bit-exact with the declared numpy policy. No Python-emitted kernel; no fallback. | | `rocm` | `rocm_ebm_loss_compiled` | `rocm_ebm_loss_compiled` | `native_gpu` | `hip_runtime` | ROCm EBM/diffusion loss artifact runs score_matching / denoising / implicit / contrastive_divergence / persistent_cd / ddpm_noise_pred / vlb / load_balance: the diff/square and reductions run on the gfx1151 binary + reduce kernels, the structure on the host. ROCm mirror of x86_ebm_loss. f32. | | `rocm` | `rocm_exotic_attn_compiled` | `rocm_exotic_attn_compiled` | `native_gpu` | `hip_runtime` | ROCm exotic-attention artifact composes the COMPILER-GENERATED WMMA flash_attn kernel with the WMMA GEMM kernel (MLA latent projections) + an elementwise gate — gated_attention, mla_decode, mla_decode_fused — and routes mla_decode_step through the DK1 absorbed-latent decode kernel against stdlib.attention.mla_decode_step. f16/bf16, f32 softmax+accumulate. | | `rocm` | `rocm_fft_compiled` | `rocm_fft_compiled` | `native_gpu` | `hip_runtime` | ROCm FFT artifact runs fft / ifft / rfft / irfft over any axis length on the shipping mixed-radix Stockham/Bluestein package on gfx1151, with batched rows, r2c/c2r pack-unpack, and plan-owned normalization. The O(N^2) compiler DFT is retained only as an explicitly named diagnostic/oracle. complex64/f32. | @@ -318,6 +321,7 @@ nvidia_sm80, nvidia_sm90, nvidia_sm100, rocm_gfx90a, rocm_gfx940, rocm_gfx942, r | `apple_value_target_ir` | Apple CPU value-call dispatch — invokes the C ABI symbol named in a tessera_apple.cpu.call value op (Value Target IR sprint; CPU cholesky executable) | | `cpu_autodiff_paired_llvm_jit` | Compiler-generated paired backward compiled through MLIR/LLVM and invoked through libtessera_jit | | `cpu_clifford_llvm_jit` | Geometric-algebra products lowered by the Clifford dialect (GradeFusion + batched ExpandProductTable) inside libtessera_jit and executed through MLIR/LLVM | +| `cpu_ebm_langevin_llvm_jit` | The EBM quadratic-energy Langevin loop: the paired autodiff pass derives the gradient, the EBM lowering emits the step with Philox noise, the whole scf.for compiles as one function in libtessera_jit | | `jit_cpu_numpy` | JIT CPU fallback via the numpy reference path | | `native_cpu` | x86 AMX / native CPU runtime via the C runtime ABI | | `nvidia_adafactor_bwd_compiled` | NVIDIA SM120 full/factored Adafactor VJP with deterministic reductions | @@ -333,6 +337,7 @@ nvidia_sm80, nvidia_sm90, nvidia_sm100, rocm_gfx90a, rocm_gfx940, rocm_gfx942, r | `nvidia_deltanet_compiled` | NVIDIA GPU (consumer Blackwell sm_120) native single-launch f32 DeltaNet recurrence with gate, beta, decay, erase, and modified variants | | `nvidia_dense_krylov_compiled` | Single cooperative-grid SM120 dense-operator CG or restarted GMRES package with resident Arnoldi state and deterministic multi-CTA reductions | | `nvidia_dequant_gemm_compiled` | NVIDIA GPU (consumer Blackwell sm_120) fused int4/int8 per-group dequantization inside grouped GEMM | +| `nvidia_ebm_langevin_native_compiled` | The EBM Langevin loop as one cooperative sm_120 kernel from the row-program emitter over the compiler-derived gradient | | `nvidia_fft_compiled` | NVIDIA GPU (consumer Blackwell sm_120) canonical C2C/R2C/C2R cuFFT package with reusable plans and explicit caller-owned device workspace | | `nvidia_flash_attn_bwd_compiled` | NVIDIA GPU (consumer Blackwell sm_120) compiler-emitted Flash Attention VJP; f32/fp16/bf16 storage and f32 dQ/dK/dV accumulation with GQA/MQA accumulation, mask/bias/soft-cap derivatives, launched through runtime.launch() | | `nvidia_flash_attn_compiled` | NVIDIA GPU (consumer Blackwell sm_120) compiler-emitted Flash Attention forward; f32/fp16/bf16 storage, f32 accumulation, MHA/GQA/MQA with online softmax, causal/sliding-window masks, dense additive bias and logit soft-cap, launched through runtime.launch() | @@ -402,6 +407,7 @@ nvidia_sm80, nvidia_sm90, nvidia_sm100, rocm_gfx90a, rocm_gfx940, rocm_gfx942, r | `rocm_dspark_draft_block_compiled` | DS2 DSpark draft-block ROCm compiler path — compiler-generated fused HIP/ROCDL draft-block kernel (generate-rocm-dspark-draft-block-kernel) with DS1 oracle fallback when ROCm hardware is unavailable. f32/i64 | | `rocm_ebm_compute_compiled` | AMD GPU RDNA EBM compute (energy_quadratic / inner_step / refinement / self_verify) — diff/square/reduce on the gfx1151 binary + reduce kernels, host structure. ROCm mirror of x86_ebm_compute. f32 | | `rocm_ebm_langevin_compiled` | AMD GPU RDNA EBM Langevin sampling — y − η·grad + noise_scale·z with z drawn on-device from Philox-4x32-10 Box-Muller via the COMPILER-GENERATED gfx1151 kernel. f32 | +| `rocm_ebm_langevin_native_compiled` | The EBM Langevin loop as one cooperative gfx1151/gfx1201 kernel from the row-program emitter over the compiler-derived gradient (one block per row, lanes per feature, K steps and Philox inside the kernel) | | `rocm_ebm_loss_compiled` | AMD GPU RDNA EBM/diffusion loss (score_matching / denoising / implicit / contrastive_divergence / persistent_cd / ddpm_noise_pred / vlb / load_balance) — diff/square + reductions on the gfx1151 binary + reduce kernels, host structure. ROCm mirror of x86_ebm_loss. f32 | | `rocm_exotic_attn_compiled` | AMD GPU RDNA exotic-attention compositions (gated_attention, mla_decode, mla_decode_fused, mla_decode_step absorbed-latent decode) built by COMPOSING the COMPILER-GENERATED WMMA flash_attn kernel + the WMMA GEMM kernel (MLA latent projections) + an elementwise gate, plus the DK1 generated absorbed-latent ROCm decode kernel; f16/bf16 storage, f32 softmax+accumulate. The block-sparse deepseek variant stays artifact_only | | `rocm_fft_compiled` | AMD GPU RDNA spectral FFT (fft / ifft / rfft / irfft) through the shipping gfx1151 mixed-radix Stockham/Bluestein package, with batched rows, r2c/c2r pack-unpack, and plan-owned normalization. The compiler-generated direct DFT is diagnostic only. complex64/f32 | diff --git a/python/tessera/compiler/execution_matrix.py b/python/tessera/compiler/execution_matrix.py index 44a8fecc2..32e9e0fcb 100644 --- a/python/tessera/compiler/execution_matrix.py +++ b/python/tessera/compiler/execution_matrix.py @@ -122,6 +122,17 @@ class BackwardComposition: "nvidia_clifford_native_compiled": "Clifford ops lowered by the dialect's " "ExpandProductTable into a per-thread kernel, " "packaged by the arena pipeline and launched on sm_120", + "rocm_ebm_langevin_native_compiled": "The EBM Langevin loop as one cooperative " + "gfx1151/gfx1201 kernel from the row-program emitter over the " + "compiler-derived gradient (one block per row, lanes per feature, " + "K steps and Philox inside the kernel)", + "nvidia_ebm_langevin_native_compiled": "The EBM Langevin loop as one cooperative " + "sm_120 kernel from the row-program emitter over the " + "compiler-derived gradient", + "cpu_ebm_langevin_llvm_jit": "The EBM quadratic-energy Langevin loop: the paired " + "autodiff pass derives the gradient, the EBM lowering " + "emits the step with Philox noise, the whole scf.for " + "compiles as one function in libtessera_jit", "cpu_clifford_llvm_jit": "Geometric-algebra products lowered by the Clifford " "dialect (GradeFusion + batched ExpandProductTable) " "inside libtessera_jit and executed through MLIR/LLVM", @@ -1254,6 +1265,23 @@ class BackwardComposition: # to KNOWN_EXECUTORS, (3) adding an ExecutionRow here. `launch()` picks it up # automatically; the dashboard regenerates; the drift test enforces it. _MATRIX: dict[tuple[str, str], ExecutionRow] = { + ("cpu", "cpu_ebm_langevin_llvm_jit"): ExecutionRow( + target="cpu", compiler_path="cpu_ebm_langevin_llvm_jit", + execution_kind="native_cpu", executable=True, + executor_id="cpu_ebm_langevin_llvm_jit", runtime_status="success", + reason="W4-PRODUCT-1 quadratic energy loop (2026-09-16): E(y, x) = " + "0.5*||x - y||^2 is a Graph IR function marked for reverse-mode; " + "inside libtessera_jit the paired autodiff pass derives @E__bwd, " + "the EBM lowering turns each tessera_ebm.langevin_step into " + "y - eta*grad + sqrt(2 eta T)*z with Philox-4x32-10 / Box-Muller " + "noise generated in a linalg.generic, and the K-step scf.for " + "compiles as one function: no per-step host gradient or noise " + "transfers. Bit-exact with the declared numpy policy on the M1 " + "Max and the x86 hosts. No fallback.", + execution_mode="mlir_llvm_jit", op_family="ebm_langevin", + evidence_target="cpu_x86_64", + numerical_fixture="tests/unit/test_ebm_native_langevin.py", + proof_build="llvm23-core+ebm+jit"), ("cpu", "cpu_clifford_llvm_jit"): ExecutionRow( target="cpu", compiler_path="cpu_clifford_llvm_jit", execution_kind="native_cpu", executable=True, @@ -2948,6 +2976,22 @@ class BackwardComposition: "row to a power of two and flips for descending). f32, matches " "numpy.", execution_mode="hip_runtime"), + ("rocm", "rocm_ebm_langevin_native_compiled"): ExecutionRow( + target="rocm", compiler_path="rocm_ebm_langevin_native_compiled", + execution_kind="native_gpu", executable=True, + executor_id="rocm_ebm_langevin_native_compiled", runtime_status="success", + reason="EBM quadratic-energy Langevin loop on gfx1151/gfx1201 (2026-09-16): the " + "paired autodiff pass derives the gradient, the EBM lowering emits the " + "step with Philox noise, tessera-to-linalg + the row-program emitter " + "turn the K-step scf.for into ONE cooperative kernel (one block per row, " + "one lane per feature, state in registers across steps, ordered shared-" + "memory reductions), and the native storage package launches it. One " + "launch per loop; no per-step host gradient or noise transfer; bit-exact " + "with the declared numpy policy. No Python-emitted kernel; no fallback.", + execution_mode="hip_runtime", op_family="ebm_langevin", + device_proof="device_verified_abi", evidence_target="rocm_gfx1151", + numerical_fixture="tests/unit/test_ebm_native_langevin_gpu.py", + proof_build="llvm23-core+ebm+arena"), ("rocm", "rocm_clifford_native_compiled"): ExecutionRow( target="rocm", compiler_path="rocm_clifford_native_compiled", execution_kind="native_gpu", executable=True, @@ -3846,6 +3890,20 @@ class BackwardComposition: evidence_target="nvidia_sm120", numerical_fixture="tests/device/nvidia/test_conv2d.py", proof_build="cuda13.3+sm120"), + ("nvidia_sm120", "nvidia_ebm_langevin_native_compiled"): ExecutionRow( + target="nvidia_sm120", compiler_path="nvidia_ebm_langevin_native_compiled", + execution_kind="native_gpu", executable=True, + executor_id="nvidia_ebm_langevin_native_compiled", runtime_status="success", + reason="EBM quadratic-energy Langevin loop on sm_120 (2026-09-16): the K-step " + "scf.for over the compiler-derived gradient becomes one cooperative " + "kernel through the row-program emitter (one block per row, one lane " + "per feature, Philox inside the kernel), launched through the native " + "storage package. One launch per loop; bit-exact with the declared " + "numpy policy. No Python-emitted kernel; no fallback.", + execution_mode="cuda_runtime", op_family="ebm_langevin", + device_proof="device_verified_abi", evidence_target="nvidia_sm120", + numerical_fixture="tests/unit/test_ebm_native_langevin_gpu.py", + proof_build="llvm23-core+ebm+arena"), ("nvidia_sm120", "nvidia_clifford_native_compiled"): ExecutionRow( target="nvidia_sm120", compiler_path="nvidia_clifford_native_compiled", execution_kind="native_gpu", executable=True, diff --git a/python/tessera/compiler/native_clifford_gpu.py b/python/tessera/compiler/native_clifford_gpu.py index cc61b569e..9347c919c 100644 --- a/python/tessera/compiler/native_clifford_gpu.py +++ b/python/tessera/compiler/native_clifford_gpu.py @@ -16,16 +16,13 @@ """ from __future__ import annotations -import ctypes as ct import inspect import math import os import shutil import threading from pathlib import Path -from types import SimpleNamespace -import numpy as np from .native_gpu_storage import _run, build_native_gpu_storage, replay_arena_ir from .native_gpu_tensor import IndexSpec, TensorSpec @@ -174,96 +171,14 @@ def bind_clifford_gpu(op: str, shape, *, compiler, llvm_bin, backend, chip, alge # Execution from host arrays. # --------------------------------------------------------------------------- -_CONTEXT_LOCK = threading.Lock() -_CONTEXTS: dict[str, bool] = {} +from .native_host_program import HostArrayProgram, ensure_device_context # noqa: E402,F401 -def ensure_device_context(backend: str) -> None: - """Make device 0 current for this process once (the bound package launches - on the caller's active context and never creates one).""" - with _CONTEXT_LOCK: - if _CONTEXTS.get(backend): - return - cuda = backend == "nvidia" - driver = ct.CDLL("libcuda.so.1" if cuda else "libamdhip64.so") - P = ct.c_void_p - - def call(name, types, *args): - fn = getattr(driver, name) - fn.argtypes, fn.restype = types, ct.c_int - status = fn(*args) - if status: - raise RuntimeError(f"clifford GPU lane: {name} failed with status {status}") - if cuda: - call("cuInit", [ct.c_uint], 0) - context = P() - call("cuDevicePrimaryCtxRetain", [ct.POINTER(P), ct.c_int], ct.byref(context), 0) - call("cuCtxSetCurrent", [P], context) - else: - call("hipInit", [ct.c_uint], 0) - call("hipSetDevice", [ct.c_int], 0) - _CONTEXTS[backend] = True - - -class CliffordDeviceProgram: - """One packaged Clifford op, runnable from host arrays. - - Device buffers are allocated per call through the bound package's own - driver handle and freed before returning; the output comes back as a new - host array. Correctness evidence only -- per-call transfers are not a - performance path. - """ +class CliffordDeviceProgram(HostArrayProgram): + """A packaged Clifford op runnable from host arrays (see HostArrayProgram).""" def __init__(self, binding, op: str): - self.binding, self.op = binding, op - self.package = binding.package - self.bound = binding.package.bind() - binding._bound = self.bound - cuda = self.package.backend == "nvidia" - P, S = ct.c_void_p, ct.c_size_t - - def bind(cu, hip, types): - fn = getattr(self.bound._driver, cu if cuda else hip) - fn.argtypes, fn.restype = types, ct.c_int - return fn - self._alloc = bind("cuMemAlloc_v2", "hipMalloc", [ct.POINTER(P), S]) - self._free = bind("cuMemFree_v2", "hipFree", [P]) - self._to_device = bind("cuMemcpyHtoD_v2", "hipMemcpyHtoD", [P, P, S]) - self._to_host = bind("cuMemcpyDtoH_v2", "hipMemcpyDtoH", [P, P, S]) - self.specs = tuple(s for s in binding.specs if isinstance(s, TensorSpec)) - self._lock = threading.RLock() - - def run(self, *arrays) -> np.ndarray: - inputs = [np.ascontiguousarray(np.asarray(a, dtype=np.float32)) for a in arrays] - if len(inputs) != len(self.specs) - 1: - raise ValueError(f"clifford {self.op} takes {len(self.specs) - 1} operand(s)") - for spec, value in zip(self.specs[:-1], inputs, strict=True): - if value.shape != tuple(spec.shape): - raise ValueError(f"clifford GPU lane admits {spec.name} of shape {tuple(spec.shape)}") - out_spec = self.specs[-1] - out = np.empty(tuple(out_spec.shape), np.float32) - pointers: list[ct.c_void_p] = [] - with self._lock: - try: - views = [] - for value in [*inputs, out]: - pointer = ct.c_void_p() - self.bound._check(self._alloc(ct.byref(pointer), max(value.nbytes, 4))) - pointers.append(pointer) - views.append(SimpleNamespace(__cuda_array_interface__=dict( - version=3, shape=value.shape, typestr=" None: + """Make device 0 current for this process once (the bound package launches + on the caller's active context and never creates one).""" + with _CONTEXT_LOCK: + if _CONTEXTS.get(backend): + return + cuda = backend == "nvidia" + driver = ct.CDLL("libcuda.so.1" if cuda else "libamdhip64.so") + P = ct.c_void_p + + def call(name, types, *args): + fn = getattr(driver, name) + fn.argtypes, fn.restype = types, ct.c_int + status = fn(*args) + if status: + raise RuntimeError(f"native host program: {name} failed with status {status}") + if cuda: + call("cuInit", [ct.c_uint], 0) + context = P() + call("cuDevicePrimaryCtxRetain", [ct.POINTER(P), ct.c_int], ct.byref(context), 0) + call("cuCtxSetCurrent", [P], context) + else: + call("hipInit", [ct.c_uint], 0) + call("hipSetDevice", [ct.c_int], 0) + _CONTEXTS[backend] = True + + +class HostArrayProgram: + """One bound native tensor call, runnable from host arrays; the spec order + is the ABI order (inputs first, then writable outputs).""" + def __init__(self, binding, name: str = "program"): + self.binding, self.name = binding, name + self.package = binding.package + self.bound = binding.package.bind() + binding._bound = self.bound + cuda = self.package.backend == "nvidia" + P, S = ct.c_void_p, ct.c_size_t + + def bind(cu, hip, types): + fn = getattr(self.bound._driver, cu if cuda else hip) + fn.argtypes, fn.restype = types, ct.c_int + return fn + self._alloc = bind("cuMemAlloc_v2", "hipMalloc", [ct.POINTER(P), S]) + self._free = bind("cuMemFree_v2", "hipFree", [P]) + self._to_device = bind("cuMemcpyHtoD_v2", "hipMemcpyHtoD", [P, P, S]) + self._to_host = bind("cuMemcpyDtoH_v2", "hipMemcpyDtoH", [P, P, S]) + self.specs = tuple(s for s in binding.specs if isinstance(s, TensorSpec)) + self.inputs = tuple(s for s in self.specs if not s.writable) + self.outputs = tuple(s for s in self.specs if s.writable) + self._lock = threading.RLock() + + def run(self, *arrays): + """Return the output array (or a tuple of them, in ABI order).""" + if len(arrays) != len(self.inputs): + raise ValueError(f"{self.name} takes {len(self.inputs)} operand(s)") + host: list[np.ndarray] = [] + for spec, value in zip(self.inputs, arrays, strict=True): + kind, _ = _DTYPES[spec.dtype] + array = np.ascontiguousarray(np.asarray(value, dtype=kind)) + if array.shape != tuple(spec.shape): + raise ValueError(f"{self.name} admits {spec.name} of shape {tuple(spec.shape)}") + host.append(array) + outs = [np.empty(tuple(spec.shape), _DTYPES[spec.dtype][0]) for spec in self.outputs] + pointers: list[ct.c_void_p] = [] + with self._lock: + try: + views = [] + for spec, value in zip(self.specs, [*host, *outs], strict=True): + pointer = ct.c_void_p() + self.bound._check(self._alloc(ct.byref(pointer), max(value.nbytes, 4))) + pointers.append(pointer) + views.append(SimpleNamespace(__cuda_array_interface__=dict( + version=3, shape=value.shape, typestr=_DTYPES[spec.dtype][1], data=(pointer.value, False)))) + for value, pointer in zip(host, pointers[:len(host)], strict=True): + self.bound._check(self._to_device(pointer, value.ctypes.data, value.nbytes)) + self.binding(*views, 1) + self.bound._check(self.bound._sync()) + for value, pointer in zip(outs, pointers[len(host):], strict=True): + self.bound._check(self._to_host(value.ctypes.data, pointer, value.nbytes)) + finally: + for pointer in reversed(pointers): + if pointer.value: + self._free(pointer) + return outs[0] if len(outs) == 1 else tuple(outs) + + def close(self): + self.bound.close() diff --git a/python/tessera/compiler/native_row_program.py b/python/tessera/compiler/native_row_program.py new file mode 100644 index 000000000..d90abc486 --- /dev/null +++ b/python/tessera/compiler/native_row_program.py @@ -0,0 +1,118 @@ +"""A ``[rows, features]`` tensor row program as one cooperative device kernel. + +The compiler's ``tessera-row-program-to-gpu`` pass turns a function over +``[rows, features]`` tensors (parallel ``linalg.generic`` bodies, feature-axis +``linalg.reduce``, small uniform integer vectors, ``scf.for``) into one +``gpu.func``: one block per row, one lane per feature, loop-carried state in +registers, ordered shared-memory reductions. This module runs that pass through +one ``tessera-opt`` invocation, attaches the tensor contract, packages the +kernel with ``build_native_gpu_storage`` and returns a host-array runner. It is +the seam every row-shaped domain loop (EBM Langevin first) packages through; +no Python-emitted arithmetic. +""" +from __future__ import annotations + +import inspect +import re +import threading +from pathlib import Path +from typing import Sequence + +from .native_gpu_storage import _run, build_native_gpu_storage, replay_arena_ir +from .native_gpu_tensor import IndexSpec, TensorSpec +from .native_host_program import HostArrayProgram, ensure_device_context +from .native_storage_contract import attach_tensor_contract, generate_tensor_binding + +ROW_PROGRAM_PASS = "--tessera-row-program-to-gpu" +MAX_FEATURES = 1024 + + +def row_program_kernel(module_text: str, *, entry: str, backend: str, compiler, passes: Sequence[str] = ()) -> tuple[str, int]: + """Run ``passes`` then the row-program emitter; return (kernel text, lanes).""" + if backend not in ("nvidia", "rocm"): + raise ValueError("row program device route targets nvidia or rocm") + kernel = _run(Path(compiler), *passes, f"{ROW_PROGRAM_PASS}=backend={backend} entry={entry}", source=module_text) + block = re.search(r"known_block_size = array", kernel) + if block is None or kernel.count("gpu.func ") != 1 or "gpu.func @row_program(" not in kernel: + raise ValueError("row program device route: the compiler did not produce the row-program kernel") + return kernel, int(block[1]) + + +def row_program_device_source(module_text: str, *, entry: str, specs: Sequence[TensorSpec], rows: int, + backend: str, compiler, passes: Sequence[str] = ()) -> tuple[str, tuple]: + """The packaged kernel source with its tensor contract (grid = rows, block = lanes).""" + kernel, lanes = row_program_kernel(module_text, entry=entry, backend=backend, compiler=compiler, passes=passes) + full = tuple(specs) + (IndexSpec("scratch", 1, 1),) + return attach_tensor_contract(kernel, full, grid=(int(rows), 1, 1), block=(lanes, 1, 1)), full + + +def bind_row_program(source: str, specs, *, compiler, llvm_bin, backend: str, chip: str): + """Package for (backend, chip); replay must agree; return the native tensor call.""" + compiler = Path(compiler) + package = build_native_gpu_storage(source, compiler=compiler, llvm_bin=Path(llvm_bin), backend=backend, chip=chip) + if package.arena_ir != replay_arena_ir(compiler, source): + raise ValueError("row program device route: native replay disagrees") + return generate_tensor_binding(package, inspect.Signature([ + inspect.Parameter(s.name, inspect.Parameter.POSITIONAL_ONLY) for s in specs])) + + +_PROGRAMS: dict = {} +_LOCK = threading.Lock() + + +def row_program_device(module_text: str, *, entry: str, specs: Sequence[TensorSpec], rows: int, backend: str, + chip: str, compiler, llvm_bin, passes: Sequence[str] = (), name: str = "row program") -> HostArrayProgram: + """Compile once per process and return the host-array runner for the program.""" + key = (module_text, entry, tuple(passes), backend, chip, str(compiler), str(llvm_bin)) + with _LOCK: + program = _PROGRAMS.get(key) + if program is None: + ensure_device_context(backend) + source, full = row_program_device_source(module_text, entry=entry, specs=specs, rows=rows, + backend=backend, compiler=compiler, passes=passes) + program = _PROGRAMS[key] = HostArrayProgram( + bind_row_program(source, full, compiler=compiler, llvm_bin=llvm_bin, backend=backend, chip=chip), name) + return program + + +def row_normalize_module(rows: int, features: int) -> str: + """``y = x / sqrt(sum_f x^2)`` per row: the smallest program that exercises the + ordered shared-memory reduction (a feature-axis ``linalg.reduce`` whose + result is broadcast back over the lanes).""" + r, f = int(rows), int(features) + return f"""#map = affine_map<(d0, d1) -> (d0, d1)> +#row = affine_map<(d0, d1) -> (d0, 0)> +module {{ + func.func @row_normalize(%x: tensor<{r}x{f}xf32>) -> tensor<{r}x{f}xf32> {{ + %zero = arith.constant 0.000000e+00 : f32 + %e = tensor.empty() : tensor<{r}x{f}xf32> + %sq = linalg.generic {{indexing_maps = [#map, #map, #map], iterator_types = ["parallel", "parallel"]}} + ins(%x, %x : tensor<{r}x{f}xf32>, tensor<{r}x{f}xf32>) outs(%e : tensor<{r}x{f}xf32>) {{ + ^bb0(%a: f32, %b: f32, %o: f32): + %m = arith.mulf %a, %b : f32 + linalg.yield %m : f32 + }} -> tensor<{r}x{f}xf32> + %re = tensor.empty() : tensor<{r}xf32> + %init = linalg.fill ins(%zero : f32) outs(%re : tensor<{r}xf32>) -> tensor<{r}xf32> + %s = linalg.reduce ins(%sq : tensor<{r}x{f}xf32>) outs(%init : tensor<{r}xf32>) dimensions = [1] + (%in: f32, %acc: f32) {{ + %add = arith.addf %in, %acc : f32 + linalg.yield %add : f32 + }} + %n = linalg.generic {{indexing_maps = [affine_map<(d0) -> (d0)>, affine_map<(d0) -> (d0)>], iterator_types = ["parallel"]}} + ins(%s : tensor<{r}xf32>) outs(%re : tensor<{r}xf32>) {{ + ^bb0(%v: f32, %o: f32): + %sq2 = math.sqrt %v : f32 + linalg.yield %sq2 : f32 + }} -> tensor<{r}xf32> + %nx = tensor.expand_shape %n [[0, 1]] output_shape [{r}, 1] : tensor<{r}xf32> into tensor<{r}x1xf32> + %y = linalg.generic {{indexing_maps = [#map, #row, #map], iterator_types = ["parallel", "parallel"]}} + ins(%x, %nx : tensor<{r}x{f}xf32>, tensor<{r}x1xf32>) outs(%e : tensor<{r}x{f}xf32>) {{ + ^bb0(%a: f32, %d: f32, %o: f32): + %q = arith.divf %a, %d : f32 + linalg.yield %q : f32 + }} -> tensor<{r}x{f}xf32> + return %y : tensor<{r}x{f}xf32> + }} +}} +""" diff --git a/python/tessera/compiler/native_storage_contract.py b/python/tessera/compiler/native_storage_contract.py index fcc71b3d3..f6c92c858 100644 --- a/python/tessera/compiler/native_storage_contract.py +++ b/python/tessera/compiler/native_storage_contract.py @@ -23,8 +23,14 @@ def attach_tensor_contract(source: str, specs, *, grid, block) -> str: encoded = json.dumps(data, sort_keys=True, separators=(',', ':'), allow_nan=False) encoded = encoded.replace('\\', '\\5C').replace('"', '\\22') result, count = re.subn(r'(?m)^module \{', lambda _: f'module attributes {{{ATTRIBUTE} = "{encoded}"}} {{', source) + if count == 0: + # A compiler-produced module may already carry provenance attributes + # (e.g. tessera.row_program.*); the contract joins that dictionary so + # the provenance survives into the package. + result, count = re.subn(r'(?m)^module attributes \{', lambda _: f'module attributes {{{ATTRIBUTE} = "{encoded}", ', + source) if count != 1: - raise ValueError('tensor manifest producer requires one plain top-level module') + raise ValueError('tensor manifest producer requires one top-level module') return result diff --git a/python/tessera/ebm/native_langevin.py b/python/tessera/ebm/native_langevin.py new file mode 100644 index 000000000..2cd2b2df5 --- /dev/null +++ b/python/tessera/ebm/native_langevin.py @@ -0,0 +1,254 @@ +"""The traceable quadratic energy loop through the MLIR/LLVM backbone +(W4-PRODUCT-1 / AD-SOLVER-IFT-1 acceptance, 2026-09-16). + +The energy is a Graph IR function, ``E(y, x) = 0.5 * sum((x - y)^2, axis=-1)``, +marked ``tessera.autodiff = "reverse"``. Inside ``libtessera_jit`` the paired +autodiff pass derives ``@E__bwd`` (the compiler's gradient, not a hand +formula), the EBM lowering turns each ``tessera_ebm.langevin_step`` into +``y - eta * grad + sqrt(2 eta T) * z`` with ``z`` drawn on the device side of +the lane from Philox-4x32-10 / Box-Muller inside a ``linalg.generic``, and the +whole ``scf.for`` of K steps compiles to one function: no per-step host +gradient or noise transfers. ``reference_langevin_loop`` is the declared RNG +policy mirrored bit-for-bit in numpy (the standalone ``philox_4x32_10``). + +Envelope: f32, static ``[rows, features]`` shapes, euclidean manifold, key as +two int64 words. No fallback: a library without the EBM lane raises. +""" +from __future__ import annotations + +import math +import threading +from typing import Sequence + +import numpy as np + +from tessera import _jit_boundary as jb + +_BYTE_BOUND = 64 * 1024 * 1024 + + +def has_native_langevin() -> bool: + """Whether libtessera_jit was built with the EBM lane (TESSERA_BUILD_EBM_BACKEND).""" + if not jb.is_available(): + return False + lib = jb._load() + sym = getattr(lib, "tessera_jit_has_ebm", None) + if sym is None: + return False + import ctypes + sym.restype, sym.argtypes = ctypes.c_int, [] + return bool(sym()) + + +def _check(shape, eta, temperature, steps): + shape = tuple(int(d) for d in shape) + if len(shape) != 2 or any(d <= 0 for d in shape) or math.prod(shape) * 4 > _BYTE_BOUND: + raise ValueError("native langevin admits a positive static [rows, features] f32 state within the buffer bound") + if not (isinstance(eta, (int, float)) and math.isfinite(eta) and eta > 0): + raise ValueError("native langevin requires eta > 0") + if not (isinstance(temperature, (int, float)) and math.isfinite(temperature) and temperature >= 0): + raise ValueError("native langevin requires temperature >= 0") + if type(steps) is not int or steps < 1 or steps > 1 << 20: + raise ValueError("native langevin requires 1 <= steps <= 2**20") + return shape + + +def langevin_loop_module(shape, *, eta: float, temperature: float, steps: int) -> str: + """The Graph-level program: quadratic energy + a K-step Langevin loop.""" + rows, features = _check(shape, eta, temperature, steps) + st = f"tensor<{rows}x{features}xf32>" + en = f"tensor<{rows}xf32>" + return f"""module {{ + func.func @quadratic_energy(%y: {st}, %x: {st}) -> {en} + attributes {{tessera.autodiff = "reverse"}} {{ + %d = "tessera.sub"(%x, %y) : ({st}, {st}) -> {st} + %sq = "tessera.mul"(%d, %d) : ({st}, {st}) -> {st} + %s = "tessera.reduce"(%sq) {{axis = 1 : i64, kind = "sum"}} : ({st}) -> {en} + %half = arith.constant dense<5.000000e-01> : {en} + %e = "tessera.mul"(%s, %half) : ({en}, {en}) -> {en} + return %e : {en} + }} + func.func @tessera_jit_ebm_energy(%y: {st}, %x: {st}) -> {en} {{ + %e = "tessera_ebm.energy"(%x, %y) {{ energy_fn = @quadratic_energy }} : ({st}, {st}) -> {en} + return %e : {en} + }} + func.func @tessera_jit_ebm_langevin_loop(%y0: {st}, %x: {st}, %key0: tensor<2xi64>) -> ({st}, tensor<2xi64>) {{ + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %steps = arith.constant {steps} : index + %r:2 = scf.for %t = %c0 to %steps step %c1 iter_args(%y = %y0, %key = %key0) -> ({st}, tensor<2xi64>) {{ + %n:2 = "tessera_ebm.langevin_step"(%y, %key, %x) + {{ energy_fn = @quadratic_energy, eta = {float(eta)!r} : f64, temperature = {float(temperature)!r} : f64, manifold = "euclidean" }} + : ({st}, tensor<2xi64>, {st}) -> ({st}, tensor<2xi64>) + scf.yield %n#0, %n#1 : {st}, tensor<2xi64> + }} + return %r#0, %r#1 : {st}, tensor<2xi64> + }} +}} +""" + + +def _require(): + if not has_native_langevin(): + raise jb.TesseraJitError("libtessera_jit was built without the EBM lane " + "(configure with -DTESSERA_BUILD_EBM_BACKEND=ON)") + + +def native_quadratic_energy(y, x) -> np.ndarray: + """Per-row ``0.5 * ||x - y||^2`` through the lane (the energy op lowers to a call).""" + _require() + y = np.ascontiguousarray(y, dtype=np.float32); x = np.ascontiguousarray(x, dtype=np.float32) + if y.shape != x.shape: + raise jb.TesseraJitError("energy requires equal state and context shapes") + shape = _check(y.shape, 1.0, 0.0, 1) + handle = jb.compile_module(langevin_loop_module(shape, eta=1.0, temperature=0.0, steps=1)) + try: + out = np.empty((shape[0],), np.float32) + jb.invoke(handle, "tessera_jit_ebm_energy", [y, x], out) + return out + finally: + jb.destroy(handle) + + +def native_langevin_loop(y0, x, key: Sequence[int], *, eta: float, temperature: float, steps: int): + """K Langevin steps on the quadratic energy as one compiled function. + + Returns ``(y_K, next_key)``; ``key`` is two int64 words (S4 RNGKey words). + """ + _require() + y0 = np.ascontiguousarray(y0, dtype=np.float32); x = np.ascontiguousarray(x, dtype=np.float32) + if y0.shape != x.shape: + raise jb.TesseraJitError("langevin requires equal state and context shapes") + shape = _check(y0.shape, eta, temperature, steps) + key_arr = np.ascontiguousarray(np.asarray(key, dtype=np.int64).reshape(2)) + handle = jb.compile_module(langevin_loop_module(shape, eta=eta, temperature=temperature, steps=steps)) + try: + out = np.empty(shape, np.float32) + next_key = np.empty((2,), np.int64) + jb.invoke(handle, "tessera_jit_ebm_langevin_loop", [y0, x, key_arr], [out, next_key]) + return out, next_key + finally: + jb.destroy(handle) + + +def reference_langevin_loop(y0, x, key: Sequence[int], *, eta: float, temperature: float, steps: int): + """The declared policy in numpy, bit-for-bit: Philox key = (lo32, hi32) of + key[0]; counter = (flat index, lo32, hi32 of key[1], 0); z = sqrt(-2 ln u0) + cos(2 pi u1) with u = (word + 0.5) * 2^-32 on words 0 and 1 (computed in + f64, rounded to f32 once); next key = (key[0], key[1] + 1); the gradient of + the quadratic energy is y - x.""" + from tessera.compiler.philox import philox_4x32_10 + y = np.ascontiguousarray(y0, dtype=np.float32).copy(); x = np.ascontiguousarray(x, dtype=np.float32) + shape = _check(y.shape, eta, temperature, steps) + k = [int(v) for v in np.asarray(key, dtype=np.int64).reshape(2)] + scale = math.sqrt(2.0 * float(eta) * float(temperature)) + philox_key = np.array([k[0] & 0xFFFFFFFF, (k[0] >> 32) & 0xFFFFFFFF], np.uint32) + for _ in range(steps): + grad = y - x + y = (y - np.float32(eta) * grad).astype(np.float32) + if scale > 0.0: + s0, s1 = np.uint32(k[1] & 0xFFFFFFFF), np.uint32((k[1] >> 32) & 0xFFFFFFFF) + z = np.empty(math.prod(shape), np.float32) + for i in range(z.size): + words = philox_4x32_10(np.array([np.uint32(i), s0, s1, np.uint32(0)], np.uint32), philox_key) + u0 = (float(words[0]) + 0.5) * 2.0 ** -32 + u1 = (float(words[1]) + 0.5) * 2.0 ** -32 + z[i] = np.float32(math.sqrt(-2.0 * math.log(u0)) * math.cos(2.0 * math.pi * u1)) + y = (y + np.float32(scale) * z.reshape(shape)).astype(np.float32) + k[1] += 1 + return y, np.array(k, np.int64) + + +def package_ebm_langevin_cpu(shape, *, eta: float, temperature: float, steps: int): + """A runtime artifact for the loop (row ``cpu`` / ``cpu_ebm_langevin_llvm_jit``).""" + shape = _check(shape, eta, temperature, steps) + _require() + from tessera.runtime import RuntimeArtifact + return RuntimeArtifact(metadata={ + "target": "cpu", "compiler_path": "cpu_ebm_langevin_llvm_jit", "executable": True, + "kernel_id": f"ebm_langevin_quadratic_{shape[0]}x{shape[1]}_k{steps}", + "op": "ebm_langevin_loop", "shape": shape, "eta": float(eta), "temperature": float(temperature), + "steps": int(steps), "dtype": "f32", + }) + + +# --------------------------------------------------------------------------- +# Device route (EBM_NATIVE_LOOP_ARCHITECTURE.md, realized as the row-program +# emitter): the same lowered loop becomes one cooperative kernel -- one block +# per row, one lane per feature, the K steps and the Philox draw inside the +# kernel, state in registers -- packaged by build_native_gpu_storage. +# --------------------------------------------------------------------------- + +DEVICE_ENTRY = "tessera_jit_ebm_langevin_loop" +_DEVICE_PIPELINE = ("--tessera-autodiff-paired", "--tessera-ebm-canonicalize", "--tessera-ebm-lower-langevin", + "--tessera-to-linalg", "--inline", "--convert-elementwise-to-linalg", "--canonicalize", "--cse") + + +def langevin_device_source(shape, *, eta: float, temperature: float, steps: int, backend: str, compiler): + """Run the whole chain in one tessera-opt invocation and attach the tensor contract.""" + from tessera.compiler.native_gpu_tensor import TensorSpec + from tessera.compiler.native_row_program import MAX_FEATURES, row_program_device_source + rows, feats = _check(shape, eta, temperature, steps) + if feats > MAX_FEATURES: + raise ValueError(f"langevin device route admits at most {MAX_FEATURES} features per row (one lane each)") + specs = (TensorSpec("y0", "fp32", (rows, feats), False), TensorSpec("x", "fp32", (rows, feats), False), + TensorSpec("key", "int64", (2,), False), TensorSpec("y", "fp32", (rows, feats), True), + TensorSpec("next_key", "int64", (2,), True)) + source, full = row_program_device_source( + langevin_loop_module((rows, feats), eta=eta, temperature=temperature, steps=steps), entry=DEVICE_ENTRY, + specs=specs, rows=rows, backend=backend, compiler=compiler, passes=_DEVICE_PIPELINE) + if "tessera_ebm." in source.split("gpu.func @row_program(", 1)[1]: + raise ValueError("langevin device route: an EBM op survived the lowering") + return source, full + + +def bind_ebm_langevin_gpu(shape, *, eta, temperature, steps, compiler, llvm_bin, backend, chip): + """Package the loop for (backend, chip) and return its native tensor call.""" + from tessera.compiler.native_row_program import bind_row_program + source, specs = langevin_device_source(shape, eta=eta, temperature=temperature, steps=steps, + backend=backend, compiler=compiler) + return bind_row_program(source, specs, compiler=compiler, llvm_bin=llvm_bin, backend=backend, chip=chip) + + +_PROGRAMS: dict = {} +_PROGRAM_LOCK = threading.Lock() + + +def ebm_langevin_program(shape, *, eta, temperature, steps, backend, chip, compiler, llvm_bin): + """Compile (once per process) and return the device program for the loop.""" + from tessera.compiler.native_host_program import HostArrayProgram, ensure_device_context + key = (tuple(shape), float(eta), float(temperature), int(steps), backend, chip, str(compiler), str(llvm_bin)) + with _PROGRAM_LOCK: + program = _PROGRAMS.get(key) + if program is None: + ensure_device_context(backend) + binding = bind_ebm_langevin_gpu(shape, eta=eta, temperature=temperature, steps=steps, + compiler=compiler, llvm_bin=llvm_bin, backend=backend, chip=chip) + program = _PROGRAMS[key] = HostArrayProgram(binding, "ebm langevin loop") + return program + + +def native_langevin_loop_device(y0, x, key, *, eta, temperature, steps, backend, chip, compiler, llvm_bin): + """K Langevin steps as one device launch; returns (y_K, next_key).""" + program = ebm_langevin_program(np.asarray(y0).shape, eta=eta, temperature=temperature, steps=steps, + backend=backend, chip=chip, compiler=compiler, llvm_bin=llvm_bin) + out, next_key = program.run(y0, x, np.asarray(key, dtype=np.int64).reshape(2)) + return out, next_key + + +def package_ebm_langevin_native(shape, *, eta: float, temperature: float, steps: int, target: str): + """A runtime artifact for the device route (rows ``rocm`` / + ``rocm_ebm_langevin_native_compiled``, ``nvidia_sm120`` / + ``nvidia_ebm_langevin_native_compiled``); compiled at launch on the owning host.""" + if target not in ("rocm", "nvidia_sm120"): + raise ValueError("langevin native device route targets rocm or nvidia_sm120") + shape = _check(shape, eta, temperature, steps) + from tessera.runtime import RuntimeArtifact + path = "rocm_ebm_langevin_native_compiled" if target == "rocm" else "nvidia_ebm_langevin_native_compiled" + return RuntimeArtifact(metadata={ + "target": target, "compiler_path": path, "executable": True, + "kernel_id": f"ebm_langevin_native_{shape[0]}x{shape[1]}_k{steps}", + "op": "ebm_langevin_loop", "shape": shape, "eta": float(eta), "temperature": float(temperature), + "steps": int(steps), "dtype": "f32", + }) + diff --git a/python/tessera/runtime.py b/python/tessera/runtime.py index ddd6e3b68..b4cdc30ee 100644 --- a/python/tessera/runtime.py +++ b/python/tessera/runtime.py @@ -32839,6 +32839,9 @@ def _executor_table(): "native_cpu": _execute_cpu_native_or_jit, "cpu_autodiff_paired_llvm_jit": _execute_cpu_autodiff_paired, "cpu_clifford_llvm_jit": _execute_cpu_clifford_llvm_jit, + "cpu_ebm_langevin_llvm_jit": _execute_cpu_ebm_langevin_llvm_jit, + "rocm_ebm_langevin_native_compiled": _execute_ebm_langevin_native_gpu, + "nvidia_ebm_langevin_native_compiled": _execute_ebm_langevin_native_gpu, "rocm_clifford_native_compiled": _execute_clifford_native_gpu, "nvidia_clifford_native_compiled": _execute_clifford_native_gpu, "jit_cpu_numpy": _execute_jit_cpu_artifact, @@ -33672,6 +33675,56 @@ def _execute_cpu_autodiff_paired(artifact: RuntimeArtifact, args: Any) -> Any: return tuple(outputs) +def _execute_ebm_langevin_native_gpu(artifact: RuntimeArtifact, args: Any) -> Any: + """Launch the K-step Langevin loop as one cooperative device kernel (the + row-program emitter over the compiler-derived gradient). No fallback.""" + import numpy as np + from .ebm.native_langevin import native_langevin_loop_device + from .compiler.scheduled_matmul import find_tessera_opt + from .compiler.llvm_tools import llvm_bin_dir + + metadata = artifact.metadata or {} + shape = tuple(int(d) for d in metadata.get("shape", ())) + target = str(metadata.get("target", "")) + if len(shape) != 2 or len(args) != 3: + raise ValueError("ebm langevin native artifact requires a [rows, features] shape and (y0, x, key) operands") + if target == "rocm": + backend, chip = "rocm", _rocm_chip() + elif target == "nvidia_sm120": + backend, chip = "nvidia", "sm_120" + else: + raise ValueError("ebm langevin native route targets rocm or nvidia_sm120") + compiler, llvm_bin = find_tessera_opt(), llvm_bin_dir() + if compiler is None or llvm_bin is None: + raise RuntimeError("ebm langevin native route requires tessera-opt and the matched LLVM tools") + y0, x = (np.ascontiguousarray(np.asarray(v, dtype=np.float32)) for v in args[:2]) + if y0.shape != shape or x.shape != shape: + raise ValueError(f"ebm langevin native artifact admits operands of shape {shape}") + return native_langevin_loop_device(y0, x, args[2], eta=float(metadata["eta"]), temperature=float(metadata["temperature"]), + steps=int(metadata["steps"]), backend=backend, chip=chip, compiler=compiler, llvm_bin=llvm_bin) + + +def _execute_cpu_ebm_langevin_llvm_jit(artifact: RuntimeArtifact, args: Any) -> Any: + """Launch the compiled K-step quadratic-energy Langevin loop (W4-PRODUCT-1). + + Metadata owns the shape, eta, temperature and step count; the lane derives + the gradient and draws the noise. Operands: (y0, x, key). No fallback. + """ + import numpy as np + from .ebm.native_langevin import native_langevin_loop + + metadata = artifact.metadata or {} + shape = tuple(int(d) for d in metadata.get("shape", ())) + if len(shape) != 2 or len(args) != 3: + raise ValueError("ebm langevin artifact requires a [rows, features] shape and (y0, x, key) operands") + y0, x = (np.ascontiguousarray(np.asarray(v, dtype=np.float32)) for v in args[:2]) + if y0.shape != shape or x.shape != shape: + raise ValueError(f"ebm langevin artifact admits operands of shape {shape}") + out, next_key = native_langevin_loop(y0, x, args[2], eta=float(metadata["eta"]), + temperature=float(metadata["temperature"]), steps=int(metadata["steps"])) + return out, next_key + + def _execute_cpu_clifford_llvm_jit(artifact: RuntimeArtifact, args: Any) -> Any: """Launch a geometric product through the Clifford lane of libtessera_jit. diff --git a/src/compiler/ir/AdjointInterface.cpp b/src/compiler/ir/AdjointInterface.cpp index 16c2283b0..411ca6ed0 100644 --- a/src/compiler/ir/AdjointInterface.cpp +++ b/src/compiler/ir/AdjointInterface.cpp @@ -322,6 +322,27 @@ llvm::SmallVector AddOp::buildAdjoint( return {outputCotangents[0], outputCotangents[0]}; } +// d(lhs - rhs): the cotangent passes through to lhs and is negated for rhs. +// Until 2026-09-16 `sub` carried only a tangent rule, so any energy or loss +// written with a subtraction (the quadratic energy 0.5*||x - y||^2 first) +// stopped reverse-mode with AUTODIFF_OP_NOT_DIFFERENTIABLE. The negation is +// `0 - dy` on a static zero of the operand type; dynamic operand shapes are +// refused (empty cotangent) rather than guessed. +llvm::SmallVector SubOp::buildAdjoint( + mlir::OpBuilder &builder, mlir::ValueRange outputCotangents) { + if (outputCotangents.size() != 1 || !outputCotangents[0]) + return {mlir::Value(), mlir::Value()}; + mlir::Value dy = outputCotangents[0]; + auto rhsTy = mlir::dyn_cast(getRhs().getType()); + if (!rhsTy || !rhsTy.hasStaticShape()) + return {dy, mlir::Value()}; + auto zeroAttr = mlir::DenseElementsAttr::get( + rhsTy, builder.getZeroAttr(rhsTy.getElementType())); + mlir::Value zero = builder.create(getLoc(), zeroAttr); + auto dRhs = builder.create(getLoc(), rhsTy, zero, dy); + return {dy, dRhs.getResult()}; +} + llvm::SmallVector MulOp::buildAdjoint( mlir::OpBuilder &builder, mlir::ValueRange outputCotangents) { if (outputCotangents.size() != 1 || !outputCotangents[0]) diff --git a/src/compiler/ir/TesseraOps.td b/src/compiler/ir/TesseraOps.td index 40e144151..25cb3f7cc 100644 --- a/src/compiler/ir/TesseraOps.td +++ b/src/compiler/ir/TesseraOps.td @@ -1626,7 +1626,8 @@ def Tessera_AddOp : Op< // shared name→arith-kind table. sub/mul are the minimum needed before // softmax/normalization arrive (softmax = max → sub → exp → div). def Tessera_SubOp : Op]> { + [Pure, DeclareOpInterfaceMethods, + DeclareOpInterfaceMethods]> { let summary = "Elementwise tensor subtraction (a - b)"; let arguments = (ins TensorType:$lhs, TensorType:$rhs); let results = (outs TensorType:$result); diff --git a/src/solvers/ebm/CMakeLists.txt b/src/solvers/ebm/CMakeLists.txt index 9ad06bf55..6c57521a3 100644 --- a/src/solvers/ebm/CMakeLists.txt +++ b/src/solvers/ebm/CMakeLists.txt @@ -48,6 +48,7 @@ add_mlir_library(TesseraEBM lib/Passes/FuseEnergyGrad.cpp lib/Passes/CheckpointInnerLoop.cpp lib/Passes/PipelineCandidates.cpp + lib/Passes/LowerLangevin.cpp DEPENDS TesseraEBMTableGen @@ -55,8 +56,12 @@ add_mlir_library(TesseraEBM LINK_LIBS PUBLIC MLIRIR + MLIRArithDialect MLIRFuncDialect + MLIRLinalgDialect + MLIRMathDialect MLIRSCFDialect + MLIRTensorDialect MLIRPass MLIRTransforms ) @@ -64,6 +69,10 @@ add_mlir_library(TesseraEBM add_executable(ts-ebm-opt tools/ts-ebm-opt.cpp) target_link_libraries(ts-ebm-opt PRIVATE TesseraEBM + MLIRArithDialect + MLIRLinalgDialect + MLIRMathDialect + MLIRTensorDialect MLIRIR MLIRPass MLIRSupport diff --git a/src/solvers/ebm/include/tessera/EBM/EBMPasses.h b/src/solvers/ebm/include/tessera/EBM/EBMPasses.h index b3f96e48b..648336a04 100644 --- a/src/solvers/ebm/include/tessera/EBM/EBMPasses.h +++ b/src/solvers/ebm/include/tessera/EBM/EBMPasses.h @@ -12,5 +12,8 @@ std::unique_ptr createEBMCanonicalizePass(); std::unique_ptr createEBMFuseEnergyGradPass(); std::unique_ptr createEBMCheckpointInnerLoopPass(); std::unique_ptr createEBMPipelineCandidatesPass(); +// 2026-09-16 — the first EBM lowering: energy / inner_step / langevin_step +// (euclidean) to arith/linalg over the compiler-derived gradient. +std::unique_ptr createEBMLowerLangevinPass(); } // namespace tessera diff --git a/src/solvers/ebm/include/tessera/EBM/EBMPasses.td b/src/solvers/ebm/include/tessera/EBM/EBMPasses.td index a590dfda0..8b2d5c733 100644 --- a/src/solvers/ebm/include/tessera/EBM/EBMPasses.td +++ b/src/solvers/ebm/include/tessera/EBM/EBMPasses.td @@ -40,3 +40,10 @@ def EBMPipelineCandidates : Pass<"tessera-ebm-pipeline-candidates", "mlir::Modul "[EBM6 stub] Map the K-candidate dimension across streams / devices " "for self_verify parallelism."; } + +def EBMLowerLangevin : Pass<"tessera-ebm-lower-langevin", "mlir::ModuleOp"> { + let summary = + "Lower energy / inner_step / langevin_step (euclidean) to arith/linalg over " + "the compiler-derived gradient @E__bwd with on-device Philox noise."; +} + diff --git a/src/solvers/ebm/lib/Dialect/EBM/EBMOps.td b/src/solvers/ebm/lib/Dialect/EBM/EBMOps.td index acf4f9a09..e9b2ebdfa 100644 --- a/src/solvers/ebm/lib/Dialect/EBM/EBMOps.td +++ b/src/solvers/ebm/lib/Dialect/EBM/EBMOps.td @@ -140,6 +140,10 @@ def EBM_LangevinStepOp : EBMOp<"langevin_step", []> { let arguments = (ins AnyType:$state, AnyType:$rng_key, + // Captured context the energy closes over (2026-09-16): `energy_fn` is + // E(state, captures...) -> energies, so a quadratic energy's target `x` + // rides here instead of being baked into the module. + Variadic:$captures, FlatSymbolRefAttr:$energy_fn, F64Attr:$eta, F64Attr:$temperature, @@ -148,7 +152,7 @@ def EBM_LangevinStepOp : EBMOp<"langevin_step", []> { ); let results = (outs AnyType:$new_state, AnyType:$next_key); let assemblyFormat = [{ - $state `,` $rng_key `energy_fn` $energy_fn `eta` $eta + $state `,` $rng_key (`captures` $captures^)? `energy_fn` $energy_fn `eta` $eta `temperature` $temperature `manifold` $manifold (`grade` $grade^)? attr-dict `:` functional-type(operands, results) }]; diff --git a/src/solvers/ebm/lib/Passes/LowerLangevin.cpp b/src/solvers/ebm/lib/Passes/LowerLangevin.cpp new file mode 100644 index 000000000..52d621440 --- /dev/null +++ b/src/solvers/ebm/lib/Passes/LowerLangevin.cpp @@ -0,0 +1,280 @@ +//===- LowerLangevin.cpp -----------------------------------------*- C++ -*-===// +// +// EBMLowerLangevinPass (2026-09-16, W4-PRODUCT-1 / AD-SOLVER-IFT-1): the first +// *lowering* in the EBM dialect. Every pass before it annotated; this one +// turns the energy program into arithmetic the shared backbone executes: +// +// %e = tessera_ebm.energy(%x, %y) {energy_fn = @E} +// -> func.call @E(%y, %x) (E is E(state, captures...)) +// %y1 = tessera_ebm.inner_step(%y, %g) {eta} +// -> %y - eta * %g +// %y1, %k1 = tessera_ebm.langevin_step(%y, %key, captures...) +// {energy_fn = @E, eta, temperature, manifold = "euclidean"} +// -> %g = call @E__bwd(%y, captures..., ones)#0 (the compiler's own +// reverse-mode gradient of E w.r.t. the state; the paired +// autodiff pass must have produced @E__bwd -- this pass never +// differentiates by hand and refuses when the symbol is absent) +// %z = Philox-4x32-10 / Box-Muller standard normals, one per +// element, generated inside a linalg.generic on the device +// side of whatever backend consumes the loop (no host RNG) +// %y1 = %y - eta * %g + sqrt(2 * eta * T) * %z +// %k1 = %key + [0, 1] +// +// Declared RNG policy (mirrored bit-for-bit by +// python/tessera/ebm/native_langevin.py::reference_langevin_loop): +// philox key = (lo32(key[0]), hi32(key[0])) +// philox counter = (flat element index, lo32(key[1]), hi32(key[1]), 0) +// z = sqrt(-2 ln u0) * cos(2 pi u1), u = (word + 0.5) * 2^-32 on words 0, 1 +// next key = (key[0], key[1] + 1) -- one stream per step +// Temperature 0 emits no noise at all (pure gradient descent), so a T = 0 loop +// is exactly y - eta * g per step. +// +// Envelope: static ranked f32 state, key `tensor<2xi64>`, manifold +// "euclidean" only ("sphere" / "bivector" fail closed with a diagnostic). +// +//===----------------------------------------------------------------------===// +#include "tessera/EBM/EBMPasses.h" + +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/Dialect/Func/IR/FuncOps.h" +#include "mlir/Dialect/Linalg/IR/Linalg.h" +#include "mlir/Dialect/Math/IR/Math.h" +#include "mlir/Dialect/Tensor/IR/Tensor.h" +#include "mlir/IR/Builders.h" +#include "mlir/IR/BuiltinAttributes.h" +#include "mlir/IR/BuiltinOps.h" +#include "mlir/IR/PatternMatch.h" +#include "mlir/Pass/Pass.h" +#include "mlir/Transforms/GreedyPatternRewriteDriver.h" + +#include + +using namespace mlir; + +namespace tessera { +namespace { + +constexpr StringRef kEnergyOp = "tessera_ebm.energy"; +constexpr StringRef kInnerStepOp = "tessera_ebm.inner_step"; +constexpr StringRef kLangevinOp = "tessera_ebm.langevin_step"; + +static RankedTensorType staticF32(Value v) { + auto ty = dyn_cast(v.getType()); + if (!ty || !ty.hasStaticShape() || !ty.getElementType().isF32()) return nullptr; + return ty; +} + +static Value splat(OpBuilder &b, Location loc, RankedTensorType ty, double value) { + auto attr = DenseElementsAttr::get(ty, b.getFloatAttr(ty.getElementType(), value)); + return b.create(loc, attr); +} + +// energy(x, y) {energy_fn = @E} -> call @E(y, x) +struct LowerEnergy : public RewritePattern { + LowerEnergy(MLIRContext *ctx) : RewritePattern(kEnergyOp, 1, ctx) {} + LogicalResult matchAndRewrite(Operation *op, PatternRewriter &rewriter) const override { + auto fn = op->getAttrOfType("energy_fn"); + if (!fn) return failure(); + auto module = op->getParentOfType(); + auto callee = module.lookupSymbol(fn.getValue()); + if (!callee || callee.isExternal()) { + op->emitError("EBM lowering: energy_fn @") << fn.getValue() << " must be defined in this module"; + return failure(); + } + if (callee.getNumArguments() != 2 || callee.getNumResults() != 1 || + callee.getArgument(0).getType() != op->getOperand(1).getType() || + callee.getArgument(1).getType() != op->getOperand(0).getType() || + callee.getResultTypes()[0] != op->getResult(0).getType()) { + op->emitError("EBM lowering: energy_fn must be E(state, context) -> energies matching the op types"); + return failure(); + } + rewriter.replaceOpWithNewOp(op, callee, ValueRange{op->getOperand(1), op->getOperand(0)}); + return success(); + } +}; + +// inner_step(y, g) {eta} -> y - eta * g +struct LowerInnerStep : public RewritePattern { + LowerInnerStep(MLIRContext *ctx) : RewritePattern(kInnerStepOp, 1, ctx) {} + LogicalResult matchAndRewrite(Operation *op, PatternRewriter &rewriter) const override { + auto eta = op->getAttrOfType("eta"); + auto ty = staticF32(op->getOperand(0)); + if (!eta || !ty || op->getOperand(1).getType() != ty) { + op->emitError("EBM lowering: inner_step requires static f32 state/grad of one type and eta"); + return failure(); + } + Location loc = op->getLoc(); + Value scaled = rewriter.create(loc, op->getOperand(1), splat(rewriter, loc, ty, eta.getValueAsDouble())); + rewriter.replaceOpWithNewOp(op, op->getOperand(0), scaled); + return success(); + } +}; + +// One Philox-4x32-10 evaluation on scalar i32 words; returns the four output words. +static SmallVector philox4x32(OpBuilder &b, Location loc, ArrayRef ctr, ArrayRef key) { + Type i32 = b.getI32Type(); + auto c = [&](uint32_t v) { return b.create(loc, b.getIntegerAttr(i32, (int64_t)v)).getResult(); }; + Value m0 = c(0xD2511F53u), m1 = c(0xCD9E8D57u), w0 = c(0x9E3779B9u), w1 = c(0xBB67AE85u); + Value c0 = ctr[0], c1 = ctr[1], c2 = ctr[2], c3 = ctr[3], k0 = key[0], k1 = key[1]; + for (int round = 0; round < 10; ++round) { + auto p0 = b.create(loc, m0, c0); + auto p1 = b.create(loc, m1, c2); + Value n0 = b.create(loc, b.create(loc, p1.getHigh(), c1), k0); + Value n1 = p1.getLow(); + Value n2 = b.create(loc, b.create(loc, p0.getHigh(), c3), k1); + Value n3 = p0.getLow(); + c0 = n0; c1 = n1; c2 = n2; c3 = n3; + k0 = b.create(loc, k0, w0); + k1 = b.create(loc, k1, w1); + } + return {c0, c1, c2, c3}; +} + +// word (u32) -> (word + 0.5) * 2^-32 in f64, as the reference does. +static Value uniform(OpBuilder &b, Location loc, Value word) { + Type f64 = b.getF64Type(); + Value asF = b.create(loc, f64, word); + Value half = b.create(loc, b.getF64FloatAttr(0.5)); + Value scale = b.create(loc, b.getF64FloatAttr(std::ldexp(1.0, -32))); + return b.create(loc, b.create(loc, asF, half), scale); +} + +// langevin_step(y, key, captures...) -> (y - eta*grad + scale*noise, key + [0, 1]) +struct LowerLangevin : public RewritePattern { + LowerLangevin(MLIRContext *ctx) : RewritePattern(kLangevinOp, 1, ctx) {} + LogicalResult matchAndRewrite(Operation *op, PatternRewriter &rewriter) const override { + auto fn = op->getAttrOfType("energy_fn"); + auto eta = op->getAttrOfType("eta"); + auto temperature = op->getAttrOfType("temperature"); + auto manifold = op->getAttrOfType("manifold"); + if (!fn || !eta || !temperature || !manifold) return failure(); + if (manifold.getValue() != "euclidean") { + op->emitError("EBM lowering: manifold \"") << manifold.getValue() + << "\" has no native integrator yet; only \"euclidean\" lowers"; + return failure(); + } + if (eta.getValueAsDouble() <= 0.0 || temperature.getValueAsDouble() < 0.0) { + op->emitError("EBM lowering: langevin_step requires eta > 0 and temperature >= 0"); + return failure(); + } + Value state = op->getOperand(0), key = op->getOperand(1); + auto stateTy = staticF32(state); + auto keyTy = dyn_cast(key.getType()); + if (!stateTy || !keyTy || keyTy.getShape() != ArrayRef{2} || !keyTy.getElementType().isInteger(64)) { + op->emitError("EBM lowering: langevin_step requires a static f32 state and a tensor<2xi64> key"); + return failure(); + } + if (op->getResult(0).getType() != stateTy || op->getResult(1).getType() != keyTy) { + op->emitError("EBM lowering: langevin_step results must match its state and key types"); + return failure(); + } + auto module = op->getParentOfType(); + auto backward = module.lookupSymbol((fn.getValue() + "__bwd").str()); + if (!backward || backward.isExternal()) { + op->emitError("EBM lowering: the compiler-derived gradient @") << fn.getValue() + << "__bwd is absent; run tessera-autodiff-paired on an energy_fn marked " + "tessera.autodiff = \"reverse\" before this pass"; + return failure(); + } + // @E__bwd(state, captures..., cotangent) -> (dstate, dcaptures...) + SmallVector captures(op->getOperands().begin() + 2, op->getOperands().end()); + if (backward.getNumArguments() != captures.size() + 2 || backward.getNumResults() < 1 || + backward.getArgument(0).getType() != stateTy || backward.getResultTypes()[0] != stateTy) { + op->emitError("EBM lowering: @") << fn.getValue() << "__bwd must be (state, captures..., cotangent) -> (dstate, ...)"; + return failure(); + } + for (auto [i, capture] : llvm::enumerate(captures)) + if (backward.getArgument(i + 1).getType() != capture.getType()) { + op->emitError("EBM lowering: capture ") << i << " type disagrees with @" << fn.getValue() << "__bwd"; + return failure(); + } + auto cotTy = dyn_cast(backward.getArgument(captures.size() + 1).getType()); + if (!cotTy || !cotTy.hasStaticShape() || !cotTy.getElementType().isF32()) { + op->emitError("EBM lowering: the energy cotangent must be a static f32 tensor"); + return failure(); + } + Location loc = op->getLoc(); + SmallVector args(captures.size() + 2); + args[0] = state; + for (auto [i, capture] : llvm::enumerate(captures)) args[i + 1] = capture; + args[captures.size() + 1] = splat(rewriter, loc, cotTy, 1.0); // dE/dE = 1 per energy + Value grad = rewriter.create(loc, backward, args).getResult(0); + Value step = rewriter.create(loc, grad, splat(rewriter, loc, stateTy, eta.getValueAsDouble())); + Value next = rewriter.create(loc, state, step); + const double noiseScale = std::sqrt(2.0 * eta.getValueAsDouble() * temperature.getValueAsDouble()); + if (noiseScale > 0.0) { + // Standard normals per element from Philox on (flat index, key[1]), key[0]. + Type i32 = rewriter.getI32Type(), i64 = rewriter.getI64Type(); + Value c0i = rewriter.create(loc, 0); + Value c1i = rewriter.create(loc, 1); + Value key0 = rewriter.create(loc, key, ValueRange{c0i}); + Value key1 = rewriter.create(loc, key, ValueRange{c1i}); + Value c32 = rewriter.create(loc, rewriter.getIntegerAttr(i64, 32)); + Value k0 = rewriter.create(loc, i32, key0); + Value k1 = rewriter.create(loc, i32, rewriter.create(loc, key0, c32)); + Value s0 = rewriter.create(loc, i32, key1); + Value s1 = rewriter.create(loc, i32, rewriter.create(loc, key1, c32)); + Value zero32 = rewriter.create(loc, rewriter.getIntegerAttr(i32, 0)); + Value init = rewriter.create(loc, stateTy.getShape(), stateTy.getElementType()); + const int64_t rank = stateTy.getRank(); + SmallVector maps{rewriter.getMultiDimIdentityMap(rank)}; + SmallVector iterators(rank, utils::IteratorType::parallel); + auto generic = rewriter.create( + loc, TypeRange{stateTy}, ValueRange{}, ValueRange{init}, maps, iterators, + [&](OpBuilder &b, Location l, ValueRange) { + // flat index = sum(idx[d] * stride[d]) in row-major order. + Value flat = b.create(l, 0); + int64_t stride = 1; + for (int64_t d = rank - 1; d >= 0; --d) { + Value idx = b.create(l, d); + Value s = b.create(l, stride); + flat = b.create(l, flat, b.create(l, idx, s)); + stride *= stateTy.getShape()[d]; + } + Value flat32 = b.create(l, i32, flat); + auto words = philox4x32(b, l, {flat32, s0, s1, zero32}, {k0, k1}); + Value u0 = uniform(b, l, words[0]), u1 = uniform(b, l, words[1]); + Value minusTwo = b.create(l, b.getF64FloatAttr(-2.0)); + Value twoPi = b.create(l, b.getF64FloatAttr(2.0 * M_PI)); + Value r = b.create(l, b.create(l, minusTwo, b.create(l, u0))); + Value z = b.create(l, r, b.create(l, b.create(l, twoPi, u1))); + b.create(l, b.create(l, b.getF32Type(), z).getResult()); + }); + Value noise = rewriter.create(loc, generic.getResult(0), splat(rewriter, loc, stateTy, noiseScale)); + next = rewriter.create(loc, next, noise); + } + auto bump = DenseElementsAttr::get(keyTy, ArrayRef{0, 1}); + Value nextKey = rewriter.create(loc, key, rewriter.create(loc, bump)); + rewriter.replaceOp(op, {next, nextKey}); + return success(); + } +}; + +struct EBMLowerLangevinPass + : public PassWrapper> { + MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(EBMLowerLangevinPass) + StringRef getArgument() const final { return "tessera-ebm-lower-langevin"; } + StringRef getDescription() const final { + return "Lower tessera_ebm.energy / inner_step / langevin_step (euclidean) to " + "arith/linalg over the compiler-derived gradient (@E__bwd) with " + "on-device Philox noise."; + } + void getDependentDialects(DialectRegistry ®istry) const override { + registry.insert(); + } + void runOnOperation() override { + RewritePatternSet patterns(&getContext()); + patterns.add(&getContext()); + if (failed(applyPatternsGreedily(getOperation(), std::move(patterns)))) + signalPassFailure(); + } +}; + +} // namespace + +std::unique_ptr createEBMLowerLangevinPass() { + return std::make_unique(); +} + +} // namespace tessera diff --git a/src/solvers/ebm/test/ir/passes/lower_langevin_quadratic.mlir b/src/solvers/ebm/test/ir/passes/lower_langevin_quadratic.mlir new file mode 100644 index 000000000..f923f598d --- /dev/null +++ b/src/solvers/ebm/test/ir/passes/lower_langevin_quadratic.mlir @@ -0,0 +1,64 @@ +// RUN: ts-ebm-opt --tessera-ebm-lower-langevin %s | FileCheck %s +// +// The first EBM lowering (2026-09-16): energy / inner_step / langevin_step +// become arith + linalg over the compiler-derived gradient @E__bwd. The +// gradient function here is written by hand only because this driver has no +// tessera dialect; in the JIT the paired autodiff pass produces it. + +module { + // E(state, context) -> per-row energies; @E__bwd(state, context, cot) -> (dstate, dcontext) + func.func @E(%y: tensor<4x8xf32>, %x: tensor<4x8xf32>) -> tensor<4xf32> { + %z = arith.constant dense<0.0> : tensor<4xf32> + return %z : tensor<4xf32> + } + func.func @E__bwd(%y: tensor<4x8xf32>, %x: tensor<4x8xf32>, %cot: tensor<4xf32>) -> (tensor<4x8xf32>, tensor<4x8xf32>) { + %d = arith.subf %y, %x : tensor<4x8xf32> + return %d, %d : tensor<4x8xf32>, tensor<4x8xf32> + } + + // CHECK-LABEL: func.func @sample + // CHECK: %[[ONES:.*]] = arith.constant dense<1.000000e+00> : tensor<4xf32> + // CHECK: %[[G:.*]]:2 = call @E__bwd(%{{.*}}, %{{.*}}, %[[ONES]]) + // CHECK: arith.mulf %[[G]]#0, %{{.*}} : tensor<4x8xf32> + // CHECK: arith.subf + // CHECK: linalg.generic + // CHECK: linalg.index + // CHECK: arith.mului_extended + // CHECK: math.log + // CHECK: math.cos + // CHECK: arith.truncf + // CHECK: arith.addf {{.*}} : tensor<4x8xf32> + // CHECK: arith.addi %{{.*}}, %{{.*}} : tensor<2xi64> + // CHECK-NOT: tessera_ebm. + func.func @sample(%y: tensor<4x8xf32>, %x: tensor<4x8xf32>, %key: tensor<2xi64>) -> (tensor<4x8xf32>, tensor<2xi64>) { + %r:2 = "tessera_ebm.langevin_step"(%y, %key, %x) + { energy_fn = @E, eta = 0.1 : f64, temperature = 0.5 : f64, manifold = "euclidean" } + : (tensor<4x8xf32>, tensor<2xi64>, tensor<4x8xf32>) -> (tensor<4x8xf32>, tensor<2xi64>) + return %r#0, %r#1 : tensor<4x8xf32>, tensor<2xi64> + } + + // Temperature 0 draws no noise at all: pure gradient descent, no linalg. + // CHECK-LABEL: func.func @descent + // CHECK: call @E__bwd + // CHECK-NOT: linalg.generic + // CHECK: arith.subf + // CHECK-NOT: tessera_ebm. + func.func @descent(%y: tensor<4x8xf32>, %x: tensor<4x8xf32>, %key: tensor<2xi64>) -> tensor<4x8xf32> { + %r:2 = "tessera_ebm.langevin_step"(%y, %key, %x) + { energy_fn = @E, eta = 0.1 : f64, temperature = 0.0 : f64, manifold = "euclidean" } + : (tensor<4x8xf32>, tensor<2xi64>, tensor<4x8xf32>) -> (tensor<4x8xf32>, tensor<2xi64>) + return %r#0 : tensor<4x8xf32> + } + + // energy -> call E(state, context); inner_step -> y - eta * g. + // CHECK-LABEL: func.func @evaluate + // CHECK: call @E(%arg1, %arg0) + // CHECK: arith.mulf + // CHECK: arith.subf + // CHECK-NOT: tessera_ebm. + func.func @evaluate(%x: tensor<4x8xf32>, %y: tensor<4x8xf32>, %g: tensor<4x8xf32>) -> (tensor<4xf32>, tensor<4x8xf32>) { + %e = "tessera_ebm.energy"(%x, %y) { energy_fn = @E } : (tensor<4x8xf32>, tensor<4x8xf32>) -> tensor<4xf32> + %s = "tessera_ebm.inner_step"(%y, %g) { eta = 0.25 : f64 } : (tensor<4x8xf32>, tensor<4x8xf32>) -> tensor<4x8xf32> + return %e, %s : tensor<4xf32>, tensor<4x8xf32> + } +} diff --git a/src/solvers/ebm/test/ir/passes/lower_langevin_rejects.mlir b/src/solvers/ebm/test/ir/passes/lower_langevin_rejects.mlir new file mode 100644 index 000000000..ff7ad3f8f --- /dev/null +++ b/src/solvers/ebm/test/ir/passes/lower_langevin_rejects.mlir @@ -0,0 +1,25 @@ +// RUN: ts-ebm-opt --tessera-ebm-lower-langevin %s 2>&1 | FileCheck %s +// +// Fail closed: a manifold without a native integrator, and a langevin_step +// whose energy has no compiler-derived gradient in the module. + +module { + func.func @E__bwd(%y: tensor<4x8xf32>, %cot: tensor<4xf32>) -> tensor<4x8xf32> { + return %y : tensor<4x8xf32> + } + func.func @sphere(%y: tensor<4x8xf32>, %key: tensor<2xi64>) -> tensor<4x8xf32> { + %r:2 = "tessera_ebm.langevin_step"(%y, %key) + { energy_fn = @E, eta = 0.1 : f64, temperature = 0.5 : f64, manifold = "sphere" } + : (tensor<4x8xf32>, tensor<2xi64>) -> (tensor<4x8xf32>, tensor<2xi64>) + return %r#0 : tensor<4x8xf32> + } + func.func @no_gradient(%y: tensor<4x8xf32>, %key: tensor<2xi64>) -> tensor<4x8xf32> { + %r:2 = "tessera_ebm.langevin_step"(%y, %key) + { energy_fn = @Missing, eta = 0.1 : f64, temperature = 0.5 : f64, manifold = "euclidean" } + : (tensor<4x8xf32>, tensor<2xi64>) -> (tensor<4x8xf32>, tensor<2xi64>) + return %r#0 : tensor<4x8xf32> + } +} + +// CHECK-DAG: manifold "sphere" has no native integrator yet +// CHECK-DAG: the compiler-derived gradient @Missing__bwd is absent diff --git a/src/solvers/ebm/tools/ts-ebm-opt.cpp b/src/solvers/ebm/tools/ts-ebm-opt.cpp index 8850d04ae..08c9b5e90 100644 --- a/src/solvers/ebm/tools/ts-ebm-opt.cpp +++ b/src/solvers/ebm/tools/ts-ebm-opt.cpp @@ -17,6 +17,9 @@ #include "tessera/EBM/EBMDialect.h" #include "tessera/EBM/EBMPasses.h" #include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/Dialect/Linalg/IR/Linalg.h" +#include "mlir/Dialect/Math/IR/Math.h" +#include "mlir/Dialect/Tensor/IR/Tensor.h" #include "mlir/Dialect/Func/IR/FuncOps.h" #include "mlir/Dialect/SCF/IR/SCF.h" #include "mlir/IR/DialectRegistry.h" @@ -38,11 +41,15 @@ int main(int argc, char **argv) { registry.insert(); registry.insert(); registry.insert(); + registry.insert(); + registry.insert(); + registry.insert(); registerPass(tessera::createEBMCanonicalizePass); registerPass(tessera::createEBMFuseEnergyGradPass); registerPass(tessera::createEBMCheckpointInnerLoopPass); registerPass(tessera::createEBMPipelineCandidatesPass); + registerPass(tessera::createEBMLowerLangevinPass); // W0.2 — `checkpoint-inner-loop` is NOT in the default pipeline. // diff --git a/src/transforms/include/Tessera/Transforms/Passes.h b/src/transforms/include/Tessera/Transforms/Passes.h index 166329543..08c2e2edb 100644 --- a/src/transforms/include/Tessera/Transforms/Passes.h +++ b/src/transforms/include/Tessera/Transforms/Passes.h @@ -179,6 +179,10 @@ std::unique_ptr createAutodiffPass(); // Body: src/transforms/lib/AutodiffPairedPass.cpp std::unique_ptr createAutodiffPairedPass(); std::unique_ptr createNativeTapeToGPUPass(); +// RowProgramToGPUPass -- a [rows, features] tensor row program (parallel linalg +// bodies, feature-axis reductions, uniform integer vectors, scf.for) becomes one +// cooperative kernel in the native storage ABI. Body: RowProgramToGPUPass.cpp. +std::unique_ptr createRowProgramToGPUPass(); // AutodiffForwardPass — AD-FWD-CORE-1. Emits a separate paired JVP: // @f__jvp(inputs..., input_tangents...) -> (primals..., output_tangents...) diff --git a/src/transforms/lib/CMakeLists.txt b/src/transforms/lib/CMakeLists.txt index 8b45fe6d5..81edd2cbe 100644 --- a/src/transforms/lib/CMakeLists.txt +++ b/src/transforms/lib/CMakeLists.txt @@ -40,6 +40,7 @@ add_library(TesseraPasses # ODS scaffold lives at src/compiler/ir/include/Tessera/AdjointInterface.td. AutodiffPass.cpp NativeTapeToGPUPass.cpp + RowProgramToGPUPass.cpp AutodiffPairedPass.cpp AutodiffForwardPass.cpp diff --git a/src/transforms/lib/Passes.cpp b/src/transforms/lib/Passes.cpp index 267f362de..6837a57aa 100644 --- a/src/transforms/lib/Passes.cpp +++ b/src/transforms/lib/Passes.cpp @@ -512,6 +512,7 @@ void registerTesseraPasses() { // Phase 2 — paired forward/backward autodiff (separate @f__bwd function). ::mlir::registerPass([]() { return createAutodiffPairedPass(); }); ::mlir::registerPass([]() { return createNativeTapeToGPUPass(); }); + ::mlir::registerPass([]() { return createRowProgramToGPUPass(); }); // AD-FWD-CORE-1 — paired forward-mode JVP through TangentInterface. ::mlir::registerPass([]() { return createAutodiffForwardPass(); }); ::mlir::registerPass([]() { return createAutodiffHvpPreparePass(); }); diff --git a/src/transforms/lib/RowProgramToGPUPass.cpp b/src/transforms/lib/RowProgramToGPUPass.cpp new file mode 100644 index 000000000..3783f483c --- /dev/null +++ b/src/transforms/lib/RowProgramToGPUPass.cpp @@ -0,0 +1,634 @@ +//===- RowProgramToGPUPass.cpp ---------------------------------*- C++ -*-===// +// +// RowProgramToGPUPass (2026-09-16, EBM_NATIVE_LOOP_ARCHITECTURE.md slice G2): +// a *row program* -- a tensor-level function over `[rows, features]` tensors +// built from parallel linalg.generic bodies, single-axis linalg.reduce over +// the feature axis, small uniform integer vectors (RNG keys) and scf.for +// loops carrying those tensors -- becomes one cooperative GPU kernel: +// +// * one block per row, one lane per feature (block = next power of two +// >= features, inactive lanes masked); +// * every `[rows, features]` value is one f32 register per lane, every +// `[rows]` / `[rows, 1]` value is a block-uniform scalar, every +// `[n] xi64` value (n <= 8) is n uniform scalars, every non-tensor value +// is cloned as is; +// * a linalg.generic body is cloned once per lane (or once per block for +// row results) with `linalg.index 0/1` mapped to (row, lane); +// * a linalg.reduce over the feature axis is an *ordered* reduction: +// lanes write to workgroup memory, a barrier, the block leader folds the +// combiner body over features 0..F-1 in order, a barrier, all lanes read +// the result -- the same order the CPU lane's sequential loop uses, so +// reduced quantities agree to the same rounding on both; +// * scf.for loops carry their tensors as per-lane / uniform registers, so a +// K-step sampling loop never leaves the kernel. +// +// The kernel uses the native GPU storage ABI (`!llvm.ptr<1>` per tensor plus +// one index scratch, `tile.alloc_shared` marker for the arena sizer), so +// `build_native_gpu_storage` packages it for gfx1151/gfx1201/sm_120 and +// `replay_arena_ir` validates it. The input module is retained verbatim in +// `tessera.row_program.source` for that replay. +// +// This is the device route for the EBM Langevin loop (the compiler-derived +// gradient is exactly such a row program after tessera-to-linalg) and for any +// future row-wise program; nothing in it knows about energies. +// +// Fails closed on anything outside the contract (dynamic shapes, calls, +// tensors of other ranks, non-parallel generics, reductions over other axes, +// unsupported ops) with a diagnostic naming the op. +// +//===----------------------------------------------------------------------===// +#include "Tessera/Dialect/Tile/TileDialect.h" +#include "Tessera/IR/TesseraOps.h" +#include "Tessera/Transforms/Passes.h" + +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/Dialect/Func/IR/FuncOps.h" +#include "mlir/Dialect/GPU/IR/GPUDialect.h" +#include "mlir/Dialect/LLVMIR/LLVMDialect.h" +#include "mlir/Dialect/Linalg/IR/Linalg.h" +#include "mlir/Dialect/Math/IR/Math.h" +#include "mlir/Dialect/MemRef/IR/MemRef.h" +#include "mlir/Dialect/SCF/IR/SCF.h" +#include "mlir/Dialect/Tensor/IR/Tensor.h" +#include "mlir/IR/Builders.h" +#include "mlir/IR/BuiltinOps.h" +#include "mlir/IR/IRMapping.h" +#include "mlir/Parser/Parser.h" +#include "mlir/Pass/Pass.h" +#include "llvm/ADT/DenseMap.h" +#include "llvm/Support/raw_ostream.h" + +#include + +using namespace mlir; + +// The emitter's numeric contract is IEEE arithmetic. `math.sqrt` on the NVVM +// route lowers to libdevice's __nv_sqrtf, which takes the approximate path +// (MUFU.SQRT, 1 ulp) because MLIR's pipeline leaves the NVVM reflect flag for +// precise sqrt unset (measured 2026-09-16: one row of the row-normalization +// proof differed by 1 ulp on sm_120 only). convert-gpu-to-nvvm marks the LLVM +// math intrinsics illegal, so on NVIDIA the kernel calls libdevice's +// rounding-explicit __nv_fsqrt_rn / __nv_dsqrt_rn (correctly rounded whatever +// the reflect flags say); on ROCm llvm.intr.sqrt lowers to AMDGPU's correctly +// rounded expansion. +static void pinSqrt(gpu::GPUFuncOp kernel, StringRef backend) { + SmallVector ops; + kernel->walk([&](math::SqrtOp op) { ops.push_back(op); }); + if (ops.empty()) return; + auto gpuModule = kernel->getParentOfType(); + for (math::SqrtOp op : ops) { + OpBuilder b(op); + Value r; + if (backend == "nvidia") { + bool f64 = op.getType().isF64(); + StringRef name = f64 ? "__nv_dsqrt_rn" : "__nv_fsqrt_rn"; + auto fn = gpuModule.lookupSymbol(name); + if (!fn) { + OpBuilder mb(gpuModule.getBodyRegion()); + fn = LLVM::LLVMFuncOp::create(mb, op.getLoc(), name, + LLVM::LLVMFunctionType::get(op.getType(), {op.getType()})); + } + r = LLVM::CallOp::create(b, op.getLoc(), fn, ValueRange{op.getOperand()}).getResult(); + } else { + r = LLVM::SqrtOp::create(b, op.getLoc(), op.getType(), op.getOperand()); + } + op.replaceAllUsesWith(r); + op.erase(); + } +} + + +namespace tessera { +namespace { + +enum class Kind { Lane, Row, Uniform, Scalar, None }; + +struct Slot { + Kind kind = Kind::None; + SmallVector values; // 1 for Lane/Row/Scalar, n for Uniform +}; + +struct Emitter { + ModuleOp module; + func::FuncOp entry; + OpBuilder b; + Location loc; + int64_t rows = 0, feats = 0, lanes = 0; + Value row, lane, active, leader, shared, fzero; + SmallVector pointers; // one per entry argument then per result + DenseMap slots; + LLVM::LLVMArrayType sharedTy; + bool broken = false; + + Emitter(ModuleOp m, func::FuncOp f) : module(m), entry(f), b(m.getContext()), loc(f.getLoc()) {} + + InFlightDiagnostic fail(Operation *op, const Twine &msg) { + broken = true; + return op->emitError("row program: ") << msg; + } + + // --- classification ---------------------------------------------------- + Kind classify(Type type, Operation *at) { + auto tensor = dyn_cast(type); + if (!tensor) return Kind::Scalar; + if (!tensor.hasStaticShape()) { fail(at, "dynamic tensor shapes are not admitted"); return Kind::None; } + Type elem = tensor.getElementType(); + if (tensor.getRank() == 0) return Kind::Scalar; + if (tensor.getRank() == 2 && elem.isF32()) { + if (tensor.getDimSize(0) == rows && tensor.getDimSize(1) == feats) return Kind::Lane; + if (tensor.getDimSize(0) == rows && tensor.getDimSize(1) == 1) return Kind::Row; + fail(at, "rank-2 tensors must be [rows, features] or [rows, 1]"); + return Kind::None; + } + if (tensor.getRank() == 1) { + if (elem.isF32() && tensor.getDimSize(0) == rows) return Kind::Row; + if (elem.isInteger() && tensor.getDimSize(0) >= 1 && tensor.getDimSize(0) <= 8) return Kind::Uniform; + fail(at, "rank-1 tensors must be f32 [rows] or an integer vector of at most 8 elements"); + return Kind::None; + } + fail(at, "unsupported tensor type"); + return Kind::None; + } + + Slot &slot(Value v, Operation *at) { + auto it = slots.find(v); + if (it == slots.end()) { + fail(at, "value has no device mapping (produced by an unsupported op?)"); + static Slot dead; + return dead; + } + return it->second; + } + + // Scalar broadcast of a slot for use inside a body: Lane/Row/Scalar -> the one value. + Value scalarOf(Value v, Operation *at) { + Slot &s = slot(v, at); + if (s.values.size() != 1) { fail(at, "expected a single-valued operand"); return {}; } + return s.values[0]; + } + + // --- device helpers ---------------------------------------------------- + Value idx(int64_t v) { return arith::ConstantIndexOp::create(b, loc, v); } + Value i64(Value index) { return arith::IndexCastOp::create(b, loc, b.getI64Type(), index); } + Value gep(Value base, Type elem, Value index64) { + return LLVM::GEPOp::create(b, loc, base.getType(), elem, base, ValueRange{index64}); + } + Value load(Value base, Type elem, Value index) { return LLVM::LoadOp::create(b, loc, elem, gep(base, elem, i64(index))); } + void store(Value value, Value base, Value index) { + LLVM::StoreOp::create(b, loc, value, gep(base, value.getType(), i64(index))); + } + Value laneIndex() { // row * feats + lane + return arith::AddIOp::create(b, loc, arith::MulIOp::create(b, loc, row, idx(feats)), lane); + } + Value sharedPtr(Value index) { + Value zero = arith::ConstantOp::create(b, loc, b.getI64IntegerAttr(0)); + return LLVM::GEPOp::create(b, loc, shared.getType(), sharedTy, shared, ValueRange{zero, i64(index)}); + } + Value zeroOf(Type t) { + if (auto ft = dyn_cast(t)) return arith::ConstantOp::create(b, loc, b.getFloatAttr(ft, 0.0)); + return arith::ConstantOp::create(b, loc, b.getIntegerAttr(t, 0)); + } + + // Load an argument tensor into a slot. + Slot loadArgument(Value pointer, Type type, Operation *at) { + Slot s; + s.kind = classify(type, at); + auto tensor = dyn_cast(type); + Type elem = tensor ? tensor.getElementType() : type; + switch (s.kind) { + case Kind::Lane: { + auto guarded = scf::IfOp::create(b, loc, TypeRange{elem}, active, /*withElse=*/true); + { + OpBuilder::InsertionGuard g(b); + b.setInsertionPointToStart(&guarded.getThenRegion().front()); + scf::YieldOp::create(b, loc, load(pointer, elem, laneIndex())); + b.setInsertionPointToStart(&guarded.getElseRegion().front()); + scf::YieldOp::create(b, loc, zeroOf(elem)); + } + s.values.push_back(guarded.getResult(0)); + break; + } + case Kind::Row: + s.values.push_back(load(pointer, elem, row)); + break; + case Kind::Uniform: + for (int64_t k = 0; k < tensor.getDimSize(0); ++k) s.values.push_back(load(pointer, elem, idx(k))); + break; + case Kind::Scalar: + if (tensor) s.values.push_back(load(pointer, elem, idx(0))); + else { fail(at, "non-tensor arguments are not admitted"); } + break; + case Kind::None: + break; + } + return s; + } + + void storeResult(Slot &s, Value pointer, Operation *at) { + switch (s.kind) { + case Kind::Lane: { + auto guarded = scf::IfOp::create(b, loc, active, /*withElse=*/false); + OpBuilder::InsertionGuard g(b); + b.setInsertionPointToStart(&guarded.getThenRegion().front()); + store(s.values[0], pointer, laneIndex()); + break; + } + case Kind::Row: + case Kind::Uniform: + case Kind::Scalar: { + auto guarded = scf::IfOp::create(b, loc, leader, /*withElse=*/false); + OpBuilder::InsertionGuard g(b); + b.setInsertionPointToStart(&guarded.getThenRegion().front()); + if (s.kind == Kind::Row || s.kind == Kind::Scalar) store(s.values[0], pointer, s.kind == Kind::Row ? row : idx(0)); + else for (auto [k, v] : llvm::enumerate(s.values)) store(v, pointer, idx(k)); + break; + } + case Kind::None: + fail(at, "result has no device mapping"); + } + } + + // --- ops ---------------------------------------------------------------- + // Clone a scalar body region once with `args` bound to its block arguments. + Value cloneBody(Region ®ion, ArrayRef args, Operation *at, bool inLane) { + Block &body = region.front(); + if (body.getNumArguments() != args.size()) { fail(at, "body arity disagrees"); return {}; } + IRMapping map; + for (auto [arg, value] : llvm::zip(body.getArguments(), args)) map.map(arg, value); + Value result; + for (Operation &op : body) { + if (auto index = dyn_cast(op)) { + if (index.getDim() == 0) map.map(index.getResult(), row); + else if (index.getDim() == 1 && inLane) map.map(index.getResult(), lane); + else { fail(at, "linalg.index outside the (row, lane) mapping"); return {}; } + continue; + } + if (auto yield = dyn_cast(op)) { + if (yield.getNumOperands() != 1) { fail(at, "single-result bodies only"); return {}; } + result = map.lookupOrDefault(yield.getOperand(0)); + continue; + } + if (op.getNumRegions() != 0) { fail(&op, "nested regions inside bodies are not admitted"); return {}; } + for (Type t : op.getResultTypes()) + if (isa(t)) { fail(&op, "tensor-typed ops inside bodies are not admitted"); return {}; } + // Operands defined outside the body (hoisted constants, loop scalars) + // resolve through their device slots, never through the source module. + for (Value operand : op.getOperands()) { + if (map.contains(operand)) continue; + if (isa(operand.getType())) { fail(&op, "tensor operand captured inside a body"); return {}; } + Value v = scalarOf(operand, &op); + if (!v) return {}; + map.map(operand, v); + } + b.clone(op, map); + } + return result; + } + + bool isIdentityMap(AffineMap m, int64_t rank) { return m.isIdentity() && m.getNumDims() == rank; } + // (d0, d1) -> (d0, 0) or (d0, d1) -> (d0): a row-uniform read. + bool isRowMap(AffineMap m) { + if (m.getNumDims() != 2 || m.getNumResults() < 1 || m.getNumResults() > 2) return false; + if (m.getResult(0) != getAffineDimExpr(0, m.getContext())) return false; + if (m.getNumResults() == 2) { + auto c = dyn_cast(m.getResult(1)); + return c && c.getValue() == 0; + } + return true; + } + + LogicalResult emitGeneric(linalg::GenericOp op) { + if (op.getNumResults() != 1) return fail(op, "one-result generics only"); + Kind out = classify(op.getResult(0).getType(), op); + if (out == Kind::None) return failure(); + for (auto it : op.getIteratorTypesArray()) + if (it != utils::IteratorType::parallel) return fail(op, "reductions must be linalg.reduce, not a reduction generic"); + auto maps = op.getIndexingMapsArray(); + const int64_t rank = cast(op.getResult(0).getType()).getRank(); + if (!isIdentityMap(maps.back(), rank)) return fail(op, "the output map must be the identity"); + const bool laneBody = out == Kind::Lane; + auto perElement = [&](int64_t element) -> LogicalResult { + SmallVector args; + for (auto [i, input] : llvm::enumerate(op.getInputs())) { + Slot &s = slot(input, op); + AffineMap m = maps[i]; + if (s.kind == Kind::Lane) { + if (!laneBody || !isIdentityMap(m, 2)) return fail(op, "a [rows, features] input needs a lane body and an identity map"); + args.push_back(s.values[0]); + } else if (s.kind == Kind::Row) { + if (!(isRowMap(m) || (rank == 1 && isIdentityMap(m, 1)) || (rank == 2 && isIdentityMap(m, 2) && !laneBody))) + return fail(op, "a row input must be read with a row map"); + args.push_back(s.values[0]); + } else if (s.kind == Kind::Scalar) { + args.push_back(s.values[0]); + } else if (s.kind == Kind::Uniform) { + if (out != Kind::Uniform || !isIdentityMap(m, 1)) return fail(op, "uniform vectors combine only elementwise with uniform vectors"); + args.push_back(s.values[element]); + } else return failure(); + } + // outs block arguments: the init value if it has one, else zero. + for (Value init : op.getOutputs()) { + auto it = slots.find(init); + if (it != slots.end() && !it->second.values.empty()) + args.push_back(it->second.kind == Kind::Uniform ? it->second.values[element] : it->second.values[0]); + else args.push_back(zeroOf(cast(init.getType()).getElementType())); + } + Value r = cloneBody(op.getRegion(), args, op, laneBody); + if (!r) return failure(); + slots[op.getResult(0)].values.push_back(r); + return success(); + }; + slots[op.getResult(0)].kind = out; + if (out == Kind::Uniform) { + const int64_t n = cast(op.getResult(0).getType()).getDimSize(0); + for (int64_t k = 0; k < n; ++k) if (failed(perElement(k))) return failure(); + return success(); + } + return perElement(0); + } + + LogicalResult emitReduce(linalg::ReduceOp op) { + if (op.getNumResults() != 1 || op.getInputs().size() != 1) return fail(op, "single-input reductions only"); + auto dims = op.getDimensions(); + if (dims.size() != 1 || dims[0] != 1) return fail(op, "reductions must be over the feature axis (dimension 1)"); + Slot &in = slot(op.getInputs()[0], op); + if (in.kind != Kind::Lane) return fail(op, "the reduced operand must be [rows, features]"); + Kind out = classify(op.getResult(0).getType(), op); + if (out != Kind::Row) return fail(op, "the reduction result must be [rows]"); + Value init = scalarOf(op.getInits()[0], op); + if (!init) return failure(); + Type elem = init.getType(); + // lanes -> shared memory, in order. + { + auto guarded = scf::IfOp::create(b, loc, active, /*withElse=*/false); + OpBuilder::InsertionGuard g(b); + b.setInsertionPointToStart(&guarded.getThenRegion().front()); + LLVM::StoreOp::create(b, loc, in.values[0], sharedPtr(lane)); + } + gpu::BarrierOp::create(b, loc); + auto folded = scf::IfOp::create(b, loc, TypeRange{elem}, leader, /*withElse=*/true); + { + OpBuilder::InsertionGuard g(b); + b.setInsertionPointToStart(&folded.getThenRegion().front()); + auto loop = scf::ForOp::create(b, loc, idx(0), idx(feats), idx(1), ValueRange{init}); + { + OpBuilder::InsertionGuard g2(b); + b.setInsertionPointToStart(loop.getBody()); + Value element = LLVM::LoadOp::create(b, loc, elem, sharedPtr(loop.getInductionVar())); + Value acc = cloneBody(op.getRegion(), {element, loop.getRegionIterArgs()[0]}, op, /*inLane=*/false); + if (!acc) return failure(); + scf::YieldOp::create(b, loc, acc); + } + LLVM::StoreOp::create(b, loc, loop.getResult(0), sharedPtr(idx(0))); + scf::YieldOp::create(b, loc, loop.getResult(0)); + b.setInsertionPointToStart(&folded.getElseRegion().front()); + scf::YieldOp::create(b, loc, init); + } + gpu::BarrierOp::create(b, loc); + Value result = LLVM::LoadOp::create(b, loc, elem, sharedPtr(idx(0))); + gpu::BarrierOp::create(b, loc); // shared memory is free for the next reduction + Slot s; s.kind = Kind::Row; s.values.push_back(result); + slots[op.getResult(0)] = s; + return success(); + } + + LogicalResult emitFor(scf::ForOp op) { + for (Value v : {op.getLowerBound(), op.getUpperBound(), op.getStep()}) + if (isa(v.getType())) return fail(op, "loop bounds must be scalars"); + SmallVector flat; + SmallVector layout; + for (Value init : op.getInitArgs()) { + Slot &s = slot(init, op); + if (s.kind == Kind::None) return failure(); + layout.push_back(s); + flat.append(s.values.begin(), s.values.end()); + } + auto loop = scf::ForOp::create(b, loc, scalarOf(op.getLowerBound(), op), scalarOf(op.getUpperBound(), op), + scalarOf(op.getStep(), op), flat); + { + OpBuilder::InsertionGuard g(b); + b.setInsertionPointToStart(loop.getBody()); + Slot iv; iv.kind = Kind::Scalar; iv.values.push_back(loop.getInductionVar()); + slots[op.getInductionVar()] = iv; + unsigned cursor = 0; + for (auto [arg, shape] : llvm::zip(op.getRegionIterArgs(), layout)) { + Slot s; s.kind = shape.kind; + for (unsigned k = 0; k < shape.values.size(); ++k) s.values.push_back(loop.getRegionIterArgs()[cursor++]); + slots[arg] = s; + } + auto yield = cast(op.getBody()->getTerminator()); + if (failed(emitBlock(*op.getBody(), /*terminator=*/yield))) return failure(); + SmallVector yielded; + for (auto [v, shape] : llvm::zip(yield.getOperands(), layout)) { + Slot &s = slot(v, op); + if (s.kind != shape.kind || s.values.size() != shape.values.size()) return fail(op, "loop-carried value changed its mapping"); + yielded.append(s.values.begin(), s.values.end()); + } + scf::YieldOp::create(b, loc, yielded); + } + unsigned cursor = 0; + for (auto [result, shape] : llvm::zip(op.getResults(), layout)) { + Slot s; s.kind = shape.kind; + for (unsigned k = 0; k < shape.values.size(); ++k) s.values.push_back(loop.getResult(cursor++)); + slots[result] = s; + } + return success(); + } + + LogicalResult emitConstant(arith::ConstantOp op) { + Slot s; + if (auto dense = dyn_cast(op.getValue())) { + s.kind = classify(op.getType(), op); + if (s.kind == Kind::None) return failure(); + Type elem = dense.getElementType(); + if (dense.isSplat()) { + Attribute v = dense.getSplatValue(); + Value c = arith::ConstantOp::create(b, loc, elem, cast(v)); + const int64_t n = s.kind == Kind::Uniform ? cast(op.getType()).getDimSize(0) : 1; + for (int64_t k = 0; k < n; ++k) s.values.push_back(c); + } else { + if (s.kind != Kind::Uniform) return fail(op, "non-splat constants must be small uniform integer vectors"); + for (Attribute v : dense.getValues()) + s.values.push_back(arith::ConstantOp::create(b, loc, elem, cast(v))); + } + } else { + s.kind = Kind::Scalar; + s.values.push_back(b.clone(*op.getOperation())->getResult(0)); + } + slots[op.getResult()] = s; + return success(); + } + + LogicalResult emitOp(Operation &op) { + if (auto c = dyn_cast(op)) return emitConstant(c); + if (isa(op)) { slots[op.getResult(0)] = Slot{classify(op.getResult(0).getType(), &op), {}}; return success(); } + if (auto fill = dyn_cast(op)) { + Slot s; s.kind = classify(fill.getResult(0).getType(), &op); + Value v = scalarOf(fill.getInputs()[0], &op); + if (!v) return failure(); + const int64_t n = s.kind == Kind::Uniform ? cast(fill.getResult(0).getType()).getDimSize(0) : 1; + for (int64_t k = 0; k < n; ++k) s.values.push_back(v); + slots[fill.getResult(0)] = s; + return success(); + } + if (auto g = dyn_cast(op)) return emitGeneric(g); + if (auto r = dyn_cast(op)) return emitReduce(r); + if (isa(op)) { + Slot &in = slot(op.getOperand(0), &op); + Kind out = classify(op.getResult(0).getType(), &op); + if (in.kind != Kind::Row || out != Kind::Row) return fail(&op, "reshapes are admitted only between [rows] and [rows, 1]"); + slots[op.getResult(0)] = in; + return success(); + } + if (auto extract = dyn_cast(op)) { + Slot &in = slot(extract.getTensor(), &op); + if (in.kind != Kind::Uniform || extract.getIndices().size() != 1) return fail(&op, "tensor.extract is admitted on uniform vectors only"); + auto c = extract.getIndices()[0].getDefiningOp(); + if (!c || c.value() < 0 || c.value() >= (int64_t)in.values.size()) return fail(&op, "uniform vector extracts need a constant in-range index"); + Slot s; s.kind = Kind::Scalar; s.values.push_back(in.values[c.value()]); + slots[extract.getResult()] = s; + return success(); + } + if (auto loop = dyn_cast(op)) return emitFor(loop); + if (isa(op)) return fail(&op, "calls must be inlined before the row-program lowering"); + if (op.getNumRegions() != 0) return fail(&op, "unsupported region-carrying op"); + // Scalar op: clone with its operands' single values. + IRMapping map; + for (Value operand : op.getOperands()) { + if (isa(operand.getType())) return fail(&op, "tensor operand on an unsupported op"); + Value v = scalarOf(operand, &op); + if (!v) return failure(); + map.map(operand, v); + } + for (Type t : op.getResultTypes()) + if (isa(t)) return fail(&op, "tensor result on an unsupported op"); + Operation *cloned = b.clone(op, map); + for (auto [r, c] : llvm::zip(op.getResults(), cloned->getResults())) { + Slot s; s.kind = Kind::Scalar; s.values.push_back(c); + slots[r] = s; + } + return success(); + } + + LogicalResult emitBlock(Block &block, Operation *terminator) { + for (Operation &op : block) { + if (&op == terminator) break; + if (failed(emitOp(op)) || broken) return failure(); + } + return success(); + } + + LogicalResult run(StringRef backend) { + // Shape from the first [rows, features] argument. + for (Type t : entry.getArgumentTypes()) + if (auto tt = dyn_cast(t)) + if (tt.getRank() == 2 && tt.hasStaticShape() && tt.getElementType().isF32()) { rows = tt.getDimSize(0); feats = tt.getDimSize(1); break; } + if (rows <= 0 || feats <= 0) return entry.emitError("row program: no [rows, features] f32 argument"); + if (feats > 1024) return entry.emitError("row program: at most 1024 features per row (one lane each)"); + lanes = 1; while (lanes < feats) lanes *= 2; + if (!entry.getBody().hasOneBlock()) return entry.emitError("row program: the entry must have one block"); + // Calls are inlined upstream: `--inline` after `--tessera-to-linalg` (the + // Tessera dialect has no inliner interface, so inlining before the linalg + // lowering is refused by the inliner; measured 2026-09-16). + bool calls = false; + entry.walk([&](func::CallOp) { calls = true; }); + if (calls) return entry.emitError("row program: inline every call before lowering (run --inline after --tessera-to-linalg)"); + + // Retain the source for replay, then build the kernel skeleton textually + // (the same shape the SSD and native-tape routes use) and fill its body. + std::string source; { llvm::raw_string_ostream os(source); module.print(os); } + const unsigned count = entry.getNumArguments() + entry.getNumResults(); + std::string text; llvm::raw_string_ostream s(text); + s << "module attributes {tessera.row_program.source = "; + StringAttr::get(module.getContext(), source).print(s); + s << ", tessera.row_program.entry = \"" << entry.getName() << "\", tessera.row_program.rows = " << rows + << " : i64, tessera.row_program.features = " << feats << " : i64, tessera.row_program.backend = \"" << backend + << "\", tessera.autodiff.temporary_bytes = " << lanes * 8 << " : i64} {\n gpu.module @native_row {\n"; + s << " llvm.mlir.global private @row_reduction() {addr_space = 3 : i32, alignment = 16 : i64} : !llvm.array<" << lanes << " x f32>\n"; + s << " gpu.func @row_program("; + for (unsigned i = 0; i < count; ++i) s << "%p" << i << ": !llvm.ptr<1>, "; + s << "%scratch: index) kernel attributes {known_block_size = array} {\n"; + s << " %marker = memref.alloca(%scratch) : memref\n \"tile.alloc_shared\"(%marker) : (memref) -> ()\n"; + s << " gpu.return\n }\n }\n}\n"; + s.flush(); + OwningOpRef lowered = parseSourceString(text, module.getContext()); + if (!lowered) return entry.emitError("row program: kernel skeleton failed to parse"); + gpu::GPUFuncOp kernel; + lowered->walk([&](gpu::GPUFuncOp f) { kernel = f; }); + Operation *ret = kernel.getBody().front().getTerminator(); + b.setInsertionPoint(ret); + for (unsigned i = 0; i < count; ++i) pointers.push_back(kernel.getArgument(i)); + sharedTy = LLVM::LLVMArrayType::get(b.getF32Type(), lanes); + row = gpu::BlockIdOp::create(b, loc, gpu::Dimension::x); + lane = gpu::ThreadIdOp::create(b, loc, gpu::Dimension::x); + active = arith::CmpIOp::create(b, loc, arith::CmpIPredicate::ult, lane, idx(feats)); + leader = arith::CmpIOp::create(b, loc, arith::CmpIPredicate::eq, lane, idx(0)); + LLVM::GlobalOp global; + lowered->walk([&](LLVM::GlobalOp g) { global = g; }); + shared = LLVM::AddressOfOp::create(b, loc, global); + fzero = zeroOf(b.getF32Type()); + for (auto [i, arg] : llvm::enumerate(entry.getArguments())) { + Slot s = loadArgument(pointers[i], arg.getType(), entry); + if (broken || s.kind == Kind::None) return failure(); + slots[arg] = s; + } + auto terminator = cast(entry.getBody().front().getTerminator()); + if (failed(emitBlock(entry.getBody().front(), terminator))) return failure(); + for (auto [r, result] : llvm::enumerate(terminator.getOperands())) { + Slot &s = slot(result, terminator); + storeResult(s, pointers[entry.getNumArguments() + r], terminator); + if (broken) return failure(); + } + pinSqrt(kernel, backend); + // Replace the module contents with the kernel module. + module->setAttrs((*lowered)->getAttrs()); + module.getBody()->clear(); + module.getBody()->getOperations().splice(module.getBody()->end(), lowered->getBody()->getOperations()); + return success(); + } +}; + +struct RowProgramToGPUPass : public PassWrapper> { + MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(RowProgramToGPUPass) + RowProgramToGPUPass() = default; + RowProgramToGPUPass(const RowProgramToGPUPass &other) : PassWrapper(other) {} + Option backend{*this, "backend", llvm::cl::desc("nvidia or rocm"), llvm::cl::init("nvidia")}; + Option entryName{*this, "entry", llvm::cl::desc("The row-program entry function (default: the only function)"), llvm::cl::init("")}; + StringRef getArgument() const final { return "tessera-row-program-to-gpu"; } + StringRef getDescription() const final { + return "Lower a [rows, features] tensor row program (parallel linalg bodies, feature-axis " + "reductions, uniform integer vectors, scf.for) to one cooperative GPU kernel in the " + "native storage ABI: one block per row, one lane per feature, ordered reductions."; + } + void getDependentDialects(DialectRegistry &r) const override { + // The emitted skeleton names `tile.alloc_shared` (the arena sizer marker): + // the tile dialect must be declared here or an assertions-enabled MLIR + // refuses to load it inside the pass manager (Tajasarus, 2026-09-16). + // The provenance attributes are tessera.*-prefixed: parsing them loads the + // Tessera dialect, which an input without Tessera ops never loaded. + r.insert(); + } + void runOnOperation() override { + ModuleOp module = getOperation(); + if (backend != "nvidia" && backend != "rocm") { module.emitError("row program: backend must be nvidia or rocm"); return signalPassFailure(); } + func::FuncOp entry; + SmallVector functions(module.getOps()); + if (entryName.empty()) { + if (functions.size() != 1) { module.emitError("row program: name the entry when the module has several functions"); return signalPassFailure(); } + entry = functions[0]; + } else { + for (auto f : functions) if (f.getName() == entryName) entry = f; + if (!entry) { module.emitError("row program: entry function not found: ") << entryName; return signalPassFailure(); } + } + Emitter emitter(module, entry); + if (failed(emitter.run(backend))) signalPassFailure(); + } +}; + +} // namespace + +std::unique_ptr createRowProgramToGPUPass() { return std::make_unique(); } + +} // namespace tessera diff --git a/src/transforms/lib/TesseraToLinalgPass.cpp b/src/transforms/lib/TesseraToLinalgPass.cpp index ce65b6b86..af17225c7 100644 --- a/src/transforms/lib/TesseraToLinalgPass.cpp +++ b/src/transforms/lib/TesseraToLinalgPass.cpp @@ -986,6 +986,117 @@ struct ReduceLowering : public RewritePattern { // selected extremum for max/min; a second sum reduction counts ties and the // final generic divides the incoming cotangent equally among them. Mean reads // a dynamic axis extent directly from the original input. +// `tessera.unsqueeze` / `tessera.broadcast` (2026-09-16). The sum / static-mean +// reduce adjoint decomposes into exactly these two ops, and until now neither +// had a linalg lowering, so no reduce-sum gradient could reach the CPU JIT +// lane -- the EBM quadratic energy (0.5 * sum((x - y)^2)) was the first to +// hit it. Static shapes only; anything else is left for the residual check +// to report rather than guessed. +struct UnsqueezeLowering : public RewritePattern { + UnsqueezeLowering(MLIRContext *ctx) + : RewritePattern("tessera.unsqueeze", /*benefit=*/1, ctx) {} + + LogicalResult matchAndRewrite(Operation *op, + PatternRewriter &rewriter) const override { + if (op->getNumOperands() != 1 || op->getNumResults() != 1) + return failure(); + auto inTy = dyn_cast(op->getOperand(0).getType()); + auto outTy = dyn_cast(op->getResult(0).getType()); + if (!inTy || !outTy || !inTy.hasStaticShape() || !outTy.hasStaticShape() || + inTy.getElementType() != outTy.getElementType()) + return rewriter.notifyMatchFailure(op, "static ranked tensors required"); + auto axesAttr = op->getAttrOfType("axes"); + if (!axesAttr) + return rewriter.notifyMatchFailure(op, "missing axes"); + const int64_t outRank = outTy.getRank(); + llvm::SmallVector inserted(outRank, false); + for (Attribute a : axesAttr) { + auto ia = dyn_cast(a); + if (!ia) return rewriter.notifyMatchFailure(op, "axes must be integers"); + int64_t axis = ia.getInt(); + if (axis < 0) axis += outRank; + if (axis < 0 || axis >= outRank || inserted[axis]) + return rewriter.notifyMatchFailure(op, "axis out of range or repeated"); + inserted[axis] = true; + } + if (outRank != inTy.getRank() + (int64_t)axesAttr.size()) + return rewriter.notifyMatchFailure(op, "result rank != input rank + axes"); + // Expected result shape: input dims in order with 1s at the inserted axes. + llvm::SmallVector expected; + int64_t next = 0; + for (int64_t d = 0; d < outRank; ++d) + expected.push_back(inserted[d] ? 1 : inTy.getShape()[next++]); + if (outTy.getShape() != ArrayRef(expected)) + return rewriter.notifyMatchFailure(op, "result shape disagrees with axes"); + if (inTy.getRank() == 0) { + // A scalar tensor expands through a splat-free reshape. + rewriter.replaceOpWithNewOp( + op, outTy, op->getOperand(0), + rewriter.create( + op->getLoc(), rewriter.getIndexTensorAttr(expected))); + return success(); + } + // Reassociation: each input dim absorbs the unit dims inserted before it + // (leading unit dims join the first input dim; trailing ones the last). + llvm::SmallVector reassociation(inTy.getRank()); + int64_t group = 0; + llvm::SmallVector pending; + for (int64_t d = 0; d < outRank; ++d) { + if (inserted[d]) { pending.push_back(d); continue; } + reassociation[group].append(pending.begin(), pending.end()); + pending.clear(); + reassociation[group].push_back(d); + ++group; + } + reassociation.back().append(pending.begin(), pending.end()); + rewriter.replaceOpWithNewOp(op, outTy, op->getOperand(0), reassociation); + return success(); + } +}; + +struct BroadcastLowering : public RewritePattern { + BroadcastLowering(MLIRContext *ctx) + : RewritePattern("tessera.broadcast", /*benefit=*/1, ctx) {} + + LogicalResult matchAndRewrite(Operation *op, + PatternRewriter &rewriter) const override { + if (op->getNumOperands() != 1 || op->getNumResults() != 1) + return failure(); + auto inTy = dyn_cast(op->getOperand(0).getType()); + auto outTy = dyn_cast(op->getResult(0).getType()); + if (!inTy || !outTy || !inTy.hasStaticShape() || !outTy.hasStaticShape() || + inTy.getElementType() != outTy.getElementType()) + return rewriter.notifyMatchFailure(op, "static ranked tensors required"); + const int64_t rank = outTy.getRank(); + if (inTy.getRank() != rank) + return rewriter.notifyMatchFailure(op, "same-rank broadcast required (unsqueeze first)"); + // Input index for output dim d: d when the extents agree, 0 when the + // input extent is 1 (a stride-0 read); anything else is not a broadcast. + llvm::SmallVector inExprs; + for (int64_t d = 0; d < rank; ++d) { + if (inTy.getShape()[d] == outTy.getShape()[d]) + inExprs.push_back(rewriter.getAffineDimExpr(d)); + else if (inTy.getShape()[d] == 1) + inExprs.push_back(rewriter.getAffineConstantExpr(0)); + else + return rewriter.notifyMatchFailure(op, "extent neither equal nor 1"); + } + Location loc = op->getLoc(); + Value init = rewriter.create(loc, outTy.getShape(), outTy.getElementType()); + llvm::SmallVector maps{ + AffineMap::get(rank, 0, inExprs, rewriter.getContext()), + rewriter.getMultiDimIdentityMap(rank)}; + llvm::SmallVector iterators(rank, utils::IteratorType::parallel); + auto generic = rewriter.create( + loc, TypeRange{outTy}, ValueRange{op->getOperand(0)}, ValueRange{init}, maps, iterators, + [](OpBuilder &b, Location l, ValueRange args) { + b.create(l, args[0]); + }); + rewriter.replaceOp(op, generic.getResult(0)); + return success(); + } +}; + struct ReduceBackwardLowering : public RewritePattern { ReduceBackwardLowering(MLIRContext *ctx) : RewritePattern("tessera.reduce_backward", /*benefit=*/1, ctx) {} @@ -3073,6 +3184,8 @@ class TesseraToLinalgPass patterns.add(ctx); patterns.add(ctx); patterns.add(ctx); + patterns.add(ctx); + patterns.add(ctx); patterns.add(ctx); patterns.add(ctx); patterns.add(ctx); diff --git a/tests/tessera-ir/lit.cfg.py b/tests/tessera-ir/lit.cfg.py index 8e6e355fc..2f65fd8f6 100644 --- a/tests/tessera-ir/lit.cfg.py +++ b/tests/tessera-ir/lit.cfg.py @@ -129,6 +129,11 @@ def _opt_help_contains(needle: str) -> bool: if _opt_help_contains("tessera-lower-to-apple_gpu"): config.available_features.add("tessera-apple-backend") +# The EBM solver dialect + its lowering and the row-program emitter register in +# tessera-opt only when TESSERA_BUILD_EBM_BACKEND=ON (2026-09-16). +if _opt_help_contains("tessera-ebm-lower-langevin") and _opt_help_contains("tessera-row-program-to-gpu"): + config.available_features.add("tessera-ebm") + # The x86 executable pass is always registered so it can fail closed with a # useful rebuild diagnostic. Help-text probing therefore cannot distinguish a # build that actually registered TesseraX86Dialect. Probe the dialect directly diff --git a/tests/tessera-ir/phase2_autodiff/autodiff_paired_sub.mlir b/tests/tessera-ir/phase2_autodiff/autodiff_paired_sub.mlir new file mode 100644 index 000000000..412d4403c --- /dev/null +++ b/tests/tessera-ir/phase2_autodiff/autodiff_paired_sub.mlir @@ -0,0 +1,24 @@ +// RUN: tessera-opt --tessera-autodiff-paired %s | FileCheck %s +// +// `tessera.sub` carried only a tangent rule until 2026-09-16, so any energy +// or loss written with a subtraction stopped reverse-mode at +// AUTODIFF_OP_NOT_DIFFERENTIABLE. The adjoint passes the cotangent to lhs and +// negates it (0 - dy on a static zero) for rhs. + +module { + func.func @quadratic(%y: tensor<4x8xf32>, %x: tensor<4x8xf32>) -> tensor<4xf32> + attributes {tessera.autodiff = "reverse"} { + %d = "tessera.sub"(%x, %y) : (tensor<4x8xf32>, tensor<4x8xf32>) -> tensor<4x8xf32> + %sq = "tessera.mul"(%d, %d) : (tensor<4x8xf32>, tensor<4x8xf32>) -> tensor<4x8xf32> + %s = "tessera.reduce"(%sq) {axis = 1 : i64, kind = "sum"} : (tensor<4x8xf32>) -> tensor<4xf32> + %half = arith.constant dense<5.000000e-01> : tensor<4xf32> + %e = "tessera.mul"(%s, %half) : (tensor<4xf32>, tensor<4xf32>) -> tensor<4xf32> + return %e : tensor<4xf32> + } + // CHECK-LABEL: func.func @quadratic__bwd + // CHECK-SAME: (%[[Y:.*]]: tensor<4x8xf32>, %[[X:.*]]: tensor<4x8xf32>, %[[COT:.*]]: tensor<4xf32>) -> (tensor<4x8xf32>, tensor<4x8xf32>) + // CHECK-NOT: tessera.custom_adjoint_call + // CHECK: %[[ZERO:.*]] = arith.constant dense<0.000000e+00> : tensor<4x8xf32> + // CHECK: %[[DY:.*]] = tessera.sub %[[ZERO]], %[[DX:.*]] : (tensor<4x8xf32>, tensor<4x8xf32>) -> tensor<4x8xf32> + // CHECK: return %[[DY]], %[[DX]] +} diff --git a/tests/tessera-ir/phase2_autodiff/row_program_to_gpu_langevin.mlir b/tests/tessera-ir/phase2_autodiff/row_program_to_gpu_langevin.mlir new file mode 100644 index 000000000..223cb2a97 --- /dev/null +++ b/tests/tessera-ir/phase2_autodiff/row_program_to_gpu_langevin.mlir @@ -0,0 +1,62 @@ +// RUN: tessera-opt --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=backend=nvidia entry=loop" %s | FileCheck %s +// RUN: not tessera-opt --tessera-autodiff-paired --tessera-ebm-canonicalize --tessera-ebm-lower-langevin --tessera-to-linalg --convert-elementwise-to-linalg --canonicalize --cse "--tessera-row-program-to-gpu=backend=nvidia entry=loop" %s 2>&1 | FileCheck %s --check-prefix=CALL +// REQUIRES: tessera-ebm +// +// The EBM Langevin loop through the compiler alone (2026-09-16): the paired +// autodiff pass derives the gradient of the quadratic energy, the EBM lowering +// emits the step with Philox-4x32-10 / Box-Muller noise, tessera-to-linalg +// and inlining leave a [rows, features] row program, and the row-program +// emitter folds it into ONE cooperative gpu.func: one block per row, one lane +// per feature (8 features -> 8 lanes), the three-step scf.for carrying the +// state and both key words in registers, guarded loads and stores over +// !llvm.ptr<1> arguments, and no linalg/tensor op left. The second RUN skips +// --inline: the emitter refuses a call with a diagnostic instead of guessing. + +module { + func.func @energy(%y: tensor<4x8xf32>, %x: tensor<4x8xf32>) -> tensor<4xf32> + attributes {tessera.autodiff = "reverse"} { + %d = "tessera.sub"(%x, %y) : (tensor<4x8xf32>, tensor<4x8xf32>) -> tensor<4x8xf32> + %sq = "tessera.mul"(%d, %d) : (tensor<4x8xf32>, tensor<4x8xf32>) -> tensor<4x8xf32> + %s = "tessera.reduce"(%sq) {axis = 1 : i64, kind = "sum"} : (tensor<4x8xf32>) -> tensor<4xf32> + %half = arith.constant dense<5.000000e-01> : tensor<4xf32> + %e = "tessera.mul"(%s, %half) : (tensor<4xf32>, tensor<4xf32>) -> tensor<4xf32> + return %e : tensor<4xf32> + } + func.func @loop(%y0: tensor<4x8xf32>, %x: tensor<4x8xf32>, %key0: tensor<2xi64>) -> (tensor<4x8xf32>, tensor<2xi64>) { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %steps = arith.constant 3 : index + %r:2 = scf.for %t = %c0 to %steps step %c1 iter_args(%y = %y0, %key = %key0) -> (tensor<4x8xf32>, tensor<2xi64>) { + %n:2 = "tessera_ebm.langevin_step"(%y, %key, %x) + { energy_fn = @energy, eta = 0.1 : f64, temperature = 0.7 : f64, manifold = "euclidean" } + : (tensor<4x8xf32>, tensor<2xi64>, tensor<4x8xf32>) -> (tensor<4x8xf32>, tensor<2xi64>) + scf.yield %n#0, %n#1 : tensor<4x8xf32>, tensor<2xi64> + } + return %r#0, %r#1 : tensor<4x8xf32>, tensor<2xi64> + } +} + +// CHECK: module attributes { +// CHECK-SAME: tessera.row_program.backend = "nvidia" +// CHECK-SAME: tessera.row_program.entry = "loop" +// CHECK-SAME: tessera.row_program.features = 8 +// CHECK-SAME: tessera.row_program.rows = 4 +// CHECK: gpu.module @native_row +// CHECK: llvm.mlir.global private @row_reduction +// CHECK: gpu.func @row_program(%{{.*}}: !llvm.ptr<1>, %{{.*}}: !llvm.ptr<1>, %{{.*}}: !llvm.ptr<1>, %{{.*}}: !llvm.ptr<1>, %{{.*}}: !llvm.ptr<1>, %{{.*}}: index) kernel attributes {known_block_size = array} +// CHECK: gpu.block_id +// CHECK: gpu.thread_id +// CHECK: scf.if +// CHECK: llvm.load +// CHECK: scf.for {{.*}} iter_args({{.*}}) -> (f32, i64, i64) +// CHECK: arith.mului_extended +// CHECK: math.log +// CHECK: math.cos +// CHECK: scf.yield {{.*}} : f32, i64, i64 +// CHECK: llvm.store +// CHECK: gpu.return +// CHECK-NOT: linalg. +// CHECK-NOT: tensor. +// CHECK-NOT: tessera_ebm. + +// CALL: inline every call before lowering diff --git a/tests/tessera-ir/phase_f5/row_program_to_gpu_reduce.mlir b/tests/tessera-ir/phase_f5/row_program_to_gpu_reduce.mlir new file mode 100644 index 000000000..0ddac6ea3 --- /dev/null +++ b/tests/tessera-ir/phase_f5/row_program_to_gpu_reduce.mlir @@ -0,0 +1,82 @@ +// RUN: tessera-opt "--tessera-row-program-to-gpu=backend=rocm entry=row_normalize" %s | FileCheck %s +// RUN: tessera-opt "--tessera-row-program-to-gpu=backend=nvidia entry=row_normalize" %s | FileCheck %s --check-prefix=NV +// RUN: not tessera-opt "--tessera-row-program-to-gpu=backend=rocm entry=missing" %s 2>&1 | FileCheck %s --check-prefix=MISSING +// RUN: not tessera-opt "--tessera-row-program-to-gpu=backend=rocm entry=too_wide" %s 2>&1 | FileCheck %s --check-prefix=WIDE +// +// The row-program emitter's reduction path (2026-09-16): a feature-axis +// linalg.reduce becomes an ordered shared-memory fold — every lane writes its +// element into @row_reduction (address space 3), a barrier, the block leader +// folds the combiner over the lanes in index order (the declared reduction +// order, matched by a sequential f32 fold on the host), a barrier, the result +// is broadcast through the same array, a barrier. 6 features -> 8 lanes with +// the two spare lanes masked. Programs outside the envelope fail closed. + +#map = affine_map<(d0, d1) -> (d0, d1)> +#row = affine_map<(d0, d1) -> (d0, 0)> +module { + func.func @row_normalize(%x: tensor<3x6xf32>) -> tensor<3x6xf32> { + %zero = arith.constant 0.000000e+00 : f32 + %e = tensor.empty() : tensor<3x6xf32> + %sq = linalg.generic {indexing_maps = [#map, #map, #map], iterator_types = ["parallel", "parallel"]} + ins(%x, %x : tensor<3x6xf32>, tensor<3x6xf32>) outs(%e : tensor<3x6xf32>) { + ^bb0(%a: f32, %b: f32, %o: f32): + %m = arith.mulf %a, %b : f32 + linalg.yield %m : f32 + } -> tensor<3x6xf32> + %re = tensor.empty() : tensor<3xf32> + %init = linalg.fill ins(%zero : f32) outs(%re : tensor<3xf32>) -> tensor<3xf32> + %s = linalg.reduce ins(%sq : tensor<3x6xf32>) outs(%init : tensor<3xf32>) dimensions = [1] + (%in: f32, %acc: f32) { + %add = arith.addf %in, %acc : f32 + linalg.yield %add : f32 + } + %n = linalg.generic {indexing_maps = [affine_map<(d0) -> (d0)>, affine_map<(d0) -> (d0)>], iterator_types = ["parallel"]} + ins(%s : tensor<3xf32>) outs(%re : tensor<3xf32>) { + ^bb0(%v: f32, %o: f32): + %sq2 = math.sqrt %v : f32 + linalg.yield %sq2 : f32 + } -> tensor<3xf32> + %nx = tensor.expand_shape %n [[0, 1]] output_shape [3, 1] : tensor<3xf32> into tensor<3x1xf32> + %y = linalg.generic {indexing_maps = [#map, #row, #map], iterator_types = ["parallel", "parallel"]} + ins(%x, %nx : tensor<3x6xf32>, tensor<3x1xf32>) outs(%e : tensor<3x6xf32>) { + ^bb0(%a: f32, %d: f32, %o: f32): + %q = arith.divf %a, %d : f32 + linalg.yield %q : f32 + } -> tensor<3x6xf32> + return %y : tensor<3x6xf32> + } + func.func @too_wide(%x: tensor<2x2048xf32>) -> tensor<2x2048xf32> { + return %x : tensor<2x2048xf32> + } +} + +// CHECK: tessera.row_program.features = 6 +// CHECK: tessera.row_program.rows = 3 +// CHECK: llvm.mlir.global private @row_reduction() {addr_space = 3 : i32{{.*}} : !llvm.array<8 x f32> +// CHECK: gpu.func @row_program(%{{.*}}: !llvm.ptr<1>, %{{.*}}: !llvm.ptr<1>, %{{.*}}: index) kernel attributes {known_block_size = array} +// CHECK: llvm.mlir.addressof @row_reduction +// CHECK: arith.mulf +// CHECK: llvm.store {{.*}} : f32, !llvm.ptr<3> +// CHECK: gpu.barrier +// CHECK: scf.if +// CHECK: scf.for {{.*}} to %c6 +// CHECK: llvm.load {{.*}} : !llvm.ptr<3> -> f32 +// CHECK: arith.addf +// CHECK: llvm.store {{.*}} : f32, !llvm.ptr<3> +// CHECK: gpu.barrier +// CHECK: llvm.load {{.*}} : !llvm.ptr<3> -> f32 +// CHECK: gpu.barrier +// CHECK: llvm.intr.sqrt +// CHECK: arith.divf +// CHECK: llvm.store {{.*}} : f32, !llvm.ptr<1> +// CHECK: gpu.return +// CHECK-NOT: linalg. +// CHECK-NOT: tensor. + +// NV: llvm.func @__nv_fsqrt_rn(f32) -> f32 +// NV: gpu.func @row_program( +// NV: llvm.call @__nv_fsqrt_rn({{.*}}) : (f32) -> f32 +// NV-NOT: math.sqrt + +// MISSING: entry +// WIDE: 1024 diff --git a/tests/unit/test_autodiff_ledger.py b/tests/unit/test_autodiff_ledger.py index 40abc1dfb..df1f3133f 100644 --- a/tests/unit/test_autodiff_ledger.py +++ b/tests/unit/test_autodiff_ledger.py @@ -42,7 +42,10 @@ def test_native_and_placeholder_adjoints_are_disjoint_and_grounded() -> None: # tanh/sigmoid's W5 closed forms, comparison-backed ReLU, and the shared # normalization-statistics formulas. Nothing else. assert native == { - "add", "broadcast", "expand", "flatten", "gelu", "huber_loss", + # 2026-09-16 (EBM quadratic energy): `sub` gained its adjoint — + # dy to lhs, (0 - dy) to rhs — so an energy or loss written with + # a subtraction no longer stops reverse-mode. + "add", "sub", "broadcast", "expand", "flatten", "gelu", "huber_loss", "mae_loss", "mse_loss", "mul", "matmul", "permute", "reshape", "squeeze", "transpose", "unsqueeze", "view", "reduce", "sgd", "silu", "smooth_l1_loss", diff --git a/tests/unit/test_ebm_native_langevin.py b/tests/unit/test_ebm_native_langevin.py new file mode 100644 index 000000000..8027df4c6 --- /dev/null +++ b/tests/unit/test_ebm_native_langevin.py @@ -0,0 +1,89 @@ +"""The traceable quadratic energy loop through the MLIR/LLVM backbone +(W4-PRODUCT-1 / AD-SOLVER-IFT-1 acceptance, 2026-09-16). + +Acceptance from the GA/EBM review: native forward and gradient agree with +independent formulas, fixed-key samples agree with the declared policy, and +the complete loop executes without per-step host gradient transfers. Here the +gradient is the compiler's (paired autodiff on the Graph IR energy), the noise +is Philox drawn inside the compiled loop, and one JIT invocation runs all K +steps. Skips only when libtessera_jit lacks the EBM lane. +""" +from __future__ import annotations + +import numpy as np +import pytest + +from tessera import _jit_boundary as jb +from tessera.ebm import native_langevin as nl + +pytestmark = pytest.mark.skipif( + not nl.has_native_langevin(), + reason="libtessera_jit built without the EBM lane (TESSERA_BUILD_EBM_BACKEND=ON)", +) + + +def _pair(shape, seed): + rng = np.random.default_rng(seed) + return rng.standard_normal(shape).astype(np.float32), rng.standard_normal(shape).astype(np.float32) + + +def test_energy_matches_the_independent_formula(): + y, x = _pair((5, 7), 1) + np.testing.assert_allclose(nl.native_quadratic_energy(y, x), 0.5 * np.sum((x - y) ** 2, axis=1), rtol=1e-5, atol=1e-6) + + +def test_gradient_step_matches_the_independent_formula(): + """Temperature 0: y1 = y0 - eta * (y0 - x), with the gradient derived by + the compiler from the energy, not written by hand.""" + y, x = _pair((6, 4), 2) + out, key = nl.native_langevin_loop(y, x, [7, 9], eta=0.25, temperature=0.0, steps=1) + np.testing.assert_allclose(out, y - 0.25 * (y - x), rtol=1e-6, atol=1e-6) + assert list(key) == [7, 10] + + +def test_descent_converges_to_the_context(): + y, x = _pair((3, 8), 3) + out, _ = nl.native_langevin_loop(y, x, [1, 1], eta=0.5, temperature=0.0, steps=40) + np.testing.assert_allclose(out, x, atol=1e-5) + + +@pytest.mark.parametrize("shape,steps", [((4, 8), 1), ((4, 8), 5), ((3, 5), 12)]) +def test_fixed_key_samples_match_the_declared_policy(shape, steps): + """The whole loop runs in one compiled call and its samples are bit-exact + with the numpy mirror of the declared Philox / Box-Muller policy.""" + y, x = _pair(shape, 4 + steps) + before = jb.invocation_count() + out, key = nl.native_langevin_loop(y, x, [0x1234ABCD9876, 42], eta=0.1, temperature=0.7, steps=steps) + assert jb.invocation_count() == before + 1 # one native call for all K steps + expect, expect_key = nl.reference_langevin_loop(y, x, [0x1234ABCD9876, 42], eta=0.1, temperature=0.7, steps=steps) + assert list(key) == list(expect_key) + np.testing.assert_allclose(out, expect, rtol=1e-5, atol=1e-5) + assert not np.allclose(out, y - 0.1 * (y - x)) # noise actually entered + + +def test_different_keys_give_different_samples_and_same_key_repeats(): + y, x = _pair((4, 8), 9) + a, _ = nl.native_langevin_loop(y, x, [1, 0], eta=0.1, temperature=1.0, steps=2) + b, _ = nl.native_langevin_loop(y, x, [2, 0], eta=0.1, temperature=1.0, steps=2) + c, _ = nl.native_langevin_loop(y, x, [1, 0], eta=0.1, temperature=1.0, steps=2) + assert not np.allclose(a, b) and np.array_equal(a, c) + + +def test_runtime_launch_reports_the_native_cpu_row(): + from tessera import runtime as rt + y, x = _pair((4, 8), 11) + result = rt.launch(nl.package_ebm_langevin_cpu((4, 8), eta=0.1, temperature=0.3, steps=3), (y, x, [5, 6])) + assert result["ok"] and result["execution_kind"] == "native_cpu" + assert result["compiler_path"] == "cpu_ebm_langevin_llvm_jit" + out, key = result["output"] + expect, expect_key = nl.reference_langevin_loop(y, x, [5, 6], eta=0.1, temperature=0.3, steps=3) + np.testing.assert_allclose(np.asarray(out), expect, rtol=1e-5, atol=1e-5) + assert list(np.asarray(key)) == list(expect_key) + + +@pytest.mark.parametrize("kwargs", [dict(eta=0.0, temperature=1.0, steps=1), dict(eta=0.1, temperature=-1.0, steps=1), + dict(eta=0.1, temperature=1.0, steps=0)]) +def test_out_of_envelope_requests_are_refused(kwargs): + y, x = _pair((2, 2), 0) + with pytest.raises(ValueError): + nl.native_langevin_loop(y, x, [0, 0], **kwargs) diff --git a/tests/unit/test_ebm_native_langevin_gpu.py b/tests/unit/test_ebm_native_langevin_gpu.py new file mode 100644 index 000000000..196586257 --- /dev/null +++ b/tests/unit/test_ebm_native_langevin_gpu.py @@ -0,0 +1,194 @@ +"""The EBM Langevin loop as one cooperative device kernel (2026-09-16). + +Host-free half (any host with tessera-opt built with the EBM backend): one +tessera-opt invocation runs paired autodiff, the EBM lowering, tessera-to- +linalg, inlining and the row-program emitter; the result is a single +gpu.func with guarded per-lane loads, the K-step scf.for carrying the state +and both key words in registers, Philox inside the loop, and no tensor or +linalg op left. Device half (owning ROCm / sm_120 host): the loop is one +launch, bit-exact with the declared numpy policy for 1/5/12 steps, and +`runtime.launch` reports the new rows. +""" +from __future__ import annotations + +import os +import re + +import numpy as np +import pytest + +from tessera.compiler.llvm_tools import llvm_bin_dir +from tessera.compiler.scheduled_matmul import find_tessera_opt +from tessera.ebm import native_langevin as nl +from tests._support.environment import (native_storage_target, nvidia_gpu_is_plausibly_present, + rocm_gpu_is_plausibly_present) + + +def _compiler(): + tool = find_tessera_opt() + if tool is None: + pytest.skip("tessera-opt required") + import subprocess + help_text = subprocess.run([str(tool), "--help"], capture_output=True, text=True).stdout + if "--tessera-ebm-lower-langevin" not in help_text or "--tessera-row-program-to-gpu" not in help_text: + pytest.skip("tessera-opt built without the EBM backend or the row-program emitter") + return tool + + +def _kernel(shape=(4, 8), steps=3, temperature=0.7, backend="nvidia"): + tool = _compiler() + source, specs = nl.langevin_device_source(shape, eta=0.1, temperature=temperature, steps=steps, + backend=backend, compiler=tool) + body = source.split("gpu.func @row_program(", 1)[1] + return source, specs, body + + +def test_single_driver_chain_emits_one_cooperative_kernel(): + source, specs, body = _kernel() + assert source.count("gpu.func ") == 1 and "kernel attributes {known_block_size = array" in source + assert "tessera_ebm." not in body and "linalg." not in body and "tensor." not in body and "func.call" not in body + assert "gpu.block_id" in body and "gpu.thread_id" in body + # Guarded lane loads of y0 and x, the two key words, then ONE loop carrying + # (state f32, key i64, key i64), then guarded stores of y and the key. + assert body.count("llvm.load") == 4 and body.count("llvm.store") == 3 + loops = re.findall(r"scf\.for [^\n]*iter_args\(([^)]*)\) -> \(([^)]*)\)", body) + assert len(loops) == 1 and loops[0][1] == "f32, i64, i64" + # Philox-4x32-10 inside the loop body (CSE folds the constant-keyed round). + assert body.count("arith.mului_extended") >= 18 and "math.log" in body and "math.cos" in body + assert [s.name for s in specs] == ["y0", "x", "key", "y", "next_key", "scratch"] + assert '\\22grid\\22:[4,1,1]' in source and '\\22block\\22:[8,1,1]' in source + # The quadratic gradient is elementwise, so this energy needs no cross-lane + # reduction: the emitter places no barrier at all. + assert "gpu.barrier" not in body + + +def test_zero_temperature_kernel_has_no_noise_and_no_barriers(): + _, _, body = _kernel(temperature=0.0) + assert "math.log" not in body and "arith.mului_extended" not in body and "gpu.barrier" not in body + + +def test_kernel_replays_through_the_arena_pipeline_for_both_backends(): + from tessera.compiler.native_gpu_storage import replay_arena_ir + tool = _compiler() + for backend in ("nvidia", "rocm"): + source, _, _ = _kernel(shape=(3, 5), backend=backend) + arena = replay_arena_ir(tool, source) + assert "__tessera_shared_bytes_" in arena and "gpu.func @row_program" in arena + assert "known_block_size = array" in arena # 5 features -> 8 lanes + + +def test_row_reduction_program_emits_the_ordered_shared_fold(): + """The quadratic gradient is elementwise, so the Langevin kernel carries no + reduction; the emitter's reduction path is proven on a row-normalization + program (feature-axis linalg.reduce broadcast back over the lanes).""" + from tessera.compiler.native_row_program import row_normalize_module, row_program_kernel + tool = _compiler() + kernel, lanes = row_program_kernel(row_normalize_module(3, 6), entry="row_normalize", backend="rocm", compiler=tool) + body = kernel.split("gpu.func @row_program(", 1)[1] + assert lanes == 8 and body.count("gpu.barrier") == 3 + assert "llvm.mlir.addressof @row_reduction" in body and body.count("!llvm.ptr<3>") >= 4 + assert "linalg." not in body and "tensor." not in body and "llvm.intr.sqrt" in body and "math.sqrt" not in body + # NVIDIA: convert-gpu-to-nvvm outlaws the LLVM math intrinsics, so the + # emitter calls libdevice's rounding-explicit __nv_fsqrt_rn instead. + nv, _ = row_program_kernel(row_normalize_module(3, 6), entry="row_normalize", backend="nvidia", compiler=tool) + nv_body = nv.split("gpu.func @row_program(", 1)[1] # the provenance attribute keeps the source text + assert "llvm.call @__nv_fsqrt_rn" in nv_body and "math.sqrt" not in nv_body and "llvm.intr.sqrt" not in nv_body + + +@pytest.mark.parametrize("shape,message", [((4, 2000), "1024 features"), ((4,), "rows, features")]) +def test_out_of_envelope_shapes_are_refused(shape, message): + tool = _compiler() + with pytest.raises(ValueError, match=message): + nl.langevin_device_source(shape, eta=0.1, temperature=0.5, steps=1, backend="nvidia", compiler=tool) + + +# --- device half ----------------------------------------------------------- + +def _device_lane(): + target = native_storage_target() + if target is None: + pytest.skip("no CUDA or ROCm toolchain on this host") + backend, chip = target + present = nvidia_gpu_is_plausibly_present() if backend == "nvidia" else rocm_gpu_is_plausibly_present() + if not present: + pytest.skip(f"no {backend} device on this host") + if backend == "rocm" and chip == "gfx1201" and os.environ.get("TESSERA_GFX1201_DEVICE_PROOF") != "1": + pytest.skip("gfx1201 device proof requires TESSERA_GFX1201_DEVICE_PROOF=1") + tool, llvm = _compiler(), llvm_bin_dir() + if llvm is None: + pytest.skip("matched LLVM tools required") + return backend, chip, tool, llvm + + +def _pair(shape, seed): + rng = np.random.default_rng(seed) + return rng.standard_normal(shape).astype(np.float32), rng.standard_normal(shape).astype(np.float32) + + +@pytest.mark.parametrize("shape,steps", [((4, 8), 1), ((6, 8), 5), ((3, 5), 12), ((9, 100), 4)]) +def test_device_loop_is_bit_exact_with_the_declared_policy(shape, steps): + backend, chip, tool, llvm = _device_lane() + y, x = _pair(shape, 40 + steps) + out, key = nl.native_langevin_loop_device(y, x, [0x1234ABCD9876, 42], eta=0.1, temperature=0.7, steps=steps, + backend=backend, chip=chip, compiler=tool, llvm_bin=llvm) + expect, expect_key = nl.reference_langevin_loop(y, x, [0x1234ABCD9876, 42], eta=0.1, temperature=0.7, steps=steps) + assert list(key) == list(expect_key) + np.testing.assert_allclose(out, expect, rtol=1e-5, atol=1e-5) + assert not np.allclose(out, y - 0.1 * (y - x)) + + +def test_device_descent_matches_the_independent_formula(): + backend, chip, tool, llvm = _device_lane() + y, x = _pair((5, 7), 3) + out, _ = nl.native_langevin_loop_device(y, x, [1, 1], eta=0.25, temperature=0.0, steps=1, + backend=backend, chip=chip, compiler=tool, llvm_bin=llvm) + np.testing.assert_allclose(out, y - 0.25 * (y - x), rtol=1e-6, atol=1e-6) + + +def test_device_loop_is_one_launch_and_launch_row_reports_native_gpu(): + from tessera import runtime as rt + backend, chip, tool, llvm = _device_lane() + y, x = _pair((4, 8), 11) + target = "rocm" if backend == "rocm" else "nvidia_sm120" + artifact = nl.package_ebm_langevin_native((4, 8), eta=0.1, temperature=0.3, steps=6, target=target) + result = rt.launch(artifact, (y, x, [5, 6])) + assert result["ok"] and result["execution_kind"] == "native_gpu" + assert result["compiler_path"] == f"{'rocm' if backend == 'rocm' else 'nvidia'}_ebm_langevin_native_compiled" + out, key = result["output"] + expect, expect_key = nl.reference_langevin_loop(y, x, [5, 6], eta=0.1, temperature=0.3, steps=6) + np.testing.assert_allclose(np.asarray(out), expect, rtol=1e-5, atol=1e-5) + assert list(np.asarray(key)) == list(expect_key) + # The whole 6-step loop was one kernel launch: the packaged program's + # binding launched once for this call (per-call launches are counted by + # the bound package). + program = nl.ebm_langevin_program((4, 8), eta=0.1, temperature=0.3, steps=6, backend=backend, chip=chip, + compiler=tool, llvm_bin=llvm) + before = getattr(program.bound, "launch_count", None) + program.run(y, x, [5, 6]) + if before is not None: + assert program.bound.launch_count == before + 1 + + +def _sequential_row_normalize(x): + """The declared reduction order: lanes folded in index order in f32.""" + out = np.empty_like(x) + for r in range(x.shape[0]): + acc = np.float32(0.0) + for v in x[r]: + acc = np.float32(np.float32(v * v) + acc) + out[r] = x[r] / np.float32(np.sqrt(acc)) + return out + + +@pytest.mark.parametrize("shape", [(3, 6), (5, 8), (2, 100), (7, 1024)]) +def test_row_reduction_program_is_bit_exact_with_the_declared_order_on_device(shape): + from tessera.compiler.native_gpu_tensor import TensorSpec + from tessera.compiler.native_row_program import row_normalize_module, row_program_device + backend, chip, tool, llvm = _device_lane() + x = np.random.default_rng(shape[1]).standard_normal(shape).astype(np.float32) + program = row_program_device(row_normalize_module(*shape), entry="row_normalize", rows=shape[0], + specs=(TensorSpec("x", "fp32", shape, False), TensorSpec("y", "fp32", shape, True)), + backend=backend, chip=chip, compiler=tool, llvm_bin=llvm, name="row normalize") + out = program.run(x) + np.testing.assert_array_equal(np.asarray(out), _sequential_row_normalize(x)) + np.testing.assert_allclose(np.linalg.norm(np.asarray(out), axis=1), 1.0, rtol=1e-5, atol=1e-5) diff --git a/tools/tessera-jit/CMakeLists.txt b/tools/tessera-jit/CMakeLists.txt index 1be884320..7c2195f5a 100644 --- a/tools/tessera-jit/CMakeLists.txt +++ b/tools/tessera-jit/CMakeLists.txt @@ -45,6 +45,13 @@ add_library(tessera_jit SHARED tessera_jit.cpp) # tessera->linalg, so a `tessera_clifford.geo_product` on `[..., dim]` tensors # executes through MLIR/LLVM here (W6.4, 2026-09-16) rather than through a # Python-emitted kernel. Absent the backend the symbol reports no support. +if(TESSERA_BUILD_EBM_BACKEND) + target_compile_definitions(tessera_jit PRIVATE TESSERA_JIT_HAVE_EBM=1) + target_include_directories(tessera_jit PRIVATE + ${CMAKE_SOURCE_DIR}/src/solvers/ebm/include + ${CMAKE_BINARY_DIR}/src/solvers/ebm) + target_link_libraries(tessera_jit PRIVATE TesseraEBM) +endif() if(TESSERA_BUILD_CLIFFORD_BACKEND) target_compile_definitions(tessera_jit PRIVATE TESSERA_JIT_HAVE_CLIFFORD=1) target_include_directories(tessera_jit PRIVATE diff --git a/tools/tessera-jit/tessera_jit.cpp b/tools/tessera-jit/tessera_jit.cpp index 0891d9444..625791451 100644 --- a/tools/tessera-jit/tessera_jit.cpp +++ b/tools/tessera-jit/tessera_jit.cpp @@ -105,6 +105,10 @@ #include "tessera/Clifford/CliffordDialect.h" #include "tessera/Clifford/CliffordPasses.h" #endif +#ifdef TESSERA_JIT_HAVE_EBM +#include "tessera/EBM/EBMDialect.h" +#include "tessera/EBM/EBMPasses.h" +#endif using namespace mlir; @@ -323,7 +327,16 @@ void maybeTrace(PassManager &pm) { // copy is correct for ANY producer, and for our identity-layout C-contiguous // boundary (§12.4) `memref.copy` lowers to a `memcpy` intrinsic — no runtime // symbol, negligible cost. +// +// Intra-module callers (2026-09-16): a compiler-derived gradient reached the +// lane as `call @E__bwd` from the EBM Langevin lowering, and every such call +// site must follow its callee's new signature -- the caller allocates the +// out-buffers (static memrefs), passes them trailing, and reads them back as +// the former results. A callee whose result is not a static memref is left +// untouched together with its callers. Before this, the rewrite silently left +// callers with the old arity and the verifier failed after bufferization. LogicalResult rewriteResultsToOutParams(ModuleOp module) { + llvm::StringMap> rewritten; for (auto fn : module.getOps()) { if (fn.getNumResults() == 0 || fn.isExternal()) continue; @@ -333,6 +346,18 @@ LogicalResult rewriteResultsToOutParams(ModuleOp module) { }); if (!allMemref) continue; + bool called = false; + module.walk([&](func::CallOp call) { + if (call.getCallee() == fn.getName()) called = true; + }); + if (called && !llvm::all_of(fn.getResultTypes(), [](Type t) { + return cast(t).hasStaticShape(); + })) { + fn.emitError("tessera_jit: a called function must return static memrefs " + "for the DPS rewrite to allocate its out-buffers at call sites"); + return failure(); + } + rewritten[fn.getName()] = SmallVector(fn.getResultTypes()); Block &entry = fn.getBody().front(); auto retOp = cast(entry.getTerminator()); @@ -351,6 +376,25 @@ LogicalResult rewriteResultsToOutParams(ModuleOp module) { fn.setType( FunctionType::get(fn.getContext(), entry.getArgumentTypes(), {})); } + // Follow every call site of a rewritten callee. + SmallVector calls; + module.walk([&](func::CallOp call) { + if (rewritten.count(call.getCallee())) calls.push_back(call); + }); + for (func::CallOp call : calls) { + OpBuilder b(call); + SmallVector operands(call.getOperands()); + SmallVector outs; + for (Type t : rewritten[call.getCallee()]) { + Value buffer = memref::AllocOp::create(b, call.getLoc(), cast(t)); + operands.push_back(buffer); + outs.push_back(buffer); + } + func::CallOp::create(b, call.getLoc(), call.getCallee(), TypeRange{}, operands); + for (auto [result, buffer] : llvm::zip(call.getResults(), outs)) + result.replaceAllUsesWith(buffer); + call.erase(); + } return success(); } @@ -685,6 +729,24 @@ LogicalResult buildAndRunPipeline(ModuleOp module) { // arith that one-shot bufferization consumes like any other tensor loop. pm1a.addPass(tessera::createCliffordGradeFusionPass()); pm1a.addPass(tessera::createCliffordExpandProductTablePass(/*expandRotorSandwich=*/true)); +#endif +#ifdef TESSERA_JIT_HAVE_EBM + // Energy-based models: the paired autodiff pass derives @E__bwd for every + // energy marked tessera.autodiff = "reverse" that has no paired partner yet + // (modules that arrive already paired -- the autodiff CPU rows -- carry + // `tessera.autodiff.paired` and must not have their __bwd regenerated); + // the EBM lowering then turns energy / inner_step / langevin_step into + // arith + linalg (Philox noise in a linalg.generic) so a whole sampling + // loop compiles as one function. + bool needsPairing = false; + module.walk([&](func::FuncOp fn) { + if (fn->hasAttr("tessera.autodiff") && !fn->hasAttr("tessera.autodiff.paired")) + needsPairing = true; + }); + if (needsPairing) + pm1a.addPass(tessera::createAutodiffPairedPass()); + pm1a.addPass(tessera::createEBMCanonicalizePass()); + pm1a.addPass(tessera::createEBMLowerLangevinPass()); #endif pm1a.nest().addPass(tessera::createTesseraToLinalgPass()); // Elementwise arith/math ops ON TENSORS (e.g. the paired autodiff pass's @@ -871,6 +933,14 @@ const char *tessera_jit_last_error(void) { return g_lastError.c_str(); } // Whether this library was built with the Clifford dialect lane (CMake option // TESSERA_BUILD_CLIFFORD_BACKEND). Callers needing geometric products check // this instead of compiling and reading a parse error. +int tessera_jit_has_ebm(void) { +#ifdef TESSERA_JIT_HAVE_EBM + return 1; +#else + return 0; +#endif +} + int tessera_jit_has_clifford(void) { #ifdef TESSERA_JIT_HAVE_CLIFFORD return 1; @@ -891,6 +961,9 @@ void *tessera_jit_compile(const char *mlir_text) { tessera::registerTesseraDialects(registry); #ifdef TESSERA_JIT_HAVE_CLIFFORD registry.insert(); +#endif +#ifdef TESSERA_JIT_HAVE_EBM + registry.insert(); #endif registry.insert hsaco serialization (no mlir-opt shell-out). @@ -417,6 +427,20 @@ int main(int argc, char **argv) { tessera::solver::registerTesseraLinalgSolverPipeline(); #endif +#ifdef TESSERA_HAVE_EBM + ::mlir::registerPass([]() { return tessera::createEBMCanonicalizePass(); }); + ::mlir::registerPass([]() { return tessera::createEBMFuseEnergyGradPass(); }); + ::mlir::registerPass([]() { return tessera::createEBMCheckpointInnerLoopPass(); }); + ::mlir::registerPass([]() { return tessera::createEBMPipelineCandidatesPass(); }); + ::mlir::registerPass([]() { return tessera::createEBMLowerLangevinPass(); }); +#endif +#ifdef TESSERA_HAVE_CLIFFORD + ::mlir::registerPass([]() { return tessera::createCliffordAnnotateAlgebraPass(); }); + ::mlir::registerPass([]() { return tessera::createCliffordExpandProductTablePass(); }); + ::mlir::registerPass([]() { return tessera::createCliffordGradeFusionPass(); }); + ::mlir::registerPass([]() { return tessera::createCliffordRotorSandwichFoldPass(); }); +#endif + #ifdef TESSERA_HAVE_SCALING_RESILIENCE mlir::tessera::sr::registerPasses(); #endif @@ -480,6 +504,9 @@ int main(int argc, char **argv) { // link surface), not the full conversion umbrella. mlir::registerConvertLinalgToParallelLoopsPass(); mlir::bufferization::registerBufferizationPasses(); + // Tensor-level arith on tensors becomes linalg.generic for the row-program + // emitter (the JIT runs the same pass in-process). + mlir::registerConvertElementwiseToLinalgPass(); mlir::registerGpuMapParallelLoopsPass(); mlir::registerGpuKernelOutliningPass(); mlir::registerConvertParallelLoopToGpuPass(); @@ -519,6 +546,15 @@ int main(int argc, char **argv) { // attributes and every loaded dialect's ConvertToLLVM interface must be // registered before MlirOptMain creates its context. mlir::registerAllExtensions(registry); + // The LLVM dialect *promises* a DialectInlinerInterface that + // registerAllExtensions does not provide. `--inline` builds its interface + // collection over every loaded dialect, and the row-program emitter loads + // llvm as a dependent dialect before the inliner runs, so without this + // registration an assertions-enabled MLIR aborts with "checking for an + // interface ... promised by dialect 'llvm' but never implemented" while an + // NDEBUG build silently proceeds (found on Tajasarus 2026-09-16; the + // Decision #19 standing lesson, third instance). + mlir::LLVM::registerInlinerInterface(registry); #ifdef TESSERA_HAVE_CORE_TESSERA_IR tessera::registerTesseraDialects(registry); @@ -544,6 +580,12 @@ int main(int argc, char **argv) { #ifdef TESSERA_HAVE_SOLVERS tessera::solver::registerTesseraLinalgSolverDialect(registry); #endif +#ifdef TESSERA_HAVE_EBM + registry.insert(); +#endif +#ifdef TESSERA_HAVE_CLIFFORD + registry.insert(); +#endif #ifdef TESSERA_HAVE_COLLECTIVES tessera::collective::registerCollectiveDialect(registry); #endif