[WebGPU EP] Add PRelu support - #30512
Conversation
|
Azure Pipelines: There may be pipelines that require an authorized user to comment /azp run to run. |
|
@microsoft-github-policy-service agree |
|
@hariharans29 can you please take a look and help trigger the CI workflows? Thanks! |
There was a problem hiding this comment.
Pull request overview
This PR adds ONNX PRelu operator support to the WebGPU Execution Provider by implementing the kernel via the existing BinaryElementwise infrastructure and registering it for the relevant ONNX opset ranges, with targeted tests and JS Web test-suite enablement.
Changes:
- Implement WebGPU
PReluusing a WGSLselect()expression inbinary_elementwise_ops.cc. - Register
PRelukernels for WebGPU across opsets 7–8, 9–15, and 16 in the WebGPU EP kernel-create table. - Add WebGPU-only float16 activation tests for opset 9 and 16, and enable corresponding
test_prelu_*entries in the JS web test list.
Reviewed changes
Copilot reviewed 4 out of 4 changed files in this pull request and generated 1 comment.
| File | Description |
|---|---|
| onnxruntime/test/providers/cpu/activation/activation_op_test.cc | Adds WebGPU-only float16 PRelu tests for opsets 9 and 16 to cover non-default opset registrations. |
| onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc | Registers PRelu in the WebGPU kernel create-info function table for opsets 7–8, 9–15, and 16. |
| onnxruntime/core/providers/webgpu/math/binary_elementwise_ops.cc | Adds the WebGPU PRelu kernel implementation and its versioned kernel registrations. |
| js/web/test/suite-test-list.jsonc | Enables ONNX backend tests test_prelu_broadcast and test_prelu_example in one suite block. |
| WEBGPU_BINARY_IMPL(PRelu, "select(b * a, a, a >= vec4<input_a_element_t>(0))") | ||
| WEBGPU_BINARY_VERSIONED_KERNEL(PRelu, 7, 8, PRelu, WebGpuSupportedFloatTypes()) | ||
| WEBGPU_BINARY_VERSIONED_KERNEL(PRelu, 9, 15, PRelu, WebGpuSupportedNumberTypes()) | ||
| WEBGPU_BINARY_KERNEL(PRelu, 16, PRelu, WebGpuSupportedNumberTypes()) |
There was a problem hiding this comment.
Narrowed both to WebGpuSupportedFloatTypes(). Flagging that this deviates from #29859, which specifies WebGpuSupportedNumberTypes() for opsets 9-15 and 16.
On the schema point: the float-only constraint applies to opset 7-8 only. Opset 9 widened T to tensor(float16), tensor(float), tensor(double), tensor(uint32), tensor(uint64), tensor(int32), tensor(int64), and the constraint description changed from "Constrain input and output types to float tensors" to "float/int tensors". PRelu-16 adds bfloat16 on top. Per-version constraints are rendered at onnx.ai/onnx/operators/onnx__PRelu.html.
So the original 7-8 vs 9+ split was tracking that opset 9 change, per the issue. Note too that WebGpuSupportedFloatTypes() does not match the schema either, since every PRelu opset including 7 permits tensor(double) and WGSL has no f64. No registration here matches the schema exactly.
I have narrowed anyway since onnxruntime/core/providers/cpu/math/element_wise_ops.h records a decision to keep PRelu float-only:
// Currently only registered for float. Integer types would require addressing
// overflow in multiplication and avoiding 0.0*inf=NaN in branchless formulations.
static_assert(std::is_floating_point_v<T>,
"PRelu is only supported for floating-point types.");
It also makes the registered surface and the tests line up exactly. Every registration is now exercised: opset 7-8 by the existing ActivationOpTest.PRelu* cases plus the newly enabled test_prelu_example and test_prelu_broadcast node tests, opset 9-15 by PRelu_Float16_Opset9_WebGpu, and opset 16 by PRelu_Float16_Opset16_WebGpu. Keeping WebGpuSupportedNumberTypes() would leave int32 and uint32 with no coverage anywhere, since CPU PRelu is float-only and cannot back a shared test.
Happy to open a follow up issue to restore the integer types with int32 and uint32 coverage, which would also need a decision on overflow semantics for the multiply. Happy to revert this narrowing if you would rather the PR match #29859 as written.
There was a problem hiding this comment.
I'm personally OK with the narrow set of types, but it's preferrable to support more types. The AI comment was interesting. Good job on checking the schema.
|
@hariharans29 would you be able to take a look, or help trigger the CI pipelines? Only the CLA check has run so far. Thanks! |
Appreciate your contibution. Please allow us some time to review PRs as some of us are juggling multiple areas. Do give us a shout if something takes more than a few days. As for the CIs, they did run the first time 
They haven't run the second time (after your recent commits). I will trigger them now. |
Description
Add PRelu support for the WebGPU EP: kernel implementation plus registrations for
opsets 7-8, 9-15 and 16. Opset 6 is not registered, matching the CPU and CUDA EPs.
Coverage: the existing ActivationOpTest.PRelu* cases fan out to every registered EP,
so they exercise the new kernel on the scalar/no-broadcast and general-broadcast paths
in BinaryElementwise, plus the inf/NaN edge cases. Those run at OpTester's default
opset (7), so two new float16 cases cover the 9-15 and 16 registrations; the opset-9
case uses a last dim of 4, which is also the only coverage of the vectorized path.
Both use float16 because the CPU EP's PRelu is float-only, so they cannot silently
pass via CPU fallback.
Motivation and Context
Fixes #29859
Without it, PRelu falls back to the CPU EP and Memcpy nodes get inserted around it
when it sits between WebGPU nodes. That costs a GPU→CPU→GPU round trip per inference,
and causes a session-init failure when graph capture is enabled — the scenario named
in the issue.