Merge branch 'main' of github.com:Accelergy-Project/timeloop-python

Author: gilbertmike

pytimeloop/fastfusion/mapper/mapper_snowcat.py

@@ -30,7 +30,8 @@ def mapper(
     explore_glb_uneven,
     spec,
     tmp_path,
-    ffmt: bool=False
+    ffmt: bool=False,
+    ffmt_refetch_weights: bool=True,
 ):
     logger.info(f"Calling mapper for {spec}")
 
@@ -65,7 +66,8 @@ def mapper(
         explore_glb_uneven=explore_glb_uneven,
         spec=spec,
         energy_dict=energy_dict,
-        ffmt=ffmt
+        ffmt=ffmt,
+        ffmt_refetch_weights=ffmt_refetch_weights,
     )
 
     generated_data = {}
@@ -132,12 +134,12 @@ def detect_similar_einsums(workload, analyzer, separated_einsums=None):
             found = False
             for from_einsum in ref_to_to_einsums:
                 rank_renaming, tensor_renaming = is_equivalent(from_einsum,
-                                                            einsum,
-                                                            workload,
-                                                            analyzer)
+                                                               einsum,
+                                                               workload,
+                                                               analyzer)
                 if rank_renaming is not None:
                     ref_to_to_einsums[from_einsum][einsum] = (rank_renaming,
-                                                                tensor_renaming)
+                                                              tensor_renaming)
                     found = True
                     break
             if not found:
@@ -158,6 +160,14 @@ def convert_rank_to_group_renaming(ref_to_to_einsums, equiv_ranks):
 
 
 def get_ffmt_separated_einsums(workload):
+    einsum_id_to_name = workload.einsum_id_to_name()
+    if len(einsum_id_to_name) == 1:
+        return [{0}]
+    elif len(einsum_id_to_name) == 2:
+        return [{0}, {1}]
+    elif len(einsum_id_to_name) == 3:
+        return [{0}, {1}, {2}]
+
     first_einsum = {0}
     second_einsum = {1}
     last_einsum = {max(workload.einsum_id_to_name().keys())}

pytimeloop/fastfusion/mapper/per_einsum_mapper.py

@@ -687,4 +687,11 @@ def make_temporal_fors_with_smallest_tile(original, ranks):
         mapping = original.copy()
         for r in ordered_ranks:
             mapping.add_temporal(r, tile_shape=1)
+        yield mapping
+
+def make_temporal_fors_in_order(original, ranks):
+    for i in range(len(ranks)+1):
+        mapping = original.copy()
+        for r in ranks[:i]:
+            mapping.add_temporal(r)
         yield mapping

pytimeloop/fastfusion/mapper/per_einsum_mapper_snowcat.py

@@ -1,9 +1,5 @@
 from copy import deepcopy
 from collections import defaultdict
-from collections.abc import Callable, Set
-from itertools import combinations, product, permutations
-from functools import reduce
-from operator import or_, mul
 
 from joblib import Parallel, delayed
 
@@ -27,7 +23,8 @@ def per_einsum_mapper_snowcat(
     explore_glb_uneven,
     einsums_to_explore,
     energy_dict,
-    ffmt=False
+    ffmt=False,
+    ffmt_refetch_weights=True,
 ):
     data = {}
     for einsum_id in einsums_to_explore:
@@ -70,13 +67,15 @@ def per_einsum_mapper_snowcat(
                                             intermediate_tensors,
                                             tensor_to_relevant_ranks,
                                             einsum_id,
-                                            workload)
+                                            workload,
+                                            refetch_weights=ffmt_refetch_weights)
 
         n_jobs=32
         parallelized_spaces, task_spaces = \
             split_dependent_product(n_split_min=n_jobs, spaces=subspaces)
 
         partial_mappings = list(dependent_product(parallelized_spaces))
+        partial_mappings = [x if isinstance(x, tuple) else (x,) for x in partial_mappings]
 
         def per_worker_exploration(*args):
             analyzer = LooptreeWorkloadDependencyAnalyzer(workload)
@@ -123,12 +122,5 @@ def per_worker_exploration(*args):
             for k, v in res.items():
                 data[einsum_id][k] += v
 
-        print(einsum_id)
-        for k, v in data[einsum_id].items():
-            min_metric = float("inf")
-            for m in v:
-                min_metric = min(min_metric, m["Offchip_Ac"])
-            print(min_metric)
-
     return data
 

pytimeloop/fastfusion/mapper/per_einsum_subspaces/snowcat.py

