accelforge 0.0.1__py3-none-any.whl

accelforge/frontend/model.py

@@ -0,0 +1,14 @@
+from accelforge.frontend.mapper.metrics import Metrics
+from accelforge.util._basetypes import ParsableModel
+
+
+class Model(ParsableModel):
+    """Configuration for the model."""
+
+    metrics: Metrics = Metrics.all_metrics()
+    """
+    Metrics to evaluate.
+
+    If using spec to call mapper, leave this configuration as is. The mapper
+    will make necessary configurations.
+    """

accelforge/frontend/renames.py

@@ -0,0 +1,150 @@
+import copy
+from typing import Annotated, Any, TypeAlias
+from accelforge.util._basetypes import (
+    ParsableList,
+    ParsableModel,
+    ParsesTo,
+    TryParseTo,
+    _PostCall,
+)
+from accelforge._version import assert_version, __version__
+from accelforge.util._parse_expressions import ParseError
+from accelforge.util._setexpressions import InvertibleSet
+
+TensorName: TypeAlias = str
+RankVariable: TypeAlias = str
+Rank: TypeAlias = str
+EinsumName: TypeAlias = str
+
+
+class Rename(ParsableModel):
+    """
+    A rename of something into something else.
+    """
+
+    name: str
+    """ The name of the thing to be renamed. This is a string representing the new name."""
+
+    source: TryParseTo[InvertibleSet[TensorName | RankVariable]]
+    """ The source of the rename. This is a set expression that can be parsed, yielding
+    a set that can be referenced by the new name. """
+
+    expected_count: ParsesTo[int] | None = None
+    """
+    The expected count of the source set expression. If this is set, then the source
+    expression must resolve to the expected count or an error will be raised. Otherwise,
+    any count (including zero for an empty set) is allowed.
+    """
+
+    def _parse_expressions(self, symbol_table: dict[str, Any], *args, **kwargs):
+        parsed, symbol_table = super()._parse_expressions(symbol_table, *args, **kwargs)
+        expected_count = parsed.expected_count
+        if (
+            expected_count is not None
+            and isinstance(parsed.source, InvertibleSet)
+            and len(parsed.source) != expected_count
+        ):
+            parsed, symbol_table = super()._parse_expressions(
+                symbol_table, *args, **kwargs
+            )
+            raise ParseError(
+                f"Expected count is {parsed.expected_count}, but got "
+                f"{len(parsed.source)}: {parsed.source}",
+                source_field="source",
+            )
+        return parsed, symbol_table
+
+
+def rename_list_factory(rename_list: list | dict) -> "RenameList":
+    if isinstance(rename_list, list):
+        return RenameList(rename_list)
+
+    if not isinstance(rename_list, dict):
+        raise TypeError(
+            f"Expected a list or dict, got {type(rename_list)}: {rename_list}"
+        )
+
+    return RenameList(
+        Rename(name=k, source=v, expected_count=None) for k, v in rename_list.items()
+    )
+
+
+class RenameList(ParsableList[Rename]):
+    """A list of renames."""
+
+    def __dict__(self) -> dict[str, Any]:
+        return {r.name: r.source for r in self}
+
+    def _parse_expressions(self, symbol_table: dict[str, Any], *args, **kwargs):
+
+        cur_symbol_table = symbol_table.copy()
+
+        class PostCallRenameList(_PostCall[Rename]):
+            def __call__(self, field, value, parsed, symbol_table):
+                symbol_table[parsed.name] = parsed.source
+                return parsed
+
+        new, _ = super()._parse_expressions(
+            cur_symbol_table, *args, **kwargs, post_calls=(PostCallRenameList(),)
+        )
+        return new, symbol_table
+
+
+class EinsumRename(ParsableModel):
+    """
+    Renames for a single Einsum.
+    """
+
+    name: EinsumName
+    """ The name of the Einsum. Set this to "default" to apply the renames to all
+    Einsums, unless overridden. Overriding is specific to a single name, so every rename
+    in the default must be overridden independently. """
+
+    tensor_accesses: ParsableList[Rename] = ParsableList()
+    """ Renames for the tensor accesses of this Einsum. This may be given either as a
+    dictionary ``{new_name: source_set_expression}`` expressions, or as a list of
+    dictionaries, each one having the structure ``{name: new_name, source:
+    source_set_expression, expected_count: 1}``, where expected count is optional for
+    each and may be set to any integer. """
+
+    rank_variables: ParsableList[Rename] = ParsableList()
+    """ Renames for the rank variables of this Einsum. This may be given either as a
+    dictionary ``{new_name: source_set_expression}`` expressions, or as a list of
+    dictionaries, each one having the structure ``{name: new_name, source:
+    source_set_expression, expected_count: 1}``, where expected count is optional for
+    each and may be set to any integer. """
+
+    def __init__(self, *args, **kwargs) -> None:
+        if "tensor_accesses" in kwargs:
+            kwargs["tensor_accesses"] = rename_list_factory(kwargs["tensor_accesses"])
+        if "rank_variables" in kwargs:
+            kwargs["rank_variables"] = rename_list_factory(kwargs["rank_variables"])
+        super().__init__(*args, **kwargs)
+
+
+class Renames(ParsableModel):
+    # version: Annotated[str, assert_version] = __version__
+    einsums: list[EinsumRename] = list()
+    """
+    Renames for a workload. The Einsum list is a list of EinsumRename objects, and
+    renames will be applied to Einsums whose names match the EinsumRename.name. If an
+    EinsumRename is named "default", then its renames are applied to every Einsum unless
+    overridden. Overriding is specific to a single name, so every rename in the default
+    must be overridden independently.
+    """
+
+    def get_renames_for_einsum(self, einsum_name: EinsumName) -> EinsumRename:
+        if einsum_name not in self.einsums:
+            rename = EinsumRename(name=einsum_name)
+        else:
+            rename = copy.deepcopy(self.einsums[einsum_name])
+        for einsum in self.einsums:
+            if einsum.name != "default":
+                continue
+            for tensor_rename in einsum.tensor_accesses:
+                if tensor_rename.name not in rename.tensor_accesses:
+                    rename.tensor_accesses.append(tensor_rename)
+            for rank_variable_rename in einsum.rank_variables:
+                if rank_variable_rename.name not in rename.rank_variables:
+                    rename.rank_variables.append(rank_variable_rename)
+        return rename

