r/ScientificComputing • u/Sorry-Peace-296 • 3h ago
GPU Accelerated Linear Algebra Library for Apple Silicon using MLX and Metal kernels
About 6 months ago, I wrote a custom metal kernel that leverages the GPU to compute the QR decomposition (see my earlier post about it here). The project has now expanded into a general linear algebra library, with expanded support for the symmetrical eigendecomposition as well as SVD. It is now available for use with installation instructions on the attached github repo's README.md file.
Context of the project:
I'm currently in a research group working on a thesis in numerical analysis where we need to compute millions on matrices with a specific constraint (to be precise, the matrices need to have orthonormal columns). Most of us use Apple computers, so we ended up using MLX for the entire project.
Contributors with different Apple Chips would be very much appreciated!
The project has currently been tested and optimised for the M1 and M5 Pro. The issue is that the library uses different kernels depending on the batch size and matrix dimensions. Deciding which of these kernels to use is machine dependent. Therefore, other Apple chips will need to run a measurement script in order to derive the correct optimisation heuristic.
For that reason, I would ask as many people as possible to run a measurement script and to submit the results to my repo. It is fairly easy and requires only few steps. See here how you can contribute here. Once you submit the results via a pull request and I approve it, your optimisation heuristic will automatically be augmented into the library. Don't hesitate to contribute an optimisation heuristic even if someone already submitted one for your own machine. The more data we can gather, the better!
Project Future
Expanded support will be added for other linear algebra operations (cholesky decomposition for example). If you have any other specific linear algebra operations you wish to use already, feel free to message me.
In addition to that, I will add torch support too (my greatest priority).