Julia access to Apple GPU with MLX, and or Metal Performance Shaders (MPS)?

Is there any on-going effort to provide a Julia interface to the following?

  • Metal Performance Shaders (MPS) is a PyTorch framework backend for GPU training acceleration, providing scripts and capabilities to set up and run operations on Mac.

  • MLX is an array framework for machine learning research on Apple silicon, by Apple machine learning research. MLX provides FFTs, Linear Algebra operations Cholesky, inv, qr, svd, and includes a list of of examples:

Or calling them via their Python interfaces using PythonCall.jl is the only (good-enough?) way? MLX has a C API MLX-C

Apologies if this has been discussed already, but I could not locate any related information here.

3 Likes

MPS is not a PyTorch Framework. MPS are implementations of common operations as Metal (GPU) kernels - one sub set of which is implementing NN/CNN operations. In addition, PyTorch has a MPS back-end which uses the MPS kernels to accelerate PyTorch operations etc.

MPS NN/CNN-support via Metal.jl (and NNlib.jl, and in turn Flux.jl) is so far tracked in MPS - Support for Convolutional Neural Network kernels · Issue #210 · JuliaGPU/Metal.jl · GitHub.

If I am not mistaken, there has not yet been any public applications of MLX in Julia.

1 Like

That API seems the way to go, to wrap (with a JLL) if not already done, should be simple, you could also call without a JLL.

High-level API to that would be of help, e.g. to Wisper, not just for Apple hardware, i.e. use that code there, and other alternatives elsewhere.

It’s good enough for most people calling Python with PythonCall.jl for stuff like that, though might be better to call some such (cross-platform) Rust wrapper or other, if not done in Julia.

MLX has a Python API that closely follows NumPy. MLX also has fully featured C++, C, and Swift APIs, which closely mirror the Python API. MLX has higher-level packages like mlx.nn and mlx.optimizers with APIs that closely follow PyTorch to simplify building more complex models.

1 Like

Of course, you are right.

Took a stab at building MLX from source: WIP Added MLX by stemann · Pull Request #9761 · JuliaPackaging/Yggdrasil · GitHub

(but it may need to be a re-packaging of the official Python wheels instead - as building the metal backend likely requires the macOS/Windows-only metal compiler)

2 Likes

Additional WIP:

2 Likes