Mojo First Steps

TL;DR

• Mojo is a new, Python-like JIT-complied programming language from Chris Lattner of LLVM (and Swift) fame.
• Read this Chris’s excellent series of blog posts for some background.
• The basic idea is to have the convenience as ease of use of Triton, but supporting non-CUDA backends such as AMD’s ROCm and Apple’s Metal. SIMD support is also excellent.
• Once Python interop becomes more mature, Mojo should be a good choice for accelearating performance-critical Python code.
• Have a look at Get started with Mojo as well as Get started with GPU programming.

State Of Metal Support

• Metal (MPS) support is fairly new to Mojo. It was announced in this post from Brad Larson back in September.
• Under the covers Mojo makes use of metal-cpp, the C++ drop-in alternative to Metal’s Objective C headers.
• Run xcrun -find metallib. It should return something like /var/run/com.apple.security.cryptexd/mnt/com.apple.MobileAsset.MetalToolchain-v17.2.54.0.Nf6rON/Metal.xctoolchain/usr/bin/metallib.
• If you get an error, launch XCode, navigate to XCode->Settings->Components and click the Get button next to Metal Toolchain under Other Components. After the Metal Toolchain finishes downloading, run xcrun -find metallib again to confirm everything works.

First Steps

• Install Pixi: brew install pixi or curl -fsSL https://pixi.sh/install.sh | sh. Pixi is like uv, but it lets you use Conda packages.
• Clone the Mojo repo.
• cd examples/mojo/gpu_functions (there is a pixi.toml there that specifies the necessary deps).
• Use your favorite editor to create a file in that directory called accelerator.mojo, with the following contents:

from sys import has_accelerator
from gpu.host import DeviceContext

def main():
    @parameter
    if not has_accelerator():
        print("No compatible GPU found")
    else:
        # Get the context for the attached GPU
        ctx = DeviceContext()
        print("Found GPU:", ctx.name())

• Run pixi run mojo accelerator.mojo. This will JIT-compile and run accelerator.mojo. On a Mac, the output should be something like “Found GPU: Apple M3 Max”.
• You can also use mojo to compile standalone executables, like so: pixi run mojo build mandelbrot.mojo. They can then be run directly: ./mandelbrot.
• The rest of the examples in gpu-functions should also work, with the exception of reduction.mojo. That one is too NVIDIA-specifc.
• For a more involved example demonstrating Mojo language constructs such as Lists, structs and modules, as well as calling Python from Mojo, cd examples/mojo/life and pixi run mojo run lifev3.mojo (this is just the example from Get started with Mojo).
• For more information about Mojo language basics, see this page.