accelforge 0.0.1__py3-none-any.whl

accelforge/_deprecate/compatibility_util.py

@@ -0,0 +1,181 @@
+import itertools
+
+from collections.abc import Iterable, Set
+
+from fastfusion.frontend.spec import Spec
+from fastfusion.frontend.workload import EinsumName, TensorName
+from fastfusion.mapper.FFM._join_pmappings.compatibility import Compatibility
+from fastfusion.mapper.FFM._join_pmappings.sim import SIM
+
+
+DO_PRINT = False
+
+
+def myprint(*args, **kwargs):
+    if DO_PRINT:
+        print(*args, **kwargs)
+
+
+def sims2untiled_compats(
+    einsum2sims: dict[EinsumName, Iterable[SIM]],
+) -> dict[EinsumName, set[Compatibility]]:
+    return {
+        einsum_name: {sim.compatibility.clear_loop_bounds() for sim in sims}
+        for einsum_name, sims in einsum2sims.items()
+    }
+
+
+def join_compatibilities(
+    einsum2compatibilities: dict[EinsumName, Iterable[Compatibility]],
+    spec: Spec = None,
+) -> dict[EinsumName, set[Compatibility]]:
+    """
+    Return dict from Einsum name to compatibilities (without tile shape)
+    that will ever contribute to full mappings.
+
+    CONTRACT FOR MAPPINGS GETTING TO THIS POINT: see `join_pmappings.join_sims`
+    """
+    for einsum_name, compats in einsum2compatibilities.items():
+        if sum(len(c) for c in compats) == 0:
+            raise ValueError(f"No pmappings for {einsum_name}")
+
+    if len(einsum2compatibilities) == 0:
+        raise ValueError("Nothing to join")
+
+    for einsum_name, per_einsum_compats in einsum2compatibilities.items():
+        if not per_einsum_compats:
+            raise ValueError(f"No compatibility for {einsum_name}")
+
+    compat2einsum2original: dict[
+        Compatibility, dict[EinsumName, set[Compatibility]]
+    ] = {}
+    for einsum_name, per_einsum_compats in einsum2compatibilities.items():
+        for compat in per_einsum_compats:
+            einsum2original = compat2einsum2original.setdefault(compat, {})
+            original = einsum2original.setdefault(einsum_name, set())
+            original.add(compat)
+
+    compatibilities = list(einsum2compatibilities.items())
+
+    einsum2tensor_names = {
+        einsum_name: spec.workload.einsums[einsum_name].tensor_names
+        for einsum_name in einsum2compatibilities
+    }
+
+    einsum2important_compatibilities = {}
+
+    # while-loop states
+    assert len(compatibilities) > 0
+    left_einsum, all_left_compats = compatibilities.pop(0)
+    left_tensors = einsum2tensor_names[left_einsum]
+
+    while compatibilities:
+        right_einsum, all_right_compats = compatibilities.pop(0)
+
+        right_tensors = einsum2tensor_names[right_einsum]
+        live_tensors = set.union(
+            set(), *(einsum2tensor_names[e] for e, _ in compatibilities)
+        )
+
+        grouped_left_compats = group_left(all_left_compats, right_tensors)
+        grouped_right_compats = group_right(all_right_compats, left_tensors)
+
+        combined = combine_left_and_right_compats(
+            compat2einsum2original,
+            grouped_left_compats,
+            grouped_right_compats,
+            live_tensors,
+        )
+
+        if DO_PRINT:
+            print_reverse_unmatched(grouped_left_compats, grouped_right_compats)
+
+        if not combined:
+            raise ValueError("No match found for any group")
+
+        # update while-loop states
+        all_left_compats = combined
+        left_einsum = right_einsum
+        left_tensors |= right_tensors
+
+    einsum2important_compatibilities: dict[EinsumName, set[Compatibility]] = {}
+    for compat in combined:
+        for einsum, original in compat2einsum2original[compat].items():
+            important_compats = einsum2important_compatibilities.setdefault(
+                einsum, set()
+            )
+            important_compats.update(original)
+    return einsum2important_compatibilities
+
+
+def combine_left_and_right_compats(
+    compat2einsum2original: dict[Compatibility, dict[EinsumName, set[Compatibility]]],
+    grouped_left_compats: dict[Compatibility, Iterable[Compatibility]],
+    grouped_right_compats: dict[Compatibility, Iterable[Compatibility]],
+    live_tensors: set[TensorName],
+):
+    combined: list[Compatibility] = []
+    for left_key, left_compats in grouped_left_compats.items():
+        myprint(f"Left key {left_key}")
+
+        compatible_right_compats = grouped_right_compats.get(left_key, [])
+
+        if len(compatible_right_compats) == 0:
+            if DO_PRINT:
+                for l in left_compats:
+                    print(f"\tNo match for {l}")
+            continue
+
+        for l, r in itertools.product(left_compats, compatible_right_compats):
+            if l.tags.are_compatible_with(r.tags):
+                merged = l.merge_next(r, live_tensors)
+                combined.append(merged)
+
+                einsum2original = compat2einsum2original.setdefault(merged, {})
+
+                left_einsum2original = compat2einsum2original[l]
+                right_einsum2original = compat2einsum2original[r]
+
+                einsums = set(left_einsum2original) | set(right_einsum2original)
+                for einsum in einsums:
+                    einsum2original.setdefault(einsum, set()).update(
+                        left_einsum2original.get(einsum, set())
+                        | right_einsum2original.get(einsum, set())
+                    )
+
+                myprint(f"\t{l}\n\t<-->\n\t{r}")
+                myprint(f"\t-->\n\t{merged}")
+    return combined
+
+
+def print_reverse_unmatched(
+    grouped_left_compats: dict[Compatibility, Iterable[Compatibility]],
+    grouped_right_compats: dict[Compatibility, Iterable[Compatibility]],
+):
+    for right_key, right_compats in grouped_right_compats.items():
+        if right_key not in grouped_left_compats:
+            for r in right_compats:
+                print(f"\tREVERSE: No match for {r} using {right_key}")
+
+
+def group_left(
+    left_compatibilities: Iterable[Compatibility],
+    right_tensors: Set[TensorName],
+) -> dict[Compatibility, set[Compatibility]]:
+    grouped_compats = {}
+    for compat in left_compatibilities:
+        key = compat.clear_dead_tensors(right_tensors, keep_loops=True, drop_tags=True)
+        grouped_compats.setdefault(key, set()).add(compat)
+    return grouped_compats
+
+
+def group_right(
+    right_compatibilities: Iterable[Compatibility],
+    left_tensors: Set[TensorName],
+) -> dict[Compatibility, set[Compatibility]]:
+    grouped_compats = {}
+    for compat in right_compatibilities:
+        key = compat.clear_dead_tensors(left_tensors, keep_loops=True, drop_tags=True)
+        for per_loop_key in key.all_n_loops():
+            grouped_compats.setdefault(per_loop_key, set()).add(compat)
+    return grouped_compats

