1: #include <petsc/private/logimpl.h>
  2: #include <petsc/private/loghandlerimpl.h>
  3: #include <petscpkg_version.h>
  4: #include <petscdevice.h>
  5: #if PETSC_PKG_CUDA_VERSION_GE(10, 0, 0)
  6:   #include <nvtx3/nvToolsExt.h>
  7: #else
  8:   #include <nvToolsExt.h>
  9: #endif

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

 16:   PetscFunctionBegin;
 17:   if (PetscDeviceInitialized(PETSC_DEVICE_CUDA)) {
 18:     PetscCall(PetscLogHandlerGetState(handler, &state));
 19:     PetscCall(PetscLogStateEventGetInfo(state, event, &info));
 20:     (void)nvtxRangePushA(info.name);
 21:   }
 22:   PetscFunctionReturn(PETSC_SUCCESS);
 23: }

 25: static PetscErrorCode PetscLogHandlerEventEnd_NVTX(PetscLogHandler handler, PetscLogEvent event, PetscObject o1, PetscObject o2, PetscObject o3, PetscObject o4)
 26: {
 27:   PetscFunctionBegin;
 28:   if (PetscDeviceInitialized(PETSC_DEVICE_CUDA)) (void)nvtxRangePop();
 29:   PetscFunctionReturn(PETSC_SUCCESS);
 30: }

 32: /*MC
 33:   PETSCLOGHANDLERNVTX - PETSCLOGHANDLERNVTX = "nvtx" -  A
 34:   `PetscLogHandler` that creates an NVTX range (which appears in Nvidia Nsight
 35:   profiling) for each PETSc event.

 37:   Options Database Keys:
 38: + -log_nvtx   - start an nvtx log handler manually
 39: - -log_nvtx 0 - stop the nvtx log handler from starting automatically in `PetscInitialize()` in a program run within an nsys profiling session (see Note)

 41:   Level: developer

 43:   Note:
 44:   If `PetscInitialize()` detects the environment variable `NSYS_PROFILING_SESSION_ID` (which is defined by `nsys
 45:   profile`) or `NVPROF_ID` (which is defined by `nvprof`) an instance of this log handler will automatically be
 46:   started.

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

 51: PETSC_INTERN PetscErrorCode PetscLogHandlerCreate_NVTX(PetscLogHandler handler)
 52: {
 53:   PetscFunctionBegin;
 54:   handler->ops->eventbegin = PetscLogHandlerEventBegin_NVTX;
 55:   handler->ops->eventend   = PetscLogHandlerEventEnd_NVTX;
 56:   PetscCall(PetscInfo(handler, "nvtx log handler created\n"));
 57:   PetscFunctionReturn(PETSC_SUCCESS);
 58: }