Resolve the state-dependent EOS switch at build time under case optimization

A kernel is allocated registers for the worst path through its call graph, so the state-dependent EOS chain costs occupancy in every kernel that can reach it even when the runtime branch is never taken. On gfx90a that took the HLLC kernel from 142 to 250 VGPRs, three waves per SIMD down to two, and cost 17 to 27 percent grind on the Riemann-solver benchmarks with no such EOS in the case. Baking any_state_dependent_eos lets the compiler drop the chain outright: case-optimized, the kernel registers are now identical to master. Generic builds are unchanged, since a generic binary must still carry every family. fluid_pp%eos stays in the namelist while the flag is baked, so the initializer checks the two agree rather than silently running the wrong branch.

Author: sbryngelson

src/common/m_global_parameters_common.fpp

@@ -62,9 +62,13 @@ module m_global_parameters_common
     real(wp), allocatable, dimension(:) :: mg_mu_maxs  !< Compression at which a cubic Hugoniot fit turns over
     real(wp), allocatable, dimension(:) :: jwl_as, jwl_bs, jwl_r1s, jwl_r2s
     real(wp), allocatable, dimension(:) :: vinet_k0s, vinet_k0ps
-    logical :: any_state_dependent_eos                 !< True when some fluid's coefficients vary with density; set at init
+    !> any_state_dependent_eos is declared with the case-optimization block above: a parameter when the case is baked in, so the
+    !! compiler drops the whole state-dependent chain from kernels that never need it.
     $:GPU_DECLARE(create='[eoss, rho0s, t0s, gruneisen0s, gruneisen_as, mg_c0s, mg_ss, mg_s2s, mg_s3s, mg_mu_maxs, jwl_as, &
-                  & jwl_bs, jwl_r1s, jwl_r2s, vinet_k0s, vinet_k0ps, any_state_dependent_eos]')
+                  & jwl_bs, jwl_r1s, jwl_r2s, vinet_k0s, vinet_k0ps]')
+    #:if not MFC_CASE_OPTIMIZATION
+        $:GPU_DECLARE(create='[any_state_dependent_eos]')
+    #:endif
     !> @}
 
     !> @name Fluids participating in shear and bulk viscosity

src/common/m_variables_conversion.fpp

@@ -254,6 +254,7 @@ contains
         logical, optional, intent(in) :: enforce_density_floor, preserve_qbmm_number
         integer, optional, intent(in) :: lagrange_beta_index
         logical                       :: allocate_mixture_fields
+        logical                       :: state_dependent  !< Whether this case's fluids need a density-dependent EOS
 
         allocate_mixture_fields = .false.
         if (present(store_mixture_fields)) allocate_mixture_fields = store_mixture_fields
@@ -280,7 +281,7 @@ contains
         @:ALLOCATE(qvps    (1:num_fluids))
         @:ALLOCATE(Gs_vc     (1:num_fluids))
 
-        any_state_dependent_eos = .false.
+        state_dependent = .false.
         do i = 1, num_fluids
             gammas(i) = fluid_pp(i)%gamma
             isentrope_n(i) = f_isentrope_exponent(gammas(i))
@@ -335,11 +336,21 @@ contains
                 gruneisen0s(i) = fluid_pp(i)%vinet_gruneisen
                 gruneisen_as(i) = fluid_pp(i)%vinet_gruneisen_a
             end select
-            if (f_is_state_dependent(i)) any_state_dependent_eos = .true.
+            if (f_is_state_dependent(i)) state_dependent = .true.
         end do
+        #:if MFC_CASE_OPTIMIZATION
+            ! Baked in at build time, so a case that changed its EOS family since the build would silently
+            ! run the wrong branch. The namelist still carries fluid_pp%eos, so check the two agree.
+            @:PROHIBIT(state_dependent .neqv. any_state_dependent_eos, &
+                       & "This case's equations of state do not match the ones  this case-optimized binary was built for. Rebuild.")
+        #:else
+            any_state_dependent_eos = state_dependent
+        #:endif
         $:GPU_UPDATE(device='[gammas, isentrope_n, pi_infs, isentrope_B, cvs, qvs, qvps, Gs_vc, eoss, rho0s, t0s, gruneisen0s, &
-                     & gruneisen_as, mg_c0s, mg_ss, mg_s2s, mg_s3s, mg_mu_maxs, jwl_as, jwl_bs, jwl_r1s, jwl_r2s, vinet_k0s, &
-                     & vinet_k0ps, any_state_dependent_eos]')
+                     & gruneisen_as, mg_c0s, mg_ss, mg_s2s, mg_s3s, mg_mu_maxs, jwl_as, jwl_bs, jwl_r1s, jwl_r2s, vinet_k0s, vinet_k0ps]')
+        #:if not MFC_CASE_OPTIMIZATION
+            $:GPU_UPDATE(device='[any_state_dependent_eos]')
+        #:endif
 
         @:ALLOCATE(Res_vc(1:2, 1:max(1, Re_size_max)))
         Res_vc = dflt_real

toolchain/mfc/case.py

@@ -371,6 +371,12 @@ def __get_sim_fpp(self, print: bool) -> str:
             else:
                 num_vels = num_dims
 
+            # Baking this lets the compiler drop the state-dependent EOS chain entirely. Left in the
+            # call graph it costs registers, and so occupancy, in every kernel that can reach it.
+            eos_state_dependent = {3, 4, 5}  # Mie-Gruneisen, JWL, Vinet; see eos_* in m_constants.fpp
+            num_fluids_case = int(self.params.get("num_fluids", 1))
+            any_state_dependent_eos = 1 if any(int(self.params.get(f"fluid_pp({f})%eos", 1)) in eos_state_dependent for f in range(1, num_fluids_case + 1)) else 0
+
             mhd = 1 if self.params.get("mhd", "F") == "T" else 0
             relativity = 1 if self.params.get("relativity", "F") == "T" else 0
             viscous = 1 if self.params.get("viscous", "F") == "T" else 0
@@ -404,6 +410,7 @@ def __get_sim_fpp(self, print: bool) -> str:
 #:set igr_pres_lim          = {igr_pres_lim}
 #:set igr_order             = {self.params.get("igr_order", 3)}
 #:set viscous               = {viscous}
+#:set any_state_dependent_eos = {any_state_dependent_eos}
 """
 
         else:

toolchain/mfc/params/generators/fortran_gen.py

@@ -283,8 +283,9 @@ def generate_constants_fpp() -> str:
     ("muscl_polyn", "integer", "Degree of the MUSCL polynomials"),
     ("weno_num_stencils", "integer", "Number of stencils for WENO reconstruction"),
     ("wenojs", "logical", "WENO-JS (default)"),
+    ("any_state_dependent_eos", "logical", "Some fluid's coefficients vary with density"),
 ]
-COMMON_CASE_OPT_EXTRA_NAMES = {"num_dims", "num_vels", "weno_polyn", "muscl_polyn"}
+COMMON_CASE_OPT_EXTRA_NAMES = {"num_dims", "num_vels", "weno_polyn", "muscl_polyn", "any_state_dependent_eos"}
 
 _CASE_OPT_DECL_COL = 24  # '::' alignment for case-opt declarations