Generate CSEL node coditions

Author: hzhou

maint/gen_coll.py

@@ -16,7 +16,7 @@ def main():
     func_list = load_C_func_list(binding_dir, silent=True)
 
     # Loading coll_algorithms.txt. It sets -
-    #   - G.restrictions: a list of boolean conditions that can be used as restrictions
+    #   - G.conditions: a list of conditions that can be used as restrictions and in JSON tuning files
     #   - G.algos:        a two level array: [func-commkind][algo]
     load_coll_algos("src/mpi/coll/coll_algorithms.txt")
     # Prepare a one level algo array for conveninece -
@@ -47,12 +47,13 @@ def main():
     for a in G.coll_names:
         add_sched_auto_prototypes(a)
 
-    # initialize MPIR_Coll_algo_table[algo_id] -> algo_fn
+    # initialize MPIR_Coll_algo_table and MPIR_Coll_algo_names
     dump_MPII_Coll_algo_init()
-    # initialize MPIR_Coll_cvar_table
-    dump_MPII_Coll_cvar_init()
-    # create csel container from parsing json
-    dump_MPII_Create_container()
+    # initialize MPIR_Coll_cvar_table and MPIR_Coll_type_names
+    dump_MPII_Coll_type_init()
+    # parsing routines for loading JSONs
+    dump_MPII_Csel_parse_container()
+    dump_MPII_Csel_parse_operator()
     # routines for checking algorithm CVARs
     dump_MPIR_Coll_cvar_to_algo_id()
     dump_MPIR_Coll_init_algo_container()
@@ -62,6 +63,8 @@ def main():
     dump_MPIR_Csel_coll_type_e()
     # enum for algorithm id
     dump_MPIR_Csel_container_type_e()
+    # enum CSEL_NODE_TYPE
+    dump_MPIR_Csel_node_type_e()
     # algorithm container struct
     dump_MPII_Csel_container()
 
@@ -162,19 +165,25 @@ def add_sched_auto_prototypes(coll_name):
     if not re.match(r'(scan|exscan|neighbor_)', coll_name):
         add_prototype("int MPIR_I%s_inter_sched_auto(%s)" % (coll_name, params))
 
-def dump_MPII_Coll_cvar_init():
+def dump_MPII_Coll_type_init():
     G.out.append("")
-    decl = "void MPII_Coll_cvar_init(void)"
+    decl = "void MPII_Coll_type_init(void)"
     add_prototype(decl)
     G.out.append(decl)
     dump_open('{')
     for a in G.coll_names:
         for is_blocking in (True, False):
-            G.out.append("MPIR_Coll_cvar_table[%s * 2] = %s;" % (coll_type(a, is_blocking), cvar_name(a, is_blocking, "intra")))
-            if not re.match(r'(scan|exscan|neighbor_)', a):
-                G.out.append("MPIR_Coll_cvar_table[%s * 2 + 1] = %s;" % (coll_type(a, is_blocking), cvar_name(a, is_blocking, "inter")))
-            else:
-                G.out.append("MPIR_Coll_cvar_table[%s * 2 + 1] = 0;" % (coll_type(a, is_blocking)))
+            for commkind in ("intra", "inter"):
+                if commkind == "inter" and re.match(r'(scan|exscan|neighbor_)', a):
+                    # CVARs for these inter-collective does not exist
+                    G.out.append("MPIR_Coll_cvar_table[%s] = 0;" % (coll_type(a, is_blocking, commkind)))
+                else:
+                    G.out.append("MPIR_Coll_cvar_table[%s] = %s;" % (coll_type(a, is_blocking, commkind), cvar_name(a, is_blocking, commkind)))
+
+    for a in G.coll_names:
+        for is_blocking in (True, False):
+            for commkind in ("intra", "inter"):
+                G.out.append("MPIR_Coll_type_names[%s] = \"%s\";" % (coll_type(a, is_blocking, commkind), a + '-' + commkind))
     dump_close('}')
 
 def dump_MPII_Coll_algo_init():
