Exploring the scalable matrix extension of the Apple M4 processor

Apple's M4 chip, first shipped in the 2024 iPad Pro, is the first publicly available device to support ARM's scalable matrix extension (SME), replacing the proprietary AMX instruction set Apple has used in its devices since 2019. An independent developer investigation is microbenchmarking the M4 SME hardware to measure peak compute throughput for common matrix and vector operations and memory transfer rates, while a team at the University of Jena is building a high-performance GEMM implementation on the same extension. The project currently runs only on an external M4 iPad Pro paired with a Mac, with the author noting that M4 Macs are expected soon and inviting pull requests.

The 2024 iPad Pro, featuring the Apple M4 chip, is the first device available to the public that supports ARM's scalable matrix extension SME . Although Apple has included a matrix accelerator in its devices since 2019, it used a proprietary instruction set called AMX<sup 1 user-content-fn-1-10b5e5c28ae4286f4e64d1da8b0c4e3f </sup . The developers had to use Apple-provided numerical frameworks to take advantage of the hardware. The M4 chip with SME changes this by allowing programmers to write low-level code that directly targets the matrix hardware, bringing potential performance improvements to scientific and machine learning algorithms. My investigation aims to explore the M4 SME hardware and study the SME instruction set. I am particularly interested in how these features can be used to accelerate vector operations. The initial experiments study the peak compute throughput for common matrix and vector operations and the memory transfer rates using various methods. My methodology relies on generating multiple variants of microbenchmark code over a wide range of parameters and studying the results see the Python code in tools/ for how the benchmarks are generated . I expect to extend this project as new aspects of the hardware are tested. Pull requests and commentary are welcome Another in-depth investigation of M4 SME is provided by a team at the University of Jena, who are building a high-performance GEMM implementation <sup 2 user-content-fn-2-10b5e5c28ae4286f4e64d1da8b0c4e3f </sup . It is an excellent read if you are interested in the topic. - An M4 iPad Pro with Developer Mode activated and paired with your Mac - Apple Xcode with installed Command Line Tools - XcodeGen https://github.com/yonaskolb/XcodeGen brew install xcodegen - Python3 with PyYAML package to generate the tests - Quarto https://quarto.org +R/tidyverse to render the reports Note: this project currently only runs on an external M4 iPad Pro. As M4 Macs are expected to be available soon, it would be desirable to run these tests locally. Pull requests are welcome After downloading or cloning the project, perform the following steps setup a paired M4 iPad and the code signing information ./setup compile the project and install it on the iPad make build run the tests — this will take a while make run copy the benchmark results make copy results After the testing is done, the generated JSON reports are copied from the iPad and placed in the results/ folder. If you have R and Quarto installed, you can also build the R-Markdown reports using make reports . By default, both single-core and multi-core tests are executed. This can take a long time. If you are only interested in peak single-core rates, you can change the second line in src/tests.swift from let multiCoreTests = true to false and rebuild. Footnotes 1. https://github.com/corsix/amx https://github.com/corsix/amx - description of Apple proprietary AMX extensions ↩ user-content-fnref-1-10b5e5c28ae4286f4e64d1da8b0c4e3f 2. https://scalable.uni-jena.de/opt/sme/ https://scalable.uni-jena.de/opt/sme/ ↩ user-content-fnref-2-10b5e5c28ae4286f4e64d1da8b0c4e3f