xref: /petsc/src/sys/logging/handler/impls/nvtx/lognvtx.c (revision bcd4bb4a4158aa96f212e9537e87b40407faf83e)
1 #include <petsc/private/logimpl.h> /*I "petscsys.h" I*/
2 #include <petsc/private/loghandlerimpl.h>
3 #include <petscdevice.h>
4 #if PETSC_PKG_CUDA_VERSION_GE(10, 0, 0)
5   #include <nvtx3/nvToolsExt.h>
6 #else
7   #include <nvToolsExt.h>
8 #endif
9 
10 static PetscErrorCode PetscLogHandlerEventBegin_NVTX(PetscLogHandler handler, PetscLogEvent event, PetscObject o1, PetscObject o2, PetscObject o3, PetscObject o4)
11 {
12   PetscLogState     state;
13   PetscLogEventInfo info;
14 
15   PetscFunctionBegin;
16   if (PetscDeviceInitialized(PETSC_DEVICE_CUDA)) {
17     PetscCall(PetscLogHandlerGetState(handler, &state));
18     PetscCall(PetscLogStateEventGetInfo(state, event, &info));
19     (void)nvtxRangePushA(info.name);
20   }
21   PetscFunctionReturn(PETSC_SUCCESS);
22 }
23 
24 static PetscErrorCode PetscLogHandlerEventEnd_NVTX(PetscLogHandler handler, PetscLogEvent event, PetscObject o1, PetscObject o2, PetscObject o3, PetscObject o4)
25 {
26   PetscFunctionBegin;
27   if (PetscDeviceInitialized(PETSC_DEVICE_CUDA)) (void)nvtxRangePop();
28   PetscFunctionReturn(PETSC_SUCCESS);
29 }
30 
31 /*MC
32   PETSCLOGHANDLERNVTX - PETSCLOGHANDLERNVTX = "nvtx" -  A
33   `PetscLogHandler` that creates an NVTX range (which appears in Nvidia Nsight
34   profiling) for each PETSc event.
35 
36   Options Database Keys:
37 + -log_nvtx   - start an nvtx log handler manually
38 - -log_nvtx 0 - stop the nvtx log handler from starting automatically in `PetscInitialize()` in a program run within an nsys profiling session (see Note)
39 
40   Level: developer
41 
42   Note:
43   If `PetscInitialize()` detects the environment variable `NSYS_PROFILING_SESSION_ID` (which is defined by `nsys
44   profile`) or `NVPROF_ID` (which is defined by `nvprof`) an instance of this log handler will automatically be
45   started.
46 
47 .seealso: [](ch_profiling), `PetscLogHandler`
48 M*/
49 
50 PETSC_INTERN PetscErrorCode PetscLogHandlerCreate_NVTX(PetscLogHandler handler)
51 {
52   PetscFunctionBegin;
53   handler->ops->eventbegin = PetscLogHandlerEventBegin_NVTX;
54   handler->ops->eventend   = PetscLogHandlerEventEnd_NVTX;
55   PetscCall(PetscInfo(handler, "nvtx log handler created\n"));
56   PetscFunctionReturn(PETSC_SUCCESS);
57 }
58