SIGN IN SIGN UP

aten.avg_pool2d (#3770)

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

## The Operator
`nn.Module` invocations of [`torch.nn.AvgPool2d`](https://pytorch.org/docs/stable/generated/torch.nn.AvgPool2d.html) get compiled to `aten.avg_pool2d.default` in the Edge Dialect, which carries the following signature.
```
- func: avg_pool2d(Tensor self, int[2] kernel_size, int[2] stride=[], int[2] padding=0, bool ceil_mode=False, bool count_include_pad=True, int? divisor_override=None) -> Tensor
```

## Implementation
This is a full C-packing implementation including dynamic shape support. We start with [LiteInterpreter's `avg_pool2d.glsl` logic](https://github.com/pytorch/pytorch/blob/9257a0698b57acc5607ee6fe31a16fdd93af1731/aten/src/ATen/native/vulkan/glsl/avg_pool2d.glsl), which is incomplete, and cover `ceil_mode=True`,  `count_include_pad=True`, and `divisor_override` cases for full support. As a result, the divisor's computation is now a bit complex. If needed, we can simplify it into separate shaders in the future.
ghstack-source-id: 228476264

Reviewed By: copyrightly

Differential Revision: D57918523

fbshipit-source-id: 8069c4a2dcc5d46da7221d58661e57bf2055b521
J
Jorge Pineda committed
a463f0be2462d64b399e33229bacfcd723e593d9
Parent: 8c8d965
Committed by Facebook GitHub Bot <facebook-github-bot@users.noreply.github.com> on 5/31/2024, 8:46:14 PM