Update jetson handling in cuda configuration

wenkangwei · Jun 24, 2021 · 8fe0e8b · 8fe0e8b
1 parent 8f7fcd1
commit 8fe0e8b
Showing 1 changed file with 7 additions and 4 deletions.
diff --git a/third_party/gpus/find_cuda_config.py b/third_party/gpus/find_cuda_config.py
@@ -73,6 +73,9 @@ class ConfigError(Exception):
   pass
 
 
+def _is_aarch64():
+  return platform.machine() == "aarch64"
+
 def _is_linux():
   return platform.system() == "Linux"
 
@@ -343,7 +346,7 @@ def get_header_version(path):
 
 def _find_cusolver_config(base_paths, required_version, cuda_version):
 
-  if _at_least_version(cuda_version, "11.0"):
+  if (_at_least_version(cuda_version, "11.0") or (_is_aarch64 and _at_least_version(cuda_version, "10.2"))):
 
     def get_header_version(path):
       version = (
@@ -372,7 +375,7 @@ def get_header_version(path):
 
 def _find_curand_config(base_paths, required_version, cuda_version):
 
-  if _at_least_version(cuda_version, "11.0"):
+  if (_at_least_version(cuda_version, "11.0") or (_is_aarch64 and _at_least_version(cuda_version, "10.2"))):
 
     def get_header_version(path):
       version = (
@@ -401,7 +404,7 @@ def get_header_version(path):
 
 def _find_cufft_config(base_paths, required_version, cuda_version):
 
-  if _at_least_version(cuda_version, "11.0"):
+  if (_at_least_version(cuda_version, "11.0") or (_is_aarch64 and _at_least_version(cuda_version, "10.2"))):
 
     def get_header_version(path):
       version = (
@@ -454,7 +457,7 @@ def get_header_version(path):
 
 def _find_cusparse_config(base_paths, required_version, cuda_version):
 
-  if _at_least_version(cuda_version, "11.0"):
+  if (_at_least_version(cuda_version, "11.0") or (_is_aarch64 and _at_least_version(cuda_version, "10.2"))):
 
     def get_header_version(path):
       version = (