accelforge/frontend/spec copy.py

@@ -0,0 +1,230 @@
+from accelforge.frontend.mapper.mapper import Mapper
+from accelforge.frontend.renames import EinsumName, Renames
+from accelforge.util._parse_expressions import ParseError, ParseExpressionsContext
+from accelforge.frontend.arch import Compute, Leaf, Component, Arch, Fanout
+
+from accelforge.frontend.workload import Workload
+from accelforge.frontend.variables import Variables
+from accelforge.frontend.config import Config, get_config
+from accelforge.frontend.mapping import Mapping
+from accelforge.frontend.model import Model
+import hwcomponents
+
+from typing import Any, Dict, Optional, Self
+from accelforge.util._basetypes import ParsableModel
+from pydantic import Field
+
+
+class Spec(ParsableModel):
+    """The top-level spec of all of the inputs to this package."""
+
+    arch: Arch = Arch()
+    """ The hardware architecture being used. """
+
+    mapping: Mapping = Mapping()
+    """ How the workload is programmed onto the architecture. Do not specify this if
+    you'd like the mapper to generate a mapping for you. """
+
+    workload: Workload = Workload()
+    """ The program to be run on the arch. """
+
+    variables: Variables = Variables()
+    """ Variables that can be referenced in other places in the spec. """
+
+    config: Config = Field(default_factory=get_config)
+    """ Configuration settings. """
+
+    renames: Renames = Renames()
+    """ Aliases for tensors in the workload so that they can be called
+    by canonical names when writing architecture constraints. For example, workload
+    tensors may be renamed to "input", "output", and "weight"."""
+
+    mapper: Mapper = Mapper()
+    """ Configures the mapper used to map the workload onto the architecture. """
+
+    model: Model = Model()
+    """Configures the model used to evaluate mappings."""
+
+    def _parse_expressions(
+        self,
+        symbol_table: Optional[Dict[str, Any]] = None,
+        **kwargs,
+    ) -> Self:
+        """
+        Parse all string expressions in the spec into concrete values.
+
+        Parameters
+        ----------
+        symbol_table : dict, optional
+            Optional pre-populated symbols to seed parsing; a shallow copy is made and
+            augmented with ``spec`` and ``variables``.
+        kwargs : dict, optional
+            Additional keyword arguments forwarded to the base
+            ``ParsableModel._parse_expressions``.
+
+        Returns
+        -------
+        Self
+            The parsed specification.
+
+        Raises
+        ------ParseError
+            If any field fails to parse; the error is annotated with the field path.
+        """
+        symbol_table = {} if symbol_table is None else symbol_table.copy()
+        symbol_table["spec"] = self
+        with ParseExpressionsContext(self):
+            try:
+                parsed_variables, _ = self.variables._parse_expressions(
+                    symbol_table, **kwargs
+                )
+            except ParseError as e:
+                e.add_field("Spec().variables")
+                raise e
+            symbol_table.update(parsed_variables)
+            symbol_table["variables"] = parsed_variables
+            parsed_spec, _ = super()._parse_expressions(symbol_table, **kwargs)
+            return parsed_spec, symbol_table
+
+    def calculate_component_area_energy_latency_leak(
+        self,
+        area: bool = True,
+        energy: bool = True,
+        latency: bool = True,
+        leak: bool = True,
+        einsum_name: EinsumName | None = None,
+    ) -> "Spec":
+        """
+        Populates per-component area, energy, latency, and/or leak power. For each
+        component, populates the ``attributes.area``, ``attributes.total_area``,
+        ``attributes.leak_power`` and ``attributes.total_leak_power``. Additionally, for
+        each action of each component, populates the ``arguments.energy`` and
+        ``arguments.latency`` fields. Extends the ``component_modeling_log`` field with
+        log messages. Also populates the ``component_model`` attribute for each
+        component if not already set.
+
+        Some architectures may reference variables that depend on the workload. In that
+        case, an Einsum name can be provided to populate those symbols with the Einsum's
+        symbols from the workload.
+
+        Parameters
+        ----------
+        area : bool, optional
+            Whether to compute and populate area entries.
+        energy : bool, optional
+            Whether to compute and populate energy entries.
+        latency : bool, optional
+            Whether to compute and populate latency entries.
+        leak : bool, optional
+            Whether to compute and populate leak power entries.
+        symbol_table : dict[str, Any], optional
+            Symbol table to use for parsing expressions.
+        """
+        if not area and not energy and not latency and not leak:
+            return self
+
+        models = hwcomponents.get_models(
+            self.config.component_models,
+            include_installed=self.config.use_installed_component_models,
+        )
+
+        components = set()
+        if not getattr(self, "_parsed", False):
+            self = self._parse_expressions()
+        else:
+            self = self.copy()
+
+        for arch in self._get_flattened_architecture():
+            fanout = 1
+            for component in arch:
+                fanout *= component.get_fanout()
+                if component.name in components or isinstance(component, Fanout):
+                    continue
+                assert isinstance(component, Component)
+                components.add(component.name)
+                orig: Component = self.arch.find(component.name)
+                if area:
+                    c = component.calculate_area(models)
+                    orig.attributes.area = c.attributes.area
+                    orig.attributes.total_area = c.attributes.area * fanout
+                if energy:
+                    c = component.calculate_action_energy(models)
+                    for a in c.actions:
+                        orig_action = orig.actions[a.name]
+                        orig_action.arguments.energy = a.arguments.energy
+                if latency:
+                    c = component.calculate_action_latency(models)
+                    for a in c.actions:
+                        orig_action = orig.actions[a.name]
+                        orig_action.arguments.latency = a.arguments.latency
+                if leak:
+                    c = component.calculate_leak_power(models)
+                    orig.attributes.leak_power = c.attributes.leak_power
+                    orig.attributes.total_leak_power = c.attributes.leak_power * fanout
+                orig.component_modeling_log.extend(c.component_modeling_log)
+                orig.component_model = c.component_model
+
+        return self
+
+    def _get_flattened_architecture(
+        self,
+        compute_node: str | Compute | None = None,
+        symbol_table: dict[str, Any] | None = None,
+    ) -> list[list[Leaf]] | list[Leaf]:
+        """
+        Return the architecture as paths of ``Leaf`` instances from the highest-level
+        node to each ``Compute`` node. Parses arithmetic expressions in the
+        architecture for each one. If a symbol table is provided, it will be used to
+        parse the expressions.
+
+        Parameters
+        ----------
+        compute_node : str or Compute, optional
+            Optional compute node (name or ``Compute``) to restrict results to a single
+            compute node.
+        symbol_table : dict[str, Any], optional
+            Symbol table to use for parsing expressions.
+
+        Returns
+        -------
+            - If ``compute_node`` is ``None``: list of lists of ``Leaf`` for all compute
+              nodes.
+            - Otherwise: a single-item list containing the list of ``Leaf`` for the
+              requested node.
+
+        Raises
+        ------
+        AssertionError
+            If the spec has not been parsed.
+        ParseError
+            If there are duplicate names or the requested compute node cannot be found.
+        """
+        # Assert that we've been parsed
+        assert getattr(
+            self, "_parsed", False
+        ), "Spec must be parsed before getting flattened architecture"
+        all_leaves = self.arch.get_nodes_of_type(Leaf)
+        found_names = set()
+        for leaf in all_leaves:
+            if leaf.name in found_names:
+                raise ParseError(f"Duplicate name in architecture: {leaf.name}")
+            found_names.add(leaf.name)
+
+        found = []
+        if compute_node is None:
+            compute_nodes = [c.name for c in self.arch.get_nodes_of_type(Compute)]
+        else:
+            compute_nodes = [
+                compute_node.name if isinstance(compute_node, Compute) else compute_node
+            ]
+
+        for c in compute_nodes:
+            variables = {**self.variables.model_dump(), **(symbol_table or {})}
+            found.append(self.arch._flatten(variables, c))
+            if found[-1][-1].name != c:
+                raise ParseError(f"Compute node {c} not found in architecture")
+
+        return found if compute_node is None else [found[0]]
+
+
+Specification = Spec

