Fix BinaryOp broadcasting for packed dim (#2653)
Summary:
Pull Request resolved: https://github.com/pytorch/executorch/pull/2653
As copyrightly pointed out, broadcasting was not working properly for the example below. I root caused the to confusion between `sizes()` vs `gpu_sizes()` once again! These concepts are explained in https://github.com/pytorch/executorch/pull/2520
We should use the CPU size, not the GPU size to detect when we should broadcast across the packed-dim texel's elements.
# Example
Given inputs `torch.ones(2, 3)` and `torch.ones(2, 1)` and `GPUMemoryLayout::WIDTH_PACKED`, we have CPU widths 3 and 1, respectively. These are aligned up to GPU widths 4 and 4, and hence we were failing to broadcast along the packed-dim texel's elements.
## torch.ones(2, 3)
```
(2, 3) = (H, W) = sizes
[[1 1 1]
[1 1 1]]
-> (W, H) = (3, 2) → (4, 2) = gpu_sizes
-> extents = (1, 2)
[1 1 1 0] [1 1 1 0]
```
## torch.ones(2, 1)
```
(2, 1) = (H, W) = sizes
[[1]
[1]]
-> (W, H) = (1, 2) → (4, 2) = gpu_sizes
-> extents = (1, 2)
[1 0 0 0] [1 0 0 0]
-> (broadcast from this change)
[1 1 1 1] [1 1 1 1]
```
## torch.ones(2, 3) + torch.ones(2, 1)
Ignore the final element of each texel as it's just padding we never read.
```
No broadcast:
[1 1 1 0] [1 1 1 0] + [1 0 0 0] [1 0 0 0] = [2 1 1 0] [2 1 1 0]
Broadcast:
[1 1 1 0] [1 1 1 0] + [1 1 1 1] [1 1 1 1] = [2 2 2 1] [2 2 2 1]
```
# Cleanup
Remove unneeded `check_broadcastable()` since this is caught earlier in the PyTorch compiler pipeline. For example, `torch.ones(2, 3) + torch.ones(2, 2)` triggers this error:
```
TorchRuntimeError: Failed running call_function <built-in function add>(*(FakeTensor(..., size=(2, 3)), FakeTensor(..., size=(2, 2))), **{}):
Attempting to broadcast a dimension of length 2 at -1! Mismatching argument at index 1 had torch.Size([2, 2]); but expected shape should be broadcastable to [2, 3]
```
bypass-github-export-checks
Reviewed By: SS-JIA
Differential Revision: D55278527
fbshipit-source-id: abb8a83924370b21dbbabdd5f1f4af8f502edc1f J
Jorge Pineda committed
25c5b67d77397e82b4b056e4b9adfc2c291bb84a
Parent: 04d568d
Committed by Facebook GitHub Bot <facebook-github-bot@users.noreply.github.com>
on 3/27/2024, 10:47:49 PM