fix gpu templates for real though (#1077)

Author: sbryngelson

toolchain/mfc/args.py

@@ -59,7 +59,7 @@ def add_common_arguments(p: argparse.ArgumentParser, mask = None):
         if "m" not in mask:
             for f in dataclasses.fields(config):
                 if f.name == 'gpu':
-                    p.add_argument(f"--{f.name}", action="store", nargs='?', const= gpuConfigOptions.ACC.value,default=gpuConfigOptions.ACC.value, dest=f.name, choices=[e.value for e in gpuConfigOptions], help=f"Turn the {f.name} option to OpenACC or OpenMP.")
+                    p.add_argument(f"--{f.name}", action="store", nargs='?', const= gpuConfigOptions.ACC.value,default=gpuConfigOptions.NONE.value, dest=f.name, choices=[e.value for e in gpuConfigOptions], help=f"Turn the {f.name} option to OpenACC or OpenMP.")
                     p.add_argument(f"--no-{f.name}", action="store_const", const = gpuConfigOptions.NONE.value, dest=f.name, help=f"Turn the {f.name} option OFF.")
                     continue
                 p.add_argument(   f"--{f.name}", action="store_true",               help=f"Turn the {f.name} option ON.")

toolchain/mfc/run/run.py

@@ -7,7 +7,7 @@
 
 from ..build   import get_targets, build, REQUIRED_TARGETS, SIMULATION
 from ..printer import cons
-from ..state   import ARG, ARGS, CFG
+from ..state   import ARG, ARGS, CFG, gpuConfigOptions
 from ..common  import MFCException, isspace, file_read, does_command_exist
 from ..common  import MFC_TEMPLATE_DIR, file_write, system, MFC_ROOT_DIR
 from ..common  import format_list_to_string, file_dump_yaml
@@ -99,6 +99,20 @@ def __generate_job_script(targets, case: input.MFCInputFile):
             'HIP_VISIBLE_DEVICES':  gpu_ids
         })
 
+    # Compute GPU mode booleans for templates
+    gpu_mode = ARG('gpu')
+
+    # Validate gpu_mode is one of the expected values
+    valid_gpu_modes = {e.value for e in gpuConfigOptions}
+    if gpu_mode not in valid_gpu_modes:
+        raise MFCException(
+            f"Invalid GPU mode '{gpu_mode}'. Must be one of: {', '.join(sorted(valid_gpu_modes))}"
+        )
+
+    gpu_enabled = gpu_mode != gpuConfigOptions.NONE.value
+    gpu_acc = gpu_mode == gpuConfigOptions.ACC.value
+    gpu_mp = gpu_mode == gpuConfigOptions.MP.value
+
     content = __get_template().render(
         **{**ARGS(), 'targets': targets},
         ARG=ARG,
@@ -107,7 +121,10 @@ def __generate_job_script(targets, case: input.MFCInputFile):
         MFC_ROOT_DIR=MFC_ROOT_DIR,
         SIMULATION=SIMULATION,
         qsystem=queues.get_system(),
-        profiler=shlex.join(__profiler_prepend())
+        profiler=shlex.join(__profiler_prepend()),
+        gpu_enabled=gpu_enabled,
+        gpu_acc=gpu_acc,
+        gpu_mp=gpu_mp
     )
 
     file_write(__job_script_filepath(), content)

toolchain/templates/bridges2.mako

@@ -14,7 +14,7 @@
 % if account:
 #SBATCH --account="${account}"
 % endif
-% if gpu:
+% if gpu_enabled:
 #SBATCH --gpu-bind=verbose,closest
 #SBATCH --gres=gpu:v100-16:${tasks_per_node}
 % endif
@@ -31,7 +31,7 @@ ${helpers.template_prologue()}
 
 ok ":) Loading modules:\n"
 cd "${MFC_ROOT_DIR}"
-. ./mfc.sh load -c b -m ${'g' if gpu else 'c'}
+. ./mfc.sh load -c b -m ${'g' if gpu_enabled else 'c'}
 cd - > /dev/null
 echo
 

toolchain/templates/carpenter-cray.mako

@@ -25,7 +25,7 @@ ${helpers.template_prologue()}
 
 ok ":) Loading modules:\n"
 cd "${MFC_ROOT_DIR}"
-. ./mfc.sh load -c cc -m ${'g' if gpu else 'c'}
+. ./mfc.sh load -c cc -m ${'g' if gpu_enabled else 'c'}
 cd - > /dev/null
 echo
 

toolchain/templates/carpenter.mako

@@ -25,7 +25,7 @@ ${helpers.template_prologue()}
 
 ok ":) Loading modules:\n"
 cd "${MFC_ROOT_DIR}"
