Building Mojo CPU and AMD GPU Packages#

In this tutorial, we will build CPU and AMD GPU variants of a Mojo package with Pixi. For background on building packages, see the Getting Started guide and the Mojo backend documentation.

Why is This Useful?#

Build variants let us distribute CPU and GPU implementations of the same package without changing its name or version. Pixi can then select the AMD GPU build on compatible hardware and use the CPU build otherwise. We will publish both builds to a local channel and install them in a separate workspace, without a source dependency or build backend in the consumer.

Let's Get Started#

We will create a mojo-matmul package that multiplies two matrices. This example needs Linux x86-64, glibc 2.34 or newer, and an x86-64-v3-capable CPU for both builds. The GPU build targets gfx1103, such as the Radeon 780M. See rich platforms for hardware declarations, including other AMD targets.

Create these files in a new directory:

.
├── package/
│   ├── main.mojo
│   └── pixi.toml
└── pixi.toml

First copy the matrix multiplication program.

from std.collections import List
from std.math import ceildiv
from std.memory import Pointer
from std.sys import get_defined_string

from max.gpu import global_idx
from max.gpu.host import DeviceContext

comptime ACCELERATOR = get_defined_string["mojo_accelerator"]()
comptime SIDE = 1024
comptime BLOCK_SIZE = 256


def dot[lhs_origin: MutOrigin, rhs_origin: MutOrigin](
    lhs: Pointer[Float32, lhs_origin],
    rhs: Pointer[Float32, rhs_origin],
    row: Int,
    col: Int,
) -> Float32:
    var value = Float32(0)
    for k in range(SIDE):
        value += lhs[unsafe_offset=row * SIDE + k] * rhs[unsafe_offset=k * SIDE + col]
    return value


def multiply(
    lhs: Pointer[Float32, MutAnyOrigin],
    rhs: Pointer[Float32, MutAnyOrigin],
    output: Pointer[Float32, MutAnyOrigin],
):
    var i = global_idx.x
    if i < SIDE * SIDE:
        output[unsafe_offset=i] = dot(lhs, rhs, i // SIDE, i % SIDE)


def main() raises:
    var lhs = List[Float32](length=SIDE * SIDE, fill=1)
    var rhs = List[Float32](length=SIDE * SIDE, fill=2)
    var output = List[Float32](unsafe_uninit_length=SIDE * SIDE)
    var result = output.unsafe_ptr()
    for i in range(SIDE):
        lhs[i * SIDE + i] += 1
        rhs[i * SIDE + i] += 1

    comptime if ACCELERATOR == "cpu":
        for row in range(SIDE):
            for col in range(SIDE):
                result[unsafe_offset=row * SIDE + col] = dot(
                    lhs.unsafe_ptr(), rhs.unsafe_ptr(), row, col
                )
    elif ACCELERATOR == "gfx1103":
        var ctx = DeviceContext()
        var device_lhs = ctx.enqueue_create_buffer[DType.float32](SIDE * SIDE)
        var device_rhs = ctx.enqueue_create_buffer[DType.float32](SIDE * SIDE)
        var device_output = ctx.enqueue_create_buffer[DType.float32](SIDE * SIDE)
        ctx.enqueue_copy(device_lhs, lhs.unsafe_ptr())
        ctx.enqueue_copy(device_rhs, rhs.unsafe_ptr())
        ctx.enqueue_function[multiply](
            device_lhs, device_rhs, device_output,
            grid_dim=ceildiv(SIDE * SIDE, BLOCK_SIZE), block_dim=BLOCK_SIZE,
        )
        ctx.enqueue_copy(result, device_output)
        ctx.synchronize()

    for i in range(SIDE * SIDE):
        var expected = Float32(2 * SIDE + 3 + Int(i // SIDE == i % SIDE))
        if output[i] != expected:
            raise Error("Matrix multiplication mismatch at index " + String(i))
    print("path=", ACCELERATOR, " side=", SIDE, " validation=ok", sep="")

There is a lot of math involved, but you don't have to understand all of it to profit from this tutorial. The important bit is that different code is executed if ACCELERATOR == "cpu" than if ACCELERATOR == "gfx1103".

The CPU branch uses a plain, single-threaded loop with normal compiler optimizations. The GPU branch copies the inputs to device buffers, launches one thread per output cell, copies the result back, and synchronizes before reading it. Both branches multiply the same matrices, so we expect the same result but different performance.

Next, define the producer workspace and package:

[workspace]
channels = ["https://conda.modular.com/max", "https://prefix.dev/conda-forge"]
platforms = [{ platform = "linux-64", glibc = "2.34", archspec = "x86_64_v3" }]
preview = ["pixi-build"]

[workspace.build-variants]
mojo_accelerator = ["cpu", "gfx1103"]

[package]
name = "mojo-matmul"
version = "0.1.0"

[package.build.backend]
name = "pixi-build-mojo"
version = "*"

[[package.build.config.bins]]
name = "matmul"

[package.build-dependencies]
mojo-compiler = "==1.2.0"

[package.host-dependencies]
max-core = "==26.7.0"

The mojo_accelerator build variant produces two packages with the same name and version but different build hashes.

You can create a local conda channel like this:

pixi publish --path ./package/pixi.toml --target-channel channel

--target-channel creates an indexed local channel containing both builds. The producer has no dependency on that channel, so publishing works before it exists.

Install and run#

Create the consumer manifest beside package/:

[workspace]
channels = ["./channel", "https://conda.modular.com/max", "https://prefix.dev/conda-forge"]
platforms = [
  { name = "amd", platform = "linux-64", glibc = "2.34", archspec = "x86_64_v3", amdgpu = "gfx1103" },
  { name = "cpu", platform = "linux-64", glibc = "2.34", archspec = "x86_64_v3" },
]

[dependencies]
mojo-matmul = "*"

[tasks]
matmul = "matmul"

The two rich platforms describe separate solves for the same linux-64 subdir. amd comes first, so Pixi selects it automatically on compatible hardware, otherwise it selects cpu. Both require archspec = "x86_64_v3" because the prebuilt Modular runtime needs that CPU baseline. amdgpu = "gfx1103" requires exactly that ISA; even newer ISAs do not match. If the GPU build is suitable, typically the CPU is suitable as well. To prioritize the GPU build in that case, pixi-build-mojo makes sure to down prioritize the CPU build.

Run from the consumer directory:

pixi run matmul

On a compatible Radeon 780M, this prints:

path=gfx1103 side=1024 validation=ok

To run the CPU build explicitly:

pixi run --platform cpu matmul

path=cpu side=1024 validation=ok

Optional: measure whole-program time#

On a compatible AMD machine, install Hyperfine and measure each implementation:

pixi global install hyperfine
pixi run --platform cpu hyperfine --warmup 1 --runs 5 matmul
pixi run --platform amd hyperfine --warmup 1 --runs 5 matmul

On the machine the tutorial author used, the GPU version was 4.8 times faster than the CPU version.