From 5c9ec64f6f1c956793904bf6316833c3c0e44f68 Mon Sep 17 00:00:00 2001
From: Matthew Thompson <matthew.thompson@nasa.gov>
Date: Tue, 20 Aug 2024 11:31:28 -0400
Subject: [PATCH 1/3] Update shell scripts for better spack integration

---
 scripts/fv3.j     |  6 ++++-
 scripts/fv3_setup | 66 ++++++++++++++++++++++++++++-------------------
 2 files changed, 45 insertions(+), 27 deletions(-)

diff --git a/scripts/fv3.j b/scripts/fv3.j
index b7f1946..f2e0ba6 100755
--- a/scripts/fv3.j
+++ b/scripts/fv3.j
@@ -132,7 +132,11 @@ set RUN_CMD = "$GEOSBIN/esma_mpirun -np "
 
 setenv ARCH `uname`
 source $GEOSBIN/g5_modules
-setenv LD_LIBRARY_PATH ${LD_LIBRARY_PATH}:${BASEDIR}/${ARCH}/lib:${GEOSDIR}/lib
+setenv @LD_LIBRARY_PATH_CMD ${LD_LIBRARY_PATH}:${GEOSDIR}/lib
+# We only add BASEDIR to the @LD_LIBRARY_PATH_CMD if BASEDIR is defined (i.e., not running with Spack)
+if ( $?BASEDIR ) then
+    setenv @LD_LIBRARY_PATH_CMD ${@LD_LIBRARY_PATH_CMD}:${BASEDIR}/${ARCH}/lib
+endif
 module list
 
 #########################################
diff --git a/scripts/fv3_setup b/scripts/fv3_setup
index 4895f13..76381f3 100755
--- a/scripts/fv3_setup
+++ b/scripts/fv3_setup
@@ -38,8 +38,16 @@ setenv NODE `uname -n`
 
 if ($ARCH == Darwin) then
    set FINDPATH = realpath
+   set PRELOAD_COMMAND = 'DYLD_INSERT_LIBRARIES'
+   set LD_LIBRARY_PATH_CMD = 'DYLD_LIBRARY_PATH'
+   # On macOS we seem to need to call mpirun directly and not use esma_mpirun
+   # For some reason SIP does not let the libraries be preloaded
+   set RUN_CMD = 'mpirun -np '
 else
    set FINDPATH = 'readlink -f'
+   set PRELOAD_COMMAND = 'LD_PRELOAD'
+   set LD_LIBRARY_PATH_CMD = 'LD_LIBRARY_PATH'
+   set RUN_CMD = '$GEOSBIN/esma_mpirun -np '
 endif
 set FV3SETUP = `$FINDPATH $0`
 set BINDIR   = `dirname $FV3SETUP`
@@ -64,9 +72,11 @@ endif
 # path to the sandbox. If we are not, then we set it to an empty
 # string.
 #
-# We also set the REAL_BIND_PATH environment variable. This is used
-# to set the bind paths for Singularity runs. REAL_BIND_PATH is the
-# physical path to the NOBACKUP directory
+# We also set the REAL_BIND_PATH and BASE_BIND_PATH environment
+# variables. These are used to set the bind paths for Singularity
+# runs. REAL_BIND_PATH is the physical path to the NOBACKUP directory
+# and BASE_BIND_PATH is the physical path to the /gpfsm directory
+# (needed for boundary conditions)
 
 setenv KEYFILE ".singularity.d"
 setenv singstat 0
@@ -85,34 +95,25 @@ while ($singstat == 0)
       endif
    endif
 end
-
 if ($singstat == 1) then
    setenv SINGULARITY_SANDBOX $SINGPATH
    setenv REAL_BIND_PATH `realpath $NOBACKUP`
-
-   set USING_SINGULARITY = TRUE
-   set SINGULARITY_BUILD = ""
-   set NATIVE_BUILD = "#DELETE"
+   setenv BASE_BIND_PATH "/gpfsm"
 else
    setenv SINGULARITY_SANDBOX ""
    setenv REAL_BIND_PATH ""
