1: #include <petsc/private/logimpl.h>
  2: #include <petsc/private/loghandlerimpl.h>
  3: #include <petscpkg_version.h>
  4: #include <petscdevice_hip.h>

  6: #if PetscDefined(HAVE_ROCTX)

  8:   #if PETSC_PKG_HIP_VERSION_GE(6, 4, 0)
  9:     #include <rocprofiler-sdk-roctx/roctx.h>

 11: static PetscErrorCode PetscLogHandlerEventsPause_ROCTX(PetscLogHandler h)
 12: {
 13:   PetscFunctionBegin;
 14:   /* Pause all profiling */
 15:   PetscInt err = roctxProfilerPause(0);
 16:   PetscCheck(err == 0, PETSC_COMM_SELF, PETSC_ERR_GPU, "Failed to pause ROCTX profiler");
 17:   PetscFunctionReturn(PETSC_SUCCESS);
 18: }

 20: static PetscErrorCode PetscLogHandlerEventsResume_ROCTX(PetscLogHandler h)
 21: {
 22:   PetscFunctionBegin;
 23:   /* Resume all profiling */
 24:   PetscInt err = roctxProfilerResume(0);
 25:   PetscCheck(err == 0, PETSC_COMM_SELF, PETSC_ERR_GPU, "Failed to resume ROCTX profiler");
 26:   PetscFunctionReturn(PETSC_SUCCESS);
 27: }

 29:   #elif PETSC_PKG_HIP_VERSION_GE(6, 0, 0)
 30:     #include <roctracer/roctx.h>
 31:   #elif PETSC_PKG_HIP_VERSION_GE(5, 0, 0)
 32:     #include <roctx.h>
 33:   #endif

 35: #endif

 37: #if PetscDefined(HAVE_ROCTX)

 39: static PetscErrorCode PetscLogHandlerEventBegin_ROCTX(PetscLogHandler handler, PetscLogEvent event, PetscObject o1, PetscObject o2, PetscObject o3, PetscObject o4)
 40: {
 41:   PetscLogState     state;
 42:   PetscLogEventInfo info;

 44:   PetscFunctionBegin;
 45:   if (PetscDeviceInitialized(PETSC_DEVICE_HIP)) {
 46:     PetscCall(PetscLogHandlerGetState(handler, &state));
 47:     PetscCall(PetscLogStateEventGetInfo(state, event, &info));
 48:     (void)roctxRangePush(info.name);
 49:   }
 50:   PetscFunctionReturn(PETSC_SUCCESS);
 51: }

 53: static PetscErrorCode PetscLogHandlerEventEnd_ROCTX(PetscLogHandler handler, PetscLogEvent event, PetscObject o1, PetscObject o2, PetscObject o3, PetscObject o4)
 54: {
 55:   PetscFunctionBegin;
 56:   if (PetscDeviceInitialized(PETSC_DEVICE_HIP)) {
 57:     (void)roctxRangePush("StreamSync0");
 58:     /* Sync the default stream to ensure proper timing within event*/
 59:     PetscCallHIP(hipDeviceSynchronize());
 60:     (void)roctxRangePop();
 61:     (void)roctxRangePop();
 62:   }
 63:   PetscFunctionReturn(PETSC_SUCCESS);
 64: }

 66: #else

 68: static PetscErrorCode PetscLogHandlerEventBegin_ROCTX(PetscLogHandler handler, PetscLogEvent event, PetscObject o1, PetscObject o2, PetscObject o3, PetscObject o4)
 69: {
 70:   PetscFunctionBegin;
 71:   PetscFunctionReturn(PETSC_SUCCESS);
 72: }

 74: static PetscErrorCode PetscLogHandlerEventEnd_ROCTX(PetscLogHandler handler, PetscLogEvent event, PetscObject o1, PetscObject o2, PetscObject o3, PetscObject o4)
 75: {
 76:   PetscFunctionBegin;
 77:   PetscFunctionReturn(PETSC_SUCCESS);
 78: }

 80: #endif

 82: /*MC
 83:   PETSCLOGHANDLERROCTX - PETSCLOGHANDLERROCTX = "roctx" -  A `PetscLogHandler` that creates an ROCTX range (which appears in rocprof profiling) for each PETSc event.

 85:   Options Database Keys:
 86: + -log_roctx   - start an roctx log handler manually
 87: - -log_roctx 0 - stop the roctx log handler from starting automatically in `PetscInitialize()` in a program run within a rocprof profiling session

 89:   Level: developer

 91: .seealso: [](ch_profiling), `PetscLogHandler`
 92: M*/

 94: PETSC_INTERN PetscErrorCode PetscLogHandlerCreate_ROCTX(PetscLogHandler handler)
 95: {
 96:   PetscFunctionBegin;
 97:   handler->ops->eventbegin = PetscLogHandlerEventBegin_ROCTX;
 98:   handler->ops->eventend   = PetscLogHandlerEventEnd_ROCTX;
 99: #if PetscDefined(HAVE_ROCTX)
100:   #if PETSC_PKG_HIP_VERSION_GE(6, 4, 0)
101:   PetscCall(PetscObjectComposeFunction((PetscObject)handler, "PetscLogHandlerEventsPause_C", PetscLogHandlerEventsPause_ROCTX));
102:   PetscCall(PetscObjectComposeFunction((PetscObject)handler, "PetscLogHandlerEventsResume_C", PetscLogHandlerEventsResume_ROCTX));
103:   #endif
104: #endif
105:   PetscCall(PetscInfo(handler, "roctx log handler created\n"));
106:   PetscFunctionReturn(PETSC_SUCCESS);
107: }