cuda: Also read CUDART_VERSION from cuda_runtime_api.h

1 files changed, 25 insertions, 2 deletions

diff --git a/mesonbuild/dependencies/cuda.py b/mesonbuild/dependencies/cuda.py
index c962cae..3b2612b 100644
--- a/mesonbuild/dependencies/cuda.py
+++ b/mesonbuild/dependencies/cuda.py
@@ -15,6 +15,8 @@
 import glob
 import re
 import os
+import typing as T
+from .._pathlib import Path
 
 from .. import mlog
 from .. import mesonlib
@@ -149,11 +151,15 @@ class CudaDependency(ExternalDependency):
     toolkit_version_regex = re.compile(r'^CUDA Version\s+(.*)$')
     path_version_win_regex = re.compile(r'^v(.*)$')
     path_version_nix_regex = re.compile(r'^cuda-(.*)$')
+    cudart_version_regex = re.compile(r'#define\s+CUDART_VERSION\s+([0-9]+)')
 
-    def _cuda_toolkit_version(self, path):
+    def _cuda_toolkit_version(self, path: str) -> str:
         version = self._read_toolkit_version_txt(path)
         if version:
             return version
+        version = self._read_cuda_runtime_api_version(path)
+        if version:
+            return version
 
         mlog.debug('Falling back to extracting version from path')
         path_version_regex = self.path_version_win_regex if self._is_windows() else self.path_version_nix_regex
@@ -168,7 +174,24 @@ class CudaDependency(ExternalDependency):
 
         return '0.0'
 
-    def _read_toolkit_version_txt(self, path):
+    def _read_cuda_runtime_api_version(self, path_str: str) -> T.Optional[str]:
+        path = Path(path_str)
+        for i in path.rglob('cuda_runtime_api.h'):
+            raw = i.read_text()
+            m = self.cudart_version_regex.search(raw)
+            if not m:
+                continue
+            try:
+                vers_int = int(m.group(1))
+            except ValueError:
+                continue
+            # use // for floor instead of / which produces a float
+            major = vers_int // 1000                  # type: int
+            minor = (vers_int - major * 1000) // 10   # type: int
+            return '{}.{}'.format(major, minor)
+        return None
+
+    def _read_toolkit_version_txt(self, path: str) -> T.Optional[str]:
         # Read 'version.txt' at the root of the CUDA Toolkit directory to determine the tookit version
         version_file_path = os.path.join(path, 'version.txt')
         try: