[ExecuTorch][WebGPU] Add conv_with_clamp op (et_vk.conv_with_clamp)#21217
[ExecuTorch][WebGPU] Add conv_with_clamp op (et_vk.conv_with_clamp)#21217JCNTH wants to merge 1 commit into
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/21217
Note: Links to docs will display an error until the docs builds have been completed. ❗ 1 Active SEVsThere are 1 currently active SEVs. If your PR is affected, please view them below: ❌ 46 New Failures, 3 Unrelated FailuresAs of commit e64fee7 with merge base 266e0dc ( NEW FAILURES - The following jobs have failed:
FLAKY - The following jobs failed but were likely due to flakiness present on trunk:
BROKEN TRUNK - The following job failed but were present on the merge base:👉 Rebase onto the `viable/strict` branch to avoid these failures
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
This PR needs a
|
psiddh
left a comment
There was a problem hiding this comment.
Approving full WebGPU stack
Stack from ghstack (oldest at bottom):
Problem: The WebGPU delegate has no
et_vk.conv_with_clamp— a fused fp32 2D convolution + output clamp (e.g. aConv2dfollowed byrelu6). It delegates throughVulkanPartitioner(op_registry.pytagset_vk.conv_with_clamp.default;nn.Conv2d+F.relu6fuses to it) but had no WebGPU handler.Solution: Port
et_vk.conv_with_clamp(fp32, groups==1), mirroring the Vulkanconvhandler (impl/Convolution.cpp:837) + the eagerclamp(convolution(...), output_min, output_max)(custom_ops_lib.py). Direct windowed conv reusing the stagedq8ta_conv2dwindowing structure but fp32 (no int8 unpack/dequant/requant) + a clamp epilogue: per NCHW output cellout[n,oc,oh,ow] = clamp(bias[oc] + Σ_{ic,kh,kw} weight[oc,ic,kh,kw] * input[n,ic, oh*sh-ph+kh*dh, ow*sw-pw+kw*dw], output_min, output_max), OOB taps skipped.Implementation:
conv_with_clamp/ConvWithClamp.cppregisterset_vk.conv_with_clamp.default->conv_with_clamp_impl. Args[in, weight, bias?, stride, padding, dilation, transposed, output_padding, groups, output_min?, output_max?, out]; stride/padding/dilation as int-pairs,output_min/output_maxasScalar?(Double/Int -> value, absent -> -/+inf soclampis a passthrough).conv_with_clamp.wgsl: one thread per output element, RAW[OC,IC,Kh,Kw]weight (the fp32 buffer path gets the raw constant, like conv1d_dw/pw — not the q8ta im2col[OC,Kh*Kw*IC]), 2D-folded dispatch.ConvWithClampParams(80 bytes) matches the WGSL Params. Guards: in/weight/out 4D + fp32,weight.dims[1]==IC, numel<=u32, groups==1 + not-transposed, all fail-loud.Constraints / divergences from the Vulkan reference: buffer re-derivation of Vulkan's texture conv (
conv2d_prepack_weights.glsl) — the naive one-thread-per-output form is forced by the buffer-only backend. Scoped togroups==1+ not-transposed (fail-loud otherwise). Note: overlaps the generalconvolutionop — this handles the clamped variant only.Differential Revision: D112257644