Merge branch 'main' into hops

Author: gilbertmike

accelforge/frontend/arch/_flattened_arch.py

@@ -29,6 +29,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

@@ -66,6 +66,16 @@ class FFM(EvalableModel):
     max_fused_loops_per_rank_variable: int = 1
     """ The maximum number of fused loops in a pmapping for a given rank variable. """
 
+    tiling_coarseness: float = 1
+    """
+    When enumerating tile shapes, a tile shape is omitted if it is less than
+    tiling_coarseness times the next-inner tile shape. Set to 1 to disable. For example,
+    tiling_coarseness=2 means each enumerated tile shape must be at least 2x the
+    next-inner tile shape. Recommended to use 1<=tiling_coarseness<=2
+    explore_imperfect_temporal_loops=True and/or explore_imperfect_spatial_loops=True,
+    to get good mappings without exploring too many tile shapes.
+    """
+
     max_fused_loops: float | int = float("inf")
     """ The maximum total number of fused loops in a pmapping. """
 
@@ -139,6 +149,9 @@ class FFM(EvalableModel):
 
     Only "simple" rank variables (those appearing alone and not inside an expression in
     any tensor access) may have imperfect loop bounds.
+    
+    See tiling_coarseness for how to reduce the number of tile shapes explored when this
+    is set.
     """
 
     explore_imperfect_temporal_loops: bool = False
@@ -154,6 +167,9 @@ class FFM(EvalableModel):
 
     Only "simple" rank variables (those appearing alone and not inside an expression in
     any tensor access) may have imperfect loop bounds.
+    
+    See tiling_coarseness for how to reduce the number of tile shapes explored when this
+    is set.
     """
 
     prioritize_reuse_of_unfused_tensors: bool = False
@@ -216,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

@@ -724,21 +724,47 @@ def _factorize_imperfect(n: int) -> np.ndarray:
     return np.array(sorted(oset(factors)))
 
 
-def get_possible_factor_sizes(n: int, imperfect: bool, inner_size: int) -> list[int]:
+@lru_cache(maxsize=10000)
+def get_possible_factor_sizes(
+    outer_size: int, imperfect: bool, inner_size: int, coarseness: float = 1
+) -> list[int]:
     # Int to protect from numpy overflows
-    n = int(n)
-    inner_size = int(inner_size)
-    if not imperfect:
-        factors = _factorize(math.ceil(n / inner_size))
-        return factors * inner_size
+    outer_size, inner_size = int(outer_size), int(inner_size)
+
+    factors = set()
+    n_tiles = set()
+
+    def _try_admit(n: int) -> bool:
+        if not imperfect and outer_size % n != 0:
+            return False
+        if n > outer_size:
+            return False
+
+        # This check: Avoid redundant tile shapes that lead to the same number of tiles,
+        # which causes the same accesses (from same average shape) but worse latency &
+        # memory usage (from worse long pole shape).
+        cur_n_tiles = math.ceil(outer_size / n)
+        if cur_n_tiles not in n_tiles:
+            n_tiles.add(cur_n_tiles)
+            # new_n calculation: Grab the smallest-possible tile shape that would get us
+            # this number of tiles, which is the best (see above comment).
+            new_n = math.ceil(outer_size / cur_n_tiles)
+            factors.add(new_n)
+        return True
+
+    n = inner_size
+    while n <= outer_size:
+        # If we successfully admit an n, don't try another n until we hit our coarseness
+        # target.
+        if _try_admit(n):
+            n = max(n + 1, math.ceil(n * coarseness))
+        else:
+            n += 1
+
+    # One more in case coarseness jumped the max value
+    _try_admit(outer_size)
 
-    factors = _factorize_imperfect(n)
-    if n % inner_size == 0:
-        perfects = _factorize(math.ceil(n / inner_size)) * inner_size
-        assert all(
-            f in factors for f in perfects
-        ), f"perfects: {perfects} not in factors: {factors}"
-    return factors[factors >= inner_size]
+    return np.array(sorted(factors))
 
 
 def append_vector(matrix: np.ndarray, vector: np.ndarray):
@@ -1470,6 +1496,7 @@ def eval_objective(
                 raise RuntimeError("BUG: both inner and outer tiles are unknown")
 
             symbol_imperfect = _symbol_is_imperfect(symbol)
+            coarseness = job.spec_one_einsum.mapper.tiling_coarseness
             # Use inner size and outer size to generate choices
             if inner_tiles_type in oset(
                 ["set", "unknown"]
@@ -1480,7 +1507,7 @@ def eval_objective(
                 ]
             ):
                 factors = get_possible_factor_sizes(
-                    outer_size, symbol_imperfect, inner_size
+                    outer_size, symbol_imperfect, inner_size, coarseness
                 )
                 choices.append(append_vector(choices_enumerated, factors))
             elif inner_tiles_type == "enumerated":
@@ -1491,7 +1518,7 @@ def eval_objective(
                         np.where(choices_enumerated[:, i] == inner_choice)
                     ]
                     factors = get_possible_factor_sizes(
-                        outer_size, symbol_imperfect, inner_choice
+                        outer_size, symbol_imperfect, inner_choice, coarseness
                     )
                     choices.append(append_vector(partition, factors))
             else:
@@ -1503,7 +1530,7 @@ def eval_objective(
                         np.where(choices_enumerated[:, i] == outer_choice)
                     ]
                     factors = get_possible_factor_sizes(
-                        outer_choice, symbol_imperfect, inner_size
+                        outer_choice, symbol_imperfect, inner_size, coarseness
                     )
                     choices.append(append_vector(partition, factors))
         elif what_tiles_symbol.is_initial_tile_shape(symbol):

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(

docs/source/guide/timeloop_compare.rst

@@ -44,6 +44,15 @@ The following are features **suppported by AccelForge, not Timeloop**:
   <https://github.com/Accelergy-Project/accelergy>`_. Hwcomponents is a successor to
   Accelergy, and makes it significantly easier to write custom component models.
 - **Easy Installation**: AccelForge and dependencies are pip installable.
+- **Mapper for imperfectly-factorized loop levels**: The mapper can explore imperfect
+  mappings for all loop levels. We note two differences between our implementations and
+  Timeloop's. **(1)** We permit the outermost loop for each rank variable to have a
+  residual, while Timeloop permits a residual on every loop. **(2)** We only run
+  imperfect mapping for simple rank variables which do not participate in indexing
+  expressions, while Timeloop allows every rank variable to be imperfect. Our
+  constraints allow us to find good mappings while keeping the search space tractable,
+  while Timeloop may only search the imperfect space with additional user-written
+  constraints.
 
 The following are features **supported by Timeloop, work-in-progress in AccelForge**:
 
@@ -53,10 +62,6 @@ The following are features **supported by Timeloop, work-in-progress in AccelFor
 - **Layout Support**: Support for the costs of how data is laid out in memory.
 - **Skew in Model**: User-defined spatial skews that can split data between multiple
   multiple spatial instances.
-- **Mapper for $\sim 1-2$ imperfectly-factorized loop levels**: The mapper can explore
-  imperfect mappings for a handful of loop levels, but becomes intractible for more loop
-  levels. This, as well as full imperfect mapping, is present in AccelForge, but the API
-  is not yet implemented.
 - **ISL-Based Model**: Timeloop supports high-fidelity modeling with ISL to generate the
   accesses for every spatial instance, while AccelForge uses a simpler and faster
   analytical model.
@@ -66,8 +71,6 @@ Timeloop**:
 
 - **Skew in Mapping**: Automatically deriving the best skew for a given workload and
   architecture.
-- **Mapper for all imperfectly-factorized loop levels**: The mapper can explore
-  imperfect mappings for all loop levels.
 - **Einsum-Level Spatial Parallelism**: The mapper can parallelize Einsums across
   spatial instances, rather than executing them sequentially.
 

tests/test_network.py

@@ -196,9 +196,7 @@ def test_hierarchical(self):
                 PE_TILE * (PE_TILE - 1)
             )
             # 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(