vibe kappa shear

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

src/equation_of_state/MOM_EOS.F90

@@ -9,6 +9,7 @@ module MOM_EOS
 use MOM_EOS_linear, only : linear_EOS, avg_spec_vol_linear
 use MOM_EOS_linear, only : int_density_dz_linear, int_spec_vol_dp_linear
 use MOM_EOS_Wright, only : buggy_Wright_EOS, avg_spec_vol_buggy_Wright
+use MOM_EOS_Wright, only : calculate_density_derivs_elem_buggy_Wright_loc
 use MOM_EOS_Wright, only : int_density_dz_wright, int_spec_vol_dp_wright
 use MOM_EOS_Wright_full, only : Wright_full_EOS, avg_spec_vol_Wright_full
 use MOM_EOS_Wright_full, only : int_density_dz_wright_full, int_spec_vol_dp_wright_full
@@ -17,6 +18,7 @@ module MOM_EOS
 use MOM_EOS_Jackett06, only : Jackett06_EOS
 use MOM_EOS_UNESCO, only : UNESCO_EOS
 use MOM_EOS_Roquet_rho, only : Roquet_rho_EOS
+use MOM_EOS_Roquet_rho, only : calculate_density_derivs_elem_Roquet_rho_loc
 use MOM_EOS_Roquet_SpV, only : Roquet_SpV_EOS
 use MOM_EOS_TEOS10, only : TEOS10_EOS
 use MOM_EOS_TEOS10, only : gsw_sp_from_sr, gsw_pt_from_ct, gsw_sr_from_sp, gsw_ct_from_pt
@@ -46,6 +48,8 @@ module MOM_EOS
 public calculate_density_elem
 public calculate_density
 public calculate_density_derivs
+public calculate_density_derivs_elem_loc
+public get_EOS_form_and_scaling
 public calculate_density_second_derivs
 public calculate_spec_vol
 public calculate_specific_vol_derivs
@@ -992,6 +996,54 @@ subroutine calculate_density_derivs_1d(T, S, pressure, drho_dT, drho_dS, EOS, do
 
 end subroutine calculate_density_derivs_1d
 
+!> Device-callable dispatcher for density derivatives at a single point, in mks units,
+!! selecting the equation-of-state form at runtime by integer id (no polymorphic dispatch)
+!! so it can be called from inside a do concurrent / target region by whole-column GPU
+!! kernels. Unit rescaling (EOS%*_to_* factors) and any `scale` factor are the caller's
+!! responsibility, exactly as in calculate_density_derivs_1d. Forms without a device-callable
+!! _loc kernel are not handled here; a device-using module must FATAL at init on a GPU build
+!! before reaching this with an unsupported form.
+subroutine calculate_density_derivs_elem_loc(form_of_EOS, T, S, pressure, drho_dT, drho_dS)
+  integer, intent(in)  :: form_of_EOS !< The equation of state form (EOS_ROQUET_RHO, EOS_WRIGHT, ...)
+  real,    intent(in)  :: T           !< Temperature in the EOS kernel's mks units [degC]
+  real,    intent(in)  :: S           !< Salinity in the EOS kernel's mks units [ppt or g kg-1]
+  real,    intent(in)  :: pressure    !< Pressure [Pa]
+  real,    intent(out) :: drho_dT     !< Partial derivative of density wrt temperature [kg m-3 degC-1]
+  real,    intent(out) :: drho_dS     !< Partial derivative of density wrt salinity [kg m-3 ppt-1]
+  !$omp declare target
+
+  select case (form_of_EOS)
+    case (EOS_ROQUET_RHO)
+      call calculate_density_derivs_elem_Roquet_rho_loc(T, S, pressure, drho_dT, drho_dS)
+    case (EOS_WRIGHT)
+      call calculate_density_derivs_elem_buggy_Wright_loc(T, S, pressure, drho_dT, drho_dS)
+    case default
+      drho_dT = 0.0 ; drho_dS = 0.0
+  end select
+
+end subroutine calculate_density_derivs_elem_loc
+
+!> Return the equation-of-state form id and the unit-rescaling factors held in an EOS_type.
+!! Lets a caller (e.g. a whole-column GPU kernel) reproduce, host-side, the unit conversion
+!! and rescaling that calculate_density_derivs_1d applies around the mks _loc kernels, without
+!! needing access to the private components of EOS_type.
+subroutine get_EOS_form_and_scaling(EOS, form_of_EOS, kg_m3_to_R, C_to_degC, S_to_ppt, RL2_T2_to_Pa)
+  type(EOS_type), intent(in)  :: EOS          !< Equation of state structure
+  integer,        intent(out) :: form_of_EOS  !< The equation of state form id (EOS_ROQUET_RHO, ...)
+  real,           intent(out) :: kg_m3_to_R   !< Factor converting kg m-3 to the internal
+                                              !! density unit R [R m3 kg-1 ~> 1]
+  real,           intent(out) :: C_to_degC    !< Factor converting the temperature unit to degC [degC C-1 ~> 1]
+  real,           intent(out) :: S_to_ppt     !< Factor converting the salinity unit to ppt [ppt S-1 ~> 1]
+  real,           intent(out) :: RL2_T2_to_Pa !< Factor converting the pressure unit to Pa [Pa T2 R-1 L-2 ~> 1]
+
+  form_of_EOS  = EOS%form_of_EOS
+  kg_m3_to_R   = EOS%kg_m3_to_R
+  C_to_degC    = EOS%C_to_degC
+  S_to_ppt     = EOS%S_to_ppt
+  RL2_T2_to_Pa = EOS%RL2_T2_to_Pa
+
+end subroutine get_EOS_form_and_scaling
+
 
 !> Calls the appropriate subroutine to calculate density derivatives for 1-D array inputs.
 subroutine calculate_density_derivs_2d(T, S, pressure, drho_dT, drho_dS, EOS, dom)

src/equation_of_state/MOM_EOS_Roquet_rho.F90

@@ -10,6 +10,10 @@ module MOM_EOS_Roquet_rho
 implicit none ; private
 
 public Roquet_rho_EOS
+! Exposed as a device-callable (declare target) elemental so whole-column GPU kernels
+! can compute density derivatives in-region without polymorphic dispatch.
+public calculate_density_derivs_elem_Roquet_rho_loc
+!$omp declare target(calculate_density_derivs_elem_Roquet_rho_loc)
 
 real, parameter :: Pa2kb  = 1.e-8 !< Conversion factor between Pa and kbar [kbar Pa-1]
 !>@{ Parameters in the Roquet_rho (Roquet density) equation of state

src/equation_of_state/MOM_EOS_Wright.F90

@@ -12,6 +12,10 @@ module MOM_EOS_Wright
 implicit none ; private
 
 public buggy_Wright_EOS
+! Exposed as a device-callable (declare target) elemental so whole-column GPU kernels
+! can compute density derivatives in-region without polymorphic dispatch.
+public calculate_density_derivs_elem_buggy_Wright_loc
+!$omp declare target(calculate_density_derivs_elem_buggy_Wright_loc)
 public int_density_dz_wright, int_spec_vol_dp_wright
 public avg_spec_vol_buggy_Wright
 public set_params_buggy_Wright