accelforge/_deprecate/layerdeduplication/__init__.py

@@ -0,0 +1,2 @@
+from .group_similar_einsums import group_similar_einsums
+from .grouped_einsums import GroupOfSimilarEinsums, Id, Name

accelforge/_deprecate/layerdeduplication/group_similar_einsums.py

@@ -0,0 +1,160 @@
+from collections import defaultdict
+from collections.abc import Iterable
+from itertools import permutations, product
+
+from pytimeloop.bindings.looptree import LooptreeWorkload, LooptreeDependencyAnalyzer
+
+from pytimeloop.looptree.mapping_utilities import get_intermediate_tensors
+from fastfusion.util._frozenset import fzs
+
+from .grouped_einsums import GroupOfSimilarEinsums, Id
+
+
+def group_similar_einsums(
+    einsum_ids: Iterable[int],
+    workload: LooptreeWorkload,
+    analyzer: LooptreeDependencyAnalyzer,
+) -> list[GroupOfSimilarEinsums[Id]]:
+    """
+    Groups similar Einsums in `einsum_ids`.
+    """
+    grouped_einsums: list[GroupOfSimilarEinsums[Id]] = []
+    for einsum_id in einsum_ids:
+        found = False
+        for einsum_group in grouped_einsums:
+            einsum_ref_id = einsum_group.reference_einsum
+            rank_renaming, tensor_renaming = is_equivalent(
+                einsum_ref_id, einsum_id, workload, analyzer
+            )
+            if rank_renaming is not None:
+                einsum_group.add_similar_einsum(
+                    einsum_id, rank_renaming, tensor_renaming
+                )
+                found = True
+                break
+
+        if not found:
+            grouped_einsums.append(GroupOfSimilarEinsums(einsum_id, workload))
+    return grouped_einsums
+
+
+def is_equivalent(
+    einsum_id1: int,
+    einsum_id2: int,
+    workload: LooptreeWorkload,
+    analyzer: LooptreeDependencyAnalyzer,
+) -> tuple[dict[int, int], dict[int, int]]:
+    """
+    Determines whether two Einsums are equivalent in tensor shapes and
+    tensor indexing expressions.
+
+    If the two Einsums are equivalent, the rank and tensor renamings are
+    returned.
+
+    Returns:
+      If the two Einsums are equivalent, the function returns two dicts,
+      `rank_renaming` and `tensor_renaming`, representing how to rename
+      ranks (tensors) of `einsum_id1` to `einsum_id2`.
+
+      Otherwise, a tuple `(None, None)` is returned.
+    """
+    einsum1_ranks = workload.einsum_ospace_dimensions(einsum_id1)
+    einsum2_ranks = workload.einsum_ospace_dimensions(einsum_id2)
+
+    if len(einsum1_ranks) != len(einsum2_ranks):
+        return None, None
+
+    einsum1_input_tensors = workload.tensors_read_by_einsum(einsum_id1)
+    einsum1_output_tensor = workload.tensors_written_by_einsum(einsum_id1)
+    einsum2_input_tensors = workload.tensors_read_by_einsum(einsum_id2)
+    einsum2_output_tensor = workload.tensors_written_by_einsum(einsum_id2)
+
+    if einsum1_output_tensor is None:
+        einsum1_output_tensor = set()
+    if einsum2_output_tensor is None:
+        einsum2_output_tensor = set()
+
+    intermediate_tensors = get_intermediate_tensors(workload)
+
+    all_tensor_properties = []
+    all_tensors = [
+        (einsum1_input_tensors, einsum1_output_tensor),
+        (einsum2_input_tensors, einsum2_output_tensor),
+    ]
+    for input_tensors, output_tensors in all_tensors:
+        tensor_properties = defaultdict(set)
+        for tensor in input_tensors:
+            tensor_properties[tensor].add("input")
+        for tensor in output_tensors:
+            tensor_properties[tensor].add("output")
+        for tensor in tensor_properties:
+            if tensor in intermediate_tensors:
+                tensor_properties[tensor].add("intermediate")
+        tensor_properties = {
+            tensor: fzs(properties) for tensor, properties in tensor_properties.items()
+        }
+        all_tensor_properties.append(tensor_properties)
+
+    property_to_tensors = defaultdict(lambda: (set(), set()))
+    for i, tensor_properties in enumerate(all_tensor_properties):
+        for tensor, property in tensor_properties.items():
+            tensor_sets = property_to_tensors[property]
+            tensor_sets[i].add(tensor)
+
+    # Check if we can rename tensors in einsum1 to einsum2
+    for tensor_renaming in tensor_renamings(property_to_tensors):
+        # Check if we can rename einsum1 ranks to create einsum2
+        for renamed_ranks in permutations(einsum2_ranks):
+            rank_renaming = {r1: r2 for r1, r2 in zip(einsum1_ranks, renamed_ranks)}
+            if not _shape_is_equivalent(rank_renaming, workload):
+                continue
+
+            if not _dependency_is_equivalent(
+                einsum_id1, einsum_id2, rank_renaming, tensor_renaming, analyzer
+            ):
+                continue
+
+            return rank_renaming, tensor_renaming
+    return None, None
+
+
+def tensor_renamings(property_to_tensors):
+    for tensors_of_1, tensors_of_2 in property_to_tensors.values():
+        if len(tensors_of_1) != len(tensors_of_2):
+            return
+
+    all_tensors_of_1 = [
+        t for tensors_of_1, _ in property_to_tensors.values() for t in tensors_of_1
+    ]
+    permutations_of_tensor_2_by_property = []
+    for _, tensors_of_2 in property_to_tensors.values():
+        permutations_of_tensor_2_by_property.append(permutations(tensors_of_2))
+    for permutation_of_2 in product(*permutations_of_tensor_2_by_property):
+        permutation_of_2 = tuple(t for tupl in permutation_of_2 for t in tupl)
+        renaming = dict(zip(all_tensors_of_1, permutation_of_2))
+        yield renaming
+
+
+def _shape_is_equivalent(rank_renaming, workload):
+    for r1, r2 in rank_renaming.items():
+        r1_shape = workload.get_rank_shape(r1)
+        r2_shape = workload.get_rank_shape(r2)
+        if r1_shape != r2_shape:
+            return False
+    return True
+
+
+def _dependency_is_equivalent(
+    einsum_id1, einsum_id2, rank_renaming, tensor_renaming, analyzer
+):
+    for t1, t2 in tensor_renaming.items():
+        for r1, r2 in rank_renaming.items():
+            r1_relevant_to_t1 = analyzer.einsum_dim_is_directly_relevant_to_tensor(
+                einsum_id1, r1, t1
+            )
+            r2_relevant_to_t2 = analyzer.einsum_dim_is_directly_relevant_to_tensor(
+                einsum_id2, r2, t2
+            )
+            if r1_relevant_to_t1 != r2_relevant_to_t2:
+                return False
+    return True

