From 26a0c420527a2f5deba521c6d562be8cb5861f37 Mon Sep 17 00:00:00 2001
From: Rahul Mahajan <aerorahul@users.noreply.github.com>
Date: Tue, 10 Dec 2024 17:21:21 -0500
Subject: [PATCH] update config.resources in GEFS system to allow use of
 traditional threading if desired

---
 parm/config/gefs/config.resources | 63 +++++++++++++++++++++++--------
 1 file changed, 47 insertions(+), 16 deletions(-)

diff --git a/parm/config/gefs/config.resources b/parm/config/gefs/config.resources
index a730ea401c..e1b9a036de 100644
--- a/parm/config/gefs/config.resources
+++ b/parm/config/gefs/config.resources
@@ -110,22 +110,33 @@ case ${step} in
       [[ "${DO_ICE}" == "YES" ]] && nthreads_cice6=1
     fi
 
-    # PETS for the atmosphere dycore
-    (( FV3PETS = ntasks_fv3 * nthreads_fv3 ))
-    echo "FV3 using (nthreads, PETS) = (${nthreads_fv3}, ${FV3PETS})"
+    # FV3
+    if [[ "${USE_ESMF_THREADING:-}" == "YES" ]]; then
+      (( FV3THREADS = nthreads_fv3 ))
+      (( FV3PETS = ntasks_fv3 * nthreads_fv3 ))
+    else
+      (( FV3THREADS = UFS_THREADS ))
+      (( FV3PETS = ntasks_fv3 ))
+    fi
+    echo "FV3 using (nthreads, PETS) = (${FV3THREADS}, ${FV3PETS})"
 
-    # PETS for quilting
+    # Write grid component
+    QUILTPETS=0; QUILTTHREADS=0
     if [[ "${QUILTING:-}" == ".true." ]]; then
-      (( QUILTPETS = ntasks_quilt * nthreads_fv3 ))
+      if [[ "${USE_ESMF_THREADING:-}" == "YES" ]]; then
+        (( QUILTTHREADS = nthreads_fv3 ))
+        (( QUILTPETS = ntasks_quilt * nthreads_fv3 ))
+      else
+        (( QUILTTHREADS = UFS_THREADS ))
+        (( QUILTPETS = ntasks_quilt ))
+      fi
       (( WRTTASK_PER_GROUP = WRTTASK_PER_GROUP_PER_THREAD ))
       export WRTTASK_PER_GROUP
-    else
-      QUILTPETS=0
     fi
-    echo "QUILT using (nthreads, PETS) = (${nthreads_fv3}, ${QUILTPETS})"
+    echo "QUILT using (nthreads, PETS) = (${QUILTTHREADS}, ${QUILTPETS})"
 
     # Total PETS for the atmosphere component
-    ATMTHREADS=${nthreads_fv3}
+    ATMTHREADS=${FV3THREADS}
     (( ATMPETS = FV3PETS + QUILTPETS ))
     export ATMPETS ATMTHREADS
     echo "FV3ATM using (nthreads, PETS) = (${ATMTHREADS}, ${ATMPETS})"
@@ -133,6 +144,7 @@ case ${step} in
     # Total PETS for the coupled model (starting w/ the atmosphere)
     NTASKS_TOT=${ATMPETS}
 
+    # Mediator
     # The mediator PETS can overlap with other components, usually it lands on the atmosphere tasks.
     # However, it is suggested limiting mediator PETS to 300, as it may cause the slow performance.
     # See https://docs.google.com/document/d/1bKpi-52t5jIfv2tuNHmQkYUe3hkKsiG_DG_s6Mnukog/edit
@@ -143,6 +155,7 @@ case ${step} in
     export MEDPETS MEDTHREADS
     echo "MEDIATOR using (threads, PETS) = (${MEDTHREADS}, ${MEDPETS})"
 
+    # GOCART
     CHMPETS=0; CHMTHREADS=0
     if [[ "${DO_AERO_FCST}" == "YES" ]]; then
       # GOCART shares the same grid and forecast tasks as FV3 (do not add write grid component tasks).
@@ -153,28 +166,46 @@ case ${step} in
     fi
     export CHMPETS CHMTHREADS
 
+    # Waves
     WAVPETS=0; WAVTHREADS=0
     if [[ "${DO_WAVE}" == "YES" ]]; then
-      (( WAVPETS = ntasks_ww3 * nthreads_ww3 ))
-      (( WAVTHREADS = nthreads_ww3 ))
+      if [[ "${USE_ESMF_THREADING:-}" == "YES" ]]; then
+        (( WAVTHREADS = nthreads_ww3 ))
+        (( WAVPETS = ntasks_ww3 * nthreads_ww3 ))
+      else
+        (( WAVTHREADS = UFS_THREADS ))
+        (( WAVPETS = ntasks_ww3 ))
+      fi
       echo "WW3 using (threads, PETS) = (${WAVTHREADS}, ${WAVPETS})"
       (( NTASKS_TOT = NTASKS_TOT + WAVPETS ))
     fi
     export WAVPETS WAVTHREADS
 
+    # Ocean
     OCNPETS=0; OCNTHREADS=0
     if [[ "${DO_OCN}" == "YES" ]]; then
-      (( OCNPETS = ntasks_mom6 * nthreads_mom6 ))
-      (( OCNTHREADS = nthreads_mom6 ))
+      if [[ "${USE_ESMF_THREADING:-}" == "YES" ]]; then
+        (( OCNTHREADS = nthreads_mom6 ))
+        (( OCNPETS = ntasks_mom6 * nthreads_mom6 ))
+      else
+        (( OCNTHREADS = UFS_THREADS ))
+        (( OCNPETS = ntasks_mom6 ))
+      fi
       echo "MOM6 using (threads, PETS) = (${OCNTHREADS}, ${OCNPETS})"
       (( NTASKS_TOT = NTASKS_TOT + OCNPETS ))
     fi
     export OCNPETS OCNTHREADS
 
+    # Ice
     ICEPETS=0; ICETHREADS=0
     if [[ "${DO_ICE}" == "YES" ]]; then
-      (( ICEPETS = ntasks_cice6 * nthreads_cice6 ))
-      (( ICETHREADS = nthreads_cice6 ))
+      if [[ "${USE_ESMF_THREADING:-}" == "YES" ]]; then
+        (( ICETHREADS = nthreads_cice6 ))
+        (( ICEPETS = ntasks_cice6 * nthreads_cice6 ))
+      else
+        (( ICETHREADS = UFS_THREADS ))
+        (( ICEPETS = ntasks_cice6 ))
+      fi
       echo "CICE6 using (threads, PETS) = (${ICETHREADS}, ${ICEPETS})"
       (( NTASKS_TOT = NTASKS_TOT + ICEPETS ))
     fi
@@ -184,7 +215,7 @@ case ${step} in
 
     declare -x "ntasks"="${NTASKS_TOT}"
     declare -x "threads_per_task"="${UFS_THREADS}"
-    declare -x "tasks_per_node"="${max_tasks_per_node}"
+    tasks_per_node=$(( max_tasks_per_node / threads_per_task ))
 
     case "${CASE}" in
       "C48" | "C96" | "C192")