From 9102444e59929a919a24037501cd56d9a359033d Mon Sep 17 00:00:00 2001 From: gstoner Date: Thu, 25 Jun 2026 19:46:38 -0700 Subject: [PATCH] =?UTF-8?q?feat(s2):=20elementwise=20unary-math=20kernels?= =?UTF-8?q?=20(exp/log/sqrt/erf/=E2=80=A6)=20on=20gfx1151=20+=20AVX-512?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The S2 scalar-math / stability family gets a real optimized device kernel on both devices we have hardware for, pairing with the reduction family. ROCm (gfx1151): `tessera_rocm.unary` directive + `generate-rocm-unary-kernel` pass → a flat per-element kernel (one thread/element), the unary sibling of the activation lane. Covers exp/log/sqrt/rsqrt/reciprocal/abs(absolute)/sign/erf/ tanh/sigmoid/log1p/expm1/softplus (softplus stable: log1p(exp(-|x|))+max(x,0)); transcendentals lower through math → ROCDL. New runtime.launch() lane `rocm_unary_compiled`, dispatched by op name; f16/bf16/f32 storage, f32 compute. Validated on gfx1151 vs numpy across kind × dtype × shape (143 tests). x86 (AVX-512): `avx512_unary_f32.cpp` — the algebraic subset (sqrt/rsqrt/ reciprocal/abs/neg/sign) with direct intrinsics (no polynomial approx) + scalar reference + on-device test (31/31). The transcendentals stay numpy-reference on CPU (no fused x86 claim). Wiring: executor table + execution_matrix KNOWN_EXECUTORS/_MATRIX row + 13 _ROCM_COMPILED manifest entries (status `compiled`) + fixtures + ROCM_AUDIT Landed entry + regenerated dashboards (runtime_abi/execution_matrix/test_cov). Co-Authored-By: Claude Opus 4.8 --- docs/audit/backend/rocm/ROCM_AUDIT.md | 15 ++ docs/audit/generated/runtime_abi.csv | 2 + docs/audit/generated/runtime_abi.md | 4 +- .../generated/runtime_execution_matrix.csv | 1 + .../generated/runtime_execution_matrix.md | 2 + docs/audit/generated/test_coverage.csv | 30 +-- docs/audit/generated/test_coverage.md | 30 +-- python/tessera/compiler/backend_manifest.py | 17 ++ python/tessera/compiler/execution_matrix.py | 21 ++ python/tessera/runtime.py | 112 +++++++++ .../include/TesseraROCM/IR/TesseraROCMOps.td | 23 ++ .../include/TesseraROCM/Passes.h | 1 + .../lib/Conversion/CMakeLists.txt | 1 + .../Conversion/GenerateROCMUnaryKernel.cpp | 233 ++++++++++++++++++ .../lib/Conversion/Passes.cpp | 1 + .../tessera_x86_backend/CMakeLists.txt | 4 + .../src/kernels/avx512_unary_f32.cpp | 84 +++++++ .../tessera_x86_backend/tests/test_unary.cpp | 91 +++++++ tests/unit/test_rocm_unary_compiled.py | 148 +++++++++++ 19 files changed, 788 insertions(+), 32 deletions(-) create mode 100644 src/compiler/codegen/Tessera_ROCM_Backend/lib/Conversion/GenerateROCMUnaryKernel.cpp create mode 100644 src/compiler/codegen/tessera_x86_backend/src/kernels/avx512_unary_f32.cpp create mode 100644 src/compiler/codegen/tessera_x86_backend/tests/test_unary.cpp create mode 100644 tests/unit/test_rocm_unary_compiled.py diff --git a/docs/audit/backend/rocm/ROCM_AUDIT.md b/docs/audit/backend/rocm/ROCM_AUDIT.md index 5ba765053..9e6e08d53 100644 --- a/docs/audit/backend/rocm/ROCM_AUDIT.md +++ b/docs/audit/backend/rocm/ROCM_AUDIT.md @@ -341,6 +341,21 @@ become the on-silicon **oracle** the compiled path validates against. GPU-free codegen gate. The CPU half (AVX-512) landed in the x86 backend (`avx512_reduce_f32.cpp`) — so the reduction family now has a real optimized kernel on both devices we have hardware for. Status `compiled`. +- **Elementwise unary math** (S2 scalar-math / stability family, 2026-06-25): + a `tessera_rocm.unary` directive + `generate-rocm-unary-kernel` pass emitting + a flat per-element kernel (one thread per element), the unary sibling of the + activation lane. Covers `exp`/`log`/`sqrt`/`rsqrt`/`reciprocal`/`abs` + (`absolute`)/`sign`/`erf`/`tanh`/`sigmoid`/`log1p`/`expm1`/`softplus` + (softplus stable: `log1p(exp(-|x|)) + max(x,0)`); transcendentals lower + through the `math` → ROCDL path. New `runtime.launch()` lane + `rocm_unary_compiled`, dispatched by op name; f16/bf16/f32 storage, f32 + compute. Validated on gfx1151 vs numpy across kind × dtype × shape incl. + rank-3 (`test_rocm_unary_compiled.py`) + a GPU-free codegen gate. The CPU + half landed in the x86 backend as an AVX-512 kernel for the **algebraic + subset** (`sqrt`/`rsqrt`/`reciprocal`/`abs`/`neg`/`sign`, direct intrinsics, + no polynomial approx — `avx512_unary_f32.cpp`, validated standalone); the + transcendentals stay numpy-reference on CPU (no fused x86 claim). Status + `compiled`. - **rmsnorm / layer_norm** (2026-06-25): the row-reduction siblings of the softmax kernel — a `tessera_rocm.norm` directive + `generate-rocm-norm-kernel` pass (one workgroup per row). rmsnorm tree-reduces Σx² in one pass; layer_norm diff --git a/docs/audit/generated/runtime_abi.csv b/docs/audit/generated/runtime_abi.csv index 0d837b572..52258024d 100644 --- a/docs/audit/generated/runtime_abi.csv +++ b/docs/audit/generated/runtime_abi.csv @@ -578,8 +578,10 @@ x86,tessera_x86_amx_gemm_s8s8_s32,amx_gemm_s8s8_s32,,src/compiler/codegen/tesser x86,tessera_x86_amx_gemm_s8s8_s32,amx_gemm_s8s8_s32,,src/compiler/codegen/tessera_x86_backend/src/kernels/amx_gemm_int8.cpp x86,tessera_x86_avx512_gemm_bf16,avx512_gemm,bf16,src/compiler/codegen/tessera_x86_backend/src/kernels/avx512_gemm_bf16.cpp x86,tessera_x86_avx512_reduce_f32,avx512_reduce,f32,src/compiler/codegen/tessera_x86_backend/src/kernels/avx512_reduce_f32.cpp +x86,tessera_x86_avx512_unary_f32,avx512_unary,f32,src/compiler/codegen/tessera_x86_backend/src/kernels/avx512_unary_f32.cpp x86,tessera_x86_avx512_vnni_gemm_u8s8_s32,avx512_vnni_gemm_u8s8_s32,,src/compiler/codegen/tessera_x86_backend/src/kernels/avx512_vnni_gemm_int8.cpp x86,tessera_x86_epilogue_bias_fp32,epilogue_bias_fp32,,src/compiler/codegen/tessera_x86_backend/src/kernels/epilogue.cpp x86,tessera_x86_epilogue_bias_gelu_fp32,epilogue_bias_gelu_fp32,,src/compiler/codegen/tessera_x86_backend/src/kernels/epilogue.cpp x86,tessera_x86_reference_gemm_bf16,reference_gemm,bf16,src/compiler/codegen/tessera_x86_backend/src/kernels/avx512_gemm_bf16.cpp x86,tessera_x86_reference_reduce_f32,reference_reduce,f32,src/compiler/codegen/tessera_x86_backend/src/kernels/avx512_reduce_f32.cpp +x86,tessera_x86_reference_unary_f32,reference_unary,f32,src/compiler/codegen/tessera_x86_backend/src/kernels/avx512_unary_f32.cpp diff --git a/docs/audit/generated/runtime_abi.md b/docs/audit/generated/runtime_abi.md index cf8434328..d08251ae5 100644 --- a/docs/audit/generated/runtime_abi.md +++ b/docs/audit/generated/runtime_abi.md @@ -4,7 +4,7 @@ Human-readable view. The canonical machine-readable artifact is `runtime_abi.csv ## Headline -- **328** unique `extern "C" tessera_*` C ABI symbols across all backends. +- **330** unique `extern "C" tessera_*` C ABI symbols across all backends. - **6 / 6** core runtime headers present. - **134** Apple GPU kernel families with per-dtype variants. @@ -26,7 +26,7 @@ Human-readable view. The canonical machine-readable artifact is `runtime_abi.csv | `apple` | 304 | | `nvidia` | 4 | | `rocm` | 10 | -| `x86` | 10 | +| `x86` | 12 | ## Apple GPU kernel families × dtype matrix diff --git a/docs/audit/generated/runtime_execution_matrix.csv b/docs/audit/generated/runtime_execution_matrix.csv index 8e6abc3c5..7272092e2 100644 --- a/docs/audit/generated/runtime_execution_matrix.csv +++ b/docs/audit/generated/runtime_execution_matrix.csv @@ -19,4 +19,5 @@ rocm,rocm_reduce_compiled,native_gpu,1,rocm_reduce_compiled,success,hip_runtime, rocm,rocm_rope_compiled,native_gpu,1,rocm_rope_compiled,success,hip_runtime,"ROCm rope artifact runs the COMPILER-GENERATED interleaved-pair rotary-position-embedding kernel (one workgroup per row): tessera-opt generates + serializes the kernel to hsaco in-process, then HIP loads + launches it." rocm,rocm_silu_mul_compiled,native_gpu,1,rocm_silu_mul_compiled,success,hip_runtime,"ROCm silu_mul artifact runs the COMPILER-GENERATED flat 2-operand elementwise SwiGLU gate-multiply silu(a)·b (one thread per element): tessera-opt generates + serializes the kernel to hsaco in-process, then HIP loads + launches it." rocm,rocm_softmax_compiled,native_gpu,1,rocm_softmax_compiled,success,hip_runtime,"ROCm softmax artifact runs the COMPILER-GENERATED RDNA row-reduction kernel (stable softmax over the last axis, one workgroup per row, LDS tree-reduce): tessera-opt generates + serializes the kernel to hsaco in-process, then HIP loads + launches it. The first non-matmul/non-WMMA compiled ROCm kernel." +rocm,rocm_unary_compiled,native_gpu,1,rocm_unary_compiled,success,hip_runtime,"ROCm unary artifact runs the COMPILER-GENERATED flat elementwise unary-math kernel (S2 scalar-math/stability: exp/log/sqrt/rsqrt/reciprocal/abs/sign/erf/tanh/sigmoid/log1p/expm1/softplus, one thread per element): tessera-opt generates + serializes the kernel to hsaco in-process, then HIP loads + launches it. Dispatched by op name." rocm,rocm_wmma,native_gpu,1,rocm_wmma,success,hip_runtime,"ROCm matmul via the hand-written RDNA WMMA GEMM (tessera_rocm_wmma_gemm_{f16,bf16} C ABI symbol, HIPRTC-compiled for the device arch). Now the reference ORACLE + availability fallback for the compiled lane (rocm_compiled) — still directly selectable by stamping compiler_path=""rocm_wmma""." diff --git a/docs/audit/generated/runtime_execution_matrix.md b/docs/audit/generated/runtime_execution_matrix.md index 538a8e07c..f4facbb1e 100644 --- a/docs/audit/generated/runtime_execution_matrix.md +++ b/docs/audit/generated/runtime_execution_matrix.md @@ -33,6 +33,7 @@ Single source of truth for what `runtime.launch()` does with each `(target, comp | `rocm` | `rocm_rope_compiled` | `rocm_rope_compiled` | `native_gpu` | `hip_runtime` | ROCm rope artifact runs the COMPILER-GENERATED interleaved-pair rotary-position-embedding kernel (one workgroup per row): tessera-opt generates + serializes the kernel to hsaco in-process, then HIP loads + launches it. | | `rocm` | `rocm_silu_mul_compiled` | `rocm_silu_mul_compiled` | `native_gpu` | `hip_runtime` | ROCm silu_mul artifact runs the COMPILER-GENERATED flat 2-operand elementwise SwiGLU gate-multiply silu(a)·b (one thread per element): tessera-opt generates + serializes the kernel to hsaco in-process, then HIP loads + launches it. | | `rocm` | `rocm_softmax_compiled` | `rocm_softmax_compiled` | `native_gpu` | `hip_runtime` | ROCm softmax artifact runs the COMPILER-GENERATED RDNA row-reduction kernel (stable softmax over the last axis, one workgroup per row, LDS tree-reduce): tessera-opt generates + serializes the kernel to hsaco in-process, then HIP loads + launches it. The first non-matmul/non-WMMA compiled ROCm kernel. | +| `rocm` | `rocm_unary_compiled` | `rocm_unary_compiled` | `native_gpu` | `hip_runtime` | ROCm unary artifact runs the COMPILER-GENERATED flat elementwise unary-math kernel (S2 scalar-math/stability: exp/log/sqrt/rsqrt/reciprocal/abs/sign/erf/tanh/sigmoid/log1p/expm1/softplus, one thread per element): tessera-opt generates + serializes the kernel to hsaco in-process, then HIP loads + launches it. Dispatched by op name. | | `rocm` | `rocm_wmma` | `rocm_wmma` | `native_gpu` | `hip_runtime` | ROCm matmul via the hand-written RDNA WMMA GEMM (tessera_rocm_wmma_gemm_{f16,bf16} C ABI symbol, HIPRTC-compiled for the device arch). Now the reference ORACLE + availability fallback for the compiled lane (rocm_compiled) — still directly selectable by stamping compiler_path="rocm_wmma". | ## Targets with no executable row @@ -67,4 +68,5 @@ nvidia_sm80, nvidia_sm90, nvidia_sm100, rocm_gfx90a, rocm_gfx940, rocm_gfx942, r | `rocm_rope_compiled` | AMD GPU RDNA rotary-position-embedding the Tessera compiler GENERATES (generate-rocm-rope-kernel -> ROCDL -> hsaco, in-process via tessera-opt), then HIP loads + launches it. Interleaved-pair RoPE over [M, D] (one workgroup per row); f32/f16/bf16 | | `rocm_silu_mul_compiled` | AMD GPU RDNA SwiGLU gate-multiply the Tessera compiler GENERATES (generate-rocm-silu-mul-kernel -> ROCDL -> hsaco, in-process via tessera-opt), then HIP loads + launches it. Flat 2-operand elementwise silu(a)·b (one thread per element); the standalone analog of the fused SwiGLU gate-multiply; f32/f16/bf16 storage, f32 compute | | `rocm_softmax_compiled` | AMD GPU RDNA row-reduction softmax the Tessera compiler GENERATES (generate-rocm-softmax-kernel -> ROCDL -> hsaco, in-process via tessera-opt), then HIP loads + launches it. Stable softmax over the last axis (one workgroup per row, LDS tree-reduce); the first non-matmul/non-WMMA compiled ROCm kernel. f32/f16/bf16 storage, f32 reduce | +| `rocm_unary_compiled` | AMD GPU RDNA flat elementwise unary-math kernel the Tessera compiler GENERATES (generate-rocm-unary-kernel -> ROCDL -> hsaco, in-process via tessera-opt), then HIP loads + launches it — the S2 scalar-math / stability family (exp/log/sqrt/rsqrt/reciprocal/abs/sign/erf/tanh/sigmoid/log1p/expm1/softplus), one thread per element, dispatched by op name; f32/f16/bf16 storage, f32 compute | | `rocm_wmma` | AMD GPU RDNA WMMA matrix-core GEMM via the shipped libtessera_rocm_gemm.so tessera_rocm_wmma_gemm_{f16,bf16} C ABI symbol (HIPRTC-compiled for the device arch; f16/bf16 storage, f32 accumulate) | diff --git a/docs/audit/generated/test_coverage.csv b/docs/audit/generated/test_coverage.csv index 6880ced53..a4fcc1fab 100644 --- a/docs/audit/generated/test_coverage.csv +++ b/docs/audit/generated/test_coverage.csv @@ -1,6 +1,6 @@ op,python_refs,lit_refs,negative_refs,total_refs,is_thinly_tested,dtype_variants,bucket,reason -abs,2,0,0,2,0,,directly_tested,2 direct test references -absolute,1,0,0,1,1,,structural_only,unclassified — defaults to structural_only +abs,3,0,0,3,0,bf16 f32,directly_tested,3 direct test references +absolute,2,0,0,2,0,bf16 f32,directly_tested,2 direct test references acos,1,0,0,1,1,,covered_by_family,category default for 'elementwise' adafactor,4,0,0,4,0,fp32 fp64,directly_tested,4 direct test references adam,7,5,0,12,0,fp32 fp64,directly_tested,12 direct test references @@ -171,11 +171,11 @@ ema_update,1,0,0,1,1,,structural_only,category default for 'grad_transform' empty_state_tree,0,0,0,0,1,,structural_only,category default for 'state_tree' eq,0,0,0,0,1,,covered_by_family,category default for 'comparison' equiprob_band_partition,0,0,0,0,1,,structural_only,unclassified — defaults to structural_only -erf,1,0,0,1,1,,covered_by_family,category default for 'elementwise' +erf,2,0,0,2,0,bf16 f32,directly_tested,2 direct test references erfc,1,0,0,1,1,,covered_by_family,category default for 'elementwise' -exp,4,0,0,4,0,,directly_tested,4 direct test references +exp,5,0,0,5,0,bf16 f32,directly_tested,5 direct test references expand,1,0,0,1,1,,structural_only,unclassified — defaults to structural_only -expm1,1,0,0,1,1,,covered_by_family,category default for 'elementwise' +expm1,2,0,0,2,0,bf16 f32,directly_tested,2 direct test references factorized_matmul,4,0,0,4,0,,directly_tested,4 direct test references factorized_pos_emb,6,0,1,6,0,,directly_tested,6 direct test references fake_quantize,3,0,0,3,0,,directly_tested,3 direct test references @@ -247,8 +247,8 @@ lion,4,0,0,4,0,fp32 fp64,directly_tested,4 direct test references load_balance_loss,2,0,0,2,0,,directly_tested,2 direct test references load_sharded,3,0,1,3,0,,directly_tested,3 direct test references load_state,14,0,2,14,0,,directly_tested,14 direct test references -log,3,0,0,3,0,,directly_tested,3 direct test references -log1p,2,0,0,2,0,,directly_tested,2 direct test references +log,4,0,0,4,0,bf16 f32,directly_tested,4 direct test references +log1p,3,0,0,3,0,bf16 f32,directly_tested,3 direct test references log_cosh_loss,1,0,0,1,1,,covered_by_family,category default for 'loss' log_softmax,5,5,0,10,0,,directly_tested,10 direct test references logical_and,1,0,0,1,1,,structural_only,unclassified — defaults to structural_only @@ -345,7 +345,7 @@ quantize_int8,3,0,0,3,0,,directly_tested,3 direct test references quantize_nvfp4,11,0,2,11,0,fp16 fp32 fp4_e2m1 fp6_e2m3 fp8_e4m3 fp8_e5m2 int8 nvfp4,directly_tested,11 direct test references quantized_matmul,5,0,0,5,0,f16,directly_tested,5 direct test references rearrange,2,0,0,2,0,,directly_tested,2 direct test references -reciprocal,1,0,0,1,1,,structural_only,unclassified — defaults to structural_only +reciprocal,2,0,0,2,0,bf16 f32,directly_tested,2 direct test references reduce,86,0,7,86,0,f32 fp16 fp32 fp4_e2m1 fp6_e2m3 fp8_e4m3 nvfp4,directly_tested,86 direct test references reduce_scatter,0,4,0,4,0,,directly_tested,4 direct test references relu,105,20,9,125,0,bf16 f16 f32 f64 fp32 fp8_e4m3 int8,directly_tested,125 direct test references @@ -381,7 +381,7 @@ rope,19,10,0,29,0,bf16,directly_tested,29 direct test references rope_merge,2,0,1,2,0,fp16 fp32,directly_tested,2 direct test references rope_split,8,0,1,8,0,fp16 fp32,directly_tested,8 direct test references round,1,0,0,1,1,,structural_only,unclassified — defaults to structural_only -rsqrt,2,0,0,2,0,,directly_tested,2 direct test references +rsqrt,3,0,0,3,0,bf16 f32,directly_tested,3 direct test references safetensors_export,1,0,0,1,1,,structural_only,category default for 'aot' save_sharded,1,0,1,1,1,,structural_only,category default for 'serialization' save_state,6,0,2,6,0,,directly_tested,6 direct test references @@ -399,9 +399,9 @@ seq2seq_loss,3,0,0,3,0,,directly_tested,3 direct test references sgd,6,0,0,6,0,,directly_tested,6 direct test references shard_map,0,0,0,0,1,,structural_only,category default for 'sharding' sharded_dataset,2,0,0,2,0,,directly_tested,2 direct test references -sigmoid,4,0,0,4,0,bf16 f16 f32 fp32,directly_tested,4 direct test references +sigmoid,5,0,0,5,0,bf16 f16 f32 fp32,directly_tested,5 direct test references sigmoid_safe,4,0,0,4,0,,directly_tested,4 direct test references -sign,2,0,0,2,0,,directly_tested,2 direct test references +sign,4,0,1,4,0,bf16 f32,directly_tested,4 direct test references silu,84,2,5,86,0,bf16 f16 f32 f64 fp16 fp32 fp4_e2m1 fp6_e2m3 fp8_e4m3 nvfp4,directly_tested,86 direct test references silu_mul,14,13,0,27,0,bf16 fp32,directly_tested,27 direct test references simple_rnn_cell,5,0,0,5,0,,directly_tested,5 direct test references @@ -410,9 +410,9 @@ sinh,1,0,0,1,1,,covered_by_family,category default for 'elementwise' slice,11,0,0,11,0,,directly_tested,11 direct test references smooth_l1_loss,1,0,0,1,1,,covered_by_family,category default for 'loss' softcap,7,0,1,7,0,,directly_tested,7 direct test references -softmax,97,37,9,134,0,bf16 f16 f32 fp16 fp32 fp4_e2m1 fp6_e2m3 fp8_e4m3 nvfp4,directly_tested,134 direct test references +softmax,98,37,10,135,0,bf16 f16 f32 fp16 fp32 fp4_e2m1 fp6_e2m3 fp8_e4m3 nvfp4,directly_tested,135 direct test references softmax_safe,4,4,1,8,0,bf16,directly_tested,8 direct test references -softplus,2,0,0,2,0,,directly_tested,2 direct test references +softplus,3,0,0,3,0,bf16 f32,directly_tested,3 direct test references sort,3,0,0,3,0,,directly_tested,3 direct test references spectral_conv,5,0,0,5,0,fp32,directly_tested,5 direct test references spectral_filter,3,0,0,3,0,fp32,directly_tested,3 direct test references @@ -420,7 +420,7 @@ spectral_norm,2,0,0,2,0,,directly_tested,2 direct test references split,1,0,0,1,1,,structural_only,unclassified — defaults to structural_only spmm_coo,3,0,0,3,0,,directly_tested,3 direct test references spmm_csr,2,0,0,2,0,,directly_tested,2 direct test references -sqrt,4,0,0,4,0,,directly_tested,4 direct test references +sqrt,5,0,0,5,0,bf16 f32,directly_tested,5 direct test references squeeze,1,0,0,1,1,,structural_only,unclassified — defaults to structural_only stablehlo_export,2,0,0,2,0,,directly_tested,2 direct test references stack,2,0,0,2,0,,directly_tested,2 direct test references @@ -437,7 +437,7 @@ svd,8,5,0,13,0,bf16 f16 f32 fp16 fp32,directly_tested,13 direct test references switch,0,0,0,0,1,,structural_only,category default for 'control_flow' take,1,0,0,1,1,,structural_only,unclassified — defaults to structural_only tan,1,0,0,1,1,,covered_by_family,category default for 'elementwise' -tanh,4,0,1,4,0,bf16 f16 f32,directly_tested,4 direct test references +tanh,5,0,1,5,0,bf16 f16 f32,directly_tested,5 direct test references tile,1,0,0,1,1,,structural_only,unclassified — defaults to structural_only tile_view,2,0,0,2,0,,directly_tested,2 direct test references tiny_attention_conformance,0,0,0,0,1,,structural_only,category default for 'conformance' diff --git a/docs/audit/generated/test_coverage.md b/docs/audit/generated/test_coverage.md index f8340b682..1b193043a 100644 --- a/docs/audit/generated/test_coverage.md +++ b/docs/audit/generated/test_coverage.md @@ -7,11 +7,11 @@ Generated from `python/tessera/compiler/test_coverage_audit.py`. Don't edit by ## Headline - **474** ops in `primitive_coverage` registry. -- **2872** total Python-test references, **825** total lit-fixture references. +- **2888** total Python-test references, **825** total lit-fixture references. - **107** ops have **zero** references in either test surface. -- **197** ops have ≤1 reference ("thinly tested"). +- **193** ops have ≤1 reference ("thinly tested"). - **62** ops have ≥10 references ("well tested"). -- **84** ops have at least one associated `pytest.raises` negative test. +- **85** ops have at least one associated `pytest.raises` negative test. ## Top 20 most-tested ops @@ -19,7 +19,7 @@ Generated from `python/tessera/compiler/test_coverage_audit.py`. Don't edit by |----|--------:|---------:|------:|----:|--------| | `matmul` | 391 | 203 | 594 | 17 | `bf16`, `f16`, `f32`, `f64` … | | `flash_attn` | 100 | 57 | 157 | 8 | `bf16`, `f32`, `fp16`, `fp32` … | -| `softmax` | 97 | 37 | 134 | 9 | `bf16`, `f16`, `f32`, `fp16` … | +| `softmax` | 98 | 37 | 135 | 10 | `bf16`, `f16`, `f32`, `fp16` … | | `relu` | 105 | 20 | 125 | 9 | `bf16`, `f16`, `f32`, `f64` … | | `add` | 95 | 13 | 108 | 8 | `bf16`, `f16`, `f32`, `f64` … | | `mul` | 85 | 4 | 89 | 7 | `bf16`, `f16`, `f32`, `f64` … | @@ -40,11 +40,10 @@ Generated from `python/tessera/compiler/test_coverage_audit.py`. Don't edit by ## Thinly-tested ops (≤1 reference) -These **197** ops have at most one test reference across the whole test surface. Many will be legitimate — variant aliases, structural ops, or category rollups that inherit coverage from a parent family — but each one is a candidate for explicit per-op test coverage. +These **193** ops have at most one test reference across the whole test surface. Many will be legitimate — variant aliases, structural ops, or category rollups that inherit coverage from a parent family — but each one is a candidate for explicit per-op test coverage. | Op | py refs | lit refs | total | |----|--------:|---------:|------:| -| `absolute` | 1 | 0 | 1 | | `acos` | 1 | 0 | 1 | | `aot_export` | 0 | 0 | 0 | | `aot_load` | 0 | 0 | 0 | @@ -104,8 +103,9 @@ These **197** ops have at most one test reference across the whole test surface. | `dataset_map` | 0 | 0 | 0 | | `dataset_prefetch` | 0 | 0 | 0 | | `dataset_repeat` | 0 | 0 | 0 | +| `dataset_shuffle` | 0 | 0 | 0 | -_(137 additional thinly-tested ops omitted; see `collect_op_test_coverage()` for the full list.)_ +_(133 additional thinly-tested ops omitted; see `collect_op_test_coverage()` for the full list.)_ --- @@ -117,12 +117,12 @@ Companion to `test_coverage_by_op.md`. That dashboard says **which** ops are th ## Headline -**197** ops have ≤1 direct test reference. They break down as: +**193** ops have ≤1 direct test reference. They break down as: | Bucket | Count | Meaning | |--------|------:|---------| -| `covered_by_family` | 75 | Tested via a parent op or family wrapper | -| `structural_only` | 116 | Registry/metadata/wrapper; no direct numerical test meaningful | +| `covered_by_family` | 73 | Tested via a parent op or family wrapper | +| `structural_only` | 114 | Registry/metadata/wrapper; no direct numerical test meaningful | | `needs_direct_test` | 2 | **Actionable test debt** — real primitive without direct test | | `hardware_gated` | 4 | Blocked on real device hardware (Phase G/H/I) | | `deprecated_or_internal` | 0 | Not public test debt | @@ -147,7 +147,7 @@ These **4** ops need real device hardware (Phase G/H/I). They cannot be tested | `ebm_sphere_langevin_sample` | manifold Langevin needs real GPU mesh (Phase G) | | `ebm_sphere_langevin_step` | manifold Langevin needs real GPU mesh (Phase G) | -## `covered_by_family` — 75 ops +## `covered_by_family` — 73 ops Tested through a parent op or family wrapper. Sample (first 30): @@ -184,15 +184,14 @@ Tested through a parent op or family wrapper. Sample (first 30): | `denoising_score_matching_loss` | category default for 'loss' | | `dz` | exercised by complex differential tests | -_(45 additional family-covered ops omitted; see `classify_thinly_tested()` for the full list.)_ +_(43 additional family-covered ops omitted; see `classify_thinly_tested()` for the full list.)_ -## `structural_only` — 116 ops +## `structural_only` — 114 ops Registry/metadata/wrapper ops; direct numerical tests not meaningful. Sample (first 30): | Op | reason | |----|--------| -| `absolute` | unclassified — defaults to structural_only | | `aot_export` | category default for 'aot' | | `aot_load` | category default for 'aot' | | `associative_scan` | category default for 'control_flow' | @@ -222,5 +221,6 @@ Registry/metadata/wrapper ops; direct numerical tests not meaningful. Sample (f | `dataset_batch` | category default for 'data' | | `dataset_checkpoint` | category default for 'data' | | `dataset_filter` | category default for 'data' | +| `dataset_interleave` | category default for 'data' | -_(86 additional structural ops omitted.)_ +_(84 additional structural ops omitted.)_ diff --git a/python/tessera/compiler/backend_manifest.py b/python/tessera/compiler/backend_manifest.py index ac8c07b68..f5feaae3a 100644 --- a/python/tessera/compiler/backend_manifest.py +++ b/python/tessera/compiler/backend_manifest.py @@ -953,6 +953,19 @@ def primitive_is_complete(entries: tuple["BackendKernelEntry", ...]) -> bool: "of the fused SwiGLU gate-multiply. Executes via " "runtime.launch() (rocm_silu_mul_compiled).", }, + # S2 scalar-math / stability family — flat per-element unary math kernel + # (generate-rocm-unary-kernel), the unary sibling of the activation lane. + # Executes via runtime.launch() (rocm_unary_compiled). f32/f16/bf16, f32 + # compute. The transcendentals lower through the math → ROCDL path. + **{op: { + "dtypes": ("fp32", "fp16", "bf16"), + "feature_flags": ("elementwise",), + "notes": f"Standalone elementwise {op} — flat per-element unary-math " + "kernel (generate-rocm-unary-kernel). Executes via " + "runtime.launch() (rocm_unary_compiled).", + } for op in ("exp", "log", "sqrt", "rsqrt", "reciprocal", "absolute", + "sign", "erf", "tanh", "sigmoid", "log1p", "expm1", + "softplus")}, "rope": { "dtypes": ("fp32", "fp16", "bf16"), "feature_flags": ("elementwise",), @@ -1149,6 +1162,10 @@ def primitive_is_complete(entries: tuple["BackendKernelEntry", ...]) -> bool: ("gelu", "rocm"): "tests/unit/test_rocm_activation_compiled.py", ("silu", "rocm"): "tests/unit/test_rocm_activation_compiled.py", ("silu_mul", "rocm"): "tests/unit/test_rocm_silu_mul_compiled.py", + **{(op, "rocm"): "tests/unit/test_rocm_unary_compiled.py" + for op in ("exp", "log", "sqrt", "rsqrt", "reciprocal", "absolute", + "sign", "erf", "tanh", "sigmoid", "log1p", "expm1", + "softplus")}, ("rope", "rocm"): "tests/unit/test_rocm_rope_compiled.py", ("alibi", "rocm"): "tests/unit/test_rocm_alibi_compiled.py", ("batched_gemm", "rocm"): "tests/unit/test_rocm_matmul_family_compiled.py", diff --git a/python/tessera/compiler/execution_matrix.py b/python/tessera/compiler/execution_matrix.py index 15a27a4d8..f03750d44 100644 --- a/python/tessera/compiler/execution_matrix.py +++ b/python/tessera/compiler/execution_matrix.py @@ -127,6 +127,14 @@ class ExecutionRow: "launches it. Standalone gelu / silu / relu (one " "thread per element), dispatched by op name; " "f32/f16/bf16 storage, f32 compute", + "rocm_unary_compiled": "AMD GPU RDNA flat elementwise unary-math kernel the " + "Tessera compiler GENERATES (generate-rocm-unary-" + "kernel -> ROCDL -> hsaco, in-process via tessera-" + "opt), then HIP loads + launches it — the S2 scalar-" + "math / stability family (exp/log/sqrt/rsqrt/" + "reciprocal/abs/sign/erf/tanh/sigmoid/log1p/expm1/" + "softplus), one thread per element, dispatched by op " + "name; f32/f16/bf16 storage, f32 compute", "rocm_silu_mul_compiled": "AMD GPU RDNA SwiGLU gate-multiply the Tessera " "compiler GENERATES (generate-rocm-silu-mul-kernel " "-> ROCDL -> hsaco, in-process via tessera-opt), then " @@ -353,6 +361,19 @@ class ExecutionRow: "element): tessera-opt generates + serializes the kernel to hsaco " "in-process, then HIP loads + launches it. Dispatched by op name.", execution_mode="hip_runtime"), + # Standalone elementwise unary math (exp/log/sqrt/erf/…) — the S2 scalar-math + # / stability family, flat per-element; the unary sibling of activation. + ("rocm", "rocm_unary_compiled"): ExecutionRow( + target="rocm", compiler_path="rocm_unary_compiled", + execution_kind="native_gpu", executable=True, + executor_id="rocm_unary_compiled", runtime_status="success", + reason="ROCm unary artifact runs the COMPILER-GENERATED flat elementwise " + "unary-math kernel (S2 scalar-math/stability: exp/log/sqrt/rsqrt/" + "reciprocal/abs/sign/erf/tanh/sigmoid/log1p/expm1/softplus, one " + "thread per element): tessera-opt generates + serializes the " + "kernel to hsaco in-process, then HIP loads + launches it. " + "Dispatched by op name.", + execution_mode="hip_runtime"), # SwiGLU gate-multiply silu(a)·b — flat 2-operand elementwise. vs numpy. ("rocm", "rocm_silu_mul_compiled"): ExecutionRow( target="rocm", compiler_path="rocm_silu_mul_compiled", diff --git a/python/tessera/runtime.py b/python/tessera/runtime.py index c2a5b30b8..767cc9a19 100644 --- a/python/tessera/runtime.py +++ b/python/tessera/runtime.py @@ -3108,6 +3108,117 @@ def _mr(p, size): return o.reshape(x.shape) +# ───────────────────────────────────────────────────────────────────────────── +# ROCm COMPILED unary-math lane (2026-06-25) — S2 scalar-math / stability family +# A flat per-element kernel applying a pointwise unary math fn (exp/log/sqrt/erf/ +# …), the unary sibling of the activation lane. ``compiler_path = +# "rocm_unary_compiled"``. f32/f16/bf16 storage, f32 compute. ONE hsaco per +# (chip, kind, dtype), cached. +# ───────────────────────────────────────────────────────────────────────────── +_rocm_unary_hsaco_cache: dict[tuple[str, str, str], bytes] = {} + +#: op_name → unary kernel kind (the codegen `kind` attr). +_ROCM_UNARY_OPS: dict[str, str] = { + "tessera.exp": "exp", "tessera.log": "log", "tessera.sqrt": "sqrt", + "tessera.rsqrt": "rsqrt", "tessera.reciprocal": "reciprocal", + "tessera.absolute": "abs", "tessera.abs": "abs", "tessera.sign": "sign", + "tessera.erf": "erf", "tessera.tanh": "tanh", "tessera.sigmoid": "sigmoid", + "tessera.log1p": "log1p", "tessera.expm1": "expm1", + "tessera.softplus": "softplus", +} + + +def _build_compiled_unary_hsaco(kind: str, dtype: str = "f32") -> bytes: + chip = _rocm_chip() + directive = ( + 'module {\n' + f' "tessera_rocm.unary"() {{name = "u", kind = "{kind}", ' + f'dtype = "{dtype}"}} : () -> ()\n}}\n') + return _build_rocm_elementwise_hsaco( + "generate-rocm-unary-kernel", directive, _rocm_unary_hsaco_cache, + (chip, kind, dtype)) + + +def _execute_rocm_compiled_unary(artifact: RuntimeArtifact, args: Any) -> Any: + """The ``target="rocm"`` unary-math lane: run the COMPILER-GENERATED flat + elementwise kernel for a standalone unary math fn over any-shape input (one + thread per element). f32/f16/bf16 storage, f32 compute.""" + import numpy as np + + metadata = artifact.metadata or {} + arg_names = list(metadata.get("arg_names") or []) + ops = list(metadata.get("ops") or []) + op_name = str(ops[0].get("op_name", "")) if len(ops) == 1 else "" + if len(ops) != 1 or op_name not in _ROCM_UNARY_OPS: + raise ValueError( + "rocm_unary_compiled executor handles exactly one of " + f"{tuple(_ROCM_UNARY_OPS)}; got {[o.get('op_name') for o in ops]!r}") + kind = _ROCM_UNARY_OPS[op_name] + op = ops[0] + operand_names = [str(n) for n in op.get("operands", [])] + if len(operand_names) < 1: + raise ValueError("unary requires one operand") + values = _bind_launch_args(args, arg_names) + x = _as_numpy(values[operand_names[0]]) + n = int(np.prod(x.shape)) if x.ndim else 1 + if n <= 0: + return np.array(x, copy=True) + + store: Any + if x.dtype == np.float32: + dtype_tag, store, esz = "f32", np.float32, 4 + elif x.dtype == np.float16: + dtype_tag, store, esz = "f16", np.float16, 2 + else: + bf16 = _bfloat16_dtype() + if bf16 is not None and x.dtype == bf16: + dtype_tag, store, esz = "bf16", bf16, 2 + else: + raise ValueError(f"rocm unary lane handles f32/f16/bf16; got {x.dtype}") + + hsaco = _build_compiled_unary_hsaco(kind, dtype_tag) + hip = _load_hip_for_launch() + if hip is None: + raise _RocmCompiledUnavailable("libamdhip64.so not loadable") + if hip.hipInit(0) != 0: + raise _RocmCompiledUnavailable("rocm unary: hipInit failed") + mod = ctypes.c_void_p() + if hip.hipModuleLoadData(ctypes.byref(mod), hsaco) != 0: + raise _RocmCompiledUnavailable("rocm unary: no usable AMD GPU") + fn = ctypes.c_void_p() + if hip.hipModuleGetFunction(ctypes.byref(fn), mod, b"u") != 0: + raise RuntimeError("rocm unary: kernel symbol 'u' not found") + + xc = np.ascontiguousarray(x, dtype=store).reshape(-1) + o = np.zeros(n, dtype=store) + dx, do = ctypes.c_void_p(), ctypes.c_void_p() + for dev in (dx, do): + if hip.hipMalloc(ctypes.byref(dev), esz * n) != 0: + raise RuntimeError("rocm unary: hipMalloc failed") + hip.hipMemcpy(dx, xc.ctypes.data_as(ctypes.c_void_p), esz * n, 1) + + def _mr(p, size): + return [ctypes.c_void_p(p.value), ctypes.c_void_p(p.value), + ctypes.c_int64(0), ctypes.c_int64(size), ctypes.c_int64(1)] + + launch_args = _mr(dx, n) + _mr(do, n) + [ctypes.c_int64(n)] + arr = (ctypes.c_void_p * len(launch_args))() + for i, val in enumerate(launch_args): + arr[i] = ctypes.cast(ctypes.byref(val), ctypes.c_void_p) + gx = (n + _GRID_BLOCKDIM - 1) // _GRID_BLOCKDIM + rc = hip.hipModuleLaunchKernel(fn, gx, 1, 1, _GRID_BLOCKDIM, 1, 1, + 0, None, arr, None) + if rc != 0: + for dev in (dx, do): + hip.hipFree(dev) + raise RuntimeError(f"rocm unary: kernel launch failed rc={rc}") + hip.hipDeviceSynchronize() + hip.hipMemcpy(o.ctypes.data_as(ctypes.c_void_p), do, esz * n, 2) + for dev in (dx, do): + hip.hipFree(dev) + return o.reshape(x.shape) + + # ───────────────────────────────────────────────────────────────────────────── # ROCm COMPILED silu_mul lane (2026-06-25) — SwiGLU gate-multiply # silu_mul(a, b) = silu(a)·b, a flat 2-operand elementwise kernel (sibling of the @@ -3760,6 +3871,7 @@ def _executor_table(): "rocm_norm_compiled": _execute_rocm_compiled_norm, "rocm_reduce_compiled": _execute_rocm_compiled_reduce, "rocm_activation_compiled": _execute_rocm_compiled_activation, + "rocm_unary_compiled": _execute_rocm_compiled_unary, "rocm_silu_mul_compiled": _execute_rocm_compiled_silu_mul, "rocm_alibi_compiled": _execute_rocm_compiled_alibi, "rocm_matmul_family_compiled": _execute_rocm_compiled_matmul_family, diff --git a/src/compiler/codegen/Tessera_ROCM_Backend/include/TesseraROCM/IR/TesseraROCMOps.td b/src/compiler/codegen/Tessera_ROCM_Backend/include/TesseraROCM/IR/TesseraROCMOps.td index 82ce10bcb..87d1dbcdc 100644 --- a/src/compiler/codegen/Tessera_ROCM_Backend/include/TesseraROCM/IR/TesseraROCMOps.td +++ b/src/compiler/codegen/Tessera_ROCM_Backend/include/TesseraROCM/IR/TesseraROCMOps.td @@ -422,6 +422,29 @@ def ROCM_ReduceKernelOp : TesseraROCM_Op<"reduce"> { let assemblyFormat = "attr-dict"; } +def ROCM_UnaryKernelOp : TesseraROCM_Op<"unary"> { + let summary = "Elementwise unary math kernel (exp / log / sqrt / erf / …)"; + let description = [{ + The directive `generate-rocm-unary-kernel` expands into a flat elementwise + gpu kernel applying a pointwise unary math function over N elements (one + thread per element, strided grid) — the standalone S2 scalar-math / + stability family, the unary sibling of the activation lane. All kinds + compute in **f32** regardless of storage dtype; the transcendentals lower + through the `math` → ROCDL path: + + * algebraic: `sqrt`, `rsqrt`, `reciprocal`, `abs`, `neg`, `sign` + * transcend.: `exp`, `log`, `erf`, `tanh`, `sigmoid` + * stability: `log1p`, `expm1`, `softplus` + + `N` is a runtime `index` arg; `dtype` is the f16/bf16/f32 storage. + Validated vs numpy on gfx1151. + }]; + let arguments = (ins StrAttr:$name, + DefaultValuedAttr:$kind, + DefaultValuedAttr:$dtype); + let assemblyFormat = "attr-dict"; +} + def ROCM_BufferLoadOp : TesseraROCM_Op<"buffer_load", [MemoryEffects<[MemRead]>]> { let summary = "AMD buffer load with native out-of-bounds handling"; let description = [{ diff --git a/src/compiler/codegen/Tessera_ROCM_Backend/include/TesseraROCM/Passes.h b/src/compiler/codegen/Tessera_ROCM_Backend/include/TesseraROCM/Passes.h index 9a5610261..7c0ed1d3a 100644 --- a/src/compiler/codegen/Tessera_ROCM_Backend/include/TesseraROCM/Passes.h +++ b/src/compiler/codegen/Tessera_ROCM_Backend/include/TesseraROCM/Passes.h @@ -24,6 +24,7 @@ std::unique_ptr createGenerateROCMRopeKernelPass(); std::unique_ptr createGenerateROCMSoftmaxKernelPass(); std::unique_ptr createGenerateROCMNormKernelPass(); std::unique_ptr createGenerateROCMReduceKernelPass(); +std::unique_ptr createGenerateROCMUnaryKernelPass(); std::unique_ptr createLowerROCMAsyncCopyToLoopPass(); void buildTesseraROCMBackendPipeline(mlir::OpPassManager &pm); void registerTesseraROCMPasses(); diff --git a/src/compiler/codegen/Tessera_ROCM_Backend/lib/Conversion/CMakeLists.txt b/src/compiler/codegen/Tessera_ROCM_Backend/lib/Conversion/CMakeLists.txt index c3b610308..f0a639ce0 100644 --- a/src/compiler/codegen/Tessera_ROCM_Backend/lib/Conversion/CMakeLists.txt +++ b/src/compiler/codegen/Tessera_ROCM_Backend/lib/Conversion/CMakeLists.txt @@ -15,6 +15,7 @@ add_mlir_library(TesseraROCMConversion GenerateROCMSoftmaxKernel.cpp GenerateROCMNormKernel.cpp GenerateROCMReduceKernel.cpp + GenerateROCMUnaryKernel.cpp LowerROCMAsyncCopyToLoop.cpp Passes.cpp LINK_LIBS PUBLIC diff --git a/src/compiler/codegen/Tessera_ROCM_Backend/lib/Conversion/GenerateROCMUnaryKernel.cpp b/src/compiler/codegen/Tessera_ROCM_Backend/lib/Conversion/GenerateROCMUnaryKernel.cpp new file mode 100644 index 000000000..1078ad858 --- /dev/null +++ b/src/compiler/codegen/Tessera_ROCM_Backend/lib/Conversion/GenerateROCMUnaryKernel.cpp @@ -0,0 +1,233 @@ +//===- GenerateROCMUnaryKernel.cpp - elementwise unary math kernel -------===// +// +// Expands a `tessera_rocm.unary` directive into a flat elementwise gpu kernel +// applying a pointwise unary math function over N elements (one thread per +// element, strided grid) — the standalone S2 scalar-math / stability family, +// the unary sibling of the activation lane: +// +// algebraic: sqrt, rsqrt, reciprocal, abs, neg, sign +// transcend.: exp, log, erf, tanh, sigmoid +// stability: log1p, expm1, softplus (softplus stable: log1p(exp(-|x|))+max(x,0)) +// +// Computes in f32 regardless of storage dtype; the transcendentals lower +// through convert-math-to-rocdl. N is a runtime index arg. Validated vs a +// numpy reference on gfx1151. +//===----------------------------------------------------------------------===// + +#include "TesseraROCM/Passes.h" + +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/Dialect/GPU/IR/GPUDialect.h" +#include "mlir/Dialect/Math/IR/Math.h" +#include "mlir/Dialect/MemRef/IR/MemRef.h" +#include "mlir/Dialect/SCF/IR/SCF.h" +#include "mlir/IR/Builders.h" +#include "mlir/IR/BuiltinOps.h" +#include "mlir/IR/BuiltinTypes.h" +#include "mlir/Pass/Pass.h" +#include "llvm/ADT/StringSet.h" +#include "llvm/ADT/StringSwitch.h" + +using namespace mlir; + +namespace { + +static constexpr int64_t BD = 256; + +enum class Un { + Exp, Log, Sqrt, Rsqrt, Reciprocal, Abs, Neg, Sign, + Erf, Tanh, Sigmoid, Log1p, Expm1, Softplus +}; + +static Value cst(OpBuilder &b, Location loc, Type f32, float v) { + return b.create(loc, f32, b.getF32FloatAttr(v)); +} + +void emitUnaryBody(OpBuilder &b, Location loc, gpu::GPUFuncOp f, Type storeTy, + Un un) { + Type f32 = b.getF32Type(); + bool isF32 = storeTy.isF32(); + auto slt = arith::CmpIPredicate::slt; + auto ogt = arith::CmpFPredicate::OGT; + auto olt = arith::CmpFPredicate::OLT; + + b.setInsertionPointToStart(&f.getBody().front()); + Value X = f.getArgument(0), O = f.getArgument(1), N = f.getArgument(2); + + Value bid = b.create(loc, gpu::Dimension::x); + Value tid = b.create(loc, gpu::Dimension::x); + Value cBD = b.create(loc, BD); + Value gid = b.create(loc, b.create(loc, bid, cBD), + tid); + Value inb = b.create(loc, slt, gid, N); + auto ifo = b.create(loc, inb, /*withElse=*/false); + b.setInsertionPointToStart(ifo.thenBlock()); + + Value raw = b.create(loc, X, ValueRange{gid}); + Value x = isF32 ? raw : b.create(loc, f32, raw); + Value one = cst(b, loc, f32, 1.0f); + Value zero = cst(b, loc, f32, 0.0f); + Value y; + switch (un) { + case Un::Exp: + y = b.create(loc, x); + break; + case Un::Log: + y = b.create(loc, x); + break; + case Un::Sqrt: + y = b.create(loc, x); + break; + case Un::Rsqrt: + y = b.create(loc, x); + break; + case Un::Reciprocal: + y = b.create(loc, one, x); + break; + case Un::Abs: + y = b.create(loc, x); + break; + case Un::Neg: + y = b.create(loc, x); + break; + case Un::Sign: { + Value pos = b.create(loc, ogt, x, zero); + Value neg = b.create(loc, olt, x, zero); + Value negOne = cst(b, loc, f32, -1.0f); + Value s = b.create(loc, neg, negOne, zero); + y = b.create(loc, pos, one, s); + break; + } + case Un::Erf: + y = b.create(loc, x); + break; + case Un::Tanh: + y = b.create(loc, x); + break; + case Un::Sigmoid: { + Value e = b.create(loc, b.create(loc, x)); + y = b.create(loc, one, + b.create(loc, one, e)); + break; + } + case Un::Log1p: + y = b.create(loc, x); + break; + case Un::Expm1: + y = b.create(loc, x); + break; + case Un::Softplus: { + // Stable: log1p(exp(-|x|)) + max(x, 0) + Value ax = b.create(loc, x); + Value e = b.create(loc, b.create(loc, ax)); + Value lp = b.create(loc, e); + Value mx = b.create(loc, x, zero); + y = b.create(loc, lp, mx); + break; + } + } + Value sv = isF32 ? y : b.create(loc, storeTy, y); + b.create(loc, sv, O, ValueRange{gid}); + + b.setInsertionPointToEnd(&f.getBody().front()); + b.create(loc); +} + +struct GenerateROCMUnaryKernelPass + : PassWrapper> { + MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(GenerateROCMUnaryKernelPass) + + StringRef getArgument() const final { return "generate-rocm-unary-kernel"; } + StringRef getDescription() const final { + return "Expand a tessera_rocm.unary directive into a flat elementwise unary " + "math gpu kernel (compiler-generated)"; + } + void getDependentDialects(DialectRegistry ®istry) const final { + registry.insert(); + } + + void runOnOperation() override { + ModuleOp module = getOperation(); + SmallVector directives; + module.walk([&](Operation *op) { + if (op->getName().getStringRef() == "tessera_rocm.unary") + directives.push_back(op); + }); + for (Operation *op : directives) { + auto nameAttr = op->getAttrOfType("name"); + if (!nameAttr) { + op->emitError("tessera_rocm.unary missing name"); + return signalPassFailure(); + } + StringRef kindStr = "exp"; + if (auto a = op->getAttrOfType("kind")) + kindStr = a.getValue(); + Un un = llvm::StringSwitch(kindStr) + .Case("exp", Un::Exp) + .Case("log", Un::Log) + .Case("sqrt", Un::Sqrt) + .Case("rsqrt", Un::Rsqrt) + .Case("reciprocal", Un::Reciprocal) + .Case("abs", Un::Abs) + .Case("neg", Un::Neg) + .Case("sign", Un::Sign) + .Case("erf", Un::Erf) + .Case("tanh", Un::Tanh) + .Case("sigmoid", Un::Sigmoid) + .Case("log1p", Un::Log1p) + .Case("expm1", Un::Expm1) + .Case("softplus", Un::Softplus) + .Default(Un::Exp); + static const llvm::StringSet<> kValid = { + "exp", "log", "sqrt", "rsqrt", "reciprocal", "abs", "neg", + "sign", "erf", "tanh", "sigmoid", "log1p", "expm1", + "softplus"}; + if (!kValid.contains(kindStr)) { + op->emitError("generate-rocm-unary-kernel: unknown kind '") + << kindStr << "' (exp/log/sqrt/rsqrt/reciprocal/abs/neg/sign/erf/" + "tanh/sigmoid/log1p/expm1/softplus)"; + return signalPassFailure(); + } + OpBuilder b(module.getBodyRegion()); + b.setInsertionPointToEnd(module.getBody()); + Location loc = op->getLoc(); + std::string kname = nameAttr.getValue().str(); + + Type storeTy = b.getF32Type(); + if (auto a = op->getAttrOfType("dtype")) { + StringRef dt = a.getValue(); + if (dt == "f16" || dt == "float16") + storeTy = b.getF16Type(); + else if (dt == "bf16" || dt == "bfloat16") + storeTy = b.getBF16Type(); + else if (dt != "f32" && dt != "float32") { + op->emitError("generate-rocm-unary-kernel: dtype must be f32, f16, " + "or bf16 (got '") + << dt << "')"; + return signalPassFailure(); + } + } + + auto gpuMod = b.create(loc, kname + "_mod"); + b.setInsertionPointToStart(&gpuMod.getBodyRegion().front()); + Type idxTy = b.getIndexType(); + auto memTy = MemRefType::get({ShapedType::kDynamic}, storeTy); + // (X, O : memref, N : index) + auto fnTy = b.getFunctionType({memTy, memTy, idxTy}, {}); + auto gpuFunc = b.create(loc, kname, fnTy); + gpuFunc->setAttr(gpu::GPUDialect::getKernelFuncAttrName(), + b.getUnitAttr()); + OpBuilder body(gpuFunc.getContext()); + emitUnaryBody(body, loc, gpuFunc, storeTy, un); + op->erase(); + } + } +}; + +} // namespace + +std::unique_ptr +mlir::tessera_rocm::createGenerateROCMUnaryKernelPass() { + return std::make_unique(); +} diff --git a/src/compiler/codegen/Tessera_ROCM_Backend/lib/Conversion/Passes.cpp b/src/compiler/codegen/Tessera_ROCM_Backend/lib/Conversion/Passes.cpp index eaf10c511..279e07aaa 100644 --- a/src/compiler/codegen/Tessera_ROCM_Backend/lib/Conversion/Passes.cpp +++ b/src/compiler/codegen/Tessera_ROCM_Backend/lib/Conversion/Passes.cpp @@ -44,6 +44,7 @@ void registerTesseraROCMPasses() { registerPass([]() { return createGenerateROCMSoftmaxKernelPass(); }); registerPass([]() { return createGenerateROCMNormKernelPass(); }); registerPass([]() { return createGenerateROCMReduceKernelPass(); }); + registerPass([]() { return createGenerateROCMUnaryKernelPass(); }); registerPass([]() { return createLowerROCMAsyncCopyToLoopPass(); }); PassPipelineRegistration<> pipeline( "tessera-rocm-backend", diff --git a/src/compiler/codegen/tessera_x86_backend/CMakeLists.txt b/src/compiler/codegen/tessera_x86_backend/CMakeLists.txt index 5fed0a090..c245766b4 100644 --- a/src/compiler/codegen/tessera_x86_backend/CMakeLists.txt +++ b/src/compiler/codegen/tessera_x86_backend/CMakeLists.txt @@ -40,6 +40,7 @@ add_library(tessera_x86_backend src/kernels/amx_gemm_int8.cpp src/kernels/avx512_vnni_gemm_int8.cpp src/kernels/avx512_reduce_f32.cpp + src/kernels/avx512_unary_f32.cpp src/kernels/epilogue.cpp src/runtime/amx_runtime.cpp ) @@ -57,3 +58,6 @@ target_link_libraries(test_gemm_reference_tails PRIVATE tessera_x86_backend) add_executable(test_reduce tests/test_reduce.cpp) target_link_libraries(test_reduce PRIVATE tessera_x86_backend) + +add_executable(test_unary tests/test_unary.cpp) +target_link_libraries(test_unary PRIVATE tessera_x86_backend) diff --git a/src/compiler/codegen/tessera_x86_backend/src/kernels/avx512_unary_f32.cpp b/src/compiler/codegen/tessera_x86_backend/src/kernels/avx512_unary_f32.cpp new file mode 100644 index 000000000..9a060323e --- /dev/null +++ b/src/compiler/codegen/tessera_x86_backend/src/kernels/avx512_unary_f32.cpp @@ -0,0 +1,84 @@ +// AVX-512 elementwise unary math kernels (f32) for the Tessera x86 backend. +// +// Applies a pointwise unary math fn over a flat [n] f32 buffer, producing +// out[n]. This is the optimized CPU lane for the algebraic subset of the +// S2 scalar-math / numeric-helper family — the AVX-512 analog of the ROCm +// `generate-rocm-unary-kernel` lane, so these primitives get a REAL vectorized +// CPU kernel rather than only the numpy reference. A scalar reference is +// provided alongside for on-device validation (the test compares the two + a +// hand-computed expectation). +// +// Covered here are the ops with a DIRECT AVX-512 intrinsic (no polynomial +// approximation required): +// +// kind 0 = sqrt _mm512_sqrt_ps +// kind 1 = rsqrt 1 / sqrt(x) (full-precision div, not rsqrt14 approx) +// kind 2 = reciprocal 1 / x +// kind 3 = abs _mm512_abs_ps +// kind 4 = neg 0 - x +// kind 5 = sign (x>0) - (x<0), sign(0)=0 (NaN -> 0, finite-domain use) +// +// The transcendentals (exp/log/erf/tanh/…) lower through the ROCm math->ROCDL +// path on GPU; on CPU they remain the numpy reference for now (no fused x86 +// claim). 16 f32 lanes per __m512; the tail (n % 16) is handled scalar. + +#include +#include +#include + +namespace { +constexpr int kSqrt = 0; +constexpr int kRsqrt = 1; +constexpr int kRecip = 2; +constexpr int kAbs = 3; +constexpr int kNeg = 4; +constexpr int kSign = 5; + +inline float scalar_unary(float v, int kind) { + switch (kind) { + case kSqrt: return std::sqrt(v); + case kRsqrt: return 1.0f / std::sqrt(v); + case kRecip: return 1.0f / v; + case kAbs: return std::fabs(v); + case kNeg: return -v; + case kSign: return (v > 0.0f) ? 1.0f : (v < 0.0f ? -1.0f : 0.0f); + default: return v; + } +} +} // namespace + +extern "C" void tessera_x86_reference_unary_f32(const float* X, int64_t n, + float* out, int kind) { + for (int64_t i = 0; i < n; ++i) out[i] = scalar_unary(X[i], kind); +} + +extern "C" void tessera_x86_avx512_unary_f32(const float* X, int64_t n, + float* out, int kind) { + const int64_t vstep = 16; // f32 lanes per __m512 + const __m512 one = _mm512_set1_ps(1.0f); + const __m512 zero = _mm512_setzero_ps(); + const __m512 negOne = _mm512_set1_ps(-1.0f); + int64_t i = 0; + for (; i + vstep <= n; i += vstep) { + __m512 v = _mm512_loadu_ps(X + i); + __m512 y; + switch (kind) { + case kSqrt: y = _mm512_sqrt_ps(v); break; + case kRsqrt: y = _mm512_div_ps(one, _mm512_sqrt_ps(v)); break; + case kRecip: y = _mm512_div_ps(one, v); break; + case kAbs: y = _mm512_abs_ps(v); break; + case kNeg: y = _mm512_sub_ps(zero, v); break; + case kSign: { + // (x>0) ? 1 : ((x<0) ? -1 : 0) — ordered compares, NaN -> 0 + __mmask16 pos = _mm512_cmp_ps_mask(v, zero, _CMP_GT_OQ); + __mmask16 neg = _mm512_cmp_ps_mask(v, zero, _CMP_LT_OQ); + y = _mm512_mask_blend_ps(pos, _mm512_mask_blend_ps(neg, zero, negOne), + one); + break; + } + default: y = v; break; + } + _mm512_storeu_ps(out + i, y); + } + for (; i < n; ++i) out[i] = scalar_unary(X[i], kind); +} diff --git a/src/compiler/codegen/tessera_x86_backend/tests/test_unary.cpp b/src/compiler/codegen/tessera_x86_backend/tests/test_unary.cpp new file mode 100644 index 000000000..77149e016 --- /dev/null +++ b/src/compiler/codegen/tessera_x86_backend/tests/test_unary.cpp @@ -0,0 +1,91 @@ +// On-device test for the AVX-512 elementwise unary-math kernel (f32). +// +// Validates tessera_x86_avx512_unary_f32 against the scalar reference AND a +// hand-computed expectation, across kinds (sqrt/rsqrt/reciprocal/abs/neg/sign) +// and lengths incl. non-multiple-of-16 tails. Runs natively on the AVX-512 +// host — the "tested + running on the key device" proof for the CPU unary lane. + +#include +#include +#include +#include +#include + +extern "C" void tessera_x86_reference_unary_f32(const float*, int64_t, float*, + int); +extern "C" void tessera_x86_avx512_unary_f32(const float*, int64_t, float*, + int); + +static int g_fail = 0; + +static double want_unary(double v, int kind) { + switch (kind) { + case 0: return std::sqrt(v); + case 1: return 1.0 / std::sqrt(v); + case 2: return 1.0 / v; + case 3: return std::fabs(v); + case 4: return -v; + case 5: return (v > 0) ? 1.0 : (v < 0 ? -1.0 : 0.0); + default: return v; + } +} + +static void check(const char* name, int kind, int64_t n) { + std::mt19937 rng(2024 + (unsigned)(n * 131 + kind * 7)); + // sqrt/rsqrt/reciprocal need a positive domain; abs/neg/sign take any sign. + bool pos = (kind == 0 || kind == 1 || kind == 2); + std::uniform_real_distribution dist(pos ? 0.05f : -3.0f, 3.0f); + std::vector x((size_t)n); + for (auto& v : x) v = dist(rng); + + std::vector ref(n), avx(n); + tessera_x86_reference_unary_f32(x.data(), n, ref.data(), kind); + tessera_x86_avx512_unary_f32(x.data(), n, avx.data(), kind); + + for (int64_t i = 0; i < n; ++i) { + double want = want_unary(x[i], kind); + float tol = 2e-5f * (1.0f + std::fabs((float)want)); + if (std::fabs(avx[i] - ref[i]) > tol || + std::fabs(avx[i] - (float)want) > tol) { + std::printf("FAIL %s kind=%d n=%lld i=%lld: x=%g avx=%g ref=%g " + "want=%g\n", name, kind, (long long)n, (long long)i, + x[i], avx[i], ref[i], (double)want); + ++g_fail; + return; + } + } + std::printf("ok %s kind=%d n=%lld\n", name, kind, (long long)n); +} + +// sign() must yield exactly 0 at +/-0 and propagate the right sign elsewhere, +// in both the vector body and the scalar tail. +static void check_sign_zero() { + std::vector x = {-2.0f, -0.0f, 0.0f, 3.0f, -1e-9f, 1e-9f}; + int64_t n = (int64_t)x.size(); + std::vector avx(n), ref(n); + tessera_x86_avx512_unary_f32(x.data(), n, avx.data(), 5); + tessera_x86_reference_unary_f32(x.data(), n, ref.data(), 5); + float expect[] = {-1.0f, 0.0f, 0.0f, 1.0f, -1.0f, 1.0f}; + for (int64_t i = 0; i < n; ++i) { + if (avx[i] != expect[i] || ref[i] != expect[i]) { + std::printf("FAIL sign_zero i=%lld: x=%g avx=%g ref=%g want=%g\n", + (long long)i, x[i], avx[i], ref[i], expect[i]); + ++g_fail; + return; + } + } + std::printf("ok sign_zero\n"); +} + +int main() { + for (int kind = 0; kind <= 5; ++kind) { + check("aligned", kind, 64); // multiple of 16 + check("tail", kind, 70); // n % 16 != 0 + check("small", kind, 5); // n < 16 (all scalar tail) + check("wide", kind, 1024); // many vector steps + check("one", kind, 1); // degenerate + } + check_sign_zero(); + std::printf(g_fail ? "\n%d FAILED\n" : "\nALL PASSED\n", g_fail); + return g_fail ? 1 : 0; +} diff --git a/tests/unit/test_rocm_unary_compiled.py b/tests/unit/test_rocm_unary_compiled.py new file mode 100644 index 000000000..707f36ffc --- /dev/null +++ b/tests/unit/test_rocm_unary_compiled.py @@ -0,0 +1,148 @@ +"""Compiler-generated elementwise unary math (exp/log/sqrt/erf/…) on gfx1151 — +the S2 scalar-math / stability family, the unary sibling of the activation lane. + +The `tessera_rocm.unary` directive expands (via `generate-rocm-unary-kernel`) +into a flat per-element kernel (one thread per element). Reachable through +`runtime.launch()` via `compiler_path="rocm_unary_compiled"`; op names +tessera.exp / log / sqrt / rsqrt / reciprocal / absolute (abs) / sign / erf / +tanh / sigmoid / log1p / expm1 / softplus; f16/bf16/f32 storage, f32 compute. + +Validated vs numpy. Skip-clean: tessera-opt not built, or no usable AMD GPU. +""" + +from __future__ import annotations + +import numpy as np +import pytest + + +def _unary_or_skip(): + from tessera import runtime as rt + if rt._tessera_opt_path() is None: + pytest.skip("tessera-opt not built (ninja -C build tessera-opt)") + if not rt._rocm_wmma_runtime_available(): + pytest.skip("no usable AMD GPU") + return rt + + +def _artifact(rt, op_name): + return rt.RuntimeArtifact(metadata={ + "target": "rocm", "compiler_path": "rocm_unary_compiled", + "executable": True, "execution_kind": "native_gpu", + "arg_names": ["x"], "output_name": "o", + "ops": [{"op_name": op_name, "result": "o", "operands": ["x"]}], + }) + + +def _np_softplus(x): + # Stable: log1p(exp(-|x|)) + max(x, 0) (matches the kernel) + return np.log1p(np.exp(-np.abs(x))) + np.maximum(x, 0.0) + + +# op_name -> (numpy reference, sampler producing a valid domain) +_DOMAIN_POS = lambda rng, shp: (rng.random(shp) * 4.0 + 0.05).astype(np.float32) +_DOMAIN_ANY = lambda rng, shp: (rng.standard_normal(shp) * 1.5).astype(np.float32) + +_CASES = { + "tessera.exp": (np.exp, _DOMAIN_ANY), + "tessera.log": (np.log, _DOMAIN_POS), + "tessera.sqrt": (np.sqrt, _DOMAIN_POS), + "tessera.rsqrt": (lambda x: 1.0 / np.sqrt(x), _DOMAIN_POS), + "tessera.reciprocal": (np.reciprocal, _DOMAIN_POS), + "tessera.absolute": (np.abs, _DOMAIN_ANY), + "tessera.abs": (np.abs, _DOMAIN_ANY), + "tessera.sign": (np.sign, _DOMAIN_ANY), + "tessera.erf": (None, _DOMAIN_ANY), # erf filled in below if scipy present + "tessera.tanh": (np.tanh, _DOMAIN_ANY), + "tessera.sigmoid": (lambda x: 1.0 / (1.0 + np.exp(-x)), _DOMAIN_ANY), + "tessera.log1p": (np.log1p, _DOMAIN_POS), + "tessera.expm1": (np.expm1, _DOMAIN_ANY), + "tessera.softplus": (_np_softplus, _DOMAIN_ANY), +} + + +def _erf_ref(): + try: + from scipy.special import erf # type: ignore + return erf + except Exception: + # vectorized math.erf fallback — no scipy dependency required + import math + return np.vectorize(math.erf) + + +@pytest.mark.parametrize("op_name", list(_CASES)) +@pytest.mark.parametrize("dtype,tol", [ + (np.float32, 2e-5), (np.float16, 4e-3), ("bf16", 3e-2), +]) +@pytest.mark.parametrize("shape", [(8, 64), (130,), (3, 5, 7)]) +def test_unary_matches_numpy(op_name, dtype, tol, shape): + rt = _unary_or_skip() + if dtype == "bf16": + dtype = pytest.importorskip("ml_dtypes").bfloat16 + ref, sampler = _CASES[op_name] + if ref is None: + ref = _erf_ref() + rng = np.random.default_rng(11 + len(shape) + int(np.prod(shape))) + x = sampler(rng, shape).astype(dtype) + res = rt.launch(_artifact(rt, op_name), (x,)) + assert res["ok"] is True, res.get("reason") + assert res["compiler_path"] == "rocm_unary_compiled" + out = np.asarray(res["output"]).astype(np.float32) + expect = np.asarray(ref(x.astype(np.float32))).astype(np.float32) + np.testing.assert_allclose(out, expect, atol=tol, rtol=tol) + + +def test_unary_sign_at_zero(): + rt = _unary_or_skip() + x = np.array([-2.0, -0.0, 0.0, 3.0], np.float32) + res = rt.launch(_artifact(rt, "tessera.sign"), (x,)) + assert res["ok"] is True, res.get("reason") + out = np.asarray(res["output"]).astype(np.float32) + np.testing.assert_array_equal(out, np.array([-1.0, 0.0, 0.0, 1.0], np.float32)) + + +def test_unary_unknown_op_rejected(): + from tessera import runtime as rt + x = np.zeros((4, 8), np.float32) + with pytest.raises(ValueError, match="rocm_unary_compiled executor"): + rt._execute_rocm_compiled_unary(_artifact(rt, "tessera.softmax"), (x,)) + + +# ── GPU-free codegen gate (needs only tessera-opt, not a GPU) ──────────────── +import subprocess # noqa: E402 +from pathlib import Path # noqa: E402 + +_OPT = Path(__file__).resolve().parents[2] / "build/tools/tessera-opt/tessera-opt" + +_KINDS = ["exp", "log", "sqrt", "rsqrt", "reciprocal", "abs", "neg", "sign", + "erf", "tanh", "sigmoid", "log1p", "expm1", "softplus"] + + +def _opt(directive, *passes): + if not _OPT.is_file(): + pytest.skip("build tessera-opt: ninja -C build tessera-opt") + return subprocess.run([str(_OPT), "-", *passes], input=directive, + capture_output=True, text=True) + + +@pytest.mark.parametrize("kind", _KINDS) +def test_unary_codegen_and_lowers(kind): + import re + d = ('module {\n "tessera_rocm.unary"() {name = "u", ' + f'kind = "{kind}", dtype = "f32"}} : () -> ()\n}}\n') + ir = _opt(d, "--generate-rocm-unary-kernel") + assert ir.returncode == 0, ir.stderr + m = re.search(r"gpu\.func @u\(([^)]*)\)", ir.stdout) + assert m and len([a for a in m.group(1).split(",") if a.strip()]) == 3 + low = _opt(d, "--pass-pipeline=builtin.module(generate-rocm-unary-kernel," + "gpu.module(convert-scf-to-cf,convert-gpu-to-rocdl," + "reconcile-unrealized-casts))") + assert low.returncode == 0 and "llvm." in low.stdout + + +def test_unary_codegen_bad_kind_rejected(): + d = ('module {\n "tessera_rocm.unary"() {name = "u", kind = "floof"} ' + ': () -> ()\n}\n') + r = _opt(d, "--generate-rocm-unary-kernel") + assert r.returncode != 0 and "unknown kind" in r.stderr