onnxruntime
73534a26 - [WebGPU] Vectorize NHWC depthwise convolution across channels (#32498)

Commit
10 days ago
[WebGPU] Vectorize NHWC depthwise convolution across channels (#32498) ### 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`, `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 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 `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. ### 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 `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. 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.
Author
Parents
Loading