SIGN IN SIGN UP

aten.max_pool2d_with_indices (#2547)

Summary:
Pull Request resolved: https://github.com/pytorch/executorch/pull/2547

## The Operator
An `nn.Module` invocation of `torch.nn.MaxPool2d()` is represented as `aten.max_pool2d_with_indices.default` in the Edge Dialect, indpendent of `use_indices = True/False`.
```
# Return: (Tensor output, Tensor indices)
- func: max_pool2d_with_indices(Tensor self, int[2] kernel_size, int[2] stride=[], int[2] padding=0, int[2] dilation=1, bool ceil_mode=False) -> (Tensor, Tensor)
```

This is different from PT-VK where `torch.nn.MaxPool2d()` was represented as `aten.max_pool2d.default`.
```
- func: max_pool2d(Tensor self, int[2] kernel_size, int[2] stride=[], int[2] padding=0, int[2] dilation=1, bool ceil_mode=False) -> Tensor
```

The difference is we now return an additional tensor for the max indices. Still, much of the core logic is taken from [`max_pool2d.glsl`](https://github.com/pytorch/pytorch/blob/cceabe873f11c6611f627a3bb0055994952ec6b8/aten/src/ATen/native/vulkan/glsl/max_pool2d.glsl) and [`Pool.cpp`](https://github.com/pytorch/pytorch/blob/cceabe873f11c6611f627a3bb0055994952ec6b8/aten/src/ATen/native/vulkan/ops/Pool.cpp).

We provide only a `CHANNELS_PACKED` implementation.

## The Smoke Test
Given any input and kernel sizes, we fill the input tensor with increasing values, e.g.,
```
tensor([[[10., 11., 12., 13., 14., 15.],
         [16., 17., 18., 19., 20., 21.],
         [22., 23., 24., 25., 26., 27.],
         [28., 29., 30., 31., 32., 33.]]])
```
With this setup, the max number for each pool is always in the lower-right.

We use the kernel size to compute the size of the lower-right block and verify that
1. the output tensor values match the lower-right block values, and
2. the index tensor values match the lower-right block indices.
```
tensor([[[ 18., 19., 20., 21.],
         [ 24., 25., 26., 27.],
         [ 30., 31., 32., 33.]]])
tensor([[[ 8,  9, 10, 11],
         [14, 15, 16, 17],
         [20, 21, 22, 23]]])
```
ghstack-source-id: 219506264
exported-using-ghexport
bypass-github-export-checks

Reviewed By: SS-JIA

Differential Revision: D54961929

fbshipit-source-id: 277a629e973dd72d49fc059ace2d7a9a9388ab36
J
Jorge Pineda committed
a8a49a9880dde021d64b89fafb90748f5c1b999b
Parent: 8cf4d1f
Committed by Facebook GitHub Bot <facebook-github-bot@users.noreply.github.com> on 3/21/2024, 5:57:32 AM