From 16aad4dfb5cbe1ef643f78e919a15f45f107723e Mon Sep 17 00:00:00 2001
From: Min Si <msi@fb.com>
Date: Wed, 13 Apr 2022 22:01:55 -0700
Subject: [PATCH 1/2] fix ROCM_PATH check in cmake

Summary:
The original code cannot correctly check whether the user has
specified env var ROCM_PATH, thus always used the default /opt/rocm
incorrectly.

Differential Revision: D35636873

fbshipit-source-id: b6dc031f335bcbf3f1d32ee6f8b84d6a7ebb78b3
---
 cmake/Modules/Findrccl.cmake | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/cmake/Modules/Findrccl.cmake b/cmake/Modules/Findrccl.cmake
index 008d23a7a..95aa0587f 100644
--- a/cmake/Modules/Findrccl.cmake
+++ b/cmake/Modules/Findrccl.cmake
@@ -13,8 +13,7 @@
 # The path hints include CUDA_TOOLKIT_ROOT_DIR seeing as some folks
 # install RCCL in the same location as the CUDA toolkit.
 # See https://github.com/caffe2/caffe2/issues/1601
-
-if($ENV{ROCM_PATH})
+if(DEFINED ENV{ROCM_PATH})
     set(RCCL_ROOT_DIR $ENV{ROCM_PATH} CACHE PATH "Folder contains AMD RCCL")
 else()
     set(RCCL_ROOT_DIR "/opt/rocm")
@@ -26,7 +25,7 @@ find_path(RCCL_INCLUDE_DIR
   ${RCCL_INCLUDE_DIR}
   ${RCCL_ROOT_DIR}/include)
 
-if ($ENV{USE_STATIC_RCCL})
+if (DEFINED ENV{USE_STATIC_RCCL})
   message(STATUS "USE_STATIC_RCCL detected. Linking against static RCCL library")
   set(RCCL_LIBNAME "librccl_static.a")
 else()

From ac6d1b8483124e08a476902d29b46769c9f1de6d Mon Sep 17 00:00:00 2001
From: Min Si <msi@fb.com>
Date: Wed, 13 Apr 2022 22:02:10 -0700
Subject: [PATCH 2/2] build rccl module only when USE_RCCL is ON

Summary:
Current code always builds nccl module as long as USE_ROCM is ON.
However, we should build it only when USE_RCCL is ON. This patch fixes
it by adding the hipified nccl files in nccl/CMakeLists.txt which will be included
only when USE_RCCL is ON.

Differential Revision: D35636872

fbshipit-source-id: 936f61fef5808431ace35d2da916762875b65b11
---
 gloo/CMakeLists.txt      |  2 --
 gloo/nccl/CMakeLists.txt | 41 ++++++++++++++++++++++++++++------------
 2 files changed, 29 insertions(+), 14 deletions(-)

diff --git a/gloo/CMakeLists.txt b/gloo/CMakeLists.txt
index 02124e904..0b416e9cd 100644
--- a/gloo/CMakeLists.txt
+++ b/gloo/CMakeLists.txt
@@ -66,12 +66,10 @@ if(USE_ROCM)
   file(GLOB GLOO_HIP_SRCS
     "${HIPIFY_OUTPUT_DIR}/hip*.cc"
     "${HIPIFY_OUTPUT_DIR}/hip*.hip"
-    "${HIPIFY_OUTPUT_DIR}/nccl/nccl.hip"
     )
 
   file(GLOB GLOO_HIP_HDRS
     "${HIPIFY_OUTPUT_DIR}/hip*.h"
-    "${HIPIFY_OUTPUT_DIR}/nccl/nccl.h"
     )
 endif()
 
diff --git a/gloo/nccl/CMakeLists.txt b/gloo/nccl/CMakeLists.txt
index 6bceab38f..6615cbc09 100644
--- a/gloo/nccl/CMakeLists.txt
+++ b/gloo/nccl/CMakeLists.txt
@@ -1,12 +1,29 @@
-set(GLOO_NCCL_SRCS
-  "${CMAKE_CURRENT_SOURCE_DIR}/nccl.cu"
-  )
-
-set(GLOO_NCCL_HDRS
-  "${CMAKE_CURRENT_SOURCE_DIR}/nccl.h"
-  )
-
-list(APPEND GLOO_CUDA_SRCS ${GLOO_NCCL_SRCS})
-list(APPEND GLOO_CUDA_HDRS ${GLOO_NCCL_HDRS})
-set(GLOO_CUDA_SRCS ${GLOO_CUDA_SRCS} PARENT_SCOPE)
-set(GLOO_CUDA_HDRS ${GLOO_CUDA_HDRS} PARENT_SCOPE)
+if(USE_CUDA)
+    set(GLOO_NCCL_SRCS
+      "${CMAKE_CURRENT_SOURCE_DIR}/nccl.cu"
+      )
+
+    set(GLOO_NCCL_HDRS
+      "${CMAKE_CURRENT_SOURCE_DIR}/nccl.h"
+      )
+
+    list(APPEND GLOO_CUDA_SRCS ${GLOO_NCCL_SRCS})
+    list(APPEND GLOO_CUDA_HDRS ${GLOO_NCCL_HDRS})
+    set(GLOO_CUDA_SRCS ${GLOO_CUDA_SRCS} PARENT_SCOPE)
+    set(GLOO_CUDA_HDRS ${GLOO_CUDA_HDRS} PARENT_SCOPE)
+endif()
+
+if(USE_ROCM)
+    set(GLOO_NCCL_SRCS
+      "${CMAKE_CURRENT_SOURCE_DIR}/nccl.hip"
+      )
+
+    set(GLOO_NCCL_HDRS
+      "${CMAKE_CURRENT_SOURCE_DIR}/nccl.h"
+      )
+
+    list(APPEND GLOO_HIP_SRCS ${GLOO_NCCL_SRCS})
+    list(APPEND GLOO_HIP_HDRS ${GLOO_NCCL_HDRS})
+    set(GLOO_HIP_SRCS ${GLOO_HIP_SRCS} PARENT_SCOPE)
+    set(GLOO_HIP_HDRS ${GLOO_HIP_HDRS} PARENT_SCOPE)
+endif()