cuda_driver_api_gpu_enumeration_with_compute_capability.py

1from cuda import cuda, nvrtc
2import numpy as np
3
4def checkCudaErrors(result):
5    if result[0].value:
6        raise RuntimeError("CUDA error code={}({})".format(result[0].value, _cudaGetErrorString(result[0])))
7    if len(result) == 1:
8        return None
9    elif len(result) == 2:
10        return result[1]
11    else:
12        return result[1:]
13
14def _cudaGetErrorString(error):
15    res, err_str = cuda.cuGetErrorString(error)
16    if res.value == 0:
17        return err_str
18    return "Unknown Error"
19
20# Initialize CUDA Driver API
21checkCudaErrors(cuda.cuInit(0))
22
23# Get number of devices
24num_devices = checkCudaErrors(cuda.cuDeviceGetCount())
25print(f"Found {num_devices} device(s)")
26
27for i in range(num_devices):
28    device = checkCudaErrors(cuda.cuDeviceGet(i))
29    name = checkCudaErrors(cuda.cuDeviceGetName(256, device))
30    
31    # Get Compute Capability
32    major = checkCudaErrors(cuda.cuDeviceGetAttribute(cuda.CUdevice_attribute.CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR, device))
33    minor = checkCudaErrors(cuda.cuDeviceGetAttribute(cuda.CUdevice_attribute.CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MINOR, device))
34    
35    print(f"Device {i}: {name.decode('utf-8')}")
36    print(f"  Compute Capability: {major}.{minor}")