fix paddle.version.cuda_arch() (#68173)

yuanlehome · web-flow · commit d3915c807d1b · 2024-09-12T23:19:33.000-07:00
* fix paddle.version.cuda_arch()

* fix paddle.version.cuda_arch()
diff --git a/cmake/cuda.cmake b/cmake/cuda.cmake
@@ -244,6 +244,11 @@ function(select_nvcc_arch_flags out_variable out_arch_bin)
   string(REPLACE ";" " " nvcc_archs_readable "${nvcc_archs_readable}")
   string(REGEX MATCHALL "[0-9()]+" nvcc_archs_bin_list "${nvcc_archs_bin_list}")
   string(JOIN "," nvcc_real_archs ${nvcc_archs_bin_list})
+
+  set(COMPILED_CUDA_ARCHS
+      "${nvcc_real_archs}"
+      CACHE STRING "Specify compiled cuda archs.")
+
   set(${out_variable}
       ${nvcc_flags}
       PARENT_SCOPE)
@@ -256,9 +261,6 @@ function(select_nvcc_arch_flags out_variable out_arch_bin)
   set(${out_arch_bin}
       ${cuda_arch_bin}
       PARENT_SCOPE)
-  set(COMPILED_CUDA_ARCHS
-      ${cuda_arch_bin}
-      CACHE STRING "Specify compiled cuda archs.")
 endfunction()
 
 message(STATUS "CUDA detected: " ${CMAKE_CUDA_COMPILER_VERSION})
diff --git a/python/setup.py.in b/python/setup.py.in
@@ -123,7 +123,8 @@ def get_cinn_version() -> str:
 def get_cuda_archs() -> list[int]:
     compiled_cuda_archs = '@COMPILED_CUDA_ARCHS@'
     if isinstance(compiled_cuda_archs, str):
-        return [int(arch) for arch in compiled_cuda_archs.split(";") if arch]
+        compiled_cuda_archs = re.findall(r'\d+', compiled_cuda_archs)
+        return [int(arch) for arch in compiled_cuda_archs]
     else:
         return []
 
diff --git a/setup.py b/setup.py
@@ -444,7 +444,8 @@ def get_cinn_version() -> str:
 def get_cuda_archs() -> list[int]:
     compiled_cuda_archs = env_dict.get("COMPILED_CUDA_ARCHS")
     if isinstance(compiled_cuda_archs, str):
-        return [int(arch) for arch in compiled_cuda_archs.split(";") if arch]
+        compiled_cuda_archs = re.findall(r'\d+', compiled_cuda_archs)
+        return [int(arch) for arch in compiled_cuda_archs]
     else:
         return []
 
