cuda_runtime_api_device_properties_and_driver_version_query.py

1from cuda import cuda, cudart
2
3def main():
4    # Initialize the CUDA Runtime API
5    # Note: Many runtime APIs initialize the context implicitly on the first call.
6    
7    # Get the number of available CUDA devices
8    status, device_count = cudart.cudaGetDeviceCount()
9    if status != cuda.CUresult.CUDA_SUCCESS:
10        print(f"Error getting device count: {status}")
11        return
12
13    print(f"Found {device_count} CUDA device(s).")
14
15    for i in range(device_count):
16        # Get properties for each device
17        status, prop = cudart.cudaGetDeviceProperties(i)
18        if status == cuda.CUresult.CUDA_SUCCESS:
19            print(f"Device {i}: {prop.name.decode('utf-8')}")
20            print(f"  Compute Capability: {prop.major}.{minor}")
21            print(f"  Total Global Memory: {prop.totalGlobalMem / (1024**3):.2f} GB")
22        else:
23            print(f"Could not retrieve properties for device {i}")
24
25    # Get the CUDA Driver version
26    status, driver_version = cudart.cudaDriverGetVersion()
27    if status == cuda.CUresult.CUDA_SUCCESS:
28        print(f"CUDA Driver Version: {driver_version}")
29
30if __name__ == "__main__":
31    main()