Simanneal baseline updates

Author: tanner-andrulis

accelforge/_deprecate/_simanneal/wrappers.py

@@ -201,7 +201,7 @@ def join_pmappings(
                     )
                     for _ in range(n_threads)
                 ],
-                n_jobs=get_n_parallel_jobs()
+                n_jobs=get_n_parallel_jobs(),
             )
             results = pd.concat([r[0] for r in results_and_trackers])
             break

accelforge/frontend/mapping/mapping.py

@@ -593,10 +593,14 @@ def _merge(self, other: "TensorHolder") -> "TensorHolder":
                 f"Components do not match: {self.component} != {other.component}"
             )
 
+        if self._lower != other._lower:
+            raise ValueError(f"Lower flags do not match: {self._lower} != {other._lower}")
+
         new = type(self)(
             tensors=self.tensors + other.tensors,
             component=self.component,
             component_object=self.component_object,
+            _lower=self._lower,
         )
         return new
 
@@ -1185,7 +1189,7 @@ def compact_str(self) -> str:
         if prev is not None:
             result.append(prev)
 
-        return " ".join(node.compact_str() for node in result)
+        return " ".join(f"{node.compact_str()}" for node in result)
 
     def _get_single_tensor_mapping(
         self,
@@ -1519,18 +1523,33 @@ class Mapping(Nested):
     def remove_reservations(self):
         self.nodes = [n for n in self.nodes if not isinstance(n, Reservation)]
 
-    def split_loop_with_multiple_rank_variables(self):
+    def split_reservations(self):
         new_nodes = []
         for node in self.nodes:
-            if isinstance(node, Loop) and isinstance(node.rank_variable, set):
-                for rank_variable in node.rank_variable:
+            if isinstance(node, Reservation):
+                for purpose in node.purposes:
                     new_node = copy.copy(node)
-                    new_node.rank_variable = rank_variable
+                    new_node.purposes = [purpose]
                     new_nodes.append(new_node)
             else:
                 new_nodes.append(node)
         self.nodes = new_nodes
 
+    def split_loop_with_multiple_rank_variables(self, stride_and_halo: dict[str, tuple[int, int]]):
+        new_nodes = []
+        for node in self.nodes:
+            if isinstance(node, Loop) and isinstance(node.rank_variable, set):
+                ranks_in_stride_and_halo = {
+                    r for r in node.rank_variable if r in stride_and_halo
+                }
+                assert len(ranks_in_stride_and_halo) == 1, f"Expected 1 rank in stride and halo, got {len(ranks_in_stride_and_halo)}"
+                new_node = copy.copy(node)
+                new_node.rank_variable = ranks_in_stride_and_halo.pop()
+                new_nodes.append(new_node)
+            else:
+                new_nodes.append(node)
+        self.nodes = new_nodes
+
     def split_tensor_holders_with_multiple_tensors(self):
         new_nodes = []
         for node in self.nodes:

accelforge/mapper/_simanneal2/simanneal.py

@@ -233,8 +233,9 @@ 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,
-                            mixable_ranks=self.mapspace_globals.mixable_ranks
+                            s.compatibility,
+                            live_tensors,
+                            mixable_ranks=self.mapspace_globals.mixable_ranks,
                         ),
                         drop_valid_reservations=True,
                         delay=False,
@@ -334,7 +335,7 @@ def _join_pmappings(
         aliased_tensors=spec.workload.get_tensor_copies(),
         objective_function=objective_function,
         tracker=tracker,
-        mixable_ranks=spec.workload._get_ranks_that_share_indexing_rank_variables()
+        mixable_ranks=spec.workload._get_ranks_that_share_indexing_rank_variables(),
     )
 
     mappings = []
@@ -404,13 +405,17 @@ def join_pmappings(
     pop_size_per_thread = population_size // get_n_parallel_jobs()
 
     # Multiply by the number of einsums
-    print(f'Multiplying scale by {len(pmappings.einsum2pmappings)} for number of einsums')
+    print(
+        f"Multiplying scale by {len(pmappings.einsum2pmappings)} for number of einsums"
+    )
     tracker._scale_by *= len(pmappings.einsum2pmappings)
 
     # 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')
+    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():
@@ -444,15 +449,17 @@ def join_pmappings(
             # 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)))
+                    subsets.add(
+                        tuple(subset[j] for j in range(len(subset)) if i & (1 << j))
+                    )
 
             n_total += len(subsets) * n_pmappings * get_n_tile_shapes(s)
 
     # 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')
+    print(f"Multiplying scale by {n_total / n_evaluated} for number of total pmappings")
     tracker._scale_by *= n_total / n_evaluated
 
     def parallel_join(

accelforge/mapper/_simanneal2/tracking.py

@@ -22,7 +22,7 @@ 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}')
+            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: