nvidia_cupti_callback_intercept_cuda_runtime_api_profiling.py

1import sys
2from cuda import cuda, cupti
3
4# Helper to check for errors
5def checkCudaErrors(result):
6    if isinstance(result, cuda.CUresult):
7        if result != cuda.CUresult.CUDA_SUCCESS:
8            raise RuntimeError(f"CUDA Error: {result}")
9    elif isinstance(result, cupti.CUptiResult):
10        if result != cupti.CUptiResult.CUPTI_SUCCESS:
11            raise RuntimeError(f"CUPTI Error: {result}")
12
13def callback_func(userdata, domain, cbid, cbdata):
14    # Check if the callback is for a CUDA Runtime API
15    if domain == cupti.CUpti_CallbackDomain.CUPTI_CB_DOMAIN_RUNTIME_API:
16        cb_info = cupti.CUpti_CallbackData(cbdata)
17        # Check when the API enters
18        if cb_info.callbackSite == cupti.CUpti_ApiCallbackSite.CUPTI_API_EXIT:
19            print(f"Intercepted exit of: {cb_info.functionName}")
20
21def main():
22    # Initialize CUDA
23    checkCudaErrors(cuda.cuInit(0))
24    device_id = 0
25    _, device = cuda.cuDeviceGet(device_id)
26    _, context = cuda.cuCtxCreate(0, device)
27
28    # Subscribe to CUPTI callbacks
29    subscriber = cupti.CUpti_SubscriberHandle(0)
30    checkCudaErrors(cupti.cuptiSubscribe(subscriber, callback_func, None))
31    
32    # Enable runtime API domain for the subscriber
33    checkCudaErrors(cupti.cuptiEnableDomain(1, subscriber, cupti.CUpti_CallbackDomain.CUPTI_CB_DOMAIN_RUNTIME_API))
34
35    print("Running a simple CUDA operation...")
36    # This will trigger the callback
37    _, free, total = cuda.cuMemGetInfo()
38    print(f"Free memory: {free}, Total memory: {total}")
39
40    # Cleanup
41    checkCudaErrors(cupti.cuptiUnsubscribe(subscriber))
42    checkCudaErrors(cuda.cuCtxDestroy(context))
43    print("Quickstart finished successfully.")
44
45if __name__ == "__main__":
46    main()