SIGN IN SIGN UP

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