One XDNA2 Kernel, Many Tensor Shapes
Earlier this year, we published Hello XDNA!, our guide to programming the neural processing units in AMD Ryzen AI chips. We have now updated its XDNA2 example so that key dimensions are supplied at runtime instead of fixed in the assembly code. The same assembly implementation can therefore handle a range of supported tensor shapes.
New step-by-step diagrams make the design easier to follow. They show how the kernel divides the calculation, shares its limited on-chip storage, and overlaps arithmetic with data movement. We also replaced the earlier fixed-size instruction schedule after discovering a subtle conflict at the accumulator write port that could silently discard a write.
Beyond the kernel example, the update expands and corrects our inferred instruction-set tables, fixes the BFP16 format description, improves the assembly listings, and addresses smaller errors and terminology issues throughout the site. Read the updated XDNA guide, check out the full changelog, or explore the sources and benchmarks.
Hello XDNA!
AMD’s Ryzen AI chips contain XDNA neural processing units, unique spatial dataflow architectures with VLIW cores and dedicated matrix instructions. We are documenting these microarchitectures in detail: from XDNA1’s BF16 4×8×4 matrix operations at 256 FLOPs/cycle to XDNA2’s BFP16 8×8×8 operations at 1024 FLOPs/cycle. Our website covers the ISA, register files, operation latencies, and hand-optimized assembly kernels. The kernels achieve 398 BF16 GFLOPS (86% of peak) on an XDNA1 compute tile and 1760 BFP16 GFLOPS (95% of peak) on an XDNA2 tile. If you are interested in low-level NPU programming, check out the sources as well.
SME Kernels
The Scalable Matrix Extension (SME) made its debut in the M4 system-on-a-chip in the 2024 iPad Pro. Since then, more products have become available with SME support. Following our initial SME sprint, we upstreamed an SME code generator for tensor processing primitives to the LIBXSMM library.
Small matrix-matrix multiplications are one of the supported primitives, and the code generation can be tested with a few commands:
git clone https://github.com/libxsmm/libxsmm.git
cd libxsmm; make -j BLAS=0
cd samples/xgemm; make -j
./gemm_kernel F32 F32 F32 F32 512 512 512 512 512 512 \
1 1 0 0 0 1 0 0 0 nopf nobr 0 1 10000 0
On a 2024 Mac mini with an M4, this results in a performance of about 1833 GFLOPS in FP32 arithmetic:
------------------------------------------------
RUNNING (512x512) X (512x512)^T = (512x512)
a:F32, b:F32, comp:F32, c:F32, BR=1
------------------------------------------------
[...]
1.464460s for libxsmm
1832.998967 GFLOPS for libxsmm
max. error: 0.000000
------------------------------------------------
[...]
For more details, see our SME web page, our SME paper, and our paper on tensor processing primitives.