@@ -189,10 +198,16 @@ def dump_MPII_Coll_algo_init():
         if a['func-commkind'] != 'general':
             algo_funcname += "_cnt"
         G.out.append("MPIR_Coll_algo_table[%s] = %s;" % (idx, algo_funcname))
+    for a in G.algo_list:
+        algo_funcname = get_algo_funcname(a)
+        idx = algo_id(algo_funcname)
+        if a['func-commkind'] != 'general':
+            algo_funcname += "_cnt"
+        G.out.append("MPIR_Coll_algo_names[%s] = \"%s\";" % (idx, algo_funcname))
     dump_close('}')
 
 
-def dump_MPII_Create_container():
+def dump_MPII_Csel_parse_container():
     G.out.append("")
     def dump_json_foreach_open():
         dump_open("json_object_object_foreach(obj, key, val) {")
@@ -243,7 +258,7 @@ def dump_parse_params():
                 G.out.append("    break;")
         dump_close('}') # switch
 
-    decl = "void *MPII_Create_container(struct json_object *obj)"
+    decl = "void *MPII_Csel_parse_container(struct json_object *obj)"
     add_prototype(decl)
     G.out.append(decl)
     dump_open('{')
@@ -256,6 +271,44 @@ def dump_parse_params():
     G.out.append("return (void *) cnt;")
     dump_close('}')
 
+def dump_MPII_Csel_parse_operator():
+    decl = "int MPII_Csel_parse_operator(const char *ckey, MPIR_Csel_node_s *csel_node)"
+    add_prototype(decl)
+    G.out.append(decl)
+    dump_open('{')
+    dump_open("if (ckey[0] == '!') {")
+    G.out.append("csel_node->condition.negate = true;")
+    G.out.append("ckey++;")
+    dump_else()
+    G.out.append("csel_node->condition.negate = false;")
+    dump_close('}')
+
+    if_clase = "if"
+    for a in G.conditions:
+        cond = a
+        has_thresh = False
+        if RE.match(r'(.+)\(thresh\)', a):
+            cond = RE.m.group(1)
+            has_thresh = True
+        n = len(cond)
+        if has_thresh:
+            G.out.append("%s (strncmp(ckey, \"%s\", %d) == 0) {" % (if_clase, cond, n))
+            G.out.append("    csel_node->type = %s;" % condition_id(cond))
+            G.out.append("    MPIR_Assert(ckey[%d] == '(');" % n)
+            G.out.append("    csel_node->condition.thresh = atoi(ckey + %d);" % (n + 1))
+        else:
+            G.out.append("%s (strcmp(ckey, \"%s\") == 0) {" % (if_clase, cond))
+            G.out.append("    csel_node->type = %s;" % condition_id(cond))
+        if_clase = "} else if"
+    G.out.append("} else {")
+    G.out.append("    MPIR_Assert(0);")
+    G.out.append("    return MPI_ERR_OTHER;")
+    G.out.append("}")
+
+    G.out.append("")
+    G.out.append("return MPI_SUCCESS;")
+    dump_close('}')
+
 def dump_MPIR_Coll_cvar_to_algo_id():
     G.out.append("")
     def dump_cvar_cases(name, commkind):
@@ -330,19 +383,17 @@ def dump_MPIR_Coll_init_algo_container():
 def dump_MPIR_Coll_check_algo_restriction():
     G.out.append("")
     def dump_check_restriction(restriction):
-        u = "coll_sig->%s" % coll_name
         r = restriction
         negate = False
         if restriction.startswith('!'):
             r = restriction[1:]
             negate = True
+        if RE.match(r'.*\(.*\)', r):
+            raise Exception("Threshold condition %s cannot be used as a restriction" % r)
 
         cond = None
-        if r in G.restrictions:
-            if G.restrictions[r].startswith('MPIR_COLL_ATTR__'):
-                cond = "(coll_sig->attr & %s)" % G.restrictions[r]
-            else:
-                cond = "%s(coll_sig)" % G.restrictions[r]
+        if r in G.conditions:
+            cond = "%s(coll_sig)" % G.conditions[r]
         else:
             raise Exception("Restriction %s not listed" % restriction)
 
