nvidia_cupti_gpu_activity_tracing_kernel_memcpy_quickstart.py

1import sys
2from cuda import cuda, cupti
3
4def checkCudaErrors(result):
5    if result[0] != cuda.CUresult.CUDA_SUCCESS:
6        raise RuntimeError(f"CUDA Error: {result[0]}")
7    if len(result) == 1:
8        return None
9    return result[1:]
10
11def cupti_record_callback(buffer_id, size, record):
12    while True:
13        record_ptr = cupti.cuptiActivityGetNextRecord(record, size)
14        if record_ptr[0] != cupti.CUptiResult.CUPTI_SUCCESS:
15            break
16        
17        activity = record_ptr[1]
18        if activity.kind == cupti.CUpti_ActivityKind.CUPTI_ACTIVITY_KIND_KERNEL:
19            print(f"Kernel Launch: {activity.name}")
20        elif activity.kind == cupti.CUpti_ActivityKind.CUPTI_ACTIVITY_KIND_MEMCPY:
21            print(f"Memcpy: {activity.copyKind}")
22
23def main():
24    # Initialize CUDA Driver API
25    checkCudaErrors(cuda.cuInit(0))
26    device = checkCudaErrors(cuda.cuDeviceGet(0))
27    context = checkCudaErrors(cuda.cuCtxCreate(0, device))
28
29    # Enable Activity Tracing
30    cupti.cuptiActivityEnable(cupti.CUpti_ActivityKind.CUPTI_ACTIVITY_KIND_KERNEL)
31    cupti.cuptiActivityEnable(cupti.CUpti_ActivityKind.CUPTI_ACTIVITY_KIND_MEMCPY)
32
33    # Register callbacks for buffer handling
34    # (Note: In a full implementation, you'd provide buffer allocation/request functions)
35    
36    print("CUPTI Activity Tracing enabled. Perform GPU operations here...")
37    
38    # Force a flush of activity records
39    cupti.cuptiActivityFlushAll(0)
40
41    # Cleanup
42    checkCudaErrors(cuda.cuCtxDestroy(context))
43    print("Quickstart complete.")
44
45if __name__ == "__main__":
46    main()