New burn-wgpu backend

[antimora]

Feature description

The new upcoming GPU API standard WebGPU is being finalized and available for use. The single API makes possible to access compute/graphical capabilities of various GPU cards. The API is also being standardize and available on browsers.

The WGPU is WebGPU implementation in Rust. Here are its features:

wgpu is a safe and portable graphics library for Rust based on the WebGPU API. It is suitable for general purpose graphics and compute on the GPU.

Applications using wgpu run natively on Vulkan, Metal, DirectX 11/12, and OpenGL ES; and browsers via WebAssembly on WebGPU and WebGL2.

Feature motivation

It would be a great addition to Burn's backend family. wgpu backend will enable acceleration on machines without needing to install addition binaries to access GPU power.

(Optional) Suggest a Solution

WONNX project already did a heavy lifting and writing code all needed operators for burn's backend. So we need to port them for burn.

[antimora]

Pinged WONNX team to clarify their license. It seems it's a dual license as burn's (said in the cargo.toml file).

[antimora]

Implementation status page of WebGPU: https://github.com/gpuweb/gpuweb/wiki/Implementation-Status

[antimora]

Found an interesting example of wgpu usage with webgl feature flag: https://github.com/simbleau/nbody-wasm-sim/blob/main/Cargo.toml#L30. This is useful for web cases when WebGPU is not available in a browser (currently only select browsers have access to WebGPU). Probably wgpu translates wgsl to glsl behind the scene.

It seems wgpu can t

[antimora]

WebGPU is available on Chrome: https://developer.chrome.com/blog/webgpu-release/

Someone run some benchmark: https://github.com/FL33TW00D/wgpu-mm

[antimora]

Good documentation on wgpu: https://sotrh.github.io/learn-wgpu/

[antimora]

Related and interested project: https://lib.rs/crates/custos

[antimora]

https://webgpufundamentals.org/

WebGPU fundamentals

[antimora]

https://github.com/mlc-ai/web-llm

[ksf]

As in #287 the issue of kernels for OpenCL was raised, I'd like to point into the general direction of arrayfire, 3-clause BSD. It's a C++ project but the CL kernels, are, well, CL. On the rust side there's ocl.

While wgpu definitely is more portable OpenCL has some memory addressing features and also semantic guarantees (precision and the like) that just aren't present in graphics-centered APIs. Whether that's an issue probably depends more on what people want to do with burn than burn itself.

OpenCL availability should also be decent, at least under linux, especially when mesa 23.1 releases (it's almost done), promoting rusticl from an experimental to stable feature. Nvidia and AMD don't give OpenCL much love, Apple right-out axed it, but it's also far from dead.

[nathanielsimard]

Thanks for the insight, it's good to know. WebGPU is good for portability by being compiled to wasm but it may not be the best option to provide the fastest backend. OpenCL is normaly slower on Nvidia devices, so maybe the best option is not to aim for a single backend that compiles to every platform, but provide multiple backends more specialized for each platform.

[antimora]

@nathanielsimard has the first PR for the backend: #376

(if anyone is interested reviewing it)

[antimora]

One way you can gauge completion is by searching todo for missing OP implementations in burn-wgpu like this: https://github.com/search?q=repo%3Aburn-rs%2Fburn%20burn-wgpu%20todo&type=code

[antimora]

OK. A working first version of burn-wgpu (v0.8.0) has been released thanks to @nathanielsimard and @louisfd

https://github.com/burn-rs/burn/releases/tag/v0.8.0