Merge pull request #6 from asjad2401/feat/backend-dispatch

feat: GPU inference via wgpu (Metal on macOS, Vulkan elsewhere)

Author: asjad2401

.github/workflows/ci.yml

@@ -125,3 +125,38 @@ jobs:
 
       - name: Run ResNet-18 tests
         run: .venv/bin/python tests/test_resnet.py --bench
+
+  wgpu-tests:
+    name: wgpu tests (software Vulkan)
+    runs-on: ubuntu-latest
+    needs: build
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Install Rust
+        uses: dtolnay/rust-toolchain@stable
+
+      - name: Cache cargo
+        uses: Swatinem/rust-cache@v2
+
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: "3.12"
+
+      # no GPU on hosted runners: lavapipe is Mesa's CPU implementation of Vulkan,
+      # enough to check the GPU code path end to end (not its speed)
+      - name: Install lavapipe
+        run: sudo apt-get update && sudo apt-get install -y mesa-vulkan-drivers
+
+      - name: Create virtualenv
+        run: python -m venv .venv
+
+      - name: Install Python deps
+        run: .venv/bin/pip install maturin numpy onnx onnxruntime
+
+      - name: Build extension
+        run: source .venv/bin/activate && maturin develop --release
+
+      - name: Run op edge-case tests on wgpu
+        run: BURN_BACKEND=wgpu .venv/bin/python tests/test_ops.py

CHANGELOG.md

@@ -7,13 +7,34 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
-### Performance
-- Static weights that onnx-ir passes as anonymous (unnamed) constants — including all
-  conv weights — are now pre-loaded once instead of converted on every forward pass.
-- BatchNormalization directly after a Conv2d is folded into the conv's weight and bias
-  at load time and removed from the graph. ResNet-18 (Model Zoo v1) at batch 8: 279 → 103 ms.
+### Added
+- GPU inference: `load_onnx(path, backend="wgpu")`, compiled to Metal on macOS and
+  Vulkan on Linux / Windows. ResNet-18 on an Apple M-series GPU: 6.1 ms at batch 1,
+  35.6 ms at batch 8 (~3x faster than ONNX Runtime on CPU). `available_backends()`
+  lists the backends in the build.
+- CI job running the op tests on wgpu through lavapipe (software Vulkan).
+- ReduceMean, ReduceSum, ReduceMax, ReduceMin (needed by the torch dynamo exporter's
+  global average pooling).
+- `tests/test_resnet.py`: ResNet-18 end to end vs ONNX Runtime on Model Zoo and
+  torchvision (legacy + dynamo) exports, with optional `--bench`; run in CI.
+- `tests/test_ops.py`: per-op edge-case tests against ONNX Runtime, run in CI.
+- ONNX interpreter support for Conv2d, fused BatchNormalization, MaxPool2d,
+  AveragePool2d, and GlobalAveragePool.
+- GitHub Actions CI: `cargo build`, `cargo fmt` / `cargo clippy` (non-blocking),
+  and Python tests via `maturin develop`.
+- Branch protection on `main`: pull requests required, `cargo build` and
+  `python tests (maturin)` checks required to merge.
+- ONNX interpreter with Gemm, Linear, Relu, Sigmoid, Tanh, Gelu, Softmax,
+  LogSoftmax, Add, Sub, Mul, Div, Reshape, Flatten, and Transpose ops.
+- numpy ↔ Burn tensor bridge (`roundtrip`), one copy in each direction.
+- Initial maturin/pyo3 project scaffold, `burn-flex` + `onnx-ir` dependencies.
 
 ### Changed
+- All ops now run on burn-dispatch's `Dispatch` backend; the device is chosen per model.
+  The unused `burn` umbrella dependency was replaced by `burn-dispatch`.
+- Backend failures that happen on a worker thread (e.g. no GPU adapter) now report the
+  worker's panic message instead of a downstream `RecvError`, and model loading is
+  guarded too, so they surface as `RuntimeError`.
 - Errors now surface as normal Python exceptions instead of `PanicException`:
   - `TypeError` for non-array / non-float32 inputs (with a `.astype(np.float32)` hint).
   - `ValueError` for wrong input count, rank, or a static dim that doesn't match the model.
@@ -22,7 +43,17 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Rust panic output on stderr is suppressed (the message is in the exception);
   set `RUST_BACKTRACE=1` to get it back.
 
+### Performance
+- Static weights that onnx-ir passes as anonymous (unnamed) constants — including all
+  conv weights — are now pre-loaded once instead of converted on every forward pass.
+- BatchNormalization directly after a Conv2d is folded into the conv's weight and bias
+  at load time and removed from the graph. ResNet-18 (Model Zoo v1) at batch 8: 279 → 103 ms.
+
 ### Fixed