-
-   set USING_SINGULARITY = FALSE
-   set SINGULARITY_BUILD = "#DELETE"
-   set NATIVE_BUILD = ""
+   setenv BASE_BIND_PATH ""
 endif
 
 cd $ORIGPATH
 
-# Set TMPDIR to /tmp due to issues with heredocs in Singularity sandboxes
-# -----------------------------------------------------------------------
-
-setenv TMPDIR /tmp
-
 #######################################################################
 #                   Test for Command Line Flags
 #######################################################################
 
+# Set default behavior of switches
+set LINKX = FALSE
+set EXE_VERB = "copied"
 set SINGULARITY_RUN_SITE = ""
 if ($singstat == 1) then
    set USING_SINGULARITY = TRUE
@@ -131,6 +132,7 @@ if ($singstat == 1) then
 else
    set USING_SINGULARITY = FALSE
 endif
+
 while ( $#argv > 0 )
    set arg = $argv[1]
    shift argv
@@ -140,6 +142,17 @@ while ( $#argv > 0 )
       case --[Cc][Oo][Ll][Oo][Rr]:
          goto SETCOLOR
 
+      # Symlink StandAlone_FV3_Dycore.x
+      case --link:
+         set LINKX = TRUE
+         set EXE_VERB = "linked"
+         breaksw
+
+      # Avoid Symlink StandAlone_FV3_Dycore.x
+      case --nolink:
+         set LINKX = FALSE
+         breaksw
+
       # Here any string not above will trigger USAGE
       case -[Hh]:
       case --[Hh][Ee][Ll][Pp]:
@@ -342,6 +355,7 @@ else if ( $SITE == 'NAS' ) then
    echo "   ${C2}sky (Skylake)${CN} (default)"
    echo "   ${C2}cas (Cascade Lake)${CN}"
    echo "   ${C2}rom (AMD Rome)${CN}"
+   echo "   ${C2}mil (AMD Milan)${CN}"
    echo " "
    echo " NOTE Due to how FV3 is compiled by default, Sandy Bridge"
    echo "      and Ivy Bridge are not supported by current GEOS"
@@ -356,7 +370,8 @@ else if ( $SITE == 'NAS' ) then
        $MODEL != 'bro' & \
        $MODEL != 'sky' & \
        $MODEL != 'cas' & \
-       $MODEL != 'rom' ) goto ASKPROC
+       $MODEL != 'rom' & \
+       $MODEL != 'mil' ) goto ASKPROC
 
    # Some processors have weird names at NAS
    # ---------------------------------------
@@ -367,13 +382,11 @@ else if ( $SITE == 'NAS' ) then
       set MODEL = 'cas_ait'
    else if ($MODEL == rom) then
       set MODEL = 'rom_ait'
+   else if ($MODEL == mil) then
+      set MODEL = 'mil_ait'
    endif
 
-   if (     $MODEL == 'san') then
-      set NCPUS_PER_NODE = 16
-   else if ($MODEL == 'ivy') then
-      set NCPUS_PER_NODE = 20
-   else if ($MODEL == 'has') then
+   if ($MODEL == 'has') then
       set NCPUS_PER_NODE = 24
    else if ($MODEL == 'bro') then
       set NCPUS_PER_NODE = 28
@@ -383,6 +396,8 @@ else if ( $SITE == 'NAS' ) then
       set NCPUS_PER_NODE = 40
    else if ($MODEL == 'rom_ait') then
       set NCPUS_PER_NODE = 128
+   else if ($MODEL == 'mil_ait') then
+      set NCPUS_PER_NODE = 128
    endif
 
 else if( $SITE == 'AWS' | $SITE == 'Azure' ) then
@@ -780,7 +795,6 @@ EOF
 else if( $MPI_STACK == intelmpi ) then
 
 cat > $EXPDIR/SETENV.commands << EOF
-#setenv MPS_STAT_LEVEL 4
 setenv I_MPI_ADJUST_ALLREDUCE 12
 setenv I_MPI_ADJUST_GATHERV 3
 
