Added debug metrics

Author: tanner-andrulis

pytimeloop/fastfusion/mapper/mapper.py

@@ -21,7 +21,7 @@
 from pytimeloop.fastfusion.mapper.logging import make_queue_and_listener
 from pytimeloop.fastfusion.mapper.per_einsum_mapper import get_top_loop_jobs, mapper_place_fusion_level
 from pytimeloop.fastfusion.sim import Tiling, Loop, TensorStorage
-from pytimeloop.fastfusion.pareto import LOGSTRING, MAPPING, STATS, DICT_COLUMNS
+from pytimeloop.fastfusion.pareto import LOGSTRING, MAPPING, STATS, DICT_COLUMNS, TENSORS
 from pytimeloop.fastfusion.mapper.process_results import Metrics
 
 from pytimeloop.timeloopfe.v4 import Ert
@@ -153,8 +153,12 @@ def _convert_stats(from_einsum: int, to_einsum: int, stats, rank_renaming, tenso
     stats = deepcopy(stats)
     for s in stats:
         for d in DICT_COLUMNS:
-            s[d][to_einsum] = s[d].pop(from_einsum)
-        s[MAPPING][to_einsum] = s[MAPPING][to_einsum].rename(rank_renaming, tensor_renaming)
+            if d in s:
+                s[d][to_einsum] = s[d].pop(from_einsum)
+        if MAPPING in s:
+            s[MAPPING][to_einsum] = s[MAPPING][to_einsum].rename(rank_renaming, tensor_renaming)
+        if TENSORS in s:
+            s[TENSORS][to_einsum] = [t.rename(rank_renaming, tensor_renaming) for t in s[TENSORS][to_einsum]]
     return stats
 
 

pytimeloop/fastfusion/mapper/mapper_snowcat.py

@@ -138,9 +138,12 @@ def _convert_stats(from_einsum: int, to_einsum: int, stats, rank_renaming, tenso
     stats = deepcopy(stats)
     for s in stats:
         for d in DICT_COLUMNS:
-            s[d][to_einsum] = s[d].pop(from_einsum)
-        s[MAPPING][to_einsum] = s[MAPPING][to_einsum].rename(rank_renaming, tensor_renaming)
-        s[TENSORS][to_einsum] = [t.rename(rank_renaming, tensor_renaming) for t in s[TENSORS][to_einsum]]
+            if d in s:
+                s[d][to_einsum] = s[d].pop(from_einsum)
+        if MAPPING in s:
+            s[MAPPING][to_einsum] = s[MAPPING][to_einsum].rename(rank_renaming, tensor_renaming)
+        if TENSORS in s:
+            s[TENSORS][to_einsum] = [t.rename(rank_renaming, tensor_renaming) for t in s[TENSORS][to_einsum]]
     return stats
 
 

pytimeloop/fastfusion/mapper/process_results.py

@@ -30,6 +30,7 @@ class Metrics(Flag):
     # OCCUPANCY = auto()
     OFF_CHIP_ACCESSES = auto()
     OP_INTENSITY = auto()
+    DEBUG = auto()
 
     @classmethod
     def all_metrics(cls):
@@ -99,8 +100,9 @@ def record_storage(node):
             if storage.tensor_id not in found_tensors:
                 found_tensors.add(storage.tensor_id)
                 backing_storages.append(storage)
-
-        logstring.append(f"Strg({node['dspace']} in {node['target']})")
+                
+        if Metrics.DEBUG in metrics:
+            logstring.append(f"Strg({node['dspace']} in {node['target']})")
 
     def record_loop(node):
         nonlocal cur_idx
@@ -118,7 +120,8 @@ def record_loop(node):
         )
         ranks_remaining[node["rank"]] = tile_shape
         full_tiling.append(loop)
-        logstring.append(f"{node['type'][0].upper()}{node['rank']} size {tile_shape}")
+        if Metrics.DEBUG in metrics:
+            logstring.append(f"{node['type'][0].upper()}{node['rank']} size {tile_shape}")
 
     logstring = []
     full_tiling = []
@@ -157,16 +160,17 @@ def record_loop(node):
     if Metrics.ENERGY in metrics:
         results["Energy"] = energy
 
-    offchip_ac = 0
-    for (level, tensor, einsum), count in accesses.items():
-        if level == 0:
-            offchip_ac += count
-        logstring.append(f"Ac_{level}_{tensor}={count:.2e}")
-
     if Metrics.OFF_CHIP_ACCESSES in metrics:
+        offchip_ac = 0
+        for (level, tensor, einsum), count in accesses.items():
+            if level == 0:
+                offchip_ac += count
         results["Offchip Accesses"] = offchip_ac
-
-    logstring.append(f"{result.fanout}")
+        if Metrics.DEBUG in metrics:
+            logstring.append(f"Ac_{level}_{tensor}={count:.2e}")
+        
+    if Metrics.DEBUG in metrics:
+        logstring.append(f"{result.fanout}")
 
     # Only record non-backing reservations. We'll reserve backing storage later
     # when we free the tensors & we know all operations for which the tensor must
@@ -179,26 +183,25 @@ def record_loop(node):
             results[key] += r.tile_size
         # logstring.append(f"{r}")
 
-    if Metrics.LATENCY in metrics:
+    if Metrics.LATENCY in metrics and Metrics.DEBUG in metrics:
         logstring.append(f"L={results['Latency']:.2e}")
 
-    if Metrics.ENERGY in metrics:
+    if Metrics.ENERGY in metrics and Metrics.DEBUG in metrics:
         logstring.append(f"E={results['Energy']:.2e}")
 
     if Metrics.OP_INTENSITY in metrics:
         results["Op_Intensity"] = result.op_intensity[1]
 
-    logstring.append(f"Results: {results}")
-    results[LOGSTRING] = {einsum_id: str(logstring)}
-    results[MAPPING] = {einsum_id: tiling_full}
-    results[TENSORS] = {einsum_id: backing_storages}
-    results[STATS] = {
-        einsum_id: {k: v for k, v in results.items() if k not in RESERVED_COLUMNS}
-    }
-    results[IN_PROGRESS_STATS] = {einsum_id: {}}
-    results[MAPPING_HASH] = {einsum_id: hash((einsum_id, tiling_compatibility))}
-    results[TAGS] = {einsum_id: tiling_compatibility.tags}
+    if metrics.DEBUG in metrics:
+        logstring.append(f"Results: {results}")
+        results[LOGSTRING] = {einsum_id: str(logstring)}
+        results[STATS] = {einsum_id: {k: v for k, v in results.items() if k not in RESERVED_COLUMNS}}
+        results[TAGS] = {einsum_id: tiling_compatibility.tags}
+        results[MAPPING_HASH] = {einsum_id: hash((einsum_id, tiling_compatibility))}
+        results[IN_PROGRESS_STATS] = {einsum_id: {k: v for k, v in results.items() if k not in RESERVED_COLUMNS}}
+        results[TENSORS] = {einsum_id: backing_storages}
 
+    results[MAPPING] = {einsum_id: tiling_full}
     key = (tiling_compatibility, fzs(results.keys()))
 
     is_pareto = True

pytimeloop/fastfusion/pareto.py

@@ -218,10 +218,12 @@ def draw_looptree(row: pd.DataFrame, live_tensors: set[int]):
 
 def check_correctness(data: pd.DataFrame, live_tensors: set[int]):
     from pytimeloop.fastfusion.plot.looptree import tilings2looptree
+    from pytimeloop.fastfusion.sim import TensorStorage
 
     def fail(index):
         draw_looptree(data.iloc[index], live_tensors)
-        all_tensors = set(t for tn in r[TENSORS].values() for t in tn)
+        all_tensors = set(t for tn in r[MAPPING].values() for t in tn.tensors)
+        all_tensors = TensorStorage.get_backing_stores(all_tensors)
         for t in sorted(all_tensors):
             print(f"{t.__repr__()},")
 
@@ -230,8 +232,6 @@ def fail(index):
         looptree = tilings2looptree(
             r[MAPPING],
             r[STATS],
-            r[TENSORS],
-            r[IN_PROGRESS_STATS],
             skip_backing_tensors_in_right_branch=live_tensors,
             still_live_tensors=live_tensors,
         )
@@ -248,8 +248,6 @@ def fail(index):
                 looptree = tilings2looptree(
                     r[MAPPING],
                     r[STATS],
-                    r[TENSORS],
-                    r[IN_PROGRESS_STATS],
                     skip_backing_tensors_in_right_branch=live_tensors,
                     still_live_tensors=live_tensors,
                 )
@@ -337,6 +335,8 @@ def merge_cross(
         df = makepareto(df)
 
     for k in DICT_COLUMNS:
+        if k not in left.columns:
+            continue
         c0, c1 = k + MERGE_SUFFIXES[0], k + MERGE_SUFFIXES[1]
         df[k] = (
             df.apply(lambda row: {**row[c0], **row[c1]}, axis=1) if len(df) > 0 else []
@@ -345,10 +345,9 @@ def merge_cross(
 
     cols = [c for c in df.columns if c not in DICT_COLUMNS]
 
-    if True:
+    if IN_PROGRESS_STATS in df.columns:
         first_row = df.iloc[0]
-        tensors = first_row[TENSORS]
-        einsums = list(tensors.keys())
+        einsums = list(first_row[IN_PROGRESS_STATS].keys())
         last = einsums[-1]
         for i, r in df[cols].iterrows():
             df.at[i, IN_PROGRESS_STATS][last] = r.to_dict()
@@ -422,11 +421,12 @@ def alloc(self, resource_name: str, size: int, above_loop_index: int):
     def add_tensor(self, tensor):
         if len(self.data) == 0:
             return
-        last_einsum = list(self.data.iloc[0][TENSORS].keys())[-1]
-        if tensor in self.data[TENSORS].iloc[0][last_einsum]:
-            return
-        for t in self.data[TENSORS]:
-            t[last_einsum].append(tensor)
+        if TENSORS in self.data:
+            last_einsum = list(self.data.iloc[0][TENSORS].keys())[-1]
+            if tensor in self.data[TENSORS].iloc[0][last_einsum]:
+                return
+            for t in self.data[TENSORS]:
+                t[last_einsum].append(tensor)
 
     def copy(self) -> "Pareto":
         return Pareto(self.data.copy())

pytimeloop/fastfusion/plot/interactive.py

@@ -8,18 +8,18 @@
 import pandas as pd
 import plotly.express as px
 
-from pytimeloop.fastfusion.sim import Loop, Tiling
+from pytimeloop.fastfusion.sim import Loop, TensorStorage, Tiling
 from pytimeloop.fastfusion.util import expfmt
 from pytimeloop.fastfusion.plot.looptree import tilings2svg
 from pytimeloop.fastfusion.pareto import MAPPING, STATS, TENSORS, IN_PROGRESS_STATS, MAPPING_HASH
 
 import pandas as pd
 
 def mapping2svg(mapping: pd.Series):
-    return SVG(tilings2svg(mapping[MAPPING], mapping[STATS], mapping[TENSORS], mapping[IN_PROGRESS_STATS]))
+    return SVG(tilings2svg(mapping[MAPPING], mapping.get(STATS, None)))
 
 def mapping2svg(mapping: pd.Series):
-    return SVG(tilings2svg(mapping[MAPPING], mapping[STATS], mapping[TENSORS], mapping[IN_PROGRESS_STATS]))
+    return SVG(tilings2svg(mapping[MAPPING], mapping.get(STATS, None)))
 
 def diplay_mappings_on_fig(fig: plotly.graph_objs.FigureWidget, data: dict[str, pd.DataFrame]):
     # fig = go.FigureWidget(fig)
@@ -33,13 +33,10 @@ def display_mapping(trace, points, selector):
         d = data[trace.name]
         index = points.point_inds[0]
         display(mapping2svg(d.iloc[index]))
-        all_tensors = set(
-            t for tn in d.iloc[index][TENSORS].values() for t in tn
-        )
-        for t in sorted(all_tensors):
+        backing_tensors = set(t for tn in d.iloc[index][MAPPING].values() for t in tn.tensors)
+        backing_tensors = TensorStorage.get_backing_stores(backing_tensors)
+        for t in sorted(backing_tensors):
             print(f"{t.__repr__()},")
-        for k, v in d.iloc[index][MAPPING_HASH].items():
-            print(f"{k}: {v},")
         for t in sorted(list(d.iloc[index][MAPPING].values())[-1].tensors):
             print(f"{t.__repr__()},")
         for v in d.iloc[index][MAPPING].values():
@@ -53,13 +50,10 @@ def display_mapping_2(trace, points, selector):
         d = data[trace.name]
         index = points.point_inds[0]
         display(mapping2svg(d.iloc[index]))
-        all_tensors = set(
-            t for tn in d.iloc[index][TENSORS].values() for t in tn
-        )
-        for t in sorted(all_tensors):
+        backing_tensors = set(t for tn in d.iloc[index][MAPPING].values() for t in tn.tensors)
+        backing_tensors = TensorStorage.get_backing_stores(backing_tensors)
+        for t in sorted(backing_tensors):
             print(f"{t.__repr__()},")
-        for k, v in d.iloc[index][MAPPING_HASH].items():
-            print(f"{k}: {v},")
         for t in sorted(list(d.iloc[index][MAPPING].values())[-1].tensors):
             print(f"{t.__repr__()},")
         for v in d.iloc[index][MAPPING].values():

pytimeloop/fastfusion/plot/looptree.py

@@ -80,21 +80,23 @@ def get_all_storages(self, _entry: bool = True) -> list[TensorStorage]:
     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):
+def tilings2looptree(mappings: dict[str, Tiling], stats: dict[str, Any], skip_backing_tensors_in_right_branch: Iterable[str] = (), still_live_tensors: set[str] = (), shared_loop_index: int = -1):
     prev_tilings = []
     root = Node()
     einsum_ids = list(mappings.keys())
+    
     assert set(einsum_ids) == set(stats.keys())
-    assert set(einsum_ids) == set(backing_tensors.keys())
+    
+    
 
     # If a tensor appears in non-back-to-back Einsums, then we need to store it for
     # all Einsums in between
-    tensors_lifetimes = {e: [] for e in backing_tensors}
-    all_tensors = set().union(*[set(t) for t in backing_tensors.values()])
+    tensors_lifetimes = {e: [] for e in einsum_ids}
+    all_tensors = set().union(*[set(t.tensors) for t in mappings.values()])
     backers = TensorStorage.get_backing_stores(all_tensors)
     for t in all_tensors:
-        first_appearance = min(i for i, ts in enumerate(backing_tensors.values()) if t in ts)
-        last_appearance = max(i for i, ts in enumerate(backing_tensors.values()) if t in ts)
+        first_appearance = min(i for i, ts in enumerate(mappings.values()) if t in ts.tensors)
+        last_appearance = max(i for i, ts in enumerate(mappings.values()) if t in ts.tensors)
         if t.tensor_id in still_live_tensors:
             last_appearance = len(einsum_ids) - 1
         for i, l in enumerate(tensors_lifetimes.values()):
@@ -151,13 +153,13 @@ def tilings2looptree(mappings: dict[str, Tiling], stats: dict[str, Any], backing
 
     return root
 
-def tilings2svg(mappings: dict[str, Tiling], stats: dict[str, Any], tensors: dict[str, list[TensorStorage]], partial_stats: dict[str, Any]):
-    root = tilings2looptree(mappings, stats, tensors, partial_stats)
+def tilings2svg(mappings: dict[str, Tiling], stats: dict[str, Any], ):
+    root = tilings2looptree(mappings, stats)
     graph = pydot.Dot(graph_type="digraph", ranksep="0.2", nodesep="0.2")
     root.to_pydot(graph)
     return graph.create_svg()
 
-def tilings2yaml(mappings: dict[str, Tiling], stats: dict[str, Any], tensors: dict[str, list[TensorStorage]], partial_stats: dict[str, Any]):
-    root = tilings2looptree(mappings, stats, tensors, partial_stats)
+def tilings2yaml(mappings: dict[str, Tiling], stats: dict[str, Any]):
+    root = tilings2looptree(mappings, stats)
     return root.to_yaml()