Turn nccl_stub into a normal target (#3037)

CMakeLists.txt

@@ -155,6 +155,7 @@ target_link_options(mlx PUBLIC ${SANITIZER_LINK_FLAGS})
 
 if(MLX_BUILD_CUDA)
   enable_language(CUDA)
+  find_package(CUDAToolkit REQUIRED)
 endif()
 
 if(MLX_BUILD_METAL)

mlx/backend/cuda/CMakeLists.txt

@@ -202,7 +202,6 @@ FetchContent_MakeAvailable(nvtx3)
 target_link_libraries(mlx PUBLIC $<BUILD_INTERFACE:nvtx3-cpp>)
 
 # Make cuda runtime APIs available in non-cuda files.
-find_package(CUDAToolkit REQUIRED)
 target_include_directories(mlx PRIVATE ${CUDAToolkit_INCLUDE_DIRS})
 
 # Use cublasLt.

mlx/distributed/nccl/CMakeLists.txt

@@ -9,17 +9,17 @@ if(MLX_BUILD_CUDA)
       STATUS
         "NCCL not found, using stubs. To run distributed with NCCL backend, install NCCL."
     )
-
-    include(ExternalProject)
-    ExternalProject_Add(
+    file(
+      DOWNLOAD
+      "https://raw.githubusercontent.com/NVIDIA/nccl/refs/tags/v2.27.5-1/src/nccl.h.in"
+      "${CMAKE_CURRENT_BINARY_DIR}/nccl.h")
+    add_library(nccl_stub OBJECT
+                ${CMAKE_CURRENT_SOURCE_DIR}/nccl_stub/nccl_stubs.cpp)
+    target_include_directories(
       nccl_stub
-      SOURCE_DIR "${CMAKE_CURRENT_SOURCE_DIR}/nccl_stub"
-      BUILD_COMMAND ${CMAKE_COMMAND} --build .
-      INSTALL_COMMAND "")
-    set(NCCL_PATH
-        "${CMAKE_CURRENT_BINARY_DIR}/nccl_stub-prefix/src/nccl_stub-build/")
-    target_link_libraries(mlx PRIVATE ${NCCL_PATH}/libnccl.so)
-    target_include_directories(mlx PRIVATE ${NCCL_PATH})
+      PRIVATE ${CUDAToolkit_INCLUDE_DIRS}
+      PUBLIC ${CMAKE_CURRENT_BINARY_DIR})
+    target_link_libraries(mlx PRIVATE $<BUILD_INTERFACE:nccl_stub>)
   endif()
 else()
   target_sources(mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/no_nccl.cpp)

mlx/distributed/nccl/nccl_stub/CMakeLists.txt

@@ -1,14 +1 @@
-cmake_minimum_required(VERSION 3.25)
 
-project(nccl LANGUAGES C CXX)
-
-file(
-  DOWNLOAD
-  "https://raw.githubusercontent.com/NVIDIA/nccl/refs/tags/v2.27.5-1/src/nccl.h.in"
-  "${CMAKE_CURRENT_BINARY_DIR}/nccl.h")
-
-add_library(nccl SHARED nccl_stubs.cpp)
-set_target_properties(nccl PROPERTIES SOVERSION 2)
-find_package(CUDAToolkit REQUIRED)
-target_include_directories(nccl PRIVATE ${CUDAToolkit_INCLUDE_DIRS})
-target_include_directories(nccl PRIVATE ${CMAKE_CURRENT_BINARY_DIR})