[UR] Generate structure type along with function type

scripts/generate_ids.py

@@ -9,6 +9,8 @@
 from fileinput import FileInput
 import util
 import yaml
+import re
+import copy
 
 ENUM_NAME = '$x_function_t'
 
@@ -18,36 +20,99 @@ class quoted(str):
 def quoted_presenter(dumper, data):
     return dumper.represent_scalar('tag:yaml.org,2002:str', data, style='"')
 
+def get_registry_header():
+    return {'type': 'header', 'desc': quoted('Intel $OneApi Unified Runtime function registry'), 'ordinal': quoted(-1)}
+
 def write_registry(data, path):
-    header = {'type': 'header', 'desc': quoted('Intel $OneApi Unified Runtime function registry'), 'ordinal': quoted(-1)}
-    data.insert(0, header)
     with open(path, 'w') as fout:
         yaml.add_representer(quoted, quoted_presenter)
         yaml.dump_all(data, fout,
             default_flow_style=False,
             sort_keys=False,
             explicit_start=True)
 
+def find_type_in_specs(specs, type):
+    return [obj for s in specs for obj in s['objects'] if obj['name'] == type][0]
+
+def get_max_enum(enum):
+    return int(max(enum['etors'], key=lambda x : int(x['value']))['value'])
+
+def copy_and_strip_prefix_from_enums(enum, prefix):
+    cpy = copy.deepcopy(enum)
+    for etor in cpy['etors']:
+        etor['name'] = etor['name'][len(prefix):]
+    return cpy
+
+
+def generate_function_type(specs, meta, update_fn) -> dict:
+    existing_function_type = find_type_in_specs(specs, '$x_function_t')
+    existing_etors = {etor['name'] : etor['value'] for etor in existing_function_type['etors']}
+    max_etor = get_max_enum(existing_function_type)
+    functions = [obj['class'][len('$x'):] + obj['name'] for s in specs for obj in s['objects'] if obj['type'] == 'function']
+    registry = list()
+    for fname in functions:
+        etor_name = "$X_FUNCTION_" + util.to_snake_case(fname).upper()
+        id = existing_etors.get(etor_name)
+        if id is None:
+            max_etor += 1
+            id = max_etor
+        registry.append({
+            'name': etor_name, 
+            'desc': f'Enumerator for $x{fname}', 
+            'value': str(id)}
+            )
+    registry = sorted(registry, key=lambda x : int(x['value']))
+    existing_function_type['etors'] = registry
+    update_fn(existing_function_type, meta)
+
+    ## create a copy to write back to registry.yml
+    return copy_and_strip_prefix_from_enums(existing_function_type, '$X_FUNCTION_')
+
+
+def generate_structure_type(specs, meta, refresh_fn) -> dict:
+    structure_type = find_type_in_specs(specs, '$x_structure_type_t')
+    extended_structs = [obj for s in specs for obj in s['objects'] if re.match(r"struct|union", obj['type']) and 'base' in obj]
+    max_enum = get_max_enum(structure_type)
+
+    structure_type_etors = list()
+    for struct in extended_structs:
+        # skip experimental enumerations
+        if struct['name'].startswith('$x_exp_'):
+            continue
+
+        etor = [mem for mem in struct['members'] if mem['name'] == 'stype'][0]['init']
+
+        # try and match the etor
+        matched_etor = [e for e in structure_type['etors'] if e['name'] == etor]
+
+        out_etor = {
+            'name': etor,
+            'desc': struct['name']
+        }
+
+        # if no match exists we assign it a new value
+        if len(matched_etor) == 0:
+            max_enum += 1
+            out_etor['value'] = str(max_enum)
+        else:
+            out_etor['value'] = matched_etor[0]['value']
+
+        structure_type_etors.append(out_etor)
+
+    structure_type_etors = sorted(structure_type_etors, key = lambda x : int(x['value']))
+    structure_type['etors'] = structure_type_etors
+    refresh_fn(structure_type, meta)
+
+    ## create a copy to write back to registry.yml
+    return copy_and_strip_prefix_from_enums(structure_type, '$X_STRUCTURE_TYPE_')
 
-def generate_registry(path, specs):
+def generate_registry(path, specs, meta, update_fn):
     try:
