BF16 Component-Product Emulation of FP32 and FP64 GEMM on Intel AMX
Bing Cui, Yu Liu
原文摘要
Modern CPUs increasingly integrate high-throughput matrix engines optimized for low-precision AI workloads, while many scientific computing applications still rely on FP32 and FP64 GEMM to meet their numerical accuracy requirements. This paper presents a method that combines BF16 matrix products on Intel Advanced Matrix Extensions (AMX) to emulate real square FP32 and FP64 GEMM. For FP32 GEMM, each element in both input matrices is decomposed into three BF16 components, and six selected component products are evaluated to target FP32-level accuracy comparable to oneMKL SGEMM. For FP64 inputs within the supported BF16 exponent range, the method uses a simplified fixed six-slice Ozaki decomposition. Each retained BF16 product is first produced in FP32, then widened and accumulated in FP64. Four product-count settings retain 6, 10, 15, or 21 component products, exposing the accuracy-performance tradeoff relative to oneMKL DGEMM. The FP32 schedule reuses component data and retains partial sums in AMX tiles. AMX-FP32 outperforms oneMKL SGEMM at tested orders $N\ge4096$ on 64 threads and also achieves gains for smaller matrices at lower thread counts. AMX-FP64 variants retaining fewer component products can outperform oneMKL DGEMM at sufficiently large orders; more products improve accuracy at additional cost.
cs.MS · cs.DC