Fix the missing libcuda linking (pytorch#2792)

Summary: - Fix the missing `libcuda` linking Pull Request resolved: pytorch#2792 Reviewed By: spcyppt Differential Revision: D59142451 Pulled By: q10 fbshipit-source-id: 844be1113b1db925c8eb7bc5c7b850d898b70d0f
q10 · Jun 28, 2024 · f8fe8f4 · f8fe8f4
1 parent 565ab30
commit f8fe8f4
Show file tree

Hide file tree

Showing 5 changed files with 19 additions and 22 deletions.
diff --git a/.github/workflows/fbgemm_gpu_ci_cuda.yml b/.github/workflows/fbgemm_gpu_ci_cuda.yml
@@ -101,8 +101,7 @@ jobs:
 
     # Install via PIP to avoid defaulting to the CPU variant if the GPU variant of the day is not ready
     - name: Install PyTorch Nightly
-      # TODO: Change back to rolling nightly torch once we figure out the cause for undefined symbol: cuTensorMapEncodeTiled
-      run: . $PRELUDE; install_pytorch_pip $BUILD_ENV nightly/2.5.0.dev20240621 cuda/${{ matrix.cuda-version }}
+      run: . $PRELUDE; install_pytorch_pip $BUILD_ENV nightly cuda/${{ matrix.cuda-version }}
 
     - name: Collect PyTorch Environment Info
       if: ${{ success() || failure() }}
@@ -192,8 +191,7 @@ jobs:
 
     # Install via PIP to avoid defaulting to the CPU variant if the GPU variant of the day is not ready
     - name: Install PyTorch Nightly
-      # TODO: Change back to rolling nightly torch once we figure out the cause for undefined symbol: cuTensorMapEncodeTiled
-      run: . $PRELUDE; install_pytorch_pip $BUILD_ENV nightly/2.5.0.dev20240621 cuda/${{ matrix.cuda-version }}
+      run: . $PRELUDE; install_pytorch_pip $BUILD_ENV nightly cuda/${{ matrix.cuda-version }}
 
     - name: Collect PyTorch Environment Info
       if: ${{ success() || failure() }}

diff --git a/cmake/modules/CudaSetup.cmake b/cmake/modules/CudaSetup.cmake
@@ -33,3 +33,14 @@ endif()
 if(NVML_LIB_PATH)
   message(STATUS "Found NVML_LIB_PATH: ${NVML_LIB_PATH}")
 endif()
+
+# The libcuda.so path was previously set by PyTorch CMake, but the setup has
+# been removed from the PyTorch codebase, see:
+# https://github.com/pytorch/pytorch/pull/128801
+set(CUDA_DRIVER_LIBRARIES "${CUDA_cuda_driver_LIBRARY}" CACHE FILEPATH "")
+
+BLOCK_PRINT(
+  "CUDA Driver Path"
+  ""
+  "CUDA_DRIVER_LIBRARIES=${CUDA_DRIVER_LIBRARIES}"
+)
diff --git a/fbgemm_gpu/FbgemmGpu.cmake b/fbgemm_gpu/FbgemmGpu.cmake
@@ -683,12 +683,8 @@ set_target_properties(fbgemm_gpu_py PROPERTIES PREFIX "")
 # Link to PyTorch
 target_link_libraries(fbgemm_gpu_py
   ${TORCH_LIBRARIES}
-  ${NCCL_LIBRARIES})
-
-# Append -lcuda flag, since TORCH_LIBRARIES doesn't include this
-if(NOT FBGEMM_CPU_ONLY AND NOT USE_ROCM)
-  set_target_properties(fbgemm_gpu_py PROPERTIES LINK_FLAGS "-lcuda")
-endif()
+  ${NCCL_LIBRARIES}
+  ${CUDA_DRIVER_LIBRARIES})
 
 # Link to NVML
 if(NVML_LIB_PATH)

diff --git a/fbgemm_gpu/experimental/example/CMakeLists.txt b/fbgemm_gpu/experimental/example/CMakeLists.txt
@@ -53,12 +53,8 @@ target_include_directories(fbgemm_gpu_experimental_example_py PRIVATE
 
 target_link_libraries(fbgemm_gpu_experimental_example_py
     ${TORCH_LIBRARIES}
-    ${NCCL_LIBRARIES})
-
-if(NOT FBGEMM_CPU_ONLY AND NOT USE_ROCM)
-  set_target_properties(fbgemm_gpu_experimental_example_py PROPERTIES
-    LINK_FLAGS "-lcuda")
-endif()
+    ${NCCL_LIBRARIES}
+    ${CUDA_DRIVER_LIBRARIES})
 
 # Remove `lib` from the output artifact name
 set_target_properties(fbgemm_gpu_experimental_example_py PROPERTIES PREFIX "")

diff --git a/fbgemm_gpu/experimental/gen_ai/CMakeLists.txt b/fbgemm_gpu/experimental/gen_ai/CMakeLists.txt
@@ -104,14 +104,10 @@ target_include_directories(fbgemm_gpu_experimental_gen_ai_py PRIVATE
   ${TORCH_INCLUDE_DIRS}
   ${NCCL_INCLUDE_DIRS})
 
-if(NOT FBGEMM_CPU_ONLY AND NOT USE_ROCM)
-  set_target_properties(fbgemm_gpu_experimental_gen_ai_py PROPERTIES
-    LINK_FLAGS "-lcuda")
-endif()
-
 target_link_libraries(fbgemm_gpu_experimental_gen_ai_py
   ${TORCH_LIBRARIES}
-  ${NCCL_LIBRARIES})
+  ${NCCL_LIBRARIES}
+  ${CUDA_DRIVER_LIBRARIES})
 
 # Remove `lib` from the output artifact name
 set_target_properties(fbgemm_gpu_experimental_gen_ai_py PROPERTIES PREFIX "")