accelforge 0.0.1__py3-none-any.whl

accelforge/__init__.py

@@ -0,0 +1,21 @@
+from accelforge.frontend import arch
+from accelforge.frontend import config
+from accelforge.frontend import mapping
+from accelforge.frontend import renames
+from accelforge.frontend import spec
+from accelforge.frontend import variables
+from accelforge.frontend import workload
+from accelforge.frontend.spec import Spec, Spec
+from accelforge.mapper.FFM import Metrics
+from accelforge.util import set_n_parallel_jobs
+from accelforge.util import LiteralString
+import accelforge.mapper as mapper
+from accelforge.examples import examples
+
+from accelforge.frontend.variables import Variables
+from accelforge.frontend.arch import Arch
+from accelforge.frontend.config import Config
+from accelforge.frontend.mapping import Mapping
+from accelforge.frontend.renames import Renames
+from accelforge.frontend.spec import Spec
+from accelforge.frontend.workload import Workload

accelforge/_accelerated_imports.py

@@ -0,0 +1,16 @@
+import os
+
+os.environ["ACCELFORGE_ACCELERATED_IMPORTS"] = "0"
+
+if os.environ.get("ACCELFORGE_ACCELERATED_IMPORTS", "0") == "1":
+    import cudf as pd
+    import cupy as np
+    import cupy as scipy
+
+    ACCELERATED = True
+else:
+    import pandas as pd
+    import numpy as np
+    import scipy
+
+    ACCELERATED = False

accelforge/_deprecate/_simanneal/evalmapping.py

