Run black + EDP mapper metric

Author: tanner-andrulis

accelforge/frontend/mapper/metrics.py

@@ -37,16 +37,33 @@ class Metrics(Flag):
     Memory usage broken down by tensor and Einsum.
     """
 
+    ENERGY_DELAY_PRODUCT = auto()
+    """The product of energy and latency. """
+
     @classmethod
     def all_metrics(cls):
         return reduce(or_, iter(cls), cls.LATENCY)
 
     def includes_leak_energy(self) -> bool:
         """Returns True if the metrics include leak energy, either alone or as part of
         total energy. False otherwise."""
-        return self & (Metrics.ENERGY | Metrics.LEAK_ENERGY)
+        return self & (
+            Metrics.ENERGY | Metrics.LEAK_ENERGY | Metrics.ENERGY_DELAY_PRODUCT
+        )
 
     def includes_dynamic_energy(self) -> bool:
         """Returns True if the metrics include dynamic energy, either alone or as part
         of total energy. False otherwise."""
-        return self & (Metrics.ENERGY | Metrics.DYNAMIC_ENERGY)
+        return self & (
+            Metrics.ENERGY | Metrics.DYNAMIC_ENERGY | Metrics.ENERGY_DELAY_PRODUCT
+        )
+
+    def includes_latency(self) -> bool:
+        """Returns True if the metrics include latency, either alone or as part of
+        energy-delay product. False otherwise."""
+        return self & (Metrics.LATENCY | Metrics.ENERGY_DELAY_PRODUCT)
+
+    def includes_energy(self) -> bool:
+        """Returns True if the metrics include total energy, either alone or as part of
+        energy-delay product. False otherwise."""
+        return self & (Metrics.ENERGY | Metrics.ENERGY_DELAY_PRODUCT)

accelforge/mapper/FFM/_join_pmappings/join_pmappings.py

@@ -62,81 +62,65 @@ 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"]
+    return df
+
+
 class OptimalityThresholder:
     def __init__(
         self,
         prev_solutions: Mappings,
         _pmapping_row_filter_function: Callable[[pd.DataFrame], np.ndarray],
-        aggregator: str,
         print_progress: bool,
+        metrics: Metrics,
     ):
-        compare_to = prev_solutions.data
+        self.metrics = metrics
+        compare_to = _apply_edp_columns(prev_solutions.data.copy(), metrics)
         compare_cols = [c for c in compare_to.columns if col_used_in_pareto(c)]
         self._pmapping_row_filter_function = _pmapping_row_filter_function
-        self.aggregator = aggregator
-
-        if self.aggregator in ("prod", "sum"):
-            objective_cols = [c for c in compare_cols if is_objective_col(c)]
-            self._agg_cols = objective_cols
-            if objective_cols:
-                values = np.column_stack([compare_to[c].values for c in objective_cols])
-                if self.aggregator == "prod":
-                    agg = np.prod(values, axis=1)
-                else:
-                    agg = np.sum(values, axis=1)
-                self._agg_threshold = agg.min()
-            else:
-                self._agg_threshold = float("inf")
-            if print_progress:
-                label = "product" if self.aggregator == "prod" else "sum"
-                print(
-                    f"Filtering out pmappings with {label} > "
-                    f"{self._agg_threshold:.2e}"
-                )
-        else:  # "any"
-            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))
-            else:
-                chosen_indices = np.round(np.arange(len(compare_to)))
+        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))
+        else:
+            chosen_indices = np.round(np.arange(len(compare_to)))
+
+        self.compare_to: list[dict[str, float]] = []
+        if print_progress:
+            print(f"Filtering out pmappings worse than the following:")
 
-            self.compare_to: list[dict[str, float]] = []
+        for i in chosen_indices.astype(int):
+            self.compare_to.append({c: compare_to[c].iloc[i] for c in compare_cols})
             if print_progress:
-                print(f"Filtering out pmappings worse than the following:")
-
-            for i in chosen_indices.astype(int):
-                self.compare_to.append({c: compare_to[c].iloc[i] for c in compare_cols})
-                if print_progress:
-                    print(
-                        "\t"
-                        + "    ".join(
-                            f"{k}={float(v):.2e}"
-                            for k, v in self.compare_to[-1].items()
-                        )
+                print(
+                    "\t"
+                    + "    ".join(
+                        f"{k}={float(v):.2e}" for k, v in self.compare_to[-1].items()
                     )
+                )
 
     def __call__(self, mapping: pd.DataFrame) -> bool:
         nondominated_by_all = np.ones(len(mapping), dtype=bool)
 
-        if self.aggregator in ("prod", "sum"):
-            cols_present = [c for c in self._agg_cols if c in mapping.columns]
-            if cols_present:
-                values = np.column_stack([mapping[c].values for c in cols_present])
-                if self.aggregator == "prod":
-                    agg = np.prod(values, axis=1)
+        edp_mapping = _apply_edp_columns(
+            mapping.copy(), self.metrics, return_only_objectives=True
+        )
+
+        for c in self.compare_to:
+            nondominated = np.zeros(len(edp_mapping), dtype=bool)
+            for k, v in c.items():
+                if k not in edp_mapping.columns:
+                    nondominated |= True
                 else:
-                    agg = np.sum(values, axis=1)
-                nondominated_by_all = agg <= self._agg_threshold
-        else:  # "any"
-            for c in self.compare_to:
-                nondominated = np.zeros(len(mapping), dtype=bool)
-                for k, v in c.items():
-                    if k not in mapping.columns:
-                        nondominated |= True
-                    else:
-                        nondominated |= mapping[k] <= v
-                nondominated_by_all &= nondominated
+                    nondominated |= edp_mapping[k] <= v
+            nondominated_by_all &= nondominated
 
         if self._pmapping_row_filter_function is not None:
             nondominated_by_all &= self._pmapping_row_filter_function(mapping)
@@ -235,8 +219,8 @@ def join_strategy_2(
                 filter_func = OptimalityThresholder(
                     joined,
                     _pmapping_row_filter_function,
-                    spec.mapper._metric_aggregator,
                     print_progress,
+                    metrics,
                 )
         except Exception as e:
             if i == len(thresholds) - 1:
@@ -356,6 +340,8 @@ def clean_compress_and_join_pmappings(
 
     joined = decompress_pmappings(joined, decompress_data)
 
+    _apply_edp_columns(joined.data, metrics)
+
     for einsum_name in einsum2pmappings:
         col = f"{einsum_name}<SEP>{MAPPING_COLUMN}"
         joined.data[col] = joined.data[col].apply(

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

@@ -2055,7 +2055,7 @@ def _clean_energy_columns(df: dict, metrics: Metrics):
     # The model outputs separated dynamic energy and leak energy because it's easier for
     # tile shape exploration. Combine them if needed and generate the total energy
     # column.
-    if metrics & Metrics.ENERGY:
+    if metrics.includes_energy():
         leak = df.pop("Total<SEP>leak_energy")
         dynamic = df.pop("Total<SEP>dynamic_energy")
         df["Total<SEP>energy"] = leak + dynamic

accelforge/mapper/_simanneal2/simanneal.py

@@ -714,6 +714,8 @@ def _make_mapspace_globals(
         objective_function = lambda x: x["Total<SEP>latency"]
     elif objective == (Metrics.ENERGY | Metrics.LATENCY):
         objective_function = lambda x: x["Total<SEP>energy"] * x["Total<SEP>latency"]
+    elif objective == Metrics.ENERGY_DELAY_PRODUCT:
+        objective_function = lambda x: x["Total<SEP>energy"] * x["Total<SEP>latency"]
     else:
         raise ValueError(f"Unknown objective {objective}")
     mg = MapspaceGlobals(

accelforge/model/_looptree/reuse/symbolic/symbolic/_network.py

@@ -52,7 +52,9 @@ def accumulate_child_result(
             component_object = find_component_object(
                 network.component, info.job.flattened_arch
             )
-            workload_bpv = info.job.einsum.tensor_accesses[network.tensor].bits_per_value
+            workload_bpv = info.job.einsum.tensor_accesses[
+                network.tensor
+            ].bits_per_value
             bits_per_value = component_object.bits_per_value.get(
                 network.tensor, workload_bpv
             )
@@ -62,10 +64,13 @@ def accumulate_child_result(
             else:
                 actions_per_value = bits_per_value
             volume = (
-                compute_dense_tile_occupancy(projection, child_shape) * actions_per_value
+                compute_dense_tile_occupancy(projection, child_shape)
+                * actions_per_value
             )
 
-            if info.job.spec_one_einsum.arch.is_above(node.component, network.component):
+            if info.job.spec_one_einsum.arch.is_above(
+                node.component, network.component
+            ):
                 continue
 
             relevancy = info.tensor_to_relevancy[network.tensor][node.rank_variable]
@@ -103,4 +108,4 @@ def accumulate_child_result(
             else:
                 raise RuntimeError(f"unhandled relevancy type {relevancy}")
 
-        return self.overall_max_hops
+        return self.overall_max_hops

accelforge/model/_looptree/reuse/symbolic/symbolic/_symbolic.py

@@ -634,12 +634,7 @@ def handle_repeated_value(repeated_shape):
             )
 
             network_analyzer.accumulate_child_result(
-                child_result,
-                info,
-                shape_repeats,
-                einsum_name,
-                child_shape,
-                node
+                child_result, info, shape_repeats, einsum_name, child_shape, node
             )
 
         for einsum, child_steps in child_result.temporal_steps.items():

accelforge/model/run_model.py

@@ -192,7 +192,7 @@ def run_model(
     for key, count in simple_actions.items():
         actions_df[action2col(key)] = count.total * n_instances
 
-    if metrics & Metrics.LATENCY:
+    if metrics.includes_latency():
         df["Total<SEP>latency"] = overall_latency * n_instances
         # df[f"latency<SEP>compute"] = comp_latency * n_instances
         # For first latency, we'll follow the convention of treating compute

examples/arches/compute_in_memory/components/dac_c2c_r2r.py

@@ -602,4 +602,4 @@ def __init__(
         self.leak_power_scale *= 2  # Two sides
         self.area_scale *= 2  # Two sides
         self.energy_scale *= 2  # Two sides
-        # Voltage is still the full range because we're pulling from a VDD supply
+        # Voltage is still the full range because we're pulling from a VDD supply

examples/arches/compute_in_memory/components/zero_comparator.py

@@ -29,11 +29,7 @@ class ZeroComparator(ComponentModel):
     priority = 0.5
 
     def __init__(
-        self,
-        n_comparators: int,
-        n_bits: int,
-        tech_node: str,
-        voltage: float = 0.85
+        self, n_comparators: int, n_bits: int, tech_node: str, voltage: float = 0.85
     ):
         self.n_comparators = n_comparators
         self.n_bits = n_bits
@@ -81,7 +77,6 @@ def __init__(
                 energy_scale_function=quadratic,
                 latency_scale_function=reciprocal,
                 leak_power_scale_function=linear,
-
             )
             subcomponent.leak_power_scale *= 0.02  # Low-leakage technology
 

tests/not_working/test_network.py

@@ -92,7 +92,7 @@ def test_hierarchical_1d(self):
             * (KN / MAC_TILE)
             * M_TILE
             * KN  # temporal for n1 in mapping
-            * (MAC_TILE-1)  # multicast along X-axis of MacArray
+            * (MAC_TILE - 1)  # multicast along X-axis of MacArray
             * BITS_PER_VALUE,
         )
         self.assertEqual(