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..c2c6fa7 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,115 @@ 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. +/// +/// 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; + 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" + ); +}