accelforge 0.0.1__py3-none-any.whl

accelforge/mapper/FFM/_make_pmappings/make_pmappings_from_templates/make_pmappings_from_templates.py

@@ -0,0 +1,407 @@
+import copy
+import math
+import pandas as pd
+from uuid import UUID
+from collections import defaultdict
+
+import sympy
+from accelforge.frontend.mapping import Loop, Mapping, Spatial, Temporal
+from accelforge.frontend.workload import EinsumName
+from accelforge.mapper.FFM._join_pmappings.compatibility import (
+    Compatibility,
+)
+from accelforge.mapper.FFM._join_pmappings.pmapping_dataframe import (
+    MAPPING_COLUMN,
+    PmappingDataframe,
+    col2nameloop,
+    col_used_in_pareto,
+    is_reservation_col,
+    makepareto,
+    tensor2col,
+    col2nameloop,
+    is_reservation_col,
+    nameloop2col,
+)
+
+from accelforge.frontend.mapper.metrics import Metrics
+from accelforge.mapper.FFM._make_pmappings.make_pmappings_from_templates.make_tile_shapes import (
+    make_tile_shapes,
+    IMPERFECT,
+)
+from accelforge.mapper.FFM._join_pmappings.pmapping_group import PmappingGroup
+from accelforge.mapper.FFM._make_pmappings.pmapper_job import (
+    Job,
+    SameCompatibilityJobs,
+)
+from accelforge.mapper.FFM._pareto_df.df_convention import (
+    is_fused_loop_col,
+    is_n_iterations_col,
+)
+from accelforge.util._mathfuncs import _count_factorizations
+
+
+def shift_reservations_by_null_loop_indices(
+    mappings: pd.DataFrame, null_loop_indices: set[int]
+):
+    prev = copy.deepcopy(mappings)  # TODO: Is this needed?
+    target2newabovename = {}
+    dropcols = []
+    for c in mappings.columns:
+        if not is_reservation_col(c):
+            continue
+        name, above = col2nameloop(c)
+        new_above = above - sum(above > i for i in null_loop_indices)
+        target = nameloop2col(name, new_above)
+        if target in target2newabovename:
+            if above > target2newabovename[target][1]:
+                dropcols.append(nameloop2col(*target2newabovename[target]))
+                target2newabovename[target] = (name, above)
+            else:
+                dropcols.append(c)
+        else:
+            target2newabovename[target] = (name, above)
+
+    mappings.drop(columns=dropcols, inplace=True)
+    renames = {}
+    for target, (name, above) in target2newabovename.items():
+        renames[nameloop2col(name, above)] = target
+    mappings.rename(columns=renames, inplace=True)
+    if len(mappings.columns) != len(mappings.columns.unique()):
+        shift_reservations_by_null_loop_indices(prev, null_loop_indices)
+        raise ValueError(f"Duplicate columns: {mappings.columns}")
+    assert len(mappings.columns) == len(mappings.columns.unique())
+    return mappings
+
+
+def get_equivalent_pmappings(
+    pmapping_group: PmappingGroup, reservation_levels: set[int]
+) -> list[PmappingGroup]:
+    equivalent_permutations = pmapping_group.compatibility.make_equivalent_permutations(
+        reservation_levels
+    )
+    result = [PmappingGroup(c, None) for c in equivalent_permutations]
+    return result
+
+
+def mapping2fused_loop_cols(mapping: Mapping, einsum_name: EinsumName):
+    cols = []
+    for loop in [l for l in mapping.nodes if isinstance(l, Loop) and l._fused]:
+        if loop.tile_shape is not None:
+            cols.append(loop.tile_shape)
+        elif loop.tile_pattern is not None:
+            cols.append(loop.tile_pattern.tile_shape)
+            cols.append(loop.tile_pattern.initial_tile_shape)
+        else:
+            raise ValueError(f"Can't find tile shape or tile pattern for loop {loop}")
+        return [f"{einsum_name}<SEP>{c}" if isinstance(c, str) else c for c in cols]
+
+
+def get_fused_loop_indices(
+    df: pd.DataFrame,
+    compatibility: Compatibility,
+    einsum_name: EinsumName,
+    return_as_int: bool = False,
+) -> pd.Series | int:
+    result = []
+
+    loops = compatibility.loops
+    for i, loop in enumerate(loops):
+        col = loop.tile_pattern.calculated_n_iterations
+        assert col is not None, f"Loop {loop} has no calculated n_iterations"
+        if isinstance(col, str):
+            col = df[f"{einsum_name}<SEP>{col}"]
+        elif isinstance(col, sympy.Symbol):
+            col = df[f"{einsum_name}<SEP>{col.name}"]
+        result.append(col != 1)
+
+    if return_as_int:
+        n = 0
+        for b in result:
+            n = n * 2 + b
+        return n
+    else:
+        r2 = []
+        for b in result:
+            if len(b.unique()) > 1:
+                raise ValueError(f"This won't work if there's more than one")
+            r2.append(b.iloc[0])
+        return r2
+
+
+def _count_loops(job: Job) -> tuple[list[int], list[int], dict[str, int]]:
+    nodes = job.mapping.nodes
+    temporal_n_loops = []
+    spatial_n_loops = []
+    rv_spatial_count = defaultdict(int)
+    rv_temporal_count = defaultdict(int)
+    cur_n_loops = 0
+    spatial_dim = None
+
+    def pop_loop():
+        nonlocal cur_n_loops
+        if cur_n_loops >= 1:
+            if spatial_dim is not None:
+                spatial_n_loops.append(cur_n_loops)
+            else:
+                temporal_n_loops.append(cur_n_loops)
+        cur_n_loops = 0
+
+    for node in nodes:
+        cur_spatial_dim = None
+        if isinstance(node, Spatial):
+            cur_spatial_dim = node.name
+            rv_spatial_count[node.rank_variable] += 1
+        if cur_spatial_dim != spatial_dim:
+            pop_loop()
+            spatial_dim = cur_spatial_dim
+        if isinstance(node, Loop):
+            cur_n_loops += 1
+            if isinstance(node, Temporal):
+                rv_temporal_count[node.rank_variable] += 1
+        else:
+            pop_loop()
+    pop_loop()
+    return temporal_n_loops, spatial_n_loops, rv_spatial_count, rv_temporal_count
+
+
+def multiply_n_pmappings_by_permutations(n_pmappings: int, job: Job) -> int:
+    option = job.spec.mapper.ffm._count_option_for_mapsapce_size_evaluation
+    # if option == "normal":
+    #     return n_pmappings
+
+    temporal_n_loops, spatial_n_loops, rv_spatial_count, rv_temporal_count = (
+        _count_loops(job)
+    )
+
+    rv = {k: v for k, v in job.rank_variable_bounds.items()}
+
+    if "non_helpful_tile_shapes" in option:
+        rv_temporal_count = {r: len(temporal_n_loops) for r in rv.keys()}
+
+    if "non_helpful_loops_for_loop_orders" in option:
+        for i in range(len(temporal_n_loops)):
+            temporal_n_loops[i] = len(rv)
+
+    # Count number of tile shapes
+    rv2loops = {r: rv_spatial_count[r] + rv_temporal_count[r] for r in rv}
+    n_factorizations = math.prod(
+        _count_factorizations(b, rv2loops[r], imperfect=IMPERFECT)
+        for r, b in rv.items()
+    )
+    n_temporal_loop_orders = math.prod(math.factorial(n) for n in temporal_n_loops)
+
+    n = n_factorizations
+
+    # assert n >= n_pmappings, f"n_pmappings: {n_pmappings} > n: {n}"
+
+    if "redundant_loop_orders" in option:
+        # job.mapping._n_loop_orders is the number of permutations that we actually
+        # evaluate. Don't want to double count them.
+        n *= n_temporal_loop_orders / job.mapping._n_loop_orders
+
+    # assert n >= n_pmappings, f"n_pmappings: {n_pmappings} > n: {n}"
+
+    return n
+
+
+def assert_all_jobs_have_same_symbols(
+    jobs_with_similar_compatibilities: SameCompatibilityJobs,
+):
+    iteration2symbols = []
+    for j in jobs_with_similar_compatibilities:
+        for t in j.compatibility.tensors:
+            for i, l in enumerate(t.loops):
+                if len(iteration2symbols) <= i:
+                    iteration2symbols.append(set())
+                iteration2symbols[i].add(l.tile_pattern.calculated_n_iterations)
+    assert all(
+        len(s) == 1 for s in iteration2symbols
+    ), "All jobs must have the same symbols for compatibility n_iterations"
+
+
+def make_pmappings_from_templates(
+    jobs_with_similar_compatibilities: SameCompatibilityJobs,
+) -> tuple[EinsumName, list[PmappingGroup], dict[UUID, Mapping], SameCompatibilityJobs]:
+    jwsc = jobs_with_similar_compatibilities
+
+    results = []
+
+    for job in jobs_with_similar_compatibilities:
+        try:
+            result, tensor2mapping = make_tile_shapes(job)
+        except Exception as e:
+            e.add_note(f"Einsum {jwsc.einsum_name} compatibility {job.compatibility}")
+            raise
+        job.compatibility = job.compatibility.populate_loops()
+
+        # Ctrl-F for CONTIGUOUS_ITERATION_SPACE_DISCUSSION TODO: Turn tensor2pmapping
+        # into per-tensor compatibility
+
+        # This changes the pmapping count to include superfluous permutations
+        # TODO: Add a multiplier for the permutations that we include in the fusion
+        # piece, which are NOT known to be superfluous
+
+        # prev = job.spec.mapper.ffm._count_option_for_mapsapce_size_evaluation
+        # job.spec.mapper.ffm._count_option_for_mapsapce_size_evaluation = "redundant_loop_orders_and_irrelevant_loops"
+        # a = multiply_n_pmappings_by_permutations(job.n_total_pmappings, job)
+        # job.spec.mapper.ffm._count_option_for_mapsapce_size_evaluation = "redundant_loop_orders"
+        # b = multiply_n_pmappings_by_permutations(job.n_total_pmappings, job)
+
+        # if a < b:
+        #     job.spec.mapper.ffm._count_option_for_mapsapce_size_evaluation = "redundant_loop_orders_and_irrelevant_loops"
+        #     a = multiply_n_pmappings_by_permutations(job.n_total_pmappings, job)
+        #     job.spec.mapper.ffm._count_option_for_mapsapce_size_evaluation = "redundant_loop_orders"
+        #     b = multiply_n_pmappings_by_permutations(job.n_total_pmappings, job)
+        #     assert False
+
+        # job.spec.mapper.ffm._count_option_for_mapsapce_size_evaluation = prev
+        job.n_total_pmappings = multiply_n_pmappings_by_permutations(
+            job.n_total_pmappings, job
+        )
+
+        result[MAPPING_COLUMN] = job.job_id
+        cols_to_drop = []
+        for col in result.columns:
+            if is_reservation_col(col):
+                resource = col2nameloop(col)[0]
+                if resource in job.memories_track_pmappings_only:
+                    cols_to_drop.append(col)
+                assert resource not in job.ignored_resources, "Should have been ignored"
+        result.drop(columns=cols_to_drop, inplace=True)
+        results.append(result)
+
+    fusable_tensors = jwsc.fusable_tensors
+    einsum_name = jwsc.einsum_name
+    metrics = jwsc.metrics
+    limit_capacity_drop_valid_reservations = not (Metrics.RESOURCE_USAGE & metrics)
+    compatibility = jwsc.compatibility
+
+    # Creating a PmappingDataframe fills in reservation columns since different pmappings
+    # have different ones.
+    next_shared_loop_index = compatibility.n_loops - 1
+    df = PmappingDataframe.concat(
+        [
+            PmappingDataframe(
+                r,
+                skip_pareto=True,
+                next_shared_loop_index=next_shared_loop_index,
+                n_total_pmappings=1,  # Unused for now, just making an initial Pareto
+                n_valid_pmappings=1,  # Unused for now, just making an initial Pareto
+                ignored_resources=job.ignored_resources,
+                # False because we may have lifetimes that stretch through this Einsum
+                # due to data dependencies, not loops
+                limit_capacity_drop_valid_reservations=False,
+            )
+            for r in results
+        ],
+        skip_pareto=True,
+    ).data
+    if df.empty:
+        return einsum_name, [], {}, jobs_with_similar_compatibilities
+
+    tensor_cols = [tensor2col(tensor) for tensor in fusable_tensors]
+    df.columns = [
+        c if col_used_in_pareto(c) or c in tensor_cols else f"{einsum_name}<SEP>{c}"
+        for c in df.columns
+    ]
+
+    fused_loop_cols = [
+        f"{einsum_name}<SEP>{c}"
+        for c in compatibility.symbols()
+        if not is_n_iterations_col(c)
+    ]
+
+    job0 = next(iter(jobs_with_similar_compatibilities))
+
+    # Pareto prune
+    df = makepareto(df, split_by_cols=fused_loop_cols).copy()
+
+    jobs_passed_pareto = sorted(df[f"{einsum_name}<SEP>{MAPPING_COLUMN}"].unique())
+    pmapping_objects = {
+        job.job_id: job.mapping
+        for job in jobs_with_similar_compatibilities
+        if job.job_id in jobs_passed_pareto
+    }
+
+    assert_all_jobs_have_same_symbols(jobs_with_similar_compatibilities)
+    # Otherwise, following logic fails
+
+    df["fused_loop_indices"] = get_fused_loop_indices(
+        df, job0.compatibility, einsum_name, return_as_int=True
+    )
+    groups = list(df.groupby(["fused_loop_indices"]))
+    total_pmappings_per_group = sum(
+        j.n_total_pmappings for j in jobs_with_similar_compatibilities
+    ) / len(groups)
+    valid_pmappings_per_group = sum(
+        j.n_valid_pmappings for j in jobs_with_similar_compatibilities
+    ) / len(groups)
+
+    pmapping_groups = []
+    for _, mappings in groups:
+        compatibility = jwsc.compatibility
+        fused_loop_indices = []
+
+        for i, f in enumerate(
+            get_fused_loop_indices(
+                mappings, compatibility, einsum_name, return_as_int=False
+            )
+        ):
+            if f:
+                fused_loop_indices.append(i)
+
+        null_loop_indices = tuple(
+            i for i in range(compatibility.n_loops) if i not in fused_loop_indices
+        )
+
+        dropcols = ["fused_loop_indices"]
+        mappings = mappings.drop(columns=dropcols)
+
+        compatibility = compatibility.drop_loop_indices(null_loop_indices)
+
+        symbol_renames, compatibility = compatibility.make_fused_loop_symbols(
+            einsum_name
+        )
+        for k, v in symbol_renames.items():
+            mappings[v] = mappings[f"{einsum_name}<SEP>{k}"]
+        shift_reservations_by_null_loop_indices(mappings, null_loop_indices)
+
+        symbols = compatibility.symbols()
+        dropcols = [
+            c for c in mappings.columns if is_fused_loop_col(c) and c not in symbols
+        ]
+        mappings = mappings.drop(columns=dropcols)
+
+        energy_cols = [c for c in mappings.columns if "Total<SEP>energy" in c]
+        if (mappings[energy_cols] < 0).any(axis=None):
+            mapping_with_negative_energy = mappings[
+                (mappings[energy_cols] < 0).any(axis=1)
+            ]
+            msg = ""
+            for _, row in mapping_with_negative_energy.iterrows():
+                for k, v in row.items():
+                    msg += f"{k}: {v}\n"
+                msg += "\n"
+            raise RuntimeError(f"negative energy:\n{msg}")
+
+        # Skip pareto because we already did it above
+        next_shared_loop_index_this_group = compatibility.n_loops - 1
+        partial_mappings = PmappingDataframe(
+            mappings,
+            next_shared_loop_index=next_shared_loop_index_this_group,
+            n_total_pmappings=total_pmappings_per_group,
+            n_valid_pmappings=valid_pmappings_per_group,
+            skip_pareto=next_shared_loop_index_this_group == next_shared_loop_index,
+            ignored_resources=job.ignored_resources,
+            # False because we may have lifetimes that stretch through this Einsum
+            # due to data dependencies, not loops
+            limit_capacity_drop_valid_reservations=False,
+        )
+        pmapping_groups.append(PmappingGroup(compatibility, partial_mappings))
+
+    return (
+        einsum_name,
+        pmapping_groups,
+        pmapping_objects,
+        jobs_with_similar_compatibilities,
+    )