@@ -0,0 +1,271 @@
+from collections import defaultdict
+import itertools
+import time
+from fastfusion._accelerated_imports import pd
+from fastfusion.mapper.FFM._join_pmappings.sim import PmappingGroup, Loop, Compatibility
+from fastfusion.mapper.FFM._join_pmappings.pmapping_group import PmappingDataframe
+from fastfusion.mapper.simanneal.mapspaceglobals import MapspaceGlobals
+from fastfusion.util._frozenset import fzs
+
+
+def mapping2sims(einsum_to_result: Compatibility):
+    r = {}
+    for einsum_name, compat_dict in einsum_to_result.items():
+        r[einsum_name] = [paretofy(k, v) for k, v in compat_dict.items()]
+    return list(r.values())
+
+
+def get_possible_translations(
+    t: Compatibility,
+    pairwise_equivalent_rank_variables: dict[str, set[str]],
+    full_equivalent_rank_variables: dict[str, set[str]],
+    right_rank_variables: set[str],
+):
+    # Fused ranks should be transitive, but if a fused loop indexes into two
+    # different ranks in the next Einsum, we can't fuse becuase it will tile in
+    # multiple directions.
+    #
+    # The first union checks what loops we CAN fuse with in the next Einsum. The
+    # second union checks what loops MUST index into in the next
+    #
+    # Einsum. If we alias into multiple ranks, we can't fuse. Otherwise, try out
+    # each possible rank.
+    def translate_loop(l: Loop):
+        compatible_rank_variables = (
+            set.union(
+                *(full_equivalent_rank_variables[n] for n in l.rank_variable_names)
+            )
+            & right_rank_variables
+        )
+        pairwise_compatible_rank_variables = (
+            set.union(
+                *(pairwise_equivalent_rank_variables[n] for n in l.rank_variable_names)
+            )
+            & right_rank_variables
+        )
+        if len(pairwise_compatible_rank_variables) > 1:
+            return
+        for n in compatible_rank_variables:
+            yield Loop(fzs((n,)), l.bound, l.is_spatial)
+
+    for loops in itertools.product(*map(translate_loop, t.loops)):
+        yield t.update(loops=loops)
+
+
+prev_time = 0
+total_time = defaultdict(int)
+
+
+def init_print_time():
+    global prev_time, total_time
+    prev_time = time.time()
+    total_time = defaultdict(int)
+
+
+def print_time(what: str):
+    global prev_time
+    t = time.time() - prev_time
+    print(f"{what}: {t:.2f} seconds")
+    total_time[what] += t
+    prev_time = time.time()
+
+
+def print_total_time():
+    print(f"\n======== Total time ========")
+    for k, v in total_time.items():
+        print(f"{k}: {v:.2f} seconds")
+    total = sum(total_time.values())
+    if total > 60:
+        print(f"\nTotal: {total:.2f} seconds ({total/60:.2f} minutes)")
+    else:
+        print(f"\nTotal: {total:.2f} seconds")
+    print(f"============================\n")
+
+
+class PmappingsOneEinsum:
+    def __init__(self, einsum_name: str, pm_group_list: list[PmappingGroup]):
+        self.einsum_name: str = einsum_name
+        self.pmapping_groups: list[PmappingGroup] = pm_group_list
+        self.tensor_names: set[str] = set(pm_group_list[0].tensor_names)
+
+    def __getitem__(self, i):
+        return self.pmapping_groups[i]
+
+
+def make_full_equivalent_rank_variables(pairwise_equivalent_rank_variables):
+    full_equivalent_rank_variables = {
+        k: set(v) for k, v in pairwise_equivalent_rank_variables.items()
+    }
+    changed = True
+    while changed:
+        changed = False
+        for r in full_equivalent_rank_variables:
+            for r2 in list(full_equivalent_rank_variables[r]):
+                for r3 in list(full_equivalent_rank_variables[r2]):
+                    if r3 in full_equivalent_rank_variables[r]:
+                        continue
+                    changed = True
+                    full_equivalent_rank_variables[r].add(r3)
+    return full_equivalent_rank_variables
+
+
+def quick_join(
+    pmapping_groups: dict[str, PmappingGroup],
+    mapspace_globals: MapspaceGlobals,
+):
+    resource2capacity = mapspace_globals.resource2capacity
+    pairwise_equivalent_rank_variables = mapspace_globals.pairwise_equivalent_ranks
+    aliased_tensors = mapspace_globals.aliased_tensors
+    full_equivalent_rank_variables = mapspace_globals.full_equivalent_ranks
+
+    n_mappings = {}
+    runtime = {}
+    nbuckets = []
+
+    n_evaluations = 0
+
+    pmapping_groups = list(pmapping_groups.items())
+
+    init_print_time()
+
+    pmapping_groups = [PmappingsOneEinsum(*s) for s in pmapping_groups]
+
+    if not pmapping_groups:
+        raise ValueError("No PmappingGroups to join")
+
+    # ======================================================================
+    # Initial consolidate and group all PmappingGroups
+    # ======================================================================
+    for i, sim_holder in enumerate(pmapping_groups):
+        right_tensors = set.union(
+            set(), *[s.tensor_names for s in pmapping_groups[i + 1 :]]
+        )
+        if i == 0:
+            sim_holder.pmapping_groups = PmappingGroup.left_consolidate(
+                sim_holder.pmapping_groups,
+                right_tensors,
+            )
+            continue
+        t0 = time.time()
+        left_tensors = set.union(set(), *[s.tensor_names for s in pmapping_groups[:i]])
+        live_tensors = right_tensors
+        shared_tensors = left_tensors & sim_holder.tensor_names
+        sim_holder.pmapping_groups = sorted(
+            sim_holder.pmapping_groups, key=lambda x: len(x.mappings.data), reverse=True
+        )
+        sim_holder.pmapping_groups = PmappingGroup.right_consolidate(
+            sim_holder.pmapping_groups,
+            live_tensors,
+            shared_tensors,
+        )
+        sim_holder.pmapping_groups = PmappingGroup.combine_combineable(
+            sim_holder.pmapping_groups,
+            left_tensors | right_tensors,
+        )
+        if i > 0:
+            sim_holder.pmapping_groups = PmappingGroup.group_right(
+                sim_holder.pmapping_groups, left_tensors, drop_tags=True
+            )
+        einsum, prev_einsum = sim_holder.einsum_name, pmapping_groups[i - 1].einsum_name
+        runtime[f"{prev_einsum} → {einsum}"] = time.time() - t0
+        t0 = time.time()
+
+    n_iterations = 0
+    total_iterations = len(pmapping_groups)
+
+    def grab_sim_holder() -> (
+        tuple[dict[Compatibility, list[PmappingGroup]], str, set[str]]
+    ):
+        nonlocal n_iterations
+        n_iterations += 1
+        holder = pmapping_groups.pop(0)
+        return holder.pmapping_groups, holder.einsum_name, holder.tensor_names
+
+    if pmapping_groups:
+        left, left_einsum, left_tensors = grab_sim_holder()
+
+    partial_mapping_size = 1
+    while pmapping_groups:
+        t0 = time.time()
+        # ======================================================================
+        # Grab new Einsum from the right. Record logging data and find still
+        # tensors that will be live after this Einsum.
+        # ======================================================================
+        nbuckets.append(len(left))
+        # nmappings.append(sum(len(s.mappings.data) for s in left))
+        right, right_einsum, right_tensors = grab_sim_holder()
+        right_rank_variables = mapspace_globals.einsum2ranks[right_einsum]
+
+        partial_mapping_size += 1
+
+        live_tensors = set.union(set(), *[s.tensor_names for s in pmapping_groups])
+        shared_tensors = set(left_tensors) & set(right_tensors)
+        live_tensors_with_right = live_tensors | right_tensors
+
+        # ======================================================================
+        # Clean up the previously-combined PmappingGroups. Consolidate, combine, group
+        # them into buckets.
+        # ======================================================================
+
+        left = PmappingGroup.combine_combineable(
+            left,
+            live_tensors | right_tensors,
+        )
+
+        # Group left and right into buckets
+        left = PmappingGroup.group_left(left, right_tensors, drop_tags=True)
+
+        # ======================================================================
+        # Remove dead tensors from left and right. This happens after grouping
+        # because we only reserve space for shared tensors after it's dead. This
+        # is in case the tensor lifetime extends beyond the Einsums for which it
+        # is used.
+        # ======================================================================
+        PmappingGroup.remove_dead_tensors(
+            [s for lr in [left, right] for v in lr.values() for s in v], live_tensors
+        )
+
+        # ======================================================================
+        # Merge the left and right buckets.
+        # ======================================================================
+        combined: list[PmappingGroup] = []
+        for k in left:
+            for k_translated in get_possible_translations(
+                k,
+                pairwise_equivalent_rank_variables,
+                full_equivalent_rank_variables,
+                right_rank_variables,
+            ):
+                for a, b in itertools.product(left[k], right.get(k_translated, [])):
+                    if a.compatibility.tags.are_compatible_with(b.compatibility.tags):
+                        combined.append(
+                            a.merge_next(
+                                b,
+                                live_tensors,
+                                live_tensors_with_right,
+                                aliased_tensors,
+                                resource2capacity,
+                                delay=False,
+                            )
+                        )
+
+        if not combined:
+            raise ValueError("No match found for any group")
+
+        # ======================================================================
+        # Update left for the next iteration.
+        # =================================================================
+        left = combined
+        left_einsum = right_einsum
+        left_tensors |= right_tensors
+
+    # ======================================================================
+    # Final consolidate and group
+    # ======================================================================
+    t0 = time.time()
+    left = PmappingGroup.left_consolidate(left, None)
+    s_final = PmappingGroup.combine_combineable(left, set(), drop_tags=True)
+    assert len(s_final) == 1
+    mappings = s_final[0].mappings
+
+    return mappings

