AMX Benchmarks
This document coalesces data about real-world performance of the Apple AMX coprocessor. The motivating use case is electronic structure calculations (DFT simulations), which use complex-valued matrix multiplications and eigendecompositions. Interleaved complex numbers incur additional overhead compared to split complex numbers, but BLAS only accepts the interleaved format. This format underutilizes both NEON and AMX units.
Table of Contents
LOBPCG
TODO: Rewrite this, now that I understand what's happening in more detail. Real-space algorithms may be the only feasible ones.
According a recent research paper (2023), the LOBPCG iterations can be partitioned into stages with different precisions. The first iterations use single precision, while the last iterations use double precision. The researchers used a consumer RTX A4000 with a 1:32 ratio of FP64:FP32 compute power. They achieved 6-8x speedup over GPU FP64 and negligible accuracy loss. A host CPU was used for ZHEEV, but that operation only consumed 3-10% of the total time.
For Apple silicon, the best adaptation of this algorithm would use CPU and GPU simultaneously. The AMX would not perform the majority of operations, but its presence would still be important. Below is a tentative illustration of the scheme*:
- 65% of iterations: GPU FP32 (CHEMM) + GPU FP32 (CHEEV)
- 15% of iterations: GPU FP32 (CHEMM) + GPU double-single (ZHEEV)
- 15% of iterations: AMX FP32 (CHEMM) + NEON FP64 (ZHEEV)
- 5% of iterations: AMX FP64 (ZHEMM) + NEON FP64 (ZHEEV)
*De-interleaves the complex multiplications (CHEMM, ZHEMM) into four separate multiplications of their real and complex parts (SGEMM, DGEMM). This improves ALU utilization with the AMX and
simdgroup_matrix.
Using 75% of the performance cores' NEON, all of the AMX's FP64 GEMM compute, and all of the GPU's eFP64, the M1 Max could reach 1658 GFLOPS FP64. This is 4.3x faster than 100% of the performance cores' NEON alone and 2.8x faster than the GPU's eFP64 alone. However, using all of that simultaneously may cause thermal throttling, decreasing performance by up to 1.5x.
In another scheme, the AMX would perform most of the computations. Matrix sizes used for GEMM exceed the matrix sizes used for ZHEEV. ZHEEV is kn^3, where n is the number of valence electrons. Meanwhile, GEMM is kLn^2, where L is the number of grid cells. There are significantly more grid cells than valence electrons, by multiple orders of magnitude.
- 65% of iterations: AMX FP32 (CHEMM) + NEON FP32 (CHEEV)
- 30% of iterations: AMX FP32 (CHEMM) + NEON FP64 (ZHEEV)
- 5% of iterations: AMX FP64 (ZHEMM) + NEON FP64 (ZHEEV)
New scheme:
| AMX Vector | NEON Vector | GPU Matrix | GPU Vector | AMX Matrix | |
|---|---|---|---|---|---|
| Max Clock @ Full Utilization | 3.228 GHz | 3.132 GHz | 1.296 GHz | 1.296 GHz | 3.228 GHz |
| Max Observed Power | ~12 W ??? | 43.9 W | 52 W | 51.1 W | ~12 W ??? |
| Max Observed GFLOPS F32 | TBD | 655 | 9258 | 10400 | 2746 |
| Max Observed GFLOPS F64 | TBD | 352 | 0 | 0 | 700 |
| Max Theoretical GFLOPS FFMA32 | 413 | 801 | 9437 | 10617 | 3305 |
| Max Theoretical GFLOPS FDIV32 | 0 | 200 | 0 | 884 | 0 |
| Max Theoretical GFLOPS FSQRT32 | 0 | 200 | 0 | 663 | 0 |
| Max Theoretical GFLOPS FFMA64 | 206 | 400 | TBD | 589 | 826 |
| Max Theoretical GFLOPS FDIV64 | 0 | 100 | 0 | 183 | 0 |
| Max Theoretical GFLOPS FSQRT64 | 0 | 100 | 0 | 189 | 0 |
Linear Algebra Benchmark: GFLOPS/k
GFLOPS is not a plural noun. GFLOPS is a rate: (G)Billion (FL)Floating Point (OP)Operations per (S)Second. The term GFLOPS/second is often used to remove ambiguity, except that translates to GFLOP/second/second. Data throughput is a measure of speed - speed requires units of velocity, not acceleration. Therefore, this repository uses the original term GFLOPS.
GFLOPs is a plural noun. Occasionally, I use GFLOPs to specify the number of floating-point operations required for a linear algebra operation. The capitalization of s will distinguish the metric from GFLOPS. There are not many other concise, consistent ways to describe both of these terms.
TODO: Explain O(kn^3), uncertainty in computational complexity, universal measure of time-to-solution (agnostic of precision or algorithm), why I used GFLOPS/0.25k for complex-valued operations to normalize for ALU utilization
GFLOPS/k = (matrix dimension)^3 / (time to solution)
Imagine a processor has 1000 GFLOPS, uses 10 watts.
OpenBLAS GEMM: real GLOPS/k = 800, but complex GFLOPS/k = 190
Real has 80% ALU / 8.0 watts.
Complex has 76% ALU / 7.6 watts, not 19% ALU / 1.9 watts.
Both operations have ~80% ALU and ~8 watts.
However, GFLOPS/0.25k = 4 * (GFLOPS/k) ~ 760
80% ALU is much closer to 76%, and shows that complex is 4% slower,
but not because it requires more computations. Also, you would think
it's 4% **faster**, because it has **more** arithmetic intensity.
GFLOPS/0.25k is a fairer, more insightful comparison.
k_complex = 4k_real
k_real = 0.25k_complex
Real: GFLOPS = GFLOPS/k * k_real
Complex: GFLOPS = GFLOPS/0.25k * 0.25k_complex
Non-hybrid algorithms (all on one processor, either the CPU cores, AMX units, or GPU cores)
| Operation | kreal | OpenBLAS GFLOPS/k | Accelerate GFLOPS/k | Metal GFLOPS/k | NEON % | AMX % | GPU % | Max GFLOPS |
|---|---|---|---|---|---|---|---|---|
| SGEMM | 2 | 362.2 | 1327.4 | 4629.0 | 84.4% | 85.4% | 87.2% | 9258.0 |
| DGEMM | 2 | 176.2 | 337.9 | - | 90.7% | 87.0% | - | 675.8 |
| ZGEMM | 2 | 148.4 | 223.6 | - | 76.4% | 57.6% | - | 447.2 |
| SSYEV | TBD | 4.54 | 12.9 | - | TBD | TBD | - | TBD |
| DSYEV | TBD | 4.57 | 7.74 | - | TBD | TBD | - | TBD |
| ZHEEV | TBD | 6.76 | 5.48 | - | TBD | TBD | - | TBD |
| SPOTRF | ||||||||
| DPOTRF | ||||||||
| ZPOTRF | ||||||||
| STRSM | ||||||||
| DTRSM | ||||||||
| ZTRSM |
GFLOPS/k for each operation used in quantum chemistry. This metric compares each operation's execution speed regardless of the algorithm used to perform it, or the formula used to estimate GFLOPS. Complex-valued operations use GFLOPS/0.25k to directly compare ALU utilization to real-valued operations. For every operation listed so far, complex-valued versions are slower because they must de-interleave the numbers before processing them.
ZHEEV achieved maximum performance on Accelerate with MRRR. All other eigendecompositions use the divide and conquer algorithm. Although OpenBLAS beats Accelerate with asymptotically large matrices, Accelerate is faster for the matrix sizes typically encountered in DFT.
Related Work
| ISA Documentation | Performance Documentation | OSS GEMM Libraries | |
|---|---|---|---|
| Apple AMX | corsix/amx | philipturner/amx-benchmarks | xrq-phys/blis_apple |
| Apple GPU | dougallj/applegpu | philipturner/metal-benchmarks | philipturner/metal-flash-attention |