@@ -24,8 +24,8 @@ def off_chip_storage(mapping):
     def fused_temporal_fors(mapping, unfused_tensors):
         for partial_mapping in make_temporal_fors(mapping, all_ranks):
             # for partial_mapping in make_temporal_fors(mapping, all_ranks):
-                for partial_mapping in make_temporal_fors_with_smallest_tile(mapping, all_ranks):
-                    yield partial_mapping, unfused_tensors
+            for partial_mapping in make_temporal_fors_with_smallest_tile(partial_mapping, all_ranks):
+                yield partial_mapping, unfused_tensors
 
 
     def glb_storage(mapping, unfused_tensors):

pytimeloop/fastfusion/mapper/per_einsum_subspaces/snowcat_ffmt.py

@@ -1,10 +1,12 @@
-from pytimeloop.fastfusion.mapper.per_einsum_mapper import LinearMapping, make_storage, make_temporal_fors, make_temporal_fors_with_smallest_tile
+from pytimeloop.fastfusion.mapper.per_einsum_mapper import LinearMapping, make_storage, make_temporal_fors, make_temporal_fors_with_smallest_tile, make_temporal_fors_in_order
 
 def make_ffmt_subspaces(tensors,
                         intermediate_tensors,
                         tensor_to_relevant_ranks,
                         einsum_id,
-                        workload):
+                        workload,
+                        refetch_weights: bool=True):
+
     def off_chip_storage(mapping):
         off_chip_must_retain = tensors - intermediate_tensors
         off_chip_can_retain = intermediate_tensors
@@ -25,44 +27,80 @@ def off_chip_storage(mapping):
     M = all_ranks[0]
     N = all_ranks[1]
     K = all_ranks[2]
-
-    if einsum_id == 0:
-        allowed_fused_ranks = all_ranks
-    elif einsum_id == 1:
-        allowed_fused_ranks = {M, K}
-    elif einsum_id == max(workload.einsum_id_to_name().keys()):
-        allowed_fused_ranks = {M, N}
-    else:
-        allowed_fused_ranks = {M}
+    weight_tensor = None
+    input_tensor = None
+    for tensor_id in workload.tensors_read_by_einsum(einsum_id):
+        if tensor_to_relevant_ranks[tensor_id] == {K, N}:
+            weight_tensor = tensor_id
+        elif tensor_to_relevant_ranks[tensor_id] == {M, K}:
+            input_tensor = tensor_id
+    assert weight_tensor is not None
+    assert input_tensor is not None
+    output_tensor = next(iter(workload.tensors_written_by_einsum(einsum_id)))
+    non_weight_tensor = tensors - {weight_tensor}
 
     def fused_temporal_fors(mapping, unfused_tensors):
-        for partial_mapping in make_temporal_fors(mapping, allowed_fused_ranks):
-            # for partial_mapping in make_temporal_fors(mapping, all_ranks):
-                for partial_mapping in make_temporal_fors_with_smallest_tile(mapping, all_ranks):
-                    yield partial_mapping, unfused_tensors
+        if input_tensor in unfused_tensors:
+            allowed_fused_ranks = [M, N, K]
+        elif output_tensor in unfused_tensors:
+            allowed_fused_ranks = [M, N]
+        else:
+            allowed_fused_ranks = [M, K]
+        for partial_mapping in make_temporal_fors_in_order(mapping, allowed_fused_ranks):
+            yield partial_mapping, unfused_tensors
 
 
-    def glb_storage(mapping, unfused_tensors):
+    def glb_storage_io(mapping, unfused_tensors):
         glb_fused_tensors = intermediate_tensors - unfused_tensors
         yield from make_storage(
             mapping,
             level=1,
-            must_retain_tensors=tensors,
+            must_retain_tensors=non_weight_tensor,
             can_retain_tensors=set(),
             must_fully_reuse_tensors=glb_fused_tensors,
             tensor_to_relevant_ranks=tensor_to_relevant_ranks,
-            explore_uneven=True,
-            add_split_at_tensors=glb_fused_tensors
+            explore_uneven=False,
+            add_split_at_tensors=glb_fused_tensors,
+            return_retained_tensors=True,
         )
 
-    def mac(mapping):
+    def intra_temporal_fors(mapping, _):
+         for partial_mapping in make_temporal_fors_with_smallest_tile(mapping,
+                                                                      {K, N}):
+              yield partial_mapping, _ 
+
+    def glb_storage_weights(mapping, _):
+         yield from make_storage(
+              mapping,
+              level=1,
+              must_retain_tensors={weight_tensor},
+              can_retain_tensors=set(),
+              tensor_to_relevant_ranks=tensor_to_relevant_ranks,
+              explore_uneven=False,
+              return_retained_tensors=True,
+         )
+
+    def mac(mapping, _):
             mapping.add_compute(einsum_id, 2)
             yield mapping
 
