12 releases (7 breaking)

0.15.0 Oct 28, 2024
0.13.2 May 3, 2024
0.12.1 Feb 1, 2024
0.11.1 Dec 4, 2023
0.8.0 Jul 25, 2023

#710 in Machine learning

Download history 1470/week @ 2024-07-24 2985/week @ 2024-07-31 1884/week @ 2024-08-07 1693/week @ 2024-08-14 1686/week @ 2024-08-21 1636/week @ 2024-08-28 1475/week @ 2024-09-04 1538/week @ 2024-09-11 1359/week @ 2024-09-18 1633/week @ 2024-09-25 1295/week @ 2024-10-02 1563/week @ 2024-10-09 1448/week @ 2024-10-16 1814/week @ 2024-10-23 1683/week @ 2024-10-30 1320/week @ 2024-11-06

6,576 downloads per month
Used in 23 crates (3 directly)

MIT/Apache

2MB
45K SLoC

Burn WGPU Backend

Burn WGPU backend

Current Crates.io Version license

This crate provides a WGPU backend for Burn using the wgpu.

The backend supports Vulkan, Metal, DirectX11/12, OpenGL, WebGPU.

Usage Example

#[cfg(feature = "wgpu")]
mod wgpu {
    use burn_autodiff::Autodiff;
    use burn_wgpu::{Wgpu, WgpuDevice};
    use mnist::training;

    pub fn run() {
        let device = WgpuDevice::default();
        training::run::<Autodiff<Wgpu<f32, i32>>>(device);
    }
}

Configuration

You can set BURN_WGPU_MAX_TASKS to a positive integer that determines how many computing tasks are submitted in batches to the graphics API.

Alternative SPIR-V backend

When targeting Vulkan, the spirv feature flag can be enabled to enable the SPIR-V compiler backend, which performs significantly better than WGSL. This is especially true for matrix multiplication, where SPIR-V can make use of TensorCores and run at f16 precision. This isn't currently supported by WGSL. The compiler can also be selected at runtime by setting the corresponding generic parameter to either SpirV or Wgsl.

Platform Support

Option CPU GPU Linux MacOS Windows Android iOS WASM
Metal No Yes No Yes No No Yes No
Vulkan Yes Yes Yes Yes Yes Yes Yes No
OpenGL No Yes Yes Yes Yes Yes Yes No
WebGpu No Yes No No No No No Yes
Dx11/Dx12 No Yes No No Yes No No No

Dependencies

~37–72MB
~1.5M SLoC