+- Broadcasting between operands of different rank (`[N, C] + [C]`, Gemm's rank-1 `C`,
+  `[N, C, H, W] * [C, 1, 1]`) is now done numpy-style. It only worked by accident on
+  Flex; GPU backends rejected it.
+- Gemm `alpha`/`beta` and BatchNorm `epsilon` use scalar ops instead of 1-element tensors.
 - BatchNormalization on non-4D inputs (e.g. `[N, C]` after a Linear) no longer
   broadcasts to a 4D output; the pre-fused fast path is only used for NCHW.
 - Fortran-ordered inputs (e.g. `x.T`) were read in memory order, silently scrambling data.
@@ -36,21 +67,4 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Unsupported ops now fail at `load_onnx` with a list of the missing ops, instead of
   being skipped with a warning and failing (or silently misbehaving) at inference time.
 
-### Added
-- ReduceMean, ReduceSum, ReduceMax, ReduceMin (needed by the torch dynamo exporter's
-  global average pooling).
-- `tests/test_resnet.py`: ResNet-18 end to end vs ONNX Runtime on Model Zoo and
-  torchvision (legacy + dynamo) exports, with optional `--bench`; run in CI.
-- `tests/test_ops.py`: per-op edge-case tests against ONNX Runtime, run in CI.
-- ONNX interpreter support for Conv2d, fused BatchNormalization, MaxPool2d,
-  AveragePool2d, and GlobalAveragePool.
-- GitHub Actions CI: `cargo build`, `cargo fmt` / `cargo clippy` (non-blocking),
-  and Python tests via `maturin develop`.
-- Branch protection on `main`: pull requests required, `cargo build` and
-  `python tests (maturin)` checks required to merge.
-- ONNX interpreter with Gemm, Linear, Relu, Sigmoid, Tanh, Gelu, Softmax,
-  LogSoftmax, Add, Sub, Mul, Div, Reshape, Flatten, and Transpose ops.
-- Zero-copy numpy ↔ Burn tensor bridge (`roundtrip`).
-- Initial maturin/pyo3 project scaffold, `burn-flex` + `onnx-ir` dependencies.
-
 [Unreleased]: https://github.com/asjad2401/burn-python/commits/main

Cargo.toml

@@ -10,14 +10,26 @@ crate-type = ["cdylib"]
 [dependencies]
 pyo3   = { version = "0.28", features = ["extension-module"] }
 numpy  = "0.28"
-burn   = { version = "0.21", features = ["flex"] }
 burn-flex    = "0.21"
+# listed directly so its default features (autotune, fusion) are on
+burn-wgpu    = "0.21"
 burn-backend = "0.21"
 onnx-ir      = "0.21"
 cubecl-common = "0.10"
 bytes     = "1"
 bytemuck  = { version = "1", features = ["extern_crate_alloc"] }
 
+# All ops go through burn-dispatch: Flex on CPU, wgpu on GPU. The GPU path is compiled
+# to each platform's native shading language — Metal (MSL) on macOS, Vulkan (SPIR-V)
+# elsewhere — which is several times faster than the portable WGSL path (no access to
+# matrix hardware). burn-dispatch allows only one wgpu flavor per build, hence per target.
+# wgpu loads the GPU driver at runtime, so CPU-only machines still work with "flex".
+[target.'cfg(target_os = "macos")'.dependencies]
+burn-dispatch = { version = "0.21", default-features = false, features = ["std", "flex", "metal"] }
+
+[target.'cfg(not(target_os = "macos"))'.dependencies]
+burn-dispatch = { version = "0.21", default-features = false, features = ["std", "flex", "vulkan"] }
+
 [profile.release]
 lto = true
 codegen-units = 1

README.md

