cuda_core_driver_api_init_device_count_and_name.py

1from cuda import cuda, cudart
2
3def main():
4    # Initialize the CUDA Driver API
5    err, = cuda.cuInit(0)
6    if err != cuda.CUresult.CUDA_SUCCESS:
7        raise RuntimeError(f"CUDA initialization failed with error: {err}")
8
9    # Get the number of available CUDA devices
10    err, count = cuda.cuDeviceGetCount()
11    if err != cuda.CUresult.CUDA_SUCCESS:
12        raise RuntimeError(f"Failed to get device count with error: {err}")
13    
14    print(f"Number of devices found: {count}")
15
16    for i in range(count):
17        # Get handle for the device
18        err, device = cuda.cuDeviceGet(i)
19        if err != cuda.CUresult.CUDA_SUCCESS:
20            continue
21
22        # Get device name
23        err, name = cuda.cuDeviceGetName(128, device)
24        if err != cuda.CUresult.CUDA_SUCCESS:
25            continue
26            
27        print(f"Device {i}: {name.decode('utf-8')}")
28
29if __name__ == "__main__":
30    main()