PaperlayerView exact revision on arXiv

Source-linked AI summary

Demystifying ARM SME to Optimize General Matrix Multiplications

Chencheng Deng, Weiling Yang, Jianbin Fang, Dezun Dong

arXiv:2512.21473v1cs.DC

TL;DR

Existing SME-enabled GEMM libraries do not fully exploit the hardware, particularly for large matrices. MpGEMM combines SME-aware characterization with cache-aware partitioning, packing, and specialized kernels, achieving an average 1.23× speedup over Apple Accelerate on DeepSeek and LLaMA workloads.

  • Problem

    Existing SME-enabled GEMM libraries do not fully exploit SME, particularly for large matrices, despite GEMM's importance across high-performance computing and deep learning.

    Source: §I · p. 1, C. Motivation: Limitations of Existing… · p. 2

  • Method

    MpGEMM uses SME-guided cache-aware partitioning, dual-matrix packing with on-the-fly transposition and online packing, and specialized micro-kernels for multiple-precision GEMM.

    Source: §I · pp. 1–2, §VIII · p. 11

  • Results

    1.23× average speedup over vendor-optimized Apple Accelerate on real-world DeepSeek and LLaMA GEMM workloads.

    Source: §I · p. 2

  • Takeaways & Limitations

    The optimizations extend to mixed-precision GEMM with lower-precision inputs and higher-precision outputs, further improving the compute-to-memory ratio.

    Source: §I · pp. 1–2

    The identified shortcoming is specific to current SME-enabled GEMM libraries, particularly for large matrices.

    Source: C. Motivation: Limitations of Existing… · p. 2

Abstract

from arXiv · show

General Matrix Multiplication (GEMM) is a critical kernel in high-performance computing and deep learning. While modern architectures like ARM's Scalable Matrix Extension (SME) introduce dedicated hardware for matrix operations, existing linear algebra libraries fail to fully exploit its potential, particularly for large matrices. This paper presents MpGEMM, an open-source library that leverages key architectural features of SME to optimize GEMM across multiple precisions. Through a systematic characterization of SME, we derive optimization guidelines that inform our design. MpGEMM employs cache-aware partitioning, efficient data packing with on-the-fly transposition, and specialized micro-kernels that utilize multi-vector loads and all available tile registers. Evaluated on an Apple M4 Pro with real-world workloads from DeepSeek and LLaMA, MpGEMM achieves an average speedup of 1.23x over the vendor-optimized Apple Accelerate library and significantly outperforms other open-source alternatives.

I. INTRODUCTION

GEMM libraries on ARM SME leave performance untapped, especially for large matrices; MpGEMM characterizes the architecture and applies cache-aware partitioning, packing, and specialized micro-kernels. On an Apple M4 Pro, it outperforms Apple Accelerate on DeepSeek and LLaMA workloads and extends the design to mixed precision.

I. INTRODUCTION

  • Current SME GEMM libraries leave substantial performance untapped, particularly for large matrices, because of cache-utilization and memory-bandwidth bottlenecks.Simple loop structures poorly match the shared L2 cache, while existing packing and micro-kernel approaches underuse multi-vector loads.Source: §I · p. 1
  • MpGEMM combines analytical-model-guided cache-aware partitioning, dual-matrix packing, on-the-fly transposition, online packing, and micro-kernels using four-Z-register loads and all ZA tiles.Its design also dynamically distributes tasks across available SME units.Source: §I · pp. 1–2
  • 1.23× average speedup over Apple Accelerate, with additional average speedups of 3.96× over LIBXSMM, 4.69× over KleidiAI, and 5.7× over OpenBLAS.Evaluation used an Apple M4 Pro and real-world GEMM workloads from DeepSeek and LLaMA.Source: §I · p. 2
  • The optimizations extend to mixed-precision GEMMs with lower-precision inputs and higher-precision outputs, reaching up to 94% of SME peak performance.The paper identifies mixed precision as a way to improve the compute-to-memory ratio.Source: §I · p. 2, A. ARM Scalable Matrix Extension (SME) · p. 2

A. ARM Scalable Matrix Extension (SME)

  • SME builds on SVE and supports matrix operations through outer-product instructions and a dedicated two-dimensional ZA accumulator.SME2 adds multi-vector instructions for simultaneous operations on multiple scalable vector registers.Source: §I · p. 1, A. ARM Scalable Matrix Extension (SME) · p. 2

B. GEMM Algorithm and its Mapping to SME

GEMM can be computed through inner-product or outer-product traversal; the outer-product pattern aligns naturally with SME’s FMOPA instruction. The paper identifies incomplete use of SME hardware, especially for large matrices, as a limitation of existing implementations.

