1: #include <petsc/private/logimpl.h>
  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

 10: static PetscErrorCode PetscLogHandlerEventBegin_NVTX(PetscLogHandler handler, PetscLogEvent event, PetscObject o1, PetscObject o2, PetscObject o3, PetscObject o4)
 11: {
 12:   PetscLogState     state;
 13:   PetscLogEventInfo info;

 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: }

 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: }

 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.

 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)

 40:   Level: developer

 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.

 47: .seealso: [](ch_profiling), `PetscLogHandler`
 48: M*/

 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: }