Skip to main content

Wgpu

Type Alias Wgpu 

Source
pub type Wgpu = Fusion<CubeBackend<WgpuRuntime<AutoCompiler>>>;
Expand description

Tensor backend that uses the wgpu crate for executing GPU compute shaders.

This backend can target multiple graphics APIs, including:

  • Vulkan on Linux, Windows, and Android.
  • OpenGL on Linux, Windows, and Android.
  • DirectX 12 on Windows.
  • Metal on Apple hardware.
  • WebGPU on supported browsers and wasm runtimes.

The selected graphics API is chosen automatically at runtime, and the appropriate shader compiler (WGSL, SPIR-V or MSL) is dispatched via AutoCompiler. When the target API is known ahead of time, prefer the dedicated Vulkan, WebGpu or Metal backend aliases (enabled by their respective Cargo features), which lock the compiler at compile time and avoid the runtime dispatch.

To configure the wgpu backend, eg. to select what graphics API to use or what memory strategy to use, you have to manually initialize the runtime. For example:

fn custom_init() {
    let device = Default::default();
    burn::backend::wgpu::init_setup::<burn::backend::wgpu::graphics::Vulkan>(
        &device,
        Default::default(),
    );
}

will mean the given device (in this case the default) will be initialized to use Vulkan as the graphics API. It’s also possible to use an existing wgpu device, by using init_device.

§Notes

When the fusion feature flag is enabled (the default), this backend uses burn_fusion to compile and optimize streams of tensor operations for improved performance. You can disable the fusion feature flag to remove that functionality, which might be necessary on wasm for now.

Aliased Type§

pub struct Wgpu { /* private fields */ }