lambdaclass / metal_playground

A repo for learning how to parallelize computations in the GPU using Apple's Metal, in Rust.

Metal playground in rust

Made for learning how to parallelize computations in the GPU using Apple's Metal, in Rust, via the metal crate.

Overview

The source code will contain documented examples of use, growing in complexity, aimed at a final objective of parallelize a Fast Fourier Transform algorithm.

Contents

1. dotprod: learn the basics of Metal and metal-rs, implement a simple dot product between two vectors represented as uint arrays.
2. matrixprod: a more complex example to learn about grid size and thread groups, implement a product between square matrices.
3. memory: example to show how to shared memory between CPU and GPU. This example creates a vector to be modified from the GPU.

Pre-requisites

• XCode
• Rust

Running the examples

To run the examples, use the following command:

make example EXAMPLE={example}

To re-build all the necessary .metallib files, you can use

make compile_metal EXAMPLE={example}

where {example} is the name of the example to run in both commands.

References

• Apple's Metal documentation: we recommend to start with "Performing Calculations on a GPU". Note that these docs are in Swift/Obj-C.
• miniSTARK: A minimal STARK library built in Rust and gpu-accelerated with Metal.
• metal-rs examples

metal_playground's Issues

Make FP class generic enough to accept any `P` value

This issue can be divided into two tasks:

1. Remove M and R2 constants and update operations so that results are still valid. This will be less efficent than the original implementation.
2. Convert FP class to a template that takes P as an argument: template <uint32_t P> class Fp. If you want examples of a template class, see: https://github.com/andrewmilson/ministark/blob/main/gpu-poly/src/metal/felt_u128.h.metal#L50-L54