[ExecuTorch][WebGPU] Add comparison ops (aten.eq/lt/le/gt/ge.Tensor -> bool)#21219
[ExecuTorch][WebGPU] Add comparison ops (aten.eq/lt/le/gt/ge.Tensor -> bool)#21219JCNTH wants to merge 1 commit into
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/21219
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 88d1e70 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 elementwise comparison ops (
aten.eq/lt/le/gt/ge.Tensor). They delegate throughVulkanPartitioner(all 5 tagged, verified) and output BOOL tensors, but the backend had no handler and no bool-output path.Solution: Port all 5 comparisons via one shared handler + WGSL kernel (op-code switch), mirroring Vulkan
impl/BinaryOp.cpp+glsl/binary_op_buffer.yaml. Output is a 1-byte bool tensor (vk_datatype_size(BOOL)=1,is_int=true,is_int8=false); the kernel packs 4 bool bytes peru32word (reusing the int8 buffer-binding idiom), andcopy_outputsraw-copies it (dst==map, no widen). Per element:out = (a OP b) ? 1 : 0with OP in {==, <, <=, >, >=}.Implementation:
compare/Compare.cppregistersaten.eq/lt/le/gt/ge.Tensor->compare_impl(graph, args, op)(op 0=eq..4=ge).compare.wgsl: one thread per output word (4 bool bytes), 2D-folded dispatch, op switch.CompareParams(16 bytes: num_elements, op) matches the WGSL Params. Guards: fp32 inputs, 1-byte bool output, same-shape (in/out numel equal),numel%4==0(bool packs 4/word AND gates the readback map) — all fail-loud. Resize hook recomputes numel/dispatch + re-applies the %4 guard + cross-checks both operands.Constraints / divergences from the Vulkan reference: (1)
numel%4==0(bool-output packing; fail-loud otherwise, mirroring the q8ta N%4 output-pack constraint). (2) same-shape only (flat kernel; broadcast is export-smoke, likeadd/minimum). (3)eqis torch-EXACTa==b— a DELIBERATE deviation from Vulkan's floatabs(X-Y)<1e-5: the delegate serves torch-exact model semantics and the op-test golden is torch eager;lt/le/gt/gemirror Vulkan's exact comparators.Differential Revision: D112257588