burn-mpsgraph

Apple Metal Performance Shaders Graph (MPSGraph) backend for the Burn deep learning framework.

Runs tensor computations on Apple GPUs (M1, M2, M3, M4 and later) by dispatching directly to MPSGraph — Apple's graph-based GPU compute engine that powers Core ML and Create ML.

Quick start

# Cargo.toml
[dependencies]
burn-mpsgraph = "0.0.1"
burn           = "0.21.0-pre.2"

use burn::tensor::{Distribution, Tensor};
use burn_mpsgraph::prelude::*;

type B = MpsGraph;

let device = MpsGraphDevice::default();

let a: Tensor<B, 2> = Tensor::random([128, 64], Distribution::Default, &device);
let b: Tensor<B, 2> = Tensor::random([64, 256], Distribution::Default, &device);
let c = a.matmul(b);          // runs on Apple GPU
let data = c.into_data();     // copy result back to CPU

Design

Burn Tensor API
      │
      ▼
FloatTensorOps / IntTensorOps / BoolTensorOps / ModuleOps
      │
      ▼
bridge.rs ── builds an MPSGraph per op, feeds MTLBuffer tensors,
             runs synchronously inside an ObjC autorelease pool
      │
      ▼
ffi.rs ── raw objc_msgSend calls to Metal, Foundation, MPS, MPSGraph
      │
      ▼
Apple GPU (Metal)

GPU-resident tensors

Each MpsGraphTensor wraps a retained MTLBuffer pointer.
On Apple Silicon the buffer uses shared memory — CPU and GPU share the same physical pages with no copy overhead.
Data is only serialised to Vec<u8> when you call into_data().

No objc2 dependency

All Apple framework calls go through direct objc_msgSend FFI in ffi.rs. This removes the entire objc2 / objc2-foundation / objc2-metal family from the dependency tree.

Supported operations

Supported dtypes

Benchmarks

Benchmarks run on Apple M3 Pro against:

• burn-wgpu with the metal feature (wgpu → MSL → Metal)
• burn-ndarray (CPU reference)

Matmul (f32, square matrices)

At 4096×4096 both GPU backends are ~14× faster than CPU.

Element-wise add

Note: For small ops the per-graph-execution overhead (~2 ms) dominates. This is the primary area for future improvement — see Roadmap.

Examples

cargo run --example basic      # arithmetic, reductions, math, shapes
cargo run --example inference  # 2-layer MLP with softmax
cargo run --example conv       # conv2d + maxpool2d feature extractor

Seeded RNG

MpsGraph::seed(&device, 42);
let t: Tensor<B, 2> = Tensor::random([4, 4], Distribution::Default, &device);

The seed is stored globally per process. Two calls with the same seed produce identical tensors.

Platform requirements

The build.rs links these frameworks automatically: Foundation, Metal, MetalPerformanceShaders, MetalPerformanceShadersGraph.

Roadmap

• Lazy / deferred execution — accumulate ops into one graph before running. This would eliminate the ~2 ms per-op overhead and make element-wise chains as fast as native MPSGraph programs.
• MTLBuffer pool — reuse freed buffers instead of allocating new ones.
• Async execution — use encodeToCommandBuffer to pipeline work.
• BF16 arithmetic — currently storage-only.
• F64 via emulation — two F32 ops for reduce_sum / mean.
• Multi-device — map MpsGraphDevice { index } to discrete GPUs on Mac Pro.

License

Licensed under either of

• MIT License
• Apache License, Version 2.0

at your option.