Tridiagonal Solvers: MLX
Using the integrated GPU to solve tridiagonal systems on Apple Silicon
Introduction
In the previous articles we explored three different algorithms for solving tridiagonal systems. In a single threaded implementation we would expect Thomas to be the fastest solver, followed by CR, and then PCR simply because of the number of mathematical operations each is doing. However, our benchmarks revealed that the reality on modern hardware is a bit more complicated. Even on a single core, implementations can leverage SIMD capabilities to get effective parallelization of some operations in CR and PCR. In addition, we can spread work across multiple cores in CR and PCR to increase throughput. As a result, on the MacBook Pro M5 used to generate these benchmarks, the Numba CR and Parallel Numba CR implementations were fastest at large \(N\).
The Apple M-series chips have additional capability to delegate computational work to the GPU for even more parallelization. In this article we see if utilizing this advanced hardware can give us a performance boost over the CPU-based solvers.
Apple Silicon and Unified Memory
Apple’s M-series chips represent a significant departure from conventional heterogeneous workstation architectures. Rather than housing a discrete GPU connected to the CPU via a relatively narrow PCIe bus, Apple Silicon integrates high-performance CPU cores, a high-throughput GPU, and hardware neural accelerators onto a single piece of silicon using a Unified Memory Architecture (UMA).
Under UMA, the CPU and GPU share a single, wide, high-bandwidth memory pool without physical segregation into separate system RAM and dedicated VRAM. For researchers running scientific simulations, this architecture offers compelling advantages. Data can be passed between CPU-based analysis and GPU-accelerated solvers without copy overhead and the GPU can access the system’s entire memory pool rather than being constrained by discrete VRAM limits.
The MLX Framework
To harness the capabilities of Apple Silicon for numerical and machine learning workloads, Apple’s machine learning research team introduced MLX. MLX is an open-source array framework designed specifically for Apple Silicon and Metal, combining the ergonomics of standard Python scientific packages with hardware-level optimizations. Some key features of MLX include a familiar NumPy-like API, unified memory utilization, lazy evaluation, and the ability to compile computational graphs into optimized Metal GPU kernels via mx.compile.
Accelerating Tridiagonal Solvers with MLX
Algorithms like Cyclic Reduction are naturally suited to fine-grained parallelism. At each reduction step \(s\), odd-indexed equations are decoupled and updated independently from their neighbors with a stride of \(2^s\). The high memory bandwidth of Apple Silicon combined with the ability to compile multi-step updates into fused GPU kernels provides a promising avenue for performance gains over CPU implementations.
There is also a secondary benefit to implementing a native MLX tridiagonal solver. In hybrid physics-ML models, such as when neural networks learn unknown parameters or closure terms within an implicit PDE solver, the tridiagonal system must be solved repeatedly during training. Because MLX natively supports automatic differentiation, an MLX-based solver allows gradients to be backpropagated directly through the linear solve on the GPU, avoiding the latency of transferring data between traditional CPU numerical libraries (like SciPy) and machine learning frameworks.
Results
The implementations generated for this study are now available as a standalone python library called tridiag.
We implemented MLX versions of Cyclic Reduction and Parallel Cyclic Reduction and ran them against the same benchmarks we used for the CPU implementations. At first glance the results in Figure 1 are disappointing. For smaller \(N\) the sequential, CPU-based implementations are orders of magnitude faster than the GPU or parallel versions. At very large \(N\) (i.e. \(N > 10^6\)) the MLX implementations of CR finally start to beat the reference SciPy implementation. This is encouraging, but tempered by the fact that the CR implementation using compiled, cross-core CPU parallelization (CR Parallel Numba) remains the fastest tested at the largest \(N\). In these tests, the PCR implementations never beat the baseline and do not appear to have performance advantages.
Conclusions
While parallel algorithms like Cyclic Reduction are theoretically well suited for GPUs, our benchmarks demonstrate that solving a single tridiagonal system in isolation fails to exploit the strengths of Apple Silicon’s GPU. Several architectural factors explain this behavior:
- GPU Under-Utilization: At each step of CR, the active problem size is halved (\(N \to N/2 \to N/4 \dots \to 1\)). Even for large systems, the later reduction stages and early back-substitution stages only operate on dozens or hundreds of equations, leaving the vast majority of the GPU’s execution cores idle.
- Dispatch Latency and Synchronization: CR requires \(2 \log_2 N\) sequential stages separated. On the GPU, the cumulative overhead of dispatching kernels and synchronizing device memory across dozens of stages can add up. On the CPU, Numba executes without kernel launch overhead directly out of high-speed L1/L2 caches.
- Non-Coalesced Memory Access: Tridiagonal solvers have low arithmetic intensity. As the stride expands during reduction, threads access memory non-contiguously. This likely breaks memory coalescing, causing the GPU to fetch entire cache lines to retrieve only a fraction of the data. PCR attempts to maximize parallelism, but its \(O(N \log_2 N)\) memory traffic overwhelms the memory bus.
The Path Forward: Batched Solvers
These results highlight an important caveat in that our benchmark evaluated the latency of solving a single 1D system. In practical engineering simulations (e.g. 2D/3D Alternating Direction Implicit schemes) researchers rarely solve a single system in isolation. Instead, they solve thousands of independent tridiagonal systems simultaneously.
Rather than measuring the time to solve one system of size \(N\) sequentially, a more representative benchmark for GPU architectures is measuring batched throughput across \(B\) systems of size \(N\). Batching preserves constant, massive concurrency (\(B \times N\)), enables perfectly coalesced memory access across the batch dimension, and amortizes kernel launch latency over thousands of systems. We plan to explore batched MLX solvers in a follow-up study.
Finally, even for single systems, MLX solvers offer a distinct advantage for machine learning practitioners: end-to-end differentiability. In differentiable physics simulations and hybrid physics-ML models, where an implicit numerical solve lives directly inside a neural network’s forward pass, control loop, or parameter estimation loss, keeping the solve on the GPU enables backpropagation through the linear system without costly memory transfers.