cupti_callback_subscription_for_cuda_runtime_api_profiling.py

1import sys
2from cuda import cuda, cupti
3
4def callback_handler(userdata, domain, cbid, cbdata):
5    """
6    Callback function that CUPTI calls for subscribed events.
7    """
8    # Check if the callback is for a runtime API
9    if domain == cupti.CUpti_CallbackDomain.CUPTI_CB_DOMAIN_RUNTIME_API:
10        cb_info = cupti.CUpti_CallbackData(cbdata)
11        # Print the name of the function being called
12        print(f"CUPTI Callback: {cb_info.functionName} called.")
13
14def main():
15    # Initialize the CUDA Driver API
16    err, = cuda.cuInit(0)
17    if err != cuda.CUresult.CUDA_SUCCESS:
18        print(f"CUDA Initialization failed: {err}")
19        return
20
21    # Subscribe to CUPTI callbacks
22    # This creates a subscriber that listens to the Runtime API domain
23    err, subscriber = cupti.cuptiSubscribe(0, callback_handler, None)
24    if err != cupti.CUpti_Result.CUPTI_SUCCESS:
25        print(f"CUPTI Subscribe failed: {err}")
26        return
27
28    # Enable callbacks for the Runtime API domain
29    err = cupti.cuptiEnableDomain(1, subscriber, cupti.CUpti_CallbackDomain.CUPTI_CB_DOMAIN_RUNTIME_API)
30    if err != cupti.CUpti_Result.CUPTI_SUCCESS:
31        print(f"CUPTI EnableDomain failed: {err}")
32        return
33
34    print("CUPTI Profiling initialized. Running dummy CUDA operations...")
35
36    # Perform a simple CUDA operation to trigger callbacks (e.g., getting device count)
37    err, count = cuda.cuDeviceGetCount()
38    
39    # Clean up
40    cupti.cuptiUnsubscribe(subscriber)
41    print("CUPTI Profiling finished.")
42
43if __name__ == "__main__":
44    main()