Speed improvements

Author: tanner-andrulis

pytimeloop/fastfusion/filter_mappings.py

@@ -73,9 +73,8 @@ def get_ffmt_tag_mha(
         to_try += [(other_ranks[4:], (3, 4))]
         tags.append("FFMT_LAST")
     else: # Middle Einsum in a chain
-        if einsum_id == "AV":
-            a, b = b, a
-            other_ranks[-2], other_ranks[-1] = other_ranks[-1], other_ranks[-2]
+        a, b = b, a
+        other_ranks[-2], other_ranks[-1] = other_ranks[-1], other_ranks[-2]
         to_try += [(other_ranks[:4], (3, 4))]
         tags.append("FFMT_MIDDLE")
 

pytimeloop/fastfusion/mapper/per_einsum_mapper_snowcat.py

@@ -6,12 +6,15 @@
 
 from combinatorics.dependent_product import dependent_product
 from combinatorics.splitter import split_dependent_product
+import pandas as pd
 
 from pytimeloop.fastfusion.fastmodel import compile_mapping
 from pytimeloop.fastfusion.mapper.constraints import *
 from pytimeloop.fastfusion.mapper.per_einsum_mapper import explore_tile_shape, process_result, get_hardware_levels
 from pytimeloop.fastfusion.mapper.per_einsum_subspaces.snowcat import make_subspaces
 from pytimeloop.fastfusion.mapper.per_einsum_subspaces.snowcat_ffmt import make_ffmt_subspaces
+from pytimeloop.fastfusion.pareto import Pareto, makepareto
+from pytimeloop.fastfusion.sim import SIM
 from pytimeloop.fastfusion.util import parallel
 from pytimeloop.looptree.equivalent_ranks import EquivalentGroups
 from pytimeloop.looptree.mapping_utilities import get_intermediate_tensors
@@ -159,7 +162,7 @@ def per_worker_exploration(*args):
                     output_tensors=output_tensors,
                     tag_with=tag_with,
                 )
-        return einsum_id, result
+        return einsum_id, {k: makepareto(pd.DataFrame(v).fillna(0)) for k, v in result.items()}
 
 
     # # for pm in partial_mappings:
@@ -211,13 +214,21 @@ def per_einsum_mapper_snowcat(
             dataflow_constraint=dataflow_constraint,
             metrics=metrics,
             tag_with=tag_with,
-        ) 
+        )
     )
     data = {einsum_id: defaultdict(list) for einsum_id in einsums_to_explore}
 
-    for einsum_id, result in parallel(jobs, return_as="generator_unordered", pbar="Generating Single-Einsum Mappings"):
+    for einsum_id, result in parallel(jobs, pbar="Generating Single-Einsum Mappings"):
         d = data[einsum_id]
         for k, v in result.items():
-            d[k[0]] += v
+            d[k[0]].append(v)
 
-    return data
+    def makesim(einsum_id, tiling, data):
+        return einsum_id, SIM(tiling, Pareto(pd.concat(data).fillna(0), skip_pareto=len(data) == 1))
+            
+    data2 = defaultdict(list)
+    jobs = [delayed(makesim)(einsum_id, tiling, data) for einsum_id, tilings in data.items() for tiling, data in tilings.items()]
+    for einsum_id, sim in parallel(jobs, pbar="Generating SIMs"):
+        data2[einsum_id].append(sim)
+    
+    return data2

pytimeloop/fastfusion/mapper/per_einsum_subspaces/snowcat.py

@@ -42,6 +42,8 @@ def fused_temporal_fors(mapping, unfused_tensors):
 
     def glb_storage(mapping, unfused_tensors):
         glb_fused_tensors = intermediate_tensors - unfused_tensors
