Turn on compile with NVSHMEM

kwen2501 · kwen2501 · commit bc080d49d36d · 2025-05-28T16:21:24.000-07:00
ghstack-source-id: c745e1a Pull Request resolved: #154538
diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -265,6 +265,8 @@ cmake_dependent_option(USE_RCCL "Use RCCL" ON USE_NCCL OFF)
 cmake_dependent_option(USE_STATIC_NCCL "Use static NCCL" OFF "USE_NCCL" OFF)
 cmake_dependent_option(USE_SYSTEM_NCCL "Use system-wide NCCL" OFF "USE_NCCL"
                        OFF)
+cmake_dependent_option(USE_NVSHMEM "Use NVSHMEM" ON
+                       "USE_CUDA OR USE_ROCM;UNIX;NOT APPLE" OFF)
 option(USE_NNAPI "Use NNAPI" OFF)
 option(USE_NNPACK "Use NNPACK" ON)
 cmake_dependent_option(USE_NUMA "Use NUMA. Only available on Linux." ON "LINUX"
diff --git a/caffe2/CMakeLists.txt b/caffe2/CMakeLists.txt
@@ -985,22 +985,18 @@ elseif(USE_CUDA)
     target_compile_definitions(torch_cuda PRIVATE USE_NCCL)
   endif()
 
-  # Use env var for these for now for prototyping purposes
-  set(USE_NVSHMEM $ENV{USE_NVSHMEM} CACHE BOOL "Whether to build with NVSHMEM support")
+  # Compile with NVSHMEM
+  # Default value of `USE_NVSHMEM` is set in CMakeLists.txt under root, to ON.
   # If user has specified NVSHMEM_HOME, we use it;
   # Otherwise, NVSHMEM_HOME is auto detected in tools/setup_helpers/cmake.py
   if($ENV{NVSHMEM_HOME})
     set(NVSHMEM_HOME $ENV{NVSHMEM_HOME} CACHE PATH "Path to NVSHMEM build dir")
+  else()
+    message(WARNING "NVSHMEM_HOME not found. Building without NVSHMEM.")
   endif()
 
-  if(USE_NVSHMEM AND NOT DEFINED NVSHMEM_HOME)
-    message(WARNING "USE_NVSHMEM set to 1 but NVSHMEM_HOME not found. Please run `pip install nvidia-nvshmem-<version>`, or set NVSHMEM_HOME to the NVSHMEM build dir")
-    # Disable nvshmem if NVSHMEM_HOME is not found
-    set(USE_NVSHMEM FALSE CACHE BOOL "Whether to build with NVSHMEM support")
-  endif()
-
-  if(USE_NVSHMEM)
-    message("Building with NVSHMEM support:  '${NVSHMEM_HOME}'")
+  if(USE_NVSHMEM AND DEFINED NVSHMEM_HOME)
+    message(STATUS "Building with NVSHMEM support:  '${NVSHMEM_HOME}'")
     set(NVSHMEM_INCLUDE_DIR "${NVSHMEM_HOME}/include")
     set(NVSHMEM_LIB_DIR "${NVSHMEM_HOME}/lib")
 
diff --git a/cmake/Summary.cmake b/cmake/Summary.cmake
@@ -172,6 +172,7 @@ function(caffe2_print_configuration_summary)
   if(${USE_NCCL})
     message(STATUS "    USE_SYSTEM_NCCL     : ${USE_SYSTEM_NCCL}")
   endif()
+  message(STATUS "  USE_NVSHMEM           : ${USE_NVSHMEM}")
   message(STATUS "  USE_NNPACK            : ${USE_NNPACK}")
   message(STATUS "  USE_NUMPY             : ${USE_NUMPY}")
   message(STATUS "  USE_OBSERVERS         : ${USE_OBSERVERS}")
