Skip to content

MOM_cpu_clock_infra: emit NVTX ranges for MOM6 cpu clocks #79

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
JorgeG94 wants to merge 1 commit into
base: dev/gpu
Choose a base branch
from
+68 −0
Open

MOM_cpu_clock_infra: emit NVTX ranges for MOM6 cpu clocks #79

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
34 changes: 34 additions & 0 deletions config_src/infra/FMS1/MOM_cpu_clock_infra.F90
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,12 @@
!> Wraps the MPP cpu clock functions
!!
!! The functions and constants should be accessed via mom_cpu_clock
!!
!! Compiling with -DMOM_USE_NVTX additionally emits an NVTX range around every MOM6 cpu
!! clock, so each existing cpu_clock_id() name becomes a named range in an nsys timeline
!! with no call-site changes. It requires nvfortran and the NVTX library
!! (-DMOM_USE_NVTX ... -cudalib=nvtx). Undefined by default: a normal build compiles
!! exactly as before and links no extra library.
module MOM_cpu_clock_infra

! These interfaces and constants from MPP/FMS will not be directly exposed outside of this module
Expand All @@ -18,9 +24,22 @@ module MOM_cpu_clock_infra
use mpp_mod, only : MPP_CLOCK_ROUTINE => CLOCK_ROUTINE
use mpp_mod, only : MPP_CLOCK_LOOP => CLOCK_LOOP
use mpp_mod, only : MPP_CLOCK_INFRA => CLOCK_INFRA
#ifdef MOM_USE_NVTX
use nvtx, only : nvtxStartRange, nvtxEndRange
#endif

implicit none ; private

#ifdef MOM_USE_NVTX
!> The largest clock handle for which an NVTX range name is retained.
integer, parameter :: MAX_NVTX_CLOCKS = 4096
!> The NVTX range name for each clock handle, recorded by cpu_clock_id(). An empty entry
!! means no range is emitted for that handle. cpu_clock_begin() and cpu_clock_end() test
!! the same condition, so starts and ends stay balanced for handles that were never named
!! or that fall outside the table.
character(len=64), dimension(MAX_NVTX_CLOCKS) :: nvtx_clock_names = ""
#endif

! Public entities
public :: cpu_clock_id, cpu_clock_begin, cpu_clock_end
public :: CLOCK_COMPONENT, CLOCK_SUBCOMPONENT, CLOCK_MODULE_DRIVER, CLOCK_MODULE
Expand Down Expand Up @@ -60,6 +79,12 @@ module MOM_cpu_clock_infra
subroutine cpu_clock_begin(id)
integer, intent(in) :: id !< Handle for clock

#ifdef MOM_USE_NVTX
! Opened before, and closed after, the mpp clock so the NVTX range encloses it.
if (id > 0 .and. id <= MAX_NVTX_CLOCKS) then
if (len_trim(nvtx_clock_names(id)) > 0) call nvtxStartRange(trim(nvtx_clock_names(id)))
endif
#endif
call mpp_clock_begin(id)

end subroutine cpu_clock_begin
Expand All @@ -69,6 +94,11 @@ subroutine cpu_clock_end(id)
integer, intent(in) :: id !< Handle for clock

call mpp_clock_end(id)
#ifdef MOM_USE_NVTX
if (id > 0 .and. id <= MAX_NVTX_CLOCKS) then
if (len_trim(nvtx_clock_names(id)) > 0) call nvtxEndRange
endif
#endif

end subroutine cpu_clock_end

Expand Down Expand Up @@ -96,6 +126,10 @@ integer function cpu_clock_id(name, sync, grain)
endif

cpu_clock_id = mpp_clock_id(name, flags=clock_flags, grain=grain)
#ifdef MOM_USE_NVTX
if (cpu_clock_id > 0 .and. cpu_clock_id <= MAX_NVTX_CLOCKS) &
nvtx_clock_names(cpu_clock_id) = name
#endif
end function cpu_clock_id

end module MOM_cpu_clock_infra
34 changes: 34 additions & 0 deletions config_src/infra/FMS2/MOM_cpu_clock_infra.F90
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,12 @@
!> Wraps the MPP cpu clock functions
!!
!! The functions and constants should be accessed via mom_cpu_clock
!!
!! Compiling with -DMOM_USE_NVTX additionally emits an NVTX range around every MOM6 cpu
!! clock, so each existing cpu_clock_id() name becomes a named range in an nsys timeline
!! with no call-site changes. It requires nvfortran and the NVTX library
!! (-DMOM_USE_NVTX ... -cudalib=nvtx). Undefined by default: a normal build compiles
!! exactly as before and links no extra library.
module MOM_cpu_clock_infra

! These interfaces and constants from MPP/FMS will not be directly exposed outside of this module
Expand All @@ -18,9 +24,22 @@ module MOM_cpu_clock_infra
use mpp_mod, only : MPP_CLOCK_ROUTINE => CLOCK_ROUTINE
use mpp_mod, only : MPP_CLOCK_LOOP => CLOCK_LOOP
use mpp_mod, only : MPP_CLOCK_INFRA => CLOCK_INFRA
#ifdef MOM_USE_NVTX
use nvtx, only : nvtxStartRange, nvtxEndRange
#endif

implicit none ; private

#ifdef MOM_USE_NVTX
!> The largest clock handle for which an NVTX range name is retained.
integer, parameter :: MAX_NVTX_CLOCKS = 4096
!> The NVTX range name for each clock handle, recorded by cpu_clock_id(). An empty entry
!! means no range is emitted for that handle. cpu_clock_begin() and cpu_clock_end() test
!! the same condition, so starts and ends stay balanced for handles that were never named
!! or that fall outside the table.
character(len=64), dimension(MAX_NVTX_CLOCKS) :: nvtx_clock_names = ""
#endif

! Public entities
public :: cpu_clock_id, cpu_clock_begin, cpu_clock_end
public :: CLOCK_COMPONENT, CLOCK_SUBCOMPONENT, CLOCK_MODULE_DRIVER, CLOCK_MODULE
Expand Down Expand Up @@ -60,6 +79,12 @@ module MOM_cpu_clock_infra
subroutine cpu_clock_begin(id)
integer, intent(in) :: id !< Handle for clock

#ifdef MOM_USE_NVTX
! Opened before, and closed after, the mpp clock so the NVTX range encloses it.
if (id > 0 .and. id <= MAX_NVTX_CLOCKS) then
if (len_trim(nvtx_clock_names(id)) > 0) call nvtxStartRange(trim(nvtx_clock_names(id)))
endif
#endif
call mpp_clock_begin(id)

end subroutine cpu_clock_begin
Expand All @@ -69,6 +94,11 @@ subroutine cpu_clock_end(id)
integer, intent(in) :: id !< Handle for clock

call mpp_clock_end(id)
#ifdef MOM_USE_NVTX
if (id > 0 .and. id <= MAX_NVTX_CLOCKS) then
if (len_trim(nvtx_clock_names(id)) > 0) call nvtxEndRange
endif
#endif

end subroutine cpu_clock_end

Expand Down Expand Up @@ -96,6 +126,10 @@ integer function cpu_clock_id(name, sync, grain)
endif

cpu_clock_id = mpp_clock_id(name, flags=clock_flags, grain=grain)
#ifdef MOM_USE_NVTX
if (cpu_clock_id > 0 .and. cpu_clock_id <= MAX_NVTX_CLOCKS) &
nvtx_clock_names(cpu_clock_id) = name
#endif
end function cpu_clock_id

end module MOM_cpu_clock_infra
Loading