73 lines
3.2 KiB
Python
73 lines
3.2 KiB
Python
#!/usr/bin/env python3
|
|
"""Probe the installed toolkit instead of maintaining a CUDA/compiler version table."""
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import os
|
|
from pathlib import Path
|
|
import shlex
|
|
import subprocess
|
|
import tempfile
|
|
|
|
|
|
def check_toolchain(nvcc: str, cc: str, cxx: str, *, standard: str = '17',
|
|
c_args: list[str] | None = None,
|
|
cxx_args: list[str] | None = None) -> str:
|
|
version = subprocess.run([nvcc, '--version'], text=True, capture_output=True, check=True).stdout.strip()
|
|
# Do not permit ambient flags to bypass the installed toolkit's version checks.
|
|
env = dict(os.environ)
|
|
for key in ('NVCC_PREPEND_FLAGS', 'NVCC_APPEND_FLAGS', 'NVCC_CCBIN'):
|
|
env.pop(key, None)
|
|
with tempfile.TemporaryDirectory(prefix='mfem-cuda-check-') as directory:
|
|
root = Path(directory)
|
|
source = root / 'probe.cu'
|
|
source.write_text('''#include <cuda_runtime.h>
|
|
#include <thrust/device_vector.h>
|
|
#include <cmath>
|
|
__global__ void probe(int *p) { *p = 1; }
|
|
int main() { int *p = nullptr; cudaMalloc(&p, sizeof(int));
|
|
probe<<<1, 1>>>(p); return (int)cudaFree(p); }
|
|
''')
|
|
for label, host, flags in [('C', cc, c_args or []), ('C++', cxx, cxx_args or [])]:
|
|
command = [nvcc, '--compiler-bindir', host, '--std=c++' + standard,
|
|
str(source), '-o', str(root / ('probe.exe' if os.name == 'nt' else 'probe'))]
|
|
if label == 'C':
|
|
command.append('-c')
|
|
for flag in flags:
|
|
command += ['-Xcompiler', flag]
|
|
result = subprocess.run(command, env=env, text=True, capture_output=True)
|
|
if result.returncode:
|
|
raise RuntimeError(
|
|
f'CUDA compatibility check failed for the selected {label} compiler: {host}\n'
|
|
f'{version}\nCommand: {shlex.join(command)}\n'
|
|
f'{result.stdout}{result.stderr}\n'
|
|
'Select a C/C++ toolchain supported by this CUDA toolkit using CC and CXX '
|
|
'(for example CC=gcc-15 CXX=g++-15), or install a compatible toolkit. '
|
|
'Use a fresh Meson build directory when changing compilers; --reconfigure '
|
|
'does not change cached compilers. For a CPU build use -Dmfem_cuda=disabled. '
|
|
'The probe compiles and links but does not require a running GPU.'
|
|
)
|
|
return f'CUDA compatibility checks passed: C={cc}, C++={cxx}, CUDA C++{standard}\n{version}'
|
|
|
|
|
|
def main() -> int:
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument('--nvcc', required=True)
|
|
parser.add_argument('--cc', required=True)
|
|
parser.add_argument('--cxx', required=True)
|
|
parser.add_argument('--standard', default='17')
|
|
parser.add_argument('--c-arg', action='append', default=[])
|
|
parser.add_argument('--cxx-arg', action='append', default=[])
|
|
args = parser.parse_args()
|
|
try:
|
|
print(check_toolchain(args.nvcc, args.cc, args.cxx, standard=args.standard,
|
|
c_args=args.c_arg, cxx_args=args.cxx_arg))
|
|
except (RuntimeError, OSError, subprocess.SubprocessError) as error:
|
|
print(error)
|
|
return 1
|
|
return 0
|
|
|
|
|
|
if __name__ == '__main__':
|
|
raise SystemExit(main())
|