facebookresearch / shumai

Fast Differentiable Tensor Library in JavaScript and TypeScript with Bun + Flashlight
https://facebookresearch.github.io/shumai
MIT License
1.13k stars 26 forks source link

WebGPU Backend #122

Open bwasti opened 1 year ago

bwasti commented 1 year ago

This will enable browser support. We'll need to shim some files:

It doesn't make sense to support anything besides WebGPU at this point. WASM + SIMD is around 15-20x slower on my machine[1]. Although WebGL is more widely supported today, it doesn't have the compute features needed for efficient modern ML (transformers etc) and will likely be a deprecated backend for other frameworks when WebGPU comes online.

[1]: In chrome canary, with Unsafe webGPU enabled try models here: https://tensorflow.github.io/tfjs/e2e/benchmarks/local-benchmark/index.html

josephrocca commented 1 year ago

This is exciting! A cool side effect is that this will add Deno support, since Deno supports WebGPU (since v1.8) just like the browser. I guess it may make sense to eventually support Deno via FFI too as their FFI performance catches up to Bun, but even if it does get added, WebGPU has its own set of advantages (safety, isomorphism, etc.), so it may even be preferred over FFI in some circumstances despite having somewhat lower performance.

Looks like WebGPU support may actually land (in Chromium, at least) early/mid next year - looking forward to having a play around with this and building some demos/benchmarks/etc when it does.

julien-blanchon commented 8 months ago

I'd love to see that happen, @bwasti are you still working on Shumai?

Maybe brr.js 👀