Merge pull request #91 from pfxuan/fix-cuda-arch

Offload TORCH_CUDA_ARCH_LIST
pierotofy · May 7, 2024 · e861818 · e861818
2 parents bc44867 + 645cfce
commit e861818
Show file tree

Hide file tree

Showing 6 changed files with 9 additions and 14 deletions.
diff --git a/.github/workflows/cuda/Linux-env.sh b/.github/workflows/cuda/Linux-env.sh
@@ -7,43 +7,36 @@ case ${1} in
     export CUDA_HOME=/usr/local/cuda-12.1
     export LD_LIBRARY_PATH=${CUDA_HOME}/lib64:${LD_LIBRARY_PATH}
     export PATH=${CUDA_HOME}/bin:${PATH}
-    export TORCH_CUDA_ARCH_LIST="7.0;7.5" # 3.5;5.0+PTX;6.0;8.0;8.6
     ;;
   cu118)
     export CUDA_HOME=/usr/local/cuda-11.8
     export LD_LIBRARY_PATH=${CUDA_HOME}/lib64:${LD_LIBRARY_PATH}
     export PATH=${CUDA_HOME}/bin:${PATH}
-    export TORCH_CUDA_ARCH_LIST="7.0;7.5" # 3.5;5.0+PTX;6.0;8.0;8.6
     ;;
   cu117)
     export CUDA_HOME=/usr/local/cuda-11.7
     export LD_LIBRARY_PATH=${CUDA_HOME}/lib64:${LD_LIBRARY_PATH}
     export PATH=${CUDA_HOME}/bin:${PATH}
-    export TORCH_CUDA_ARCH_LIST="7.0;7.5" # 3.5;5.0+PTX;6.0;8.0;8.6
     ;;
   cu116)
     export CUDA_HOME=/usr/local/cuda-11.6
     export LD_LIBRARY_PATH=${CUDA_HOME}/lib64:${LD_LIBRARY_PATH}
     export PATH=${CUDA_HOME}/bin:${PATH}
-    export TORCH_CUDA_ARCH_LIST="7.0;7.5" # 3.5;5.0+PTX;6.0;8.0;8.6
     ;;
   cu115)
     export CUDA_HOME=/usr/local/cuda-11.5
     export LD_LIBRARY_PATH=${CUDA_HOME}/lib64:${LD_LIBRARY_PATH}
     export PATH=${CUDA_HOME}/bin:${PATH}
-    export TORCH_CUDA_ARCH_LIST="7.0;7.5" # 3.5;5.0+PTX;6.0;8.0;8.6
     ;;
   cu113)
     export CUDA_HOME=/usr/local/cuda-11.3
     export LD_LIBRARY_PATH=${CUDA_HOME}/lib64:${LD_LIBRARY_PATH}
     export PATH=${CUDA_HOME}/bin:${PATH}
-    export TORCH_CUDA_ARCH_LIST="7.0;7.5" # 3.5;5.0+PTX;6.0;8.0;8.6
     ;;
   cu102)
     export CUDA_HOME=/usr/local/cuda-10.2
     export LD_LIBRARY_PATH=${CUDA_HOME}/lib64:${LD_LIBRARY_PATH}
     export PATH=${CUDA_HOME}/bin:${PATH}
-    export TORCH_CUDA_ARCH_LIST="7.0;7.5" # 3.5;5.0+PTX;6.0;
     ;;
   *)
     ;;

diff --git a/.github/workflows/ubuntu.yml b/.github/workflows/ubuntu.yml
@@ -27,7 +27,6 @@ jobs:
       CL: /MP
       CMAKE_GENERATOR: Ninja
       CMAKE_GENERATOR_PLATFORM: ${{matrix.arch}}
-      TORCH_CUDA_ARCH_LIST: '7.0;7.5'
     steps:
       - name: Checkout Repository
         uses: actions/checkout@v4

diff --git a/.github/workflows/windows.yml b/.github/workflows/windows.yml
@@ -33,7 +33,6 @@ jobs:
       CL: /MP
       CMAKE_GENERATOR: ${{matrix.generator}}
       CMAKE_GENERATOR_PLATFORM: ${{matrix.arch}}