@@ -15,6 +15,21 @@ x = np.random.randn(1, 3, 224, 224).astype(np.float32)
 output = model([x])[0]
 ```
 
+## Backends
+
+```python
+burn.available_backends()                          # ['flex', 'wgpu']
+model = burn.load_onnx("model.onnx", backend="wgpu")  # run on the GPU
+```
+
+| backend | runs on | notes |
+|---------|---------|-------|
+| `"flex"` (default) | CPU | pure-Rust [burn-flex](https://github.com/tracel-ai/burn) backend |
+| `"wgpu"` | GPU | Metal on macOS, Vulkan on Linux / Windows; first call per input shape compiles and autotunes kernels |
+
+Both use Burn's dispatch backend, so the same interpreter runs on either. If no usable
+GPU is found, `load_onnx(..., backend="wgpu")` raises a `RuntimeError` naming the problem.
+
 ## Status
 
 Early development. The numpy ↔ Burn tensor bridge is done, and the ONNX interpreter
@@ -33,15 +48,16 @@ ResNet-18 runs end to end and matches ONNX Runtime (max logit diff < 1e-5, same
 for the ONNX Model Zoo exports (`resnet18-v1-7`, `resnet18-v2-7`) and torchvision's
 `resnet18` exported with both `torch.onnx` exporters (legacy and dynamo).
 
-CPU latency vs ONNX Runtime (torchvision export, burn `flex` backend, median of 10 runs).
-Performance varies a lot by platform:
+Latency vs ONNX Runtime on CPU (torchvision export, median of 15 runs after warm-up).
+Numbers vary a lot by platform:
 
-| platform | batch | burn-python | ONNX Runtime | burn / ORT |
-|----------|------:|------------:|-------------:|-----------:|
-| Apple M-series (10 cores) | 1 | 22.6 ms  | 14.7 ms  | 1.5× slower |
-| Apple M-series (10 cores) | 8 | 105.0 ms | 118.5 ms | 0.9× (faster) |
-| Linux x86 (GitHub Actions runner) | 1 | 58.0 ms  | 18.7 ms  | 3.1× slower |
-| Linux x86 (GitHub Actions runner) | 8 | 335.7 ms | 145.9 ms | 2.3× slower |
+| platform | backend | batch 1 | batch 8 |
+|----------|---------|--------:|--------:|
+| Apple M-series (10 cores) | burn-python `wgpu` (Metal GPU) | **6.1 ms** | **35.6 ms** |
+| Apple M-series (10 cores) | burn-python `flex` (CPU) | 22.4 ms | 135.7 ms |
+| Apple M-series (10 cores) | ONNX Runtime (CPU) | 14.9 ms | 123.0 ms |
+| Linux x86 (GitHub Actions runner) | burn-python `flex` (CPU) | 58.0 ms | 335.7 ms |
+| Linux x86 (GitHub Actions runner) | ONNX Runtime (CPU) | 18.7 ms | 145.9 ms |
 
 Run `python tests/test_resnet.py --bench` to measure on your machine.
 
@@ -62,6 +78,8 @@ python tests/test_bridge.py       # numpy <-> Burn tensor bridge
 python tests/compare_ort.py       # correctness + perf vs ONNX Runtime
 python tests/test_ops.py          # op edge cases (padding, pooling, reshape) vs ORT
 python tests/test_resnet.py       # ResNet-18 end to end vs ORT (add --bench for timings)
+
+BURN_BACKEND=wgpu python tests/test_ops.py     # same tests on the GPU backend
 ```
 
 ## License

burn_python/__init__.py

@@ -1,5 +1,5 @@
 import numpy as np
-from ._burn_python import __version__, OnnxModel
+from ._burn_python import __version__, OnnxModel, available_backends
 from . import _burn_python as _ext
 
 def roundtrip(arr: np.ndarray) -> np.ndarray:
@@ -9,9 +9,11 @@ def roundtrip(arr: np.ndarray) -> np.ndarray:
     return _ext.roundtrip(arr)
 
 def load_onnx(path: str, backend: str = "flex") -> OnnxModel:
-    """Load an ONNX model for inference."""
-    if backend != "flex":
-        raise ValueError(f"unsupported backend '{backend}' — only 'flex' is available right now")
-    return _ext.load_onnx(path)
+    """Load an ONNX model for inference.
 
-__all__ = ["__version__", "load_onnx", "roundtrip", "OnnxModel"]
+    backend: "flex" (CPU, default) or "wgpu" (GPU via Metal / Vulkan / DX12).
+    See available_backends() for what this build supports.
+    """
+    return _ext.load_onnx(str(path), backend)
+
+__all__ = ["__version__", "load_onnx", "available_backends", "roundtrip", "OnnxModel"]

src/interpreter/context.rs

@@ -6,7 +6,7 @@ use std::collections::HashMap;
 use burn_backend::backend::ops::FloatTensorOps;
 use onnx_ir::ir::{Argument, ValueSource};
 
-use crate::tensor::{B, FloatPrim, default_device};
+use crate::tensor::{B, Device, FloatPrim};
 
 /// Key a static/constant argument is pre-loaded under in the weights map.
 /// onnx-ir often hands initializers over as anonymous statics (name ""), so those
@@ -25,13 +25,15 @@ pub type OpResult = Result<(), String>;
 pub struct ExecutionContext<'w> {
     tensors: HashMap<String, FloatPrim>,
     weights: &'w HashMap<String, FloatPrim>,
+    device: &'w Device,
 }
 
 impl<'w> ExecutionContext<'w> {
-    pub fn new(weights: &'w HashMap<String, FloatPrim>) -> Self {
+    pub fn new(weights: &'w HashMap<String, FloatPrim>, device: &'w Device) -> Self {
         Self {
             tensors: HashMap::new(),
             weights,
+            device,
         }
     }
 
@@ -56,10 +58,7 @@ impl<'w> ExecutionContext<'w> {
             ValueSource::Static(_) => weight_key(arg)
                 .and_then(|key| self.weights.get(&key).cloned())
                 // not pre-loaded (e.g. non-float data) — convert on the fly
-                .or_else(|| {
-                    arg.value()
-                        .map(|d| B::float_from_data(d, &default_device()))
-                }),
+                .or_else(|| arg.value().map(|d| B::float_from_data(d, self.device))),
             ValueSource::Optional => None,
         }
     }