xref: /petsc/src/sys/logging/handler/impls/nvtx/lognvtx.c (revision 586b72174521427559d8855856d614fb70566f6f)
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 
PetscLogHandlerEventBegin_NVTX(PetscLogHandler handler,PetscLogEvent event,PetscObject o1,PetscObject o2,PetscObject o3,PetscObject o4)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 
PetscLogHandlerEventEnd_NVTX(PetscLogHandler handler,PetscLogEvent event,PetscObject o1,PetscObject o2,PetscObject o3,PetscObject o4)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 
PetscLogHandlerCreate_NVTX(PetscLogHandler handler)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