B. GEMM Algorithm and its Mapping to SME

  • The outer-product method follows the K dimension, accumulating rank-1 products into C, and aligns naturally with SME’s FMOPA instruction.It offers a higher compute-to-memory ratio than the inner-product method, which computes each C element as a row-column dot product.Source: B. GEMM Algorithm and its Mapping to SME · p. 2

C. Motivation: Limitations of Existing Implementations

1) Ineffective cache hierarchy utilization:

The section identifies cache fit, packing layout, and multi-vector load granularity as constraints on large-GEMM performance, then uses Apple M4 Pro microbenchmarks to motivate optimization guidelines.

1) Ineffective cache hierarchy utilization:

  • Existing open-source libraries use simple three-level loops that can mismatch shared L2 capacity; six-level blocking tunes tiles to cache, TLB, and bandwidth.The six-level structure is guided by an analytical model to maximize cache locality after packing.Source: 1) Ineffective cache hierarchy… · p. 3

2) Suboptimal data packing method:

  • Packing both input matrices improves access stride and spatial locality, while on-the-fly transposition and first-round online packing mitigate packing overhead.The prior approach packed only one input matrix, and could degrade when the unpacked matrix exceeded shared L2 capacity.Source: 2) Suboptimal data packing method · p. 3

3) Underutilization of multi-vector load instructions:

  • Four-Z-register groups can reach 900 GB/s below an 8 MB load size, whereas a single-Z-register load delivers 230 GB/s; prior methods use only one or two registers.The proposed micro-kernel restructuring consistently uses four-Z-register groups.Source: 3) Underutilization of multi-vector… · p. 3

III. MICROBENCHMARKS

  • Apple M4 Pro microbenchmarks characterize SME performance properties to guide MPGEMM’s optimization strategies.The experiments are conducted on the Apple M4 Pro chip.Source: §III · p. 3

A. Charaterizing Computational Throughput

B. Charaterizing Memory Bandwidth

The memory-bandwidth measurements show that load grouping and working-set size matter; the resulting principles inform data partitioning and GEMM optimization on the M4 Pro.

B. Charaterizing Memory Bandwidth

  • Four-Z-register groups reach 900 GB/s at 256 bytes, compared with 375 GB/s for one register and 625 GB/s for two registers at their respective group sizes.As bus pressure increases, one- and two-register group bandwidth falls to 230 GB/s and 460 GB/s, respectively.Source: B. Charaterizing Memory Bandwidth · p. 4
  • Beyond 8 MB of loaded data, bandwidth drops to about 230 GB/s and falls to 145 GB/s beyond 16 MB; write-back peaks at 230 GB/s.The passage identifies 8 MB as the effective cacheable working set for optimal bandwidth and calls for four-register loads and submatrices within that footprint.Source: B. Charaterizing Memory Bandwidth · p. 4

C. Architectural Implications and Optimization Guidelines

IV. OPTIMIZING GENERAL MATRIX MULTIPLICATION

  • The design discussion uses row-major FP32 GEMM on a platform with SVL=512 bits as a running example and says the methodology generalizes to other platforms, precisions, and storage orders.Source: §IV · p. 4

A. Design Overview

MPGEMM combines multilevel blocking and cache-aware packing with SME-specific data transformations to improve cache locality and reduce packing overhead.

A. Design Overview

  • MPGEMM uses a multilevel blocking strategy built on the classical Goto algorithm to exploit the memory hierarchy.The design is depicted as an algorithm with multiple loop levels.Source: A. Design Overview · p. 4, Figure 5 · p. 5
  • The algorithm packs A and B into contiguous buffers, then applies outer-product micro-kernels to update ZA-resident C blocks; it parallelizes M and N, not reduction dimension K.Tiling parameters mc, nc, and kc are guided by an analytical model; both packing routines and micro-kernels use native SME assembly.Source: A. Design Overview · pp. 4–5

B. Cache-Aware Partitioning and Packing

  • Analytical tiling guides partitioning across cache and TLB constraints, with the working blocks kept under 8 MB and the buffers aligned to cache lines.The design reserves L2 space for Bc and C so Bc is not evicted during the next L3 iteration.Source: B. Cache-Aware Partitioning and Packing · p. 5
  • The L2 compute-to-memory ratio is maximized subject to an L2-size constraint after TLB-aware selection of kc and fixed micro-kernel tile sizes.The stated micro-kernel dimensions are mr=16 and nr=64.Source: B. Cache-Aware Partitioning and Packing · p. 5
  • ZA horizontal and vertical slices transpose A during packing into column-major panels, while B’s first-round packing overlaps memory operations with FMOPA computation.The A transformation avoids inefficient gather operations; four Z registers are loaded per group during packing and computation.Source: B. Cache-Aware Partitioning and Packing · pp. 5–6, Figure 6 · p. 6

