accelforge 0.0.1__py3-none-any.whl

accelforge/mapper/FFM/mappings.py

@@ -0,0 +1,510 @@
+from collections import defaultdict
+from accelforge.frontend import arch
+from accelforge.frontend.spec import Spec
+from accelforge.frontend.workload import EinsumName
+from accelforge._accelerated_imports import pd
+from accelforge.frontend.workload import TensorName
+from accelforge.mapper.FFM._make_pmappings.make_pmappings import (
+    get_num_computes,
+    get_per_tensor_size,
+)
+from typing import Union
+import numpy as np
+
+
+class Mappings:
+    """
+    A collection of mappings and their evaluation results, generated by
+    :func:`~accelforge.mapper.FFM.join_pmappings`.
+
+    Attributes
+    ----------
+    spec:
+        The specification used to generate the mappings.
+    einsum_names:
+        The names of the Einsums in these mappings.
+    data:
+        A DataFrame containing the mappings and their evaluation results. Column names
+        in the dataframe are are string separated by "<SEP>", such as
+        "Total<SEP>energy".
+    total_mappings:
+        The total number of mappings that have been explored in order to get the
+        mappings in the dataframe, equal to the full mapspace size.
+    valid_mappings:
+        The number of valid mappings that have been explored in order to get the
+        mappings in the dataframe, equal to the valid mapspace size.
+    flattened_arches:
+        A dictionary of (EinsumName, Compute Name) to lists of architecture nodes. These
+        contain the parsed and flattened architecture node for that particular Einsum
+        and compute combination.
+    parsed_specs:
+        A dictionary of Einsum names to parsed specifications. These contain the parsed
+        specification for that particular Einsum.
+    """
+
+    def __init__(
+        self,
+        spec: Spec,
+        einsum_names: list[EinsumName],
+        data: pd.DataFrame,
+        total_mappings: int,
+        valid_mappings: int,
+        flattened_arches: dict[(EinsumName, str), list[arch.Leaf]],
+        parsed_specs: dict[EinsumName, Spec],
+    ):
+        self.spec: Spec = spec
+        self.einsum_names: list[EinsumName] = einsum_names
+        self.data: pd.DataFrame = data
+        self.total_mappings: int = total_mappings
+        self.valid_mappings: int = valid_mappings
+        self.flattened_arches: dict[(EinsumName, str), list[arch.Leaf]] = (
+            flattened_arches
+        )
+        self.parsed_specs: dict[EinsumName, Spec] = parsed_specs
+
+    def num_computes(self, einsum_name: EinsumName | None = None) -> int:
+        """
+        Returns the number of computes for the given Einsum name, or total computes if
+        ``einsum_name`` is ``None``.
+        """
+        # TODO: this is not correct if there are recomputations.
+        if einsum_name is None:
+            return sum(get_num_computes(self.spec, e) for e in self.einsum_names)
+        return get_num_computes(self.spec, einsum_name)
+
+    def per_tensor_size(self) -> dict[TensorName, int]:
+        """
+        Returns a dictionary of: {Tensor name: Number of elements} for each tensor in
+        the spec.
+        """
+        return get_per_tensor_size(self.spec)
+
+    def _update(self, **kwargs):
+        data = dict(
+            spec=self.spec,
+            einsum_names=self.einsum_names,
+            data=self.data,
+            total_mappings=self.total_mappings,
+            valid_mappings=self.valid_mappings,
+            flattened_arches=self.flattened_arches,
+            parsed_specs=self.parsed_specs,
+        )
+        data.update(kwargs)
+        return Mappings(**data)
+
+    def __getitem__(self, key: str | int) -> Union[pd.Series, "Mappings"]:
+        if isinstance(key, int):
+            return self._update(data=pd.DataFrame(self.data.iloc[key]).T)
+        return self.data[key]
+
+    def __len__(self):
+        return len(self.data)
+
+    def __iter__(self):
+        return iter(self.data)
+
+    def _get_cols(self, key: str) -> list[str]:
+        found_index = None
+        found = []
+        for col in self.data.columns:
+            col = col.split("<SEP>")
+            if key not in col:
+                continue
+
+            if sum(c == key for c in col) > 1:
+                raise ValueError(
+                    f"Key {key} found multiple times in the column names. "
+                    f'Columns: "{col}"'
+                )
+
+            if found_index is not None and col.index(key) != found_index:
+                raise ValueError(
+                    f"Key {key} found at varying indexes in the column names. "
+                    f'Columns: "{col}" and "{found}"'
+                )
+            found_index = col.index(key)
+            found.append("<SEP>".join(col))
+        return found
+
+    def access(self, *keys: str) -> "Mappings":
+        """
+        Returns a new Mappings object with only the columns that contain the given keys.
+        Column names are strings separated by "<SEP>", and this method will return
+        columns with one <SEP>-separated string matching the given key. Then, for all
+        remaining columns, the key will be removed.
+
+        For example, if the columns are "Compute<SEP>Energy", and "DRAM<SEP>Energy", and
+        "DRAM<SEP>Latency", then access("Energy") will return a Mappings object with
+        columns "Compute" and "DRAM", and access("DRAM") will return a Mappings object
+        with columns "Compute" and "Latency".
+
+        If multiple keys are given, then the procedure is repeated for each key.
+
+        Parameters
+        ----------
+        keys:
+            The keys to access from the columns.
+
+        Returns
+        -------
+        A new Mappings object with only the given keys.
+        """
+        assert len(set(self.data.columns)) == len(
+            self.data.columns
+        ), "Columns must be unique"
+
+        if len(keys) != 1:
+            for k in keys:
+                self = self.access(k)
+            return self
+
+        key = keys[0]
+        col_renames = {}
+        for col in self._get_cols(key):
+            col_renames[col] = "<SEP>".join(c for c in col.split("<SEP>") if c != key)
+
+        return self._update(
+            data=self.data[list(col_renames.keys())].rename(columns=col_renames)
+        )
+
+    def _get_keys_of_length(self, length: int) -> list[str]:
+        cols = [c for c in self.columns if len(c.split("<SEP>")) == length]
+        return cols
+
+    def drop(self, *keys: str) -> "Mappings":
+        """
+        Returns a new Mappings object with the given keys dropped from all columns.
+        Column names are strings separated by "<SEP>", and this method will will drop
+        columns with one <SEP>-separated string matching the given key. Then, for all
+        remaining columns, the key will be removed.
+
+        For example, if the columns are "Compute<SEP>Energy", and "DRAM<SEP>Energy", and
+        "DRAM<SEP>Latency", then drop("Energy") will drop "Compute<SEP>Energy" and
+        "DRAM<SEP>Energy", and drop("DRAM") will drop "DRAM<SEP>Energy" and
+        "DRAM<SEP>Latency".
+
+        If multiple keys are given, then the procedure is repeated for each key.
+
+        Parameters
+        ----------
+        keys:
+            The keys to drop from the columns.
+
+        Returns
+        -------
+        A new Mappings object with the given keys dropped from all columns.
+        """
+        assert len(set(self.data.columns)) == len(
+            self.data.columns
+        ), "Columns must be unique"
+
+        if len(keys) != 1:
+            for k in keys:
+                self = self.drop(k)
+            return self
+
+        return self._update(data=self.data.drop(columns=self._get_cols(keys[0])))
+
+    def sum(self, keep_key_index: list[int] | int | None = None) -> "Mappings":
+        if len(self.data.columns) == 1:
+            return self
+
+        if isinstance(keep_key_index, int):
+            keep_key_index = [keep_key_index]
+        elif keep_key_index is None:
+            keep_key_index = []
+
+        columns = list(self.data.columns)
+        for col in self.data.columns:
+            if len(col.split("<SEP>")) != len(columns[0].split("<SEP>")):
+                raise ValueError(
+                    f"Can only sum columns with same-length keys. Try first calling "
+                    f'access("key") or drop("key") to make all columns '
+                    f"have the same number of keys."
+                )
+
+        if any(k < 0 or k >= len(columns[0].split("<SEP>")) for k in keep_key_index):
+            raise ValueError(
+                f"Keep indices must be in the range [0, {len(columns[0].split('<SEP>'))})"
+            )
+
+        target2sources = {}
+        for col in columns:
+            target = col.split("<SEP>")
+            target = "<SEP>".join(target[i] for i in keep_key_index)
+            target2sources.setdefault(target, []).append(col)
+
+        new_data = pd.DataFrame(index=self.data.index)
+        for target, sources in target2sources.items():
+            new_data[target] = self.data[sources].sum(axis=1)
+
+        return self._update(data=new_data)
+
+    @property
+    def columns(self) -> list[str]:
+        """The columns of the dataframe."""
+        return list(self.data.columns)
+
+    def to_dict(self, value_if_one_mapping: bool = True) -> dict[str, list[float]]:
+        """
+        Returns the data in this Mappings object as a dictionary. Each column in the
+        this Mappings' data becomes a key in the dictionary. Values in the dictionary
+        may be a single value if there is only one mapping, or a list of values if there
+        are multiple mappings.
+
+        Parameters
+        ----------
+        value_if_one_mapping:
+            If True and there is only one mapping, then values in the returned
+            dictionary will be a single value, rather than a list of values. Otherwise,
+            they will always be a list of values.
+
+        Returns
+        -------
+        A dictionary with the same keys as the columns of the dataframe, and values that
+        are either a single value or a list of values.
+        """
+        new = self.data.to_dict(orient="list")
+        if value_if_one_mapping and len(self) == 1:
+            new = {k: v[0] for k, v in new.items()}
+        return new
+
+    def per_compute(self, per_einsum: bool = False) -> "Mappings":
+        """
+        Returns the per-compute evaluation results by dividing all statistics by the
+        number of computes.
+
+        Parameters
+        ----------
+        per_einsum:
+            If True, then for each Einsum, statistics will be divided only by the number
+            of computes for that Einsum. This makes it clear to see per-compute stats
+            for each Einsum, but note that total energy/compute will not be a direct sum
+            of the per-compute stats then (and the same for latency and other
+            statistics).
+
+        Returns
+        -------
+        A new Mappings object with the per-compute evaluation results.
+        """
+        new_df = self.data.copy()
+        total_computes = self.num_computes()
+        for col in new_df.columns:
+            n_computes = total_computes
+            if per_einsum:
+                einsum_name = col.split("<SEP>")[0]
+                if einsum_name not in self.einsum_names and einsum_name != "Total":
+                    raise ValueError(
+                        f"Einsum name {einsum_name} not found. Ensure that all "
+                        f"columns are prefixed with the Einsum name if per_einsum "
+                        f"is True."
+                    )
+                if einsum_name != "Total":
+                    n_computes = self.num_computes(einsum_name)
+            # Check if the column can be converted to numeric
+            try:
+                new_df[col] /= n_computes
+            except (ValueError, TypeError):
+                # Skip columns that can't be converted to numeric
+                continue
+        return self._update(data=new_df)
+
+    def drop_zeros(self) -> "Mappings":
+        """
+        Returns a new Mappings object with all columns that have only zeros dropped.
+        """
+        new_df = self.data.copy()
+        new_df = new_df[(c for c in new_df.columns if (new_df[c] != 0).any())]
+        return self._update(data=new_df)
+
+    def _repr_svg_(self) -> str:
+        return self.render()
+
+    def render(self, index: int | None = None) -> str:
+        """
+        Renders the mapping as a Pydot graph. Returns an SVG string. This is only
+        supported if there is a single mapping; if there are multiple mappings, then
+        either index into this Mappings object first, or pass in an index.
+
+        Parameters
+        ----------
+        index:
+            The index of the mapping to render. If None and there are multiple mappings,
+            then an error is raised.
+
+        Returns
+        -------
+        An SVG string of the mapping.
+
+        Raises
+        ------
+        ValueError:
+            If there are multiple mappings and no index is provided.
+        """
+        if index is not None:
+            self = self[index]
+        if len(self) != 1:
+            raise ValueError(
+                f"Can only render a single mapping, but got {len(self)}. Try calling "
+                f"mappings[i].render() instead, for some integer 0 <= i < {len(self)}."
+            )
+        return self.data.iloc[0][f"Total<SEP>mapping"].render()
+
+    def energy(
+        self: "Mappings",
+        per_einsum: bool = False,
+        per_component: bool = False,
+        per_tensor: bool = False,
+        per_action: bool = False,
+        value_if_one_mapping: bool = True,
+    ) -> dict[tuple[str, ...] | str, float | list[float]] | float | list[float]:
+        """
+        Returns the energy consumed. A dictionary is returned with keys that are tuples
+        of (Einsum name, Component name, Tensor name, Action name), with any of these
+        being omitted if the corresponding parameter is not set to True. If neither of
+        the per_... parameters are set to True, a float or a list of floats is returned.
+
+        NOTE: Leak power is not per-tensor. If per_tensor is True, then the tensor name
+        for leak will be None.
+
+        Parameters
+        ----------
+        per_einsum:
+            If True, then the energy will reported per-Einsum.
+        per_component:
+            If True, then the energy will reported per-component.
+        per_tensor:
+            If True, then the energy will reported per-tensor.
+        per_action:
+            If True, then the energy will reported per-action.
+        value_if_one_mapping:
+            If True and there is only one mapping, then values in the returned
+            dictionary will be a single value, rather than a list of values. Otherwise,
+            they will always be a list of values.
+
+        Returns
+        -------
+        dict[tuple[str, ...], float | list[float]] | float | list[float]:
+            A dictionary with the energy consumed for each Einsum, Component, Tensor,
+            and Action. Keys are tuples of (Einsum name, Component name, Tensor name,
+            Action name), with any of these being omitted if the corresponding parameter
+            is not set to True. If none of the per_... parameters are set to True, a
+            float or a list of floats is returned.
+        """
+
+        energy = self.access("energy")
+
+        result = {}
+        for einsum in self.einsum_names:
+            einsum_accessed = energy.access(einsum)
+            for tensor in self.spec.workload.einsums[einsum].tensor_names:
+                tensor_accessed = einsum_accessed.access(tensor)
+                for col in tensor_accessed._get_keys_of_length(2):
+                    component, action = col.split("<SEP>")
+                    result[(einsum, component, tensor, action)] = tensor_accessed[col]
+            for col in einsum_accessed._get_keys_of_length(2):
+                component, action = col.split("<SEP>")
+                if action == "leak":
+                    result[(einsum, component, None, action)] = einsum_accessed[col]
+
+        keep_indices = []
+        for i, idx in enumerate([per_einsum, per_component, per_tensor, per_action]):
+            if idx:
+                keep_indices.append(i)
+
+        if not keep_indices:
+            v = sum(result.values())
+            if value_if_one_mapping and len(self.data) == 1:
+                return v.iloc[0]
+            return v
+
+        new_result = defaultdict(float)
+        for key, value in result.items():
+            newkey = tuple(key[i] for i in keep_indices)
+            new_result[newkey] += value
+        result = new_result
+
+        if len(keep_indices) == 1:
+            result = {k[0]: v for k, v in result.items()}
+
+        if value_if_one_mapping and len(self.data) == 1:
+            return {k: v.iloc[0] for k, v in result.items()}
+
+        return result
+
+    def latency(
+        self: "Mappings",
+        per_einsum: bool = False,
+        per_component: bool = False,
+        value_if_one_mapping: bool = True,
+    ) -> dict[tuple[str, ...] | str, float | list[float]] | float | list[float]:
+        """
+        Returns the latency consumed. A dictionary is returned with keys that are tuples
+        of (Einsum name, Component name), with either being omitted if the corresponding
+        parameter is not set to True. If neither of the per_... parameters are set to
+        True, a float or a list of floats is returned.
+
+        NOTE: If per-Einsum is False and per-component is True, then the latency of each
+        component will be summed across all Einsums. THE TOTAL LATENCY MAY BE GREATER
+        THAN THE MAX OF THE PER-COMPONENT LATENCIES. This is because different
+        components can be the bottleneck for different Einsums.
+
+        Parameters
+        ----------
+        per_einsum:
+            If True, then the latency will reported per-Einsum.
+        per_component:
+            If True, then the latency will reported per-component.
+        value_if_one_mapping:
+            If True and there is only one mapping, then values in the returned
+            dictionary will be a single value, rather than a list of values. Otherwise,
+            they will always be a list of values.
+
+        Returns
+        -------
+        dict[tuple[str, ...], float | list[float]] | float | list[float]:
+            A dictionary with the latency for each Einsum+Component pair. Keys are
+            tuples of (Einsum name, Component name), with either being omitted if the
+            corresponding parameter is not set to True. If neither of the per_...
+            parameters are set to True, a float or a list of floats is returned.
+        """
+
+        energy = self.access("latency")
+
+        result = {}
+        for einsum in self.einsum_names:
+            einsum_accessed = energy.access(einsum)
+            for component in einsum_accessed._get_keys_of_length(1):
+                result[(einsum, component)] = einsum_accessed[component]
+
+        # If not per-component, aggregate into the per-Einsum latency
+        if not per_component:
+            new_result = {}
+            for (einsum, component), value in result.items():
+                if einsum not in new_result:
+                    new_result[einsum] = value
+                else:
+                    new_result[einsum] = np.maximum(new_result[einsum], value)
+            result = new_result
+
+        if not per_einsum:
+            # Not per-Einsum and per-component: for each component, sum across Einsums
+            if per_component:
+                new_result = {}
+                for (einsum, component), value in result.items():
+                    if component not in new_result:
+                        new_result[component] = value
+                    else:
+                        new_result[component] += value
+                result = new_result
+
+            # Not per-Einsum and not per-component: sum into a single value
+            else:
+                result = np.sum(list(result.values()))
+
+        if value_if_one_mapping and len(self.data) == 1:
+            if isinstance(result, dict):
+                return {k: v.iloc[0] for k, v in result.items()}
+            return result  # Numpy sum already pulls out the number
+
+        return result