burn_wgpu

Type Alias Wgpu

Source 
pub type Wgpu<F = f32, I = i32, B = u32> = Fusion<CubeBackend<WgpuRuntime, F, I, B>>;
Available on crate feature fusion only.
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.

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

This version of the wgpu 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<F = f32, I = i32, B = u32> { /* private fields */ }