Simanneal baseline updates

Author: tanner-andrulis

accelforge/_deprecate/_simanneal/wrappers.py

@@ -14,7 +14,7 @@
 from accelforge.mapper.simanneal.simanneal import MapspaceGlobals, _fuse_sims
 from accelforge.mapper.simanneal.tracking import EvaluationsScoreTracker
 from accelforge.util._frozenset import fzs
-from accelforge.util.parallel import parallel, util
+from accelforge.util.parallel import parallel, get_n_parallel_jobs
 
 
 def mapping2sims(einsum_to_result: Compatibility):
@@ -188,7 +188,7 @@ def join_pmappings(
         flattened_architecture,
     )
 
-    n_threads = util.N_PARALLEL_PROCESSES
+    n_threads = get_n_parallel_jobs()
     while n_threads >= 1:
         try:
             results_and_trackers = parallel(
@@ -201,7 +201,7 @@ def join_pmappings(
                     )
                     for _ in range(n_threads)
                 ],
-                n_jobs=n_threads if util.PARALLELIZE else 1,
+                n_jobs=get_n_parallel_jobs()
             )
             results = pd.concat([r[0] for r in results_and_trackers])
             break

accelforge/frontend/workload.py

@@ -785,7 +785,7 @@ class Workload(EvalableModel):
     bits_per_value: EvalableDict[str, int | str] = EvalableDict()
     """
     Bits per value for each tensor. The workload-level bits_per_value is overridden if
-    bits_per_action is specified for any given tensor access. This is a dictionary of
+    bits_per_value is specified for any given tensor access. This is a dictionary of
     set expressions to bits per value for the tensors given by those expressions. For
     example, we may write "Inputs: 8" to set the bits per value to 8 for all input
     tensors, unless overridden.

…forge/_deprecate/_simanneal2/__init__.py accelforge/mapper/_simanneal2/__init__.pyaccelforge/_deprecate/_simanneal2/__init__.py renamed to accelforge/mapper/_simanneal2/__init__.py accelforge/_deprecate/_simanneal2/__init__.py renamed to accelforge/mapper/_simanneal2/__init__.py

@@ -25,6 +25,7 @@
 import joblib
 from accelforge.mapper.FFM._join_pmappings.compatibility import Compatibility
 from accelforge.mapper._simanneal2.tracking import EvaluationsScoreTracker
+from accelforge.util.parallel import get_n_parallel_jobs
 
 # Simulated annealing algorithm
 # -----------------------------
@@ -50,16 +51,16 @@ class MapspaceGlobals:
     def __init__(
         self,
         einsum2sims: dict[EinsumName, list[PmappingGroup]],
-        resource2capacity: dict[str, int],
+        mixable_ranks: dict[str, set[str]],
         aliased_tensors: dict[str, set[str]],
         objective_function: Callable[[pd.Series], float],
         tracker: EvaluationsScoreTracker,
     ) -> None:
         self.einsum2sims: dict[EinsumName, list[PmappingGroup]] = einsum2sims
-        self.resource2capacity: dict[str, int] = resource2capacity
         self.aliased_tensors: dict[str, set[str]] = aliased_tensors
         self.objective_function: Callable[[pd.Series], float] = objective_function
         self.tracker: EvaluationsScoreTracker = tracker
+        self.mixable_ranks: dict[str, set[str]] = mixable_ranks
 
 
 class SimAnnealMapping:
@@ -200,7 +201,11 @@ def _access_index(self, e: EinsumName, index_override: int | None = None):
         i %= len(data)
         return PmappingGroup(
             compatibility=s.compatibility,
-            mappings=PmappingDataframe(data.iloc[i : i + 1]),
+            mappings=PmappingDataframe(
+                data.iloc[i : i + 1],
+                n_total_pmappings=s.mappings.n_total_pmappings,
+                n_valid_pmappings=s.mappings.n_valid_pmappings,
+            ),
         )
 
     def get_score(self) -> float:
@@ -220,7 +225,6 @@ def get_score(self) -> float:
             def _merge_next(
                 left: PmappingGroup,
                 right: PmappingGroup,
-                apply_resource_limit: bool = True,
             ) -> PmappingGroup:
                 try:
                     return left.merge_next(
@@ -229,15 +233,12 @@ def _merge_next(
                         live_tensors_with_right=live_tensors | right_tensors,
                         aliased_tensors=self.mapspace_globals.aliased_tensors,
                         compatibility_joined=joined.compatibility.merge_next(
-                            s.compatibility, live_tensors
-                        ),
-                        resource2capacity=(
-                            self.mapspace_globals.resource2capacity
-                            if apply_resource_limit
-                            else None
+                            s.compatibility, live_tensors,
+                            mixable_ranks=self.mapspace_globals.mixable_ranks
                         ),
                         drop_valid_reservations=True,
                         delay=False,
+                        ignored_resources=set(),
                     )
                 except ValueError as err:
                     # print(err)
@@ -261,7 +262,6 @@ def _merge_next(
             joined_new = _merge_next(
                 joined,
                 s,
-                apply_resource_limit=False,
             )
             s.mappings._data = s.mappings.data.drop(columns=["_INDEX"])
             try:
@@ -314,14 +314,13 @@ def get_random_mapping(mapspace_globals: MapspaceGlobals) -> SimAnnealMapping:
             continue
 
 
-def join_pmappings(
+def _join_pmappings(
     pmapping_groups: dict[EinsumName, list[PmappingGroup]],
     spec: Spec,
-    resource2capacity: dict[str, int],
     tracker: EvaluationsScoreTracker,
     pop_size_per_thread: int,
 ) -> PmappingGroup:
-    objective = spec.mapper.ffm.metrics
+    objective = spec.mapper.metrics
     if objective == Metrics.ENERGY:
         objective_function = lambda x: x["Total<SEP>energy"]
     elif objective == Metrics.LATENCY:
@@ -330,13 +329,12 @@ def join_pmappings(
         objective_function = lambda x: x["Total<SEP>energy"] * x["Total<SEP>latency"]
     else:
         raise ValueError(f"Unknown objective {objective}")
-    # print(f'Resource2capacity: {resource2capacity}')
     mapspace_globals = MapspaceGlobals(
         einsum2sims=pmapping_groups,
-        resource2capacity=resource2capacity,
         aliased_tensors=spec.workload.get_tensor_copies(),
         objective_function=objective_function,
         tracker=tracker,
+        mixable_ranks=spec.workload._get_ranks_that_share_indexing_rank_variables()
     )
 
     mappings = []
@@ -347,9 +345,7 @@ def join_pmappings(
     print(f"Completed making initial population of {len(mappings)} mappings")
 
     i = 0
-    while True:
-        if i > 1e6:
-            break
+    while not tracker.finished():
         i += 1
         for i, m in enumerate(list(mappings)):
             try:
@@ -388,34 +384,34 @@ def get_n_tile_shapes(sim: PmappingGroup) -> int:
 
 
 def join_pmappings(
-    spec: Spec,
     pmappings: MultiEinsumPmappings,
     max_evaluations: int = 1,
     population_size=100,
     score_target: float | None = None,
 ) -> EvaluationsScoreTracker:
+    spec = pmappings.spec
     tracker = EvaluationsScoreTracker(
-        max_evaluations=max_evaluations / util.N_PARALLEL_PROCESSES,
+        max_evaluations=max_evaluations / get_n_parallel_jobs(),
         stop_at_score=None,
         print_period=1,
     )
 
+    compressed, decompress_data = compress_einsum2pmappings(pmappings.einsum2pmappings)
+
     if score_target is not None:
-        tracker.multiply_score_by(1 / score_target)
+        tracker._scale_score_by *= 1 / score_target
 
-    pop_size_per_thread = population_size // util.N_PARALLEL_PROCESSES
+    pop_size_per_thread = population_size // get_n_parallel_jobs()
 
     # Multiply by the number of einsums
-    tracker.multiply_scale_by(len(pmappings.einsum2pmappings))
+    print(f'Multiplying scale by {len(pmappings.einsum2pmappings)} for number of einsums')
+    tracker._scale_by *= len(pmappings.einsum2pmappings)
 
-    # Expected #pmappings before a Pareto-optimal one is found
-    # tracker.multiply_scale_by(pmappings._evaluated_pmappings_for_simanneal_baseline_compare() / pmappings.n_pareto_optimal_pmappings())
-    tracker.multiply_scale_by(
-        pmappings.n_evaluated_pmappings() / pmappings.n_pareto_optimal_pmappings()
-    )
-
-    # Normalize to the speed of the intra-Einsum pmapper
-    tracker.multiply_scale_by(1 / pmappings.n_evaluated_pmappings())
+    # We allow FFM to query n_pareto_optimal_pmappings mappings from the mapspace, so we
+    # scale by 1 / n_pareto_optimal_pmappings to get simanneal 1 evaluation = same
+    # number of mappings as FFM
+    print(f'Multiplying scale by {1 / pmappings.n_pareto_optimal_pmappings()} for Pareto-optimal mapspace size')
+    tracker._scale_by *= 1 / pmappings.n_pareto_optimal_pmappings()
 
     for einsum_name, einsum_pmappings in pmappings.einsum2pmappings.items():
         total = sum(len(p.mappings.data) for p in einsum_pmappings)
@@ -426,50 +422,56 @@ def join_pmappings(
         if total == 0:
             raise ValueError(f"Einsum {einsum_name} has no pmappings")
 
-    print(f"TODO: Populate PmappingGroups with all permutations")
-
-    compressed, decompress_data = compress_einsum2pmappings(pmappings.einsum2pmappings)
-
     permuted = {}
+    n_total = 1
+    n_evaluated = 1
     for einsum_name, einsum_sims in compressed.items():
         for s in einsum_sims:
+            n_pmappings = s.mappings.n_total_pmappings
+            n_evaluated += n_pmappings
+            subsets = set()
+
+            # Count all permutations as separate choices
             for c_perm, _ in s.compatibility.make_equivalent_permutations():
                 permuted.setdefault(einsum_name, []).append(
                     PmappingGroup(
                         compatibility=c_perm,
                         mappings=s.mappings,
                     )
                 )
+                subsets.add(tuple(c_perm.loops))
+
+            # Count mappings with fewer loops as separate choices
+            for subset in list(subsets):
+                for i in range(2 ** len(subset)):
+                    subsets.add(tuple(subset[j] for j in range(len(subset)) if i & (1 << j)))
 
-    tile_shapes = [
-        get_n_tile_shapes(s)
-        for pmapping_groups in compressed.values()
-        for s in pmapping_groups
-    ]
+            n_total += len(subsets) * n_pmappings * get_n_tile_shapes(s)
 
-    # average_tile_shapes = sum(tile_shapes) / len(tile_shapes)
-    # print(f"Average tile shapes: {average_tile_shapes}")
-    # tracker.multiply_scale_by(average_tile_shapes)
+    # The pmapper does mappings with varied permutations, subsets of loops, and tile
+    # shapes in one shot, which won't be the case if a simulated annealing mapper is
+    # used to select the top part of the mapping. So scale to account for the number
+    # that we get for free.
+    print(f'Multiplying scale by {n_total / n_evaluated} for number of total pmappings')
+    tracker._scale_by *= n_total / n_evaluated
 
     def parallel_join(
         permuted: dict[EinsumName, list[PmappingGroup]],
         spec: Spec,
-        resource2capacity: dict[str, int],
         tracker: EvaluationsScoreTracker,
         pop_size_per_thread: int,
     ) -> EvaluationsScoreTracker:
-        join_pmappings(permuted, spec, resource2capacity, tracker, pop_size_per_thread)
+        _join_pmappings(permuted, spec, tracker, pop_size_per_thread)
         return tracker
 
     trackers = util.parallel(
         joblib.delayed(parallel_join)(
             permuted,
             spec,
-            pmappings.resource2capacity,
             tracker,
             pop_size_per_thread,
         )
-        for _ in range(util.N_PARALLEL_PROCESSES)
+        for _ in range(get_n_parallel_jobs())
     )
 
     t0 = trackers[0]

…orge/_deprecate/_simanneal2/simanneal.py …celforge/mapper/_simanneal2/simanneal.pyaccelforge/_deprecate/_simanneal2/simanneal.py renamed to accelforge/mapper/_simanneal2/simanneal.py accelforge/_deprecate/_simanneal2/simanneal.py renamed to accelforge/mapper/_simanneal2/simanneal.py

@@ -20,6 +20,9 @@ def __init__(
 
     def add_evaluation(self, n_evaluations: int, best_score: float):
         self.evaluations += n_evaluations * self._scale_by
+        new_score = best_score * self._scale_score_by
+        if new_score < self.score:
+            print(f'New score {new_score} after evaluation {self.evaluations}')
         self.score = min(self.score, best_score * self._scale_score_by)
         # Same score as before, remove the last entry
         if len(self.history) > 2 and self.history[-2][1] == self.score:
@@ -32,7 +35,6 @@ def add_evaluation(self, n_evaluations: int, best_score: float):
             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()
@@ -59,9 +61,6 @@ def finished(self):
         )
         return enough_evaluations or enough_score
 
-    def multiply_scale_by(self, scale_by: float):
-        self._scale_by *= scale_by
-
     def multiply_score_by(self, scale_by: float):
         self._scale_score_by *= scale_by