Merge branch 'main' into hops

Author: gilbertmike

accelforge/frontend/mapping/mapping.py

@@ -31,6 +31,7 @@
 from pydantic import ConfigDict, Discriminator, Tag, computed_field
 import sympy
 
+from accelforge.frontend.renames import EinsumName
 from accelforge.util._basetypes import (
     # Parsing helpers for the input files.
     EvalableModel,
@@ -318,7 +319,13 @@ class Loop(MappingNode):
     rank_variable: set[RankVariable] | RankVariable
     """ The rank variable(s) iterated over in this loop. This may be a
     single rank variable, or a set of rank variables if the loop is shared between
-    multiple Einsums. """
+    multiple Einsums.
+
+    NOTE FOR DEVELOPERS: If the rank variable is a set DURING PMAPPIN GENERATION, then
+    it means that the loop is a placeholder for multiple loops in the given Einsum. They
+    will be split into multiple loops later. If the rank variable
+
+    """
 
     tile_shape: EvalsTo[sympy.Symbol | sympy.Expr | int | str] = "symbol"
     """
@@ -353,6 +360,12 @@ class Loop(MappingNode):
     tile_shape for details.
     """
 
+    _einsum_to_rank_variables: dict[EinsumName, set[RankVariable]] = {}
+    """
+    If there are multiple rank variables in this loop, this dictionary says which rank
+    variables belong to which Einsum.
+    """
+
     _calculated_n_iterations: (
         Literal["symbol"] | sympy.Symbol | sympy.Expr | int | str | None
     ) = None
@@ -1073,8 +1086,11 @@ def _pop_loop_group(groups: list[list[MappingNode]]) -> list[MappingNode]:
                 # variables and assume that rank variable translation would fix it.
                 assert len(my_loop_group) == 1 or len(other_loop_group) == 1
                 has_one, may_not_have_one = my_loop_group, other_loop_group
+
+                switched = False
                 if len(has_one) != 1:
                     has_one, may_not_have_one = other_loop_group, my_loop_group
+                    switched = True
 
                 l = copy.deepcopy(has_one.pop(0))
                 l.rank_variable = (
@@ -1093,10 +1109,21 @@ def _pop_loop_group(groups: list[list[MappingNode]]) -> list[MappingNode]:
                     f"Warning. Matching loops {l} and {l2}. Need rank variable translation here."
                 )
 
+                my_rv = l.rank_variable if isinstance(l.rank_variable, set) else set([l.rank_variable])
+                other_rv = l2.rank_variable if isinstance(l2.rank_variable, set) else set([l2.rank_variable])
+                if switched:
+                    my_rv, other_rv = other_rv, my_rv
+
                 may_not_have_one.remove(l2)
                 rv = l2.rank_variable
+                for compute in self.get_nodes_of_type(Compute):
+                    for r in my_rv:
+                        l._einsum_to_rank_variables[compute.einsum] = r
                 rv = rv if isinstance(rv, set) else set([rv])
                 l.rank_variable = l.rank_variable | rv
+                for compute in other.get_nodes_of_type(Compute):
+                    for r in other_rv:
+                        l._einsum_to_rank_variables[compute.einsum] = r
                 to_add = [l]
 
             zipped_groups.append(to_add)
@@ -1542,15 +1569,18 @@ def split_reservations(self):
                 new_nodes.append(node)
         self.nodes = new_nodes
 
-    def split_loop_with_multiple_rank_variables(self, stride_and_halo: dict[str, tuple[int, int]]):
+    def split_loop_with_multiple_rank_variables(self, einsum_name: EinsumName):
         new_nodes = []
         my_rank_variables = set(x[1] for x in stride_and_halo)
         for node in self.nodes:
             if isinstance(node, Loop) and isinstance(node.rank_variable, set):
-                ranks_in_stride_and_halo = node.rank_variable & my_rank_variables
-                assert len(ranks_in_stride_and_halo) == 1, f"Expected 1 rank in stride and halo, got {len(ranks_in_stride_and_halo)}"
+                if einsum_name not in node._einsum_to_rank_variables:
+                    raise ValueError(
+                        f"Loop {node.compact_str()} has multiple rank variables. Make "
+                        f"sure to populate _einsum_to_rank_variables for this loop."
+                    )
                 new_node = copy.copy(node)
-                new_node.rank_variable = ranks_in_stride_and_halo.pop()
+                new_node.rank_variable = node._einsum_to_rank_variables[einsum_name]
                 new_nodes.append(new_node)
             else:
                 new_nodes.append(node)

accelforge/mapper/FFM/_join_pmappings/compatibility.py

@@ -382,7 +382,6 @@ def merge_next(
         self,
         right: "Compatibility",
         live_tensors: set[str],
-        mixable_ranks: dict[Rank, set[Rank]],
     ) -> "Compatibility":
         self_freed = self.clear_dead_tensors(live_tensors)
         right_freed = right.clear_dead_tensors(live_tensors)
@@ -405,9 +404,6 @@ def merge_next(
             | right_freed.reservation_indices,
         )
 
-        if mixable_ranks is not None and not joined._is_valid(mixable_ranks):
-            raise ValueError(f"Invalid rank mixing.")
-
         return joined
 
     def has_tensor(self, *tensors: TensorReservation) -> bool:
@@ -638,20 +634,6 @@ def add_n_iteration_symbols(self) -> "Compatibility":
             tensors=fzs(t.add_n_iteration_symbols() for t in self.tensors)
         )
 
-    def _is_valid(self, mixable_ranks: dict[Rank, set[Rank]]) -> bool:
-        # Mixable ranks: Ranks that may be co-iterated by a single loop.
-        ranks_at_each_loop_index = []
-        for i in range(self.n_loops):
-            ranks_at_each_loop_index.append(
-                set(t.loops[i].rank_name for t in self.tensors if i < len(t.loops))
-            )
-
-        for ranks in ranks_at_each_loop_index:
-            for r1, r2 in itertools.combinations(ranks, 2):
-                if r1 not in mixable_ranks[r2]:
-                    return False
-        return True
-
     def clear_unrelated_columns(self, mappings: pd.DataFrame) -> "Compatibility":
         my_symbols = set(self.symbols())
         for c in my_symbols:

accelforge/mapper/FFM/_join_pmappings/join_pmappings.py

@@ -1,3 +1,4 @@
+from accelforge.mapper.FFM._join_pmappings.compatibility import Compatibility
 from collections import defaultdict
 import itertools
 import logging
@@ -281,8 +282,6 @@ def join_pmappings(
             pmapping_groups, print_progress
         )
 
-    mixable_ranks = spec.workload._get_ranks_that_share_indexing_rank_variables()
-
     aliased_tensors = spec.workload.get_tensor_copies()
 
     n_mappings = {}
@@ -478,11 +477,7 @@ def grab_einsum_pmappings() -> (
                 b: PmappingGroup
                 perm_a: list[int]
                 perm_b: list[int]
-                key_check = (
-                    id(a),
-                    id(b),
-                    tuple((pa, pb) for pa, pb in zip(perm_a, perm_b)),
-                )
+                key_check = (id(a), id(b))
                 if key_check in combined_ids:
                     continue
                 combined_ids.add(key_check)
@@ -494,7 +489,6 @@ def grab_einsum_pmappings() -> (
                     compatibility_joined = compatibility_a.merge_next(
                         compatibility_b,
                         live_tensors,
-                        mixable_ranks,
                     )
                     if DO_PRINT:
                         print(
@@ -515,6 +509,8 @@ def grab_einsum_pmappings() -> (
                         aliased_tensors,
                         compatibility_joined=compatibility_joined,
                         drop_valid_reservations=drop_valid_reservations,
+                        permuted_compatibility_left=compatibility_a,
+                        permuted_compatibility_right=compatibility_b,
                         delay=DELAY,
                         _pmapping_row_filter_function=_pmapping_row_filter_function,
                         ignored_resources=ignored_resources,
@@ -575,6 +571,7 @@ def no_match_lookahead_error(
         # Look ahead to the next Einsum and see if any of our groups will not
         # be able to merge with it. If so, we can drop them immediately.
         # ======================================================================
+        lookahead_filter = True
         if lookahead_filter:
             cur_tensors = left_tensors | right_tensors
             for next_pmapping_groups in pmgroups:
@@ -589,11 +586,15 @@ def no_match_lookahead_error(
                     ).clear_tile_patterns_and_reservation_indices()
                     for c in next_pmapping_groups.pmapping_groups
                 }
-                for k in list(combined):
-                    k_cleared = k.clear_dead_tensors(
-                        next_right_tensors
-                    ).clear_tile_patterns_and_reservation_indices()
-                    if k_cleared not in next_keys:
+                for k in list[PmappingGroup](combined):
+                    perms = k.make_equivalent_permutations()
+                    perms = [
+                        p[0]
+                        .clear_dead_tensors(next_right_tensors)
+                        .clear_tile_patterns_and_reservation_indices()
+                        for p in perms
+                    ]
+                    if not any(p in next_keys for p in perms):
                         if DO_PRINT:
                             for b, _ in combined[k]:
                                 print(

accelforge/mapper/FFM/_join_pmappings/pmapping_group.py

@@ -58,6 +58,8 @@ def merge_next(
         aliased_tensors: dict[str, set[str]],
         compatibility_joined: Compatibility,
         ignored_resources: set[str],
+        permuted_compatibility_left: Compatibility,
+        permuted_compatibility_right: Compatibility,
         drop_valid_reservations: bool = True,
         delay: bool = False,
         _pmapping_row_filter_function: Callable[[pd.Series], bool] | None = None,
@@ -88,8 +90,8 @@ def merge_next(
             live_tensors_with_right,
             still_live_reservations,
             duplicated_aliased_tensors,
-            compatibility_left=self.compatibility,
-            compatibility_right=right.compatibility,
+            compatibility_left=permuted_compatibility_left,
+            compatibility_right=permuted_compatibility_right,
             compatibility_joined=compatibility_joined,
             drop_valid_reservations=drop_valid_reservations,
             _pmapping_row_filter_function=_pmapping_row_filter_function,
@@ -218,6 +220,7 @@ def _group(
         include_permutations: bool = False,
         clear_symbolic_tile_patterns: bool = False,
         try_permute_into_equivalent: bool = False,
+        # mixable_ranks: dict[Rank, set[Rank]] = None,
     ) -> (
         dict[Compatibility, list["PmappingGroup"]]
         | dict[Compatibility, list[tuple["PmappingGroup", list[int]]]]
@@ -253,6 +256,13 @@ def clear(c: Compatibility):
                     len(k.reservation_indices) == 0
                 ), f"Extra reservation indices are not empty: {k.reservation_indices}"
 
+        # if mixable_ranks is not None:
+        #     new_grouped = {}
+        #     for c, g in grouped.items():
+        #         for c2 in c.get_equivalent_compatibilities(mixable_ranks):
+        #             new_grouped.setdefault(c2, []).extend(g)
+        #     grouped = new_grouped
+
         if try_permute_into_equivalent:
             assert not include_permutations
             new_grouped = {}
@@ -339,13 +349,14 @@ def check(tensors_to_check):
 
     @staticmethod
     def group(
-        pmapping_groups: list["PmappingGroup"], live_tensors: set[str]
+        pmapping_groups: list["PmappingGroup"], live_tensors: set[str]#, mixable_ranks: dict[Rank, set[Rank]] = None
     ) -> dict[tuple[Compatibility, ...], list[tuple["PmappingGroup", list[int]]]]:
         x = PmappingGroup._group(
             pmapping_groups,
             live_tensors,
             clear_tile_patterns_and_reservation_indices=True,
             include_permutations=True,
+            # mixable_ranks=mixable_ranks,
         )
         return x
 

accelforge/model/main.py

@@ -113,9 +113,7 @@ def evaluate_mapping(
         )
 
         pmapping.split_reservations()
-        pmapping.split_loop_with_multiple_rank_variables(
-            job.stride_and_halo[einsum_name]
-        )
+        pmapping.split_loop_with_multiple_rank_variables(job.einsum_name)
         pmapping.split_tensor_holders_with_multiple_tensors()
         _add_backing_to_tensor_holders(pmapping)