C. High-Performance Micro-Kernel Design

The micro-kernel design targets full ZA-tile use and four-register loads, with software pipelining to overlap load latency and computation.

C. High-Performance Micro-Kernel Design

  • At SVL=512 bits, 32×32, 16×64, and 64×16 configurations can use all four ZA tiles; OpenBLAS and KleidiAI use only two tiles at boundaries.LIBXSMM uses 16×64 and 64×16 kernels for edge handling.Source: C. High-Performance Micro-Kernel Design · p. 6, Figure 7 · p. 6
  • A 16×64 main micro-kernel and 64×16 edge micro-kernel use all ZA tiles and four-register loads, addressing the lower bandwidth of two-register groups.For column-major layout, the main and edge kernel shapes are reversed.Source: C. High-Performance Micro-Kernel Design · p. 6, Figure 7 · p. 6
  • The kernel loads Cr, Ar, and Br in four-Z-register batches and accumulates outer products into ZA tiles.The algorithm passage shows the main micro-kernel and its batched loads and FMOPA updates.Source: C. High-Performance Micro-Kernel Design · p. 6
  • Loop unrolling and software pipelining launch later loads early so prior FMOPA operations overlap load latency, while extra Z registers buffer data.The passage links this scheduling approach to reducing pipeline inefficiency and scheduler pressure.Source: C. High-Performance Micro-Kernel Design · p. 7

A. Design Principals

The mixed-precision design extends the FP32 workflow by adapting packing and micro-kernels to lower-precision layouts and computation. It rearranges both inputs for contiguous access and uses a four-way-load micro-kernel to update four FP32 accumulator tiles.

A. Design Principals

  • The FP16 workflow retains the FP32 GEMM structure but must adapt packing and micro-kernels to mixed-precision layouts and computation.SME supports lower-precision inputs accumulated into higher-precision outputs, including FP16→FP32.Source: A. Design Principals · p. 7

B. Packing Strategies for Lower Precision

  • FP16 packing rearranges A into a column-major layout and vertically interleaves adjacent rows of B for contiguous, aligned micro-kernel access.A is transposed on the fly; SME ZIP interleaves B values into the packed buffer.Source: B. Packing Strategies for Lower… · p. 7

C. Mixed-Precision Micro-Kernel

  • The FP16 micro-kernel uses a 16×64 tile and four FP32 ZA tiles to accumulate results from FP16 FMOPA operations.Four-Z-register loads are masked when fewer than 64 output values are required.Source: C. Mixed-Precision Micro-Kernel · p. 7
  • The micro-kernel loads packed A and B data in batches, updates four ZA tiles with FMOPA, and stores the accumulated result back to C.Register buffering and loop unrolling support instruction-level parallelism and throughput.Source: C. Mixed-Precision Micro-Kernel · p. 8

VI. EVALUATION

The evaluation uses an Apple M4 Pro and workloads derived from DeepSeek and LLaMA. It covers large, skinny, and irregular GEMMs.

1) Hardware Platform:

  • The Apple M4 Pro has a shared 16MB L2 cache per five-P-core cluster, a 4MB L2 cache shared by E-cores, and one shared SME unit per cluster.Experiments ran on macOS 15.1 using AppleClang 16.Source: 1) Hardware Platform · p. 8, Table II · p. 8

2) Workloads:

  • Workloads derived from DeepSeek and LLaMA span large and skinny matrices, with additional irregular dimensions not divisible by micro-kernel sizes.Matrices were initialized with random values.Source: 2) Workloads · p. 8, Table III · p. 8

3) Baselines:

On a single SME unit, MPGEMM outperforms Accelerate in both storage layouts and exceeds three open-source alternatives on average. The evaluation compares libraries using their supported computational modes and measures row- and column-major FP32 GEMM.

3) Baselines:

  • The baselines are Accelerate, LIBXSMM, KleidiAI, and OpenBLAS, with MPGEMM configured to match their supported modes and storage orders.Experiments were repeated five times and reported as arithmetic means after warm-up runs.Source: 3) Baselines · p. 8

B. GEMM Performance

  • The FP32 evaluation covers row-major and column-major storage on a single SME unit.The experiment evaluates gains from the GEMM optimizations described in the paper.Source: B. GEMM Performance · p. 8

