!
! ===============================================================================
! PROFILING MODULE FOR RANGES
! =================================================================================
! From UtilXlib/clocks_handler.f90 of QuantumESPRESSO (v7.2)
! 14 Sep 2023 - bellenlau@cineca  
! 
!
! This module takes origin from QuantumESPRESSO inner clocks.
! It can be used to wrap routines and invoke ranges from different libraries
! By now, we support nvtx and roctx ranges for NVIDIA and AMD architectures
!
! There are three routines:
! - init_prof() to initialize the arrays checking if the range with a given name is already defined
! - start_prof() to start the range
! - stop_prof() to stop the range
!
! start and stop take two arguments. 
!       - If only the label is used, than a range with same name is looked for. 
!                - If it is found, the same colour is used. 
!                - If it is not found, a new color for the range is defined
!       - If also a number is used, than the number directly defines the color. No checks
!
! If a range with a given name is found, the color attribute of the range is the same
! These routines can be inserted into timers to start/stop library-agnostic ranges
!
!----------------------------------------------------------------------------
MODULE myprof
  !----------------------------------------------------------------------------
  !
#if defined(__PROFILE_NVTX)
  USE nvtx
#elif defined(__PROFILE_ROCTX)
  USE roctx
#endif
  !
  IMPLICIT NONE
  !
  SAVE
  !
  INTEGER, PARAMETER :: stdout = 6    ! unit connected to standard output
  INTEGER, PARAMETER :: maxevent = 64
  REAL*8, PARAMETER :: notrunning = -1.d0
  REAL*8, PARAMETER :: running = -2.d0
  !
  REAL*8 :: prof(maxevent)
  CHARACTER(len=12) :: prof_label(maxevent)
  !
  INTEGER :: nprof = 0
  LOGICAL :: no
  !
  CONTAINS
    !
    SUBROUTINE init_prof()
      IMPLICIT NONE
      !
      INTEGER :: n, ierr
      !
      !
      DO n = 1, maxevent
        !
        prof(n) = notrunning 
        prof_label(n) = ' '
        !
      END DO
      !
      RETURN

    END SUBROUTINE
    !
    SUBROUTINE start_prof(label,ncl)
      IMPLICIT NONE
      !
      CHARACTER(len=*), INTENT(IN) :: label
      INTEGER, INTENT(IN), OPTIONAL :: ncl
      !
      CHARACTER(len=12) :: label_
      INTEGER :: n

      label_ = TRIM(label)
      !  
      IF ( PRESENT(ncl) ) THEN
         !
         ! If ncl is present, does not look for a range with the same label
         !
#if defined(__PROFILE_NVTX)
         CALL nvtxStartRange(label_, ncl)
#elif defined(__PROFILE_ROCTX)
         CALL roctxStartRange(label_)
#endif
         !
         RETURN
         !
      ENDIF
      !
      ! If ncl is not present, looks for a range with same label
      !
      DO n = 1, nprof
        !
        IF (prof_label(n) == label_) THEN
          !
          ! ... found previously defined range: use the same index (color)
          !
          IF (prof(n) /= notrunning) THEN
          !            WRITE( stdout, '("start_clock: clock # ",I2," for ",A12, &
          !                           & " already started")' ) n, label_
          ELSE
            !
            prof(n) = running
            !
#if defined(__PROFILE_NVTX)
            CALL nvtxStartRange(label_, n)
#elif defined(__PROFILE_ROCTX)
            CALL roctxStartRange(label_)
#endif
            !
          END IF
          !
          RETURN
          !
        END IF
        !
      END DO
      ! 
      ! ... found a previously defined range: add a new one
      !
      IF (nprof == maxevent) THEN
        !
        WRITE (stdout, '("start_prof(",A,"): Too many events! call ignored")') label
        !
      ELSE
        !
        nprof = nprof + 1
        prof_label(nprof) = label_
        prof(nprof) = running
        !
#if defined(__PROFILE_NVTX)
        CALL nvtxStartRange(label_, nprof)
#elif defined(__PROFILE_ROCTX)
        CALL roctxStartRange(label_)
#endif
        !
      END IF
      !
      RETURN


    END SUBROUTINE
    !
    SUBROUTINE stop_prof(label,ncl)
      IMPLICIT NONE
      !
      CHARACTER(len=*), INTENT(IN) :: label
      INTEGER, INTENT(IN), OPTIONAL :: ncl
      !
      CHARACTER(len=12) :: label_
      INTEGER :: n

      label_ = trim(label)
      !
      IF ( PRESENT(ncl) ) THEN
         !
         ! If ncl is present, closes the any open range
         !
#if defined(__PROFILE_NVTX)
         CALL nvtxEndRange
#elif defined(__PROFILE_ROCTX)
         CALL roctxEndRange
#endif
         !
         RETURN
         !
      ENDIF
      !
      ! If ncl is not present, checks for an open range with same label
      !
      DO n = 1, nprof
        !
        IF (prof_label(n) == label_) THEN
          !
          ! ... found previously defined clock : check if properly initialised,
          ! ... close the nvtx range
          !
          IF (prof(n) == notrunning) THEN
            !WRITE (stdout, '("stop_prof: range # ",I2," for ",A12, " not running")') n, label
            !
          ELSE
            !
            prof(n) = notrunning
            !
#if defined(__PROFILE_NVTX)
            CALL nvtxEndRange
#elif defined(__PROFILE_ROCTX)
         CALL roctxEndRange
#endif
            !
          END IF
          !
          RETURN
          !
        END IF
        !
      END DO
      !
      ! ... clock not found
      !
      WRITE (stdout, '("stop_prof: no range starter for ",A12," found !")') label
      !
      RETURN

    END SUBROUTINE

END MODULE myprof
!

