accelforge 0.0.1__py3-none-any.whl

accelforge/frontend/arch copy.py

@@ -0,0 +1,1544 @@
+import copy
+import math
+from numbers import Number
+import re
+from typing import (
+    Any,
+    Callable,
+    Iterator,
+    List,
+    Literal,
+    Optional,
+    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,
+    _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, _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, *args, **kwargs):
+        class PostCallArchNode(_PostCall):
+            def __call__(self, field, value, parsed, symbol_table):
+                if isinstance(parsed, Container):
+                    symbol_table.update(parsed.attributes)
+                return parsed
+
+        return super()._parse_expressions(
+            *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: str | InvertibleSet[RankVariable] | set[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. """
+
+    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())}"
+        )
+
+
+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: ParsesTo[InvertibleSet[TensorName] | set[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.
+    """
+
+    min_usage: int | float | str = 0.0
+    """ The minimum utilization 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-utilization mappings.
+    """
+
+    reuse: ParsesTo[InvertibleSet[TensorName] | set[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.
+    """
+
+    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.
+    """
+
+    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_name="tensors",
+                    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 LeafAttributes(ParsableModel):
+    pass
+
+
+class AttributesWithExtras(ParseExtras):
+    pass
+
+
+class AttributesWithEnergyLatency(AttributesWithExtras):
+    energy: ParsesTo[int | float | None] = None
+    energy_scale: ParsesTo[int | float] = 1
+    latency_scale: ParsesTo[int | float] = 1
+
+
+class ComponentAttributes(AttributesWithEnergyLatency):
+    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.
+    """
+    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()) / n_parallel_instances"
+    )
+    """
+    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).
+    """
+
+
+class FanoutAttributes(LeafAttributes):
+    model_config = ConfigDict(extra="forbid")
+
+
+class ActionArguments(AttributesWithEnergyLatency):
+    """
+    Arguments for an action of a component.
+    """
+
+    energy: ParsesTo[int | float | None] = None
+    """
+    Dynamic energy of this action. Per-action energy is multiplied by the component's
+    attributes.energy_scale and the action's arguments.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
+    attributes.latency_scale and the action's arguments.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.
+    """
+
+
+class TensorHolderActionArguments(ActionArguments):
+    bits_per_action: ParsesTo[int | float] = (
+        "1 if attributes.bits_per_action is None else attributes.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. """
+
+
+class Action(ParsableModel):
+    name: str
+    """ The name of this action. """
+
+    arguments: ActionArguments = ActionArguments()
+    """
+    The arguments for this action. Passed to the component's model to calculate the
+    energy and latency of the action.
+    """
+
+
+class TensorHolderAction(Action):
+    arguments: TensorHolderActionArguments = TensorHolderActionArguments()
+    """
+    The arguments for this action. Passed to the component's model to calculate the
+    energy and latency of the action.
+    """
+
+
+@_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`. """
+
+    attributes: LeafAttributes = LeafAttributes()
+    """ The attributes 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.
+    """
+
+    _fields_for_energy_area_latency_leak_calculation: tuple[str, ...] = (
+        "name",
+        "attributes",
+    )
+    """
+    The fields that are used to calculate the energy, area, latency, and leak power of
+    this `Leaf`.
+    """
+
+    def _parse_expressions(self, symbol_table: dict[str, Any], *args, **kwargs):
+        class PostCallLeaf(_PostCall):
+            def __call__(self, field, value, parsed, symbol_table):
+                if field == "attributes":
+                    symbol_table.update(parsed.model_dump())
+                return parsed
+
+        if "_parsing_attributes_only_" in symbol_table:
+            kwargs["fields"] = self._fields_for_energy_area_latency_leak_calculation
+
+        parsed, symbol_table = super()._parse_expressions(
+            symbol_table,
+            *args,
+            **kwargs,
+            post_calls=(PostCallLeaf(),),
+            order=("attributes",),
+        )
+        symbol_table[self.name] = self
+        return parsed, symbol_table
+
+    def get_fanout(self) -> int:
+        """The spatial fanout of this node."""
+        return int(math.prod(x.fanout for x in self.spatial))
+
+
+@_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. """
+
+    attributes: ComponentAttributes = ComponentAttributes()
+    """ The attributes of this `Component`. """
+
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    enabled: str | 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.
+    """
+
+    _fields_for_energy_area_latency_leak_calculation: tuple[str, ...] = (
+        "actions",
+        "component_class",
+        "component_model",
+        "component_modeling_log",
+    ) + Leaf._fields_for_energy_area_latency_leak_calculation
+
+    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 models, ensure that attributes.area and "
+                "attributes.leak_power are set, as well as, for each action, "
+                f"arguments.energy and arguments.latency are set.{extra_info}",
+                source_field=f"{self.name}.component_class",
+            )
+        return self.component_class
+
+    def populate_component_model(
+        self: T,
+        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
+        ----------
+        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.model_copy()
+            self.attributes = self.attributes.model_copy()
+            self.actions = type(self.actions)([a.model_copy() for a in self.actions])
+            for action in self.actions:
+                action.arguments = action.arguments.model_copy()
+
+        if self.component_model is None:
+            if models is None:
+                models = get_models()
+            estimation = get_model(
+                self.get_component_class(trying_to_calculate=trying_to_calculate),
+                self.attributes.model_dump(),
+                required_actions=list(x.name for x in self.actions),
+                models=models,
+                _return_estimation_object=True,
+            )
+            self.component_model = estimation.value
+            self.component_modeling_log.extend(estimation.messages)
+        return self
+
+    def calculate_action_energy(
+        self: T,
+        models: list[ComponentModel] | None = None,
+        in_place: bool = False,
+    ) -> T:
+        """
+        Calculates energy for each action of this component. If energy is set in the
+        arguments or attributes (with arguments taking precedence), that value will be
+        used. Otherwise, the energy will be calculated using hwcomponents. Populates,
+        for each action, the ``<action>.arguments.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
+        ----------
+        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 = self.model_copy()
+            self.attributes = self.attributes.model_copy()
+            self.actions = type(self.actions)([a.model_copy() for a in self.actions])
+            for action in self.actions:
+                action.arguments = action.arguments.model_copy()
+
+        messages = self.component_modeling_log
+
+        attributes = self.attributes
+        for action in self.actions:
+            messages.append(f"Calculating energy for {self.name} action {action.name}.")
+            args = action.arguments
+            if args.energy is not None:
+                energy = args.energy
+                messages.append(f"Setting {self.name} energy to {args.energy=}")
+            else:
+                self.populate_component_model(
+                    models,
+                    in_place=True,
+                    trying_to_calculate=f"arguments.energy for action {action.name}",
+                )
+                energy = self.component_model.try_call_arbitrary_action(
+                    action_name=action.name,
+                    _return_estimation_object=True,
+                    **{**attributes.model_dump(), **args.model_dump()},
+                )
+                messages.extend(energy.messages)
+                energy = energy.value[0]
+            if attributes.energy_scale != 1:
+                energy *= attributes.energy_scale
+                messages.append(
+                    f"Scaling {self.name} energy by {attributes.energy_scale=}"
+                )
+            if args.energy_scale != 1:
+                energy *= args.energy_scale
+                messages.append(f"Scaling {self.name} energy by {args.energy_scale=}")
+            action.arguments.energy = energy
+        return self
+
+    def calculate_leak_power(
+        self: T,
+        models: list[ComponentModel] | None = None,
+        in_place: bool = False,
+    ) -> T:
+        """
+        Calculates the leak power for this component. If leak power is set in the
+        arguments or attributes (with arguments taking precedence), that value will be
+        used. Otherwise, the leak power will be calculated using hwcomponents. Populates
+        ``attributes.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
+        ----------
+        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 = self.model_copy()
+            self.attributes = self.attributes.model_copy()
+            self.actions = type(self.actions)([a.model_copy() for a in self.actions])
+            for action in self.actions:
+                action.arguments = action.arguments.model_copy()
+
+        attributes = self.attributes
+        messages = self.component_modeling_log
+        if attributes.leak_power is not None:
+            leak_power = attributes.leak_power
+            messages.append(
+                f"Using predefined leak power value {attributes.leak_power=}"
+            )
+        else:
+            self.populate_component_model(
+                models,
+                in_place=True,
+                trying_to_calculate="attributes.leak_power",
+            )
+            leak_power = self.component_model.leak_power
+        if attributes.leak_power_scale != 1:
+            leak_power *= attributes.leak_power_scale
+            messages.append(f"Scaling leak power by {attributes.leak_power_scale=}")
+        if attributes.n_parallel_instances != 1:
+            leak_power *= attributes.n_parallel_instances
+            messages.append(f"Scaling leak power by {attributes.n_parallel_instances=}")
+        self.attributes.leak_power = leak_power
+        return self
+
+    def calculate_area(
+        self: T,
+        models: list[ComponentModel] | None = None,
+        in_place: bool = False,
+    ) -> T:
+        """
+        Calculates the area for this component. If area is set in the attributes, that
+        value will be used. Otherwise, the area will be calculated using the
+        hwcomponents library. Populates ``attributes.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
+        ----------
+        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
+        -------
+        T
+            A copy of the component with the calculated area.
+        """
+        if not in_place:
+            self = self.model_copy()
+            self.attributes = self.attributes.model_copy()
+            self.actions = type(self.actions)([a.model_copy() for a in self.actions])
+            for action in self.actions:
+                action.arguments = action.arguments.model_copy()
+
+        attributes = self.attributes
+        messages = self.component_modeling_log
+        if attributes.area is not None:
+            area = attributes.area
+            messages.append(f"Using predefined area value {attributes.area=}")
+        else:
+            self.populate_component_model(
+                models,
+                in_place=True,
+                trying_to_calculate="attributes.area",
+            )
+            area = self.component_model.area
+        if attributes.area_scale != 1:
+            area *= attributes.area_scale
+            messages.append(f"Scaling area by {attributes.area_scale=}")
+        if attributes.n_parallel_instances != 1:
+            area *= attributes.n_parallel_instances
+            messages.append(f"Scaling area by {attributes.n_parallel_instances=}")
+        self.attributes.area = area
+        return self
+
+    def calculate_action_latency(
+        self: T,
+        models: list[ComponentModel] | None = None,
+        in_place: bool = False,
+    ) -> T:
+        """
+        Calculates the latency for each action by this component. Populates the
+        ``<action>.arguments.latency`` field. Extends the ``component_modeling_log`` field with
+        log messages.
+
+        Parameters
+        ----------
+        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
+        -------
+        T
+            A copy of the component with the calculated latency for each action.
+        """
+        if not in_place:
+            self = self.model_copy()
+            self.attributes = self.attributes.model_copy()
+            self.actions = type(self.actions)([a.model_copy() for a in self.actions])
+            for action in self.actions:
+                action.arguments = action.arguments.model_copy()
+
+        messages = self.component_modeling_log
+
+        attributes = self.attributes
+        for action in self.actions:
+            messages.append(
+                f"Calculating latency for {self.name} action {action.name}."
+            )
+            args = action.arguments
+            if args.latency is not None:
+                latency = args.latency
+                messages.append(f"Setting {self.name} latency to {args.latency=}")
+            else:
+                self.populate_component_model(
+                    models,
+                    in_place=True,
+                    trying_to_calculate=f"arguments.latency for action {action.name}",
+                )
+                latency = self.component_model.try_call_arbitrary_action(
+                    action_name=action.name,
+                    _return_estimation_object=True,
+                    **{**attributes.model_dump(), **args.model_dump()},
+                )
+                messages.extend(latency.messages)
+                latency = latency.value[1]
+            if attributes.latency_scale != 1:
+                latency *= attributes.latency_scale
+                messages.append(
+                    f"Scaling {self.name} latency by {attributes.latency_scale=}"
+                )
+            if args.latency_scale != 1:
+                latency *= args.latency_scale
+                messages.append(f"Scaling {self.name} latency by {args.latency_scale=}")
+            if attributes.n_parallel_instances != 1:
+                latency /= attributes.n_parallel_instances
+                messages.append(
+                    f"Dividing {self.name} latency by {attributes.n_parallel_instances=}"
+                )
+            action.arguments.latency = latency
+        return self
+
+    def calculate_area_energy_latency_leak(
+        self: T, models: list[ComponentModel] | None = None, in_place: bool = False
+    ) -> T:
+        """
+        Calculates the area, energy, latency, and leak power for this component.
+        Populates the ``attributes.area``, ``attributes.total_area``,
+        ``attributes.leak_power``, ``attributes.total_leak_power``,
+        ``attributes.total_latency``, and ``component_modeling_log`` fields of this
+        component. Additionally, for each action, populates the
+        ``<action>.arguments.area``, ``<action>.arguments.energy``,
+        ``<action>.arguments.latency``, and ``<action>.arguments.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
+        ----------
+        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
+            The component with the calculated energy, area, and leak power.
+        """
+        if not in_place:
+            self = self.model_copy()
+            self.attributes = self.attributes.model_copy()
+            self.actions = type(self.actions)([a.model_copy() for a in self.actions])
+            for action in self.actions:
+                action.arguments = action.arguments.model_copy()
+        self.calculate_area(models, in_place=True)
+        self.calculate_action_energy(models, in_place=True)
+        self.calculate_action_latency(models, in_place=True)
+        self.calculate_leak_power(models, in_place=True)
+        return self
+
+
+class Container(Leaf):
+    """A `Container` is an abstract node in the architecture that contains other nodes.
+    For example, a P` may be a `Container` that contains `Memory`s and `Compute` units.
+    """
+
+    pass
+
+
+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():
+        if isinstance(value, Number):
+            result[key] = value
+            continue
+        result[key] = parse_expression(
+            expression=value,
+            symbol_table=symbol_table,
+            attr_name=key,
+            location=location,
+        )
+    return result
+
+
+class TensorHolderAttributes(ComponentAttributes):
+    """
+    Attributes for a `TensorHolder`. `TensorHolder`s are components that hold tensors
+    (usually `Memory`s). When specifying these attributes, it is recommended to
+    underscore-prefix attribute names. See `TODO: UNDERSCORE_PREFIX_DISCUSSION`.
+    """
+
+    bits_per_value_scale: ParsesTo[dict | int | float] = {"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 the action arguments. 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:
+        if not isinstance(self.bits_per_value_scale, dict):
+            self.bits_per_value_scale = {"All": self.bits_per_value_scale}
+
+    def _parse_expressions(self, *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 MemoryAttributes(TensorHolderAttributes):
+    """Attributes for a `Memory`."""
+
+    size: ParsesTo[int | float]
+    """ The size of this `Memory` in bits. """
+
+
+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: ParsesTo[InvertibleSet[TensorName] | set[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.
+    """
+
+    may_keep: ParsesTo[InvertibleSet[TensorName] | set[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.
+    """
+
+    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: ParsesTo[InvertibleSet[TensorName] | set[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[ParsesTo[InvertibleSet[TensorName] | set[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_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. """
+
+    attributes: TensorHolderAttributes = pydantic.Field(
+        default_factory=TensorHolderAttributes
+    )
+    """ The `TensorHolderAttributes` that describe this `TensorHolder`. """
+
+    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.
+    """
+
+    def model_post_init(self, __context__=None) -> None:
+        self._update_actions(MEMORY_ACTIONS)
+
+
+class Fanout(Leaf):
+    """
+    Creates a spatial fanout, and doesn't do anything else.
+    """
+
+    attributes: FanoutAttributes = pydantic.Field(default_factory=FanoutAttributes)
+    """ Fanout attributes. Zero energy, leak power, area, and latency. """
+
+
+class Memory(TensorHolder):
+    """A `Memory` is a `TensorHolder` that stores data over time, allowing for temporal
+    reuse."""
+
+    attributes: "MemoryAttributes" = pydantic.Field(default_factory=MemoryAttributes)
+    """ The attributes of this `Memory`. """
+
+    actions: ParsableList[TensorHolderAction] = MEMORY_ACTIONS
+    """ The actions that this `Memory` can perform. """
+
+
+class ProcessingStageAttributes(TensorHolderAttributes):
+    """Attributes for a `ProcessingStage`."""
+
+    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.
+    """
+
+
+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.
+    """
+
+    attributes: ProcessingStageAttributes = pydantic.Field(
+        default_factory=ProcessingStageAttributes
+    )
+    """ The attributes of 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 ComputeAttributes(ComponentAttributes):
+    """Attributes for a `Compute`."""
+
+    pass
+
+
+class Compute(Component):
+    actions: ParsableList[Action] = COMPUTE_ACTIONS
+    """ The actions that this `Compute` can perform. """
+
+    attributes: ComputeAttributes = pydantic.Field(default_factory=ComputeAttributes)
+    """ The attributes of this `Compute`. """
+
+    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,
+        attributes: dict,
+        compute_node: str,
+        fanout: int = 1,
+        return_fanout: bool = False,
+    ):
+        nodes = []
+
+        def _parse_node(node: Leaf, fanout: int):
+            fanout *= node.get_fanout()
+            node2 = node.model_copy()
+            node2.attributes = type(node.attributes)(
+                **{**attributes.model_dump(), **node.attributes.model_dump()}
+            )
+            nodes.append(node2)
+            return fanout
+
+        for node in self.nodes:
+            if isinstance(node, Compute) and node.name == compute_node:
+                fanout = _parse_node(node, fanout)
+                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(
+                        attributes, 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,
+        attributes: dict,
+        compute_node: str,
+        fanout: int = 1,
+        return_fanout: bool = False,
+    ):
+        nodes = []
+
+        def _parse_node(node: Leaf, fanout: int):
+            fanout *= node.get_fanout()
+            node2 = node.model_copy()
+            attrs = {**node.attributes.model_dump()}
+            if isinstance(node.attributes, AttributesWithExtras):
+                attrs = {**attributes.model_dump(), **attrs}
+            node2.attributes = type(node.attributes)(**attrs)
+            nodes.append(node2)
+            return fanout
+
+        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(
+                        attributes, 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 = _parse_node(node, fanout)
+                        break
+                elif isinstance(node, Leaf) and not isinstance(node, Container):
+                    fanout = _parse_node(node, fanout)
+                elif isinstance(node, Container):
+                    fanout *= node.get_fanout()
+                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 __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)
+
+
+class _TileShapeConstraintLambda(_ConstraintLambda):
+    def pretty_str(self) -> str:
+        return f"Tile shape {self.constraint.operator} {self.constraint.value} {self._constrained_node_str()}"
+
+
+class _LoopBoundsConstraintLambda(_ConstraintLambda):
+    def pretty_str(self) -> str:
+        return f"Loop bounds {self.constraint.operator} {self.constraint.value} {self._constrained_node_str()}"
+
+
+class _MinUtilizationConstraintLambda(_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], utilizations: 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(utilizations.shape[0], dtype=np.bool)
+
+        # Some utilizations are already above the minimum. Return those.
+        result = utilizations >= self.min_usage
+        if np.sum(result) > 0:
+            return result
+
+        # Nobody is amove the minimum. Return the best we can do.
+        max_utilization = np.max(utilizations, axis=0)
+        return utilizations == max_utilization
+
+    def pretty_str(self) -> str:
+        return f"Min utilization {self.min_usage} {self._constrained_node_str()}"
+
+
+class Arch(Hierarchical):
+    # version: Annotated[str, assert_version] = __version__
+    # """ The version of the architecture spec. """
+
+    @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.attributes.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.attributes.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, *args, **kwargs):
+        symbol_table = kwargs["symbol_table"]
+        for node in self.get_nodes_of_type(Leaf):
+            symbol_table[node.name] = node
+        return super()._parse_expressions(*args, **kwargs)
+
+    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()