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

  5: #if PetscDefined(HAVE_ROCTX)

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

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

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

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

 34: #endif

 36: #if PetscDefined(HAVE_ROCTX)

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

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

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

 65: #else

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

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

 79: #endif

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

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

 88:   Level: developer

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

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