cuda_python_driver_api_init_device_info_context_management.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: {err}")
13    
14    print(f"Number of devices: {count}")
15
16    if count > 0:
17        # Get handle for device 0
18        err, device = cuda.cuDeviceGet(0)
19        
20        # Get device name
21        err, name = cuda.cuDeviceGetName(128, device)
22        print(f"Device 0 Name: {name.decode('utf-8')}")
23
24        # Create a context for the device
25        err, context = cuda.cuCtxCreate(0, device)
26        if err == cuda.CUresult.CUDA_SUCCESS:
27            print("Successfully created CUDA context.")
28            
29            # Context cleanup
30            cuda.cuCtxDestroy(context)
31        else:
32            print(f"Failed to create context: {err}")
33
34if __name__ == "__main__":
35    main()