@@ -373,8 +424,9 @@ def dump_MPIR_Csel_coll_type_e():
     G.out2.append("typedef enum {")
     for a in G.coll_names:
         for is_blocking in (True, False):
-            G.out2.append("    %s," % coll_type(a, is_blocking))
-    G.out2.append("    %s" % coll_type("END", True))
+            for commkind in ("intra", "inter"):
+                G.out2.append("    %s," % coll_type(a, is_blocking, commkind))
+    G.out2.append("    %s" % coll_type_END())
     G.out2.append("} MPIR_Csel_coll_type_e;")
 
 def dump_MPIR_Csel_container_type_e():
@@ -383,9 +435,27 @@ def dump_MPIR_Csel_container_type_e():
     for a in G.algo_list:
         algo_funcname = get_algo_funcname(a)
         G.out2.append("    %s," % algo_id(algo_funcname))
-    G.out2.append("    %s" % algo_id("Algorithm_count"))
+    G.out2.append("    %s" % algo_id_END())
     G.out2.append("} MPIR_Csel_container_type_e;")
 
+def dump_MPIR_Csel_node_type_e():
+    G.out2.append("")
+    G.out2.append("typedef enum {")
+    for a in G.conditions:
+        G.out2.append("    %s," % condition_id(a))
+    G.out2.append("    CSEL_NODE_TYPE__OPERATOR__ANY,")
+    G.out2.append("    CSEL_NODE_TYPE__CONTAINER,")
+    G.out2.append("} MPIR_Csel_node_type_e;")
+
+def dump_MPIR_Csel_node_s():
+    G.out2.append("")
+    G.out2.append("typedef struct MPIR_Csel_node {")
+    G.out2.append("    MPIR_Csel_node_type_e type;")
+    G.out2.append("    MPI_Aint thresh;")
+    G.out2.append("    struct MPIR_Csel_node *success;")
+    G.out2.append("    struct MPIR_Csel_node *failure;")
+    G.out2.append("} MPIR_Csel_node_s;")
+
 def dump_MPII_Csel_container():
     G.out2.append("")
     def dump_algo_params():
@@ -414,19 +484,18 @@ def add_prototype(l):
 
 def load_coll_algos(algo_txt):
     G.algos = {}
-    G.restrictions = {}
+    G.conditions = {}
     with open(algo_txt) as In:
         (func_commkind, algo_list, algo) = (None, None, None)
         for line in In:
-            if RE.match(r'(\w+-(intra|inter)|general):', line):
+            if RE.match(r'(\w+-(intra|inter)|general|conditions):', line):
                 func_commkind = RE.m.group(1)
-                algo_list = []
-                G.algos[func_commkind] = algo_list
-            elif RE.match(r'restrictions:', line):
-                func_commkind = "restrictions"
-            elif func_commkind == "restrictions":
+                if func_commkind != "conditions":
+                    algo_list = []
+                    G.algos[func_commkind] = algo_list
+            elif func_commkind == "conditions":
                 if RE.match(r'\s+([\w-]+):\s*(\w+)', line):
-                    G.restrictions[RE.m.group(1)] = RE.m.group(2)
+                    G.conditions[RE.m.group(1)] = RE.m.group(2)
             elif func_commkind:
                 if RE.match(r'\s+(\w+)\s*$', line):
                     algo = {"name": RE.m.group(1), "func-commkind": func_commkind}
@@ -462,6 +531,8 @@ def dump_coll_impl(name, blocking_type):
         G.out.append("coll_sig.is_persistent = false;")
     G.out.append("coll_sig.sched = NULL;")
 
+    G.out.append("memset(&coll_sig.cache, sizeof(coll_sig.cache), 0);");
+
     phash = {}
     for p in func['parameters']:
         if p['name'] == 'comm':
@@ -470,42 +541,6 @@ def dump_coll_impl(name, blocking_type):
             phash[p['name']] = 1
             G.out.append("coll_sig.u.%s.%s = %s;" % (name, p['name'], p['name']))
 
