Add EXCL_STRIDE to XML PE configs

Author: amametjanov

CIME/Tools/pelayout

@@ -70,13 +70,13 @@ def parse_command_line(args, description):
 
     parser.add_argument(
         "--format",
-        default="%4C: %6T/%6H; %6R %6P",
+        default="%4C: %6T/%6H; %6R %6P %6X",
         help="Format the PE layout items for each component (see below)",
     )
 
     parser.add_argument(
         "--header",
-        default="Comp  NTASKS  NTHRDS  ROOTPE PSTRIDE",
+        default="Comp  NTASKS  NTHRDS  ROOTPE PSTRIDE XSTRIDE",
         help="Custom header for PE layout display",
     )
 
@@ -116,26 +116,27 @@ def get_value_as_string(case, var, attribute=None, resolved=False, subgroup=None
 
 
 ###############################################################################
-def format_pelayout(comp, ntasks, nthreads, rootpe, pstride, arg_format):
+def format_pelayout(comp, ntasks, nthreads, rootpe, pstride, xstride, arg_format):
     ###############################################################################
     """
     Format the PE layout information for each component, using a default format,
     or using the arg_format input, if it exists.
     """
-    subs = {"C": comp, "T": ntasks, "H": nthreads, "R": rootpe, "P": pstride}
+    subs = {"C": comp, "T": ntasks, "H": nthreads, "R": rootpe, "P": pstride, "X": xstride}
     layout_str = re.sub(r"%([0-9]*)C", r"{C:\1}", arg_format)
     layout_str = re.sub(r"%([-+0-9]*)T", r"{T:\1}", layout_str)
     layout_str = re.sub(r"%([-+0-9]*)H", r"{H:\1}", layout_str)
     layout_str = re.sub(r"%([-+0-9]*)R", r"{R:\1}", layout_str)
     layout_str = re.sub(r"%([-+0-9]*)P", r"{P:\1}", layout_str)
+    layout_str = re.sub(r"%([-+0-9]*)X", r"{X:\1}", layout_str)
     layout_str = layout_str.format(**subs)
     return layout_str
 
 
 # End def format_pelayout
 
 ###############################################################################
-def print_pelayout(case, ntasks, nthreads, rootpes, pstrid, arg_format, header):
+def print_pelayout(case, ntasks, nthreads, rootpes, pstrid, xstrid, arg_format, header):
     ###############################################################################
     """
     Print the PE layout information for each component, using the format,
@@ -155,6 +156,7 @@ def print_pelayout(case, ntasks, nthreads, rootpes, pstrid, arg_format, header):
                 nthreads[comp],
                 rootpes[comp],
                 pstrid[comp],
+                xstrid[comp],
                 arg_format,
             )
         )
@@ -185,15 +187,17 @@ def gather_pelayout(case):
     nthreads = {}
     rootpes = {}
     pstride = {}
+    xstride = {}
     comp_classes = case.get_values("COMP_CLASSES")
 
     for comp in comp_classes:
         ntasks[comp] = int(case.get_value("NTASKS_" + comp))
         nthreads[comp] = int(case.get_value("NTHRDS_" + comp))
         rootpes[comp] = int(case.get_value("ROOTPE_" + comp))
         pstride[comp] = int(case.get_value("PSTRID_" + comp))
+        xstride[comp] = int(case.get_value("EXCL_STRIDE_" + comp))
     # End for
-    return ntasks, nthreads, rootpes, pstride
+    return ntasks, nthreads, rootpes, pstride, xstride
 
 
 # End def gather_pelayout
@@ -274,8 +278,8 @@ def _main_func(description):
         if set_ntasks is not None:
             modify_ntasks(case, int(set_ntasks))
         # End if
-        ntasks, nthreads, rootpes, pstrid = gather_pelayout(case)
-        print_pelayout(case, ntasks, nthreads, rootpes, pstrid, arg_format, header)
+        ntasks, nthreads, rootpes, pstrid, xstrid = gather_pelayout(case)
+        print_pelayout(case, ntasks, nthreads, rootpes, pstrid, xstrid, arg_format, header)
     # End with
 
 

CIME/XML/pes.py

@@ -25,6 +25,7 @@ def find_pes_layout(self, grid, compset, machine, pesize_opts="M", mpilib=None):
         opes_nthrds = {}
         opes_rootpe = {}
         opes_pstrid = {}
+        opes_excl_stride= {}
         oother_settings = {}
         other_settings = {}
         append = {}
@@ -41,6 +42,7 @@ def find_pes_layout(self, grid, compset, machine, pesize_opts="M", mpilib=None):
                 opes_nthrds,
                 opes_rootpe,
                 opes_pstrid,
+                opes_excl_stride,
                 oother_settings,
                 append,
                 ocomments,
@@ -60,6 +62,7 @@ def find_pes_layout(self, grid, compset, machine, pesize_opts="M", mpilib=None):
             pes_nthrds,
             pes_rootpe,
             pes_pstrid,
+            pes_excl_stride,
             other_settings,
             os_append,
             comments,
@@ -68,6 +71,7 @@ def find_pes_layout(self, grid, compset, machine, pesize_opts="M", mpilib=None):
         pes_nthrds.update(opes_nthrds)
         pes_rootpe.update(opes_rootpe)
         pes_pstrid.update(opes_pstrid)
+        pes_excl_stride.update(opes_excl_stride)
         other_settings.update(oother_settings)
         os_append.update(append)
         if ocomments is not None:
@@ -80,13 +84,16 @@ def find_pes_layout(self, grid, compset, machine, pesize_opts="M", mpilib=None):
                 pes_rootpe[i] = 0
             for i in iter(pes_pstrid):
                 pes_pstrid[i] = 0
+            for i in iter(pes_excl_stride):
+                pes_pstrid[i] = 0
 
         logger.info("Pes setting: grid          is {} ".format(grid))
         logger.info("Pes setting: compset       is {} ".format(compset))
         logger.info("Pes setting: tasks       is {} ".format(pes_ntasks))
         logger.info("Pes setting: threads     is {} ".format(pes_nthrds))
         logger.info("Pes setting: rootpe      is {} ".format(pes_rootpe))
         logger.info("Pes setting: pstrid      is {} ".format(pes_pstrid))
+        logger.info("Pes setting: excl_stride is {} ".format(pes_excl_stride))
         logger.info("Pes other settings: {}".format(other_settings))
         logger.info("Pes other settings append: {}".format(os_append))
         if comments is not None:
@@ -97,6 +104,7 @@ def find_pes_layout(self, grid, compset, machine, pesize_opts="M", mpilib=None):
             pes_nthrds,
             pes_rootpe,
             pes_pstrid,
+            pes_excl_stride,
             other_settings,
             os_append,
             comments,
@@ -110,7 +118,8 @@ def _find_matches(
         compset_choice = None
         pesize_choice = None
         max_points = -1
-        pes_ntasks, pes_nthrds, pes_rootpe, pes_pstrid, other_settings, append = (
+        pes_ntasks, pes_nthrds, pes_rootpe, pes_pstrid, pes_excl_stride, other_settings, append = (
+            {},
             {},
             {},
             {},
@@ -173,6 +182,11 @@ def _find_matches(
                                                 pes_pstrid[
                                                     self.name(child).upper()
                                                 ] = int(self.text(child))
+                                        elif "excl_stride" in vid:
+                                            for child in self.get_children(root=node):
+                                                pes_excl_stride[
+                                                    self.name(child).upper()
+                                                ] = int(self.text(child))
                                         # if the value is already upper case its something else we are trying to set
                                         else:
                                             other_settings[vid] = self.text(node)
@@ -231,6 +245,9 @@ def _find_matches(
                 elif "pstrid" in vid:
                     for child in self.get_children(root=node):
                         pes_pstrid[self.name(child).upper()] = int(self.text(child))
+                elif "excl_stride" in vid:
+                    for child in self.get_children(root=node):
+                        pes_excl_stride[self.name(child).upper()] = int(self.text(child))
                 # if the value is already upper case its something else we are trying to set
                 elif vid == self.name(node):
                     text = self.text(node).strip()
@@ -250,6 +267,7 @@ def _find_matches(
             pes_nthrds,
             pes_rootpe,
             pes_pstrid,
+            pes_excl_stride,
             other_settings,
             append,
             comment,

CIME/case/case.py

@@ -1204,6 +1204,7 @@ def _setup_mach_pes(self, pecount, multi_driver, ninst, machine_name, mpilib):
         pes_nthrds = {}
         pes_rootpe = {}
         pes_pstrid = {}
+        pes_excl_stride = {}
         other = {}
         append = {}
         comment = None
@@ -1233,6 +1234,7 @@ def _setup_mach_pes(self, pecount, multi_driver, ninst, machine_name, mpilib):
                 pes_nthrds,
                 pes_rootpe,
                 pes_pstrid,
+                pes_excl_stride,
                 other,
                 append,
                 comment,
@@ -1287,17 +1289,20 @@ def _setup_mach_pes(self, pecount, multi_driver, ninst, machine_name, mpilib):
             nthrds_str = "NTHRDS_{}".format(comp_class)
             rootpe_str = "ROOTPE_{}".format(comp_class)
             pstrid_str = "PSTRID_{}".format(comp_class)
+            excl_stride_str = "EXCL_STRIDE_{}".format(comp_class)
 
             ntasks = pes_ntasks[ntasks_str] if ntasks_str in pes_ntasks else 1
             nthrds = pes_nthrds[nthrds_str] if nthrds_str in pes_nthrds else 1
             rootpe = pes_rootpe[rootpe_str] if rootpe_str in pes_rootpe else 0
             pstrid = pes_pstrid[pstrid_str] if pstrid_str in pes_pstrid else 1
+            excl_stride = pes_excl_stride[excl_stride_str] if excl_stride_str in pes_excl_stride else 0
 
             totaltasks.append((ntasks + rootpe) * nthrds)
             mach_pes_obj.set_value(ntasks_str, ntasks)
             mach_pes_obj.set_value(nthrds_str, nthrds)
             mach_pes_obj.set_value(rootpe_str, rootpe)
             mach_pes_obj.set_value(pstrid_str, pstrid)
+            mach_pes_obj.set_value(excl_stride_str, excl_stride)
 
         # Make sure that every component has been accounted for
         # set, nthrds and ntasks to 1 otherwise. Also set the ninst values here.

CIME/data/config/config_headers.xml

@@ -69,6 +69,7 @@
     ROOTPE: the global mpi task of the component root task, if negative, indicates nodes rather than tasks.
     PSTRID: the stride of MPI tasks across the global set of pes (for now set to 1)
     NINST : the number of component instances (will be spread evenly across NTASKS)
+    EXCL_STRIDE: the stride of MPI tasks owned exclusively by a component. If 0, exclusivity is disabled.
 
     for example, for NTASKS = 8, NTHRDS = 2, ROOTPE = 32, NINST  = 2
     the MPI tasks would be placed starting on global pe 32 and each pe would be threaded 2-ways

doc/source/ccs/model-configuration/variables/pes.rst

@@ -326,6 +326,8 @@ following meanings:
      -  The global MPI task of the component root task; if negative, indicates nodes rather than tasks. The root processor for each component is set relative to the MPI global communicator.
    * - PSTRID
      - The stride of MPI tasks across the global set of pes (for now set to 1). This variable is currently not used and is a placeholder for future development.
+   * - EXCL_STRIDE
+     - Stride of MPI tasks owned exclusively by a component. If 0, exclusivity is disabled.
    * - NINST
      -  The number of component instances, which are spread evenly across NTASKS.
    * - COST_PER_NODE