-      TORCH_CUDA_ARCH_LIST: '7.0;7.5'
     steps:
       - name: Checkout Repository
         uses: actions/checkout@v4
@@ -136,7 +135,6 @@ jobs:
             -DCMAKE_CXX_COMPILER_LAUNCHER:FILEPATH=ccache `
             -DCMAKE_PREFIX_PATH=${{ github.workspace }}\libtorch `
             -DOPENCV_DIR=${{ github.workspace }}\OpenCV\opencv\build `
-            -DTORCH_CUDA_ARCH_LIST="${{ env.TORCH_CUDA_ARCH_LIST }}" `
             -DOPENSPLAT_BUILD_SIMPLE_TRAINER=ON
           cmake --build . --parallel --config ${{ matrix.cmake-build-type }}
 

diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -77,7 +77,7 @@ if(GPU_RUNTIME STREQUAL "CUDA")
         endif()
         if(NOT OPENSPLAT_MAX_CUDA_COMPATIBILITY)
             if(NOT CMAKE_CUDA_ARCHITECTURES)
-                SET(CMAKE_CUDA_ARCHITECTURES 70 75)
+                SET(CMAKE_CUDA_ARCHITECTURES 70;75;80)
             endif()
         else()
             # Build for maximum compatibility
@@ -107,6 +107,11 @@ if(GPU_RUNTIME STREQUAL "CUDA")
             endforeach()
             message("Set CUDA flags: " ${CMAKE_CUDA_FLAGS})
         endif()
+        # Set torch cuda architecture list
+        set(TORCH_CUDA_ARCH_LIST ${CMAKE_CUDA_ARCHITECTURES})
+        list(TRANSFORM TORCH_CUDA_ARCH_LIST REPLACE "([0-9])([0-9])" "\\1.\\2")
+        string(REPLACE ";" " " TORCH_CUDA_ARCH_LIST "${TORCH_CUDA_ARCH_LIST}")
+        message(STATUS "** Updated TORCH_CUDA_ARCH_LIST to ${TORCH_CUDA_ARCH_LIST} **")
     endif()
 elseif(GPU_RUNTIME STREQUAL "HIP")
     set(USE_HIP ON CACHE BOOL "Use HIP for GPU acceleration")

diff --git a/Dockerfile b/Dockerfile
@@ -5,7 +5,7 @@ FROM ubuntu:${UBUNTU_VERSION}
 ARG UBUNTU_VERSION
 ARG TORCH_VERSION=2.2.1
 ARG CUDA_VERSION=12.1.1
-ARG TORCH_CUDA_ARCH_LIST=7.0;7.5
+ARG CMAKE_CUDA_ARCHITECTURES=70;75;80
 ARG CMAKE_BUILD_TYPE=Release
 
 SHELL ["/bin/bash", "-c"]
@@ -63,6 +63,6 @@ RUN source .github/workflows/cuda/Linux-env.sh cu"${CUDA_VERSION%%.*}"$(echo $CU
     -DCMAKE_BUILD_TYPE=${CMAKE_BUILD_TYPE} \
     -DCMAKE_PREFIX_PATH=/code/libtorch \
     -DCMAKE_INSTALL_PREFIX=/code/install \
-    -DTORCH_CUDA_ARCH_LIST="${TORCH_CUDA_ARCH_LIST}" \
+    -DCMAKE_CUDA_ARCHITECTURES="${CMAKE_CUDA_ARCHITECTURES}" \
     -DCUDA_TOOLKIT_ROOT_DIR=${CUDA_HOME} && \
     ninja
diff --git a/README.md b/README.md
@@ -159,7 +159,7 @@ docker build \
   --build-arg UBUNTU_VERSION=22.04 \
   --build-arg CUDA_VERSION=12.1.1 \
   --build-arg TORCH_VERSION=2.2.1 \
-  --build-arg TORCH_CUDA_ARCH_LIST="7.0;7.5" \
+  --build-arg CMAKE_CUDA_ARCHITECTURES="70;75;80" \
   --build-arg CMAKE_BUILD_TYPE=Release .
 ```