Run the same model on CPU or GPU

Write a model's maths once, and choose where it runs by choosing a handler. The model asks for operations such as "multiply these matrices" through an effect called Tensor. One handler answers with ordinary lists on the CPU. Another answers with GPU buffers on Metal or OpenCL. The model code is the same in both cases.

Tensor is a static effect: the handler is fixed when the program is compiled, so the compiler builds a separate version of the model for each handler and nothing is looked up at run time. The worked example is GPT-2 (124M), with generation and a training step.

Try it

# osprey: flavor=python

def read_product(result):
    match result:
        case Error(message=message):
            return f"product failed: {message}"
        case Success(value=buffer):
            match gpuGet(buffer, 0):
                case Error(message=message):
                    return f"read failed: {message}"
                case Success(value=value):
                    return f"{gpuDevice()}: first element = {value}"


def main():
    values = toGpu([1.0, 2.0, 3.0, 4.0])
    print(read_product(gpuMatMul(values, values, 2, 2, 2)))

OSPREY_GPU=host osprey product.py --run
# macOS with Metal:
OSPREY_GPU=metal osprey product.py --run

Run all of GPT-2 on either backend from a built checkout. The first output line names the backend:

make -C src/toolchain/targets gpt2 PROMPT="The capital of France is" TOKENS=20 BACKEND=gpu
make -C src/toolchain/targets gpt2 PROMPT="The capital of France is" TOKENS=20 BACKEND=lists

How it works

The effect lists the operations. B is the buffer type, and each handler picks its own: List<float> on the CPU, GpuBuffer<float32> on the GPU. Operations that can fail, such as a matrix product or loading weights, return Result:

export static effect Tensor<B> {
    build: fn(int, Kernel) -> Result<B, string>
    get: fn(B, int) -> float32
    length: fn(B) -> int
    zipWith: fn(B, B, Combine) -> Result<B, string>
    reduce: fn(B, float32, Combine) -> Result<float32, string>
    indexOf: fn(B, float32) -> Result<int, string>
    matMul: fn(B, B, int, int, int) -> Result<B, string>
    fromList: fn(List<float32>) -> Result<B, string>
    load: fn(string, string) -> Result<B, string>
    device: fn() -> string
}

The same declaration in Python flavor (selected operations):

# osprey: flavor=python
@static_effect
class Tensor[B]:
    def build(size: int, kernel: Kernel) -> Result[B, str]: ...
    def get(buffer: B, index: int) -> float32: ...
    def mat_mul(a: B, b: B, rows: int, inner: int, cols: int) -> Result[B, str]: ...
    def load(path: str, name: str) -> Result[B, str]: ...
    def device() -> str: ...

And the failing operations in ML flavor:

matMul : (B, B, int, int, int) => Result<B, string>
load : (string, string) => Result<B, string>

The CPU handler keeps tensors in lists. The GPU handler runs GPU kernels; its product passes the GPU's Result straight back:

matMul a b rows inner cols => answered(gpuMatMul(a, b, rows, inner, cols))
load path name => answered(gpuLoadTensor(path, name))

If a product or load fails, the model stops and returns the Error; it never carries on with a made-up tensor. One exception: device-kernel faults currently abort instead of reaching an arithmetic handler (see the GPU contract).

What the compiler checks

• Both handlers answer every operation with matching types.
• GPU kernel code is pure: a kernel that tries to log, for example, is rejected by name.
• A regression test runs a small transformer under both handlers and compares them with independent expected values.

Speed and limits

These timings ran both handlers on the CPU (OSPREY_GPU=host), not on a GPU; GPT-2's README has GPU timings.

Data stays on the GPU until it is read back. Kernels are not fused, so each one is a separate dispatch. GPU maths functions can differ from the CPU's in the last digit, so check your results against the documented tolerances.

Effects specification · GPU specification