Autogen webgpu dawn, removing wgpu-py dependency (f16 support part 1)

[wpmed92]

This PR changes the WebGPU runtime used from wgpu to Dawn (used by Chrome):

• wgpu is not closely following the WebGPU specs, it doesn't impose a limit on storage buffers per kernel, but Dawn does, so some tests had to be skipped due to more than the supported number of buffers per kernel
• the python bindings used are created by me, it's published as dawn-python (supports both Mac (Metal), and Linux (Vulkan)
• WebGPU specs doesn’t allow using the same buffer for input and output, JIT does this, and was fine on wgpu (non spec compliant). I patch it in ops_webgpu

Here is a demo of f16 stable diffusion deployed from another branch that includes the f16 support: https://wpmed92.github.io/stable-diffusion-tinygrad-f16/
I see around 30% perf gain from f16 in SD.
The f16 wgsl change is coming in a follow-up PR, I didn't want to put these two relatively big changes into 1 PR.

[geohot]

Why are there new skips and that for loop for input/output? Why wasn't this a problem in wgpu? Read the PR message, understood.

I'm worried about the long term maintenance of "dawn-python", is it possible to do this with a ctypes autogen script and put this in the tinygrad repo instead? It looks like it's doing this in the dawn-python package already. We're moving all the backends away from having Python deps.

[wpmed92]

@geohot Yeah, i created dawn-python using clang2py, which we use for autogen, i can move it into the tinygrad repo. I think it's fine to grab the underlying prebuilt binaries from my repo though. (dawn.so, dawn.dylib)

[geohot]

Yea, I'm fine with grabbing the bins from a GitHub release (like remu or ocelot)

[geohot]

Progress on adding dawn bindings to autogen?

[wpmed92]

@geohot I moved dawn-python stuff to tinygrad autogen, I have all tests passing locally on my Mac, but have to clean it up and make it pass in CI.

[geohot]

I don't think there should be a support for this. Why can't it go in the webgpu file?

[wpmed92]

@geohot

• removed support/webgpu and moved the code into ops_webgpu
• refactored callback info creation -> removed a lot of the duplications
• Fixed autogen verification
• Fixed mypy

Everything is passing, except for the line limit test.
I'm 100 lines above the limit, but there's more to save in ops_webgpu than what i saved so far, so will work on it more.

[wpmed92]

From the inital 421 lines ops_webgpu it's down to 324, need to save only 28 more.

[geohot]

Why is this a variable?

[geohot]

Why is this a variable? (multiple places)

[wpmed92]

It consistently segfaults if it's not saved into a variable before passing it in.

[github-actions]

Changes

Name                              Lines    Diff    Tokens/Line    Diff
------------------------------  -------  ------  -------------  ------
tinygrad/runtime/ops_webgpu.py      171    +113           14.2    -2.0


total lines changes: +113

[geohot]

Bumped the line limit to merge this, really happy to have that wgpu dep gone.

[Korijn]

Bumped the line limit to merge this, really happy to have that wgpu dep gone.

Hello, just checking in here as one of the maintainers of wgpu-py, would any of you mind summarizing your reasons to move away from our library? Anything we could do better? Thanks!

cc @almarklein

[wpmed92]

@Korijn
The reason for the removal has nothing to do with something wgpu-py did wrong. I love wgpu-py and used it as a reference when rewriting our runtime with Dawn (see read_buffer implementation as an example).
The removal has to do with the general direction of the project and the tinygrad philosophy: reducing complexity, and removing dependencies.
We try to not use python interface libs, i.e. libs that provide a wrapper interface around a binary. Instead, we autogen this interface with clang2py, and we only depend on the binary (Dawn in this case). As to why we chose dawn and not wgpu as the lib: dawn is more mature than wgpu, and we want our webgpu runtime to behave as close as possible to how an exported model behaves in Google Chrome, the most widely used browser.
We also removed llvmlite and use autogened llvm instead. We also use autogened opencl, etc.
As to what could be better, but it’s still not wgu-py: I think that maybe wgpu (not wgpu-py) is moving a bit slowly, f16 is still not merged although it’s close. (the pr is open since last may)
cc: @geohot