[WebGPU] Vectorize NHWC depthwise convolution across channels - #32498
[WebGPU] Vectorize NHWC depthwise convolution across channels#32498linzuojian wants to merge 2 commits into
Conversation
A depthwise conv - one output channel per group, and as many groups as there are input channels - has a 1:1 correspondence between input, weight and output channels. In NHWC the channel is innermost, so x, w and the output can all be indexed with the same vectorized channel index and one thread can carry four channels at once. The general grouped path cannot do that: its input channels do not line up with output vectors, so it has to index x one scalar channel at a time. Its `output_channels_per_group >= 4` test therefore also leaves the depthwise case scalar, even though that case is exactly the one where the alignment is free. This adds a second body to GroupedConvProgram for the depthwise form. Besides being vectorized it steps the input and weight offsets rather than rebuilding a four-dimensional index per tap: only the width index changes inside the inner loop, and it moves by one channel vector. A 3x3 depthwise conv does nine taps per output and is bound by that address arithmetic, not by its reads, which are fully cached. Because the new body addresses x and w by offset rather than by indices, it does not implicitly pull in the shape uniforms; it still reads them for the loop bounds and the strides it steps by, so ShaderUsage::UseShapeAndStride is requested explicitly. Measured on an RTX 3060 (Dawn/Vulkan) as part of a larger WebGPU optimization pass on a document-layout model: 0.16 ms out of a 16 ms inference. Tests: added two depthwise cases to ConvTest with a channel count that is a multiple of four - one with padding, so the taps that fall outside the input are covered, and one with a dilation, so the tap offsets do not advance by a single input element. Both use per-channel constant weights so a channel that reads the wrong weight or input plane is visible in the output. The full ConvTest suite passes on the WebGPU EP: 43 tests, 41 passed, 2 pre-existing skips, 0 failures.
|
Azure Pipelines: There may be pipelines that require an authorized user to comment /azp run to run. |
|
@microsoft-github-policy-service agree company="Alibaba" |
There was a problem hiding this comment.
🟢 Approval recommended
The optimization is narrowly scoped, consistent with existing componentization patterns, and is backed by new targeted correctness tests (only a minor maintainability nit was found).
Pull request overview
This PR optimizes WebGPU’s NHWC grouped-convolution path by adding a dedicated vectorized shader body for the depthwise case (1 output channel per group, group == input channels), enabling per-thread processing of multiple channels (vec2/vec4) with reduced address arithmetic. It also extends the existing Conv operator test suite with depthwise cases intended to exercise the new vectorized path (including padding and dilation).
Changes:
- Add a depthwise-specific vectorized accumulation path to
GroupedConvProgram’s shader generation. - Update WebGPU Conv dispatch to select depthwise vectorization when applicable and to pass appropriately “reduced-by-components” metadata for the input tensor.
- Add two new depthwise Conv tests (padded + bias, and dilated) with channel counts divisible by 4.
File summaries
| File | Description |
|---|---|
| onnxruntime/test/providers/cpu/nn/conv_op_test.cc | Adds depthwise Conv regression tests targeting vectorized NHWC depthwise correctness (padding + dilation). |
| onnxruntime/core/providers/webgpu/nn/grouped_conv.h | Extends GroupedConvProgram with a depthwise_vec mode flag. |
| onnxruntime/core/providers/webgpu/nn/grouped_conv.cc | Implements the depthwise vectorized shader body and selects it based on depthwise_vec_. |
| onnxruntime/core/providers/webgpu/nn/conv.cc | Detects depthwise-vectorizable cases, adjusts componentization, and wires the flag/metadata into GroupedConvProgram. |
Review details
- Files reviewed: 4/4 changed files
- Comments generated: 1
- Review effort level: Lite
💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
| // output and the input-channel loop collapses. Everything is indexed in whole channel vectors, | ||
| // which is what makes this form worth separating - the general path below must index x one scalar | ||
| // channel at a time. | ||
| std::string CalculateResultDepthwiseVec(const ShaderVariableHelper& x, const ShaderVariableHelper& w) { |
There was a problem hiding this comment.
Done in 2a2170a. Note the pre-existing CanculateResult just above has the same issue; I left it alone to keep this PR scoped, happy to fix both if you'd prefer.
It is only used in this translation unit, so there is no reason to export the symbol from the WebGPU provider binary.
|
Addressed in 2a2170a — Local run on the WebGPU EP: full |
|
CI has run. All four failures are infrastructure, none are caused by this change.
The two tool-download failures landed in the same window as several on #32497, which hit Everything that actually compiles this change passed, including Could someone with the necessary access re-run the failed jobs? I have no write access here, and I'd rather not push a no-op commit just to retrigger. |
Description
A depthwise conv — one output channel per group, and as many groups as there are input channels — has a 1:1 correspondence between input, weight and output channels. In NHWC the channel is innermost, so
x,wand the output can all be indexed with the same vectorized channel index and one thread can carry four channels at once.The general grouped path cannot do that: its input channels do not line up with output vectors, so it has to index
xone scalar channel at a time. Itsoutput_channels_per_group >= 4test therefore also leaves the depthwise case scalar, even though that case is exactly the one where the alignment is free.This adds a second body to
GroupedConvProgramfor the depthwise form. Besides being vectorized it steps the input and weight offsets rather than rebuilding a four-dimensional index per tap: only the width index changes inside the inner loop, and it moves by one channel vector. A 3×3 depthwise conv does nine taps per output and is bound by that address arithmetic, not by its reads, which are fully cached.Because the new body addresses
xandwby offset rather than by indices, it does not implicitly pull in the shape uniforms; it still reads them for the loop bounds and the strides it steps by, soShaderUsage::UseShapeAndStrideis requested explicitly.Motivation and Context
Part of a WebGPU optimization pass on a document-layout model, measured on an RTX 3060 with the Dawn/Vulkan backend: 0.16 ms out of a 16 ms inference.
Testing
Added two depthwise cases to
ConvTestwith a channel count that is a multiple of four — one with padding, so the taps that fall outside the input are covered, and one with a dilation, so the tap offsets do not advance by a single input element. Both use per-channel constant weights, so a channel that reads the wrong weight or input plane is visible in the output.The full
ConvTestsuite passes on the WebGPU EP — 43 tests, 41 passed, 2 pre-existing skips, 0 failures.Verified on Windows / MSVC / NVIDIA (Dawn Vulkan backend) only; I do not have other vendors or backends to hand, so CI is the first run on those.