Fuse the FP8 promotion into MATMUL so BF16 never reaches DDR

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

crates/virtio-accel-xdna/compiler/xdna_compile.py

@@ -19,6 +19,9 @@
      "elements": <n>, "device": "npu2", "fold_ddr_addr_offset": false}
     {"op": "MATMUL", "in_dtype": "bf16", "out_dtype": "f32", "m": <M>, "k": <K>, "n": <N>,
      "device": "npu2", "fold_ddr_addr_offset": false}
+    {"op": "MATMUL", "in_dtype": "fp8e4m3" | "fp8e5m2", "out_dtype": "f32", "m": <M>, "k": <K>,
+     "n": <N>, "tile_m": <TM>, "tile_k": <TK>, "tile_n": <TN>, "max_dim": <D>, "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": <M>, "k": <K>, "n": <N>,
      "left_zero_point": <i8>, "right_zero_point": <i8>, "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 <cstdint>
 
 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<uint16_t>((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"),

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,

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))]