-    # msg_size
-    def init_msg_size(count, datatype):
-        G.out.append("MPIR_Datatype_get_size_macro(%s, coll_sig.msg_size);" % datatype)
-        G.out.append("coll_sig.msg_size *= %s;" % count)
-
-    if 'count' in phash and 'datatype' in phash:
-        init_msg_size('count', 'datatype')
-    elif 'recvcount' in phash and 'recvtype' in phash:
-        init_msg_size('recvcount', 'recvtype')
-    elif 'recvcount' in phash and 'datatype' in phash:
-        init_msg_size('recvcount', 'datatype')
-    elif 'recvcounts' in phash and 'recvtype' in phash:
-        # allgatherv, alltoallv - should we use total/max message size, or just skip?
-        init_msg_size('recvcounts[0]', 'recvtype')
-    elif 'recvcounts' in phash and 'recvtypes' in phash:
-        # alltoallw - should we use total/max message size, or just skip?
-        init_msg_size('recvcounts[0]', 'recvtypes[0]')
-    else:
-        raise Exception("init coll_sig: unhandled coll type")
-
-
-    G.out.append("if (MPL_is_pof2(comm_ptr->local_size)) coll_sig.attr |= MPIR_COLL_ATTR__pof2;")
-
-    if 'sendbuf' in phash:
-        G.out.append("if (sendbuf == MPI_IN_PLACE) coll_sig.attr |= MPIR_COLL_ATTR__inplace;")
-
-    if 'op' in phash:
-        G.out.append("if (HANDLE_IS_BUILTIN(op)) {")
-        G.out.append("    coll_sig.attr |= MPIR_COLL_ATTR__builtin_op | MPIR_COLL_ATTR__commutative;")
-        G.out.append("} else if (MPIR_Op_is_commutative(op)) {")
-        G.out.append("    coll_sig.attr |= MPIR_COLL_ATTR__commutative;")
-        G.out.append("}")
-
-    G.out.append("MPIR_Init_coll_sig(&coll_sig);")
-    G.out.append("MPID_Init_coll_sig(&coll_sig);")
-
     # Call csel
     G.out.append("")
     G.out.append("mpi_errno = MPIR_Coll_composition_auto(&coll_sig);")
@@ -701,12 +736,15 @@ def get_func_params(func, name, blocking_type):
 
     return ', '.join(params)
 
-def coll_type(coll_name, is_blocking):
+def coll_type(coll_name, is_blocking, commkind):
     prefix = "MPIR_CSEL_COLL_TYPE"
     if is_blocking:
-        return "%s__%s" % (prefix, coll_name.upper())
+        return "%s__%s_%s" % (prefix, commkind.upper(), coll_name.upper())
     else:
-        return "%s__I%s" % (prefix, coll_name.upper())
+        return "%s__%s_I%s" % (prefix, commkind.upper(), coll_name.upper())
+
+def coll_type_END():
+    return "MPIR_CSEL_COLL_TYPE__END"
 
 def cvar_name(coll_name, is_blocking, commkind):
     if is_blocking:
@@ -722,6 +760,15 @@ def algo_id(algo_funcname):
     else:
         return "%s__%s" % (prefix, algo_funcname)
 
+def algo_id_END():
+    return "MPII_CSEL_CONTAINER_TYPE__ALGORITHM__END"
+
+def condition_id(name):
+    prefix = "CSEL_NODE_TYPE__OPERATOR__"
+    a = re.sub(r'-', '_', name)
+    a = re.sub(r'\(thresh\)$', '', a)
+    return prefix + a
+
 def algo_struct_name(algo):
     algo_funcname = get_algo_funcname(algo)
     struct_name = re.sub(r'MPIR_', '', algo_funcname).lower()

src/include/mpir_coll.h

@@ -44,9 +44,10 @@ struct MPIR_Csel_coll_sig {
     void *sched;
     enum MPIR_sched_type sched_type;
     bool is_persistent;
-    /* derived info to assist algorithm selection */
-    MPI_Aint msg_size;
-    uint64_t attr;
+
+    struct {
+        bool is_gpu;
+    } cache;
 
     union {
         struct {