[ExecuTorch][WebGPU] Add grid_priors op (et_vk.grid_priors)#21215
[ExecuTorch][WebGPU] Add grid_priors op (et_vk.grid_priors)#21215JCNTH wants to merge 2 commits into
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/21215
Note: Links to docs will display an error until the docs builds have been completed. ❌ 10 New Failures, 28 Pending, 36 Unrelated FailuresAs of commit 99b9c28 with merge base 4a26c64 ( NEW FAILURES - The following jobs have failed:
FLAKY - The following jobs failed but were likely due to flakiness present on trunk:
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.grid_priors— a detection anchor-grid op (generates per-cell coordinate shifts). It delegates throughVulkanPartitioner(op_registry.pytagset_vk.grid_priors.default; reachable via a direct custom-op call) but had no WebGPU handler.Solution: Port
et_vk.grid_priors(fp32), mirroring Vulkanimpl/GridPriors.cpp+glsl/grid_priors.glsl. The output is[H*W, 2]computed purely from the input's H/W plus scalarstride/offset— the input's VALUES are not read. Per output rowr:out[r,0] = (r%W + offset) * stride(x-shift),out[r,1] = (r/W + offset) * stride(y-shift), matching the Vulkan glsl (pos.x==0 -> pos.y%width,pos.x==1 -> pos.y/width) and the CPU eager (custom_ops_lib.pygrid_priors:stack((shift_x_per_col, shift_y_per_row))).Implementation:
grid_priors/GridPriors.cppregisterset_vk.grid_priors.default->grid_priors_impl. Args[in, stride(int), offset(float), out]; onlyin.dims[-2:](H, W) are used.grid_priors.wgsl: one thread per output element,row=idx/2, even idx ->row%W, odd ->row/W,(coord+offset)*stride; 2D-folded dispatch.GridPriorsParams(16 bytes: numel/width/stride/offset) matches the WGSL Params. Only the output (storage) + params (uniform) are bound — the input's data is intentionally unread. Resize hook recomputes numel/width + dispatch + out dims[H*W, 2]. Guards: in >= 2D, W>0, numel<=u32, fp32 out, all fail-loud.Constraints / divergences from the Vulkan reference: buffer re-derivation of Vulkan's texture (
imageStore/VEC4_T) kernel — the flat-index buffer form is forced by the backend's buffer-only design.stride(int) is promoted to f32 for the multiply (exact for detection-scale strides; matches the glsl's int->float promotion).@exported-using-ghexport
Differential Revision: D112257635
Differential Revision: D112257635