[ExecuTorch][WebGPU] Add relu op (shared unary handler; sigmoid adopts make_compute_pipeline)#20863
[ExecuTorch][WebGPU] Add relu op (shared unary handler; sigmoid adopts make_compute_pipeline)#20863JCNTH wants to merge 1 commit into
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/20863
Note: Links to docs will display an error until the docs builds have been completed. ✅ You can merge normally! (1 Unrelated Failure)As of commit a219bb5 with merge base aceeb40 ( 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
|
Stack from ghstack (oldest at bottom):
Adds
reluto the WebGPU backend via a shared elementwise-unary handler. ReLU is on the SAM2/SAM3 mask-decoder MLP path, so it is needed to delegate those decoders.Problem — The backend had no
aten.relu.defaultkernel, andsigmoid(the only prior unary op) built its compute pipeline inline rather than through the shared helper — duplicating the bind-group/dispatch boilerplate that a second unary op would repeat.Solution
sigmoidwas implemented with a bespoke inline pipeline; there was norelu.add_unary_ophelper (inruntime/ops/sigmoid/UnaryOp.cpp) builds the input/output/params binding and dispatch for any elementwise-unary WGSL;sigmoid_impland the newrelu_implare thin wrappers over it, sosigmoidnow goes through the sameutils::make_compute_pipelinepath asrelu.relu.wgslis a one-element-per-threadoutput[idx] = max(input[idx], 0.0).Implementation
add_unary_op(graph, in, out, wgsl_source, wg_size_x, op_name)centralizes: the fp32/4-byte-alignment and same-size guards,utils::clamp_workgroup_size+utils::compute_1d_workgroup_countfor the 1D dispatch, thewg_sizeoverride constant, the uniform (num_elements) viautils::make_uniform, and the three-binding pipeline viautils::make_compute_pipeline.graph.add_tensor_resize_hookrecomputesnum_elements, rewrites the uniform viawgpuQueueWriteBuffer, and updates the dispatch's workgroup count for the live shape; the graph owns the uniform buffer so the hook can rewrite it.aten.sigmoid.default -> sigmoid_implandaten.relu.default -> relu_impl.backends/vulkan/runtime/graph/ops/impl/UnaryOp.cpp(add_unary_op_node); Vulkan expressesreluasclamp(0, inf), whereas this kernel uses a directmax(x, 0.0).Constraints — fp32 only (both operands 4-byte aligned); input and output must have identical byte size (same-shape elementwise); 1D dispatch only (throws past the 65535 workgroup cap).
Co-authored-with: Claude Code.
Differential Revision: D110836664