+        last_fused_loop_idx = get_last_fused_loop_idx(mapping, intermediate_tensors)
+        # last_fused_loop_idx = None
         for partial_mapping in make_storage(mapping,
                                             level=1,
                                             must_retain_tensors=intermediate_tensors,
@@ -51,9 +53,8 @@ def glb_storage(mapping, unfused_tensors):
                                             explore_uneven=True,
                                             add_split_at_tensors=glb_fused_tensors,
                                             must_have_terminal_storage=False,
-                                            apply_lrp_after_loop_idx=None):
-            last_fused_loop_idx = get_last_fused_loop_idx(partial_mapping,
-                                                          intermediate_tensors)
+                                            apply_lrp_after_loop_idx=last_fused_loop_idx):
+            last_fused_loop_idx = get_last_fused_loop_idx(partial_mapping, intermediate_tensors)
             yield from make_storage(partial_mapping,
                                     level=1,
                                     must_retain_tensors=tensors - intermediate_tensors,

pytimeloop/fastfusion/mapper/per_einsum_subspaces/subspaces/storage.py

@@ -43,8 +43,8 @@ def make_storage(
 
                 retained_tensors = must_retain_tensors | set(also_retained_tensors)
                 mapping.add_storage(level, retained_tensors)
-                if any(t in add_split_at_tensors for t in retained_tensors):
-                    mapping.add_sequential()
+                # if any(t in add_split_at_tensors for t in retained_tensors):
+                #     mapping.add_sequential()
 
                 if return_retained_tensors:
                     yield mapping, retained_tensors
@@ -115,8 +115,8 @@ def make_storage(
         for idx, tensors_at_idx in sorted(idx_to_tensors.items(),
                                           key=lambda pair: pair[0],
                                           reverse=True):
-            if any(t in add_split_at_tensors for t in tensors_at_idx):
-                mapping.add_sequential(idx)
+            # if any(t in add_split_at_tensors for t in tensors_at_idx):
+            #     mapping.add_sequential(idx)
             mapping.add_storage(level, tensors_at_idx, idx)
             # Check for any irrelevant loops above the backing storage for a tensor
             for t in tensors_at_idx:

pytimeloop/fastfusion/mapper/process_results.py

@@ -92,7 +92,7 @@ def record_storage(node):
                 tensor_id_to_name[dspace],
                 len(full_tiling),
                 node["target"],
-                result.occupancy[(node["target"], dspace)],
+                int(result.occupancy[(node["target"], dspace)]),
             )
             all_storages.append(storage)
             if storage.tensor_id in intermediates_to_find:
@@ -178,7 +178,7 @@ def record_loop(node):
     for r in all_storages:
         r: TensorStorage
         if r not in backing_storages:
-            key = nameloop2col(r.backer_id, r.above_loop_index)
+            key = nameloop2col(r.backer_id, min(r.above_loop_index, n_fused_loops))
             results.setdefault(key, 0)
             results[key] += r.tile_size
         # logstring.append(f"{r}")

pytimeloop/fastfusion/mapper/simexplore.py

@@ -6,10 +6,9 @@
 
 import pandas as pd
 from joblib import delayed
-from tqdm import tqdm
 
 from pytimeloop.fastfusion.sim import SIM
-from pytimeloop.fastfusion.pareto import Pareto, check_correctness
+from pytimeloop.fastfusion.pareto import Pareto
 from pytimeloop.fastfusion.util import parallel, debugger_active
 
 
@@ -28,6 +27,8 @@ def mapping2sims(einsum_to_result: Mapping):
     for einsum_id, compat_dict in einsum_to_result.items():
         r[einsum_id] = [paretofy(k, v) for k, v in compat_dict.items()]
     return list(r.values())
+def paretofy(k, v):
+    return SIM(k, Pareto(pd.DataFrame(v).fillna(0)))
 
 
 prev_time = 0
@@ -84,19 +85,20 @@ def consolidate(
 
 
 def fuse_sims(
-    sims: list[SIM],
+    sims: dict[str, list[SIM]],
     resource2capacity: dict = None,
     return_nmappings_nbuckets: bool = False,
-    pre_filter: bool = True
 ):
     nmappings = []
     nbuckets = []
-    resource2capacity = resource2capacity or {}
-    sims = [s for s in sims]
 
-    for i, s in enumerate(sims):
-        print(f'SIM {i} tensors: {s[0].tensor_names}')
-        
+    sims = list(sims.items())
+    
+    for einsum_id, s in sims:
+        print(f'SIM {einsum_id} tensors: {s[0].tensor_names}')
+
+    # TODO: Lookahead by one SIM. If we're going to create a tiling that has loops
+    # that are not in the ranks of the next SIM, we should drop that tiling.
     # if pre_filter:
     #     for i in range(len(sims) - 1):
     #         left, right = sims[i], sims[i + 1]
@@ -109,29 +111,30 @@ def fuse_sims(
     #         print(f'Filtered {len(left)} -> {len(left2)} SIMs from Einsum {i}')
     #         print(f'Filtered {len(right)} -> {len(right2)} SIMs from Einsum {i + 1}')
 
-    left = sims.pop(0)
-
     init_print_time()
-    
-    if not sims:
-        sims = copy.deepcopy(sims)
+    if len(sims) == 1:
+        left = copy.deepcopy(sims[0][1])
+        sims = []
         left = consolidate(
             x=left,
             left=True,
             live_tensors=set(),
             resource2capacity=resource2capacity,
             shared_tensors=set(),
         )
-        
-    # TODO: Lookahead by one SIM. If we're going to create a tiling that has loops
-    # that are not in the ranks of the next SIM, we should drop that tiling.
 
+    n_iterations = 0
+    total_iterations = len(sims)
+    left_einsum, left = sims.pop(0)
     while sims:
+        n_iterations += 1
         nbuckets.append(len(left))
         nmappings.append(sum(len(s.mapping.data) for s in left))
 
-        right = sims.pop(0)
-        live_tensors = set.union(set(), *[s[0].tensor_names for s in sims if s])
+        right_einsum, right = sims.pop(0)
+        print(f'\nEinsum {right_einsum} ({n_iterations}/{total_iterations})')
+        
+        live_tensors = set.union(set(), *[s[0].tensor_names for _, s in sims if s])
         shared_tensors = set(left[0].tensor_names) & set(right[0].tensor_names)
 
         right_tensors = right[0].tensor_names
@@ -144,23 +147,26 @@ def fuse_sims(
             shared_tensors=shared_tensors,
         )
 
-        left = SIM.combine_combineable(left, live_tensors | right_tensors)
-        right = SIM.combine_combineable(right, live_tensors | left_tensors)
-
-        print_time("Combining")
-
         left = sorted(left, key=lambda x: len(x.mapping.data), reverse=True)
         right = sorted(right, key=lambda x: len(x.mapping.data), reverse=True)
+        lr = parallel(
+            [delayed(lambda l: l.left_consolidate(live_tensors, resource2capacity, shared_tensors))(l) for l in left] + 
+            [delayed(lambda l: l.consolidate(live_tensors, resource2capacity, shared_tensors))(l) for l in right],
+            pbar=f"Consolidating {left_einsum} <--> {right_einsum}",
+        )
+        left, right = lr[:len(left)], lr[len(left):]
+        print_time(f"Consolidating")
 
-        left = parallel([delayed(lambda l: l.left_consolidate(live_tensors, resource2capacity, shared_tensors))(l) for l in left], pbar="Left consolidate")
-        right = parallel([delayed(lambda l: l.consolidate(live_tensors, resource2capacity, shared_tensors))(l) for l in right], pbar="Right consolidate")
+        left = SIM.combine_combineable(left, live_tensors | right_tensors)
+        right = SIM.combine_combineable(right, live_tensors | left_tensors)
+        print_time(f"Combining")
 
-        print_time("Consolidating")
+        # left = parallel([delayed(lambda l: l.left_consolidate(live_tensors, resource2capacity, shared_tensors))(l) for l in left], pbar="Left consolidate")
+        # right = parallel([delayed(lambda l: l.consolidate(live_tensors, resource2capacity, shared_tensors))(l) for l in right], pbar="Right consolidate")
 
         # Group left and right into buckets
         right = SIM.group_right(right, left_tensors)
         left = SIM.group_left(left, right_tensors)
-
         print_time("Grouping")
 
         for v in list(left.values()) + list(right.values()):
@@ -196,28 +202,27 @@ def fuse_sims(
 
         if DELAY_MERGE:
             combined = sorted(combined, key=lambda x: x._predicted_mappings, reverse=True)
-            for c, mapping in zip(combined, parallel([c.mapping for c in combined], pbar='Merging mappings')):
+            for c, mapping in zip(combined, parallel([c.mapping for c in combined], pbar=f'Merging mappings {left_einsum} <--> {right_einsum}')):
                 c.mapping = mapping
 
         print_time("Mapping merging")
 
-        print(
-            f"\tCombining {sum(len(s) for s in left)}({len(left)}) x {sum(len(s) for s in right)}({len(right)}) -> {len(combined)}"
-        )
+        print(f"\tCombining {sum(len(s) for s in left)}({len(left)}) x {sum(len(s) for s in right)}({len(right)}) -> {len(combined)}")
         # if DO_PRINT:
         #     for k in right:
         #         if k not in left:
         #             for b in right[k]:
         #                 print(f"\tREVERSE: No match for {b.tiling}")
 
         left = combined
-        print(f"Number of buckets: {len(left)}")
+        left_einsum = right_einsum
+        print(f"\tNumber of buckets for Einsum {left_einsum}: {len(left)}")
         n_mappings = sum(len(s.mapping.data) for s in left)
-        print(f"Number of mappings: {n_mappings}")
-        print(f"Mappings per bucket: {n_mappings / len(left)}")
+        print(f"\tNumber of mappings for Einsum {left_einsum}: {n_mappings}")
+        print(f"\tMappings per bucket for Einsum {left_einsum}: {n_mappings / len(left)}")
 
     for s in left:
-        s.left_consolidate(set(), resource2capacity)
+        s.left_consolidate(None, resource2capacity)
     s_final = SIM.combine_combineable(left, set())[0]
     data = s_final.mapping.data
     # check_correctness(data, set())
@@ -227,7 +232,3 @@ def fuse_sims(
     if return_nmappings_nbuckets:
         return data, nmappings, nbuckets
     return data
-
-
-def paretofy(k, v):
-    return SIM(k, Pareto(pd.DataFrame(v).fillna(0)))