-    return [
-        lambda: [LinearMapping()],
-        off_chip_storage,
-        fused_temporal_fors,
-        glb_storage,
-        mac
-    ]
+    if refetch_weights:
+        return [
+            lambda: [LinearMapping()],
+            off_chip_storage,
+            fused_temporal_fors,
+            glb_storage_io,
+            intra_temporal_fors,
+            glb_storage_weights,
+            mac
+        ]
+    else:
+        return [
+            lambda: [LinearMapping()],
+            off_chip_storage,
+            glb_storage_weights,
+            fused_temporal_fors,
+            glb_storage_io,
+            intra_temporal_fors,
+            mac
+        ]

pytimeloop/fastfusion/mapper/simexplore.py

@@ -0,0 +1,89 @@
+from collections.abc import Mapping
+import itertools
+
+import numpy as np
+import pandas as pd
+from joblib import Parallel, delayed
+
+from pytimeloop.fastfusion.sim import SIM
+from pytimeloop.fastfusion.pareto import Pareto
+
+def explore_fusion(einsum_to_result: Mapping):
+
+    r2 = {}
+    for einsum_id, compat_dict in einsum_to_result.items():
+        r2[einsum_id] = Parallel(n_jobs=1)(delayed(paretofy)(k, v) for k, v in compat_dict.items())
+
+    # for einsum_id, compat_dict in result.items():
+    #     r2[einsum_id] = [SIM(k, Pareto(pd.DataFrame(v).fillna(0))) for k, v in compat_dict.items()]
+        
+    sims = list(r2.values())
+    s = sims.pop(0)
+
+
+    while sims:
+        live_tensors = set.union(set(), *[sim[0].tensor_names for sim in sims])
+        ns = sims.pop(0)
+        next_live_tensors = set.union(set(), *[sim[0].tensor_names for sim in sims])
+
+        for s2 in s:
+            s2.consolidate(live_tensors)
+
+        ns = SIM.combine_combineable(ns, next_live_tensors | s[0].tensor_names)
+        ns = SIM.group_by_left(ns, s[0].tensor_names)
+        s = SIM.combine_combineable(s, live_tensors)
+        s = SIM.group_by_right(s, live_tensors)
+
+        print("\n\n")
+        print("\n\n" + "=" * 100 + f"\n{len(sims) + 1} Remaining\n" + "=" * 100)
+
+        DO_PRINT = False
+
+        with open('s_keys.txt', 'w') as f:
+            for key in sorted(s.keys()):
+                f.write(f"{key}\n")
+
+        with open('s2_keys.txt', 'w') as f:
+            for key in sorted(ns.keys()):
+                f.write(f"{key}\n")
+
+        combined: list[SIM] = []
+        for k in s:
+            if k in ns:
+                for a, b in itertools.product(s[k], ns[k]):
+                    if DO_PRINT:
+                        print(f"\t{a.tiling_str()} <--> {b.tiling_str()}")
+                    combined.append(a.merge_next(b, set(), delay=True))
+                    # combined_keys.append()
+            elif DO_PRINT:
+                print(f"\tNo match for {s[k][0].tiling_str()}")
+
+        for c, mapping in zip(combined, Parallel(n_jobs=128)(c.mapping for c in combined)):
+            c.mapping = mapping
+
+        s = combined
+        print(f"Generated {len(s)} solutions")
+        
+    for s2 in s:
+        s2.consolidate(set())
+    s_final = SIM.combine_combineable(s, set())[0]
+    data = s_final.mapping.data
+    # Sort data by the columns "Latency" and "Energy"
+    last_level_occupancy = None
+    for i in reversed(range(3)):
+        if f"RESOURCE_1_LEVEL_{i}" not in data:
+            continue
+        if last_level_occupancy is not None:
+            non_left_cur_level_occupancy = data[f"RESOURCE_1_LEVEL_{i}"] + last_level_occupancy
+        else:
+            non_left_cur_level_occupancy = data[f"RESOURCE_1_LEVEL_{i}"]
+        left_cur_level_occupancy = data[f"RESOURCE_1_LEFT_LEVEL_{i}"]
+        last_level_occupancy = np.maximum(non_left_cur_level_occupancy,
+                                            left_cur_level_occupancy)
+    data["Occupancy"] = last_level_occupancy
+
+    return data
+
+
+def paretofy(k, v):
+    return SIM(k, Pareto(pd.DataFrame(v).fillna(0)))