Some refactors, tile size bugfix

Author: tanner-andrulis

pytimeloop/fastfusion/mapper/per_einsum_mapper_snowcat.py

@@ -215,7 +215,7 @@ def per_einsum_mapper_snowcat(
     )
     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 data for Einsums"):
+    for einsum_id, result in parallel(jobs, return_as="generator_unordered", pbar="Generating Single-Einsum Mappings"):
         d = data[einsum_id]
         for k, v in result.items():
             d[k[0]] += v

pytimeloop/fastfusion/mapper/simexplore.py

@@ -144,21 +144,23 @@ def fuse_sims(
             shared_tensors=shared_tensors,
         )
 
-        # right = consolidate(right, left=False, **args)
-        # left = consolidate(left, left=True, **args)
-
         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)
 
         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")
 
+        print_time("Consolidating")
+
         # 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()):
@@ -167,11 +169,6 @@ def fuse_sims(
                     if t not in live_tensors:
                         del s.tensors[t]
 
-        # left = {k: SIM.combine_combineable(v, live_tensors | right_tensors) for k, v in left.items()}
-        # right = {k: SIM.combine_combineable(v, live_tensors | left_tensors) for k, v in right.items()}
-
-        print_time("Consolidating")
-
         DO_PRINT = False
         DELAY_MERGE = not debugger_active()
 

pytimeloop/fastfusion/pareto.py

@@ -370,10 +370,10 @@ def einsum_ids(self):
         return fzs(self.data[LOGSTRING].iloc[0].keys())
 
     @staticmethod
-    def concat(paretos: list["Pareto"]) -> "Pareto":
+    def concat(paretos: list["Pareto"], skip_pareto: bool=False) -> "Pareto":
         return Pareto(
             pd.concat([p.data for p in paretos]).fillna(0),
-            skip_pareto=len(paretos) == 1,
+            skip_pareto=len(paretos) == 1 or skip_pareto,
         )
 
     def merge_next(

pytimeloop/fastfusion/plot/looptree.py

@@ -84,8 +84,8 @@ def tilings2looptree(mappings: dict[str, Tiling], stats: dict[str, Any], skip_ba
     prev_tilings = []
     root = Node()
     einsum_ids = list(mappings.keys())
-    
-    assert set(einsum_ids) == set(stats.keys())
+    if stats is not None:
+        assert set(einsum_ids) == set(stats.keys())
 
 
 
@@ -133,8 +133,8 @@ def tilings2looptree(mappings: dict[str, Tiling], stats: dict[str, Any], skip_ba
                 # TODO if tensor not in n.this_level or tensor not in backers:
                 if tensor not in n.this_level or tensor not in backers:
                     n.this_level.append(tensor)
-
-        root.add_stats(stats[einsum_id])
+        if stats is not None:
+            root.add_stats(stats[einsum_id])
         # for k, v in partial_stats[einsum_id].items():
         #     last_level.append(f"_PARTIAL {k}: {expfmt(v)}")
         prev_tilings.append(tiling)

pytimeloop/fastfusion/sim.py

@@ -78,11 +78,18 @@ class TensorStorage:
     tensor_id: str
     above_loop_index: int
     backer_id: str
+    # NOTE: Tile size is not included in hash or equality functions. This is
+    # because inter-Einsum comparisons care about the loops and locations of
+    # backing storages, and the tile sizes are derived from these. We don't want
+    # rounding errors in the tile size to effect our inter-Einsum comparisons.
     tile_size: int
     # n_repititions: int = 1
 
     def __tuple__(self):
         return (self.tensor_id, self.backer_id, self.above_loop_index, self.tile_size)
+    
+    def __hash__(self):
+        return hash((self.tensor_id, self.backer_id, self.above_loop_index))
 
     @property
     def ts(self):
@@ -127,7 +134,7 @@ def get_backing_stores(all_tensors: set["TensorStorage"]) -> list["TensorStorage
     def __eq__(self, value):
         if not isinstance(value, TensorStorage):
             return False
-        for to_check in ["tensor_id", "backer_id", "above_loop_index", "tile_size"]:
+        for to_check in ["tensor_id", "backer_id", "above_loop_index"]:#$, "tile_size"]:
             a, b = getattr(self, to_check), getattr(value, to_check)
             if a != "*" and b != "*" and a != b:
                 return False
@@ -369,7 +376,13 @@ def _group(
 
     @staticmethod
     def combine_combineable(sims: list["SIM"], live_tensors: set[str], allow_different_tilings: bool=False) -> list["SIM"]:
-        return parallel(delayed(SIM.concat)(s, allow_different_tilings) for s in SIM._group(sims, live_tensors).values())
+        groups = list(SIM._group(sims, live_tensors).values())
+        groups_with_one = [g[0] for g in groups if len(g) == 1]
+        others = parallel(
+            [delayed(SIM.concat)(g, allow_different_tilings) for g in groups if len(g) > 1],
+            pbar="Combining SIMs"
+        )
+        return groups_with_one + others
 
     @staticmethod
     def filter_by_tensor_storages(