burn/crates/burn-wgpu

Burn WGPU Backend

Burn WGPU backend

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.

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