@@ -798,8 +812,8 @@ setenv I_MPI_SHM_HEAP_VSIZE 512
 setenv PSM2_MEMORY large
 EOF
 
-#Testing at NCCS showed these caused crash at higher res
-#Crashed at restart read
+# Testing at NCCS showed these caused a crash at higher resolutions at
+# restart read
 if ( $USING_SINGULARITY == FALSE ) then
 
 cat >> $EXPDIR/SETENV.commands << EOF

From 0f63296e50c8eccf6b1ebb217afeafcf451f5c2f Mon Sep 17 00:00:00 2001
From: Matthew Thompson <matthew.thompson@nasa.gov>
Date: Tue, 20 Aug 2024 12:39:41 -0400
Subject: [PATCH 2/3] Update fv3_setup

---
 scripts/fv3_setup | 27 +++++++++++++++------------
 1 file changed, 15 insertions(+), 12 deletions(-)

diff --git a/scripts/fv3_setup b/scripts/fv3_setup
index 76381f3..92685ed 100755
--- a/scripts/fv3_setup
+++ b/scripts/fv3_setup
@@ -112,7 +112,6 @@ cd $ORIGPATH
 #######################################################################
 
 # Set default behavior of switches
-set LINKX = FALSE
 set EXE_VERB = "copied"
 set SINGULARITY_RUN_SITE = ""
 if ($singstat == 1) then
@@ -142,17 +141,6 @@ while ( $#argv > 0 )
       case --[Cc][Oo][Ll][Oo][Rr]:
          goto SETCOLOR
 
-      # Symlink StandAlone_FV3_Dycore.x
-      case --link:
-         set LINKX = TRUE
-         set EXE_VERB = "linked"
-         breaksw
-
-      # Avoid Symlink StandAlone_FV3_Dycore.x
-      case --nolink:
-         set LINKX = FALSE
-         breaksw
-
       # Here any string not above will trigger USAGE
       case -[Hh]:
       case --[Hh][Ee][Ll][Pp]:
@@ -704,6 +692,21 @@ endif
 if( -e $HOME/.GROUProot ) /bin/rm $HOME/.GROUProot
 echo $GROUP > $HOME/.GROUProot
 
+#######################################################################
+#                  Set some singularity variables
+#######################################################################
+
+# Copy or link GEOSgcm.x if USING_SINGULARITY is FALSE
+if ( $USING_SINGULARITY == FALSE ) then
+   # Set a couple variables for sed'ing out bits of the run script
+   set SINGULARITY_BUILD = "#DELETE"
+   set NATIVE_BUILD = ""
+else
+   # Set a couple variables for sed'ing out bits of the run script
+   set SINGULARITY_BUILD = ""
+   set NATIVE_BUILD = "#DELETE"
+endif
+
 #######################################################################
 #               Set Recommended MPI Stack Settings
 #######################################################################

From 97d83b124fa877192537dd32ff005d775ecbdb1f Mon Sep 17 00:00:00 2001
From: Matthew Thompson <matthew.thompson@nasa.gov>
Date: Tue, 20 Aug 2024 12:43:57 -0400
Subject: [PATCH 3/3] Fix setup

---
 scripts/fv3_setup | 1 +
 1 file changed, 1 insertion(+)

diff --git a/scripts/fv3_setup b/scripts/fv3_setup
index 92685ed..5cb8592 100755
--- a/scripts/fv3_setup
+++ b/scripts/fv3_setup
@@ -896,6 +896,7 @@ s?@SINGULARITY_BUILD?$SINGULARITY_BUILD?g
 s?@NATIVE_BUILD?$NATIVE_BUILD?g
 s?@SINGULARITY_SANDBOX?$SINGULARITY_SANDBOX?g
 s?@REAL_BIND_PATH?$REAL_BIND_PATH?g
+s?@LD_LIBRARY_PATH_CMD?$LD_LIBRARY_PATH_CMD?g
 
 s^@AGCM_GRIDNAME^$AGCM_GRIDNAME^g