update

pyg-team · Apr 8, 2024 · e4cee39 · e4cee39
1 parent c3009bf
commit e4cee39
Show file tree

Hide file tree

Showing 4 changed files with 15 additions and 31 deletions.
diff --git a/.github/workflows/cuda/Windows.sh b/.github/workflows/cuda/Windows.sh
@@ -37,40 +37,24 @@ case ${1} in
     ;;
 esac
 
-# Install NVIDIA drivers, see:
-# https://github.com/pytorch/vision/blob/master/packaging/windows/internal/cuda_install.bat#L99-L102
-# curl -k -L "https://ossci-windows.s3.us-east-1.amazonaws.com/builder/additional_dlls.zip" --output "/tmp/gpu_driver_dlls.zip"
-# 7z x "/tmp/gpu_driver_dlls.zip" -o"/c/Windows/System32"
-
 curl -k -L "${CUDA_URL}/${CUDA_FILE}" --output "${CUDA_FILE}"
 echo ""
 echo "Installing from ${CUDA_FILE}..."
 PowerShell -Command "Start-Process -FilePath \"${CUDA_FILE}\" -ArgumentList \"-s nvcc_${CUDA_SHORT} cuobjdump_${CUDA_SHORT} nvprune_${CUDA_SHORT} cupti_${CUDA_SHORT} cublas_dev_${CUDA_SHORT} cudart_${CUDA_SHORT} cufft_dev_${CUDA_SHORT} curand_dev_${CUDA_SHORT} cusolver_dev_${CUDA_SHORT} cusparse_dev_${CUDA_SHORT} thrust_${CUDA_SHORT} npp_dev_${CUDA_SHORT} nvrtc_dev_${CUDA_SHORT} nvml_dev_${CUDA_SHORT}\" -Wait -NoNewWindow"
 echo "Done!"
 rm -f "${CUDA_FILE}"
 
+# echo Installing NVIDIA drivers...
+# https://github.com/pytorch/vision/blob/master/packaging/windows/internal/cuda_install.bat#L99-L102
+# curl -k -L "https://ossci-windows.s3.us-east-1.amazonaws.com/builder/additional_dlls.zip" --output "/tmp/gpu_driver_dlls.zip"
+# 7z x "/tmp/gpu_driver_dlls.zip" -o"/c/Windows/System32"
+
 echo Installing NvToolsExt...
-curl -k -L https://ossci-windows.s3.us-east-1.amazonaws.com/builder/NvToolsExt.7z --output "/tmp/NvToolsExt.7z"
-7z x "/tmp/NvToolsExt.7z" -o"/tmp/NvToolsExt"
+curl -k -L https://ossci-windows.s3.us-east-1.amazonaws.com/builder/NvToolsExt.7z --output /tmp/NvToolsExt.7z
+7z x /tmp/NvToolsExt.7z -o"/tmp/NvToolsExt"
 mkdir -p "/c/Program Files/NVIDIA Corporation/NvToolsExt/bin/x64"
 mkdir -p "/c/Program Files/NVIDIA Corporation/NvToolsExt/include"
 mkdir -p "/c/Program Files/NVIDIA Corporation/NvToolsExt/lib/x64"
-echo "-------"
-ls "/tmp/NvToolsExt"
-echo "-------"
-ls "/tmp/NvToolsExt/bin"
-echo "-------"
-ls "/tmp/NvToolsExt/bin/x64"
-echo "-------"
-ls "/tmp/NvToolsExt/include"
-echo "-------"
-ls "/tmp/NvToolsExt/lib"
-echo "-------"
-ls "/tmp/NvToolsExt/lib/x64"
 cp -r /tmp/NvToolsExt/bin/x64/* "/c/Program Files/NVIDIA Corporation/NvToolsExt/bin/x64"
 cp -r /tmp/NvToolsExt/include/* "/c/Program Files/NVIDIA Corporation/NvToolsExt/include"
 cp -r /tmp/NvToolsExt/lib/x64/* "/c/Program Files/NVIDIA Corporation/NvToolsExt/lib/x64"
-export NVTOOLSEXT_PATH="/c/Program Files/NVIDIA Corporation/NvToolsExt/bin/x64"
-
-export TORCH_CUDA_ARCH_LIST="3.5;5.0+PTX;6.0;7.0;7.5;8.0;8.6"
-export TORCH_CUDA_ARCH_LIST="35;50+PTX;6.0;7.0;7.5;8.0;8.6"
diff --git a/.github/workflows/install.yml b/.github/workflows/install.yml
@@ -32,6 +32,8 @@ jobs:
           source ./.github/workflows/cuda/${{ runner.os }}-env.sh ${{ matrix.cuda-version }}
           pip install --verbose -e .
         shell: bash
+        env:
+          TORCH_CUDA_ARCH_LIST: "3.5;5.0+PTX;6.0;7.0;7.5;8.0;8.6"
 
       - name: Test imports
         run: |

diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -42,8 +42,6 @@ endif()
 if(WITH_CUDA)
   enable_language(CUDA)
   add_definitions(-DWITH_CUDA)
-  message("CUDA FLAGS HEHEHEHEHEHE")
-  message("${CMAKE_CUDA_FLAGS}")
   set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} --expt-relaxed-constexpr")
 
   if (NOT "$ENV{EXTERNAL_CUTLASS_INCLUDE_DIR}" STREQUAL "")

diff --git a/setup.py b/setup.py
@@ -66,13 +66,13 @@ def build_extension(self, ext):
             f'-DCMAKE_PREFIX_PATH={torch.utils.cmake_prefix_path}',
         ]
 
-        os.environ['TORCH_CUDA_ARCH_LIST'] = '8.0 8.6 9.0'
+        # os.environ['TORCH_CUDA_ARCH_LIST'] = '8.0 8.6 9.0'
 
-        cuda_arch_list = os.getenv('TORCH_CUDA_ARCH_LIST')
-        print("ARCH LIST")
-        print("-----------")
-        print(cuda_arch_list)
-        cmake_args.append('-DCUDA_ARCH_PTX=5.0+PTX')
+        # cuda_arch_list = os.getenv('TORCH_CUDA_ARCH_LIST')
+        # print("ARCH LIST")
+        # print("-----------")
+        # print(cuda_arch_list)
+        # cmake_args.append('-DCUDA_ARCH_PTX=5.0+PTX')
         # if WITH_CUDA and cuda_arch_list is not None:
         #     cmake_args.append(f'-DCMAKE_CUDA_ARCHITECTURES={cuda_arch_list}')
         # else: