From b2728a17562422bdc07849f24901a1250b2261b3 Mon Sep 17 00:00:00 2001 From: Aryan Ravishankar Date: Thu, 27 Aug 2026 20:28:39 -0500 Subject: [PATCH 1/2] Fuse the FP8 promotion into MATMUL so BF16 never reaches DDR MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The FP8 storage tier could only promote to a BF16 block output. That made it a poor deal for the case it exists to serve: expanding FP8 operands to feed a multiply meant materializing both BF16 tensors in DDR, so peak working set went *up* relative to just storing BF16, and the consumer paid a second submission to read them back. This admits the same explicit promotion feeding a MATMUL as one program. The graph is unchanged in meaning — two TOSA CAST operators still state the widening — but the backend performs it per L1 tile on the compute core, so the caller binds two FP8 operands and one FP32 result and no BF16 tensor is allocated or transferred. For M=K=N=256 that is 384 KiB of DDR against 640 KiB, 2.33x less DDR traffic, and one submission instead of three. Fusing is a placement choice, not a numerical one. FP8 to BF16 is exact for every encoding, and the multiply is the identical bf16->f32 `kernels.mm` the BF16 tier already uses, so results are bit-identical to running the two admitted tiers back to back. `fused_fp8_matmul_is_bit_identical_to_cast_then_matmul` checks exactly that on the NPU by running both paths and comparing the FP32 bytes, rather than against a host oracle that could drift. This is the first admitted tier with graph-interior values, so admission pins the dataflow tightly: each MATMUL operand must come from its own CAST, each CAST must consume a distinct block input, both operands must carry the same FP8 encoding (the kernel instantiates one decoder), and a promoted value that also escapes as a block output is rejected because the fused kernel never writes BF16 to DDR. The advertised FP8 capability gains MATMUL, BF16 as INTERMEDIATE, and FP32 as OUTPUT, so the descriptor covers the graph's interior as well as its boundary. The FP8 decoders are now shared between a sized emitter and the standalone CAST tier's fixed 1,024-element entry points; that tier's kernel source is unchanged byte for byte, so its cache key and on-metal evidence still hold. Co-Authored-By: Claude Opus 5 --- crates/virtio-accel-xdna/README.md | 14 + .../compiler/xdna_compile.py | 202 +++++++++- crates/virtio-accel-xdna/src/compiler.rs | 12 + crates/virtio-accel-xdna/src/lib.rs | 5 +- crates/virtio-accel-xdna/src/lower.rs | 364 +++++++++++++++++- crates/virtio-accel-xdna/tests/common/mod.rs | 62 +++ crates/virtio-accel-xdna/tests/hardware.rs | 153 +++++++- crates/virtio-accel-xdna/tests/xdna.rs | 120 +++++- 8 files changed, 910 insertions(+), 22 deletions(-) diff --git a/crates/virtio-accel-xdna/README.md b/crates/virtio-accel-xdna/README.md index bc57bcd..330bd8e 100644 --- a/crates/virtio-accel-xdna/README.md +++ b/crates/virtio-accel-xdna/README.md @@ -19,6 +19,14 @@ subset today is BF16 IDENTITY (a DMA copy), BF16 → FP32 MATMUL (the spec-manda shape, batch 1, at multiples of the tested compute tile), BF16 NHWC MAX_POOL2D, and explicit FP8E4M3/FP8E5M2 → BF16 CAST, plus exact INT8 IDENTITY, zero-point-aware INT8 → INT32 MATMUL, and signed per-tensor INT32 → INT8 RESCALE. +The FP8 storage tier has two admitted shapes. A standalone CAST promotes an FP8 tensor to a BF16 +block output. The *fused* FP8 → FP32 MATMUL admits the same explicit promotion feeding a MATMUL, +and performs it per L1 tile on the compute core instead of through DDR: the caller binds two FP8 +operands and one FP32 result, and no BF16 tensor is ever allocated or transferred. Fusing changes +only where the promotion happens — FP8 → BF16 is exact and the multiply is the same BF16 → FP32 +kernel — so results are bit-identical to running the two tiers back to back, which +`fused_fp8_matmul_is_bit_identical_to_cast_then_matmul` checks on the NPU by running both. + FP8 is a storage tier, not an arithmetic tier: the guest keeps the conversion visible in its graph, the NPU expands each value exactly, and subsequent programs use the existing BF16 compute kernels. MAX_POOL2D is @@ -163,6 +171,12 @@ without exercising the native backend. ## Performance evidence +For a promotion that immediately feeds a multiply, the fused tier is the one to reach for: the +standalone CAST writes a BF16 tensor twice the size of its input to DDR and the consumer reads it +back, so the unfused pair moves roughly 2.3x the bytes and costs an extra submission. Peak +working-set drops correspondingly — an M=K=N=256 multiply holds 384 KiB fused against 640 KiB +unfused, because the two BF16 operands never exist. + The FP8 storage conversion streams 1,024-element tiles through one AIE2P worker. That is an intentional first-tier implementation boundary, not a claim that one worker is the final throughput configuration. Program compilation happens at load time and is content-addressed; warm submission diff --git a/crates/virtio-accel-xdna/compiler/xdna_compile.py b/crates/virtio-accel-xdna/compiler/xdna_compile.py index 43ea3e4..831822a 100644 --- a/crates/virtio-accel-xdna/compiler/xdna_compile.py +++ b/crates/virtio-accel-xdna/compiler/xdna_compile.py @@ -19,6 +19,9 @@ "elements": , "device": "npu2", "fold_ddr_addr_offset": false} {"op": "MATMUL", "in_dtype": "bf16", "out_dtype": "f32", "m": , "k": , "n": , "device": "npu2", "fold_ddr_addr_offset": false} + {"op": "MATMUL", "in_dtype": "fp8e4m3" | "fp8e5m2", "out_dtype": "f32", "m": , "k": , + "n": , "tile_m": , "tile_k": , "tile_n": , "max_dim": , "device": "npu2", + "fold_ddr_addr_offset": false} # fused: FP8 in, BF16 promotion in L1, FP32 out {"op": "MATMUL", "in_dtype": "i8", "out_dtype": "i32", "m": , "k": , "n": , "left_zero_point": , "right_zero_point": , "device": "npu2", "fold_ddr_addr_offset": false} @@ -134,7 +137,7 @@ def _find_hrx_dir(prefix: Path) -> Path | None: return None -_FP8_CAST_KERNEL_SOURCE = r""" +_FP8_DECODERS = r""" #include static inline uint16_t fp8e4m3_to_bf16(uint8_t bits) { @@ -167,7 +170,10 @@ def _find_hrx_dir(prefix: Path) -> Path | None: return sign | static_cast((exponent + 112) << 7) | (fraction << 5); } -extern "C" void cast_fp8e4m3_to_bf16( +""" + +# The standalone CAST tier's two fixed 1,024-element entry points, unchanged. +_FP8_CAST_KERNEL_SOURCE = _FP8_DECODERS + r"""extern "C" void cast_fp8e4m3_to_bf16( const uint8_t *__restrict input, uint16_t *__restrict output) { #pragma clang loop vectorize(enable) interleave(enable) for (unsigned i = 0; i < 1024; ++i) { @@ -185,6 +191,25 @@ def _find_hrx_dir(prefix: Path) -> Path | None: """ +def _fp8_cast_kernel_source(symbol: str, fp8_dtype: str, length: int) -> str: + """One sized FP8-to-BF16 entry point over the shared exact decoders. + + The fused MATMUL widens whole L1 tiles, whose length is the tile geometry rather than the + CAST tier's transport line, so the loop bound is emitted rather than fixed. + """ + decoder = "fp8e4m3_to_bf16" if fp8_dtype == "fp8e4m3" else "fp8e5m2_to_bf16" + return _FP8_DECODERS + f''' +extern "C" void {symbol}( + const uint8_t *__restrict input, uint16_t *__restrict output) {{ +#pragma clang loop vectorize(enable) interleave(enable) + for (unsigned i = 0; i < {length}; ++i) {{ + output[i] = {decoder}(input[i]); + }} +}} +''' + + + def _build_identity(line_size: int, dtype: str): """Return the @iron.jit direct-DMA IDENTITY design (shim -> memtile -> shim).""" import aie.iron as iron @@ -272,6 +297,148 @@ def sequence(source, destination, in_prod, out_cons): return fp8_to_bf16 + +def _build_fp8_matmul(tile_m: int, tile_k: int, tile_n: int, fp8_dtype: str): + """Return the fused FP8 -> BF16 -> FP32 MATMUL design (`C[M,N] = A[M,K] . B[K,N]`). + + Structurally `_build_matmul`, with one difference: A and B arrive from DDR as FP8 storage + bytes, and each L1 tile is widened to BF16 by the same exact decoders the standalone CAST tier + uses. The BF16 operands exist only as core-local scratch, so the caller never allocates a BF16 + tensor and the promotion costs no DDR round trip. + + Numerically this is CAST-then-MATMUL with the intermediate never materialized: FP8 -> BF16 is + exact for every encoding, and the multiply is the identical `kernels.mm` bf16 -> f32 kernel, so + the result is bit-identical to running the two admitted tiers back to back. The graph still + states the promotion explicitly (two TOSA CAST operators); fusing is a placement choice, not a + relabeling of the arithmetic. + + The L2 -> L1 layout transform is expressed in elements, so it is dtype-agnostic; widening + afterwards is elementwise and preserves the micro-tile ordering the matmul kernel expects. + """ + import aie.iron as iron + import ml_dtypes + import numpy as np + from aie.helpers.taplib import TensorAccessPattern, TensorTiler2D + from aie.iron import ( + Buffer, + CompileTime, + In, + ObjectFifo, + Out, + Program, + Runtime, + TaskGroup, + Worker, + kernels, + ) + from aie.iron.controlflow import range_ + from aie.iron.kernel import ExternalFunction + + in_ty = ml_dtypes.bfloat16 + out_ty = np.float32 + tm, tk, tn = tile_m, tile_k, tile_n + + @iron.jit + def matmul_fp8_f32( + input0: In, + input1: In, + output: Out, + *, + M: CompileTime[int], + K: CompileTime[int], + N: CompileTime[int], + ): + matmul_kernel = kernels.mm( + dim_m=tm, dim_k=tk, dim_n=tn, + input_dtype=in_ty, output_dtype=out_ty, vectorized=True, + ) + r, s, t = matmul_kernel.mac_dims + + A_ty = np.ndarray[(M, K), np.dtype[np.uint8]] + B_ty = np.ndarray[(K, N), np.dtype[np.uint8]] + C_ty = np.ndarray[(M, N), np.dtype[out_ty]] + a_storage_ty = np.ndarray[(tm * tk,), np.dtype[np.uint8]] + b_storage_ty = np.ndarray[(tk * tn,), np.dtype[np.uint8]] + a_ty = np.ndarray[(tm * tk,), np.dtype[in_ty]] + b_ty = np.ndarray[(tk * tn,), np.dtype[in_ty]] + c_ty = np.ndarray[(tm * tn,), np.dtype[out_ty]] + + widen_a_symbol = f"widen_a_{fp8_dtype}" + widen_b_symbol = f"widen_b_{fp8_dtype}" + widen_a = ExternalFunction( + widen_a_symbol, + source_string=_fp8_cast_kernel_source(widen_a_symbol, fp8_dtype, tm * tk), + arg_types=[a_storage_ty, a_ty], + ) + widen_b = ExternalFunction( + widen_b_symbol, + source_string=_fp8_cast_kernel_source(widen_b_symbol, fp8_dtype, tk * tn), + arg_types=[b_storage_ty, b_ty], + ) + + fifo_a_l3l2 = ObjectFifo(a_storage_ty, name="A_L3L2") + tap_a = TensorTiler2D.group_tiler((tm, tk), (r, s), (tm // r, tk // s))[0] + fifo_a_l2l1 = fifo_a_l3l2.cons().forward(dims_to_stream=tap_a.transformation_dims, name="A_L2L1") + + fifo_b_l3l2 = ObjectFifo(b_storage_ty, name="B_L3L2") + tap_b = TensorTiler2D.group_tiler((tk, tn), (s, t), (tk // s, tn // t))[0] + fifo_b_l2l1 = fifo_b_l3l2.cons().forward(dims_to_stream=tap_b.transformation_dims, name="B_L2L1") + + fifo_c_l1l2 = ObjectFifo(c_ty, name="C_L1L2") + tap_c = TensorAccessPattern( + tensor_dims=(tm, tn), offset=0, + sizes=[tm // r, r, tn // t, t], strides=[r * tn, t, r * t, 1], + ) + fifo_c_l2l3 = fifo_c_l1l2.cons().forward(dims_to_stream=list(tap_c.transformation_dims), name="C_L2L3") + + # The promoted operands: core-local only, never a runtime binding. + a_scratch = Buffer(a_ty, name="A_bf16_scratch") + b_scratch = Buffer(b_ty, name="B_bf16_scratch") + + def core_fn(of_a, of_b, of_c, a_bf16, b_bf16, widen_lhs, widen_rhs, matmul): + for _ in range_(M // tm * N // tn): + elem_out = of_c.acquire(1) + for i in range_(tm * tn): + elem_out[i] = 0 + for _ in range_(K // tk): + elem_in_a = of_a.acquire(1) + elem_in_b = of_b.acquire(1) + widen_lhs(elem_in_a, a_bf16) + widen_rhs(elem_in_b, b_bf16) + matmul(a_bf16, b_bf16, elem_out) + of_a.release(1) + of_b.release(1) + of_c.release(1) + + worker = Worker( + core_fn, + [ + fifo_a_l2l1.cons(), fifo_b_l2l1.cons(), fifo_c_l1l2.prod(), + a_scratch, b_scratch, widen_a, widen_b, matmul_kernel, + ], + ) + + a_taps = TensorTiler2D.group_tiler((M, K), (tm, tk), (1, K // tk), pattern_repeat=(N // tn)) + b_tap = TensorTiler2D.group_tiler((K, N), (tk, tn), (K // tk, N // tn), tile_group_col_major=True)[0] + c_taps = TensorTiler2D.group_tiler((M, N), (tm, tn), (1, N // tn)) + + def sequence(a_src, b_src, c_dst, a_prod, b_prod, c_cons): + for tile_row in range(M // tm): + task_group = TaskGroup() + a_prod.fill(a_src, tap=a_taps[tile_row], group=task_group) + b_prod.fill(b_src, tap=b_tap, group=task_group) + c_cons.drain(c_dst, tap=c_taps[tile_row], group=task_group, wait=True) + task_group.finish() + + rt = Runtime( + sequence, + [A_ty, B_ty, C_ty, fifo_a_l3l2.prod(), fifo_b_l3l2.prod(), fifo_c_l2l3.cons()], + ) + return Program(iron.get_current_device(), rt, workers=[worker]).resolve_program() + + return matmul_fp8_f32 + + def _build_matmul(tile_m: int, tile_k: int, tile_n: int): """Return the @iron.jit BF16 -> FP32 single-core MATMUL design (`C[M,N] = A[M,K] . B[K,N]`). @@ -753,7 +920,12 @@ def build(): out_dtype = spec.get("out_dtype") m, k, n = spec.get("m"), spec.get("k"), spec.get("n") max_dim = spec.get("max_dim") - if (in_dtype, out_dtype) not in (("bf16", "f32"), ("i8", "i32")): + if (in_dtype, out_dtype) not in ( + ("bf16", "f32"), + ("i8", "i32"), + ("fp8e4m3", "f32"), + ("fp8e5m2", "f32"), + ): return _fail( workdir, "spec-rejected", f"unsupported dtype pair: {in_dtype}->{out_dtype}" ) @@ -764,7 +936,29 @@ def build(): return _fail( workdir, "spec-rejected", f"{name}={dim} must be positive and <= {max_dim}" ) - if in_dtype == "bf16": + if in_dtype in ("fp8e4m3", "fp8e5m2"): + tile_m, tile_k, tile_n = ( + spec.get("tile_m"), + spec.get("tile_k"), + spec.get("tile_n"), + ) + for name, dim, tile in (("m", m, tile_m), ("k", k, tile_k), ("n", n, tile_n)): + if not isinstance(tile, int) or tile <= 0: + return _fail(workdir, "spec-rejected", f"invalid tile_{name}: {tile}") + if dim % tile != 0: + return _fail( + workdir, + "spec-rejected", + f"{name}={dim} must be a multiple of {tile}", + ) + # FP8 operands bind one byte per element; only the FP32 result reaches DDR at full + # width. The promoted BF16 operands never leave the compute core. + input_bytes, output_bytes = [m * k, k * n], [m * n * 4] + + def build(): + return _build_fp8_matmul(tile_m, tile_k, tile_n, in_dtype).specialize(M=m, K=k, N=n) + + elif in_dtype == "bf16": tile_m, tile_k, tile_n = ( spec.get("tile_m"), spec.get("tile_k"), diff --git a/crates/virtio-accel-xdna/src/compiler.rs b/crates/virtio-accel-xdna/src/compiler.rs index f1aac44..a012f66 100644 --- a/crates/virtio-accel-xdna/src/compiler.rs +++ b/crates/virtio-accel-xdna/src/compiler.rs @@ -86,6 +86,18 @@ fn spec_json(spec: CompilerSpec) -> String { \"tile_k\":{MATMUL_TILE_K},\"tile_n\":{MATMUL_TILE_N},\ \"max_dim\":{MATMUL_MAX_DIM},{device}}}" ), + CompilerSpec::Fp8Matmul { format, m, k, n } => { + let input = match format { + Fp8Format::E4M3 => "fp8e4m3", + Fp8Format::E5M2 => "fp8e5m2", + }; + format!( + "{{\"op\":\"MATMUL\",\"in_dtype\":\"{input}\",\"out_dtype\":\"f32\",\ + \"m\":{m},\"k\":{k},\"n\":{n},\"tile_m\":{MATMUL_TILE_M},\ + \"tile_k\":{MATMUL_TILE_K},\"tile_n\":{MATMUL_TILE_N},\ + \"max_dim\":{MATMUL_MAX_DIM},{device}}}" + ) + } CompilerSpec::Int8Matmul { m, k, diff --git a/crates/virtio-accel-xdna/src/lib.rs b/crates/virtio-accel-xdna/src/lib.rs index f917a94..440f4b2 100644 --- a/crates/virtio-accel-xdna/src/lib.rs +++ b/crates/virtio-accel-xdna/src/lib.rs @@ -19,8 +19,9 @@ //! (execution-model spec, issue #85). `load_program` accepts the crate-local precompiled format //! ([`artifact`]) directly, and a TOSA artifact by admitting it and compiling it with the bounded //! aiecc helper subprocess (issue #84). The compilable TOSA subsets today are BF16 IDENTITY, -//! BF16 → FP32 MATMUL, BF16 MAX_POOL2D, explicit FP8 → BF16 storage conversion, and exact INT8 -//! IDENTITY plus zero-point-aware INT8 → INT32 MATMUL. +//! BF16 → FP32 MATMUL, BF16 MAX_POOL2D, explicit FP8 → BF16 storage conversion, the fused +//! FP8 → FP32 MATMUL that keeps that promotion on the compute core, exact INT8 IDENTITY, +//! zero-point-aware INT8 → INT32 MATMUL, and exact INT32 → INT8 RESCALE. //! Admission (`lower`) unit-tests on every host. #![cfg_attr(not(va_xdna), forbid(unsafe_code))] diff --git a/crates/virtio-accel-xdna/src/lower.rs b/crates/virtio-accel-xdna/src/lower.rs index d5e2655..d81dbe5 100644 --- a/crates/virtio-accel-xdna/src/lower.rs +++ b/crates/virtio-accel-xdna/src/lower.rs @@ -81,13 +81,24 @@ pub const XDNA_TOSA_CAPABILITY: CapabilityDescriptor = CapabilityDescriptor { }, }; +// The roles must cover every admitted tier's boundary *and* interior. The standalone CAST tier +// ends at a BF16 block output; the fused MATMUL tier keeps the same explicit BF16 promotion as a +// graph-interior value and ends at FP32, so BF16 carries `INTERMEDIATE` and FP32 carries `OUTPUT`. const FP8_STORAGE_DTYPES: &[DTypeCapability] = &[ DTypeCapability::new(DType::FP8E4M3, ValueRoles::INPUT), DTypeCapability::new(DType::FP8E5M2, ValueRoles::INPUT), - DTypeCapability::new(DType::BF16, ValueRoles::OUTPUT), + DTypeCapability::new( + DType::BF16, + ValueRoles::OUTPUT.union(ValueRoles::INTERMEDIATE), + ), + DTypeCapability::new(DType::FP32, ValueRoles::OUTPUT), ]; -const FP8_STORAGE_OPERATORS: &[OperatorCapability] = &[OperatorCapability::new(Op::CAST)]; +const FP8_STORAGE_OPERATORS: &[OperatorCapability] = &[ + OperatorCapability::new(Op::CAST), + OperatorCapability::new(Op::CONST), + OperatorCapability::constrained(Op::MATMUL, OperatorConstraints::ZERO_ZERO_POINTS), +]; /// Conservative capability boundary for explicit FP8 storage conversion. pub const XDNA_TOSA_FP8_CAPABILITY: CapabilityDescriptor = CapabilityDescriptor { @@ -213,6 +224,19 @@ pub enum CompilerSpec { /// `k`, `n` is a positive multiple of the corresponding MATMUL tile dimension and at most /// 512. The FP32 output is the TOSA-mandated accumulator (issue #82). Matmul { m: usize, k: usize, n: usize }, + /// Fused FP8 × FP8 → FP32 matrix multiply (batch 1): the graph's explicit BF16 promotion is + /// performed on the compute core, per L1 tile, instead of through DDR. + /// + /// Numerically identical to [`CompilerSpec::Fp8ToBf16`] followed by [`CompilerSpec::Matmul`] — + /// FP8 → BF16 is exact for every encoding and the multiply is the same BF16 → FP32 kernel — so + /// fusing is a placement choice, not a change of numerical contract. It removes the + /// caller-visible BF16 tensors and their DDR round trip. + Fp8Matmul { + format: Fp8Format, + m: usize, + k: usize, + n: usize, + }, /// Exact zero-point-aware INT8 × INT8 → INT32 matrix multiply (batch 1). /// /// The serialized TOSA zero points are part of the specialization and therefore also part of @@ -327,7 +351,8 @@ pub fn admit(bytes: &[u8], target: Target) -> Result { // to the two zero-points. let mut matmul = None; let mut max_pool = None; - let mut cast = None; + let mut casts: [Option; 2] = [None, None]; + let mut cast_count = 0usize; let mut rescale = None; let mut identities = 0usize; let mut constants = 0usize; @@ -337,33 +362,46 @@ pub fn admit(bytes: &[u8], target: Target) -> Result { Op::CONST => constants += 1, Op::MATMUL if matmul.is_none() => matmul = Some(*operator), Op::MAX_POOL2D if max_pool.is_none() => max_pool = Some(*operator), - Op::CAST if cast.is_none() => cast = Some(*operator), + Op::CAST if cast_count < casts.len() => { + casts[cast_count] = Some(*operator); + cast_count += 1; + } Op::RESCALE if rescale.is_none() => rescale = Some(*operator), _ => return Err(AdmitError::Unsupported), } } match ( - target, matmul, max_pool, cast, rescale, identities, constants, + target, matmul, max_pool, cast_count, rescale, identities, constants, ) { // All-IDENTITY (zero operators included: the block output then *is* the block input, and a // DMA copy is exact for it). - (XDNA_TOSA_TARGET, None, None, None, None, _, 0) => admit_identity(&analysis, block), - (XDNA_TOSA_TARGET, Some(matmul), None, None, None, 0, _) => { + (XDNA_TOSA_TARGET, None, None, 0, None, _, 0) => admit_identity(&analysis, block), + (XDNA_TOSA_TARGET, Some(matmul), None, 0, None, 0, _) => { admit_matmul(&analysis, block, matmul) } - (XDNA_TOSA_TARGET, None, Some(max_pool), None, None, 0, 0) => { + (XDNA_TOSA_TARGET, None, Some(max_pool), 0, None, 0, 0) => { admit_max_pool2d(&analysis, block, max_pool) } - (XDNA_TOSA_FP8_TARGET, None, None, Some(cast), None, 0, 0) => { - admit_fp8_to_bf16(&analysis, block, cast) + (XDNA_TOSA_FP8_TARGET, None, None, 1, None, 0, 0) => { + admit_fp8_to_bf16(&analysis, block, casts[0].ok_or(AdmitError::Unsupported)?) } - (XDNA_TOSA_INTEGER_TARGET, None, None, None, None, _, 0) => { + // Fused: both MATMUL operands are promoted from FP8 by their own explicit CAST. + (XDNA_TOSA_FP8_TARGET, Some(matmul), None, 2, None, 0, _) => admit_fp8_matmul( + &analysis, + block, + matmul, + [ + casts[0].ok_or(AdmitError::Unsupported)?, + casts[1].ok_or(AdmitError::Unsupported)?, + ], + ), + (XDNA_TOSA_INTEGER_TARGET, None, None, 0, None, _, 0) => { admit_int8_identity(&analysis, block) } - (XDNA_TOSA_INTEGER_TARGET, Some(matmul), None, None, None, 0, _) => { + (XDNA_TOSA_INTEGER_TARGET, Some(matmul), None, 0, None, 0, _) => { admit_int8_matmul(&analysis, block, matmul) } - (XDNA_TOSA_INTEGER_TARGET, None, None, None, Some(rescale), 0, 4) => { + (XDNA_TOSA_INTEGER_TARGET, None, None, 0, Some(rescale), 0, 4) => { admit_int32_to_int8_rescale(&analysis, block, rescale) } _ => Err(AdmitError::Unsupported), @@ -779,6 +817,141 @@ fn admit_matmul( Ok(CompilerSpec::Matmul { m, k, n }) } +/// Admit a fused FP8 MATMUL: `CAST(A_fp8) . CAST(B_fp8) -> FP32`, batch 1. +/// +/// This is the only admitted tier with graph-interior values, so the dataflow is pinned exactly: +/// each MATMUL operand must be produced by its own CAST, each CAST must consume a distinct block +/// input, and the promoted BF16 values must not escape as block outputs. Anything looser would let +/// the compiled kernel — which binds two FP8 inputs and one FP32 output and promotes internally — +/// stand in for a graph it does not implement. +/// +/// The promotion stays explicit in the graph, exactly as the standalone CAST tier requires; only +/// its *placement* changes, from a DDR round trip to core-local scratch. The arithmetic is +/// unchanged, so results are bit-identical to running CAST and MATMUL as separate programs. +fn admit_fp8_matmul( + analysis: &virtio_accel_tosa::TosaAnalysis<'_>, + block: virtio_accel_tosa::BlockId, + matmul: virtio_accel_tosa::OperatorId, + casts: [virtio_accel_tosa::OperatorId; 2], +) -> Result { + let inputs = analysis.operator_inputs(matmul); + let outputs = analysis.operator_outputs(matmul); + if inputs.len() != 4 || outputs.len() != 1 { + return Err(AdmitError::Unsupported); + } + if analysis.block_outputs(block) != [outputs[0]] { + return Err(AdmitError::Unsupported); + } + + // Pair each CAST with the MATMUL operand it produces; the two must cover lhs and rhs exactly + // once each, in binding order. + let mut promoted: [Option; 2] = [None, None]; + for cast in casts { + let cast_inputs = analysis.operator_inputs(cast); + let cast_outputs = analysis.operator_outputs(cast); + if cast_inputs.len() != 1 || cast_outputs.len() != 1 { + return Err(AdmitError::Unsupported); + } + let operand = if cast_outputs[0] == inputs[0] { + 0 + } else if cast_outputs[0] == inputs[1] { + 1 + } else { + // A CAST feeding anything but a MATMUL operand is graph the kernel cannot reproduce. + return Err(AdmitError::Unsupported); + }; + if promoted[operand].is_some() { + return Err(AdmitError::Unsupported); + } + // The promoted value is interior: it feeds the multiply and must not also be a block + // output, which would require materializing the BF16 tensor this tier exists to avoid. + if analysis.block_outputs(block).contains(&cast_outputs[0]) { + return Err(AdmitError::Unsupported); + } + promoted[operand] = Some(cast_inputs[0]); + } + let ([Some(lhs_storage), Some(rhs_storage)], _) = (promoted, ()) else { + return Err(AdmitError::Unsupported); + }; + + // The block's dataflow: the two FP8 storage tensors are the block inputs in binding order. One + // value feeding both operands is rejected for the same reason as the BF16 tier. + if lhs_storage == rhs_storage || analysis.block_inputs(block) != [lhs_storage, rhs_storage] { + return Err(AdmitError::Unsupported); + } + // Every CONST feeds only the zero points (operands 2 and 3). + for operator in analysis.execution_order(block) { + if analysis.operator(*operator).op() != Op::CONST { + continue; + } + for produced in analysis.operator_outputs(*operator) { + if *produced != inputs[2] && *produced != inputs[3] { + return Err(AdmitError::Unsupported); + } + } + } + + // Both operands must carry the same FP8 encoding: the compiled kernel instantiates one decoder. + let lhs_format = fp8_storage_format(analysis, lhs_storage)?; + let rhs_format = fp8_storage_format(analysis, rhs_storage)?; + if lhs_format != rhs_format { + return Err(AdmitError::Unsupported); + } + + let lhs = matmul_dims(analysis, lhs_storage, fp8_dtype(lhs_format))?; + let rhs = matmul_dims(analysis, rhs_storage, fp8_dtype(rhs_format))?; + let lhs_bf16 = matmul_dims(analysis, inputs[0], DType::BF16)?; + let rhs_bf16 = matmul_dims(analysis, inputs[1], DType::BF16)?; + let out = matmul_dims(analysis, outputs[0], DType::FP32)?; + + // Promotion is elementwise, so each CAST must preserve its operand's shape exactly. + if lhs != lhs_bf16 || rhs != rhs_bf16 { + return Err(AdmitError::Unsupported); + } + let ([1, m, k], [1, k2, n], [1, m2, n2]) = (lhs, rhs, out) else { + return Err(AdmitError::Unsupported); + }; + if k != k2 || m != m2 || n != n2 { + return Err(AdmitError::Unsupported); + } + if !tile_admissible(m, MATMUL_TILE_M) + || !tile_admissible(k, MATMUL_TILE_K) + || !tile_admissible(n, MATMUL_TILE_N) + { + return Err(AdmitError::Unsupported); + } + + Ok(CompilerSpec::Fp8Matmul { + format: lhs_format, + m, + k, + n, + }) +} + +/// The FP8 storage encoding of `value`, or `Unsupported` for any other dtype. +fn fp8_storage_format( + analysis: &virtio_accel_tosa::TosaAnalysis<'_>, + value: virtio_accel_tosa::ValueId, +) -> Result { + let AnalyzedValueKind::Tensor(tensor) = analysis.value(value).kind() else { + return Err(AdmitError::Unsupported); + }; + match tensor.dtype() { + DType::FP8E4M3 => Ok(Fp8Format::E4M3), + DType::FP8E5M2 => Ok(Fp8Format::E5M2), + _ => Err(AdmitError::Unsupported), + } +} + +/// The TOSA dtype of an FP8 storage encoding. +fn fp8_dtype(format: Fp8Format) -> DType { + match format { + Fp8Format::E4M3 => DType::FP8E4M3, + Fp8Format::E5M2 => DType::FP8E5M2, + } +} + /// The rank-3 dimensions of `value`, requiring the given dtype and every dimension statically /// positive. Dynamic (non-positive) dimensions and non-tensor or wrong-dtype values are rejected. fn matmul_dims( @@ -1591,6 +1764,171 @@ mod tests { ); } + /// Build a fused FP8 MATMUL graph: two FP8 block inputs, each promoted by its own explicit + /// CAST, multiplied to FP32. `escape` additionally exposes the promoted lhs as a block output. + fn fp8_matmul_graph( + m: i32, + k: i32, + n: i32, + lhs_dtype: DType, + rhs_dtype: DType, + alias_operands: bool, + escape: bool, + ) -> OwnedGraph<'static> { + let mut graph = OwnedGraph::new("main"); + let rhs_name = if alias_operands { "lhs_fp8" } else { "rhs_fp8" }; + graph + .push_tensor(OwnedTensor::new("lhs_fp8", vec![1, m, k], lhs_dtype)) + .push_tensor(OwnedTensor::new("lhs_bf16", vec![1, m, k], DType::BF16)) + .push_tensor(OwnedTensor::new("rhs_bf16", vec![1, k, n], DType::BF16)) + .push_tensor(OwnedTensor::constant( + "lhs_zp", + vec![1], + DType::BF16, + vec![0u8; 2], + )) + .push_tensor(OwnedTensor::constant( + "rhs_zp", + vec![1], + DType::BF16, + vec![0u8; 2], + )) + .push_tensor(OwnedTensor::new("output", vec![1, m, n], DType::FP32)); + if !alias_operands { + graph.push_tensor(OwnedTensor::new("rhs_fp8", vec![1, k, n], rhs_dtype)); + } + graph + .push_operator(OwnedOperator::new( + OperatorKind::Cast, + vec!["lhs_fp8".into()], + vec!["lhs_bf16".into()], + )) + .push_operator(OwnedOperator::new( + OperatorKind::Cast, + vec![rhs_name.into()], + vec!["rhs_bf16".into()], + )) + .push_operator(OwnedOperator::new( + OperatorKind::Const, + vec![], + vec!["lhs_zp".into()], + )) + .push_operator(OwnedOperator::new( + OperatorKind::Const, + vec![], + vec!["rhs_zp".into()], + )) + .push_operator(OwnedOperator::new( + OperatorKind::MatMul, + vec![ + "lhs_bf16".into(), + "rhs_bf16".into(), + "lhs_zp".into(), + "rhs_zp".into(), + ], + vec!["output".into()], + )) + .push_input("lhs_fp8"); + if !alias_operands { + graph.push_input("rhs_fp8"); + } else { + graph.push_input("lhs_fp8"); + } + graph.push_output("output"); + if escape { + graph.push_output("lhs_bf16"); + } + graph + } + + #[test] + fn admits_fused_fp8_matmul_for_both_encodings() { + for (dtype, format) in [ + (DType::FP8E4M3, Fp8Format::E4M3), + (DType::FP8E5M2, Fp8Format::E5M2), + ] { + let bytes = fp8_matmul_graph(32, 64, 32, dtype, dtype, false, false) + .build(XDNA_TOSA_FP8_TARGET) + .expect("build fused fp8 matmul"); + assert_eq!( + admit(&bytes, XDNA_TOSA_FP8_TARGET), + Ok(CompilerSpec::Fp8Matmul { + format, + m: 32, + k: 64, + n: 32 + }) + ); + } + } + + /// The compiled kernel instantiates exactly one decoder, so mixing encodings across the two + /// operands must not be admitted under either encoding's label. + #[test] + fn rejects_fused_fp8_matmul_with_mixed_encodings() { + let bytes = fp8_matmul_graph(32, 64, 32, DType::FP8E4M3, DType::FP8E5M2, false, false) + .build(XDNA_TOSA_FP8_TARGET) + .expect("build mixed-encoding fused matmul"); + assert_eq!( + admit(&bytes, XDNA_TOSA_FP8_TARGET), + Err(AdmitError::Unsupported) + ); + } + + /// The promoted BF16 value is graph-interior. If the graph also demands it as a block output, + /// the fused kernel cannot serve it — it never writes BF16 to DDR. + #[test] + fn rejects_fused_fp8_matmul_whose_promoted_operand_escapes() { + let bytes = fp8_matmul_graph(32, 64, 32, DType::FP8E4M3, DType::FP8E4M3, false, true) + .build(XDNA_TOSA_FP8_TARGET) + .expect("build escaping fused matmul"); + assert_eq!( + admit(&bytes, XDNA_TOSA_FP8_TARGET), + Err(AdmitError::Unsupported) + ); + } + + /// Square shape so that one FP8 tensor can legally feed both CASTs; the rejection must then + /// come from admission, not from TOSA shape validation. + #[test] + fn rejects_fused_fp8_matmul_with_one_value_feeding_both_operands() { + let bytes = fp8_matmul_graph(64, 64, 64, DType::FP8E4M3, DType::FP8E4M3, true, false) + .build(XDNA_TOSA_FP8_TARGET) + .expect("build aliased fused matmul"); + assert_eq!( + admit(&bytes, XDNA_TOSA_FP8_TARGET), + Err(AdmitError::Unsupported) + ); + // Positive control: the same shape with two distinct operands is admitted, so the + // rejection above is specifically the aliasing and not the shape. + let distinct = fp8_matmul_graph(64, 64, 64, DType::FP8E4M3, DType::FP8E4M3, false, false) + .build(XDNA_TOSA_FP8_TARGET) + .expect("build distinct fused matmul"); + assert!(admit(&distinct, XDNA_TOSA_FP8_TARGET).is_ok()); + } + + #[test] + fn rejects_fused_fp8_matmul_off_the_tested_tiling() { + for (m, k, n) in [(48, 64, 32), (32, 64, MATMUL_MAX_DIM as i32 + 32)] { + let bytes = fp8_matmul_graph(m, k, n, DType::FP8E4M3, DType::FP8E4M3, false, false) + .build(XDNA_TOSA_FP8_TARGET) + .expect("build fused matmul"); + assert_eq!( + admit(&bytes, XDNA_TOSA_FP8_TARGET), + Err(AdmitError::Unsupported) + ); + } + } + + /// The fused tier lives on the FP8 target; the BF16 target must not admit an FP8 graph. + #[test] + fn fused_fp8_matmul_is_not_admitted_on_the_bf16_target() { + let bytes = fp8_matmul_graph(32, 64, 32, DType::FP8E4M3, DType::FP8E4M3, false, false) + .build(XDNA_TOSA_FP8_TARGET) + .expect("build fused matmul"); + assert_eq!(admit(&bytes, XDNA_TOSA_TARGET), Err(AdmitError::Analysis)); + } + #[test] fn rejects_fp32_matmul_inputs() { // FP32-input MATMUL is admissible TOSA but has no compute path on this hardware. diff --git a/crates/virtio-accel-xdna/tests/common/mod.rs b/crates/virtio-accel-xdna/tests/common/mod.rs index 5d02f53..8fb3e65 100644 --- a/crates/virtio-accel-xdna/tests/common/mod.rs +++ b/crates/virtio-accel-xdna/tests/common/mod.rs @@ -142,3 +142,65 @@ pub fn poll_to_terminal( } } } + +/// A fused FP8 MATMUL graph: two FP8E4M3 block inputs, each promoted by its own explicit CAST, +/// multiplied to FP32. The promotion stays in the graph; only its placement is the backend's +/// choice. +#[allow(dead_code)] // Used by the hardware suite; this module is compiled into conformance.rs too. +pub fn fp8e4m3_matmul_tosa(m: i32, k: i32, n: i32) -> Vec { + let mut graph = OwnedGraph::new("main"); + graph + .push_tensor(OwnedTensor::new("lhs_fp8", vec![1, m, k], DType::FP8E4M3)) + .push_tensor(OwnedTensor::new("rhs_fp8", vec![1, k, n], DType::FP8E4M3)) + .push_tensor(OwnedTensor::new("lhs_bf16", vec![1, m, k], DType::BF16)) + .push_tensor(OwnedTensor::new("rhs_bf16", vec![1, k, n], DType::BF16)) + .push_tensor(OwnedTensor::constant( + "lhs_zp", + vec![1], + DType::BF16, + vec![0u8; 2], + )) + .push_tensor(OwnedTensor::constant( + "rhs_zp", + vec![1], + DType::BF16, + vec![0u8; 2], + )) + .push_tensor(OwnedTensor::new("output", vec![1, m, n], DType::FP32)) + .push_operator(OwnedOperator::new( + OperatorKind::Cast, + vec!["lhs_fp8".into()], + vec!["lhs_bf16".into()], + )) + .push_operator(OwnedOperator::new( + OperatorKind::Cast, + vec!["rhs_fp8".into()], + vec!["rhs_bf16".into()], + )) + .push_operator(OwnedOperator::new( + OperatorKind::Const, + vec![], + vec!["lhs_zp".into()], + )) + .push_operator(OwnedOperator::new( + OperatorKind::Const, + vec![], + vec!["rhs_zp".into()], + )) + .push_operator(OwnedOperator::new( + OperatorKind::MatMul, + vec![ + "lhs_bf16".into(), + "rhs_bf16".into(), + "lhs_zp".into(), + "rhs_zp".into(), + ], + vec!["output".into()], + )) + .push_input("lhs_fp8") + .push_input("rhs_fp8") + .push_output("output"); + graph + .build(XDNA_TOSA_FP8_TARGET) + .expect("build fused fp8 matmul") +} diff --git a/crates/virtio-accel-xdna/tests/hardware.rs b/crates/virtio-accel-xdna/tests/hardware.rs index 64a3cc1..c47d644 100644 --- a/crates/virtio-accel-xdna/tests/hardware.rs +++ b/crates/virtio-accel-xdna/tests/hardware.rs @@ -10,7 +10,8 @@ use std::time::{Duration, Instant}; mod common; use common::{ - bf16_identity_tosa, bf16_matmul_tosa, fp8e4m3_to_bf16_tosa, int8_matmul_tosa, poll_to_terminal, + bf16_identity_tosa, bf16_matmul_tosa, fp8e4m3_matmul_tosa, fp8e4m3_to_bf16_tosa, + int8_matmul_tosa, poll_to_terminal, }; use virtio_accel_conformance::numerics::{ @@ -2012,6 +2013,156 @@ fn tosa_bf16_matmul_compiles_to_a_wellformed_artifact() { )); } +/// The fused FP8 MATMUL must be bit-identical to running the two admitted tiers back to back. +/// +/// This is the tier's whole correctness claim: fusing changes *where* the graph's explicit BF16 +/// promotion happens (core-local scratch instead of a DDR round trip), never the arithmetic. Both +/// paths run on the NPU here and their FP32 results are compared byte for byte, so the comparison +/// cannot drift with a host-side oracle. +#[test] +fn fused_fp8_matmul_is_bit_identical_to_cast_then_matmul() { + let Some(backend) = backend() else { return }; + if !toolchain_present() { + assert!( + !hardware_required(), + "{REQUIRE_HARDWARE_ENV}=1 but VIRTIO_ACCEL_AMDXDNA_TOOLCHAIN is not configured" + ); + eprintln!("no XDNA toolchain configured; skipping TOSA execution test"); + return; + } + // Multi-tile in every dimension, and both operand element counts are multiples of the CAST + // tier's 1,024-element line so the unfused reference path is admissible too. + const M: usize = 64; + const K: usize = 128; + const N: usize = 96; + + // Finite FP8E4M3 encodings only (0x7f/0xff are NaN), including the subnormals at 0x00..=0x07. + let lhs_fp8: Vec = (0..M * K).map(|i| (i % 120) as u8).collect(); + let rhs_fp8: Vec = (0..K * N).map(|i| ((i * 7 + 3) % 120) as u8).collect(); + + let context = backend + .create_context(ContextDesc::default()) + .expect("context"); + let queue = backend + .create_queue(&context, QueueDesc::default()) + .expect("queue"); + + let load = |tosa: &[u8], target: virtio_accel_tosa::Target| { + backend + .load_program( + &context, + ArtifactRef { + format: ARTIFACT_FORMAT, + target: target.to_identity(), + payload: &Slice(tosa), + resident_bytes: u64::MAX, + }, + ) + .expect("load + compile TOSA program") + }; + let in_desc = |bytes: usize| { + BufferDesc::new( + bytes as u64, + 4096, + MemoryDomain::Shared, + BufferUsage::TRANSFER_DESTINATION | BufferUsage::PROGRAM_INPUT, + ) + .unwrap() + }; + let out_desc = |bytes: usize| { + BufferDesc::new( + bytes as u64, + 4096, + MemoryDomain::Shared, + BufferUsage::TRANSFER_SOURCE | BufferUsage::PROGRAM_OUTPUT, + ) + .unwrap() + }; + + // One binary program run: write every input, submit, poll, read the output back. + let run = |program: &_, inputs: &[&[u8]], output_len: usize| -> Vec { + let mut in_buffers = Vec::new(); + for bytes in inputs { + let (mut buffer, _) = backend + .allocate_buffer(&context, in_desc(bytes.len())) + .expect("input buffer") + .into_parts(); + backend + .write_buffer(&mut buffer, 0, &Slice(bytes)) + .expect("write input"); + in_buffers.push(buffer); + } + let (out_buffer, _) = backend + .allocate_buffer(&context, out_desc(output_len)) + .expect("output buffer") + .into_parts(); + + let mut bindings: Vec> = in_buffers + .iter() + .enumerate() + .map(|(slot, buffer)| BindingRef { + slot: slot as u32, + buffer, + range: BufferRange::new(0, inputs[slot].len() as u64).unwrap(), + access: AccessMode::Read, + }) + .collect(); + bindings.push(BindingRef { + slot: in_buffers.len() as u32, + buffer: &out_buffer, + range: BufferRange::new(0, output_len as u64).unwrap(), + access: AccessMode::Write, + }); + + let event = backend + .submit(&queue, program, &bindings, Timeout::Infinite) + .expect("submit"); + let state = poll_to_terminal(&backend, &event, Duration::from_secs(30)) + .expect("program did not complete"); + assert!(matches!(state, EventState::Complete), "got {state:?}"); + + let mut result = vec![0u8; output_len]; + backend + .read_buffer(&out_buffer, 0, &mut SliceMut(&mut result)) + .expect("read output"); + backend.destroy_event(event).expect("destroy event"); + for buffer in in_buffers { + backend.free_buffer(buffer).expect("free input"); + } + backend.free_buffer(out_buffer).expect("free output"); + result + }; + + // Fused: FP8 in, FP32 out, no BF16 tensor anywhere. + let fused_program = load( + &fp8e4m3_matmul_tosa(M as i32, K as i32, N as i32), + XDNA_TOSA_FP8_TARGET, + ); + let fused = run(&fused_program, &[&lhs_fp8, &rhs_fp8], M * N * 4); + + // Reference: promote each operand with the standalone CAST tier, then multiply in BF16. + let lhs_cast = load(&fp8e4m3_to_bf16_tosa(M * K), XDNA_TOSA_FP8_TARGET); + let lhs_bf16 = run(&lhs_cast, &[&lhs_fp8], M * K * 2); + let rhs_cast = load(&fp8e4m3_to_bf16_tosa(K * N), XDNA_TOSA_FP8_TARGET); + let rhs_bf16 = run(&rhs_cast, &[&rhs_fp8], K * N * 2); + let matmul = load( + &bf16_matmul_tosa(M as i32, K as i32, N as i32), + XDNA_TOSA_TARGET, + ); + let unfused = run(&matmul, &[&lhs_bf16, &rhs_bf16], M * N * 4); + + assert_eq!( + fused, unfused, + "fusing the promotion must not change a single result bit" + ); + + for program in [fused_program, lhs_cast, rhs_cast, matmul] { + backend.unload_program(program).expect("unload"); + } + backend.destroy_queue(queue).expect("destroy queue"); + backend.destroy_context(context).expect("destroy context"); +} + #[test] fn tosa_bf16_matmul_runs_on_the_npu() { let Some(backend) = backend() else { return }; diff --git a/crates/virtio-accel-xdna/tests/xdna.rs b/crates/virtio-accel-xdna/tests/xdna.rs index e832299..9b8839b 100644 --- a/crates/virtio-accel-xdna/tests/xdna.rs +++ b/crates/virtio-accel-xdna/tests/xdna.rs @@ -74,20 +74,28 @@ fn targets_survive_an_identity_round_trip() { } #[test] -fn fp8_capability_is_storage_conversion_only() { +fn fp8_target_never_produces_fp8() { + // FP8 is a *storage* encoding here: it is only ever consumed. Both admitted tiers read FP8 and + // write a wider type (BF16 for the standalone CAST, FP32 for the fused MATMUL), so no graph on + // this target may produce an FP8 value. MATMUL joined this surface with the fused tier; see + // `fused_fp8_matmul_capability_covers_its_graph`. assert_eq!(XDNA_TOSA_FP8_CAPABILITY.target, XDNA_TOSA_FP8_TARGET); assert!(XDNA_TOSA_FP8_CAPABILITY.supports_operator(Op::CAST)); - assert!(!XDNA_TOSA_FP8_CAPABILITY.supports_operator(Op::MATMUL)); for dtype in [ virtio_accel_tosa::DType::FP8E4M3, virtio_accel_tosa::DType::FP8E5M2, ] { assert!(XDNA_TOSA_FP8_CAPABILITY.supports_dtype(dtype, ValueRoles::INPUT)); assert!(!XDNA_TOSA_FP8_CAPABILITY.supports_dtype(dtype, ValueRoles::OUTPUT)); + assert!(!XDNA_TOSA_FP8_CAPABILITY.supports_dtype(dtype, ValueRoles::INTERMEDIATE)); } assert!( XDNA_TOSA_FP8_CAPABILITY.supports_dtype(virtio_accel_tosa::DType::BF16, ValueRoles::OUTPUT) ); + // Nothing on this target may produce or consume FP32 as anything but the fused result. + assert!( + !XDNA_TOSA_FP8_CAPABILITY.supports_dtype(virtio_accel_tosa::DType::FP32, ValueRoles::INPUT) + ); } #[test] @@ -201,3 +209,111 @@ fn compile_artifact_works_without_hrx() { assert_eq!(parsed.slot_bytes, [(ELEMENTS * 2) as u64; 2]); assert!(parsed.xclbin.starts_with(b"xclbin2")); } + +/// The fused tier's capability must cover the graph it admits, including the interior: a promoted +/// BF16 value that the descriptor does not admit as `INTERMEDIATE` would advertise a surface the +/// backend contradicts. +#[test] +fn fused_fp8_matmul_capability_covers_its_graph() { + use virtio_accel_tosa::DType; + + for dtype in [DType::FP8E4M3, DType::FP8E5M2] { + assert!(XDNA_TOSA_FP8_CAPABILITY.supports_dtype(dtype, ValueRoles::INPUT)); + } + // The promotion is graph-interior for the fused tier and a block output for the CAST tier. + assert!(XDNA_TOSA_FP8_CAPABILITY.supports_dtype(DType::BF16, ValueRoles::INTERMEDIATE)); + assert!(XDNA_TOSA_FP8_CAPABILITY.supports_dtype(DType::BF16, ValueRoles::OUTPUT)); + // The fused tier ends at the TOSA-mandated FP32 accumulator. + assert!(XDNA_TOSA_FP8_CAPABILITY.supports_dtype(DType::FP32, ValueRoles::OUTPUT)); + for op in [Op::CAST, Op::CONST, Op::MATMUL] { + assert!(XDNA_TOSA_FP8_CAPABILITY.supports_operator(op)); + } + let matmul = XDNA_TOSA_FP8_CAPABILITY + .operator(Op::MATMUL) + .expect("MATMUL capability"); + assert!( + matmul + .constraints + .contains(OperatorConstraints::ZERO_ZERO_POINTS) + ); +} + +/// The offline path compiles the fused graph to a container that binds FP8 operands directly and +/// never binds a BF16 tensor — the whole point of the tier. +#[test] +fn fused_fp8_matmul_compiles_to_a_wellformed_artifact() { + use virtio_accel_tosa::DType; + use virtio_accel_tosa_build::{OperatorKind, OwnedGraph, OwnedOperator, OwnedTensor}; + use virtio_accel_xdna::{PrecompiledArtifact, compile_artifact}; + + if std::env::var_os("VIRTIO_ACCEL_AMDXDNA_TOOLCHAIN").is_none() { + eprintln!("no XDNA toolchain configured; skipping fused offline compile test"); + return; + } + let (m, k, n) = (32i32, 64i32, 32i32); + let mut graph = OwnedGraph::new("main"); + graph + .push_tensor(OwnedTensor::new("lhs_fp8", vec![1, m, k], DType::FP8E4M3)) + .push_tensor(OwnedTensor::new("rhs_fp8", vec![1, k, n], DType::FP8E4M3)) + .push_tensor(OwnedTensor::new("lhs_bf16", vec![1, m, k], DType::BF16)) + .push_tensor(OwnedTensor::new("rhs_bf16", vec![1, k, n], DType::BF16)) + .push_tensor(OwnedTensor::constant( + "lhs_zp", + vec![1], + DType::BF16, + vec![0u8; 2], + )) + .push_tensor(OwnedTensor::constant( + "rhs_zp", + vec![1], + DType::BF16, + vec![0u8; 2], + )) + .push_tensor(OwnedTensor::new("output", vec![1, m, n], DType::FP32)) + .push_operator(OwnedOperator::new( + OperatorKind::Cast, + vec!["lhs_fp8".into()], + vec!["lhs_bf16".into()], + )) + .push_operator(OwnedOperator::new( + OperatorKind::Cast, + vec!["rhs_fp8".into()], + vec!["rhs_bf16".into()], + )) + .push_operator(OwnedOperator::new( + OperatorKind::Const, + vec![], + vec!["lhs_zp".into()], + )) + .push_operator(OwnedOperator::new( + OperatorKind::Const, + vec![], + vec!["rhs_zp".into()], + )) + .push_operator(OwnedOperator::new( + OperatorKind::MatMul, + vec![ + "lhs_bf16".into(), + "rhs_bf16".into(), + "lhs_zp".into(), + "rhs_zp".into(), + ], + vec!["output".into()], + )) + .push_input("lhs_fp8") + .push_input("rhs_fp8") + .push_output("output"); + let tosa = graph + .build(XDNA_TOSA_FP8_TARGET) + .expect("build fused fp8 matmul graph"); + let container = + compile_artifact(&tosa, XDNA_TOSA_FP8_TARGET).expect("compile fused fp8 matmul"); + let parsed = PrecompiledArtifact::parse(&container).expect("valid container"); + assert_eq!(parsed.inputs, 2); + assert_eq!(parsed.outputs, 1); + assert_eq!( + parsed.slot_bytes, + vec![(m * k) as u64, (k * n) as u64, (m * n * 4) as u64], + "FP8 operands bind one byte per element and no BF16 tensor is bound" + ); +} From 30127778137f66a9d59cb1bb437b9500ff7cf764 Mon Sep 17 00:00:00 2001 From: Aryan Ravishankar Date: Thu, 27 Aug 2026 20:32:34 -0500 Subject: [PATCH 2/2] Gate the fused offline-compile test on unix `compile_artifact` drives the helper subprocess in its own process group and is `cfg(unix)`, so the new test must carry the same gate the existing offline-compile test does. The Windows job caught this as an unresolved import. Co-Authored-By: Claude Opus 5 --- crates/virtio-accel-xdna/tests/xdna.rs | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/crates/virtio-accel-xdna/tests/xdna.rs b/crates/virtio-accel-xdna/tests/xdna.rs index 9b8839b..c2c6fa7 100644 --- a/crates/virtio-accel-xdna/tests/xdna.rs +++ b/crates/virtio-accel-xdna/tests/xdna.rs @@ -240,6 +240,10 @@ fn fused_fp8_matmul_capability_covers_its_graph() { /// The offline path compiles the fused graph to a container that binds FP8 operands directly and /// never binds a BF16 tensor — the whole point of the tier. +/// +/// Unix-only for the same reason as `compile_artifact_works_without_hrx`: the compiler helper is a +/// subprocess in its own process group, so `compile_artifact` is not offered on other platforms. +#[cfg(unix)] #[test] fn fused_fp8_matmul_compiles_to_a_wellformed_artifact() { use virtio_accel_tosa::DType;