WGPU compute

[kushalkolar]

I thought I'd make an issue to pool together ideas and prototypes on general purpose compute with WGPU.

Performance summary:

• We can achieve performance on-par with torch-cuda on sparse operations and on some dense operations (such as matrix-vector multiply).
• Dense matrix-multiply will require using cooperative math in WGPU, I needs to look into it more. This will use tensor cores (on nvidia) and the equivalent hardware on other GPUs (including AMD and Mac AFAIK).
• subgroup operations will help performance which should be coming to wgpu very soon.

Prototypes:

• sparse matrix-vector multiply combined with visualization that demonstrates a real usecase: https://github.com/kushalkolar/wgpu-kernels
• tinygrad backend made with the help of claude: https://github.com/fastplotlib/tinygrad-sandbox

Easy things to try:

• Benchmark tinygrad matmul, it might already use cooperative math.