accelforge 0.0.1__py3-none-any.whl

accelforge/_deprecate/_simanneal/tracking.py

@@ -0,0 +1,105 @@
+import time
+
+
+class EvaluationsScoreTracker:
+    def __init__(
+        self, max_evaluations: int, stop_at_score: float, print_period: int = 10
+    ):
+        self.max_evaluations = max_evaluations
+        self.stop_at_score = stop_at_score
+        self.evaluations = 0
+        self.score = float("inf")
+        self.history = [(0, float("inf"))]
+        self._scale_by = 1
+        self.print_period = print_period
+        self.prev_print_time = None
+        self.print_stopped_text = False
+        self.n_mappings = {}
+        self.runtime = {}
+
+    def add_evaluation(self, n_evaluations: int, best_score: float):
+        self.evaluations += n_evaluations * self._scale_by
+        self.score = min(self.score, best_score)
+        # Same score as before, remove the last entry
+        if len(self.history) > 2 and self.history[-2][1] == self.score:
+            self.history.pop(-1)
+        self.history.append((self.evaluations, self.score))
+
+        cur_time = time.time()
+        if (
+            self.prev_print_time is None
+            or cur_time - self.prev_print_time > self.print_period
+        ):
+            self.prev_print_time = cur_time
+            print(f"Evaluations: {self.evaluations}, Score: {self.score}")
+
+        if self.max_evaluations is not None and self.evaluations > self.max_evaluations:
+            self.clean_history()
+            if not self.print_stopped_text:
+                print(
+                    f"Stopping due to evaluations {self.evaluations} > {self.max_evaluations}"
+                )
+                self.print_stopped_text = True
+            raise StopIteration
+            return True
+        if self.stop_at_score is not None and self.score < self.stop_at_score:
+            self.clean_history()
+            if not self.print_stopped_text:
+                print(f"Stopping due to score {self.score} < {self.stop_at_score}")
+                self.print_stopped_text = True
+            raise StopIteration
+            return True
+        return False
+
+    def multiply_scale_by(self, scale_by: float):
+        self._scale_by *= scale_by
+
+    def __repr__(self):
+        return f"Evaluations: {self.evaluations}, Score: {self.score}"
+
+    def __str__(self):
+        return f"Evaluations: {self.evaluations}, Score: {self.score}"
+
+    def clean_history(self):
+        keep_indices = [0]
+        for i in range(1, len(self.history) - 1):
+            if (
+                self.history[i][1] != self.history[i - 1][1]
+                or self.history[i][1] != self.history[i + 1][1]
+            ):
+                keep_indices.append(i)
+        keep_indices.append(len(self.history) - 1)
+        self.history = [self.history[i] for i in keep_indices]
+
+    def merge_with(self, other: "EvaluationsScoreTracker"):
+        self.score = min(self.score, other.score)
+        self.evaluations += other.evaluations
+
+        i, j = 1, 1
+        history = [(0, float("inf"))]
+        cur_score = float("inf")
+        cur_evaluations = 0
+        while i < len(self.history) or j < len(other.history):
+            # Grab whichever has the lowest evaluations
+            if i < len(self.history) and (
+                j == len(other.history) or self.history[i][0] < other.history[j][0]
+            ):
+                new_evaluations = self.history[i][0] - self.history[i - 1][0]
+                new_score = self.history[i][1]
+                cur_evaluations += new_evaluations
+                cur_score = min(cur_score, new_score)
+                history.append((cur_evaluations, cur_score))
+                i += 1
+            elif j < len(other.history):
+                new_evaluations = other.history[j][0] - other.history[j - 1][0]
+                new_score = other.history[j][1]
+                cur_evaluations += new_evaluations
+                cur_score = min(cur_score, new_score)
+                history.append((cur_evaluations, cur_score))
+                j += 1
+        self.history = history
+        self.clean_history()
+
+    def increase_all_evaluations(self, n_evaluations: int):
+        self.evaluations += n_evaluations
+        self.history = [(e + n_evaluations, s) for e, s in self.history]

accelforge/_deprecate/_simanneal/wrappers.py

@@ -0,0 +1,218 @@
+from collections import defaultdict
+import copy
+import itertools
+import time
+from joblib import delayed
+from fastfusion._accelerated_imports import pd
+from fastfusion.frontend import arch
+from fastfusion.frontend.spec import Spec
+from fastfusion.mapper.FFM._join_pmappings.sim import PmappingGroup, Loop, Compatibility
+from fastfusion.mapper.FFM._join_pmappings.pmapping_group import (
+    PmappingDataframe,
+    is_reservation_col,
+)
+from fastfusion.mapper.simanneal.simanneal import MapspaceGlobals, _fuse_sims
+from fastfusion.mapper.simanneal.tracking import EvaluationsScoreTracker
+from fastfusion.util._frozenset import fzs
+from fastfusion.util.parallel import parallel, util
+
+
+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 paretofy(k, v):
+    return PmappingGroup(k, PmappingDataframe(pd.DataFrame(v).fillna(0)))
+
+
+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 get_pmappings_data(
+    pmapping_groups: dict[str, list[PmappingGroup]],
+    evaluations_tracker,
+    spec: Spec = None,
+    flattened_architecture: list[arch.Leaf] = None,
+):
+    resource2capacity = {}
+    flattened_architecture = flattened_architecture or spec.get_flattened_architecture()
+    for l in flattened_architecture:
+        if isinstance(l, arch.Memory):
+            resource2capacity[l.name] = l.size
+
+    pairwise_equivalent_rank_variables = (
+        spec.workload.get_pairwise_equivalent_rank_variables()
+    )
+
+    aliased_tensors = spec.workload.get_tensor_copies()
+
+    full_equivalent_rank_variables = make_full_equivalent_rank_variables(
+        pairwise_equivalent_rank_variables
+    )
+
+    return (
+        pmapping_groups,
+        evaluations_tracker,
+        spec,
+        flattened_architecture,
+        resource2capacity,
+        pairwise_equivalent_rank_variables,
+        aliased_tensors,
+        full_equivalent_rank_variables,
+    )
+
+
+def join_pmappings(
+    pmapping_groups: dict[str, list[PmappingGroup]],
+    evaluations_tracker: EvaluationsScoreTracker,
+    algorithm: str,
+    spec: Spec = None,
+    flattened_architecture: list[arch.Leaf] = None,
+) -> PmappingDataframe:
+    objective_function_cols = None
+    cols = next(iter(pmapping_groups.values()))[0].mappings.data.columns
+    if objective_function_cols is None:
+        objective_function_cols = [c for c in cols if "Total" in c]
+    keepcols = []
+
+    for pm_group_list in pmapping_groups.values():
+        for sim in pm_group_list:
+            for col in objective_function_cols:
+                if col not in sim.mappings.data.columns:
+                    sim.mappings.data[col] = 0
+            reservations = [
+                c for c in sim.mappings.data.columns if is_reservation_col(c)
+            ]
+            sim.mappings._data = sim.mappings.data[
+                objective_function_cols + keepcols + reservations
+            ]
+
+    mapspace_globals = MapspaceGlobals(
+        pmapping_groups,
+        spec,
+        objective_function_cols,
+        flattened_architecture,
+    )
+
+    n_threads = util.N_PARALLEL_PROCESSES
+    while n_threads >= 1:
+        try:
+            results_and_trackers = parallel(
+                [
+                    delayed(_fuse_sims)(
+                        mapspace_globals,
+                        n_threads=n_threads,
+                        evaluations_tracker=copy.deepcopy(evaluations_tracker),
+                        algorithm=algorithm,
+                    )
+                    for _ in range(n_threads)
+                ],
+                n_jobs=n_threads if util.PARALLELIZE else 1,
+            )
+            results = pd.concat([r[0] for r in results_and_trackers])
+            break
+        except OSError as e:
+            if n_threads == 1:
+                raise OSError("Failed to fuse pmapping_groups with 1 thread") from e
+            print(
+                f"Failed to fuse pmapping_groups with {n_threads} threads, trying with {n_threads // 2}"
+            )
+            n_threads //= 2
+
+    for t in results_and_trackers:
+        evaluations_tracker.merge_with(t[1])
+    return results

accelforge/_deprecate/_simanneal2/__init__.py

@@ -0,0 +1,7 @@
+from fastfusion.mapper.FFM.main import (
+    make_pmappings,
+    MultiEinsumPmappings,
+    Mappings,
+)
+from fastfusion.frontend.mapper.metrics import Metrics
+from .simanneal import join_pmappings