experiment: claude ported everything

config_src/infra/FMS1/MOM_cpu_clock_infra.F90

@@ -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
@@ -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
@@ -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
@@ -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
 
@@ -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

config_src/infra/FMS2/MOM_cpu_clock_infra.F90

@@ -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
@@ -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
@@ -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
@@ -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
 
@@ -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