Repository: pytorch
Issue: [CUDA] Add channels_last_3d support for commonly used modules
Labels: module: cuda,triaged,module: memory format

## [CUDA] Add channels_last_3d support for commonly used modules

The goal is to add `channels_last_3d`, aka NDHWC, support on CUDA devices, and improve performance on 3D model training and inference, e.g. U-Net 3D.

This task is similar to the `channels_last` (2d) support in https://github.com/pytorch/pytorch/issues/28619, but would be much less complicated. This is because most of the C++ infrastructures already have some kind of support with `at::MemoryFormat::ChannelsLast3d`. All we need to do at this moment, is

- Enable `channels_last_3d` dispatch to other library functions, e.g. cuBLAS, cuDNN.
- Write `channels_last_3d` native CUDA kernels.
- Check dispatch logic, e.g. change `input.contiguous()` to `input.contiguous(input.suggest_memory_format())`. Also need to make sure this doesn't break existing `channels_last` (2d) logic.

I've done some preliminary tests to check the existing `channels_last_3d` support in pytorch, 

<details><summary>script </summary>

<p>

```python
import torch

nets = [
    torch.nn.Conv3d(8, 4, 3),
    torch.nn.ConvTranspose3d(8, 4, 3),
    torch.nn.MaxPool3d(3),
    torch.nn.AvgPool3d(3),
    torch.nn.AdaptiveMaxPool3d(8),
    torch.nn.AdaptiveAvgPool3d(8),
    torch.nn.AdaptiveMaxPool3d(1),
    torch.nn.AdaptiveAvgPool3d(1),
    torch.nn.BatchNorm3d(8),
    torch.nn.InstanceNorm3d(8),
]

mf = torch.channels_last_3d

for net in nets:
    print('\n', net)

    a = torch.randn(4, 8, 16, 16, 16, dtype=torch.half, device='cuda').to(memory_format=mf)
    a.requires_grad_()
    net = net.to(device='cuda', memory_format=mf)

    with torch.cuda.amp.autocast():
        out = net(a)
        if not out.is_contiguous(memory_format=mf):
            print('**** output is in the wrong format ****')

        loss = out.sum()
    
    loss.backward()

    if not a.grad.is_contiguous(memory_format=mf):
        print('**** data grad is in the wrong format ****')
```

</p>
</details>

and here are our current development plans for channels_last_3d:

- [X] Enable cuDNN Convolution in `Conv3d`, `ConvTranspose3d` (#48430)
- [x] Write native kernels for `AvgPool3d`
- [x] Write native kernels for `MaxPool3d`
- [x] Write native kernels for `AdaptiveAvgPool3d`
- [x] Write native kernels for `AdaptiveMaxPool3d`
- [ ] Enable fast path for `AdaptiveAvgPool3d` when output size is 1 (similar to #43986, #53607)
- [ ] Enable fast path for `AdaptiveMaxPool3d` when output size is 1
- [x] Enable cuDNN BatchNorm3d with `CUDNN_BATCHNORM_SPATIAL_PERSISTENT` (ongoing #59129)
- [ ] Enable channels_last_3d dispatch for `InstanceNorm3d` (relevant [code](https://github.com/pytorch/pytorch/blob/ac67cda272b0856c1be3ebfebc324e0b3bab286f/aten/src/ATen/native/Normalization.cpp#L555-L557))

Further work:
- [ ] Write native kernels for `padding`
- [ ] Write native kernels for `interpolate`

Basically, a PR that enables the channels_last_3d feature needs to have 
- Unit test
- Performance and numerical accuracy check for a few shapes
- Numerical accuracy check for a large number of shapes, if needed
- Convergence check for DL model training, if needed (I personally think this only applies to the persistent cuDNN BatchNorm3d)

Contributions are welcome! Please leave a comment if you would like to see other channels_last_3d features that I've missed from the above list.

cc @VitalyFedyunin @jamesr66a @ppwwyyxx @ptrblck @ngimel