Give training its own GPU device - #526
Open
ArthurBrussee wants to merge 8 commits into
Open
Conversation
The viewer and training shared one wgpu device, so GUI work and training contended on the same queue. Training now opens its own device and the backbuffer copies the frame across, which costs a readback per displayed frame and takes GUI work off training's queue entirely. The kernels are now generic over `R: CubeRuntime` rather than hardcoding the wgpu runtime, and device setup goes through burn's dispatch layer, so adding a non-wgpu cubecl runtime is a matter of wiring one up rather than touching the kernels. Device creation is lazy: whoever needs it first opens it, which removes the init handshake the app had to perform. Also updates burn and cubecl. Both of the cubecl patches we carried are upstream now, so the fork and its 16-crate patch table are gone.
Validation copies every checked tensor back to the host, which syncs the device — so a validated render can cost several times the render. The only guard was an `args().any(|a| a == "--bench")` check, which catches divan benches and nothing else, so timing harnesses built on `cargo test` silently measured NaN scans instead of rendering. Now it is a runtime switch (`validation::set_enabled`), the timing helpers turn it off themselves, and the first readback says out loud that timings taken now are not representative.
Pins flate2 back to 1.1.9: 1.1.10 moves to miniz_oxide 0.9 while png, exr and backtrace are still on 0.8, so it lands a second copy in the tree and `cargo deny check bans` fails.
burn's cubecl runtime refactor (#5528) makes `memory_usage` return the usage directly rather than a Result. CI catches three things the local checks missed: clippy runs with `--all-features -D warnings`, so the `unimplemented!` in the Autodiff loss arm has to go (it is `unreachable!`, which is also the accurate word: the backward is hand-rolled on the inner backend), the spell check splits `flate2` into "flate", and `cargo update` refloats that pin.
Without validation compiled in nothing calls it, and the wasm job builds with -D warnings.
burn re-exports cubecl as `burn::cubecl` behind its `cubecl` feature, so nothing here needs to name the crate to use it. That leaves the handful of items burn does not re-export: CubeRuntime, FusionCubeRuntime, CubeTensor, into_contiguous and zeros_client. brush-cube already exists to host the cube glue, so it re-exports those and everything downstream takes them from there. `burn-cubecl` goes from ten manifests to one.
burn re-exports cubecl as `burn::cubecl` behind its `cubecl` feature, so the crates that only wanted cubecl itself no longer name it. The crates that write cube kernels still depend on burn-cubecl directly for CubeRuntime and friends, which is what burn's own custom-kernel example does.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Training shared the viewer's wgpu device, so GUI work and training competed for
the same queue. Training now opens its own, and the viewer copies each displayed
frame across. That costs a readback per displayed frame and takes GUI work off
training's queue entirely.
The kernels are generic over
R: CubeRuntimenow rather than naming the wgpuruntime, and device setup goes through burn's dispatch layer, so wiring up
another cubecl runtime later doesn't mean touching the kernels.
Also updates burn and cubecl, and makes validation readbacks switchable so they
stop silently landing in timing runs.