accelforge 0.0.1__py3-none-any.whl

accelforge/frontend/arch.py

@@ -0,0 +1,1642 @@
+import copy
+import itertools
+import logging
+import math
+from numbers import Number
+import re
+from typing import (
+    Any,
+    Callable,
+    Iterator,
+    List,
+    Literal,
+    Optional,
+    Self,
+    TypeVar,
+    Annotated,
+    Type,
+    Union,
+)
+from pydantic import ConfigDict, Tag
+import pydantic
+from hwcomponents import (
+    ComponentModel,
+    get_models,
+    get_model,
+)
+import pydot
+
+from accelforge.util._basetypes import (
+    ParsableModel,
+    ParsableList,
+    ParseExtras,
+    ParsesTo,
+    TryParseTo,
+    _PostCall,
+    _get_tag,
+)
+import numpy as np
+
+from accelforge.util._parse_expressions import ParseError, parse_expression
+from accelforge.util._setexpressions import InvertibleSet, eval_set_expression
+from accelforge.frontend.renames import RankVariable, TensorName
+
+from accelforge._version import assert_version, __version__
+from pydantic import Discriminator
+from accelforge.util._basetypes import _uninstantiable
+from accelforge.util.parallel import _SVGJupyterRender
+from accelforge.util._visualization import _pydot_graph
+
+T = TypeVar("T", bound="ArchNode")
+
+
+class ArchNode(ParsableModel):
+    """A node in the architecture."""
+
+    def find(self, name: str) -> "Leaf":
+        """
+        Finds a `Leaf` node with the given name.
+        :raises ValueError: If the `Leaf` node with the given name is not found.
+        """
+        if isinstance(self, Leaf) and getattr(self, "name", None) == name:
+            return self
+
+        if isinstance(self, Branch):
+            for element in self.nodes:
+                try:
+                    return element.find(name)
+                except (AttributeError, ValueError):
+                    pass
+        raise ValueError(f"Leaf {name} not found in {self}")
+
+
+class ArchNodes(ParsableList):
+    """A list of `ArchNode`s."""
+
+    def __repr__(self):
+        return f"{self.__class__.__name__}({super().__repr__()})"
+
+    def _parse_expressions(self, symbol_table: dict[str, Any], *args, **kwargs):
+        class PostCallArchNode(_PostCall):
+            def __call__(self, field, value, parsed, symbol_table):
+                symbol_table[parsed.name] = parsed
+                return parsed
+
+        for i, node in enumerate(self):
+            symbol_table[i] = node
+
+        return super()._parse_expressions(
+            symbol_table, *args, **kwargs, post_calls=(PostCallArchNode(),)
+        )
+
+
+class Comparison(ParsableModel):
+    """
+    A comparison between a rank variable's bound and a value. A comparison is performed
+    for each rank variable.
+
+    The LHS of each comparison is the loop bound of a loop that affects this rank
+    variable. The RHS is the given value.
+
+    For example, if the expression resolves to [a, b], the operator is "<=", and the
+    value is 10, and we have loops "for a0 in [0..A0)" and "for b0 in [0..B0)", then a
+    mapping is only valid if A0 <= 10 and B0 <= 10.
+    """
+
+    expression: TryParseTo[InvertibleSet[RankVariable]]
+    """ The expression to compare. This expression should resolve to a set of rank
+    variables. A comparison is performed for each rank variable independently, and the
+    result passes if and only if all comparisons pass. The LHS of each comparison is the
+    loop bound of a loop that affects this rank variable. The RHS is the given value.
+    """
+
+    operator: str
+    """ The operator to use for the comparison. Supported operators are:
+    - == (equal to)
+    - <= (less than or equal to)
+    - >= (greater than or equal to)
+    - < (less than)
+    - > (greater than)
+    - product== (product of all loop bounds is equal to)
+    - product<= (product of all loop bounds is less than or equal to)
+    - product>= (product of all loop bounds is greater than or equal to)
+    - product< (product of all loop bounds is less than)
+    - product> (product of all loop bounds is greater than)
+    """
+
+    value: ParsesTo[int]
+    """ The value to compare against. """
+
+    _str_repr: str = None
+    """ A string to print for this comparison when __str__ is called. If None, a default
+    string will be used. """
+
+    def _parse_expressions(self, *args, **kwargs):
+        result, symbol_table = super()._parse_expressions(*args, **kwargs)
+        if len(result.expression) == 1 and "product" in result.operator:
+            result.operator = result.operator.replace("product", "")
+        return result, symbol_table
+
+    # def _parse(self, symbol_table: dict[str, Any], location: str):
+    #     # if len(self) != 3:
+    #     #     raise ValueError(f"Comparison can only have 3 elements. got {len(self)}")
+    #     new = type(self)(
+    #         expression=eval_set_expression(
+    #             self.expression, symbol_table, "rank_variables", location
+    #         ),
+    #         operator=self.operator,
+    #         value=self.value,
+    #     )
+    #     if len(new.expression) == 1 and "product" in new.operator:
+    #         new.operator = new.operator.replace("product", "")
+    #     return new
+
+    def _constrained_to_one(self) -> bool:
+        return self.value == 1 and self.operator in [
+            "==",
+            "<=",
+            "product==",
+            "product<=",
+        ]
+
+    def _split_expression(self) -> List[set[RankVariable]]:
+        if "product" in self.operator:
+            return [self.expression]
+        return sorted(set((x,)) for x in self.expression)
+
+    def _to_constraint_lambda(
+        self,
+        increasing_sizes: bool,
+    ) -> Callable[[bool, np.ndarray], bool | np.ndarray]:
+        # Equal operators can only evaluate when all sizes are known
+        eq_op = lambda final: (
+            np.equal
+            if final
+            else (np.less_equal if increasing_sizes else np.greater_equal)
+        )
+
+        # If we're increasing, we can evaluate leq immediately. If we're
+        # decreasing, we can evaluate geq immediately. The other must wait
+        # until all sizes are known.
+        le_wrapper = lambda op: lambda final, sizes: (
+            op(sizes) if final or increasing_sizes else True
+        )
+        ge_wrapper = lambda op: lambda final, sizes: (
+            op(sizes) if final or not increasing_sizes else True
+        )
+
+        _all = lambda sizes: np.all(sizes, axis=1)
+        _prod = lambda sizes: np.prod(sizes, axis=1)
+
+        # fmt: off
+        operator_to_wrapper = {
+            "==":        lambda final, sizes: _all(eq_op(final)(sizes, self.value)),
+            "product==": lambda final, sizes: eq_op(final)(_prod(sizes), self.value),
+            "<=":        le_wrapper(lambda sizes: _all(sizes)  <= self.value),
+            ">=":        ge_wrapper(lambda sizes: _all(sizes)  >= self.value),
+            "<":         le_wrapper(lambda sizes: _all(sizes)  <  self.value),
+            ">":         ge_wrapper(lambda sizes: _all(sizes)  >  self.value),
+            "product<=": le_wrapper(lambda sizes: _prod(sizes) <= self.value),
+            "product>=": ge_wrapper(lambda sizes: _prod(sizes) >= self.value),
+            "product<":  le_wrapper(lambda sizes: _prod(sizes) <  self.value),
+            "product>":  ge_wrapper(lambda sizes: _prod(sizes) >  self.value),
+        }
+        # fmt: on
+
+        if self.operator in operator_to_wrapper:
+            return operator_to_wrapper[self.operator]
+        raise KeyError(
+            f"Unknown operator: {self.operator}. Known operators: {list(operator_to_wrapper.keys())}"
+        )
+
+    def __str__(self) -> str:
+        if self._str_repr is not None:
+            return self._str_repr
+        return f"({sorted(self.expression)}) {self.operator} ({self.value})"
+
+
+class Spatial(ParsableModel):
+    """A one-dimensional spatial fanout in the architecture."""
+
+    name: str
+    """
+    The name of the dimension over which this spatial fanout is occurring (e.g., X or Y).
+    """
+
+    fanout: ParsesTo[int]
+    """ The size of this fanout. """
+
+    may_reuse: TryParseTo[InvertibleSet[TensorName]] = "All"
+    """ The tensors that can be reused spatially across instances of this fanout. This
+    expression will be parsed for each mapping template. """
+
+    loop_bounds: ParsableList[Comparison] = ParsableList()
+    """ Bounds for loops over this dimension. This is a list of :class:`~.Comparison`
+    objects, all of which must be satisfied by the loops to which this constraint
+    applies.
+
+    Note: Loops may be removed if they are constrained to only one iteration.
+    """
+
+    min_usage: int | float | str = 0.0
+    """ The minimum usage of spatial instances, as a value from 0 to 1. A mapping
+    is invalid if less than this porportion of this dimension's fanout is utilized.
+    Mappers that support it (e.g., FFM) may, if no mappings satisfy this constraint,
+    return the highest-usage mappings.
+    """
+
+    reuse: TryParseTo[InvertibleSet[TensorName]] = "Nothing"
+    """ A set of tensors or a set expression representing tensors that must be reused
+    across spatial iterations. Spatial loops may only be placed that reuse ALL tensors
+    given here.
+
+    Note: Loops may be removed if they do not reuse a tensor given here and they do not
+    appear in another loop bound constraint.
+    """
+
+    usage_scale: ParsesTo[int | float | str] = 1
+    """
+    This factor scales the usage in this dimension. For example, if usage_scale is 2 and
+    10/20 spatial instances are used, then the usage will be scaled to 20/20.
+    """
+
+    power_gateable: ParsesTo[bool] = False
+    """
+    Whether this spatial fanout has power gating. If True, then unused spatial instances
+    will be power gated if not used by a particular Einsum.
+    """
+
+    # def _parse(self, symbol_table: dict[str, Any], location: str):
+    #     return type(self)(
+    #         name=self.name,
+    #         fanout=self.fanout,
+    #         may_reuse=set(
+    #             eval_set_expression(
+    #                 self.may_reuse,
+    #                 symbol_table,
+    #                 expected_space=TensorName,
+    #                 location=location + ".may_reuse",
+    #             )
+    #         ),
+    #         loop_bounds=[
+    #             x._parse(symbol_table, location + ".loop_bounds")
+    #             for x in self.loop_bounds
+    #         ],
+    #         min_usage=parse_expression(
+    #             self.min_usage,
+    #             symbol_table,
+    #             "min_usage",
+    #             location + ".min_usage",
+    #         ),
+    #         reuse=eval_set_expression(
+    #             self.reuse,
+    #             symbol_table,
+    #             "tensors",
+    #             location + ".reuse",
+    #         ),
+    #     )
+
+
+class Action(ParsableModel):
+    """
+    An action that may be performed by a component.
+    """
+
+    name: str
+    """ The name of this action. """
+
+    energy: ParsesTo[int | float | None] = None
+    """
+    Dynamic energy of this action. Per-action energy is multiplied by the component's
+    energy_scale and the action's energy_scale.
+    """
+
+    energy_scale: ParsesTo[int | float] = 1
+    """
+    The scale factor for dynamic energy of this action. Multiplies this action's energy
+    by this value.
+    """
+
+    latency: ParsesTo[int | float | None] = None
+    """
+    Latency of this action. Per-action latency is multiplied by the component's
+    latency_scale and the action's latency_scale.
+    """
+
+    latency_scale: ParsesTo[int | float] = 1
+    """
+    The scale factor for dynamic latency of this action. Multiplies this action's
+    latency by this value.
+    """
+
+    extra_attributes_for_component_model: ParseExtras = ParseExtras()
+    """ Extra attributes to pass to the component model. In addition to all attributes
+    of this action, any extra attributes will be passed to the component model as
+    arguments to the component model's action. This can be used to define attributes
+    that are known to the component model, but not accelforge, such as clock
+    frequency."""
+
+    def _attributes_for_component_model(self) -> dict[str, Any]:
+        return {
+            **self.shallow_model_dump(),
+            **self.extra_attributes_for_component_model.shallow_model_dump(),
+        }
+
+
+class TensorHolderAction(Action):
+    bits_per_action: ParsesTo[int | float] = (
+        "1 if bits_per_action is None else bits_per_action"
+    )
+    """ The number of bits accessed in this action. For example, setting bits_per_action
+    to 16 means that each call to this action yields 16 bits. """
+
+
+@_uninstantiable
+class Leaf(ArchNode):
+    """A leaf node in the architecture. This is an abstract class that represents any
+    node that is not a `Branch`."""
+
+    name: str
+    """ The name of this `Leaf`. """
+
+    spatial: ParsableList[Spatial] = ParsableList()
+    """
+    The spatial fanouts of this `Leaf`.
+
+    Spatial fanouts describe the spatial organization of components in the architecture.
+    A spatial fanout of size N for this node means that there are N instances of this
+    node. Multiple spatial fanouts lead to a multi-dimensional fanout. Spatial
+    constraints apply to the data exchange across these instances. Spatial fanouts
+    specified at this level also apply to lower-level `Leaf` nodes in the architecture.
+    """
+
+    def get_fanout(self) -> int:
+        """The spatial fanout of this node."""
+        return int(math.prod(x.fanout for x in self.spatial))
+
+
+_COMPONENT_MODEL_CACHE: dict[tuple, "Component"] = {}
+
+
+def _set_component_model_cache(key: tuple, value: "Component"):
+    while len(_COMPONENT_MODEL_CACHE) > 1000:
+        _COMPONENT_MODEL_CACHE.popitem(last=False)
+    _COMPONENT_MODEL_CACHE[key] = value
+
+
+class _ExtraAttrs(ParseExtras):
+    def _parse_expressions(self, symbol_table: dict[str, Any], *args, **kwargs):
+        if getattr(self, "_parsed", False):
+            return super()._parse_expressions(symbol_table, *args, **kwargs)
+
+        orig_symbol_table = dict(symbol_table)
+        if "arch_extra_attributes_for_all_component_models" not in orig_symbol_table:
+            raise ParseError(
+                "arch_extra_attributes_for_all_component_models is not set in the symbol "
+                "table. Was parsing called from the architecture on top?"
+            )
+
+        for k, v in orig_symbol_table[
+            "arch_extra_attributes_for_all_component_models"
+        ].items():
+            if getattr(self, k, None) is None:
+                setattr(self, k, v)
+
+        return super()._parse_expressions(symbol_table, *args, **kwargs)
+
+
+@_uninstantiable
+class Component(Leaf):
+    """A component object in the architecture. This is overridden by different
+    component types, such as `Memory` and `Compute`."""
+
+    name: str
+    """ The name of this `Component`. """
+
+    component_class: Optional[str] = None
+    """ The class of this `Component`. Used if an energy or area model needs to be
+    called for this `Component`. """
+
+    component_model: ComponentModel | None = None
+    """ The model to use for this `Component`. If not set, the model will be found with
+    `hwcomponents.get_models()`. If set, the `component_class` will be ignored. """
+
+    component_modeling_log: list[str] = []
+    """ A log of the energy and area calculations for this `Component`. """
+
+    actions: ParsableList[Action]
+    """ The actions that this `Component` can perform. """
+
+    enabled: TryParseTo[bool] = True
+    """ Whether this component is enabled. If the expression resolves to False, then
+    the component is disabled. This is parsed per-pmapping-template, so it is a function
+    of the tensors in the current Einsum. For example, you may say `len(All) >= 3` and
+    the component will only be enabled with Einsums with three or more tensors.
+    """
+
+    area: ParsesTo[int | float | None] = None
+    """
+    The area of a single instance of this component in m^2. If set, area calculations
+    will use this value.
+    """
+
+    total_area: ParsesTo[int | float | None] = None
+    """
+    The total area of all instances of this component in m^2. Do not set this value. It
+    is calculated when the architecture's area is calculated.
+    """
+
+    area_scale: ParsesTo[int | float] = 1
+    """
+    The scale factor for the area of this comxponent. This is used to scale the area of
+    this component. For example, if the area is 1 m^2 and the scale factor is 2, then
+    the area is 2 m^2.
+    """
+
+    leak_power: ParsesTo[int | float | None] = None
+    """
+    The leak power of a single instance of this component in W. If set, leak power
+    calculations will use this value.
+    """
+
+    total_leak_power: ParsesTo[int | float | None] = None
+    """
+    The total leak power of all instances of this component in W. Do not set this value.
+    It is calculated when the architecture's leak power is calculated. If instances are
+    power gated, actual leak power may be less than this value.
+    """
+
+    leak_power_scale: ParsesTo[int | float] = 1
+    """
+    The scale factor for the leak power of this component. This is used to scale the
+    leak power of this component. For example, if the leak power is 1 W and the scale
+    factor is 2, then the leak power is 2 W.
+    """
+
+    energy_scale: ParsesTo[int | float] = 1
+    """
+    The scale factor for dynamic energy of this component. For each action, multiplies
+    this action's energy. Multiplies the calculated energy of each action.
+    """
+
+    total_latency: str | int | float = "sum(*action2latency.values())"
+    """
+    An expression representing the total latency of this component in seconds. This is
+    used to calculate the latency of a given Einsum. Special variables available are the
+    following:
+
+    - `min`: The minimum value of all arguments to the expression.
+    - `max`: The maximum value of all arguments to the expression.
+    - `sum`: The sum of all arguments to the expression.
+    - `X_actions`: The number of times action `X` is performed. For example,
+      `read_actions` is the number of times the read action is performed.
+    - `X_latency`: The total latency of all actions of type `X`. For example,
+      `read_latency` is the total latency of all read actions. It is equal to the
+      per-read latency multiplied by the number of read actions.
+    - `action2latency`: A dictionary of action names to their latency.
+
+    Additionally, all component attributes are availble as variables, and all other
+    functions generally available in parsing. Note this expression is parsed after other
+    component attributes are parsed.
+
+    For example, the following expression calculates latency assuming that each read or
+    write action takes 1ns: ``1e-9 * (read_actions + write_actions)``.
+    """
+
+    latency_scale: ParsesTo[int | float] = 1
+    """
+    The scale factor for the latency of this component. This is used to scale the
+    latency of this component. For example, if the latency is 1 ns and the scale factor
+    is 2, then the latency is 2 ns. Multiplies the calculated latency of each action.
+    """
+
+    n_parallel_instances: ParsesTo[int | float] = 1
+    """
+    The number of parallel instances of this component. Increasing parallel instances
+    will proportionally increase area and leakage, while reducing latency (unless
+    latency calculation is overridden).
+    """
+
+    extra_attributes_for_component_model: _ExtraAttrs = _ExtraAttrs()
+    """ Extra attributes to pass to the component model. In addition to all attributes
+    of this component, any extra attributes will be passed to the component model. This
+    can be used to define attributes that are known to the component model, but not
+    accelforge, such as the technology node."""
+
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    def _update_actions(self, new_actions: ParsableList[Action]):
+        has_actions = set(x.name for x in self.actions)
+        for action in new_actions:
+            if action.name not in has_actions:
+                self.actions.append(action)
+
+    def get_component_class(self, trying_to_calculate: str = None) -> str:
+        """Returns the class of this `Component`.
+
+        Parameters
+        ----------
+        trying_to_parse : str, optional
+            What was trying to be calculated using this component. If provided, the
+            error message will be more specific.
+
+        :raises ParseError: If the `component_class` is not set.
+        """
+        extra_info = ""
+        if trying_to_calculate is not None:
+            extra_info = f" Occurred while trying to calculate {trying_to_calculate}."
+
+        if self.component_class is None:
+            raise ParseError(
+                f"component_class must be set to a valid string. "
+                f"Got {self.component_class}. This occurred because the model tried to "
+                "talk to hwcomponents, but was missing necessary attributes. If you do "
+                "not want to use hwcomponents component_models, ensure that area and "
+                "leak_power are set, as well as, for each action, "
+                f"energy and latency are set.{extra_info}",
+                source_field=f"{self.name}.component_class",
+            )
+        return self.component_class
+
+    def _is_dummy(self) -> bool:
+        return (
+            self.component_model is None
+            and self.component_class is not None
+            and self.component_class.lower() == "dummy"
+        )
+
+    def populate_component_model(
+        self: T,
+        component_models: list[ComponentModel] | None = None,
+        in_place: bool = False,
+        trying_to_calculate: str = None,
+    ) -> T:
+        """
+        Populates the ``component_model`` attribute with the model for this component.
+        Extends the ``component_modeling_log`` field with log messages. Uses the
+        ``component_class`` attribute to find the model and populate the
+        ``component_model`` attribute. Uses the ``hwcomponents.get_model()`` function to
+        find the model.
+
+        Parameters
+        ----------
+        component_models : list[ComponentModel] | None
+            The models to use for energy calculation. If not provided, the models will
+            be found with `hwcomponents.get_models()`.
+        in_place : bool
+            If True, the component will be modified in place. Otherwise, a copy will be
+            returned.
+        trying_to_calculate : str, optional
+            What was trying to be calculated using this component. If provided, the
+            error messages for missing component_class will be more specific.
+
+        Returns
+        -------
+        T
+            A copy of the component with the populated ``component_model`` attribute.
+        """
+        if not in_place:
+            self: Self = self._copy_for_component_modeling()
+
+        if self._is_dummy():
+            return self
+
+        if self.component_model is None:
+            if component_models is None:
+                component_models = get_models()
+            estimation = get_model(
+                self.get_component_class(trying_to_calculate=trying_to_calculate),
+                self._attributes_for_component_model(),
+                required_actions=list(x.name for x in self.actions),
+                models=component_models,
+                _return_estimation_object=True,
+            )
+            self.component_model = estimation.value
+            self.component_modeling_log.extend(estimation.messages)
+        return self
+
+    def calculate_action_energy(
+        self,
+        component_models: list[ComponentModel] | None = None,
+        in_place: bool = False,
+    ) -> Self:
+        """
+        Calculates energy for each action of this component. If energy is set in the
+        action or component (with action taking precedence), that value will be used.
+        Otherwise, the energy will be calculated using hwcomponents. Populates, for each
+        action, the ``<action>.energy`` and field. Extends the
+        ``component_modeling_log`` field with log messages.
+
+        Uses the ``component_model`` attribute, or, if not set, the ``component_class``
+        attribute to find the model and populate the ``component_model`` attribute.
+
+        Note that these methods will be called by the Spec when calculating energy and
+        area. If you call them yourself, note that string expressions may not be parsed
+        because they need the Spec's global scope. If you are sure that all necessary
+        values are present and not a result of an expression, you can call these
+        directly. Otherwise, you can call the
+        ``Spec.calculate_component_area_energy_latency_leak`` and then grab components
+        from the returned ``Spec``.
+
+        Parameters
+        ----------
+        component_models : list[ComponentModel] | None
+            The models to use for energy calculation. If not provided, the models will
+            be found with `hwcomponents.get_models()`.
+        in_place : bool
+            If True, the component will be modified in place. Otherwise, a copy will be
+            returned.
+
+        Returns
+        -------
+        T
+            A copy of the component with the calculated energy.
+        """
+        if not in_place:
+            self: Component = self._copy_for_component_modeling()
+
+        messages = self.component_modeling_log
+
+        for action in self.actions:
+            messages.append(f"Calculating energy for {self.name} action {action.name}.")
+            if action.energy is not None:
+                energy = action.energy
+                messages.append(f"Setting {self.name} energy to {action.energy=}")
+            elif self._is_dummy():
+                energy = 0
+                messages.append('Component is "Dummy". Setting energy to 0.')
+            else:
+                self.populate_component_model(
+                    component_models,
+                    in_place=True,
+                    trying_to_calculate=f"energy for action {action.name}",
+                )
+                energy = self.component_model.try_call_arbitrary_action(
+                    action_name=action.name,
+                    _return_estimation_object=True,
+                    **{
+                        **self._attributes_for_component_model(),
+                        **action._attributes_for_component_model(),
+                    },
+                )
+                messages.extend(energy.messages)
+                energy = energy.value[0]
+            if self.energy_scale != 1:
+                energy *= self.energy_scale
+                messages.append(f"Scaling {self.name} energy by {self.energy_scale=}")
+            if action.energy_scale != 1:
+                energy *= action.energy_scale
+                messages.append(f"Scaling {self.name} energy by {action.energy_scale=}")
+            action.energy = energy
+            if action.energy < 0:
+                logging.warning(
+                    f"Component {self.name} action {action.name} has negative energy: "
+                    f"{action.energy=}"
+                )
+        return self
+
+    def calculate_leak_power(
+        self,
+        component_models: list[ComponentModel] | None = None,
+        in_place: bool = False,
+    ) -> Self:
+        """
+        Calculates the leak power for this component. If leak power is set in the
+        component, that value will be used. Otherwise, the leak power will be calculated
+        using hwcomponents. Populates ``leak_power`` field. Extends the
+        ``component_modeling_log`` field with log messages.
+
+        Uses the ``component_model`` attribute, or, if not set, the ``component_class``
+        attribute to find the model and populate the ``component_model`` attribute.
+
+        Note that these methods will be called by the Spec when calculating energy and
+        area. If you call them yourself, note that string expressions may not be parsed
+        because they need the Spec's global scope. If you are sure that all necessary
+        values are present and not a result of an expression, you can call these
+        directly. Otherwise, you can call the
+        ``Spec.calculate_component_area_energy_latency_leak`` and then grab components
+        from the returned ``Spec``.
+
+        Parameters
+        ----------
+        component_models : list[ComponentModel] | None
+            The models to use for energy calculation. If not provided, the models will
+            be found with `hwcomponents.get_models()`.
+        in_place : bool
+            If True, the component will be modified in place. Otherwise, a copy will be
+            returned.
+
+        Returns
+        -------
+        Self
+            A copy of the component with the calculated energy.
+        """
+        if not in_place:
+            self: Self = self._copy_for_component_modeling()
+
+        messages = self.component_modeling_log
+        if self.leak_power is not None:
+            leak_power = self.leak_power
+            messages.append(f"Using predefined leak power value {self.leak_power=}")
+        elif self._is_dummy():
+            leak_power = 0
+            messages.append("Component is dummy. Setting leak power to 0.")
+        else:
+            self.populate_component_model(
+                component_models,
+                in_place=True,
+                trying_to_calculate="leak_power",
+            )
+            leak_power = self.component_model.leak_power
+        if self.leak_power_scale != 1:
+            leak_power *= self.leak_power_scale
+            messages.append(f"Scaling leak power by {self.leak_power_scale=}")
+        if self.n_parallel_instances != 1:
+            leak_power *= self.n_parallel_instances
+            messages.append(f"Scaling leak power by {self.n_parallel_instances=}")
+        self.leak_power = leak_power
+        if self.leak_power < 0:
+            logging.warning(
+                f"Component {self.name} has negative leak power: {self.leak_power}"
+            )
+        return self
+
+    def calculate_area(
+        self,
+        component_models: list[ComponentModel] | None = None,
+        in_place: bool = False,
+    ) -> Self:
+        """
+        Calculates the area for this component. If area is set in the component, that
+        value will be used. Otherwise, the area will be calculated using the
+        hwcomponents library. Populates ``area`` field. Extends the
+        ``component_modeling_log`` field with log messages.
+
+        Uses the ``component_model`` attribute, or, if not set, the ``component_class``
+        attribute to find the model and populate the ``component_model`` attribute.
+
+        Note that these methods will be called by the Spec when calculating energy and
+        area. If you call them yourself, note that string expressions may not be parsed
+        because they need the Spec's global scope. If you are sure that all necessary
+        values are present and not a result of an expression, you can call these
+        directly. Otherwise, you can call the
+        ``Spec.calculate_component_area_energy_latency_leak`` and then grab components
+        from the returned ``Spec``.
+
+        Parameters
+        ----------
+        component_models : list[ComponentModel] | None
+            The models to use for area calculation. If not provided, the models will be
+            found with `hwcomponents.get_models()`.
+        in_place : bool
+            If True, the component will be modified in place. Otherwise, a copy will be
+            returned.
+
+        Returns
+        -------
+        Self
+            A copy of the component with the calculated area.
+        """
+        if not in_place:
+            self: Self = self._copy_for_component_modeling()
+
+        messages = self.component_modeling_log
+        if self.area is not None:
+            area = self.area
+            messages.append(f"Using predefined area value {self.area=}")
+        elif self._is_dummy():
+            area = 0
+            messages.append("Component is dummy. Setting area to 0.")
+        else:
+            self.populate_component_model(
+                component_models,
+                in_place=True,
+                trying_to_calculate="area",
+            )
+            area = self.component_model.area
+        if self.area_scale != 1:
+            area *= self.area_scale
+            messages.append(f"Scaling area by {self.area_scale=}")
+        if self.n_parallel_instances != 1:
+            area *= self.n_parallel_instances
+            messages.append(f"Scaling area by {self.n_parallel_instances=}")
+        self.area = area
+        if self.area < 0:
+            logging.warning(f"Component {self.name} has negative area: {self.area}")
+        return self
+
+    def calculate_action_latency(
+        self,
+        component_models: list[ComponentModel] | None = None,
+        in_place: bool = False,
+    ) -> Self:
+        """
+        Calculates the latency for each action by this component. Populates the
+        ``<action>.latency`` field. Extends the ``component_modeling_log`` field with
+        log messages.
+
+        Parameters
+        ----------
+        component_models : list[ComponentModel] | None
+            The models to use for latency calculation. If not provided, the models will be
+            found with `hwcomponents.get_models()`.
+        in_place : bool
+            If True, the component will be modified in place. Otherwise, a copy will be
+            returned.
+
+        Returns
+        -------
+        Self
+            A copy of the component with the calculated latency for each action.
+        """
+        if not in_place:
+            self: Self = self._copy_for_component_modeling()
+
+        messages = self.component_modeling_log
+
+        for action in self.actions:
+            messages.append(
+                f"Calculating latency for {self.name} action {action.name}."
+            )
+            if action.latency is not None:
+                latency = action.latency
+                messages.append(f"Setting {self.name} latency to {action.latency=}")
+            elif self._is_dummy():
+                latency = 0
+                messages.append("Component is dummy. Setting latency to 0.")
+            else:
+                self.populate_component_model(
+                    component_models,
+                    in_place=True,
+                    trying_to_calculate=f"latency for action {action.name}",
+                )
+                latency = self.component_model.try_call_arbitrary_action(
+                    action_name=action.name,
+                    _return_estimation_object=True,
+                    **{
+                        **self._attributes_for_component_model(),
+                        **action._attributes_for_component_model(),
+                    },
+                )
+                messages.extend(latency.messages)
+                latency = latency.value[1]
+            if self.latency_scale != 1:
+                latency *= self.latency_scale
+                messages.append(f"Scaling {self.name} latency by {self.latency_scale=}")
+            if action.latency_scale != 1:
+                latency *= action.latency_scale
+                messages.append(
+                    f"Scaling {self.name} latency by {action.latency_scale=}"
+                )
+            if self.n_parallel_instances != 1:
+                latency /= self.n_parallel_instances
+                messages.append(
+                    f"Dividing {self.name} latency by {self.n_parallel_instances=}"
+                )
+            action.latency = latency
+            if action.latency < 0:
+                logging.warning(
+                    f"Component {self.name} action {action.name} has negative latency: "
+                    f"{action.latency}"
+                )
+        return self
+
+    def calculate_area_energy_latency_leak(
+        self,
+        component_models: list[ComponentModel] | None = None,
+        in_place: bool = False,
+        _use_cache: bool = False,
+    ) -> Self:
+        """
+        Calculates the area, energy, latency, and leak power for this component.
+        Populates the ``area``, ``total_area``, ``leak_power``, ``total_leak_power``,
+        ``total_latency``, and ``component_modeling_log`` fields of this component.
+        Additionally, for each action, populates the ``<action>.area``,
+        ``<action>.energy``, ``<action>.latency``, and ``<action>.leak_power`` fields.
+        Extends the ``component_modeling_log`` field with log messages.
+
+        Note that these methods will be called by the Spec when calculating energy and
+        area. If you call them yourself, note that string expressions may not be parsed
+        because they need the Spec's global scope. If you are sure that all necessary
+        values are present and not a result of an expression, you can call these
+        directly. Otherwise, you can call the
+        ``Spec.calculate_component_area_energy_latency_leak`` and then grab components
+        from the returned ``Spec``.
+
+        Parameters
+        ----------
+        component_models : list[ComponentModel] | None
+            The models to use for energy calculation. If not provided, the models will
+            be found with `hwcomponents.get_models()`.
+        in_place : bool
+            If True, the component will be modified in place. Otherwise, a copy will be
+            returned.
+        _use_cache : bool
+            If True, the component model will be cached and reused if the same component
+            class, attributes, and actions are provided. Note that this may return
+            copies of the same object across multiple calls.
+
+        Returns
+        -------
+        Self
+            The component with the calculated energy, area, and leak power.
+        """
+        if not in_place:
+            self: Self = self._copy_for_component_modeling()
+
+        if _use_cache and self.component_model is None:
+            cachekey = (
+                self.component_class,
+                self._attributes_for_component_model(),
+                self.actions,
+            )
+            if cachekey in _COMPONENT_MODEL_CACHE:
+                component = _COMPONENT_MODEL_CACHE[cachekey]
+                self.component_model = component.component_model
+                self.component_modeling_log = component.component_modeling_log
+                self.actions = component.actions
+                return self
+
+        self.calculate_area(component_models, in_place=True)
+        self.calculate_action_energy(component_models, in_place=True)
+        self.calculate_action_latency(component_models, in_place=True)
+        self.calculate_leak_power(component_models, in_place=True)
+        if _use_cache:
+            _set_component_model_cache(cachekey, self)
+        return self
+
+    def _attributes_for_component_model(self) -> dict[str, Any]:
+        return {
+            **self.shallow_model_dump(),
+            **self.extra_attributes_for_component_model.shallow_model_dump(),
+        }
+
+    def _copy_for_component_modeling(self) -> Self:
+        self: Component = self.model_copy()
+        self.extra_attributes_for_component_model = (
+            self.extra_attributes_for_component_model.model_copy()
+        )
+        self.actions = type(self.actions)([a.model_copy() for a in self.actions])
+        for action in self.actions:
+            action.extra_attributes_for_component_model = (
+                action.extra_attributes_for_component_model.model_copy()
+            )
+        return self
+
+
+MEMORY_ACTIONS = ParsableList[TensorHolderAction](
+    [
+        TensorHolderAction(name="read"),
+        TensorHolderAction(name="write"),
+    ]
+)
+
+
+PROCESSING_STAGE_ACTIONS = ParsableList[TensorHolderAction](
+    [
+        TensorHolderAction(name="read"),
+    ]
+)
+
+COMPUTE_ACTIONS = ParsableList(
+    [
+        Action(name="compute"),
+    ]
+)
+
+
+def _parse_tensor2bits(
+    to_parse: dict[str, Any], location: str, symbol_table: dict[str, Any]
+) -> dict[str, Any]:
+    result = {}
+    for key, value in to_parse.items():
+        key_parsed = eval_set_expression(
+            expression=key,
+            symbol_table=symbol_table,
+            expected_space=TensorName,
+            location=f"{location} key {key}",
+        ).instance
+        result[key_parsed] = parse_expression(
+            expression=value,
+            symbol_table=symbol_table,
+            attr_name=key,
+            location=location,
+        )
+
+    all = symbol_table["All"].instance
+    for k in result:
+        all -= k
+
+    if all:
+        raise ParseError(f"Missing bits_per_value_scale for {all}")
+
+    for a, b in itertools.combinations(result.keys(), 2):
+        if a & b:
+            raise ParseError(f"bits_per_value_scale for {a} and {b} overlap")
+
+    return {k2: v for k, v in result.items() for k2 in k}
+
+
+class Tensors(ParsableModel):
+    """
+    Fields that control which tensor(s) are kept in a :py:class:`~.TensorHolder` and in
+    what order their nodes may appear in the mapping.
+    """
+
+    keep: TryParseTo[InvertibleSet[TensorName]] = "<Defaults to Nothing>"
+    """
+    A set expression describing which tensors must be kept in this
+    :class:`accelforge.frontend.arch.TensorHolder`. If this is not defined, then all
+    tensors must be kept. Any tensors that are in ``back`` will also be added to
+    ``keep``.
+    """
+
+    may_keep: TryParseTo[InvertibleSet[TensorName]] = (
+        "<Nothing if keep is defined, else All>"
+    )
+    """
+    A set expression describing which tensors may optionally be kept in this
+    :class:`accelforge.frontend.arch.TensorHolder`. The mapper will explore both keeping
+    and not keeping each of these tensors. If this is not defined, then all tensors may
+    be kept.
+    """
+
+    back: TryParseTo[InvertibleSet[TensorName]] = "Nothing"
+    """
+    A set expression describing which tensors must be backed by this
+    :class:`accelforge.frontend.arch.TensorHolder`. If this is not defined, then no
+    tensors must be backed.
+    """
+
+    tile_shape: ParsableList[Comparison] = []
+    """
+    The tile shape for each rank variable. This is given as a list of
+    :class:`~.Comparison` objects, where each comparison must evaluate to True for a
+    valid mapping.
+    """
+
+    no_refetch_from_above: TryParseTo[InvertibleSet[TensorName]] = "~All"
+    """
+    The tensors that are not allowed to be refetched from above. This is given as a set
+    of :class:`~.TensorName` objects or a set expression that resolves to them. These
+    tensors must be fetched at most one time from above memories, and may not be
+    refetched across any temporal or spatial loop iterations. Tensors may be fetched in
+    pieces (if they do not cause re-fetches of any piece).
+    """
+
+    tensor_order_options: ParsableList[
+        ParsableList[TryParseTo[InvertibleSet[TensorName]]]
+    ] = ParsableList()
+    """
+    Options for the order of tensor storage nodes in the mapping. This is given as a
+    list-of-lists-of-sets. Each list-of-sets is a valid order of tensor storage nodes.
+    Order is given from highest in the mapping to lowest.
+
+    For example, an option could be [input | output, weight], which means that there is
+    no relative ordering required between input and output, but weight must be below
+    both.
+    """
+
+    force_memory_hierarchy_order: bool = True
+    """
+    If set to true, storage nodes for lower-level memories must be placed below storage
+    nodes for higher-level memories. For example, all MainMemory storage nodes must go
+    above all LocalBuffer storage nodes.
+
+    This constraint always applies to same-tensor storage nodes (e.g., MainMemory
+    reusing Output must go above LocalBuffer reusing Output); turning it off will permit
+    things like MainMemory reusing Output going above LocalBuffer reusing Input.
+
+    This is identical to the `force_memory_hierarchy_order` field in the `FFM` class,
+    but only applies to this tensor holder.
+    """
+
+    def _parse_expressions(self, *args, **kwargs):
+        self = copy.copy(self)
+        keep, may_keep = self.keep, self.may_keep
+        if may_keep == "<Nothing if keep is defined, else All>":
+            self.may_keep = "All" if keep == "<Defaults to Nothing>" else "~All"
+        if keep == "<Defaults to Nothing>":
+            self.keep = "Nothing"
+        parsed, symbol_table = super(self.__class__, self)._parse_expressions(
+            *args, **kwargs
+        )
+        if isinstance(parsed.back, InvertibleSet):
+            if isinstance(parsed.keep, InvertibleSet):
+                parsed.keep |= parsed.back
+            if isinstance(parsed.may_keep, InvertibleSet):
+                parsed.may_keep -= parsed.back
+
+        # Assert that there are no intersecting sets
+        for order in parsed.tensor_order_options:
+            for i, s0 in enumerate(order):
+                for j, s1 in enumerate(order):
+                    if i == j:
+                        continue
+                    if not isinstance(s0, InvertibleSet) or not isinstance(
+                        s1, InvertibleSet
+                    ):
+                        continue
+                    if s0 & s1:
+                        raise ValueError(
+                            f"Intersecting entries in tensor_order_options: {s0} and {s1}"
+                        )
+
+        return parsed, symbol_table
+
+    # def _parse_tensor_order_options(
+    #     self, symbol_table: dict[str, Any], location: str
+    # ) -> "Tensors":
+    #     result = type(self)(
+    #         tensor_order_options=[
+    #             [
+    #                 eval_set_expression(x, symbol_table, "tensors", location)
+    #                 for x in order_choice
+    #             ]
+    #             for order_choice in self.tensor_order_options
+    #         ],
+    #     )
+    #     # Assert that there are no intersecting sets
+    #     for order in result.tensor_order_options:
+    #         for i, s0 in enumerate(order):
+    #             for j, s1 in enumerate(order):
+    #                 if i == j:
+    #                     continue
+    #                 if s0 & s1:
+    #                     raise ValueError(
+    #                         f"Intersecting entries in dataflow constraint: {s0} and {s1}"
+    #                     )
+    #     return result
+
+    # def _parse_keep(self, symbol_table: dict[str, Any], location: str) -> "Tensors":
+    #     keep, may_keep = self.keep, self.may_keep
+    #     if may_keep == "<Nothing if keep is defined, else All>":
+    #         may_keep = "All" if keep == "<Defaults to Nothing>" else "~All"
+    #     if keep == "<Defaults to Nothing>":
+    #         keep = "Nothing"
+
+    #     may_keep_first = isinstance(keep, str) and re.findall(r"\bmay_keep\b", keep)
+    #     keep_first = isinstance(may_keep, str) and re.findall(r"\bkeep\b", may_keep)
+    #     if keep_first and may_keep_first:
+    #         raise ValueError(
+    #             f"Keep and may_keep reference each other: " f"{keep} and {may_keep}"
+    #         )
+
+    #     if may_keep_first:
+    #         may_keep = eval_set_expression(may_keep, symbol_table, "tensors", location)
+    #         symbol_table = copy.copy(symbol_table)
+    #         symbol_table["may_keep"] = may_keep
+    #         keep = eval_set_expression(keep, symbol_table, "tensors", location)
+    #         return type(self)(keep=keep, may_keep=may_keep)
+    #     else:
+    #         keep = eval_set_expression(keep, symbol_table, "tensors", location)
+    #         symbol_table = copy.copy(symbol_table)
+    #         symbol_table["keep"] = keep
+    #         may_keep = eval_set_expression(may_keep, symbol_table, "tensors", location)
+    #         return type(self)(keep=keep, may_keep=may_keep)
+
+    # def _parse_non_keep(self, symbol_table: dict[str, Any], location: str) -> "Tensors":
+    #     return type(self)(
+    #         tile_shape=[x._parse(symbol_table, location) for x in self.tile_shape],
+    #         no_refetch_from_above=eval_set_expression(
+    #             self.no_refetch_from_above, symbol_table, "tensors", location
+    #         ),
+    #         force_memory_hierarchy_order=parse_expression(
+    #             self.force_memory_hierarchy_order,
+    #             symbol_table,
+    #             "force_memory_hierarchy_order",
+    #             location,
+    #         ),
+    #     )
+
+
+@_uninstantiable
+class TensorHolder(Component):
+    """
+    A `TensorHolder` is a component that holds tensors. These are usually `Memory`s,
+    but can also be `ProcessingStage`s.
+    """
+
+    actions: ParsableList[TensorHolderAction] = MEMORY_ACTIONS
+    """ The actions that this `TensorHolder` can perform. """
+
+    tensors: Tensors = Tensors()
+    """
+    Fields that control which tensor(s) are kept in this `TensorHolder` and in what
+    order their nodes may appear in the mapping.
+    """
+
+    bits_per_value_scale: ParsesTo[dict] = {"All": 1}
+    """
+    A scaling factor for the bits per value of the tensors in this `TensorHolder`. If
+    this is a dictionary, keys in the dictionary are parsed as expressions and may
+    reference one or more tensors.
+    """
+
+    bits_per_action: ParsesTo[int | float | None] = None
+    """
+    The number of bits accessed in each of this component's actions. Overridden by
+    bits_per_action in any action of this component. If set here, acts as a default
+    value for the bits_per_action of all actions of this component.
+    """
+
+    def model_post_init(self, __context__=None) -> None:
+        self._update_actions(MEMORY_ACTIONS)
+
+    def _parse_expressions(self, *args, **kwargs):
+        # Sometimes the same component object may appear in the mapping and the
+        # architecture, in which case we don't want parsing to happen twice.
+        if getattr(self, "_parsed", False):
+            return super()._parse_expressions(*args, **kwargs)
+
+        class MyPostCall(_PostCall):
+            def __call__(self, field, value, parsed, symbol_table):
+                if field == "bits_per_value_scale":
+                    parsed = _parse_tensor2bits(
+                        parsed,
+                        location="bits_per_value_scale",
+                        symbol_table=symbol_table,
+                    )
+                return parsed
+
+        return super()._parse_expressions(*args, **kwargs, post_calls=(MyPostCall(),))
+
+
+class Fanout(Leaf):
+    """
+    Creates a spatial fanout, and doesn't do anything else.
+    """
+
+
+class Memory(TensorHolder):
+    """A `Memory` is a `TensorHolder` that stores data over time, allowing for temporal
+    reuse."""
+
+    size: ParsesTo[int | float]
+    """ The size of this `Memory` in bits. """
+
+    actions: ParsableList[TensorHolderAction] = MEMORY_ACTIONS
+    """ The actions that this `Memory` can perform. """
+
+
+class ProcessingStage(TensorHolder):
+    """A `ProcessingStage` is a `TensorHolder` that does not store data over time, and
+    therefore does not allow for temporal reuse. Use this as a toll that charges reads
+    and writes every time a piece of data moves through it.
+
+    Every write to a `ProcessingStage` is immediately written to the next `Memory`
+    (which may be above or below depending on where the write came from), and same for
+    reads.
+
+    The access counts of a `ProcessingStage` are only included in the "read" action.
+    Each traversal through the `ProcessingStage` is counted as a read. Writes are always
+    zero.
+    """
+
+    direction: Literal["up", "down", "up_and_down"]
+    """
+    The direction in which data flows through this `ProcessingStage`. If "up", then data
+    flows from below `TensorHolder`, through this `ProcessingStage` (plus paying
+    associated costs), and then to the next `TensorHolder` above it. Other data
+    movements are assumed to avoid this ProcessingStage.
+    """
+
+    actions: ParsableList[TensorHolderAction] = PROCESSING_STAGE_ACTIONS
+    """ The actions that this `ProcessingStage` can perform. """
+
+    def model_post_init(self, __context__=None) -> None:
+        self._update_actions(PROCESSING_STAGE_ACTIONS)
+
+
+class Compute(Component):
+    actions: ParsableList[Action] = COMPUTE_ACTIONS
+    """ The actions that this `Compute` can perform. """
+
+    def model_post_init(self, __context__=None) -> None:
+        self._update_actions(COMPUTE_ACTIONS)
+
+
+T = TypeVar("T")
+
+
+@_uninstantiable
+class Branch(ArchNode):
+    # nodes: ArchNodes[_InferFromTag[Compute, Memory, "Hierarchical"]] = ArchNodes()
+    nodes: ArchNodes[
+        Annotated[
+            Union[
+                Annotated[Compute, Tag("Compute")],
+                Annotated[Memory, Tag("Memory")],
+                Annotated[ProcessingStage, Tag("ProcessingStage")],
+                Annotated[Fanout, Tag("Fanout")],
+                Annotated["Parallel", Tag("Parallel")],
+                Annotated["Hierarchical", Tag("Hierarchical")],
+            ],
+            Discriminator(_get_tag),
+        ]
+    ] = ArchNodes()
+
+    def get_nodes_of_type(self, types: Type[T] | tuple[Type[T], ...]) -> Iterator[T]:
+        for node in self.nodes:
+            if isinstance(node, types):
+                yield node
+            elif isinstance(node, Branch):
+                yield from node.get_nodes_of_type(types)
+
+
+class Parallel(Branch):
+    def _flatten(
+        self,
+        compute_node: str,
+        fanout: int = 1,
+        return_fanout: bool = False,
+    ):
+        nodes = []
+
+        for node in self.nodes:
+            if isinstance(node, Compute) and node.name == compute_node:
+                fanout *= node.get_fanout()
+                nodes.append(node)
+                break
+            if isinstance(node, Branch):
+                computes = node.get_nodes_of_type(Compute)
+                if compute_node in [c.name for c in computes]:
+                    new_nodes, new_fanout = node._flatten(
+                        compute_node, fanout, return_fanout=True
+                    )
+                    nodes.extend(new_nodes)
+                    fanout *= new_fanout
+                    break
+        else:
+            raise ParseError(f"Compute node {compute_node} not found in parallel node")
+
+        return nodes, fanout if return_fanout else nodes
+
+
+class Hierarchical(Branch):
+    def _flatten(
+        self,
+        compute_node: str,
+        fanout: int = 1,
+        return_fanout: bool = False,
+    ):
+        nodes = []
+
+        for i, node in enumerate(self.nodes):
+            try:
+                if isinstance(node, (Hierarchical, Parallel)):
+                    if isinstance(node, Parallel) and i < len(self.nodes) - 1:
+                        raise ParseError(
+                            f"Parallel node {node.name} must be the last node in a "
+                            "hierarchical node"
+                        )
+                    new_nodes, new_fanout = node._flatten(
+                        compute_node, fanout, return_fanout=True
+                    )
+                    nodes.extend(new_nodes)
+                    fanout *= new_fanout
+                    if any(
+                        isinstance(n, Compute) and n.name == compute_node
+                        for n in new_nodes
+                    ):
+                        break
+                elif isinstance(node, Compute):
+                    if node.name == compute_node:
+                        fanout *= node.get_fanout()
+                        nodes.append(node)
+                        break
+                elif isinstance(node, Leaf):
+                    fanout *= node.get_fanout()
+                    nodes.append(node)
+                else:
+                    raise TypeError(f"Can't flatten {node}")
+            except ParseError as e:
+                e.add_field(node)
+                raise e
+
+        if return_fanout:
+            return nodes, fanout
+        return nodes
+
+    def render(self) -> str:
+        graph = _pydot_graph()
+        graph.add_node(pydot.Node("root", shape="box", label="TODO: Arch Render"))
+        return _SVGJupyterRender(graph.create_svg(prog="dot").decode("utf-8"))
+
+    def _repr_svg_(self) -> str:
+        return self.render()
+
+
+class _ConstraintLambda:
+    def __init__(
+        self,
+        constraint: Comparison,
+        target_mapping_nodes: list[Spatial],
+        rank_variables: set[str],
+    ):
+        self.constraint = constraint
+        self.constraint_lambda = (
+            None if constraint is None else constraint._to_constraint_lambda(True)
+        )
+        self.target_mapping_nodes = target_mapping_nodes
+        self.rank_variables = rank_variables
+        self._target_node_indices = None
+        self._target_loop_indices = None
+
+    def __repr__(self):
+        return f"_ConstraintLambda({self.constraint}, {self.target_mapping_nodes}, {self.rank_variables})"
+
+    def __call__(self, rank_variables: set[RankVariable], sizes: np.ndarray) -> bool:
+        final = self.rank_variables.issubset(rank_variables)
+        return self.constraint_lambda(final, sizes)
+
+    def _constrained_node_str(self) -> str:
+        return f"constrains {self._target_node_indices}"
+
+    def __bool__(self) -> bool:
+        return bool(self.target_mapping_nodes)
+
+    def __str__(self) -> str:
+        return self.constraint.__str__()
+
+
+class _TileShapeConstraintLambda(_ConstraintLambda):
+    def __str__(self) -> str:
+        if self.constraint._str_repr is not None:
+            return self.constraint._str_repr
+        return "tile_shape " + super().__str__()
+
+    def pretty_str(self) -> str:
+        return f"Tile shape {self.constraint.operator} {self.constraint.value} {self._constrained_node_str()}"
+
+
+class _LoopBoundsConstraintLambda(_ConstraintLambda):
+    def __str__(self) -> str:
+        if self.constraint._str_repr is not None:
+            return self.constraint._str_repr
+        return "loop_bounds " + super().__str__()
+
+    def pretty_str(self) -> str:
+        return f"Loop bounds {self.constraint.operator} {self.constraint.value} {self._constrained_node_str()}"
+
+
+class _MinUsageConstraintLambda(_ConstraintLambda):
+    def __init__(
+        self,
+        target_mapping_nodes: list[Spatial],
+        rank_variables: set[str],
+        min_usage: float,
+    ):
+        super().__init__(None, target_mapping_nodes, rank_variables)
+        self.min_usage = min_usage
+
+    def __call__(self, complete_indices: list[int], usages: np.ndarray) -> bool:
+        # final = self.rank_variables.issubset(rank_variables)
+        final = set(self._target_loop_indices).issubset(set(complete_indices))
+        if not final:
+            return np.ones(usages.shape[0], dtype=np.bool)
+
+        # Some usages are already above the minimum. Return those.
+        result = usages >= self.min_usage
+        if np.sum(result) > 0:
+            return result
+
+        # Nobody is amove the minimum. Return the best we can do.
+        max_usage = np.max(usages, axis=0)
+        return usages == max_usage
+
+    def pretty_str(self) -> str:
+        return f"Min usage {self.min_usage} {self._constrained_node_str()}"
+
+
+class Arch(Hierarchical):
+    """Top-level architecture specification."""
+
+    arch_globals_dependent_on_workload: ParseExtras = ParseExtras()
+    """
+    Attributes that are dependent on the workload. This is parsed first in the
+    architecture, and symbols here are available to the rest of the architecture.
+    """
+
+    extra_attributes_for_all_component_models: ParseExtras = ParseExtras()
+    """
+    Extra attributes to pass to all component models. This can be used to pass global
+    attributes, such as technology node or clock period, to every component model.
+    """
+
+    @property
+    def total_area(self) -> float:
+        """
+        Returns the total area of the architecture in m^2.
+
+        Returns
+        -------
+        float
+            The total area of the architecture in m^2.
+        """
+        return sum(self.per_component_total_area.values())
+
+    @property
+    def total_leak_power(self) -> float:
+        """
+        Returns the total leak power of the architecture in W.
+
+        Returns
+        -------
+        float
+            The total leak power of the architecture in W.
+        """
+        return sum(self.per_component_total_leak_power.values())
+
+    @property
+    def per_component_total_area(self) -> dict[str, float]:
+        """
+        Returns the total area used by each component in the architecture in m^2.
+
+        Returns
+        -------
+        dict[str, float]
+            A dictionary of component names to their total area in m^2.
+        """
+        area = {
+            node.name: node.total_area for node in self.get_nodes_of_type(Component)
+        }
+        for k, v in area.items():
+            if v is None:
+                raise ValueError(
+                    f"Area of {k} is not set. Please call the Spec's "
+                    "`calculate_component_area_energy_latency_leak` method before accessing this "
+                    "property."
+                )
+        return area
+
+    @property
+    def per_component_total_leak_power(self) -> dict[str, float]:
+        """
+        Returns the total leak power of each component in the architecture in W.
+
+        Returns
+        -------
+        dict[str, float]
+            A dictionary of component names to their total leak power in W.
+        """
+        leak_power = {
+            node.name: node.total_leak_power
+            for node in self.get_nodes_of_type(Component)
+        }
+        for k, v in leak_power.items():
+            if v is None:
+                raise ValueError(
+                    f"Leak power of {k} is not set. Please call the Spec's "
+                    "`calculate_component_area_energy_latency_leak` method before accessing this "
+                    "property."
+                )
+        return leak_power
+
+    def _parse_expressions(self, symbol_table: dict[str, Any], *args, **kwargs):
+        outer_st = symbol_table
+
+        class PostCallArch(_PostCall):
+            def __call__(self, field, value, parsed, symbol_table):
+                if field == "arch_globals_dependent_on_workload":
+                    parsed_dump = parsed.shallow_model_dump()
+                    symbol_table.update(parsed_dump)
+                    symbol_table["arch_globals_dependent_on_workload"] = parsed_dump
+                if field == "extra_attributes_for_all_component_models":
+                    parsed_dump = parsed.shallow_model_dump()
+                    symbol_table["arch_extra_attributes_for_all_component_models"] = (
+                        parsed_dump
+                    )
+                return parsed
+
+        cur_st = dict(symbol_table)
+
+        for node in self.get_nodes_of_type(Leaf):
+            cur_st[node.name] = node
+
+        parsed, _ = super()._parse_expressions(
+            cur_st,
+            *args,
+            **kwargs,
+            post_calls=(PostCallArch(),),
+            order=(
+                "arch_globals_dependent_on_workload",
+                "extra_attributes_for_all_component_models",
+            ),
+        )
+        return parsed, symbol_table
+
+    def __getitem__(self, name: str) -> Leaf:
+        return self.name2leaf(name)
+
+    def model_post_init(self, __context__=None) -> None:
+        # Make sure all leaf names are unique
+        leaves = {}
+        for l in self.get_nodes_of_type(Leaf):
+            n = l.name
+            leaves.setdefault(n, l)
+            assert l is leaves[n], f"Duplicate name {n} found in architecture"
+
+
+# We had to reference Hierarchical before it was defined
+Branch.model_rebuild()