89 lines
4.7 KiB
Python
89 lines
4.7 KiB
Python
"""CUDA diagnostics and compiler propagation, without requiring a toolkit/GPU."""
|
|
import os
|
|
from pathlib import Path
|
|
import subprocess
|
|
import sys
|
|
import tempfile
|
|
from types import SimpleNamespace
|
|
import unittest
|
|
from unittest.mock import patch
|
|
|
|
sys.path.insert(0, str(Path(__file__).resolve().parents[1] / 'tools'))
|
|
from check_cuda_toolchain import check_toolchain
|
|
from build_mfem_bundle import mpich_cuda_architectures, prepare_hypre_source
|
|
|
|
|
|
class CudaToolchainTest(unittest.TestCase):
|
|
def test_probe_checks_both_compilers_without_running_gpu_code(self):
|
|
calls = []
|
|
|
|
def run(command, **kwargs):
|
|
calls.append(command)
|
|
if '--version' not in command:
|
|
self.assertNotIn('NVCC_APPEND_FLAGS', kwargs['env'])
|
|
self.assertNotIn('NVCC_CCBIN', kwargs['env'])
|
|
return subprocess.CompletedProcess(command, 0, 'CUDA test toolkit', '')
|
|
|
|
with patch.dict(os.environ, {'NVCC_APPEND_FLAGS': '-allow-unsupported-compiler',
|
|
'NVCC_CCBIN': 'wrong-host'}), patch('subprocess.run', side_effect=run):
|
|
check_toolchain('nvcc', 'gcc-compatible', 'g++-compatible')
|
|
self.assertEqual(len(calls), 3)
|
|
self.assertEqual(calls[1][2], 'gcc-compatible')
|
|
self.assertIn('-c', calls[1]) # A C driver need not link the C++ runtime.
|
|
self.assertEqual(calls[2][2], 'g++-compatible')
|
|
self.assertNotIn('-c', calls[2])
|
|
|
|
def test_mpich_architectures_follow_toolkit_and_user_selection(self):
|
|
self.assertEqual(mpich_cuda_architectures('86-real;90-virtual;86', 'nvcc'), '86,90')
|
|
self.assertEqual(mpich_cuda_architectures('native', 'nvcc'), 'auto')
|
|
with patch('build_mfem_bundle.output', return_value='sm_75\nsm_80\nsm_86\nsm_90'):
|
|
self.assertEqual(mpich_cuda_architectures('all-major', 'nvcc'), '75,80,90')
|
|
self.assertEqual(mpich_cuda_architectures('all', 'nvcc'), '75,80,86,90')
|
|
with self.assertRaisesRegex(RuntimeError, 'Invalid CUDA architecture'):
|
|
mpich_cuda_architectures('garbage', 'nvcc')
|
|
|
|
def test_hypre_patch_preserves_archive_and_older_cuda_branch(self):
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
root = Path(directory)
|
|
original = root / 'original'
|
|
utilities = original / 'src' / 'utilities'
|
|
utilities.mkdir(parents=True)
|
|
header = ('#include <thrust/execution_policy.h>\n'
|
|
'#define HYPRE_THRUST_NOT(pred) thrust::not1(pred)\n'
|
|
'using namespace thrust::placeholders;\n')
|
|
for name in ('device_utils.h', '_hypre_utilities.hpp'):
|
|
(utilities / name).write_text(header)
|
|
memory = ''.join(
|
|
f'HYPRE_CUDA_CALL( cudaMemPrefetchAsync(ptr, size, {device},\n'
|
|
' hypre_HandleComputeStream(hypre_handle())) );\n'
|
|
for device in ('hypre_HandleDevice(hypre_handle())', 'cudaCpuDeviceId'))
|
|
(utilities / 'memory.c').write_text(memory)
|
|
args = SimpleNamespace(hypre_source=original, work_dir=root / 'work')
|
|
self.assertEqual(prepare_hypre_source(args, {'cuda': False}), original)
|
|
patched = prepare_hypre_source(args, {'cuda': True})
|
|
self.assertEqual((utilities / 'memory.c').read_text(), memory)
|
|
updated = (patched / 'src/utilities/memory.c').read_text()
|
|
self.assertIn('#if CUDART_VERSION >= 13000', updated)
|
|
self.assertIn('cudaMemLocationTypeHost, 0', updated)
|
|
self.assertIn('cudaMemPrefetchAsync(ptr, size, cudaCpuDeviceId,', updated)
|
|
self.assertIn('#define HYPRE_THRUST_NOT(pred) thrust::not1(pred)',
|
|
(patched / 'src/utilities/device_utils.h').read_text())
|
|
prepare_hypre_source(args, {'cuda': True})
|
|
self.assertEqual((patched / 'src/utilities/memory.c').read_text(), updated)
|
|
|
|
def test_cxx_failure_reports_toolkit_host_and_recovery(self):
|
|
results = [subprocess.CompletedProcess([], 0, 'CUDA test toolkit', ''),
|
|
subprocess.CompletedProcess([], 0, '', ''),
|
|
subprocess.CompletedProcess([], 1, '', 'unsupported Microsoft Visual Studio version')]
|
|
with patch('subprocess.run', side_effect=results):
|
|
with self.assertRaises(RuntimeError) as caught:
|
|
check_toolchain('nvcc.exe', 'cl.exe', 'cl.exe')
|
|
message = str(caught.exception)
|
|
for expected in ['C++ compiler: cl.exe', 'CUDA test toolkit',
|
|
'unsupported Microsoft Visual Studio version', 'fresh Meson build directory']:
|
|
self.assertIn(expected, message)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|