[WebGPU] Fix 1D-dispatch shader fast path

[Jianhui Dai (daijh)]

Description

ShaderHelper::Init() special-cased shaders whose dispatch grid has a single row of workgroups (dispatch_group_size_y_ == 1 && dispatch_group_size_z_ == 1) by computing global_idx directly from @builtin(global_invocation_id).x and workgroup_idx from @builtin(workgroup_id).x.

This is only correct when the workgroup itself is also 1D (workgroup_size_y == 1 && workgroup_size_z == 1). For a program dispatched as a single row of 2D/3D workgroups (e.g. workgroup_size = (8, 8, 1)), global_id.x omits the contribution of local_invocation_id.y/z that local_invocation_index folds in, so global_idx (and workgroup_idx-derived offsets) come out wrong and the shader reads/writes incorrect elements.

Remove the special-cased fast path so every non-indirect dispatch uses the general num_workgroups-based formula, which is correct regardless of workgroup or dispatch shape. Also drop the now-unnecessary is_1d_dispatch distinction from the program cache key (program_cache_key.h/.cc, webgpu_context.cc), since generated shader source no longer varies with dispatch dimensionality.

Motivation and Context

See above.

[azure-pipelines]

Azure Pipelines:
There may be pipelines that require an authorized user to comment /azp run to run.

[Jianhui Dai (daijh)]

Jiajia Qin (@qjia7) Hariharan Seshadri (@hariharans29) PTAL

[Copilot]

Pull request overview

Fixes incorrect flattened invocation indices for 1D dispatch grids using multidimensional workgroups.

Changes:

• Uses the general workgroup-flattening formula for direct dispatches.
• Removes dispatch dimensionality from program cache keys.
• Updates cache-key interfaces and call sites.

Reviewed changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated 1 comment.

File	Description
webgpu_context.cc	Removes dispatch shape from cache-key construction.
shader_helper.cc	Removes the faulty 1D shader path.
program_cache_key.h	Simplifies the cache-key API.
program_cache_key.cc	Removes the dispatch flag from generated keys.

Suppressed comments (1)

onnxruntime/core/providers/webgpu/shader_helper.cc:109

• Applying this branch to every direct dispatch removes the inexpensive global_id.x path from the many common shaders that use both a 1D grid and a 1D workgroup. Those invocations now execute the workgroup-flattening multiplications and additions even though global_id.x is already the desired index. Please retain the specialization with the corrected predicate (dispatch Y/Z and effective workgroup Y/Z must all be 1), and retain the corresponding cache-key distinction; this fixes the 2D/3D-workgroup bug without regressing the established 1D case.

  } else {
    body_ss_ << ",\n"
                "        @builtin(num_workgroups) num_workgroups : vec3<u32>) {\n";
    body_ss_ << "  let workgroup_idx = workgroup_id.z * num_workgroups[0] * num_workgroups[1] + workgroup_id.y * num_workgroups[0] + workgroup_id.x;\n"
                "  let global_idx = workgroup_idx * (workgroup_size_x * workgroup_size_y * workgroup_size_z) + local_idx;\n";

---

💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Jiajia Qin (qjia7)]

Jianhui Dai (@daijh) Why didn't we catch this kind of error earlier? Is it because all existing shaders use a 1D dispatch with a 1D workgroup size as well?

Another fixing is to also check the workgroupSize.yz except for dispatch.yz.

[Jianhui Dai (daijh)]

Jianhui Dai (Jianhui Dai (@daijh)) Why didn't we catch this kind of error earlier? Is it because all existing shaders use a 1D dispatch with a 1D workgroup size as well?

Another fixing is to also check the workgroupSize.yz except for dispatch.yz.

Our previous practice was to prefer 1D workgroup sizes, which prevented this issue from being triggered.
I recently discovered that the GPU kernel was compiling twice simply due to the dispatch size, resulting in slightly longer launch times

[Jianhui Dai (daijh)]

Another fixing is to also check the workgroupSize.yz except for dispatch.yz.

Prefer always 3D dispatch for the following reasons:

• It simplifies the shader template by covering only two conditions instead of three.
• It aligns with the indirect_dispatch shader template.
• The performance benefit of the 1D dispatch optimization is quite limited.

[Jianhui Dai (daijh)]

Linux CPU Minimal Build E2E / 1. Build Full ORT and Generate ORT Files (pull_request)

The job is failing in the CoreML provider build because uuid/uuid.h is missing:

Failing file: coremltools-src/modelpackage/src/ModelPackage.cpp
Error: fatal error: uuid/uuid.h: No such file or directory
This points to a missing Linux development dependency for libuuid in the build environment used by the Linux CPU minimal workflow. The workflow is building CoreML-related code even though it’s on Linux, and the container/self-hosted image doesn’t have the header available.