Adjust launch_bounds annotation for AMD hardware. (#17555)

Summary:
The max pooling backwards kernel is currently annotated with launch bounds (256,8).

Adjust the number of waves to 4 (4 times 64 is 256) for ROCm. This improves training performance for torchvision models by up to 15% (AlexNet) on a gfx906 GPU.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/17555

Differential Revision: D14277744

Pulled By: bddppq

fbshipit-source-id: 2a62088f7b8a87d1e350c432bf655288967c7883

diff --git a/aten/src/THCUNN/SpatialDilatedMaxPooling.cu b/aten/src/THCUNN/SpatialDilatedMaxPooling.cu
index d0213b1..e1202e0 100644 (file)
--- a/aten/src/THCUNN/SpatialDilatedMaxPooling.cu
+++ b/aten/src/THCUNN/SpatialDilatedMaxPooling.cu
@@ -48,7 +48,11 @@ __global__ void MaxPoolForward(const int nthreads, const Dtype* bottom_data,
 const int BACKWARD_THREADS = 256;
 
 template <typename Dtype, typename AccType>
+#if defined (__HIP_PLATFORM_HCC__)
+C10_LAUNCH_BOUNDS(BACKWARD_THREADS, 4)
+#else
 C10_LAUNCH_BOUNDS(BACKWARD_THREADS, 8)
+#endif
 __global__ void MaxPoolBackward(const int nthreads, const Dtype* top_diff,
     const int64_t* top_mask, const int num, const int channels,
     const int height, const int width, const int pooled_height,