Skip to content

[WebGPU] Add Gather int64 support and make kernel version numbers function params - #31714

Open
Bin Miao (miaobin) wants to merge 3 commits into
microsoft:mainfrom
miaobin:webgpu-gather-int64
Open

[WebGPU] Add Gather int64 support and make kernel version numbers function params#31714
Bin Miao (miaobin) wants to merge 3 commits into
microsoft:mainfrom
miaobin:webgpu-gather-int64

Conversation

@miaobin

Copy link
Copy Markdown
Contributor

Gather is a pure data-movement op (elements are copied, never interpreted in shader arithmetic), so int64 can be supported safely. int64 (stored as vec2) is copied losslessly via the raw storage-word path, preserving the full 64-bit value instead of the truncating i32 value type used by arithmetic kernels. int64 is only added to the "T" constraint when the enable_int64 provider option is set.

This PR also refactor the int64-migrated kernels' factory functions to take the version range as runtime function parameters instead of template parameters (StartVersion/EndVersion). This removes the explicit template instantiations (reducing binary size) and eliminates the duplication of version ranges between the instantiations and the registrations in webgpu_execution_provider.cc; the ranges now live only at the registration site. (This is a follow up fix of #31049 )

Gather is a pure data-movement op (elements are copied, never interpreted
in shader arithmetic), so int64 can be supported safely. int64 (stored as
vec2<u32>) is copied losslessly via the raw storage-word path, preserving
the full 64-bit value instead of the truncating i32 value type used by
arithmetic kernels. int64 is only added to the "T" constraint when the
enable_int64 provider option is set.

Also refactor the int64-migrated kernels' factory functions to take the
version range as runtime function parameters instead of template
parameters (StartVersion/EndVersion). This removes the explicit template
instantiations (reducing binary size) and eliminates the duplication of
version ranges between the instantiations and the registrations in
webgpu_execution_provider.cc; the ranges now live only at the registration
site.
@azure-pipelines

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

@miaobin

Copy link
Copy Markdown
Contributor Author

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

This PR extends the WebGPU Execution Provider’s Gather kernel to support int64 tensors (behind the existing enable_int64 provider option) by copying values via the raw storage-word path (vec2<u32>) so full 64-bit payloads are preserved. It also refactors several int64-gated WebGPU kernel factory functions to take opset version ranges as runtime parameters (instead of template parameters), removing explicit template instantiations and centralizing version-range ownership at the registration site in webgpu_execution_provider.cc.

Changes:

  • Add WebGPU Gather int64 support (lossless copy via storage type accessors) gated by enable_int64, plus a targeted WebGPU-only test that disables CPU fallback.
  • Register Gather via RegisterKernels() to allow conditional int64 type constraints (removing it from the static build-kernel table).
  • Refactor multiple WebGPU kernel factory APIs (Cast/Unsqueeze/Expand/Reshape/Concat/Tile/Where/ReduceSum/Add/Sub/Equal) to accept opset version numbers as function parameters and drop explicit template instantiations.

Reviewed changes

Copilot reviewed 22 out of 22 changed files in this pull request and generated no comments.

Show a summary per file
File Description
onnxruntime/test/providers/cpu/tensor/gather_op_test.cc Adds a WebGPU-only int64 Gather test gated by provider option and disables CPU fallback to ensure the WebGPU kernel runs.
onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc Moves Gather registration into RegisterKernels() and updates registrations to pass opset ranges as runtime parameters for int64-gated kernels.
onnxruntime/core/providers/webgpu/tensor/gather.h Extends GatherProgram to carry an is_int64 flag and declares new CreateGather*KernelInfo factory functions.
onnxruntime/core/providers/webgpu/tensor/gather.cc Implements lossless int64 Gather via use_storage_type=true read/write and adds kernel factory functions with conditional type constraints.
onnxruntime/core/providers/webgpu/tensor/cast.h Changes Cast kernel factory declaration to runtime opset parameters and documents the open-ended registration convention.
onnxruntime/core/providers/webgpu/tensor/cast.cc Refactors Cast kernel factory implementation to runtime opset parameters and removes explicit template instantiations.
onnxruntime/core/providers/webgpu/tensor/unsqueeze.h Updates Unsqueeze kernel factory declarations to runtime opset parameters.
onnxruntime/core/providers/webgpu/tensor/unsqueeze.cc Refactors Unsqueeze kernel factory implementations to runtime opset parameters and removes explicit template instantiations.
onnxruntime/core/providers/webgpu/tensor/expand.h Updates Expand kernel factory declarations to runtime opset parameters.
onnxruntime/core/providers/webgpu/tensor/expand.cc Refactors Expand kernel factory implementations to runtime opset parameters and removes explicit template instantiations.
onnxruntime/core/providers/webgpu/tensor/reshape.h Updates Reshape kernel factory declarations to runtime opset parameters.
onnxruntime/core/providers/webgpu/tensor/reshape.cc Refactors Reshape kernel factory implementations to runtime opset parameters and removes explicit template instantiations.
onnxruntime/core/providers/webgpu/tensor/concat.h Updates Concat kernel factory declarations to runtime opset parameters.
onnxruntime/core/providers/webgpu/tensor/concat.cc Refactors Concat kernel factory implementations to runtime opset parameters and removes explicit template instantiations.
onnxruntime/core/providers/webgpu/tensor/tile.h Updates Tile kernel factory declarations to runtime opset parameters.
onnxruntime/core/providers/webgpu/tensor/tile.cc Refactors Tile kernel factory implementations to runtime opset parameters and removes explicit template instantiations.
onnxruntime/core/providers/webgpu/tensor/where.h Updates Where kernel factory declarations to runtime opset parameters.
onnxruntime/core/providers/webgpu/tensor/where.cc Refactors Where kernel factory implementations to runtime opset parameters and removes explicit template instantiations.
onnxruntime/core/providers/webgpu/reduction/reduction_ops.h Updates ReduceSum kernel factory declarations to runtime opset parameters.
onnxruntime/core/providers/webgpu/reduction/reduction_ops.cc Refactors ReduceSum kernel factory implementations to runtime opset parameters and removes explicit template instantiations.
onnxruntime/core/providers/webgpu/math/binary_elementwise_ops.h Updates Add/Sub/Equal kernel factory declarations to runtime opset parameters.
onnxruntime/core/providers/webgpu/math/binary_elementwise_ops.cc Refactors Add/Sub/Equal kernel factory implementations to runtime opset parameters and removes explicit template instantiations.

Comment thread onnxruntime/core/providers/webgpu/tensor/gather.cc Outdated

@edgchen1 Edward Chen (edgchen1) left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

thanks for updating the other kernel info creation helpers too

use_storage_type is only honored by GetByOffset/SetByOffset for
Int64/Uint64; for every other type the parameter is ignored, so passing
is_int64_ directly collapses the previous is_int64_ ? ... : ... branching
into a single call, addressing reviewer feedback about the extra value
variable and duplicated branches.

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Copilot reviewed 22 out of 22 changed files in this pull request and generated no new comments.

Suppressed comments (1)

onnxruntime/core/providers/webgpu/tensor/cast.h:55

  • The comment implies the start_version==end_version case is only for the “latest opset” registration, but this function is also called with equal start/end for non-latest versions (e.g., opset 23). This wording is misleading and makes it unclear that the equality case is simply the open-ended SinceVersion(start_version) registration pattern.
// Create Cast kernel info with appropriate type constraints based on int64 support.
// Passing start_version == end_version registers an open-ended kernel (SinceVersion(start_version)),
// matching the "latest opset" registration; otherwise a bounded [start_version, end_version] range.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

ep:WebGPU ort-web webgpu provider

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants