ARM Matrix Multiplication with NEON Assembly on M2
Exploring ARM assembly and NEON instructions for matrix multiplication on Apple M2 chips to boost Gflops performance.
Stock photo for illustration only, not from the actual event
- Developer explores writing custom ARM assembly for matrix multiplication on an M2 MacBook Air
- Naive implementation yields 1.98 Gflops, jumping to 7.86 Gflops with NEON SIMD
- Applying tiling techniques and expanding tile sizes significantly enhances compute intensity
- Laying the groundwork for FP16 calculations and upcoming SME or SVE2 experiments
Writing neural network kernels is an area many developers revisit, and while GPU tooling like CUDA is highly advanced, ARM-based computers continue to proliferate in cloud environments. This raises the question of how to construct a matrix multiplication routine using ARM assembly language that runs efficiently. Using a MacBook Air powered by an M2 processor—which features a theoretical single-threaded f32 matmul limit of 112 GFlops—as the primary testing ground, the exploration begins with basic assembly concepts.
Performance benchmarks in this project stem from shape sweeps ranging between 64x64 and 1024x1024, running 20 iterations per step and averaging the output values for each kernel. Starting with a naive assembly implementation, the performance reached 1.98 Gflops, translating to roughly 1.8% of the theoretical maximum and demonstrating that C compilers still manage preliminary code optimizations quite effectively.
Stock photo for illustration only, not from the actual event
Shifting toward SIMD execution involves loading four values from matrix B into a single 128-bit register while broadcasting a single value from matrix A into another 128-bit register, executing four fma operations simultaneously across four output values using the fmla instruction. This adjustment brings performance up to 7.86 Gflops—roughly a fourfold improvement—with the structural caveat that N must remain divisible by four.
To increase compute intensity per load operation, classical GPU-inspired tiling techniques are introduced. Leveraging 32 available 128-bit floating-point registers allows for an 8x8 output tile utilizing 16 registers for output and another 16 for inputs. This structure executes 16 loads and 64 fmla operations per K-dimension iteration while unrolling K by 4. Furthermore, interleaving specific B data loads with fmla instructions acts as a lightweight form of double buffering, securing an additional performance gain of around 5 Gflops.
Optimizing matrix multiplication on ARM CPUs via NEON provides deep insights into vector register mechanics and memory cache management. Tiling helps mitigate memory bottleneck issues by maximizing register reuse, serving as a crucial step in extracting peak performance from Apple Silicon architectures before advancing to heavier instruction sets like SVE2 or SME.
Ultimately, transitioning to FP16 precision enables an 8x16 tile structure, doubling the number of computed values using the same register count. This concludes the initial exploration, paving the way for upcoming experiments involving integer matrix multiplication and advanced instruction sets like SME or SVE2 in future updates.
Source: Dev.to
Found something wrong in this article? Report an issue with this article
Comments
Leave a Comment