accelforge/frontend/spec.py

@@ -0,0 +1,301 @@
+from accelforge.frontend.mapper.mapper import Mapper
+from accelforge.frontend.renames import EinsumName, Renames
+from accelforge.util._parse_expressions import ParseError, ParseExpressionsContext
+from accelforge.frontend.arch import Compute, Leaf, Component, Arch, Fanout
+
+from accelforge.frontend.workload import Workload
+from accelforge.frontend.variables import Variables
+from accelforge.frontend.config import Config, get_config
+from accelforge.frontend.mapping import Mapping
+from accelforge.frontend.model import Model
+import hwcomponents
+
+from typing import Any, Dict, Optional, Self
+from accelforge.util._basetypes import ParsableModel
+from pydantic import Field
+
+
+class Spec(ParsableModel):
+    """The top-level spec of all of the inputs to this package."""
+
+    arch: Arch = Arch()
+    """ The hardware architecture being used. """
+
+    mapping: Mapping = Mapping()
+    """ How the workload is programmed onto the architecture. Do not specify this if
+    you'd like the mapper to generate a mapping for you. """
+
+    workload: Workload = Workload()
+    """ The program to be run on the arch. """
+
+    variables: Variables = Variables()
+    """ Variables that can be referenced in other places in the spec. """
+
+    config: Config = Field(default_factory=get_config)
+    """ Configuration settings. """
+
+    renames: Renames = Renames()
+    """ Aliases for tensors in the workload so that they can be called
+    by canonical names when writing architecture constraints. For example, workload
+    tensors may be renamed to "input", "output", and "weight"."""
+
+    mapper: Mapper = Mapper()
+    """ Configures the mapper used to map the workload onto the architecture. """
+
+    model: Model = Model()
+    """Configures the model used to evaluate mappings."""
+
+    def _parse_expressions(
+        self,
+        einsum_name: EinsumName | None = None,
+        **kwargs,
+    ) -> tuple[Self, dict[str, Any]]:
+        raise NotImplementedError("Call _spec_parse_expressions instead.")
+
+    def _spec_parse_expressions(
+        self,
+        einsum_name: EinsumName | None = None,
+        _parse_arch: bool = True,
+        _parse_non_arch: bool = True,
+    ) -> Self:
+        """
+        Parse all string expressions in the spec into concrete values.
+
+        Parameters
+        ----------
+        einsum_name: EinsumName | None = None
+            Optional Einsum name to populate symbols with the Einsum's symbols from the
+            workload. If None, no symbols are populated from the workload.
+
+        _parse_arch: bool = True
+            Whether to parse the architecture.
+
+        _parse_non_arch: bool = True
+            Whether to parse the non-architecture fields.
+
+        Returns
+        -------
+        Self
+            The parsed specification.
+
+        Raises
+        ------
+        ParseError
+            If any field fails to parse; the error is annotated with the field path.
+        """
+        st = {}
+        st["spec"] = self
+        with ParseExpressionsContext(self):
+            already_parsed = {}
+
+            parsed_variables = self.variables
+            if _parse_non_arch:
+                try:
+                    parsed_variables, st = self.variables._parse_expressions(st)
+                except ParseError as e:
+                    e.add_field("Spec().variables")
+                    raise e
+            already_parsed["variables"] = parsed_variables
+            st.update(parsed_variables)
+            st["variables"] = parsed_variables
+
+            parsed_renames = self.renames
+            if _parse_non_arch:
+                try:
+                    parsed_renames, st = self.renames._parse_expressions(st)
+                except ParseError as e:
+                    e.add_field("Spec().renames")
+                    raise e
+            already_parsed["renames"] = parsed_renames
+            st["renames"] = parsed_renames
+
+            parsed_workload = self.workload
+            if _parse_non_arch:
+                try:
+                    parsed_workload, st = self.workload._parse_expressions(
+                        st, renames=parsed_renames
+                    )
+                except ParseError as e:
+                    e.add_field("Spec().workload")
+                    raise e
+            already_parsed["workload"] = parsed_workload
+            st["workload"] = parsed_workload
+
+            if einsum_name is not None:
+                renames = parsed_workload.einsums[einsum_name].renames
+                st.update(**{k.name: k.source for k in renames})
+
+            if _parse_arch:
+                parsed_arch, st = self.arch._parse_expressions(st)
+            else:
+                parsed_arch = self.arch
+            st["arch"] = parsed_arch
+            already_parsed["arch"] = parsed_arch
+
+            parsed_spec, _ = super()._parse_expressions(
+                st,
+                already_parsed=already_parsed,
+            )
+            parsed_spec._parsed = True
+            return parsed_spec
+
+    def calculate_component_area_energy_latency_leak(
+        self,
+        einsum_name: EinsumName | None = None,
+        area: bool = True,
+        energy: bool = True,
+        latency: bool = True,
+        leak: bool = True,
+    ) -> "Spec":
+        """
+        Populates per-component area, energy, latency, and/or leak power. For each
+        component, populates the ``area``, ``total_area``, ``leak_power`` and
+        ``total_leak_power``. Additionally, for each action of each component, populates
+        the ``<action>.energy`` and ``<action>.latency`` fields. Extends the
+        ``component_modeling_log`` field with log messages. Also populates the
+        ``component_model`` attribute for each component if not already set.
+
+        Some architectures' attributes may depend on the workload. In that case, an
+        Einsum name can be provided to populate those symbols with the Einsum's symbols
+        from the workload.
+
+        Parameters
+        ----------
+        einsum_name: EinsumName | None = None
+            Optional Einsum name to populate symbols with the Einsum's symbols from the
+            workload. If None, and there are Einsums in the workload, the first Einsum
+            is used. If None and there are no Einsums in the workload, then no symbols
+            are populated from the workload.
+        area : bool, optional
+            Whether to compute and populate area entries.
+        energy : bool, optional
+            Whether to compute and populate energy entries.
+        latency : bool, optional
+            Whether to compute and populate latency entries.
+        leak : bool, optional
+            Whether to compute and populate leak power entries.
+        """
+        if not area and not energy and not latency and not leak:
+            return self
+
+        models = hwcomponents.get_models(
+            self.config.component_models,
+            include_installed=self.config.use_installed_component_models,
+        )
+
+        if einsum_name is None and len(self.workload.einsums) > 0:
+            einsum_name = self.workload.einsums[0].name
+
+        components = set()
+        try:
+            if not getattr(self, "_parsed", False):
+                self = self._spec_parse_expressions(einsum_name=einsum_name)
+            else:
+                self = self.copy()
+        except ParseError as e:
+            if "arch" in e.message:
+                e.add_note(
+                    "If this error seems to be caused by a missing symbol that depends on \n"
+                    "the workload, you may need to provide an appropriate einsum_name to \n"
+                    "calculate_component_area_energy_latency_leak. This may occur if the \n"
+                    "architecture depends on something in the workload.\n"
+                )
+            raise
+
+        for arch in self._get_flattened_architecture():
+            fanout = 1
+            for component in arch:
+                fanout *= component.get_fanout()
+                if component.name in components or isinstance(component, Fanout):
+                    continue
+                assert isinstance(component, Component)
+                components.add(component.name)
+                orig: Component = self.arch.find(component.name)
+                c = component
+                if area:
+                    c = c.calculate_area(models)
+                    orig.area = c.area
+                    orig.total_area = c.area * fanout
+                if energy:
+                    c = c.calculate_action_energy(models)
+                    for a in c.actions:
+                        orig_action = orig.actions[a.name]
+                        orig_action.energy = a.energy
+                if latency:
+                    c = c.calculate_action_latency(models)
+                    for a in c.actions:
+                        orig_action = orig.actions[a.name]
+                        orig_action.latency = a.latency
+                if leak:
+                    c = c.calculate_leak_power(models)
+                    orig.leak_power = c.leak_power
+                    orig.total_leak_power = c.leak_power * fanout
+                orig.component_modeling_log.extend(c.component_modeling_log)
+                orig.component_model = c.component_model
+
+        return self
+
+    def _get_flattened_architecture(
+        self,
+        compute_node: str | Compute | None = None,
+    ) -> list[list[Leaf]] | list[Leaf]:
+        """
+        Return the architecture as paths of ``Leaf`` instances from the highest-level
+        node to each ``Compute`` node. Parses arithmetic expressions in the
+        architecture for each one. If a symbol table is provided, it will be used to
+        parse the expressions.
+
+        Parameters
+        ----------
+        compute_node : str or Compute, optional
+            Optional compute node (name or ``Compute``) to restrict results to a single
+            compute node.
+
+        Returns
+        -------
+            - If ``compute_node`` is ``None``: list of lists of ``Leaf`` for all compute
+              nodes.
+            - Otherwise: a single-item list containing the list of ``Leaf`` for the
+              requested node.
+
+        Raises
+        ------
+        AssertionError
+            If the spec has not been parsed.
+        ParseError
+            If there are duplicate names or the requested compute node cannot be found.
+        """
+        # Assert that we've been parsed
+        assert getattr(
+            self, "_parsed", False
+        ), "Spec must be parsed before getting flattened architecture"
+        all_leaves = self.arch.get_nodes_of_type(Leaf)
+        found_names = set()
+        for leaf in all_leaves:
+            if leaf.name in found_names:
+                raise ParseError(f"Duplicate name in architecture: {leaf.name}")
+            found_names.add(leaf.name)
+
+        found = []
+        if compute_node is None:
+            compute_nodes = [c.name for c in self.arch.get_nodes_of_type(Compute)]
+        else:
+            compute_nodes = [
+                compute_node.name if isinstance(compute_node, Compute) else compute_node
+            ]
+
+        for c in compute_nodes:
+            found.append(self.arch._flatten(c))
+            if found[-1][-1].name != c:
+                raise ParseError(f"Compute node {c} not found in architecture")
+
+        # These can't be pickled if they use dynamically-loaded code
+        for f in found:
+            for n in f:
+                if hasattr(n, "component_model"):
+                    n.component_model = None
+
+        return found if compute_node is None else found[0]
+
+
+Specification = Spec