Support NHWC upsampling for CUDA.

[r-zenine]

🚀 Feature

Add support for NHWC images for upsampling nearest in CUDA.

pytorch/vision#1179 (comment)

Motivation

We need this work in order to match the work from This PR and have feature parity with CPU.

In addition, this would help us further support an ongoing work done with @fmassa to have efficient native image reading ops in torchvision.

Pitch

I propose the replicate work done in This PR](#34597) and add a new code path as it was done here:

https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/cpu/UpSampleKernel.cpp#L342-L351

This approach is simple. We would also stay consistent with the approach already taken. However, it's going to be hard to avoid unnecessary code duplication. I'll do my best.

@fmassa , @VitalyFedyunin Thoughts?

cc @ngimel @VitalyFedyunin @jamesr66a

[r-zenine]

Hi,
Small Update :
I see an opportunity to reuse most of the code if we manage to make this portion of the code parametrized by the memory layout.
https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/cuda/UpSampleNearest2d.cu#L46-L49

We could add a template parameter to the upsample_nearest2d_out_frame for memory layout and get the code snippet above into it a template function. We could then add an if statement somewhere around here :
https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/cuda/UpSampleNearest2d.cu#L195-L206

Thoughts ?

[VitalyFedyunin]

Parameterized looks acceptable, but please avoid parameterizing on at::MemoryFormat value, as it might change in future. I think best would be parameterising on something like bool is_channels_last

[VitalyFedyunin]

Nit: You might find this information useful https://github.com/pytorch/pytorch/wiki/Writing-memory-format-aware-operators

[r-zenine]

Hi @VitalyFedyunin.

Thanks for the resources, I will take a look and work on the changes.
I'll get back to you when I have a first PR ready for review.

[r-zenine]

Hi @VitalyFedyunin,
Thank you for the resources. I learned things :)

I have a question: I can see that for the cpu version, the parametrization is done on at::MemoryFormat, Would adding a is_channels_last bolean parameter introduce inconsistencies in the code ?

Do we add the parameter and update both CPU and GPU code ?
What are your thoughts about this ?

Thanks

[r-zenine]

Hi,

I cannot manage to put time aside to work on the issue for the moment. My schedule got crazy.
Sorry for the inconvenience.

Thanks.