accelforge/_deprecate/layerdeduplication/grouped_einsums.py

@@ -0,0 +1,84 @@
+from collections.abc import Iterable
+
+from bindings.looptree import LooptreeWorkload
+
+
+type Id = int
+type Name = str
+
+
+class GroupOfSimilarEinsums[IdOrName: Id | Name]:
+    def __init__(self, reference_einsum: Id, workload: LooptreeWorkload):
+        self.reference_einsum = reference_einsum
+        self.workload = workload
+        self.similar_einsums_to_renaming = {}
+        self.in_id = True
+
+    def add_similar_einsum(
+        self,
+        similar_einsum: IdOrName,
+        rank_renaming: IdOrName,
+        tensor_renaming: IdOrName,
+    ):
+        self.similar_einsums_to_renaming[similar_einsum] = (
+            rank_renaming,
+            tensor_renaming,
+        )
+
+    @property
+    def similar_einsums(self) -> Iterable[IdOrName]:
+        return self.similar_einsums_to_renaming.keys()
+
+    @property
+    def get_renaming(
+        self, other_einsum: IdOrName
+    ) -> tuple[dict[IdOrName, IdOrName], dict[IdOrName, IdOrName]]:
+        """Returns iterable over tuple `(rank_renaming, tensor_renaming)`"""
+        try:
+            return self.similar_einsums_to_renaming[other_einsum]
+        except Exception as e:
+            e.add_note(f"{other_einsum} not in group of similar Einsums.")
+            raise
+
+    @property
+    def similar_einsums_and_renamings(
+        self,
+    ) -> Iterable[
+        tuple[IdOrName, tuple[dict[IdOrName, IdOrName], dict[IdOrName, IdOrName]]]
+    ]:
+        """
+        Returns iterable over tuple `(similar_einsum, renaming)`
+        where `renaming` itself is `(rank_renaming, tensor_renaming).
+        """
+        return self.similar_einsums_and_renamings.items()
+
+    def convert_id_to_name(self) -> "GroupOfSimilarEinsums[Name]":
+        einsum_id_to_name = self.workload.EinsumIdToName()
+        tensor_id_to_name = self.workload.DataSpaceIdToName()
+        rank_id_to_name = self.workload.DimensionIdToName()
+
+        grouped_einsums_in_name = GroupOfSimilarEinsums(
+            einsum_id_to_name[self.reference_einsum], self.workload
+        )
+        self.in_id = False
+
+        similar_einsums_to_renamings = self.get_einsums_similar_to_reference(
+            self.reference_einsum
+        )
+        for einsum_id, renaming in similar_einsums_to_renamings.items():
+            rank_renaming, tensor_renaming = renaming
+            rank_renaming_in_names = {
+                rank_id_to_name[k]: rank_id_to_name[v] for k, v in rank_renaming.items()
+            }
+            tensor_renaming_in_names = {
+                tensor_id_to_name[k]: tensor_id_to_name[v]
+                for k, v in tensor_renaming.items()
+            }
+            grouped_einsums_in_name.add_einsum_similar_to_reference(
+                einsum_id_to_name[self.reference_einsum],
+                einsum_id_to_name[einsum_id],
+                rank_renaming_in_names,
+                tensor_renaming_in_names,
+            )
+
+        return grouped_einsums_in_name

accelforge/_deprecate/mapping_filter_tags/__init__.py

@@ -0,0 +1,2 @@
+from .ffmt import get_ffmt_tag
+from .onesplit import get_one_split_tag

accelforge/_deprecate/mapping_filter_tags/ffmt.py

@@ -0,0 +1,212 @@
+from fastfusion.frontend.mapping import Loop, Temporal
+from fastfusion.mapper.FFM.deprecate_maybe.tags import Tags
+
+from .util import get_fused_loops_per_tensor
+
+
+FFMT_VALID = "FFMT_VALID"
+FFMT_WEIGHT_UNTILED = "FFMT_WEIGHT_UNTILED"
+FFMT_WEIGHT_TILED = "FFMT_WEIGHT_TILED"
+
+
+def get_ffmt_tag(compatibility):
+    return get_ffmt_matmul_tag(compatibility)
+    if "Matmul" in einsum_name:
+        return get_ffmt_matmul_tag(compatibility)
+    else:
+        return get_ffmt_mha_tag(compatibility)
+
+
+def get_ffmt_matmul_tag(compatibility):
+    # FFMT is:
+    # - [input | output, weight]
+    # If there's >1 fused loop, they must be above the same number of loops
+    tensors = [s for s in compatibility.tensors if s.resource_name != "MainMemory"]
+    if len(tensors) <= 1:
+        return Tags((FFMT_VALID,))
+
+    allowed_n_loops = [
+        (0, 0),
+        (1, 1),
+        (1, 2),
+    ]
+
+    # If there's a B or H fused loop, add one to the allowed n_loops
+    for rank_var in "b", "h":
+        if any(rank_var in l.rank_variable for l in compatibility.loops):
+            allowed_n_loops = [(x + 1, y + 1) for x, y in allowed_n_loops]
+
+    if tuple(sorted(s.above_loop_index for s in tensors)) in [
+        (0, 0),
+        (1, 1),
+        (1, 2),
+    ]:
+        return Tags((FFMT_VALID,))
+    raise ValueError()
+
+
+def get_ffmt_mha_tag(compatibility):
+    tensors = [s for s in compatibility.tensors if s.resource_name != "MainMemory"]
+    if len(compatibility.loops) == 0:
+        return Tags((FFMT_VALID,))
+
+    # Loops have to be in the order (b, h)
+    if len(compatibility.loops) == 1:
+        return Tags((FFMT_INVALID,))
+
+    if len(set(s.above_loop_index for s in tensors)) > 1:
+        raise ValueError()
+    return Tags((FFMT_VALID,))
+
+    for tensors in compatibility.tensors:
+        if tensor.resource_name == "MainMemory":
+            continue
+        unique_loops.add(tensor.above_loop_index)
+
+    if len(unique_loops) == 0:
+        return Tags()  # unfused is compatible with anything
+
+    untiled_fused = len(unique_loops) == 1 and next(iter(unique_loops)) == 0
+    if untiled_fused:
+        return Tags((FFMT_VALID,))
+
+    min_weight_idx, max_weight_idx, max_non_weight_idx = float("inf"), 0, 0
+    max_weight_idx = 0
+    for tensor, n_loops in tensor_to_n_fused_loops.items():
+        is_weight = "Filter" in tensor.name
+        if is_weight:
+            min_weight_idx = min(min_weight_idx, n_loops)
+            max_weight_idx = max(max_weight_idx, n_loops)
+        else:
+            max_non_weight_idx = max(max_non_weight_idx, n_loops)
+
+    weight_untiled = min_weight_idx == 0 and max_weight_idx == 0
+    if weight_untiled:
+        return Tags((FFMT_VALID, FFMT_WEIGHT_UNTILED))
+    elif min_weight_idx >= max_non_weight_idx:
+        return Tags((FFMT_VALID, FFMT_WEIGHT_TILED))
+    raise ValueError()
+
+
+def get_ffmt_mha_tag(pmapping):
+    einsum_name = pmapping[-1].einsum_name
+    B, H, M, F, P, G, E, D, C, J = "bhmfpgedcj"
+    EINSUM_NAME_TO_REDUCED_RANK_OUTPUT_RANK = {
+        "Q": [D, E],
+        "K": [D, E],
+        "V": [D, F],
+        "QK": [E, P],
+        "AV": [P, F],
+        "Z": [F, G],
+        "FFA": [G, C],
+        "FFB": [C, J],
+    }
+
+    rank_var_permutation = []
+    for node in pmapping:
+        if isinstance(node, Loop):
+            if not isinstance(node, Temporal):
+                raise RuntimeError(
+                    "get_ffmt_mha_tag should not be used for "
+                    "anything other than Snowcat"
+                )
+            rank_var_permutation.append(node.rank_variable)
+
+    tensor_to_n_fused_loops = get_fused_loops_per_tensor(
+        pmapping, intermediate_tensors, "MainMemory"
+    )
+    unfused = all(
+        n is None
+        for t, n in tensor_to_n_fused_loops.items()
+        if t in intermediate_tensors
+    )
+    if einsum_name not in EINSUM_NAME_TO_REDUCED_RANK_OUTPUT_RANK:
+        if unfused:
+            return Tags((FFMT_VALID,))
+        raise ValueError()
+
+    reduced_rank, output_rank = EINSUM_NAME_TO_REDUCED_RANK_OUTPUT_RANK[einsum_name]
+
+    EINSUM_NAME_TO_INPUT_OUTPUT_TENSORS = {
+        "Q": ["I_I_to_Q_K_V", "Q_Q_to_QK"],
+        "K": ["I_I_to_Q_K_V", "K_K_to_QK"],
+        "V": ["I_I_to_Q_K_V", "V_V_to_AV"],
+        "QK": ["Q_Q_to_QK", "QK_QK_to_AV"],
+        "AV": ["QK_QK_to_AV", "AV_AV_to_Z"],
+        "Z": ["AV_AV_to_Z", "Z_Z_to_FFA"],
+        "FFA": ["Z_Z_to_FFA", "FFA_FFA_to_FFB"],
+        "FFB": ["FFA_FFA_to_FFB", "FFB_FFB_to_n"],
+    }
+
+    input_tensor, output_tensor = EINSUM_NAME_TO_INPUT_OUTPUT_TENSORS[einsum_name]
+    input_output_tensors = {input_tensor, output_tensor}
+
+    min_weight_idx = float("inf")
+    max_weight_idx = 0
+    max_non_weight_idx = 0
+    first, last = True, True
+    for tensor, n_loops in tensor_to_n_fused_loops.items():
+        if tensor.name == input_tensor and n_loops is not None:
+            first = False
+        if tensor.name == output_tensor and n_loops is not None:
+            last = False
+
+        is_weight = tensor.name not in input_output_tensors
+        if is_weight:
+            min_weight_idx = min(min_weight_idx, n_loops)
+            max_weight_idx = max(max_weight_idx, n_loops)
+        else:
+            max_non_weight_idx = max(max_non_weight_idx, n_loops)
+
+    # Rank variable order and the n_loops for (input, output)
+    prefix_choices = [([B, H], (2, 2))]
+
+    # Rank variable order and the n_loops for (input, output)
+    extra_rank_choices = [
+        ([M], (1, 1)),
+    ]
+    if first:
+        if output_rank is not None:
+            extra_rank_choices.append(([M, output_rank], (1, 2)))
+        if reduced_rank is not None and output_rank is not None:
+            extra_rank_choices.append(([M, output_rank, reduced_rank], (3, 2)))
+        if output_rank is None and reduced_rank is not None:
+            extra_rank_choices.append(([M, reduced_rank], (2, 1)))
+    elif last:
+        if output_rank is not None:
+            extra_rank_choices.append(([M, output_rank], (1, 2)))
+    else:
+        if reduced_rank is not None:
+            extra_rank_choices.append(([M, reduced_rank], (2, 1)))
+
+    for prefix_permutation, prefix_n_loops in prefix_choices:
+        for extra_permutation, extra_n_loops in extra_rank_choices:
+            permutation = prefix_permutation + extra_permutation
+            input_n_loops = prefix_n_loops[0] + extra_n_loops[0]
+            output_n_loops = prefix_n_loops[1] + extra_n_loops[1]
+            untiled_weight_idx = len(prefix_permutation)
+
+            permutation_matches = True
+            for rank_var, ref_rank_var in zip(rank_var_permutation, permutation):
+                if rank_var != ref_rank_var:
+                    permutation_matches = False
+                    break
+
+            if not permutation_matches:
+                continue
+
+            if tensor_to_n_fused_loops[input_tensor] != input_n_loops:
+                continue
+            if tensor_to_n_fused_loops[output_tensor] != output_n_loops:
+                continue
+
+            weight_untiled = (
+                min_weight_idx == untiled_weight_idx
+                and max_weight_idx == untiled_weight_idx
+            )
+            if weight_untiled:
+                return Tags((FFMT_VALID, FFMT_WEIGHT_UNTILED))
+            elif min_weight_idx >= max_non_weight_idx:
+                return Tags((FFMT_VALID, FFMT_WEIGHT_TILED))
+
+    raise ValueError()

accelforge/_deprecate/mapping_filter_tags/onesplit.py

@@ -0,0 +1,24 @@
+from fastfusion.mapper.FFM.deprecate_maybe.tags import Tags
+from fastfusion.mapper.FFM._join_pmappings.compatibility import Compatibility
+
+
+ONE_SPLIT = "ONE_SPLIT"
+NOT_ONE_SPLIT = "NOT_ONE_SPLIT"
+
+
+def get_one_split_tag(compatibility: Compatibility) -> Tags:
+    # TODO
+    unique_loops = set()
+    for tensor in compatibility.tensors:
+        if tensor.resource_name == "MainMemory":
+            continue
+        unique_loops.add(tensor.above_loop_index)
+
+    if len(unique_loops) == 0:
+        return Tags()  # unfused is compatible with anything
+
+    # Fused with both sides. Make sure that the number of loops is the same.
+    if len(unique_loops) > 1:
+        return Tags(("INVALID",))
+
+    return Tags((ONE_SPLIT, f"FUSED_LOOPS={next(iter(unique_loops))}"))

accelforge/_deprecate/mapping_filter_tags/util.py

@@ -0,0 +1,24 @@
+from fastfusion.frontend.mapping import Reservation, Loop, Mapping
+
+
+def get_fused_loops_per_tensor(
+    pmapping: Mapping, intermediate_tensors, non_fused_memory
+):
+    """
+    Returns a dictionary mapping tensor to number of fused loops or None
+    if unfused (backed in non_fused_memory).
+    """
+    tensor_to_n_fused_loops = {}
+    n_loops = 0
+    for node in pmapping.nodes:
+        if isinstance(node, Reservation):
+            tensor = node.tensor
+            if tensor not in intermediate_tensors or tensor in tensor_to_n_fused_loops:
+                continue
+            if node.component == non_fused_memory:
+                tensor_to_n_fused_loops[tensor] = None
+            else:
+                tensor_to_n_fused_loops[tensor] = n_loops
+        elif isinstance(node, Loop):
+            n_loops += 1
+    return tensor_to_n_fused_loops