From 908b451efb1b35a3ef26154bc959cae7ff29e15e Mon Sep 17 00:00:00 2001
From: Pruthvi Madugundu <mpruthvi@gmail.com>
Date: Wed, 5 Feb 2020 12:00:49 -0800
Subject: [PATCH] Enabling the nccl/rccl test for ROCM environment (#32340)

Summary:
Enabling the RCCL test on rocm by adding a temporary grace period to clean up.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/32340

Differential Revision: D19744459

Pulled By: xw285cornell

fbshipit-source-id: 1af3b64113a67f93e622d010ddd3020e5d6c8bc8
---
 .jenkins/caffe2/test.sh              | 2 +-
 caffe2/contrib/nccl/cuda_nccl_gpu.cc | 5 +++++
 2 files changed, 6 insertions(+), 1 deletion(-)

diff --git a/.jenkins/caffe2/test.sh b/.jenkins/caffe2/test.sh
index f60eaa1853a..0bc29a7ef90 100755
--- a/.jenkins/caffe2/test.sh
+++ b/.jenkins/caffe2/test.sh
@@ -85,7 +85,7 @@ fi
 EXTRA_TESTS=()
 
 # CUDA builds always include NCCL support
-if [[ "$BUILD_ENVIRONMENT" == *-cuda* ]]; then
+if [[ "$BUILD_ENVIRONMENT" == *-cuda* ]] || [[ "$BUILD_ENVIRONMENT" == *-rocm* ]]; then
   EXTRA_TESTS+=("$caffe2_pypath/contrib/nccl")
 fi
 
diff --git a/caffe2/contrib/nccl/cuda_nccl_gpu.cc b/caffe2/contrib/nccl/cuda_nccl_gpu.cc
index 2ea9a2f111a..f08e296da81 100644
--- a/caffe2/contrib/nccl/cuda_nccl_gpu.cc
+++ b/caffe2/contrib/nccl/cuda_nccl_gpu.cc
@@ -28,8 +28,13 @@ class NCCLContext {
       // get stream priorities
       int lo_pri, hi_pri;
       CUDA_ENFORCE(cudaDeviceGetStreamPriorityRange(&lo_pri, &hi_pri));
+#ifndef __HIP_PLATFORM_HCC__
       CUDA_ENFORCE(cudaStreamCreateWithPriority(
           &streams_[i], cudaStreamNonBlocking, hi_pri));
+#else
+      CUDA_ENFORCE(cudaStreamCreateWithFlags(
+          &streams_[i], cudaStreamNonBlocking));
+#endif // __HIP_PLATFORM_HCC__
       CUDA_ENFORCE(cudaEventCreateWithFlags(
           &events_[i], cudaEventDefault | cudaEventDisableTiming));
     }