accelforge/_deprecate/_simanneal/mapspaceglobals.py

@@ -0,0 +1,298 @@
+from collections import defaultdict
+import itertools
+
+from fastfusion.frontend import arch
+from fastfusion.frontend.spec import Spec
+from fastfusion.mapper.FFM._join_pmappings.join_pmappings import PmappingGroup
+from fastfusion.mapper.FFM._join_pmappings.compatibility import Loop, Compatibility
+from fastfusion.util._frozenset import fzs
+from fastfusion.mapper.FFM._join_pmappings.join_pmappings import (
+    make_full_equivalent_rank_variables,
+)
+
+
+class MapspaceGlobals:
+    def __init__(
+        self,
+        pmapping_groups: dict[str, list[PmappingGroup]],
+        spec: Spec,
+        objective_function_cols: list[str] = None,
+        flattened_architecture: list[arch.Leaf] = None,
+    ):
+        self.pmapping_groups = pmapping_groups
+        self.einsum_names = spec.workload.einsum_names
+        self.einsum2ranks = {
+            einsum_name: spec.workload.einsums[einsum_name].rank_variables
+            for einsum_name in self.einsum_names
+        }
+        self.einsum2tensors = {
+            einsum_name: spec.workload.einsums[einsum_name].tensor_names
+            for einsum_name in self.einsum_names
+        }
+        self.tensor_names = set().union(
+            *(self.einsum2tensors[e] for e in self.einsum_names)
+        )
+        self.tensor_names_used_in_multiple_einsums = (
+            spec.workload.tensor_names_used_in_multiple_einsums
+        )
+        self.pairwise_equivalent_ranks = (
+            spec.workload.get_pairwise_equivalent_rank_variables()
+        )
+        self.full_equivalent_ranks = make_full_equivalent_rank_variables(
+            self.pairwise_equivalent_ranks
+        )
+
+        self.resource2capacity = {}
+        flattened_architecture = (
+            flattened_architecture or spec.get_flattened_architecture()
+        )
+        for l in flattened_architecture:
+            if isinstance(l, arch.Memory):
+                self.resource2capacity[l.name] = l.size
+
+        self.objective_function_cols = objective_function_cols
+        self.rank_translations = self._create_rank_translations(self.einsum2ranks)
+
+        for i, (left_id, left_sims) in enumerate(pmapping_groups.items()):
+            for j, (right_id, right_sims) in enumerate(pmapping_groups.items()):
+                if i >= j:
+                    continue
+
+                left_live = self.get_live_tensors(*self.einsum_names[: i + 1])
+                right_live = self.get_live_tensors(*self.einsum_names[j:])
+                left_tensors = self.get_tensors(self.einsum_names[i])
+                right_tensors = self.get_tensors(self.einsum_names[j])
+
+                if not (left_live & right_live):
+                    continue
+                print(f"Checking {left_id} {right_id}")
+
+                right_tilings = {
+                    s.compatibility.clear_dead_tensors(
+                        live_tensors=left_live
+                    ).clear_dead_tensors(left_tensors, keep_loops=True)
+                    for s in right_sims
+                }
+                assert right_tilings, f"R {left_id} {right_id}"
+                for s in list(left_sims):
+                    for t in self.get_possible_translations(s.compatibility, right_id):
+                        t = t.clear_dead_tensors(live_tensors=right_live)
+                        t = t.clear_dead_tensors(
+                            live_tensors=right_tensors, keep_loops=True
+                        )
+                        if t in right_tilings:
+                            break
+                    else:
+                        left_sims.remove(s)
+                assert (
+                    left_sims
+                ), f"Removed all of left {left_id} while checking right {right_id}"
+
+                left_tilings = {
+                    s.compatibility.clear_dead_tensors(
+                        live_tensors=right_live
+                    ).clear_dead_tensors(right_tensors, keep_loops=True)
+                    for s in left_sims
+                }
+                assert left_tilings, f"L {left_id} {right_id}"
+                for s in list(right_sims):
+                    for t in self.get_possible_translations(s.compatibility, left_id):
+                        t = t.clear_dead_tensors(live_tensors=left_live)
+                        t = t.clear_dead_tensors(
+                            live_tensors=left_tensors, keep_loops=True
+                        )
+                        if t in left_tilings:
+                            break
+                    else:
+                        right_sims.remove(s)
+                assert (
+                    right_sims
+                ), f"Removed all of right {right_id} while checking left {left_id}"
+
+        self.tensor2possible_loops_above = self._create_tensor2possible_loops_above()
+        self.tensor2possible_loops_above_set = {
+            k: {k2: set(v2) for k2, v2 in v.items()}
+            for k, v in self.tensor2possible_loops_above.items()
+        }
+        self.tensor2memories = self._create_tensor2memories()
+        self.einsum_tiling_2_sim = self._create_einsum_tiling_2_sim()
+        self.einsum_rank_index_to_loops = self._create_einsum_rank_index_to_loops()
+        (
+            self.compatibility2leftcompatibility,
+            self.compatibility2rightcompatibility,
+            self.leftcompatibility2tiling,
+            self.rightcompatibility2tiling,
+        ) = self._create_compatibility()
+        self.size_scale = len(self.einsum2ranks)
+        n_optimal = sum(
+            len(s.mappings.data)
+            for simlist in self.pmapping_groups.values()
+            for s in simlist
+        )
+        n_pmappings = sum(
+            s.mappings.n_pmappings
+            for simlist in self.pmapping_groups.values()
+            for s in simlist
+        )
+        self.find_pmapping_scale = n_pmappings / n_optimal
+        self.aliased_tensors = spec.workload.get_tensor_copies()
+
+    def get_live_tensors(self, *einsums: str):
+        return set.union(*(self.einsum2tensors[e] for e in einsums))
+
+    def _create_compatibility(self):
+        tiling2leftcompatibility = {}
+        tiling2rightcompatibility = {}
+
+        def tilings2compatibility(tilings: list[Compatibility], live_tensors: set[str]):
+            return {t: t.clear_dead_tensors(live_tensors=live_tensors) for t in tilings}
+
+        for i, (einsum_name, pm_group_list) in enumerate(self.pmapping_groups.items()):
+            if i > 0:
+                prev_live = self.get_live_tensors(*self.einsum_names[:i])
+                tiling2leftcompatibility[einsum_name] = tilings2compatibility(
+                    [s.compatibility for s in pm_group_list],
+                    prev_live,
+                )
+            if i < len(self.pmapping_groups) - 1:
+                next_live = self.get_live_tensors(*self.einsum_names[i + 1 :])
+                tiling2rightcompatibility[einsum_name] = tilings2compatibility(
+                    [s.compatibility for s in pm_group_list],
+                    next_live,
+                )
+
+        leftcompatibility2tiling = {}
+        rightcompatibility2tiling = {}
+        for einsum_name in self.einsum_names:
+            for src, dst in (
+                (tiling2leftcompatibility, leftcompatibility2tiling),
+                (tiling2rightcompatibility, rightcompatibility2tiling),
+            ):
+                if einsum_name not in src:
+                    continue
+                dst = dst.setdefault(einsum_name, {})
+                for k, v in src[einsum_name].items():
+                    dst.setdefault(v, []).append(k)
+        return (
+            tiling2leftcompatibility,
+            tiling2rightcompatibility,
+            leftcompatibility2tiling,
+            rightcompatibility2tiling,
+        )
+
+    def _create_einsum_tiling_2_sim(self):
+        einsum_tiling_2_sim = {}
+        for e, pm_group_list in self.pmapping_groups.items():
+            cur_sims = defaultdict(list)
+            for sim in pm_group_list:
+                cur_sims[sim.compatibility].append(sim)
+            einsum_tiling_2_sim[e] = {}
+            for t, s in cur_sims.items():
+                s = PmappingGroup.concat(s)
+                einsum_tiling_2_sim[e][t] = s
+        return einsum_tiling_2_sim
+
+    def _create_tensor2possible_loops_above(self):
+        tensor2possible_loops_above = {}
+        for einsum_name, pm_group_list in self.pmapping_groups.items():
+            tensor2possible_loops_above[einsum_name] = defaultdict(set)
+            for sim in pm_group_list:
+                for tensor in sim.compatibility.tensors:
+                    tensor2possible_loops_above[einsum_name][tensor] |= set(
+                        sim.compatibility.loops[: tensor.above_loop_index]
+                    )
+        return {
+            e: {s: list(l) for s, l in d.items()}
+            for e, d in tensor2possible_loops_above.items()
+        }
+
+    def _create_tensor2memories(self):
+        tensor2memories = {}
+        for t in self.tensor_names_used_in_multiple_einsums:
+            possible_memories = []
+            for einsum_name, pm_group_list in self.pmapping_groups.items():
+                cur_memories = set()
+                if t not in pm_group_list[0].tensor_names:
+                    continue
+                for sim in pm_group_list:
+                    tensor = sim.compatibility.get_tensor_by_name(t)
+                    cur_memories.add(tensor)
+                possible_memories.append(cur_memories)
+            if possible_memories:
+                tensor2memories[t] = list(set.intersection(*possible_memories))
+            else:
+                raise ValueError(f"No memories for {t}")
+        return tensor2memories
+
+    def _create_rank_translations(self, einsum2ranks: dict[str, set[str]]):
+        rank_translations = {}
+        for einsum_name, ranks in einsum2ranks.items():
+            translations = {einsum_name2: {} for einsum_name2 in self.einsum_names}
+            for einsum_name2, ranks2 in einsum2ranks.items():
+                for rank in ranks:
+                    equiv = self.full_equivalent_ranks[rank] & ranks2
+                    translations[einsum_name2][rank] = equiv
+            rank_translations[einsum_name] = {
+                k: {k2: list(v2) for k2, v2 in v.items()}
+                for k, v in translations.items()
+            }
+        return rank_translations
+
+    def _create_full_equivalent_ranks(
+        self, pairwise_equivalent_ranks: dict[str, set[str]]
+    ):
+        full_equivalent_ranks = {
+            k: set(v) for k, v in pairwise_equivalent_ranks.items()
+        }
+        changed = True
+        while changed:
+            changed = False
+            for r in full_equivalent_ranks:
+                for r2 in list(full_equivalent_ranks[r]):
+                    for r3 in list(full_equivalent_ranks[r2]):
+                        if r3 in full_equivalent_ranks[r]:
+                            continue
+                        changed = True
+                        full_equivalent_ranks[r].add(r3)
+        return full_equivalent_ranks
+
+    def _create_einsum_rank_index_to_loops(
+        self,
+    ) -> dict[str, dict[str, dict[int, list[Loop]]]]:
+        einsum_rank_index_to_loops = {}
+        for einsum_name, pm_group_list in self.pmapping_groups.items():
+            einsum_rank_index_to_loops[einsum_name] = {}
+            for sim in pm_group_list:
+                for rank_index, loop in enumerate(sim.compatibility.loops):
+                    x = einsum_rank_index_to_loops[einsum_name].setdefault(
+                        loop.rank_variable_name, {}
+                    )
+                    x.setdefault(rank_index, []).append(loop)
+        return einsum_rank_index_to_loops
+
+    def get_tensors(self, *einsums: str):
+        return set.union(*(self.einsum2tensors[e] for e in einsums))
+
+    def get_possible_translations(self, t: Compatibility, to_einsum: str):
+        pairwise_equivalent_ranks = self.pairwise_equivalent_ranks
+        full_equivalent_ranks = self.full_equivalent_ranks
+        right_ranks = self.einsum2ranks[to_einsum]
+
+        def translate_loop(l: Loop):
+            compatible_ranks = (
+                set.union(*(full_equivalent_ranks[n] for n in l.rank_variable_names))
+                & right_ranks
+            )
+            pairwise_compatible_ranks = (
+                set.union(
+                    *(pairwise_equivalent_ranks[n] for n in l.rank_variable_names)
+                )
+                & right_ranks
+            )
+            if len(pairwise_compatible_ranks) > 1:
+                return
+            for n in compatible_ranks:
+                yield Loop(fzs((n,)), l.bound, l.is_spatial)
+
+        for loops in itertools.product(*map(translate_loop, t.loops)):
+            yield Compatibility(loops, t.tensors, t.tags)