From 4569a9920a9cb9ae04b1793bb06af0fe74365786 Mon Sep 17 00:00:00 2001 From: angst Date: Tue, 7 Jul 2026 13:31:56 -0600 Subject: [PATCH] D1/B1: bare-matmul op-kind + NVIDIA GEMM arbiter candidates (sm_120) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Unblock the D1 arbiter for plain GEMM — the candidate registry only had fused_region/attention/gated/pointwise, no bare matmul, so the emitted mma.sync GEMM lane had nowhere to plug in. * fusion_core: MatmulRegion (D = A @ B, 16-bit storage / f32 accumulate) + verify_synthesized_matmul (arbiter-only F4 oracle, dtype-rounded reference) + a _round_to_storage helper; re-exported through the fusion facade. * candidate.py: OP_MATMUL op-kind + its (verify, run_matmul) map entry + a run_matmul method on the _as_runner adapter. * runtime.py: two 2D GEMM execution helpers the candidates call — _nvidia_mma_gemm_2d (shipped libtessera_nvidia_gemm, row-major B) and _nvidia_ptx_gemm_2d (compiler-emitted ptx_emit via the launch bridge, col-major B), keyed by 16-bit dtype; + a bridge loader mirroring the shipped-GEMM one. * emit/nvidia_cuda.py: NvidiaMmaGemmShippedCandidate (Tier-3 hand-tuned) and NvidiaMmaGemmEmittedCandidate (Tier-2 emitted, aligned-only) registered under (nvidia, matmul); both F4-gated, f16 accuracy budget. Tier-priority picks the shipped lane by default (lead-safe, Decision #28); the E3 force hatch selects the emitted lane. So NVIDIA gains its Tier-3 hand-tuned GEMM candidate (previously only reachable via the jit nvidia_mma executor) next to the Tier-2 emitted lane. D2's measured loop (lets Tier-2 win where faster) is the follow-on. Live-proven on sm_120 (RTX 5070 Ti): both lanes verify + execute + match the dtype-rounded reference across bf16/f16 x 16x8x16/32x16x32/64x64x64; arbiter selects shipped by default, force selects emitted (test_nvidia_plugin.py). Co-Authored-By: Claude Opus 4.8 --- docs/audit/compiler/COMPILER_REFACTOR_PLAN.md | 22 +++- python/tessera/compiler/emit/candidate.py | 7 + python/tessera/compiler/emit/nvidia_cuda.py | 92 ++++++++++++- python/tessera/compiler/fusion.py | 4 + python/tessera/compiler/fusion_core.py | 65 ++++++++++ python/tessera/runtime.py | 121 ++++++++++++++++++ tests/unit/test_nvidia_plugin.py | 68 +++++++++- 7 files changed, 370 insertions(+), 9 deletions(-) diff --git a/docs/audit/compiler/COMPILER_REFACTOR_PLAN.md b/docs/audit/compiler/COMPILER_REFACTOR_PLAN.md index cb7829b98..ee1b25594 100644 --- a/docs/audit/compiler/COMPILER_REFACTOR_PLAN.md +++ b/docs/audit/compiler/COMPILER_REFACTOR_PLAN.md @@ -308,13 +308,21 @@ chains, small attention). Crown-jewel GEMM stays Tier 2/3. documented *instruction-encoding skeleton* (`emit_wgmma_matmul_ptx` — its own header says "NOT a complete assemblable kernel: needs smem matrix descriptors + TMA/cp.async + the full accumulator operand list"), so completing it is a real Hopper WGMMA-kernel - build (assemble-only here; execution needs Hopper) — **not** a bug fix; sm_100 - tcgen05; and wiring the emit lane as a **Tier-2 EMITTED** arbiter candidate (blocked - on a bare-matmul op-kind — the candidate registry has no matmul op today, only - fused_region/attention/gated/pointwise). Unlike C3, NVIDIA has no *fused* shipped - kernel to register as a Tier-3 `FusedRegion` candidate yet (the shipped kernel is a - pure GEMM served by the - jit `nvidia_mma` executor). + build (assemble-only here; execution needs Hopper) — **not** a bug fix; and sm_100 + tcgen05. + **B1 landed 2026-07-07 — the emit lane is now a first-class arbiter candidate.** + A bare-matmul op-kind (`candidate.OP_MATMUL` + `fusion_core.MatmulRegion` + + `verify_synthesized_matmul`) unblocks the D1 arbiter for plain GEMM. Two NVIDIA + matmul candidates register under `(nvidia, matmul)`: the **shipped** mma.sync GEMM + (`NvidiaMmaGemmShippedCandidate`, Tier-3 via `runtime._nvidia_mma_gemm_2d`) and the + **compiler-emitted** `ptx_emit` GEMM (`NvidiaMmaGemmEmittedCandidate`, Tier-2 via the + launch bridge, `runtime._nvidia_ptx_gemm_2d`). Both F4-gated by the universal oracle; + tier-priority picks the shipped lane by default (lead-safe, Decision #28), the E3 + `force` hatch selects the emitted lane. Live-proven on sm_120 (bf16/f16 × + 16x8x16/32x16x32/64x64x64). So NVIDIA now has its Tier-3 *hand-tuned* GEMM candidate + (the pure-GEMM shipped kernel, previously only reachable via the jit `nvidia_mma` + executor) alongside the Tier-2 emitted lane — D2's measured loop is the follow-on + that lets Tier-2 win where faster. - **C3 · ROCm generic synth → HIP** — **generic lane LANDED 2026-07-06**, `[MAC]` author → `[AMD]` proof. `emit/rocm_hip.py` is now a **full three-seam plugin** (parallel to x86): `RocmHipEmitter` turns a `FusedRegion` into HIP source (a diff --git a/python/tessera/compiler/emit/candidate.py b/python/tessera/compiler/emit/candidate.py index d0a37ddef..d13d47168 100644 --- a/python/tessera/compiler/emit/candidate.py +++ b/python/tessera/compiler/emit/candidate.py @@ -53,6 +53,9 @@ class Tier(IntEnum): OP_ATTENTION = "attention" OP_GATED_MATMUL = "gated_matmul" OP_POINTWISE = "pointwise" +#: A bare matmul ``D = A @ B`` (no fusion) — the op-kind for plain-GEMM candidates +#: (NVIDIA emitted mma.sync vs shipped GEMM). Keyed on ``fusion_core.MatmulRegion``. +OP_MATMUL = "matmul" #: op-kind → (verify function name, KernelRunner method) — the seam that lets the #: arbiter reuse the exact universal F4 oracle for each candidate. @@ -61,6 +64,7 @@ class Tier(IntEnum): OP_ATTENTION: ("verify_synthesized_attention", "run_fused_attention"), OP_GATED_MATMUL: ("verify_synthesized_gated", "run_gated_matmul_region"), OP_POINTWISE: ("verify_synthesized_pointwise", "run_pointwise_graph"), + OP_MATMUL: ("verify_synthesized_matmul", "run_matmul"), } @@ -180,6 +184,9 @@ def run_gated_matmul_region(self, region, *a, **k): def run_pointwise_graph(self, region, *a, **k): return self._route("run_pointwise_graph", region, a, k) + def run_matmul(self, region, *a, **k): + return self._route("run_matmul", region, a, k) + def _route(self, called, region, a, k): if called != method: raise NotImplementedError( diff --git a/python/tessera/compiler/emit/nvidia_cuda.py b/python/tessera/compiler/emit/nvidia_cuda.py index e8ebc673b..80683fa3d 100644 --- a/python/tessera/compiler/emit/nvidia_cuda.py +++ b/python/tessera/compiler/emit/nvidia_cuda.py @@ -46,6 +46,7 @@ from tessera.compiler.emit._fused_scalar_body import row_compute_body from tessera.compiler.emit.candidate import ( OP_FUSED_REGION, + OP_MATMUL, Candidate, Tier, register_candidate, @@ -61,7 +62,7 @@ register_emitter, register_runner, ) -from tessera.compiler.fusion_core import FusedRegion +from tessera.compiler.fusion_core import FusedRegion, MatmulRegion _TARGET = "nvidia" _LANG = "cuda" @@ -275,9 +276,98 @@ def run(self, region: Any, A: Any, B: Any, bias: Any = None, residual=residual) +# ── D1 matmul candidates (B1) — bare GEMM, Tier-2 emitted vs Tier-3 shipped ──── +# +# The arbiter enumerates these per (target="nvidia", op=matmul) and F4-gates each. +# Tier-priority (Decision #28) prefers the hand-tuned shipped lane by default; D2's +# measured loop lets the emitted lane win where it is faster + in accuracy budget. +# Both are 16-bit storage (bf16/f16) → f32 accumulate, so they declare the f16 +# budget the oracle honors. Off an NVIDIA GPU / without the built libs they decline +# to the reference and drop out of the enumeration. + +_GEMM_F16_ATOL = 5e-3 # 16-bit storage vs the f32 reference (Decision #28) +_GEMM_DTYPES = ("bfloat16", "float16") + + +def _aligned_2d(A: Any, B: Any) -> bool: + """A (M,K) @ B (K,N) with the emitted kernel's tile alignment (M%16,N%8,K%16).""" + import numpy as np + Aa, Ba = np.asarray(A), np.asarray(B) + if Aa.ndim != 2 or Ba.ndim != 2 or Aa.shape[1] != Ba.shape[0]: + return False + M, K = Aa.shape + _, N = Ba.shape + return M % 16 == 0 and N % 8 == 0 and K % 16 == 0 + + +class NvidiaMmaGemmShippedCandidate(Candidate): + """Tier-3 (hand-tuned): the shipped ``libtessera_nvidia_gemm`` mma.sync GEMM — + the crown-jewel lane, arbiter default until D2 measures otherwise. Serves any + (unaligned OK) bf16/f16 matmul; declines off an NVIDIA GPU.""" + + name = "nvidia_mma_gemm_shipped" + tier = Tier.HAND_TUNED + target = _TARGET + op = OP_MATMUL + accuracy_atol = _GEMM_F16_ATOL + + def available(self) -> bool: + try: + from tessera import runtime as rt + return rt._nvidia_mma_runtime_available() + except Exception: + return False + + def applies_to(self, region: Any) -> bool: + return isinstance(region, MatmulRegion) and region.dtype in _GEMM_DTYPES + + def run(self, region: Any, A: Any, B: Any, *a: Any, **k: Any) -> tuple[Any, str]: + try: + from tessera import runtime as rt + return rt._nvidia_mma_gemm_2d(A, B, region.dtype), "nvidia_mma_shipped" + except Exception: + return region.reference(A, B), "reference" + + +class NvidiaMmaGemmEmittedCandidate(Candidate): + """Tier-2 (emitted): the compiler-EMITTED ``ptx_emit`` mma.sync GEMM driven + through the launch bridge — the C2 emit lane as a first-class arbiter candidate. + Serves ALIGNED (M%16/N%8/K%16) bf16/f16 matmuls; declines (to the reference) for + ragged shapes or off an NVIDIA GPU / without the built bridge.""" + + name = "nvidia_mma_gemm_emitted" + tier = Tier.EMITTED + target = _TARGET + op = OP_MATMUL + accuracy_atol = _GEMM_F16_ATOL + + def available(self) -> bool: + try: + from tessera import runtime as rt + return (rt._load_nvidia_ptx_launch() is not None + and rt._nvidia_mma_runtime_available()) + except Exception: + return False + + def applies_to(self, region: Any) -> bool: + return isinstance(region, MatmulRegion) and region.dtype in _GEMM_DTYPES + + def run(self, region: Any, A: Any, B: Any, *a: Any, **k: Any) -> tuple[Any, str]: + if not _aligned_2d(A, B): # emitter is aligned-only (for now) + return region.reference(A, B), "reference" + try: + from tessera import runtime as rt + return rt._nvidia_ptx_gemm_2d(A, B, region.dtype), "nvidia_ptx_gemm" + except Exception: + return region.reference(A, B), "reference" + + # ── registration (import side effect, exactly like rocm_hip / x86_llvm) ──────── register_emitter(NvidiaCudaEmitter()) register_compiler(_TARGET, _nvidia_cuda_compile_fn) register_runner(NvidiaCudaRunner(), default=False) register_candidate(NvidiaGenericCudaCandidate()) +# Bare-GEMM lanes: hand-tuned shipped (Tier 3) + compiler-emitted (Tier 2). +register_candidate(NvidiaMmaGemmShippedCandidate()) +register_candidate(NvidiaMmaGemmEmittedCandidate()) diff --git a/python/tessera/compiler/fusion.py b/python/tessera/compiler/fusion.py index e8c2ecfc3..6b3c81958 100644 --- a/python/tessera/compiler/fusion.py +++ b/python/tessera/compiler/fusion.py @@ -17,6 +17,7 @@ FusedRegion as FusedRegion, FusionCost as FusionCost, GatedMatmulRegion as GatedMatmulRegion, + MatmulRegion as MatmulRegion, NormChainRegion as NormChainRegion, POINTWISE_OPS as POINTWISE_OPS, PointwiseGraphRegion as PointwiseGraphRegion, @@ -68,6 +69,7 @@ should_fuse_region as should_fuse_region, verify_synthesized_attention as verify_synthesized_attention, verify_synthesized_gated as verify_synthesized_gated, + verify_synthesized_matmul as verify_synthesized_matmul, verify_synthesized_pointwise as verify_synthesized_pointwise, verify_synthesized_region as verify_synthesized_region, ) @@ -196,6 +198,8 @@ "discover_fusable_regions", "discover_attention_regions", "GatedMatmulRegion", + "MatmulRegion", + "verify_synthesized_matmul", "synthesize_gated_matmul_msl", "run_gated_matmul_region", "should_fuse_gated", diff --git a/python/tessera/compiler/fusion_core.py b/python/tessera/compiler/fusion_core.py index e27dbba41..675f36da1 100644 --- a/python/tessera/compiler/fusion_core.py +++ b/python/tessera/compiler/fusion_core.py @@ -302,6 +302,41 @@ def reference(self, A: np.ndarray, B: np.ndarray, return out.astype(np.float32) +def _round_to_storage(x: "np.ndarray", dtype: str) -> "np.ndarray": + """Round f32 ``x`` to the 16-bit storage ``dtype`` and back to f32, so a + reference sees the same rounded operands the 16-bit GEMM kernels do.""" + a = np.ascontiguousarray(x, np.float32) + if dtype in ("float16", "f16"): + return a.astype(np.float16).astype(np.float32) + if dtype in ("bfloat16", "bf16"): + try: + import ml_dtypes + return a.astype(ml_dtypes.bfloat16).astype(np.float32) + except Exception: # round-to-nearest-even fallback + u = a.view(np.uint32).astype(np.uint64) + bits = (((u + 0x7FFF + ((u >> 16) & 1)) >> 16) << 16).astype(np.uint32) + return bits.view(np.float32) + raise ValueError(f"unsupported matmul storage dtype {dtype!r}") + + +@dataclass(frozen=True) +class MatmulRegion: + """A bare matmul ``D = A @ B`` (no fusion) — the region kind the D1 arbiter + keys a plain-GEMM candidate on (there is no matmul feature in ``FusedRegion``, + which always carries at least one fused op). ``dtype`` is the 16-bit storage + (``bfloat16``/``float16``); the accumulate is f32, matching the emitted + ``mma.sync`` and shipped GEMM kernels.""" + + dtype: str = "bfloat16" + + def reference(self, A: np.ndarray, B: np.ndarray) -> np.ndarray: + """The f32 result of ``A @ B`` with both operands rounded to ``dtype`` + first — the horizontal-oracle ground truth the GEMM candidate matches.""" + Aq = _round_to_storage(A, self.dtype) + Bq = _round_to_storage(B, self.dtype) + return (Aq @ Bq).astype(np.float32) + + # ───────────────────────────────────────────────────────────────────────────── # F2b-tiled — threadgroup-tiled synthesis for large N (the stack kernel caps at # SYNTH_MAX_N; this lifts it to SYNTH_MAX_N_TILED via dynamic threadgroup memory) @@ -934,6 +969,36 @@ def verify_synthesized_region(region: FusedRegion, *, seed: int = 0, return verdict +def verify_synthesized_matmul(region: "MatmulRegion", *, seed: int = 0, + atol: float = 1e-3, force: bool = False, + runner: KernelRunner | None = None) -> bool: + """Codegen-gated oracle for a bare-matmul candidate (see + ``verify_synthesized_region``): run the candidate's GEMM on an aligned probe + and compare to the dtype-rounded ``A @ B`` reference. Arbiter-only — the + ``runner`` is a candidate adapter exposing ``run_matmul``; a runner without it + (a real backend runner) trusts the reference, since nothing device-emitted is + in play for this op.""" + r = runner or _runner() + run = getattr(r, "run_matmul", None) + if run is None: + return True + key = (r.target, "M", region.dtype) + if not force and key in _VERIFY_CACHE: + return _VERIFY_CACHE[key] + rng = np.random.default_rng(seed) + M, N, K = 32, 16, 32 # aligned probe (M%16, N%8, K%16) + A = (rng.standard_normal((M, K)) * 0.4).astype(np.float32) + B = (rng.standard_normal((K, N)) * 0.4).astype(np.float32) + out, execution = run(region, A, B) + if execution in REFERENCE_EXECUTIONS: + verdict = True + else: + verdict = bool(np.allclose(out, region.reference(A, B), + atol=_effective_atol(r, atol))) + _VERIFY_CACHE[key] = verdict + return verdict + + def verify_synthesized_attention(region: AttentionRegion, *, seed: int = 0, atol: float = 1e-3, force: bool = False, runner: KernelRunner | None = None) -> bool: diff --git a/python/tessera/runtime.py b/python/tessera/runtime.py index 8d752479c..f380db102 100644 --- a/python/tessera/runtime.py +++ b/python/tessera/runtime.py @@ -1879,6 +1879,127 @@ def _execute_nvidia_mma_artifact(artifact: RuntimeArtifact, args: Any) -> Any: return d +# ── 2D GEMM execution helpers for the D1 arbiter's NVIDIA matmul candidates ──── +# Two lanes for a bare ``D = A @ B`` (f32 accumulate), both keyed by 16-bit +# storage dtype ("bfloat16"/"float16"): the shipped hand-tuned symbol (Tier 3) +# and the compiler-EMITTED ptx_emit kernel via the launch bridge (Tier 2). The +# arbiter (emit/candidate.py) F4-gates and selects between them. + +def _nvidia_mma_gemm_2d(A: Any, B: Any, dtype: str = "bfloat16") -> Any: + """Shipped ``libtessera_nvidia_gemm`` mma.sync GEMM on 2D ``A @ B`` -> f32. + ``B`` is row-major (the shipped convention). Raises on no lib / GPU / kernel + error (the caller declines to the reference).""" + import numpy as np + lib = _load_nvidia_gemm_runtime() + if lib is None: + raise RuntimeError("libtessera_nvidia_gemm.so not loadable") + sym = _NVIDIA_GEMM_SYMBOLS.get(dtype) + fn = getattr(lib, sym, None) if sym else None + if fn is None: + raise RuntimeError(f"shipped GEMM lacks a symbol for dtype {dtype!r}") + store = np.float16 if dtype == "float16" else _bfloat16_dtype() + if store is None: + raise RuntimeError("bfloat16 dtype unavailable (ml_dtypes not installed)") + Ac = np.ascontiguousarray(A, store) + Bc = np.ascontiguousarray(B, store) + M, K = Ac.shape + _, N = Bc.shape + D = np.zeros((M, N), np.float32) + rc = fn(Ac.ctypes.data_as(ctypes.c_void_p), Bc.ctypes.data_as(ctypes.c_void_p), + D.ctypes.data_as(ctypes.c_void_p), int(M), int(N), int(K)) + if rc != 0: + raise RuntimeError(f"shipped nvidia GEMM returned rc={rc}") + return D + + +_nvidia_ptx_launch_lib: ctypes.CDLL | None = None +_nvidia_ptx_registered: set[str] = set() + + +def _nvidia_ptx_launch_lib_path() -> Optional[Path]: + """Locate the PTX launch-bridge lib (env override -> canonical CMake build).""" + env = os.environ.get("TESSERA_NVIDIA_PTX_LAUNCH_LIB") + candidates: list[Path] = [] + if env: + candidates.append(Path(env)) + root = Path(__file__).resolve().parents[2] + candidates.append( + root / "build/src/compiler/codegen/tessera_gpu_backend_NVIDIA/runtime/cuda" + / "libtessera_nvidia_ptx_launch.so") + for c in candidates: + if c.is_file(): + return c + return None + + +def _load_nvidia_ptx_launch() -> ctypes.CDLL | None: + """Load the PTX launch bridge once (preloading libcuda globally so its weak + ``tsrRegisterGpuLauncher`` ref and cuda deps resolve). Returns None (never + raises) when the lib / CUDA deps are absent.""" + global _nvidia_ptx_launch_lib + if _nvidia_ptx_launch_lib is not None: + return _nvidia_ptx_launch_lib + path = _nvidia_ptx_launch_lib_path() + if path is None: + return None + cuda_dirs = ["/usr/lib/wsl/lib", + os.path.join(os.environ.get("CUDA_PATH", "/usr/local/cuda"), "lib64")] + for dep in ("libcuda.so.1", "libcuda.so"): + for d in cuda_dirs: + p = os.path.join(d, dep) + if os.path.isfile(p): + try: + ctypes.CDLL(p, mode=ctypes.RTLD_GLOBAL) + except OSError: + pass + break + try: + lib = ctypes.CDLL(str(path), mode=ctypes.RTLD_GLOBAL) + except OSError: + return None + lib.tessera_nvidia_ptx_register.argtypes = [ctypes.c_char_p, ctypes.c_char_p] + lib.tessera_nvidia_ptx_register.restype = ctypes.c_int + lib.tessera_nvidia_ptx_invoke.argtypes = [ + ctypes.c_char_p, ctypes.POINTER(ctypes.c_void_p), ctypes.c_size_t, + ctypes.POINTER(ctypes.c_int64), ctypes.c_size_t] + lib.tessera_nvidia_ptx_invoke.restype = ctypes.c_int + _nvidia_ptx_launch_lib = lib + return lib + + +def _nvidia_ptx_gemm_2d(A: Any, B: Any, dtype: str = "bfloat16") -> Any: + """Compiler-EMITTED mma.sync GEMM (ptx_emit) via the launch bridge on 2D + ``A @ B`` -> f32. Registers the emitted general-GEMM PTX once per dtype, then + invokes it. The emitted kernel wants ``B`` col-major (converted here) and + aligned M%16 / N%8 / K%16. Raises on no lib / GPU / launch error.""" + import numpy as np + from tessera.compiler import ptx_emit as pe + lib = _load_nvidia_ptx_launch() + if lib is None: + raise RuntimeError("libtessera_nvidia_ptx_launch.so not loadable") + edt = "f16" if dtype == "float16" else "bf16" + entry = pe.MMA_SYNC_GEMM_ENTRY[edt] + if entry not in _nvidia_ptx_registered: + ptx = pe.emit_mma_sync_gemm_ptx(dtype=edt) + if lib.tessera_nvidia_ptx_register(entry.encode(), ptx.encode()) != 0: + raise RuntimeError(f"ptx register failed for {entry}") + _nvidia_ptx_registered.add(entry) + store = np.float16 if dtype == "float16" else _bfloat16_dtype() + if store is None: + raise RuntimeError("bfloat16 dtype unavailable (ml_dtypes not installed)") + Ac = np.ascontiguousarray(A, store) # row-major A + Bc = np.asfortranarray(np.ascontiguousarray(B, store)) # col-major B storage + M, K = Ac.shape + _, N = Bc.shape + D = np.zeros((M, N), np.float32) + bufs = (ctypes.c_void_p * 3)(Ac.ctypes.data, Bc.ctypes.data, D.ctypes.data) + dims = (ctypes.c_int64 * 3)(int(M), int(N), int(K)) + rc = lib.tessera_nvidia_ptx_invoke(entry.encode(), bufs, 3, dims, 3) + if rc != 0: + raise RuntimeError(f"emitted nvidia GEMM invoke rc={rc}") + return D + + # ───────────────────────────────────────────────────────────────────────────── # Stage L4 — the COMPILED GEMM lane (the default rocm matmul execution path). # diff --git a/tests/unit/test_nvidia_plugin.py b/tests/unit/test_nvidia_plugin.py index 3beafc512..9897f6577 100644 --- a/tests/unit/test_nvidia_plugin.py +++ b/tests/unit/test_nvidia_plugin.py @@ -25,7 +25,7 @@ import tessera.compiler.fusion as F import tessera.compiler.emit.nvidia_cuda as nvidia # noqa: F401 — self-registers from tessera.compiler.emit import candidate as C -from tessera.compiler.emit.candidate import OP_FUSED_REGION, Tier +from tessera.compiler.emit.candidate import OP_FUSED_REGION, OP_MATMUL, Tier from tessera.compiler.emit.kernel_emitter import ( EmitError, SpecPolicy, get_emitter, get_runner, ) @@ -115,6 +115,36 @@ def test_nvidia_arbitrated_residual_threads_not_raises(): np.testing.assert_allclose(out, region.reference(A, B, None, res), atol=1e-2) +def test_nvidia_matmul_candidates_registered(): + # B1: the bare-matmul op-kind + the two GEMM lanes (shipped Tier-3, emitted + # Tier-2) registered under (nvidia, matmul). + cands = {c.name: c for c in C.candidates_for("nvidia", OP_MATMUL)} + assert set(cands) == {"nvidia_mma_gemm_shipped", "nvidia_mma_gemm_emitted"} + assert cands["nvidia_mma_gemm_shipped"].tier == Tier.HAND_TUNED + assert cands["nvidia_mma_gemm_emitted"].tier == Tier.EMITTED + for c in cands.values(): + assert c.op == OP_MATMUL + assert c.accuracy_atol == 5e-3 # 16-bit storage budget + assert c.applies_to(F.MatmulRegion(dtype="bfloat16")) + assert c.applies_to(F.MatmulRegion(dtype="float16")) + assert not c.applies_to(F.MatmulRegion(dtype="float32")) # not 16-bit + assert not c.applies_to(F.FusedRegion(epilogue=("relu",))) # not a matmul + + +def test_nvidia_matmul_off_gpu_arbitrates_to_reference(): + # Host-free: with no GPU the candidates are unavailable, so the arbiter finds + # no winner and run_arbitrated returns the numpy reference (never raises). + if _nvidia_cuda_live(): + pytest.skip("GPU present — covered by the live arbitration test") + region = F.MatmulRegion(dtype="bfloat16") + rng = np.random.default_rng(0) + A = rng.standard_normal((32, 32)).astype(np.float32) + B = rng.standard_normal((32, 16)).astype(np.float32) + out, tag = C.run_arbitrated(region, OP_MATMUL, "nvidia", A, B) + assert tag == "reference" + np.testing.assert_allclose(out, region.reference(A, B), atol=1e-3) + + def test_nvidia_missing_required_buffer_declines_not_segfault(): # Same NULL-deref guard as x86/ROCm: a residual/bias region without the buffer # must not launch the CUDA kernel (which would deref a null). Child process so a @@ -196,3 +226,39 @@ def test_live_nvidia_arbitrated_residual_executes(): out, tag = C.run_arbitrated(region, OP_FUSED_REGION, "nvidia", A, B, None, res) assert tag == "nvidia_cuda" np.testing.assert_allclose(out, region.reference(A, B, None, res), atol=1e-3) + + +def _nvidia_matmul_live() -> bool: + if not _nvidia_cuda_live(): + return False + try: + from tessera import runtime as rt + return rt._load_nvidia_ptx_launch() is not None + except Exception: + return False + + +@pytest.mark.slow +@pytest.mark.skipif(not _nvidia_matmul_live(), + reason="live NVIDIA GPU + shipped GEMM + PTX launch bridge required") +@pytest.mark.parametrize("dtype", ["bfloat16", "float16"]) +@pytest.mark.parametrize("shape", [(16, 8, 16), (32, 16, 32), (64, 64, 64)], + ids=lambda s: f"{s[0]}x{s[1]}x{s[2]}") +def test_live_nvidia_matmul_arbitrated(dtype, shape): + # B1: the arbiter picks the hand-tuned shipped GEMM (Tier 3) by default and + # runs it on-GPU; the E3 escape hatch forces the compiler-emitted lane (Tier 2) + # through the PTX launch bridge. Both match the dtype-rounded reference. + F.clear_verification_cache() + M, N, K = shape + region = F.MatmulRegion(dtype=dtype) + rng = np.random.default_rng(M + K) + A = (rng.standard_normal((M, K)) * 0.4).astype(np.float32) + B = (rng.standard_normal((K, N)) * 0.4).astype(np.float32) + ref = region.reference(A, B) + out, tag = C.run_arbitrated(region, OP_MATMUL, "nvidia", A, B) + assert tag == "nvidia_mma_shipped" # Tier-3 default + np.testing.assert_allclose(out, ref, atol=5e-3, rtol=0) + out2, tag2 = C.run_arbitrated(region, OP_MATMUL, "nvidia", A, B, + force="nvidia_mma_gemm_emitted") + assert tag2 == "nvidia_ptx_gemm" # Tier-2, forced (E3) + np.testing.assert_allclose(out2, ref, atol=5e-3, rtol=0)