-. ./mfc.sh load -c c -m ${'g' if gpu else 'c'}
+. ./mfc.sh load -c c -m ${'g' if gpu_enabled else 'c'}
 cd - > /dev/null
 echo
 

toolchain/templates/default.mako

@@ -48,7 +48,7 @@ if engine == 'batch':
                 (set -x; ${profiler}   \
                     jsrun --nrs          ${tasks_per_node*nodes} \
                         --cpu_per_rs   1                       \
-                        --gpu_per_rs   ${1 if gpu else 0}      \
+                        --gpu_per_rs   ${1 if gpu_enabled else 0}      \
                         --tasks_per_rs 1                       \
                         "${target.get_install_binpath(case)}")
             elif [ "$binary" == "srun" ]; then

toolchain/templates/delta.mako

@@ -14,7 +14,7 @@
 % if account:
 #SBATCH --account="${account}"
 % endif
-% if gpu:
+% if gpu_enabled:
 #SBATCH --gpus-per-node=${tasks_per_node}
 #SBATCH --mem=208G
 #SBATCH --gpu-bind=closest
@@ -32,7 +32,7 @@ ${helpers.template_prologue()}
 
 ok ":) Loading modules:\n"
 cd "${MFC_ROOT_DIR}"
-. ./mfc.sh load -c d -m ${'g' if gpu else 'c'}
+. ./mfc.sh load -c d -m ${'g' if gpu_enabled else 'c'}
 cd - > /dev/null
 echo
 

toolchain/templates/deltaai.mako

@@ -14,7 +14,7 @@
 % if account:
 #SBATCH --account="${account}"
 % endif
-% if gpu:
+% if gpu_enabled:
 #SBATCH --gpus-per-node=${tasks_per_node}
 #SBATCH --mem=208G
 #SBATCH --gpu-bind=closest
@@ -32,7 +32,7 @@ ${helpers.template_prologue()}
 
 ok ":) Loading modules:\n"
 cd "${MFC_ROOT_DIR}"
-. ./mfc.sh load -c dai -m ${'g' if gpu else 'c'}
+. ./mfc.sh load -c dai -m ${'g' if gpu_enabled else 'c'}
 cd - > /dev/null
 echo
 

toolchain/templates/frontier.mako

@@ -10,7 +10,7 @@
 #SBATCH --time=${walltime}
 #SBATCH --cpus-per-task=7
 #SBATCH -C nvme
-% if gpu != 'no':
+% if gpu_enabled:
 #SBATCH --gpus-per-task=1
 #SBATCH --gpu-bind=closest
 % endif
@@ -34,18 +34,18 @@ ${helpers.template_prologue()}
 ok ":) Loading modules:\n"
 cd "${MFC_ROOT_DIR}"
 % if engine == 'batch':
-. ./mfc.sh load -c f -m ${'g' if gpu != 'no' else 'c'}
+. ./mfc.sh load -c f -m ${'g' if gpu_enabled else 'c'}
 % endif
 cd - > /dev/null
 echo
 
-% if gpu != 'no':
+% if gpu_enabled:
     export MPICH_GPU_SUPPORT_ENABLED=1
 % else:
     export MPICH_GPU_SUPPORT_ENABLED=0
 % endif
 
-%if unified:
+% if unified:
     export CRAY_ACC_USE_UNIFIED_MEM=1
 % endif
 
@@ -66,7 +66,7 @@ ulimit -s unlimited
         % if engine == 'interactive':
                 --unbuffered --nodes ${nodes} --ntasks-per-node ${tasks_per_node} \
                 --cpus-per-task 7                                    \
-            % if gpu != 'no':
+            % if gpu_enabled:
                 --gpus-per-task 1 --gpu-bind closest                 \
             % endif
             ${profiler} "${target.get_install_binpath(case)}")

toolchain/templates/hipergator.mako

@@ -9,7 +9,7 @@
 #SBATCH --output="${name}.out"
 #SBATCH --time=${walltime}
 #SBATCH --cpus-per-task=7
-% if gpu:
+% if gpu_enabled:
 #SBATCH --gpus-per-task=1
 #SBATCH --gpu-bind=closest
 % endif
@@ -35,7 +35,7 @@ ${helpers.template_prologue()}
 ok ":) Loading modules:\n"
 cd "${MFC_ROOT_DIR}"
 % if engine == 'batch':
-. ./mfc.sh load -c h -m ${'g' if gpu else 'c'}
+. ./mfc.sh load -c h -m ${'g' if gpu_enabled else 'c'}
 % endif
 cd - > /dev/null
 echo