accelforge 0.0.1__py3-none-any.whl

accelforge/mapper/FFM/_join_pmappings/compress_pmappings.py

@@ -0,0 +1,140 @@
+from typing import Any, NamedTuple
+
+from tqdm import tqdm
+from accelforge._accelerated_imports import pd
+from accelforge.frontend.workload import EinsumName
+from accelforge.mapper.FFM._join_pmappings.pmapping_group import PmappingGroup
+from accelforge.mapper.FFM._pareto_df.df_convention import (
+    COMPRESSED_INDEX,
+    col_used_in_pareto,
+    is_fused_loop_col,
+    is_tensor_col,
+)
+from accelforge.mapper.FFM._join_pmappings.pmapping_dataframe import PmappingDataframe
+from accelforge.util.parallel import parallel, delayed
+
+
+class DecompressData(NamedTuple):
+    data: dict[EinsumName, dict[int, pd.DataFrame]]
+
+
+def _compress(
+    einsum_name: EinsumName, pmappings: PmappingGroup, start_index: int
+) -> tuple["PmappingDataframe", pd.DataFrame]:
+    data = pmappings.mappings.data
+    data.reset_index(drop=True, inplace=True)
+    data.index += start_index
+    keep_cols = [
+        c
+        for c in data.columns
+        if col_used_in_pareto(c) or is_fused_loop_col(c) or is_tensor_col(c)
+    ]
+    compress_cols = [c for c in data.columns if c not in keep_cols]
+    compressed_data = data[keep_cols].copy()
+    decompress_data = data[compress_cols].copy()
+    compressed_data[f"{einsum_name}<SEP>{COMPRESSED_INDEX}"] = data.index
+    return (
+        pmappings.mappings.update(data=compressed_data, skip_pareto=True),
+        decompress_data,
+    )
+
+
+def _compress_pmapping_list(
+    einsum_name: EinsumName, pmappings: list[PmappingGroup]
+) -> tuple[list[PmappingDataframe], dict[int, pd.DataFrame]]:
+    decompress_data = {}
+    compressed = []
+    start_index = 0
+    jobs = []
+
+    def job(start_index: int, pmappings: PmappingGroup):
+        compress, decompress = _compress(
+            einsum_name,
+            pmappings,
+            start_index,
+        )
+        compress = PmappingGroup(pmappings.compatibility, compress)
+        compressed.append(compress)
+        return compress, decompress, start_index
+
+    for pmapping in pmappings:
+        jobs.append(delayed(job)(start_index, pmapping))
+        start_index += len(pmapping.mappings.data)
+
+    for compress, decompress, start_index in parallel(jobs, n_jobs=1):
+        compressed.append(compress)
+        decompress_data[start_index] = decompress
+
+    return compressed, decompress_data  # pd.concat(decompress_data)
+
+
+def compress_einsum2pmappings(
+    einsum2pmappings: dict[EinsumName, list[PmappingGroup]],
+) -> tuple[dict[EinsumName, list[PmappingGroup]], DecompressData]:
+    decompress_data = {}
+    compressed_einsum2pmappings = {}
+
+    jobs = []
+
+    def job(einsum_name: EinsumName, pmappings: list[PmappingGroup]):
+        compressed, decompress = _compress_pmapping_list(einsum_name, pmappings)
+        return einsum_name, compressed, decompress
+
+    for einsum_name, pmappings in einsum2pmappings.items():
+        jobs.append(delayed(job)(einsum_name, pmappings))
+
+    name_order = [einsum_name for einsum_name in einsum2pmappings.keys()]
+    for einsum_name, compressed, decompress in parallel(
+        jobs, pbar="Compressing pmappings", return_as="generator_unordered"
+    ):
+        compressed_einsum2pmappings[einsum_name] = compressed
+        decompress_data[einsum_name] = decompress
+
+    compressed_einsum2pmappings = {
+        einsum_name: compressed_einsum2pmappings[einsum_name]
+        for einsum_name in name_order
+    }
+    decompress_data = {
+        einsum_name: decompress_data[einsum_name] for einsum_name in name_order
+    }
+
+    # for einsum_name, pmappings in tqdm(einsum2pmappings.items(), desc="Compressing pmappings"):
+    #     compressed, decompress = _compress_pmapping_list(einsum_name, pmappings)
+    #     compressed_einsum2pmappings[einsum_name] = compressed
+    #     decompress_data[einsum_name] = decompress
+    return compressed_einsum2pmappings, DecompressData(decompress_data)
+
+
+def decompress_pmappings(
+    pmappings: PmappingDataframe,
+    decompress_data: DecompressData,
+) -> PmappingDataframe:
+    data = pmappings.data
+    for einsum_name, decompress in decompress_data.data.items():
+        decompress_sub_dfs = []
+        # We reverse it because we have the start index of each sub df, but we need the
+        # end index so we know when to change to the next
+        decompressed_iter = reversed(decompress.items())
+        start_index, chosen = float("inf"), None
+        for i in reversed(sorted(set(data[f"{einsum_name}<SEP>{COMPRESSED_INDEX}"]))):
+            while chosen is None or i < start_index:
+                start_index, chosen = next(decompressed_iter)
+            cur_chosen = chosen[chosen.index == i]
+            assert (
+                len(cur_chosen) == 1
+            ), f"Expected 1 row, got {len(cur_chosen)}. Index {i}, start index {start_index}, chosen {list(chosen.index)}"
+            decompress_sub_dfs.append(cur_chosen)
+        data = pd.merge(
+            data,
+            pd.concat(decompress_sub_dfs),
+            left_on=f"{einsum_name}<SEP>{COMPRESSED_INDEX}",
+            right_index=True,
+            how="left",
+        )
+
+    # Remove compressed_index columns that may have been created during
+    # merge
+    compressed_index_cols = [col for col in data.columns if COMPRESSED_INDEX in col]
+    if compressed_index_cols:
+        data = data.drop(columns=compressed_index_cols)
+    return pmappings.update(data=data, skip_pareto=True)