[WebGPU] Add Gather int64 support and make kernel version numbers function params - #31714
[WebGPU] Add Gather int64 support and make kernel version numbers function params#31714Bin Miao (miaobin) wants to merge 3 commits into
Conversation
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: There may be pipelines that require an authorized user to comment /azp run to run. |
There was a problem hiding this comment.
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
Gatherint64support (lossless copy via storage type accessors) gated byenable_int64, plus a targeted WebGPU-only test that disables CPU fallback. - Register
GatherviaRegisterKernels()to allow conditionalint64type 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. |
Edward Chen (edgchen1)
left a comment
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
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.
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 )