This is very exciting! (I had suspected it would slip to 114)
WebGPU implementations are still pretty immature, but certainly enough to get started with. I've been implementing a Rust + WebGPU ML runtime for the past few months and have enjoyed writing WGSL.
That said, matmuls are still pretty handicapped in the browser (especially considering the bounds checking enforced in the browser). From my benchmarking I've struggled to hit 50% of theoretical FLOPS, which is cut down to 30% when the bounds checking comes in. (Benchmarks here: https://github.com/FL33TW00D/wgpu-mm)
I look forward to accessing shader cores as they mentioned in the post.
Burn looks interesting... so I notice in the sample code you mention:
use burn_tch::TchBackend;
I don't see "TchBackend" mentioned in the docs. Is Tch = torch and Burn can be used as a Torch (or tch-rs) wrapper? Or am I imagining too much from a miscellaneous three letter prefix?
I've been looking for a good way to do standard modern AI stuff in Rust, and was leaning towards tch-rs because it exists and seems reasonably frequently maintained, but if you have some sort of reason to use Burn instead then I'm curious to hear it.
That's right, Tch stands for torch. We use tch-rs (bindings to torch++).
The selling point of burn is that it offers the full level of deep learning framework with ability to swap backends. This is useful, for example, I can train with Torch backend on GPUs but if I want to deploy, lets say to embedded device, I can use NDArray backend, which supports pure rust code (also with various BLAS acceleration if needed). You can also use NDArray backend with Accelerate (iOS or MacOS blass). You can even compile for WASM because Burn supports inference with no_std (you can try this online demo https://burn-rs.github.io/demo).
The framework is well written and follows Rust's conventions and best practices. It's still evolving, however.
Give it a try. If you find problems, you can file a ticket or join discord chat.
It’s better to compare against an ML framework than to maximum theoretical flops because sometimes it’s not possible to reach. These models are often limited by memory bandwidth rather than flop capability.
Something more than 30 years ago, I had the privilege of working as a young compiler writer for [a supercomputer designer]'s penultimate start-up. I once naively asked him why the machine he was working on couldn't have more memory bandwidth, since the floating-point functional units were sometimes starved for operand data and it was hard to hit the peak Flop/sec figures. And his response has stuck with me ever since; basically, it makes better sense for the memory to be fully utilized, not the floating-point units, because the memory paths were way more expensive than the floating-point units. And this was something you could actually physically see through the transparent top of the system's case. I guess the lesson would be: Don't let a constraint that would be fairly cheap to overdesign be the limiting factor in a system's performance.
In my opinion, ONNX is more complex than necessary. Therefore, I opted to convert it to an intermediate representation (IR) first, which is then used to generate source code. A key advantage of this approach is the ease of merging nodes into corresponding operations, since ONNX and Burn don't share the same set of operators.
Actually WONNX also transforms to an IR first (early versions did not and simply translated the graph 1:1 to GPU shader invocations in topographically sorted order of the graph). In WONNX the IR nodes are (initially) simply (copy-on-write references to) the ONNX nodes. This IR is then optimized in various ways, including the fusion of ONNX ops (e.g. Conv+ReLU->ConvReLU). The newly inserted node still embeds an ONNX node structure to describe it but uses an internal operator.
oh cool! will this be numpy-like or will it have autograd as well? We're looking around for a web backend for shumai[1] and the former is really all we need :)
I guess all the AI type use cases are front seat here but the performance boost increase to Immersive Web (WebXR) and general moves towards webpages as 3D UI/UX for applications -- that is where I hope to see expansion into new creative territory.
what would it take to python -> wasm -> webgpu for the entire existing webgpu ecosystem (all of the libraries around neural networks, torch, yada yada)
The Apache TVM machine learning compiler has a WASM and WebGPU backend, and can import from most DNN frameworks. Here's a project running Stable Diffusion with webgpu and TVM [1].
Questions exist around post-and-pre-processing code in folks' Python stacks, with e.g. NumPy and opencv. There's some NumPy to JS transpilers out there, but those aren't feature complete or fully integrated.
WebGPU implementations are still pretty immature, but certainly enough to get started with. I've been implementing a Rust + WebGPU ML runtime for the past few months and have enjoyed writing WGSL.
I recently got a 250M parameter LLM running in the browser without much optimisation and it performs pretty well! (https://twitter.com/fleetwood___/status/1638469392794091520)
That said, matmuls are still pretty handicapped in the browser (especially considering the bounds checking enforced in the browser). From my benchmarking I've struggled to hit 50% of theoretical FLOPS, which is cut down to 30% when the bounds checking comes in. (Benchmarks here: https://github.com/FL33TW00D/wgpu-mm)
I look forward to accessing shader cores as they mentioned in the post.