diff --git a/docs/audit/compiler/COMPILER_REFACTOR_PLAN.md b/docs/audit/compiler/COMPILER_REFACTOR_PLAN.md index 3cda7f5d1..8234a37e2 100644 --- a/docs/audit/compiler/COMPILER_REFACTOR_PLAN.md +++ b/docs/audit/compiler/COMPILER_REFACTOR_PLAN.md @@ -74,8 +74,9 @@ table is the single skim surface. `โœ…` done ยท `๐ŸŸก` partial ยท `โฌœ` not star | **3** | Oracle accuracy budget (`KernelRunner.accuracy_atol`, D2 seed) | โœ… | โ€” | โ€” | | **3** | C1b x86 AOCL-DLP Tier-3 candidate (opt-in) | โ€” | โฌœ | โ€” | | **4** | C2 NVIDIA emit pipeline + launch bridge | โฌœ | โ€” | โฌœ | -| **5** | C3 ROCm emit pipeline (generic synth โ†’ HIP) | โฌœ | โฌœ | โ€” | -| **3.5** | ROCm runner โ†’ F4 gate (`emit/rocm_hip.py`, shipped kernels, f16 budget) | โœ… | โœ… gfx1151 attn | โ€” | +| **5** | C3 ROCm generic synth โ†’ HIP (`emit/rocm_hip.py`: emitter + `hipcc` + runner) | โœ… emit host-free | โœ… gfx1151 FusedRegion | โ€” | +| **5** | C3 tail โ€” drive WMMA/MFMA `Generate*` passes through the loop | โฌœ | โฌœ | โ€” | +| **3.5** | ROCm shipped-kernel โ†’ F4 gate (flash-attn, f16 budget) + shared scalar body | โœ… | โœ… gfx1151 attn | โ€” | | **6** | D1โ€“D3 arbitration + measured autotune | โฌœ | โฌœ | โฌœ | **Gate reality (softens ยง4/ยง9.2):** "Phase 0 gates everything" holds only for the @@ -256,10 +257,21 @@ chains, small attention). Crown-jewel GEMM stays Tier 2/3. **bf16-only, few-shape** emitter, and today's executing sm_120 matmul runs via the shipped `libtessera_nvidia_gemm.so`, **not** the emit path โ€” so the bridge is the long pole, ahead of broadening shapes/dtypes. -- **C3 ยท ROCm in-process emit pipeline** `[MAC]` authoring โ†’ `[AMD]` proof โ€” - `--tessera-emit-rocm`: drives the existing gfx1151 WMMA + CDNA MFMA `Generate*` - passes through the shared loop into the launch bridge; reuses the async-token - SSA model in `ROCMWaveLdsPipeline`. +- **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 + one-thread-per-row `__global__` kernel + a host-pointer C-ABI wrapper doing + H2D/launch/D2H), `_rocm_hip_compile_fn` compiles it with `hipcc + --offload-arch=` โ†’ `.so`, and `RocmHipRunner.run_fused_region` dlopens + + launches on gfx1151 (`"rocm_hip"`), F4-gated. The per-row scalar body is + **shared with the x86 C lane** (`emit/_fused_scalar_body.py`) so both stay + locked to the one `fusion_core` reference. Same NULL-buffer guard as x86. + **C3 tail (still open):** wire the existing gfx1151 WMMA + CDNA MFMA + `Generate*` passes (the hand-tuned MLIR kernel generators) through the shared + loop as Tier-3 arbiter candidates, reusing the async-token SSA model in + `ROCMWaveLdsPipeline` โ€” the generic scalar HIP kernel above is a correctness- + first middle-ground candidate, NOT a replacement for those crown-jewel lanes + (lead-safety; the D1 arbiter picks per measured latency + accuracy budget). - **C3-precursor (landed 2026-07-06): ROCm runner โ†’ F4 gate + oracle accuracy budget** `[MAC]` author โ†’ `[AMD]` proof. Ahead of the full emit pipeline, the *shipped* gfx1151 kernels are now wired into the universal F4 oracle: diff --git a/python/tessera/compiler/emit/_fused_scalar_body.py b/python/tessera/compiler/emit/_fused_scalar_body.py new file mode 100644 index 000000000..0ce8f0624 --- /dev/null +++ b/python/tessera/compiler/emit/_fused_scalar_body.py @@ -0,0 +1,107 @@ +"""Shared scalar kernel body for the ``FusedRegion`` compiled lanes (x86 C + ROCm +HIP). The per-element math is arch-agnostic โ€” the same ``expf``/``tanhf``/ +``rsqrtf`` snippets compile in host C (clang/cc, x86 lane) and HIP device code +(hipcc, ROCm lane) โ€” so both backends synthesize the *identical* body from here +and stay locked to the same ``fusion_core`` numpy reference. Only the wrapper +differs (a plain C function that loops over rows vs a ``__global__`` one thread +per row); each backend supplies its own. + +The body assumes these names are in scope: ``const float* A, B, bias, residual``, +``float* row`` (= ``out + m*N``), and ``int M, N, K, m``. It writes ``row[0..N]`` +for row ``m`` (matmul + prologue on A + epilogue chain + optional residual + +optional row reduction), matching ``FusedRegion.reference`` element for element. + +``EpilogueOp.emit(target)`` / ``ReductionOp.emit(target)`` deliberately raise for +non-Metal targets, so the op-name โ†’ C tables live here (Decision #21: no silent +wrong-language emit). +""" +from __future__ import annotations + +from typing import Any + + +def pointwise_snippet(op: str, var: str) -> str: + """C/HIP statement applying pointwise epilogue op ``op`` to lvalue ``var`` + (f32). Mirrors ``EPILOGUE_OPS[op].ref`` numerically. ``bias`` is handled by + the caller (it reads ``bias[n]``); the rest are activations valid in a + prologue too.""" + if op == "relu": + return f"{var} = {var} > 0.0f ? {var} : 0.0f;" + if op == "gelu": # tanh approximation, clamped โ€” identical to fusion_core._gelu + return ( + f"{{ float _t = 0.7978845608028654f*({var}+0.044715f*{var}*{var}*{var});" + f" _t = _t < -30.0f ? -30.0f : (_t > 30.0f ? 30.0f : _t);" + f" {var} = 0.5f*{var}*(1.0f+tanhf(_t)); }}" + ) + if op == "silu": + return f"{var} = {var} / (1.0f + expf(-{var}));" + if op == "sigmoid": + return f"{var} = 1.0f / (1.0f + expf(-{var}));" + if op == "tanh": + return f"{var} = tanhf({var});" + raise ValueError(f"no scalar snippet for pointwise op {op!r}") + + +def reduction_snippet(name: str, eps: float) -> str: + """C/HIP block reducing the length-``N`` row ``row`` in place (f32). Mirrors + ``REDUCTION_OPS[name].ref``.""" + if name == "rmsnorm": + return ( + " { float _ss = 0.0f;\n" + " for (int n = 0; n < N; ++n) _ss += row[n]*row[n];\n" + f" float _inv = 1.0f/sqrtf(_ss/(float)N + {eps!r}f);\n" + " for (int n = 0; n < N; ++n) row[n] = row[n]*_inv; }\n" + ) + if name == "softmax": + return ( + " { float _mx = -INFINITY;\n" + " for (int n = 0; n < N; ++n) if (row[n] > _mx) _mx = row[n];\n" + " float _sm = 0.0f;\n" + " for (int n = 0; n < N; ++n) { row[n] = expf(row[n]-_mx); _sm += row[n]; }\n" + " for (int n = 0; n < N; ++n) row[n] = row[n]/_sm; }\n" + ) + if name == "layer_norm": + return ( + " { float _mean = 0.0f;\n" + " for (int n = 0; n < N; ++n) _mean += row[n];\n" + " _mean /= (float)N;\n" + " float _var = 0.0f;\n" + " for (int n = 0; n < N; ++n) { float _d = row[n]-_mean; _var += _d*_d; }\n" + f" float _inv = 1.0f/sqrtf(_var/(float)N + {eps!r}f);\n" + " for (int n = 0; n < N; ++n) row[n] = (row[n]-_mean)*_inv; }\n" + ) + raise ValueError(f"no scalar snippet for reduction op {name!r}") + + +def row_compute_body(region: Any) -> str: + """The per-row compute body for row ``m`` (``float* row = out + m*N`` in scope): + matmul + prologue(A) + epilogue chain + optional residual + optional row + reduction. Shared verbatim by the x86 C function and the ROCm HIP kernel.""" + prologue = "".join( + f" {pointwise_snippet(op, 'a')}\n" for op in region.prologue + ) + epi_lines = [] + for op in region.epilogue: + if op == "bias": + epi_lines.append(" v = v + bias[n];") + else: + epi_lines.append(f" {pointwise_snippet(op, 'v')}") + epilogue = "\n".join(epi_lines) + residual = (" v = v + residual[(long)m*N + n];\n" + if region.has_residual else "") + reduction = (reduction_snippet(region.reduction, region.eps) + if region.reduction else "") + return ( + " for (int n = 0; n < N; ++n) {\n" + " float v = 0.0f;\n" + " for (int k = 0; k < K; ++k) {\n" + " float a = A[(long)m*K + k];\n" + f"{prologue}" + " v += a * B[(long)k*N + n];\n" + " }\n" + f"{epilogue}\n" + f"{residual}" + " row[n] = v;\n" + " }\n" + f"{reduction}" + ) diff --git a/python/tessera/compiler/emit/rocm_hip.py b/python/tessera/compiler/emit/rocm_hip.py index c9784b716..41668b8b7 100644 --- a/python/tessera/compiler/emit/rocm_hip.py +++ b/python/tessera/compiler/emit/rocm_hip.py @@ -1,44 +1,226 @@ -"""Workstream C โ€” ROCm gfx1151 plugin: wire the shipped hardware-verified -gfx1151 kernels into the target-agnostic F4 oracle. - -Unlike the x86 plugin (which *emits* + compiles C), ROCm's lead kernels are the -already-shipped, hardware-verified gfx1151 lanes (WMMA GEMM, compiled FA-2 -flash-attention). So this module registers a :class:`KernelRunner` **only** โ€” no -emitter, no ``compile_fn`` โ€” which lets ROCm's real kernels be gated by the same -universal F4 correctness oracle as the synthesized backends (the cross-backend -differential-equivalence superpower, Theory ยง7.5) *without* claiming a generic -ROCm emit lane (that is C3). For region kinds ROCm has no single fused GPU kernel -for yet (matmul+epilogue / gated / pointwise), it declines to the numpy reference -โ€” honest, never a mislabeled kernel (Decision #21). - -Precision: the flash-attn lane is f16 storage / f32 accumulate, so the runner -declares an f16 accuracy budget (:attr:`accuracy_atol`); the oracle widens its -tolerance to it so f16 rounding is not misread as a miscompile, while an O(1) -miscompile is still caught (Decision #28, the accuracy-budgeted arbiter). - -Runs only where a live gfx1151 + the compiled flash lane are present (probed via -``runtime._rocm_compiled_flash_attn_available``); on any other host it declines -to the reference, so authoring/tests stay host-free. +"""Workstream C3 โ€” ROCm gfx1151 codegen plugin: generic synth โ†’ HIP. + +Two lanes under one `target = "rocm"` plugin, both F4-gated on real silicon: + +* **Generic compiled lane (C3)** โ€” a full three seams for the fusable + middle ground (`FusedRegion`: matmul + prologue/epilogue/residual/reduction): + - :class:`RocmHipEmitter` (`register_emitter`) โ€” region โ†’ HIP source (a + ``__global__`` one-thread-per-row kernel + a host-pointer C-ABI wrapper), + reusing the *same* scalar body as the x86 C lane + (`_fused_scalar_body.row_compute_body`) so both stay locked to the + `fusion_core` numpy reference. + - :func:`_rocm_hip_compile_fn` (`register_compiler`) โ€” `hipcc + --offload-arch= -O3 -shared` โ†’ a `.so` the runtime dlopens. + - :meth:`RocmHipRunner.run_fused_region` โ€” H2D / launch / D2H via the shipped + lib's host-pointer ABI โ†’ `(out, "rocm_hip")`, else the reference. +* **Shipped hand-tuned lane (Tier 3)** โ€” :meth:`RocmHipRunner.run_fused_attention` + runs the shipped compiled FA-2 flash-attn kernel (not generically emitted); the + same universal oracle gates it. This is the cross-backend differential- + equivalence superpower on the lead's real kernels. + +Lead-safety: the generic HIP kernel is a correctness-first candidate for the +middle ground โ€” crown-jewel WMMA/MFMA GEMM stays first-class (the D1 arbiter +picks the generic lane only where it measures faster and in budget). Runs only +where a live gfx1151 + `hipcc` are present; everywhere else it declines to the +numpy reference so authoring/tests stay host-free. + +Precision: the flash lane is f16 storage, so the runner declares an f16 +`accuracy_atol` budget the oracle honors (Decision #28); the generic f32 HIP +kernel is comfortably within it. """ from __future__ import annotations +import ctypes +import os +import shutil +import subprocess +import tempfile from typing import Any -from tessera.compiler.emit.kernel_emitter import KernelRunner, register_runner +from tessera.compiler.emit._fused_scalar_body import row_compute_body +from tessera.compiler.emit.kernel_cache import build, register_compiler +from tessera.compiler.emit.kernel_emitter import ( + EmitError, + KernelEmitter, + KernelSource, + KernelRunner, + SpecPolicy, + bucket_key, + register_emitter, + register_runner, +) +from tessera.compiler.fusion_core import FusedRegion _TARGET = "rocm" +_LANG = "hip" +_ENTRY = "tessera_rocm_fused" _REAL_TAG = "rocm_hip" -#: f16 storage budget for the WMMA / flash lanes vs the f32 reference. Loose -#: enough for f16 rounding (measured max ~2.5e-3 on the oracle probes), tight -#: enough that an O(1) miscompile (transpose / wrong softmax / wrong scale) is -#: still caught. +#: f16 storage budget for the shipped flash lane vs the f32 reference. Loose +#: enough for f16 rounding (~2.5e-3 on the oracle probes), tight enough that an +#: O(1) miscompile is still caught. The generic f32 HIP kernel is well within it. _F16_ATOL = 5e-3 +# โ”€โ”€ HIP source synthesis (generic FusedRegion lane) โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ + +def _synthesize_fused_hip(region: FusedRegion) -> str: + """HIP source for a ``FusedRegion`` (f32): a one-thread-per-row kernel embedding + the shared scalar body, plus a host-pointer C-ABI wrapper that does H2D / + launch / D2H (same shape as the shipped ``libtessera_rocm_gemm.so`` symbols). + Dims are runtime args, so one kernel serves every shape.""" + return ( + "#include \n" + "#include \n" + f"__global__ void {_ENTRY}_kernel(const float* A, const float* B,\n" + " const float* bias, const float* residual, float* out,\n" + " int M, int N, int K) {\n" + " int m = blockIdx.x*blockDim.x + threadIdx.x;\n" + " if (m >= M) return;\n" + " float* row = out + (long)m * N;\n" + f"{row_compute_body(region)}" + "}\n" + f'extern "C" int {_ENTRY}(const float* hA, const float* hB,\n' + " const float* hbias, const float* hresidual, float* hout,\n" + " int M, int N, int K) {\n" + " size_t szA=(size_t)M*K*sizeof(float), szB=(size_t)K*N*sizeof(float),\n" + " szO=(size_t)M*N*sizeof(float);\n" + " float *dA=0,*dB=0,*dbias=0,*dres=0,*dO=0;\n" + " if (hipMalloc(&dA,szA)!=hipSuccess) return 2;\n" + " if (hipMalloc(&dB,szB)!=hipSuccess) { hipFree(dA); return 2; }\n" + " if (hipMalloc(&dO,szO)!=hipSuccess) { hipFree(dA); hipFree(dB); return 2; }\n" + " hipMemcpy(dA,hA,szA,hipMemcpyHostToDevice);\n" + " hipMemcpy(dB,hB,szB,hipMemcpyHostToDevice);\n" + " if (hbias) { hipMalloc(&dbias,(size_t)N*sizeof(float));\n" + " hipMemcpy(dbias,hbias,(size_t)N*sizeof(float),hipMemcpyHostToDevice); }\n" + " if (hresidual) { hipMalloc(&dres,szO);\n" + " hipMemcpy(dres,hresidual,szO,hipMemcpyHostToDevice); }\n" + " int t=64, b=(M+t-1)/t;\n" + f" hipLaunchKernelGGL({_ENTRY}_kernel, dim3(b), dim3(t), 0, 0,\n" + " dA,dB,dbias,dres,dO,M,N,K);\n" + " int ok = (hipDeviceSynchronize()==hipSuccess) ? 1 : 3;\n" + " if (ok==1) hipMemcpy(hout,dO,szO,hipMemcpyDeviceToHost);\n" + " hipFree(dA); hipFree(dB); hipFree(dO);\n" + " if (dbias) hipFree(dbias);\n" + " if (dres) hipFree(dres);\n" + " return ok;\n" + "}\n" + ) + + +class RocmHipEmitter(KernelEmitter): + target = _TARGET + lang = _LANG + + def can_emit(self, region: Any) -> bool: + return isinstance(region, FusedRegion) + + def emit(self, region: Any, *, spec: SpecPolicy = SpecPolicy.BUCKET, + dtype: str = "f32", dims: tuple[int, ...] | None = None) -> KernelSource: + if not isinstance(region, FusedRegion): + raise EmitError( + f"RocmHipEmitter cannot emit a region of type " + f"{type(region).__name__} (only FusedRegion; attention uses the " + "shipped flash lane)") + if spec is SpecPolicy.DYNAMIC: + raise EmitError("RocmHipEmitter does not yet support SpecPolicy.DYNAMIC") + if dtype != "f32": + raise EmitError(f"RocmHipEmitter only supports f32 so far, got {dtype!r}") + source = _synthesize_fused_hip(region) + key = bucket_key(dims, spec, dim_names=getattr(region, "dim_names", None)) + return KernelSource(source=source, entry=_ENTRY, lang=self.lang, + spec=spec, shape_key=key) + + +# โ”€โ”€ compile_fn (HIP โ†’ .so) โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ + +def _rocm_arch() -> str: + """gfx target: ``$TESSERA_ROCM_ARCH`` override, else the live device's chip, + else gfx1151 (the Strix Halo default).""" + env = os.environ.get("TESSERA_ROCM_ARCH") + if env: + return env + try: + from tessera import runtime as rt + chip = rt._rocm_chip() + if chip: + return str(chip) + except Exception: + pass + return "gfx1151" + + +def _rocm_hip_compile_fn(source: KernelSource) -> str: + """Compile the emitted HIP to a shared object with hipcc and return its path. + Raises on a missing toolchain/compile failure; ``build`` wraps in + ``CompileError`` (never a silent no-op).""" + hipcc = shutil.which("hipcc") or "/opt/rocm/bin/hipcc" + d = tempfile.mkdtemp(prefix="tessera_rocm_") + src = os.path.join(d, "kernel.hip") + so = os.path.join(d, "kernel.so") + with open(src, "w") as f: + f.write(source.source) + subprocess.run( + [hipcc, f"--offload-arch={_rocm_arch()}", "-O3", "-fPIC", "-shared", + src, "-o", so], + check=True, capture_output=True, text=True) + return so + + +# โ”€โ”€ runner (execute โ†’ (out, tag)) โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ + +_LIB_CACHE: dict[str, Any] = {} + + +def _load_entry(artifact: str): + lib = _LIB_CACHE.get(artifact) + if lib is None: + lib = ctypes.CDLL(artifact) + _LIB_CACHE[artifact] = lib + fn = getattr(lib, _ENTRY) + fn.restype = ctypes.c_int + fn.argtypes = [ctypes.c_void_p] * 5 + [ctypes.c_int] * 3 + return fn + + +def _ptr(arr): + return arr.ctypes.data_as(ctypes.c_void_p) if arr is not None else None + + class RocmHipRunner(KernelRunner): target = _TARGET accuracy_atol = _F16_ATOL + def run_fused_region(self, region: Any, A: Any, B: Any, bias: Any = None, + *args: Any, residual: Any = None, + **kwargs: Any) -> tuple[Any, str]: + import numpy as np + # Required-buffer guard BEFORE launch: the emitted HIP dereferences + # bias[n] / residual[...] whenever the region declares them, so a missing + # buffer would pass a null the kernel derefs. Route ill-formed calls + # through the reference (a clean, catchable ValueError) instead. + if (region.has_bias and bias is None) or \ + (region.has_residual and residual is None): + return region.reference(A, B, bias, residual), "reference" + try: + Af = np.ascontiguousarray(A, np.float32) + Bf = np.ascontiguousarray(B, np.float32) + M, K = Af.shape + _, N = Bf.shape + compiled = build(region, _TARGET, dtype="f32", dims=None) + fn = _load_entry(compiled.artifact) + bias_arr = (np.ascontiguousarray(bias, np.float32) + if bias is not None else None) + res_arr = (np.ascontiguousarray(residual, np.float32) + if residual is not None else None) + out = np.zeros((M, N), np.float32) + rc = fn(_ptr(Af), _ptr(Bf), _ptr(bias_arr), _ptr(res_arr), + _ptr(out), M, N, K) + if rc == 1: + return out, _REAL_TAG + except Exception: + pass + return region.reference(A, B, bias, residual), "reference" + def run_fused_attention(self, region: Any, Q: Any, K: Any, V: Any, *a: Any, **k: Any) -> tuple[Any, str]: import numpy as np @@ -61,13 +243,7 @@ def run_fused_attention(self, region: Any, Q: Any, K: Any, V: Any, except Exception: return region.reference(Q, K, V), "reference" - # No single fused GPU kernel for these yet (that is the C3 emit lane) โ€” - # decline honestly to the numpy reference. - def run_fused_region(self, region: Any, A: Any, B: Any, bias: Any = None, - *a: Any, residual: Any = None, - **k: Any) -> tuple[Any, str]: - return region.reference(A, B, bias, residual), "reference" - + # No single fused GPU kernel for these yet โ€” decline to the numpy reference. def run_gated_matmul_region(self, region: Any, A: Any, Wg: Any, Wu: Any, *a: Any, **k: Any) -> tuple[Any, str]: return region.reference(A, Wg, Wu), "reference" @@ -77,6 +253,7 @@ def run_pointwise_graph(self, region: Any, arrays: Any, return region.reference(*arrays), "reference" -# Runner only (no emitter/compile_fn โ€” ROCm's kernels are shipped, not -# synthesized here). default=False so Apple stays the active default runner. +# โ”€โ”€ registration โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ +register_emitter(RocmHipEmitter()) +register_compiler(_TARGET, _rocm_hip_compile_fn) register_runner(RocmHipRunner(), default=False) diff --git a/python/tessera/compiler/emit/x86_llvm.py b/python/tessera/compiler/emit/x86_llvm.py index 66d17277f..6a7e1c49c 100644 --- a/python/tessera/compiler/emit/x86_llvm.py +++ b/python/tessera/compiler/emit/x86_llvm.py @@ -32,6 +32,7 @@ import tempfile from typing import Any +from tessera.compiler.emit._fused_scalar_body import row_compute_body from tessera.compiler.emit.kernel_cache import build, register_compiler from tessera.compiler.emit.kernel_emitter import ( EmitError, @@ -51,77 +52,12 @@ _REAL_TAG = "x86_native" -# โ”€โ”€ op-name โ†’ C snippet tables (match the fusion_core numpy references) โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ - -def _pointwise_c(op: str, var: str) -> str: - """C statement applying pointwise epilogue op ``op`` to lvalue ``var`` (f32). - Mirrors ``EPILOGUE_OPS[op].ref`` numerically. ``bias`` is handled separately - (it reads ``bias[n]``); the rest are pure activations valid in a prologue too.""" - if op == "relu": - return f"{var} = {var} > 0.0f ? {var} : 0.0f;" - if op == "gelu": # tanh approximation, clamped โ€” identical to fusion_core._gelu - return ( - f"{{ float _t = 0.7978845608028654f*({var}+0.044715f*{var}*{var}*{var});" - f" _t = _t < -30.0f ? -30.0f : (_t > 30.0f ? 30.0f : _t);" - f" {var} = 0.5f*{var}*(1.0f+tanhf(_t)); }}" - ) - if op == "silu": - return f"{var} = {var} / (1.0f + expf(-{var}));" - if op == "sigmoid": - return f"{var} = 1.0f / (1.0f + expf(-{var}));" - if op == "tanh": - return f"{var} = tanhf({var});" - raise EmitError(f"x86: no C snippet for pointwise op {op!r}") - - -def _reduction_c(name: str, eps: float) -> str: - """C block reducing over the length-``N`` row ``row`` in place (f32). Mirrors - ``REDUCTION_OPS[name].ref``.""" - if name == "rmsnorm": - return ( - " { float _ss = 0.0f;\n" - " for (int n = 0; n < N; ++n) _ss += row[n]*row[n];\n" - f" float _inv = 1.0f/sqrtf(_ss/(float)N + {eps!r}f);\n" - " for (int n = 0; n < N; ++n) row[n] = row[n]*_inv; }\n" - ) - if name == "softmax": - return ( - " { float _mx = -INFINITY;\n" - " for (int n = 0; n < N; ++n) if (row[n] > _mx) _mx = row[n];\n" - " float _sm = 0.0f;\n" - " for (int n = 0; n < N; ++n) { row[n] = expf(row[n]-_mx); _sm += row[n]; }\n" - " for (int n = 0; n < N; ++n) row[n] = row[n]/_sm; }\n" - ) - if name == "layer_norm": - return ( - " { float _mean = 0.0f;\n" - " for (int n = 0; n < N; ++n) _mean += row[n];\n" - " _mean /= (float)N;\n" - " float _var = 0.0f;\n" - " for (int n = 0; n < N; ++n) { float _d = row[n]-_mean; _var += _d*_d; }\n" - f" float _inv = 1.0f/sqrtf(_var/(float)N + {eps!r}f);\n" - " for (int n = 0; n < N; ++n) row[n] = (row[n]-_mean)*_inv; }\n" - ) - raise EmitError(f"x86: no C snippet for reduction op {name!r}") - - def _synthesize_fused_c(region: FusedRegion) -> str: - """Emit the C source for a ``FusedRegion`` (f32). Signature is dims-invariant + """Emit the C source for a ``FusedRegion`` (f32) โ€” a plain host function that + loops over rows and embeds the shared per-row body. Signature is dims-invariant (M/N/K are runtime args), so one kernel serves every shape โ€” the arbiter/cache - key it shape-anonymously.""" - prologue = "".join( - f" {_pointwise_c(op, 'a')}\n" for op in region.prologue - ) - epi_lines = [] - for op in region.epilogue: - if op == "bias": - epi_lines.append(" v = v + bias[n];") - else: - epi_lines.append(f" {_pointwise_c(op, 'v')}") - epilogue = "\n".join(epi_lines) - residual = (" v = v + residual[(long)m*N + n];\n" - if region.has_residual else "") - reduction = _reduction_c(region.reduction, region.eps) if region.reduction else "" + key it shape-anonymously. The per-row math is shared with the ROCm HIP lane + (`_fused_scalar_body.row_compute_body`) so both stay locked to one reference.""" return ( "#include \n" f"int {_ENTRY}(const float* A, const float* B, const float* bias,\n" @@ -129,18 +65,7 @@ def _synthesize_fused_c(region: FusedRegion) -> str: " int M, int N, int K) {\n" " for (int m = 0; m < M; ++m) {\n" " float* row = out + (long)m * N;\n" - " for (int n = 0; n < N; ++n) {\n" - " float v = 0.0f;\n" - " for (int k = 0; k < K; ++k) {\n" - " float a = A[(long)m*K + k];\n" - f"{prologue}" - " v += a * B[(long)k*N + n];\n" - " }\n" - f"{epilogue}\n" - f"{residual}" - " row[n] = v;\n" - " }\n" - f"{reduction}" + f"{row_compute_body(region)}" " }\n" " return 1;\n" "}\n" diff --git a/tests/unit/test_kernel_emitter.py b/tests/unit/test_kernel_emitter.py index dce9c2170..2f1369da5 100644 --- a/tests/unit/test_kernel_emitter.py +++ b/tests/unit/test_kernel_emitter.py @@ -116,8 +116,9 @@ def test_apple_emitter_rejects_unknown_region(): def test_registry_resolves_apple_and_reports_unknown_target(): assert get_emitter("apple_gpu").target == "apple_gpu" + # A genuinely unregistered target (x86/rocm now register real emitters). with pytest.raises(EmitError, match="no KernelEmitter registered"): - emit_kernel(F.FusedRegion(epilogue=("relu",)), "rocm") + emit_kernel(F.FusedRegion(epilogue=("relu",)), "no_such_backend") def test_register_emitter_requires_target(): diff --git a/tests/unit/test_rocm_plugin.py b/tests/unit/test_rocm_plugin.py index 228488199..be20dfe8a 100644 --- a/tests/unit/test_rocm_plugin.py +++ b/tests/unit/test_rocm_plugin.py @@ -1,27 +1,30 @@ -"""Workstream C โ€” ROCm gfx1151 plugin: wire shipped kernels into the F4 oracle. +"""Workstream C3 โ€” ROCm gfx1151 plugin: generic synth โ†’ HIP + shipped-kernel gate. Three layers: -1. **Registration + decline paths (host-free)** โ€” a runner-only plugin registers - for target "rocm" (no emitter/compiler); it declines matmul-epilogue / gated / - pointwise regions to the numpy reference (ROCm has no single fused GPU kernel - for those yet โ€” that is C3). +1. **Registration + emit + decline paths (host-free)** โ€” a full three-seam plugin + for target "rocm": the emitter turns a FusedRegion into HIP source; gated / + pointwise regions (no fused GPU kernel yet) decline to the numpy reference. 2. **Accuracy-budget wiring (host-free)** โ€” the F4 oracle widens its tolerance to a runner's declared ``accuracy_atol`` so an f16 lead kernel's rounding is not misread as a miscompile, while an O(1) miscompile still is. -3. **Live attention gate (needs a live gfx1151 + compiled flash lane)** โ€” the - shipped compiled flash-attn kernel is gated by the same universal oracle: - `run_fused_attention` runs on-device ("rocm_hip") and matches the numpy - reference within the f16 budget. +3. **Live gates (needs a live gfx1151)** โ€” the generic HIP FusedRegion lane + (`hipcc` compile + launch, "rocm_hip") and the shipped compiled flash-attn + lane are both gated by the same universal oracle on-device. """ from __future__ import annotations +import os +import shutil + import numpy as np import pytest import tessera.compiler.fusion as F import tessera.compiler.emit.rocm_hip as rocm # noqa: F401 โ€” self-registers -from tessera.compiler.emit.kernel_emitter import KernelRunner, get_runner +from tessera.compiler.emit.kernel_emitter import ( + EmitError, KernelRunner, SpecPolicy, get_emitter, get_runner, +) def _rocm_flash_live() -> bool: @@ -32,7 +35,17 @@ def _rocm_flash_live() -> bool: return False -# โ”€โ”€ 1. Registration + decline paths (host-free) โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ +def _rocm_hip_live() -> bool: + if not (shutil.which("hipcc") or os.path.exists("/opt/rocm/bin/hipcc")): + return False + try: + from tessera import runtime as rt + return rt._rocm_wmma_runtime_available() + except Exception: + return False + + +# โ”€โ”€ 1. Registration + emit + decline paths (host-free) โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ def test_rocm_runner_registered_with_f16_budget(): r = get_runner("rocm") @@ -40,25 +53,69 @@ def test_rocm_runner_registered_with_f16_budget(): assert r.accuracy_atol == 5e-3 # f16 storage budget -def test_rocm_registers_no_emitter(): - # ROCm's kernels are shipped, not synthesized here โ€” emit("rocm") still raises. - from tessera.compiler.emit.kernel_emitter import EmitError, emit_kernel - with pytest.raises(EmitError, match="no KernelEmitter registered"): - emit_kernel(F.FusedRegion(epilogue=("relu",)), "rocm") +def test_rocm_emitter_registered_produces_hip(): + from tessera.compiler.emit.kernel_cache import get_compiler + src = get_emitter("rocm").emit(F.FusedRegion(epilogue=("bias", "gelu")), + dtype="f32") + assert src.lang == "hip" + assert src.entry == "tessera_rocm_fused" + assert "__global__" in src.source and 'extern "C"' in src.source + assert "hipLaunchKernelGGL" in src.source + assert callable(get_compiler("rocm")) + +def test_rocm_emitter_rejects_unsupported(): + e = get_emitter("rocm") + with pytest.raises(EmitError, match="cannot emit"): + e.emit(F.AttentionRegion()) + with pytest.raises(EmitError, match="DYNAMIC"): + e.emit(F.FusedRegion(epilogue=("relu",)), spec=SpecPolicy.DYNAMIC) + with pytest.raises(EmitError, match="f32"): + e.emit(F.FusedRegion(epilogue=("relu",)), dtype="f16") -def test_rocm_declines_non_attention_regions(): + +def test_rocm_declines_gated_and_pointwise(): + # No single fused GPU kernel for these yet โ€” always the numpy reference. r = get_runner("rocm") A = np.zeros((8, 12), np.float32) - B = np.zeros((12, 16), np.float32) - _, ex = r.run_fused_region(F.FusedRegion(epilogue=("relu",)), A, B, None) - assert ex == "reference" _, ex = r.run_gated_matmul_region(F.GatedMatmulRegion(), A, np.zeros((12, 16), np.float32), np.zeros((12, 16), np.float32)) assert ex == "reference" +def test_rocm_missing_required_buffer_declines_not_segfault(): + # Same NULL-deref guard as x86: a residual/bias region without the buffer must + # not launch the HIP kernel (which would deref a null). Child process so a + # regression is a failed assert, not a crashed session. + import subprocess + import sys + import textwrap + code = textwrap.dedent( + """ + import numpy as np + import tessera.compiler.fusion as F + import tessera.compiler.emit.rocm_hip as rocm + r = rocm.RocmHipRunner() + A = np.zeros((8, 12), np.float32) + B = np.zeros((12, 16), np.float32) + for region in (F.FusedRegion(epilogue=("relu",), residual=True), + F.FusedRegion(epilogue=("bias", "relu"))): + try: + r.run_fused_region(region, A, B, None) + raise SystemExit("expected ValueError, got a result") + except ValueError: + pass + print("ok") + """ + ) + p = subprocess.run([sys.executable, "-c", code], capture_output=True, text=True) + assert p.returncode == 0, ( + f"missing-buffer guard failed (rc={p.returncode}, -11=SIGSEGV): " + f"{p.stderr[-300:]}") + assert "ok" in p.stdout + + # โ”€โ”€ 2. Accuracy-budget wiring (host-free) โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ class _FakeF16Attn(KernelRunner): @@ -108,7 +165,38 @@ def run_fused_attention(self, region, Q, K, V, *a, **k): F.AttentionRegion(scale=0.25), runner=_Wrong(), force=True) is False -# โ”€โ”€ 3. Live attention gate (needs a live gfx1151 + compiled flash lane) โ”€โ”€โ”€โ”€โ”€โ”€โ”€ +# โ”€โ”€ 3. Live gates (need a live gfx1151) โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ + +_C3_CHAINS = [ + F.FusedRegion(epilogue=("relu",)), + F.FusedRegion(epilogue=("bias", "gelu")), + F.FusedRegion(epilogue=("silu",)), + F.FusedRegion(epilogue=("bias",), reduction="softmax"), + F.FusedRegion(epilogue=(), reduction="rmsnorm"), + F.FusedRegion(epilogue=("relu",), reduction="layer_norm"), + F.FusedRegion(epilogue=("gelu",), prologue=("relu",)), +] + + +@pytest.mark.slow +@pytest.mark.skipif(not _rocm_hip_live(), + reason="live gfx1151 + hipcc required") +@pytest.mark.parametrize("region", _C3_CHAINS, + ids=lambda r: f"{r.epilogue}/{r.reduction}/{r.prologue}") +def test_live_rocm_generic_hip_gated(region): + # C3: the generically-synthesized HIP FusedRegion kernel compiles with hipcc, + # runs on gfx1151 ("rocm_hip"), matches numpy (f32), and passes the F4 oracle. + F.clear_verification_cache() + runner = get_runner("rocm") + rng = np.random.default_rng(0) + A = rng.standard_normal((8, 12)).astype(np.float32) + B = rng.standard_normal((12, 16)).astype(np.float32) + bias = rng.standard_normal((16,)).astype(np.float32) if region.has_bias else None + out, execution = runner.run_fused_region(region, A, B, bias) + assert execution == "rocm_hip" + assert np.allclose(out, region.reference(A, B, bias), atol=1e-3) + assert F.verify_synthesized_region(region, runner=runner, force=True) is True + @pytest.mark.slow @pytest.mark.skipif(not _rocm_flash_live(),