SIGN IN SIGN UP

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