-        contents = list(util.yamlRead(path))
-        existing_registry = contents[1]['etors']
-        existing_etors = {etor["name"]: etor["value"] for etor in existing_registry}
-        max_etor = int(max(existing_registry, key = lambda x : int(x["value"]))["value"])
-        functions = [obj['class'][len('$x'):] + obj['name'] for s in specs for obj in s['objects'] if obj['type'] == 'function']
-        registry = list()
-        for fname in functions:
-            etor_name = util.to_snake_case(fname).upper()
-            id = existing_etors.get(etor_name)
-            if id is None:
-                max_etor += 1
-                id = max_etor
-            registry.append({'name': util.to_snake_case(fname).upper(), 'desc': 'Enumerator for $x'+fname, 'value': str(id)})
-        registry = sorted(registry, key=lambda x: int(x['value']))
-        wrapper = { 'name': ENUM_NAME, 'type': 'enum', 'desc': 'Defines unique stable identifiers for all functions' , 'etors': registry}
-        contents[1] =  wrapper
-        write_registry(contents[1:], path)
+        write_registry([
+            get_registry_header(), 
+            generate_function_type(specs, meta, update_fn), 
+            generate_structure_type(specs, meta, update_fn)
+            ], path)
 
     except BaseException as e:
         print("Failed to generate registry.yml... %s", e)

scripts/parse_specs.py

@@ -853,54 +853,6 @@ def sort_etors(x):
             return value
         matching_enum['etors'] = sorted(matching_enum['etors'], key=sort_etors)
 
-def _generate_structure_type_t(specs, meta, registry):
-    extended_structs = [obj for s in specs for obj in s['objects'] if re.match(r"struct|union", obj['type']) and 'base' in obj]
-    ur_structure_type_t = [obj for s in specs for obj in s['objects'] if obj['name'] == "$x_structure_type_t"][0]
-    existing_etors = ur_structure_type_t['etors']
-    max_enum = int(max(ur_structure_type_t['etors'], key= lambda x : int(x['value']))['value'])
-
-    out_etors = []
-    for struct in extended_structs:
-        # skip experimental enumerations
-        if struct['name'].startswith("$x_exp_"):
-            continue
-
-        # name of the etor
-        etor = [mem for mem in struct['members'] if mem['name'] == 'stype'][0]['init']
-
-        # try and match the etor
-        matched_etor = [e for e in existing_etors if e['name'] == etor]
-
-        # if no match exists then we have to add it
-        if len(matched_etor) == 0:
-            max_enum += 1
-            out_etors.append({
-                "name": etor,
-                "desc": struct['name'],
-                "value": str(max_enum)
-            })
-        else:
-            out_etors.append({
-                "name": etor,
-                "desc": struct['name'],
-                "value": matched_etor[0]['value']
-            })
-
-    out_etors = sorted(out_etors, key = lambda x : int(x['value']))
-    ur_structure_type_t['etors'] = out_etors
-    _refresh_enum_meta(ur_structure_type_t, meta)
-
-    ## write the result out to the yml file
-    try:
-        contents = list(util.yamlRead(registry))
-        cpy = copy.deepcopy(ur_structure_type_t)
-        for e in cpy['etors']:
-            e['name'] = e['name'][len("$X_STRUCTURE_TYPE_"):]
-        contents[2] = cpy
-        generate_ids.write_registry(contents[1:], registry)
-    except:
-        raise
-
 """
 Entry-point:
     Reads each YML file and extracts data
@@ -911,13 +863,10 @@ def parse(section, version, tags, meta, ref):
     specs = []
 
     files = util.findFiles(path, "*.yml")
-    # make sure registry is last, because it's autogenerated based on the rest of the spec
-    files = sorted(files, key=lambda f: 1 if f.endswith('registry.yml') else 0)
+    registry = [f for f in files if f.endswith('registry.yml')][0]
 
     enum_extensions = []
     for f in files:
-        if f.endswith('registry.yml'):
-            generate_ids.generate_registry(f, specs)
 
         print("Parsing %s..."%f)
         docs = util.yamlRead(f)
@@ -970,8 +919,7 @@ def parse(section, version, tags, meta, ref):
 
     specs = sorted(specs, key=lambda s: s['header']['ordinal'])
     _inline_extended_structs(specs, meta)
-    registry = [f for f in files if f.endswith('registry.yml')][0]
-    _generate_structure_type_t(specs, meta, registry)
+    generate_ids.generate_registry(registry, specs, meta, _refresh_enum_meta)
     _extend_enums(enum_extensions, specs, meta)
     _generate_extra(specs, meta)