Actual source code: lognvtx.c

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