1) Single SME Unit Performance:

  • MPGEMM achieves average speedups of 2.85× over OpenBLAS, 2.34× over KleidiAI, and 1.95× over LIBXSMM.These comparisons are reported for the single-SME-unit evaluation in Fig. 10.Source: 1) Single SME Unit Performance · p. 8, Figure 10 · p. 9
  • MPGEMM uses cache-aware six-level blocking and packs both inputs, addressing shared-L2/TLB inefficiency and baseline limitations in multi-vector loads.LIBXSMM can surpass Accelerate when unpacked data fits L2, but performance degrades when the working set exceeds it.Source: 1) Single SME Unit Performance · p. 9
  • MPGEMM outperforms Accelerate by 1.21× for row-major and 1.18× for column-major layouts on average.Fig. 10 reports the single-SME-unit performance.Source: 1) Single SME Unit Performance · p. 9, Figure 10 · p. 9

2) Multiple SME Unit Performance:

MPGEMM shows gains over Accelerate across multi-unit FP32 and FP64 GEMM, and over KleidiAI in mixed precision. The section also reports performance on irregular dimensions and high utilization of INT8 SME throughput.

2) Multiple SME Unit Performance:

  • With up to two threads per cluster, MPGEMM achieves average single-precision speedups of 1.24× and 1.22× over Accelerate for row- and column-major layouts.The M4 Pro has a limited number of SME units, so the evaluation focuses on parallel efficiency rather than scalability.Source: 2) Multiple SME Unit Performance · pp. 9–10
  • MPGEMM achieves a 1.18× speedup over Accelerate for multi-core row-major double-precision GEMM, despite FP64 throughput being one-quarter of FP32.Accelerate is the only library identified as supporting FP64 GEMM on SME.Source: 2) Multiple SME Unit Performance · p. 10

3) Irregular Matrices:

  • MPGEMM consistently outperforms alternatives on irregular dimensions using efficient packing and optimized main and edge micro-kernels.The evaluation varies M and N from 80 to 200 and sets K to 25600.Source: 3) Irregular Matrices · p. 10

C. Mixed-Precision GEMM Performance

  • FP16-input GEMM achieves an average 1.7× speedup over KleidiAI across the mixed-precision throughput tests.The paper attributes the improvement to refined partitioning and high-bandwidth four-way loads.Source: C. Mixed-Precision GEMM Performance · p. 10
  • INT8 GEMM achieves an average 2× speedup over the FP16 GEMM baseline, while SMOPA provides up to 4010 GOPS on a performance SME unit.The passage notes that KleidiAI uses NEON for INT8 packing, whose overhead is not amortized for flat GEMMs.Source: C. Mixed-Precision GEMM Performance · p. 10
  • Multi-core INT8 GEMM reaches 7856 and 7882 GOPS on tests 14 and 18, about 94% of measured SMOPA peak performance.The comparison spans FP32, FP16, and INT8 input precisions.Source: C. Mixed-Precision GEMM Performance · p. 10, Figure 14 · p. 10

D. Optimization Breakdown

MpGEMM’s cache-aware partitioning and packing and multi-vector loading improve performance, while first-round online packing provides limited benefit. Overall, MpGEMM averages a 1.23× speedup over Accelerate and significantly outperforms other open-source solutions.

D. Optimization Breakdown

  • 1.62× and 1.17× average speedups over LIBXSMM came from partitioning and packing, and multi-vector loading, respectively.The partitioning and packing convert the three-loop structure into six-level blocking, with tiling tuned for shared L2 residency.Source: D. Optimization Breakdown · pp. 10–11
  • First-round online packing had limited benefit because FMOPA cycles could not fully hide memory-write overhead.For IDs (3, 9, 15) with K = 512, packing both matrices underperformed LIBXSMM, while four-way loading and online packing together still improved MpGEMM performance.Source: D. Optimization Breakdown · p. 11

VII. RELATED WORKS

  • Goto et al.’s multi-level blocking and packing established a foundation for BLAS libraries, while LIBXSMM extended the approach with JIT micro-kernels for small matrices.The paper frames SME optimization as an area where mainstream linear algebra libraries remain suboptimal.Source: §VII · p. 11
  • Prior SME GEMM libraries leave large-scale performance opportunities; MpGEMM combines cache-aware partitioning, dual packing, and four-Z-register grouped loads.The related-work discussion contrasts these features with approaches limited to small GEMMs or lacking memory management or SME2 vector-group processing.Source: §VII · p. 11

VIII. CONCLUSION

  • 1.23× average speedup over Accelerate, with significantly better performance than other open-source solutions.The conclusion reports this result for MpGEMM’s multiple-precision GEMM optimizations on ARM SME CPUs.Source: §VIII · p. 11

Related papers

Loading 2512.21473v1…