update

pyg-team · Apr 7, 2024 · 4016fc3 · 4016fc3
1 parent bec7096
commit 4016fc3
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 5 deletions.
diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -42,6 +42,8 @@ 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
@@ -70,11 +70,12 @@ def build_extension(self, ext):
         print("ARCH LIST")
         print("-----------")
         print(cuda_arch_list)
-        if WITH_CUDA and cuda_arch_list is not None:
-            cmake_args.append(f'-DCMAKE_CUDA_ARCHITECTURES={cuda_arch_list}')
-        else:
-            cuda_arch_list = "50;60;70;75;80;86"
-            cmake_args.append(f'-DCMAKE_CUDA_ARCHITECTURES={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:
+        #     cuda_arch_list = "50;60;70;75;80;86"
+        #     cmake_args.append(f'-DCMAKE_CUDA_ARCHITECTURES={cuda_arch_list}')
 
         if CMakeBuild.check_env_flag('USE_MKL_BLAS'):
             include_dir = f"{sysconfig.get_path('data')}{os.sep}include"