fix(wgpu): clamp compute dispatch cube count to u16::MAX (NVIDIA Vulkan)#1388
Open
thewtex wants to merge 6 commits into
Open
fix(wgpu): clamp compute dispatch cube count to u16::MAX (NVIDIA Vulkan)#1388thewtex wants to merge 6 commits into
thewtex wants to merge 6 commits into
Conversation
Some adapters — notably the Vulkan backend on NVIDIA hardware — report a `max_compute_workgroups_per_dimension` larger than the limit the driver actually validates against at dispatch time. `create_server` fed that adapter-reported value straight into `HardwareProperties::max_cube_count`, so the launch logic treated an oversized `[N, 1, 1]` cube count as legal and let it reach wgpu, which rejected the dispatch with a validation error once `N` exceeded the real cap of `u16::MAX` (65_535). This bites any single large elementwise dispatch: e.g. a 52,428,800-element launch needs 204,800 workgroups along X (52_428_800 / 256), well past 65_535. Clamp the reported limit to `u16::MAX`, keeping it consistent with the cap already advertised by `<WgpuRuntime as Runtime>::max_cube_count`. The clamp is factored into a small `clamp_max_cube_count` helper with unit tests.
`register_vulkan_features` computed `storage8` from `ext_feat.buf_16` and its `uniform_and_storage_buffer16_bit_access` field — a copy-paste of the `storage16` line above it. As a result 8-bit storage capability tracked the 16-bit feature rather than the actual `PhysicalDevice8BitStorageFeatures`, so I8/U8 types could be marked storable based on the wrong device feature. Read `ext_feat.buf_8.uniform_and_storage_buffer8_bit_access` instead. Not unit-testable without a live Vulkan device; verified to compile under `--features spirv`.
The error printed for an unrecognised backend reads the `AUTO_GRAPHICS_BACKEND` env var but told the user to check `GRAPHICS_BACKEND`, sending them to the wrong variable. Use the correct name in the message.
`contiguous_strides` iterated `(0..rank - 1)`, which underflows and panics when `rank == 0` (a scalar/rank-0 shape). Use `rank.saturating_sub(1)` so a rank-0 shape yields empty strides instead of panicking. Adds unit tests for the rank-0 case and row-major correctness.
Same rank-0 underflow as the wgpu `contiguous_strides` fix: the `(0..rank - 1)` loop bound underflows and panics when `rank == 0` (a scalar/rank-0 shape). Use `rank.saturating_sub(1)` so a rank-0 shape yields empty strides. Adds unit tests for the rank-0 case and row-major correctness.
Same rank-0 underflow as the wgpu `contiguous_strides` fix: the `(0..rank - 1)` loop bound underflows and panics when `rank == 0` (a scalar/rank-0 shape). Use `rank.saturating_sub(1)` so a rank-0 shape yields empty strides. The pitched-allocation path in `ContiguousMemoryLayoutPolicy` already guards its `rank - 2` loop behind `rank > 2`, so only this helper needed the fix. Adds unit tests for the rank-0 case and row-major correctness.
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
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.
Summary
A small batch of independent correctness fixes in
cubecl-wgpu, plus the samerank-0 stride bug fixed in two sibling crates that carry a copy of the helper.
Each fix is its own commit so they can be reviewed.
Fixes
1. Clamp dispatch cube count to the enforced per-dimension limit (
cubecl-wgpu)Some adapters — notably the Vulkan backend on NVIDIA hardware — report a
max_compute_workgroups_per_dimensionlarger than the limit the driveractually validates against at dispatch time.
create_serverfed thatadapter-reported value straight into
HardwareProperties::max_cube_count, sothe launch logic treated an oversized
[N, 1, 1]cube count as legal and letit reach wgpu, which then rejected the dispatch with a validation error once
Nexceeded the real cap ofu16::MAX(65_535).This bites any single large elementwise dispatch — e.g. a 52,428,800-element
launch needs
52_428_800 / 256 = 204_800workgroups along X, well past 65_535.The fix clamps the reported value to
u16::MAXvia a newclamp_max_cube_counthelper, keeping it consistent with the cap already advertised by
<WgpuRuntime as Runtime>::max_cube_count.2. Detect 8-bit storage via the 8-bit feature struct (
cubecl-wgpu)In
register_types,storage8was computed fromext_feat.buf_16and itsuniform_and_storage_buffer16_bit_accessfield — a copy-paste of thestorage16line above it. As a result 8-bit storage capability tracked the16-bit device feature, so I8/U8 types could be marked storable based on the
wrong feature. Now reads
ext_feat.buf_8.uniform_and_storage_buffer8_bit_access.3. Name
AUTO_GRAPHICS_BACKENDin its invalid-value error (cubecl-wgpu)The error printed for an unrecognised backend reads the
AUTO_GRAPHICS_BACKENDenv var but told the user to check
GRAPHICS_BACKEND, sending them to the wrongvariable. Corrected the message.
4. Guard
contiguous_stridesagainst rank-0 shapes (cubecl-wgpu,cubecl-cpu,cubecl-runtime)contiguous_stridesiterated(0..rank - 1), which underflows and panics whenrank == 0(a scalar / rank-0 shape). Changed to(0..rank.saturating_sub(1))so a rank-0 shape yields empty strides. The same helper is duplicated in
cubecl-cpuandcubecl-runtime, fixed in all three.Notes:
cubecl-std'scontiguous_stridesuses a different, subtraction-freeimplementation and is not affected.
ContiguousMemoryLayoutPolicy(
cubecl-runtime) already guards itsrank - 2loop behindrank > 2, soonly the standalone helper needed the fix.
Testing
clamp_max_cube_count(caps inflated limits, preserveshonest ones, agrees with
WgpuRuntime::max_cube_count) and forcontiguous_strides(rank-0 returns empty, row-major correctness) in each ofthe three crates.
cargo test,cargo clippy --all-targets, andcargo fmt --checkpass forthe touched crates;
cubecl-wgpuadditionally checked under--features spirvto compile the Vulkan path.