Return if unary_op is out of bounds (#2520)
Summary: Pull Request resolved: https://github.com/pytorch/executorch/pull/2520 Eliminate additional computation when we go out of bounds, which can occur when 1. global work group size is not a perfect multiple of local work group size, and/or 2. dynamic shapes are used to reshape the logical tensor data. ## Aside on sizes and extents This is the extension of me struggling on the `max_pool2d` implementation and hence looking closer at how `gpu_sizes_ubo()` and `extents_ubo()` differ. I'm writing this summary to explain this to future me when I inevitably forget. This knowledge can be obtained by studying [`Tensor.*`](https://github.com/pytorch/pytorch/blob/cceabe873f11c6611f627a3bb0055994952ec6b8/aten/src/ATen/native/vulkan/api/Tensor.cpp). If our tensor has - `cpu_sizes`: (N, C, H, W) then - `gpu_sizes`: (W, H, C, N) but the packed-dim size is aligned up to a multiple of 4. - `extents`: Size of the actual image texture, i.e., gpu_sizes but merging C,N and dividing the packed-dim size by 4. So we obtain: 1. WIDTH_PACKED => `extents`: (W / 4, H, C*N) 2. HEIGHT_PACKED => `extents`: (W, H / 4, C*N) 3. CHANNELS_PACKED => `extents`: (W, H, C*N / 4) Hence, - for texture positions, use `extents`, - for logical coordinates, use `gpu_sizes`. ghstack-source-id: 219292658 bypass-github-export-checks Reviewed By: SS-JIA Differential Revision: D55097275 fbshipit-source-id: 15419d5d7eb260c8cf8a8ecb1f57e6253af7f309
J
Jorge Pineda committed
0f0c307136074450310722260fe06b87a6a03166
Parent: 969060b
Committed by Facebook GitHub Bot <facebook-github-bot@users.noreply.github.com>
on 3/20/2024, 12:47:18 AM