Skip to content

wgpu runtime can't honor a requested plane size (subgroup size control) #1451

Description

@SamuelBelanger

CubeCl kernels can't request a specific plane (subgroup/wave) size on the wgpu runtime. There is no working path for it today (not through WGSL, not through SPIR-V).

This matters wherever the plane size is a tuning knob rather than a constant: wave32 vs wave64 on AMD, 8/16/32 on Intel. Without it the driver picks any size in [plane_size_min, plane_size_max], and on Vulkan different planes within the same dispatch may even get different sizes.

Why it doesn't work

SPIR-V. The obvious route OpExecutionMode SubgroupSize is not usable in Vulkan. It requires the SubgroupDispatch capability, which is OpenCL-only, so the module is rejected outright rather than merely warned about. On Mesa/Intel:

SPIR-V offset 516: SPIR-V parsing FAILED:
    b->shader->info.stage == MESA_SHADER_KERNEL
spirv_to_nir failed (VK_ERROR_UNKNOWN)

Declaring the capability doesn't help. Vulkan doesn't support it either.

Vulkan sets the required subgroup size at pipeline creation instead, via VkPipelineShaderStageRequiredSubgroupSizeCreateInfo (VK_EXT_subgroup_size_control, core in 1.3). wgpu and wgpu-hal 29/30 don't expose it, and we create pipelines through the wgpu API (crates/cubecl-wgpu/src/backend/base.rs), so there's no seam to reach it from our side.

WGSL has no notion of it at all. The WebGPU spec proposal is still open: gpuweb/gpuweb#5578.

Upstream status

gfx-rs/wgpu#9523 is open and adds subgroup_size to PassthroughShaderEntryPoint which is exactly the path we already use for SPIR-V. Following or helping land that looks better than carrying a fork.

A local prototype confirms the gap is only API surface: ~50 lines across wgpu/wgpu-core/wgpu-hal, threading an Option<u32> from the compute pipeline descriptor into the Vulkan backend, and PLANE_DIM follows the request exactly (8/16/32 on an Intel Xe whose default is 16).

Impact

Any kernel-level plane size option is inert on the wgpu runtime until this lands upstream. CUDA is unaffected (warp size is always 32), so this is a Vulkan/Metal-shaped gap.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions