coreylowman / dfdx

Deep learning in Rust, with shape checked tensors and neural networks

Other

1.71k stars 99 forks source link

feat(webgpu): add tensor operations #904

Open DonIsaac opened 9 months ago

DonIsaac commented 9 months ago

What This PR Does

Starts a general binary op implementation for Webgpu. This code is based on an (old, in-progress WGPU branch)[https://github.com/DonIsaac/dfdx/tree/don/feat/wgpu2] I made a while ago.

I've (mostly) got forward down, but I think I need some help with backward

Todo

[x] Implement BinaryOpWebgpuKernel.forward()
[ ] Implement UnaryOpWebgpuKernel.forward() (may require push constants)
[x] Implement binary add
[x] Implement to_dtype
[ ] Pipeline caching (may require std)
[x] Support f16 via shader-f16 feature

Other Notes

WebGPU does not support f64 (refer to this issue). If we want to support them We'll need to use a polyfill (e.g. this one]
We may need to consider buffer usage flags when caching tensors. Right now, they're all COPY_SRC and COPY_DST, but this leaves them un-mappable without an intermediate MAP_READ buffer.

DonIsaac commented 9 months ago

Note: I've commented out todo implementations for all BinaryKernel and UnaryKernel traits, and I've commented out the top-level exports for Webgpu while I write code. This is so I can build it without getting 100+ compilation errors; I'll re-add these later as needed.