|
1 | 1 | """Collection of utilities to detect properties of the underlying architecture.""" |
2 | 2 |
|
| 3 | +from contextlib import suppress |
3 | 4 | from functools import cached_property |
4 | | -from subprocess import PIPE, Popen, DEVNULL, run |
| 5 | +from subprocess import PIPE, Popen, DEVNULL, run, CalledProcessError |
5 | 6 | from pathlib import Path |
6 | 7 | import ctypes |
7 | 8 | import re |
|
11 | 12 |
|
12 | 13 | import cpuinfo |
13 | 14 | import numpy as np |
| 15 | +from packaging.version import parse, InvalidVersion |
14 | 16 | import psutil |
15 | 17 |
|
16 | 18 | from devito.logger import warning |
@@ -553,6 +555,30 @@ def get_cuda_path(): |
553 | 555 | return None |
554 | 556 |
|
555 | 557 |
|
| 558 | +@memoized_func |
| 559 | +def get_cuda_version(): |
| 560 | + cuda_home = get_cuda_path() |
| 561 | + if cuda_home is None: |
| 562 | + nvc_version_command = ['nvcc', '--version'] |
| 563 | + else: |
| 564 | + nvc_version_command = [f'{cuda_home}/bin/nvcc', '--version'] |
| 565 | + |
| 566 | + cuda_version = None |
| 567 | + try: |
| 568 | + out = run(nvc_version_command, capture_output=True, text=True) |
| 569 | + except (FileNotFoundError, CalledProcessError): |
| 570 | + pass |
| 571 | + finally: |
| 572 | + if out.returncode == 0: |
| 573 | + start = out.stdout.find('release') |
| 574 | + start = out.stdout.find(',', start) |
| 575 | + stop = out.stdout.find('\n', start) |
| 576 | + with supress(InvalidVersion): |
| 577 | + cuda_version = parse(out.stdout[start:stop]) |
| 578 | + |
| 579 | + return cuda_version |
| 580 | + |
| 581 | + |
556 | 582 | @memoized_func |
557 | 583 | def get_advisor_path(): |
558 | 584 | """ |
|
0 commit comments