Speedup for few procs to join, EDP early prune improvements

Author: tanner-andrulis

accelforge/frontend/arch/_flattened_arch.py

@@ -20,6 +20,7 @@ class FlattenedArch:
     `frontend`, this one is intentionally *not* an
     `EvalableModel`.
     """
+
     def __init__(self, nodes: list["Leaf"]):
         self.nodes = nodes
 

accelforge/frontend/mapper/ffm.py

@@ -232,11 +232,3 @@ class FFM(EvalableModel):
     If set, append per-step runtime as JSON lines to this file. Used for ablation study
     measurements.
     """
-
-    _metric_aggregator: Literal["sum", "prod", "any"] = "any"
-    """
-    How to aggregate metrics together to determine whether one pmapping is better than
-    another. "sum" means that the metrics are added together, "prod" means that the
-    metrics are multiplied together, and "any" means that any metric being better than
-    the other is enough to consider it non-dominated.
-    """

accelforge/mapper/FFM/_join_pmappings/join_pmappings.py

@@ -34,7 +34,13 @@
     Compatibility,
 )
 from accelforge.mapper.FFM._pareto_df.df_convention import col2reservation
-from accelforge.util import _fillna_and__numeric_cast, parallel, delayed, oset
+from accelforge.util import (
+    _fillna_and__numeric_cast,
+    delayed,
+    get_n_parallel_jobs,
+    oset,
+    parallel,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -62,13 +68,16 @@ def log_total_time(self):
         logger.info(f"============================\n")
 
 
-def _apply_edp_columns(df: pd.DataFrame, metrics: Metrics) -> pd.DataFrame:
-    if not (metrics & Metrics.ENERGY_DELAY_PRODUCT):
-        return df
-    if not (metrics & Metrics.ENERGY):
-        del df["Total<SEP>energy"]
-    if not (metrics & Metrics.LATENCY):
-        del df["Total<SEP>latency"]
+def _apply_edp_columns(
+    df: pd.DataFrame, metrics: Metrics, return_only_objectives: bool = False
+) -> pd.DataFrame:
+    if metrics & Metrics.ENERGY_DELAY_PRODUCT:
+        if not (metrics & Metrics.ENERGY):
+            del df["Total<SEP>energy"]
+        if not (metrics & Metrics.LATENCY):
+            del df["Total<SEP>latency"]
+    if return_only_objectives:
+        df = df[[c for c in df.columns if col_used_in_pareto(c)]]
     return df
 
 
@@ -87,8 +96,8 @@ def __init__(
 
         compare_to = compare_to.sort_values(by=compare_cols, ascending=False)
 
-        if len(compare_to) > 10:
-            chosen_indices = np.round(np.linspace(0, len(compare_to) - 1, 10))
+        if len(compare_to) > 100:
+            chosen_indices = np.round(np.linspace(0, len(compare_to) - 1, 100))
         else:
             chosen_indices = np.round(np.arange(len(compare_to)))
 
@@ -713,6 +722,28 @@ def grab_einsum_pmappings() -> (
         left = PmappingGroup.group(left, right_tensors)
         # print_time("Grouping")
 
+        # =============================================================================
+        # If we're multiprocessing and the left side has fewer groups than the number of
+        # processes, repeatedly split the largest group in half so that the merge work
+        # below can fan out across all workers.
+        # =============================================================================
+        n_procs = get_n_parallel_jobs()
+        n_groups = sum(len(v) for v in left.values())
+        for _ in range(n_procs - n_groups):
+            best_k, best_i, best_len = None, None, -1
+            for k, vs in left.items():
+                for i, (pg, _perm) in enumerate(vs):
+                    length = len(pg.mappings.data) if pg.mappings is not None else 0
+                    if length > best_len:
+                        best_len = length
+                        best_k, best_i = k, i
+            if best_k is None or best_len < 2:
+                break  # nothing left worth splitting
+            pg, perm = left[best_k][best_i]
+            first, second = pg.split_in_half()
+            left[best_k][best_i] = (first, perm)
+            left[best_k].append((second, perm))
+
         # ======================================================================
         # Remove dead tensors from left and right. This happens after grouping because
         # we only reserve space for shared tensors after they're dead (alive is handled

accelforge/mapper/FFM/_join_pmappings/pmapping_dataframe.py

@@ -772,6 +772,24 @@ def copy(self, copy_df: bool = True) -> "PmappingDataframe":
             check_above_subset_below=False,
         )
 
+    def split_in_half(self) -> tuple["PmappingDataframe", "PmappingDataframe"]:
+        mid = len(self.data) // 2
+        half_total = self.n_total_pmappings / 2
+        half_valid = self.n_valid_pmappings / 2
+        first = self.update(
+            data=self.data.iloc[:mid].copy(),
+            skip_pareto=True,
+            n_total_pmappings=half_total,
+            n_valid_pmappings=half_valid,
+        )
+        second = self.update(
+            data=self.data.iloc[mid:].copy(),
+            skip_pareto=True,
+            n_total_pmappings=half_total,
+            n_valid_pmappings=half_valid,
+        )
+        return first, second
+
     def limit_capacity(
         self,
         next_shared_loop_index: int = None,

accelforge/mapper/FFM/_join_pmappings/pmapping_group.py

@@ -47,6 +47,13 @@ def tensor_names(self) -> set[str]:
     def copy(self) -> "PmappingGroup":
         return PmappingGroup(self.compatibility, self.mappings.copy())
 
+    def split_in_half(self) -> tuple["PmappingGroup", "PmappingGroup"]:
+        first_mappings, second_mappings = self.mappings.split_in_half()
+        return (
+            PmappingGroup(self.compatibility, first_mappings),
+            PmappingGroup(self.compatibility, second_mappings),
+        )
+
     def __len__(self) -> int:
         return len(self.mappings)
 

accelforge/mapper/FFM/_make_pmappings/make_pmappings.py

@@ -70,15 +70,11 @@ def get_n_computes(spec: Spec, einsum_name: EinsumName | None = None) -> int:
     return sum(get_operation_space_size(spec.workload, e) for e in einsums)
 
 
-def get_per_tensor_size(
-    workload: Workload, return_n_elements: bool = False
-) -> dict[TensorName, int]:
+def get_per_tensor_n_elements(workload: Workload) -> dict[TensorName, int]:
     sizes = {}
     for einsum in workload.einsums:
         for tensor in einsum.tensor_names:
             sizes[tensor] = get_tensor_size(workload, tensor)
-            if not return_n_elements:
-                sizes[tensor] *= einsum.tensor_accesses[tensor].bits_per_value
     return sizes
 
 
@@ -265,7 +261,7 @@ def get_memories_to_track(
         )
 
     tensor_sizes = {}
-    for tensor, size in get_per_tensor_size(spec.workload).items():
+    for tensor, size in get_per_tensor_n_elements(spec.workload).items():
         scale = 1
         for einsum in spec.workload.einsums_with_tensor(tensor):
             if einsum.tensor_accesses[tensor].persistent:
@@ -292,9 +288,7 @@ def get_memories_to_track(
                         .bits_per_value
                     )
                     effective_bpv = mem.bits_per_value.get(tensor, workload_bpv)
-                    usage += (
-                        tensor_sizes[tensor] * effective_bpv / (workload_bpv * mem.size)
-                    )
+                    usage += tensor_sizes[tensor] * effective_bpv / mem.size
 
         if usage <= 1:
             ignored_resources.add(memory)

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

@@ -723,6 +723,7 @@ def _factorize_imperfect(n: int) -> np.ndarray:
         factors.append(math.ceil(n / i))
     return np.array(sorted(oset(factors)))
 
+
 @lru_cache(maxsize=10000)
 def get_possible_factor_sizes(
     outer_size: int, imperfect: bool, inner_size: int, coarseness: float = 1

accelforge/mapper/FFM/mappings.py

@@ -7,7 +7,7 @@
 from accelforge.util._frozenset import oset
 from accelforge.mapper.FFM._make_pmappings.make_pmappings import (
     get_n_computes,
-    get_per_tensor_size,
+    get_per_tensor_n_elements,
 )
 from typing import Union
 from accelforge._accelerated_imports import numpy as np
@@ -123,10 +123,15 @@ def per_tensor_size(self, return_n_elements: bool = False) -> dict[TensorName, i
         assert (
             self.evaluated_specs is not None
         ), "Can't get per-tensor size if no Einsums have been mapped."
-        return get_per_tensor_size(
-            next(iter(self.evaluated_specs.values())).workload,
-            return_n_elements=return_n_elements,
-        )
+        workload = next(iter(self.evaluated_specs.values())).workload
+        sizes = get_per_tensor_n_elements(workload)
+        if not return_n_elements:
+            for einsum in workload.einsums:
+                for tensor in einsum.tensor_names:
+                    sizes[tensor] = (
+                        sizes[tensor] * einsum.tensor_accesses[tensor].bits_per_value
+                    )
+        return sizes
 
     def _update(self, **kwargs):
         data = dict(

tests/test_network.py

@@ -82,7 +82,7 @@ def test_hierarchical_1d(self):
             * (KN / MAC_TILE)  # number of used Scratchpad
             * M_TILE
             * KN  # temporal for n1 in mapping
-            * sum(i+1 for i in range(MAC_TILE))  # unicast along X-axis of MacArray
+            * sum(i + 1 for i in range(MAC_TILE))  # unicast along X-axis of MacArray
             * BITS_PER_VALUE,
         )
         # NOTE: assuming XY routing (as defined in mapping)
@@ -92,7 +92,7 @@ def test_hierarchical_1d(self):
             * (KN / MAC_TILE)
             * M_TILE
             * KN  # temporal for n1 in mapping
-            * MAC_TILE   # multicast along X-axis of MacArray
+            * MAC_TILE  # multicast along X-axis of MacArray
             * BITS_PER_VALUE,
         )
         self.assertEqual(
@@ -101,14 +101,16 @@ def test_hierarchical_1d(self):
             * (KN / MAC_TILE)
             * M_TILE
             * KN
-            * sum(i+1 for i in range(MAC_TILE))
+            * sum(i + 1 for i in range(MAC_TILE))
             * BITS_PER_VALUE,
         )
 
         self.assertEqual(
             result.data["Matmul0<SEP>action<SEP>PeArray<SEP>T0<SEP>hops"].iloc[0],
             (M / M_TILE)
-            * sum(i+1 for i in range(KN // MAC_TILE))  # unicast along X-axis of PeArray
+            * sum(
+                i + 1 for i in range(KN // MAC_TILE)
+            )  # unicast along X-axis of PeArray
             * M_TILE
             * MAC_TILE
             * BITS_PER_VALUE,
@@ -117,15 +119,16 @@ def test_hierarchical_1d(self):
         self.assertEqual(
             result.data["Matmul0<SEP>action<SEP>PeArray<SEP>T1<SEP>hops"].iloc[0],
             (M / M_TILE)
-            * KN // MAC_TILE  # multicast along X-axis of PeArray
+            * KN
+            // MAC_TILE  # multicast along X-axis of PeArray
             * M_TILE
             * KN
             * BITS_PER_VALUE,
         )
         self.assertEqual(
             result.data["Matmul0<SEP>action<SEP>PeArray<SEP>W0<SEP>hops"].iloc[0],
             (M / M_TILE)
-            * sum(i+1 for i in range(KN // MAC_TILE))  # unicast along PeArray
+            * sum(i + 1 for i in range(KN // MAC_TILE))  # unicast along PeArray
             * MAC_TILE
             * KN
             * BITS_PER_VALUE,
@@ -160,9 +163,8 @@ def test_hierarchical(self):
             * (KN / MAC_TILE) ** 2
             * M_TILE
             * (
-                sum(i+1 for i in range(MAC_TILE))  # unicasting along X
-                +
-                MAC_TILE * MAC_TILE  # multicast along Y for each column
+                sum(i + 1 for i in range(MAC_TILE))  # unicasting along X
+                + MAC_TILE * MAC_TILE  # multicast along Y for each column
             )
             * BITS_PER_VALUE,
         )
@@ -173,9 +175,10 @@ def test_hierarchical(self):
             * (KN / MAC_TILE) ** 2
             * M_TILE
             * (
-                MAC_TILE * MAC_TILE  # multicast along X (the tile is shape N1, which is MAC_TILE here)
-                +
-                MAC_TILE * sum(i+1 for i in range(MAC_TILE))  # unicasting along Y for each row
+                MAC_TILE
+                * MAC_TILE  # multicast along X (the tile is shape N1, which is MAC_TILE here)
+                + MAC_TILE
+                * sum(i + 1 for i in range(MAC_TILE))  # unicasting along Y for each row
             )
             * BITS_PER_VALUE,
         )
@@ -185,35 +188,27 @@ def test_hierarchical(self):
             * (KN / MAC_TILE) ** 2
             * M_TILE
             * (
-                MAC_TILE * sum(i+1 for i in range(MAC_TILE))  # unicast along X (the tile is shape N1, which is MAC_TILE here)
-                +
-                MAC_TILE * sum(i+1 for i in range(MAC_TILE))  # unicasting along Y for each row
+                MAC_TILE
+                * sum(
+                    i + 1 for i in range(MAC_TILE)
+                )  # unicast along X (the tile is shape N1, which is MAC_TILE here)
+                + MAC_TILE
+                * sum(i + 1 for i in range(MAC_TILE))  # unicasting along Y for each row
             )
             * BITS_PER_VALUE,
         )
 
         self.assertEqual(
             result.data["Matmul0<SEP>action<SEP>PeArray<SEP>T0<SEP>hops"].iloc[0],
-            (M / M_TILE)
-            * (
-                sum(i+1 for i in range(PE_TILE))
-                +
-                PE_TILE * PE_TILE
-            )
+            (M / M_TILE) * (sum(i + 1 for i in range(PE_TILE)) + PE_TILE * PE_TILE)
             # tile shape
-            * M_TILE
-            * MAC_TILE
-            * BITS_PER_VALUE,
+            * M_TILE * MAC_TILE * BITS_PER_VALUE,
         )
         # NOTE: assuming XY routing (as defined in mapping)
         self.assertEqual(
             result.data["Matmul0<SEP>action<SEP>PeArray<SEP>T1<SEP>hops"].iloc[0],
             (M / M_TILE)
-            * (
-                PE_TILE * PE_TILE
-                +
-                PE_TILE * sum(i+1 for i in range(PE_TILE))
-            )
+            * (PE_TILE * PE_TILE + PE_TILE * sum(i + 1 for i in range(PE_TILE)))
             * M_TILE
             * MAC_TILE
             * BITS_PER_VALUE,
@@ -222,9 +217,8 @@ def test_hierarchical(self):
             result.data["Matmul0<SEP>action<SEP>PeArray<SEP>W0<SEP>hops"].iloc[0],
             (M / M_TILE)
             * (
-                PE_TILE * sum(i+1 for i in range(PE_TILE))
-                +
-                PE_TILE * sum(i+1 for i in range(PE_TILE))
+                PE_TILE * sum(i + 1 for i in range(PE_TILE))
+                + PE_TILE * sum(i + 1 for i in range(PE_TILE))
             )
             * MAC_TILE**2
             * BITS_PER_VALUE,
@@ -252,4 +246,4 @@ def test_hierarchical(self):
                 "M_TILE": M_TILE,
             },
         )
-        result = spec.map_workload_to_arch()
+        result = spec.map_workload_to_arch()