Fix/sapling nccl hack

CMakeLists.txt

@@ -87,6 +87,10 @@ if (FF_CUDA_ARCH STREQUAL "")
   message(FATAL_ERROR "FF_CUDA_ARCH cannot be an empty string. Set it to `autodetect`, `all`, or pass one or multiple valid CUDA archs.")
 endif()
 
+if(${CMAKE_SYSTEM_NAME} MATCHES "Linux")
+  set(LIBEXT ".so")
+endif()
+
 include(cuda)
 include(cudnn)
 include(nccl)
@@ -108,10 +112,6 @@ include(any)
 #include(gtest)
 #include(fmt)
 
-if(${CMAKE_SYSTEM_NAME} MATCHES "Linux")
-  set(LIBEXT ".so")
-endif()
-
 # TODO @lockshaw remove me
 # https://discourse.nixos.org/t/get-clangd-to-find-standard-headers-in-nix-shell/11268/6
 if(CMAKE_EXPORT_COMPILE_COMMANDS)

cmake/cudnn.cmake

@@ -40,3 +40,7 @@ message( STATUS "CUDNN include : ${CUDNN_INCLUDE_DIR}" )
 else()
   message( FATAL_ERROR "CUDNN package not found -> specify search path via CUDNN_DIR variable")
 endif()
+
+add_library(cudnn INTERFACE)
+target_include_directories(cudnn SYSTEM INTERFACE ${CUDNN_INCLUDE_DIR})
+target_link_libraries(cudnn INTERFACE ${CUDNN_LIBRARY})

cmake/nccl.cmake

@@ -1,4 +1,4 @@
-set(NCCL_NAME nccl)
+set(NCCL_NAME nccl_internal)
 # set(NCCL_CUDA_ARCH "-gencode=arch=compute_${CUDA_ARCH},code=sm_${CUDA_ARCH}")
 # message("NCCL_CUDA_ARCH: ${NCCL_CUDA_ARCH}")
 
@@ -74,18 +74,19 @@ if(NCCL_URL)
   list(APPEND FLEXFLOW_EXT_LIBRARIES ${NCCL_LIB_DIR}/libnccl${LIBEXT})
   install(DIRECTORY ${NCCL_INCLUDE_DIR}/ DESTINATION include)
   install(DIRECTORY ${NCCL_LIB_DIR}/ DESTINATION lib PATTERN "pkgconfig" EXCLUDE)
-
+
+  set(NCCL_LIB "${INSTALL_DIR}/lib/libnccl${LIBEXT}")
 else()
   # Build NCCL from source
   message(STATUS "Building NCCL from source")
   list(TRANSFORM CUDA_GENCODE PREPEND "NVCC_GENCODE=" OUTPUT_VARIABLE NCCL_BUILD_NVCC_GENCODE)
-  
+
   include(ExternalProject)
   ExternalProject_Add(${NCCL_NAME}
-   SOURCE_DIR ${PROJECT_SOURCE_DIR}/deps/${NCCL_NAME}
-   PREFIX ${CMAKE_BINARY_DIR}/deps/${NCCL_NAME}
-   INSTALL_DIR ${CMAKE_BINARY_DIR}/deps/${NCCL_NAME}
-   BUILD_BYPRODUCTS ${CMAKE_BINARY_DIR}/deps/${NCCL_NAME}/lib/libnccl${LIBEXT}
+   SOURCE_DIR ${PROJECT_SOURCE_DIR}/deps/nccl
+   PREFIX ${CMAKE_BINARY_DIR}/deps/nccl
+   INSTALL_DIR ${CMAKE_BINARY_DIR}/deps/nccl
+   BUILD_BYPRODUCTS ${CMAKE_BINARY_DIR}/deps/nccl/lib/libnccl${LIBEXT}
    INSTALL_COMMAND ""
    CONFIGURE_COMMAND ""
    BUILD_COMMAND make src.build "${NCCL_BUILD_NVCC_GENCODE}" "CUDA_HOME=${CUDA_TOOLKIT_ROOT_DIR}" "BUILDDIR=${CMAKE_BINARY_DIR}/deps/${NCCL_NAME}" "CXX=${CMAKE_CXX_COMPILER}" CC="${CMAKE_CC_COMPILER}"
@@ -98,6 +99,14 @@ else()
     ${INSTALL_DIR}/include)
   list(APPEND FLEXFLOW_EXT_LIBRARIES
     ${INSTALL_DIR}/lib/libnccl${LIBEXT})
-  set_directory_properties(PROPERTIES ADDITIONAL_CLEAN_FILES "${CMAKE_BINARY_DIR}/deps/${NCCL_NAME}/lib/")
+  set_directory_properties(PROPERTIES ADDITIONAL_CLEAN_FILES "${CMAKE_BINARY_DIR}/deps/nccl/lib/")
 
+  set(NCCL_INCLUDE_DIR "${INSTALL_DIR}/include")
+  set(NCCL_LIB "${INSTALL_DIR}/lib/libnccl${LIBEXT}")
 endif()
+message("NCCL_LIB = ${NCCL_LIB}")
+
+add_library(nccl INTERFACE)
+target_include_directories(nccl SYSTEM INTERFACE ${NCCL_INCLUDE_DIR})
+#add_dependencies(nccli nccl)
+target_link_libraries(nccl INTERFACE ${NCCL_LIB})

lib/CMakeLists.txt

@@ -1,6 +1,6 @@
 add_subdirectory(pcg)
 add_subdirectory(compiler)
-# add_subdirectory(runtime)
+add_subdirectory(runtime)
 add_subdirectory(op-attrs)
 add_subdirectory(kernels)
 add_subdirectory(utils)

lib/kernels/CMakeLists.txt

@@ -23,6 +23,9 @@ target_include_directories(
 target_link_libraries(
   ${project_target}
   op-attrs
+  cuda
+  cudnn
+  nccl
 )
 
 define_ff_vars(kernels)