Mixed Precision Physics-Based Computations

[sck-at-ucy]

I am interested in exploring using MLX for mixed-precision computations for physics/egineering problems. An example usage would be discretizing the Navier-Stokes so that the convection and diffusion terms are done in float64 on the CPU and the Poisson pressure solver is done via MLX on the GPU.

What would be the best way to approach this project. Best in my mind is defined based on the following needs:

1. Best computational speed.
2. Flexibility to add features and portability moving forward.
3. Ability to easily extend for distributed computing.

I guess MLX is not likely to support directly float64 any time soon (even on the CPU streams) so I had stared exploring some simple examples using Swift and PythonKit to connect to MLX.

But with the latest release now MLX supports exporting/importing MLX functions from one API to another. Very cool and interesting and I am trying to understand is this is a better way to approach this issue.

I do not see clearly the full picture of the available options and the pros and cons and longer-term implications, thus I would like the advice of the MLX team as to which combination of tools would the best (based on the criteria above).

Some examples I can think of:

1. Swift for float64 discretization on CPU calling MLX via PythonKit
2. Swift for float64 discretization on CPU importing MLX functions via the MLX native import/export
3. Same as 1 and 2 but doing the float64 stuff directly in C++.
4. Any reason to throw Metal in the mix?

If @awni and @angeloskath think this is a workable idea, I would like to invest the time to do a basic physics case with my team and include it in MLX Projects. A few other engineering teams have asked me about this possibility, so I think if this can be done efficiently it would be quite useful for others as well.

[awni]

Another option you didn't add is doing the CPU float64 stuff in Python with NumPy. Depending on what the computational graph looks like this could be reasonably fast or not such a good idea. I'm not familiar enough to say.

and the Poisson pressure solver is done via MLX on the GPU.

If that model is pretty regular without variable shapes and control flow then using MLX import/export is a very viable option. You could do this with C++ or Swift and there I would just pick whichever you are more comfortable with.

I guess MLX is not likely to support directly float64 any time soon (even on the CPU streams)

I also don't think it's a crazy idea to add CPU-only fp64 support in MLX. It's mostly boiler plate changes that we'd need since all the CPU back-end is relatively well templatized and the accelerate API can handle double.

[sck-at-ucy]

I also don't think it's a crazy idea to add CPU-only fp64 support in MLX. It's mostly boiler plate changes that we'd need since all the CPU back-end is relatively well templatized and the accelerate API can handle double.

That would be awesome!! (❤️) It would make for a much cleaner source code and avoid conversions. I love where possible to do pure MLX. For example, if I could specify the device stream as being CPU for a given function and then cast variables received from a GPU stream as float64 and then the opposite on for functions that use GPU stream and receive variables from functions using the CPU.

and the Poisson pressure solver is done via MLX on the GPU.

I tested the idea with simple code solving the Navier-Stokes for a 2D lid-driven cavity flow. The convective-diffusive terms are one on the CPU with Swift and the Poisson on the GPU with MLX (at this point using PythonKit). It seems to work and the rough relative speeds are:

• Swift+MLX-GPU is about 7x faster than Swift on CPU-only (even though GPU load is only a few percent during run)
• MLX-GPU-only is something like 60x faster than Swift+MLX-GPU (but with float32 throughout).