Bugfixes

Author: tanner-andrulis

pytimeloop/fastfusion/filter_mappings.py

@@ -1,5 +1,5 @@
 from collections import defaultdict
-from pytimeloop.fastfusion.sim import Tag, TensorStorage, Tiling
+from pytimeloop.fastfusion.sim import TensorStorage, Tiling
 
 def get_ffmt_tag_mha(
         einsum_id: str, 
@@ -19,6 +19,8 @@ def get_ffmt_tag_mha(
         "AV": ["QK_QK_to_AV", "AV_AV_to_Z"], # NOTE: AV IS MISSING
         "Z": ["AV_AV_to_Z", "Z_Z_to_n"],
     }
+    if einsum_id not in einsum_id_to_input_output:
+        return ("FFMT_VALID",)
     a, b = einsum_id_to_input_output[einsum_id]
 
     tags = []
@@ -103,7 +105,7 @@ def get_tileflow_tag_mha(
     ):
     # Valid iff it's an even mapping
     storage2level = defaultdict(set)
-    for ts in tiling.tensors:
+    for ts in backing_storages:
         storage2level[ts.backer_id].add(ts.above_loop_index)
     if all(len(s) == 1 for s in storage2level.values()):
         return ("TILEFLOW_VALID",)

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

@@ -69,27 +69,30 @@ def make_storage(
         else:
             min_i = last_storage_idx + 1
 
-        no_irrelevant_loop = True
+        any_irrelevant_loop = False
         for i in range(min_i, len(mapping)):
             node = mapping[i]
             if node["type"] == "temporal":
                 is_relevant = node["rank"] in relevant_ranks
 
-                no_irrelevant_loop = no_irrelevant_loop and not is_relevant
+                any_irrelevant_loop = any_irrelevant_loop or not is_relevant
 
-                auto_lower = has_storage and i > apply_lrp_after_loop_idx
+                auto_lower = i > apply_lrp_after_loop_idx
+                auto_lower = False
 
                 if not auto_lower or (last_is_relevant and not is_relevant):
                     tensor_choices.append(i)
 
                 last_is_relevant = is_relevant
 
-                if not is_relevant and tensor_must_be_fully_reused:
+                if tensor_must_be_fully_reused and any_irrelevant_loop:
                     break
+                
+        auto_lower = len(mapping) > apply_lrp_after_loop_idx
+        auto_lower = False
 
         # Lowest possible storage node
-        if last_is_relevant and (not tensor_must_be_fully_reused
-                                 or no_irrelevant_loop):
+        if (not auto_lower or last_is_relevant) and not (tensor_must_be_fully_reused and any_irrelevant_loop):
             tensor_choices.append(len(mapping))
 
         if tensor_id in can_retain_tensors:

pytimeloop/fastfusion/mapper/process_results.py

@@ -142,21 +142,12 @@ def record_loop(node):
         tiling=tiling_full,
         rank_name_to_shared_name=rank_name_to_shared_name,
     )
-    # print(tiling_full)
 
     tiling_compatibility = Tiling(
         loops=tuple(full_tiling[:n_fused_loops]),
         tensors=frozenset(backing_storages),
         tags=fzs().union(*([set()] + [set(t(**tagger_args)) for t in tag_with]))
     )
-    
-    if "FFMT_VALID" in tiling_compatibility.tags:
-        print(tiling_compatibility)
-    # assert max(t.above_loop_index for t in backing_storages) == len(tiling_compatibility.loops), (
-    #     f"\n\ttiling_compatibility: {tiling_compatibility} "
-    #     f"\n\tbacking_storages: {backing_storages} "
-    #     f"\n\ttiling_full: {tiling_full}"
-    # )
 
     results = {}
 

pytimeloop/fastfusion/mapper/simexplore.py

@@ -156,6 +156,9 @@ def fuse_sims(
                 for a in left[k]:
                     print(f"\tNo match for {k} |||| {a.tiling_str()}")
 
+        if not combined:
+            print(f'No valid combinations found.')
+
         print_time("Bucket merging")
 
         f = parallel if DELAY_MERGE else lambda x: x

pytimeloop/fastfusion/pareto.py

@@ -198,17 +198,29 @@ def _free_to_loop_index(
 def paretofy_by(data: pd.DataFrame, columns: list[str]) -> pd.DataFrame:
     return data[paretoset(data[columns])].reset_index(drop=True)
 
+def draw_looptree(row: pd.DataFrame, live_tensors: set[int]):
+    from pytimeloop.fastfusion.plot.looptree import tilings2looptree
+
+    looptree = tilings2looptree(
+        row[MAPPING],
+        row[STATS],
+        row[TENSORS],
+        row[IN_PROGRESS_STATS],
+        skip_backing_tensors_in_right_branch=live_tensors,
+        still_live_tensors=live_tensors,
+    )
+    import pydot
 
+    graph = pydot.Dot(graph_type="digraph", ranksep="0.2", nodesep="0.2")
+    looptree.to_pydot(graph)
+    with open(f"test.png", "wb") as f:
+        f.write(graph.create_png())
+        
 def check_correctness(data: pd.DataFrame, live_tensors: set[int]):
     from pytimeloop.fastfusion.plot.looptree import tilings2looptree
 
-    def fail(looptree):
-        import pydot
-
-        graph = pydot.Dot(graph_type="digraph", ranksep="0.2", nodesep="0.2")
-        looptree.to_pydot(graph)
-        with open(f"test.png", "wb") as f:
-            f.write(graph.create_png())
+    def fail(index):
+        draw_looptree(data.iloc[index], live_tensors)
         all_tensors = set(t for tn in r[TENSORS].values() for t in tn)
         for t in sorted(all_tensors):
             print(f"{t.__repr__()},")

pytimeloop/fastfusion/plot/interactive.py

@@ -90,7 +90,7 @@ def plotly_show(
             fig.add_scatter(x=v[x], y=v[y], name=k, line={"shape": 'hv'})
     else:
         data.sort_values(by=[x, y], inplace=True)
-        fig.add_scatter(x=data[x], y=data[y], line={"shape": 'hv'})
+        fig.add_scatter(x=data[x], y=data[y], name="", line={"shape": 'hv'})
         data = {"" : data}
     if title is not None:
         fig.update_layout(title=title)

pytimeloop/fastfusion/plot/looptree.py

@@ -81,7 +81,7 @@ def get_backing_stores(self):
         return TensorStorage.get_backing_stores(self.get_all_storages())
 
 def tilings2looptree(mappings: dict[str, Tiling], stats: dict[str, Any], backing_tensors: dict[str, list[TensorStorage]], partial_stats: dict[str, Any], skip_backing_tensors_in_right_branch: Iterable[str] = (), still_live_tensors: set[str] = (), shared_loop_index: int = -1):
-    prev_tiling = None
+    prev_tilings = []
     root = Node()
     einsum_ids = list(mappings.keys())
     assert set(einsum_ids) == set(stats.keys())
@@ -105,9 +105,10 @@ def tilings2looptree(mappings: dict[str, Tiling], stats: dict[str, Any], backing
     for i, einsum_id in enumerate(einsum_ids):
         skip_backing_tensors = () #if i < len(einsum_ids) - 1 else skip_backing_tensors_in_right_branch
         tiling = mappings[einsum_id]
-        index = (
-            prev_tiling.shared_loop_index(tiling.tensor_names) if prev_tiling else -1
-        )
+        if not prev_tilings:
+            index = -1
+        else:
+            index = max(p.shared_loop_index(tiling.tensor_names) for p in prev_tilings)
         n = root.access_level(index)
         loops = tiling.loops[index:] if index != -1 else tiling.loops
         for l in loops:
@@ -128,13 +129,13 @@ def tilings2looptree(mappings: dict[str, Tiling], stats: dict[str, Any], backing
             for tensor in storages:
                 n = root.access_level(tensor.above_loop_index)
                 # TODO if tensor not in n.this_level or tensor not in backers:
-                if tensor not in n.this_level:
+                if tensor not in n.this_level or tensor not in backers:
                     n.this_level.append(tensor)
 
         root.add_stats(stats[einsum_id])
         # for k, v in partial_stats[einsum_id].items():
         #     last_level.append(f"_PARTIAL {k}: {expfmt(v)}")
-        prev_tiling = tiling
+        prev_tilings.append(tiling)
 
     n = root
     skip_backing_tensors_in_right_branch= set(skip_backing_tensors_in_right_branch)

pytimeloop/fastfusion/sim.py

@@ -132,38 +132,11 @@ def __eq__(self, value):
                 return False
         return True
 
-class Tag:
-    def __eq__(self, value: "Tag") -> bool:
-        # Match base Tag
-        if type(value) == Tag:
-            return True
-        # If other is a non-base Tag or anything else, let
-        # the other object handle the comparison
-        return value.__eq__(self)
-    
-    def __hash__(self) -> int:
-        return hash("Tag")
-    
-    def merge_next(self, n: "Tag") -> "Tag":
-        return n
-    
-    @staticmethod
-    def get_wildcard():
-        return Wildcard()
-
-class Wildcard(Tag):
-    def __eq__(self, value: "Tag") -> bool:
-        return True
-    
-    def __hash__(self) -> int:
-        return hash("Wildcard")
-
-
 @dataclass(frozen=True)
 class Tiling:
     loops: tuple[Loop, ...]
     tensors: fzs[TensorStorage]
-    tags: fzs[Tag] = fzs()
+    tags: fzs[Any] = fzs()
 
     @cached_property
     def tensor_names(self) -> set[str]:
@@ -209,9 +182,7 @@ def __str__(self):
         return self.__repr__()
 
     def __repr__(self):
-        if type(self.tags) == Tag:
-            return f"Tiling({self.loops}, {self.tensors})"
-        return f"Tiling({self.loops}, {self.tensors}, {self.tags})"
+        return f"Tiling({self.loops.__repr__()}, {self.tensors.__repr__()}, {self.tags.__repr__()})"
 
     def merge_next(self, n: "Tiling", live_tensors: set[str]) -> "Tiling":
         tensors = fzs(t for t in (n.tensors | self.tensors) if t.tensor_id in live_tensors)
@@ -248,7 +219,7 @@ def matches_permutation(self, permutation: list[str]) -> bool:
     def has_tensor(self, *tensors: TensorStorage) -> bool:
         return all(any(t == tensor for t in self.tensors) for tensor in tensors)
 
-    def set_tags(self, *new_tags: Tag) -> "Tiling":
+    def set_tags(self, *new_tags: Any) -> "Tiling":
         return Tiling(self.loops, self.tensors, fzs(new_tags))
 
 class SIM:
@@ -410,20 +381,12 @@ def group_right(
     ) -> dict[tuple[Tiling, ...], list["SIM"]]:
         return SIM._group(sims, live_tensors)
 
-    def set_tags(self, *tags: Tag) -> "SIM":
+    def set_tags(self, *tags: Any) -> "SIM":
         self.tiling = self.tiling.set_tags(*tags)
 
-    @attribute
-    def tags(self) -> fzs[Tag]:
+    @property
+    def tags(self) -> fzs[Any]:
         return self.tiling.tags
-    
-    @attribute
-    def tensors(self) -> fzs[TensorStorage]:
-        return self.tiling.tensors
-    
-    @attribute
-    def loops(self) -> tuple[Loop, ...]:
-        return self.tiling.loops
 
 
 import unittest

pytimeloop/fastfusion/util.py

@@ -9,10 +9,7 @@
 
 class fzs(frozenset):
     def __repr__(self):
-        try:
-            return f"fzs({', '.join(x.__repr__() for x in sorted(self))}"
-        except:
-            return f"fzs({', '.join(x.__repr() for x in self)})"
+        return f"fzs({', '.join(sorted(x.__repr__() for x in self))})"
 
     def __str__(self):
         return self.__repr__()

pytimeloop/looptree/mapping_utilities.py

@@ -62,9 +62,10 @@ def get_last_storage_node(mapping, tensor):
 
 
 def get_last_fused_loop_idx(mapping, intermediate_tensors):
-    last_i = 0
+    intermedaites_remaining = set(intermediate_tensors)
     for i, node in enumerate(mapping):
-        if node['type'] == 'storage' \
-                and any(t in node['dspace'] for t in intermediate_tensors):
-            last_i = i
-    return last_i
+        if node['type'] == 'storage':
+            intermedaites_remaining -= set(node['dspace'])
+        if not intermedaites_remaining:
+            return i
+    return float('inf')