accelforge 0.0.1__py3-none-any.whl

accelforge/mapper/FFM/pmappings.py

@@ -0,0 +1,310 @@
+import copy
+from typing import Callable
+from uuid import UUID
+from accelforge.frontend import arch
+from accelforge.frontend.spec import Spec
+from accelforge.mapper.FFM._join_pmappings.pmapping_group import PmappingGroup
+from accelforge.frontend.workload import EinsumName
+from accelforge.frontend.mapping import Mapping
+from accelforge.mapper.FFM._make_pmappings.pmapper_job import Job
+
+
+class MultiEinsumPmappings:
+    """
+    A collection of pmappings for each Einsum in a workload, generated by
+    :func:`~accelforge.mapper.FFM.make_pmappings`.
+
+    Parameters
+    ----------
+    einsum2pmappings:
+        A dictionary of Einsum names to lists of PmappingGroups. PmappingGroups contain
+        the Pareto-optimal pmappings for the Einsum.
+    pmapping_objects:
+        A dictionary of Einsum names to dictionaries of UUIDs to Mappings. The entries
+        in the PmappingGroup objects reference these pmapping objects.
+    einsum2jobs:
+        A dictionary of Einsum names to lists of Jobs that generated the pmappings.
+    can_combine_multiple_runs:
+        Whether the pmappings can be combined from multiple runs. If this is True, then
+        multiple make_pmappings calls can be summed together to get a single
+        MultiEinsumPmappings object. If this is True, the mapper may run more slowly.
+    einsums_with_pmappings_generated:
+        Einsums for which pmappings were generated (or attempted to be generated).
+    flattened_arches:
+        A dictionary of (EinsumName, Compute Name) to lists of architecture nodes. These
+        contain the parsed and flattened architecture node for that particular Einsum
+        and compute combination.
+    parsed_specs:
+        A dictionary of Einsum names to parsed specifications. These contain the parsed
+        specification for that particular Einsum.
+    """
+
+    def __init__(
+        self,
+        einsum2pmappings: dict[EinsumName, list[PmappingGroup]],
+        pmapping_objects: dict[EinsumName, dict[UUID, Mapping]],
+        einsum2jobs: dict[EinsumName, list[Job]],
+        can_combine_multiple_runs: bool,
+        einsums_with_pmappings_generated: set[EinsumName],
+        flattened_arches: dict[EinsumName, list[arch.Leaf]],
+        parsed_specs: dict[EinsumName, Spec],
+    ):
+        self.einsum2pmappings: dict[EinsumName, list[PmappingGroup]] = einsum2pmappings
+        self.pmapping_objects: dict[EinsumName, dict[UUID, Mapping]] = pmapping_objects
+        self.einsum2jobs: dict[EinsumName, list[Job]] = einsum2jobs
+        self.can_combine_multiple_runs: bool = can_combine_multiple_runs
+        self.einsums_with_pmappings_generated: set[EinsumName] = (
+            einsums_with_pmappings_generated
+        )
+        self.flattened_arches: dict[(EinsumName, str), list[arch.Leaf]] = (
+            flattened_arches
+        )
+        self.parsed_specs: dict[EinsumName, Spec] = parsed_specs
+
+    def __or__(self, other: "MultiEinsumPmappings"):
+        if not self.can_combine_multiple_runs or not other.can_combine_multiple_runs:
+            raise ValueError(
+                "Must call make_pmappings with can_combine_multiple_runs=True to combine pmappings "
+                "from multiple runs."
+            )
+        self = copy.copy(self)
+        for einsum_name, pmappings in other.einsum2pmappings.items():
+            self.einsum2pmappings.setdefault(einsum_name, []).extend(pmappings)
+        for einsum_name, jobs in other.einsum2jobs.items():
+            self.einsum2jobs.setdefault(einsum_name, []).extend(jobs)
+        self.pmapping_objects.update(other.pmapping_objects)
+        self.einsums_with_pmappings_generated.update(
+            other.einsums_with_pmappings_generated
+        )
+        self.parsed_specs.update(other.parsed_specs)
+        self.flattened_arches.update(other.flattened_arches)
+        return self
+
+    def _filter(
+        self,
+        filter_lambda: Callable[[PmappingGroup], bool],
+        einsums_with_pmappings_generated: list[EinsumName] | None = None,
+    ):
+        new_einsum2pmappings = {}
+        if einsums_with_pmappings_generated is None:
+            einsums_with_pmappings_generated = list(self.einsum2pmappings.keys())
+        for einsum_name in einsums_with_pmappings_generated:
+            new_einsum2pmappings[einsum_name] = [
+                pm for pm in self.einsum2pmappings[einsum_name] if filter_lambda(pm)
+            ]
+
+        return MultiEinsumPmappings(
+            einsum2pmappings=new_einsum2pmappings,
+            pmapping_objects=self.pmapping_objects,
+            einsum2jobs=self.einsum2jobs,
+            can_combine_multiple_runs=self.can_combine_multiple_runs,
+            einsums_with_pmappings_generated=self.einsums_with_pmappings_generated,
+            flattened_arches=self.flattened_arches,
+            parsed_specs=self.parsed_specs,
+        )
+
+    def drop_einsums(self, *einsums_with_pmappings_generated: EinsumName):
+        """
+        Removes all pmappings for the given Einsums.
+
+        Parameters
+        ----------
+        einsums_with_pmappings_generated:
+            The Einsums for which to remove pmappings.
+        """
+        for einsum_name in einsums_with_pmappings_generated:
+            del self.einsum2pmappings[einsum_name]
+            del self.pmapping_objects[einsum_name]
+            del self.einsum2jobs[einsum_name]
+            self.einsums_with_pmappings_generated.remove(einsum_name)
+
+    def pmapping_keep_rates(
+        self, per_einsum: bool = False
+    ) -> dict[EinsumName, dict[str, float]] | dict[str, float]:
+        """
+        Returns the keep rates for each cause of pmapping removal. For example, if only
+        25% of the pmappings have a valid spatial fanout, then the keep rate for the
+        spatial fanout cause will be 0.25.
+
+        Parameters
+        ----------
+        per_einsum:
+            If True, returns a dictionary of keep rates for each Einsum.
+
+        Returns
+        -------
+        A dictionary of keep rates for each cause of pmapping removal.
+        """
+        result = {}
+        einsum2npmappings = self.n_total_pmappings(per_einsum=True)
+
+        for einsum_name, jobs in self.einsum2jobs.items():
+            cur_result = result.setdefault(einsum_name, {})
+            for job in jobs:
+                for cause, keep_rate in job.pmapping_keep_rates.items():
+                    cur_result.setdefault(cause, 0)
+                    cur_result[cause] += job.n_total_pmappings * keep_rate
+
+        if per_einsum:
+            for einsum_name, npmappings in einsum2npmappings.items():
+                for cause, keep_rate in result[einsum_name].items():
+                    result[einsum_name][cause] = keep_rate / npmappings
+        else:
+            new_result = {}
+            n_total_pmappings = sum(einsum2npmappings.values())
+            for einsum_name, keep_rates in result.items():
+                for cause, keep_rate in keep_rates.items():
+                    new_result.setdefault(cause, 0)
+                    new_result[cause] += keep_rate / n_total_pmappings
+            result = new_result
+
+        return result
+
+    def n_total_pmappings(
+        self, per_einsum: bool = False
+    ) -> int | dict[EinsumName, int]:
+        """
+        Returns the number of total pmappings in the mapspace.
+
+        Parameters
+        ----------
+        per_einsum:
+            If True, returns a dictionary of total pmappings for each Einsum.
+
+        Returns
+        -------
+        The number of total pmappings in the mapspace.
+        """
+        result = {
+            einsum_name: sum(job.n_total_pmappings for job in jobs)
+            for einsum_name, jobs in self.einsum2jobs.items()
+        }
+        if per_einsum:
+            return result
+        return sum(result.values())
+
+    def n_valid_pmappings(
+        self, per_einsum: bool = False
+    ) -> int | dict[EinsumName, int]:
+        """
+        Returns the number of valid pmappings for each Einsum. A valid pmapping is one
+        that satisfies all constraints and resource usage limits.
+
+        Parameters
+        ----------
+        per_einsum:
+            If True, returns a dictionary of valid pmappings for each Einsum.
+
+        Returns
+        -------
+        The number of valid pmappings in the mapspace.
+        """
+        result = {
+            einsum_name: sum(job.n_valid_pmappings for job in jobs)
+            for einsum_name, jobs in self.einsum2jobs.items()
+        }
+        if per_einsum:
+            return result
+        return sum(result.values())
+
+    def n_pareto_optimal_pmappings(
+        self, per_einsum: bool = False
+    ) -> int | dict[EinsumName, int]:
+        """
+        Returns the number of Pareto-optimal pmappings for each Einsum. This is the
+        number of mappings that will be returned by the make_pmappings function.
+
+        Parameters
+        ----------
+        per_einsum:
+            If True, returns a dictionary of Pareto-optimal pmappings for each Einsum.
+
+        Returns
+        -------
+        The number of Pareto-optimal pmappings in the mapspace.
+        """
+        result = {
+            einsum_name: sum(len(p) for p in pmappings)
+            for einsum_name, pmappings in self.einsum2pmappings.items()
+        }
+        if per_einsum:
+            return result
+        return sum(result.values())
+
+    def n_evaluated_pmappings(
+        self, per_einsum: bool = False
+    ) -> int | dict[EinsumName, int]:
+        """
+        Returns the number of pmappings that were evaluated for each Einsum. This is
+        greater than the number of Pareto-optimal pmappings because some mappings are
+        found to be suboptimal after they have been evaluated.
+
+        Parameters
+        ----------
+        per_einsum:
+            If True, returns a dictionary of evaluated pmappings for each Einsum.
+
+        Returns
+        -------
+        The number of evaluated pmappings in the mapspace.
+        """
+
+        result = {
+            einsum_name: sum(job.n_evaluated_pmappings for job in jobs)
+            for einsum_name, jobs in self.einsum2jobs.items()
+        }
+        if per_einsum:
+            return result
+        return sum(result.values())
+
+    def n_pmapping_string(self) -> str:
+        """
+        Returns a string representation of the number of pmappings in the mapspace.
+        Printing this can help diagnose if the mapper is not finding any pmappings or
+        mappings.
+
+        Returns
+        -------
+        A string representation of the number of pmappings in the mapspace.
+        """
+        if "Total" in self.einsum2pmappings:
+            raise ValueError(
+                f"Cannot print stats for a MultiEinsumPmappings object that has "
+                f"an Einsum named 'Total'. Use a different name for the Einsum."
+            )
+
+        n_total_pmappings = self.n_total_pmappings(per_einsum=True)
+        n_valid_pmappings = self.n_valid_pmappings(per_einsum=True)
+        n_evaluated_pmappings = self.n_evaluated_pmappings(per_einsum=True)
+        n_pareto_optimal_pmappings = self.n_pareto_optimal_pmappings(per_einsum=True)
+
+        for x in (
+            n_total_pmappings,
+            n_valid_pmappings,
+            n_evaluated_pmappings,
+            n_pareto_optimal_pmappings,
+        ):
+            x["Total"] = sum(x.values())
+
+        s = []
+        for e in n_total_pmappings:
+            t = n_total_pmappings[e]
+            v = n_valid_pmappings[e]
+            ev = n_evaluated_pmappings[e]
+            p = n_pareto_optimal_pmappings[e]
+
+            def fmt(x, total: bool = True):
+                x = round(x)
+
+                def _f(y):
+                    y = round(y)
+                    return str(y) if y < 1000 else f"{y:.2e}".replace("e+", "e")
+
+                divved = _f(round(t) / x) if x != 0 else "inf"
+                return f"{_f(x)} (1/{divved})" if total else _f(x)
+
+            s.append(
+                f"{e}: {fmt(t, False)} total, {fmt(v)} valid, {fmt(ev)} evaluated, "
+                f"{fmt(p)} Pareto-Optimal"
+            )
+        return "\n".join(s)

accelforge/mapper/__init__.py

@@ -0,0 +1,4 @@
+import accelforge.mapper.FFM as FFM
+from accelforge.frontend.mapper.metrics import Metrics
+
+# import accelforge.mapper._simanneal2 as simanneal2

accelforge/model/__init__.py

@@ -0,0 +1 @@
+from accelforge.model.main import evaluate_mapping

accelforge/model/_looptree/accesses.py

@@ -0,0 +1,335 @@
+from dataclasses import dataclass
+from typing import Optional, overload
+
+# from bindings.looptree import TemporalTag, SequentialTag, PipelineTemporalTag
+
+import islpy as isl
+
+from accelforge.model._looptree.reuse.isl import IslReuseAnalysisOutput
+from accelforge.model._looptree.reuse.symbolic import (
+    BuffetStats,
+    SymbolicAnalysisOutput,
+)
+from accelforge.model._looptree.mapping_utilities import get_paths, get_leaves
+
+from accelforge.frontend.mapping import Mapping, TensorHolder, Compute
+from accelforge.frontend.workload import Workload
+
+# from pytimeloop._isl.singular import get_sum_of_pw_qpolynomial
+# from pytimeloop._isl.sum import sum_with_mask
+
+
+@dataclass(eq=True)
+class Accesses:
+    total_reads: float
+    total_writes: float
+    max_per_unit_reads: float
+    max_per_unit_writes: float
+
+
+class BufferAccesses:
+    def __init__(self):
+        self.accesses: dict[tuple, Accesses] = {}
+
+    def get_accesses(self, buffer, dspace, einsum) -> Accesses:
+        key = (buffer, dspace, einsum)
+        if key not in self.accesses:
+            self.accesses[key] = Accesses(0, 0, 0, 0)
+        return self.accesses[key]
+
+    def items(self):
+        return self.accesses.items()
+
+    def items_with_buffer(self, ref_buffer):
+        """Returns iterator similar to `items` but only for `ref_buffer`"""
+        return (
+            ((buffer, dspace, einsum), value)
+            for (buffer, dspace, einsum), value in self.accesses.items()
+            if buffer == ref_buffer
+        )
+
+    def __str__(self):
+        return repr(self.accesses)
+
+    def __repr__(self):
+        return f"BufferAccesses({repr(self.accesses)})"
+
+
+@overload
+def summarize_total_and_per_unit_actions(
+    reuse_analysis_result: IslReuseAnalysisOutput,
+) -> dict[tuple, BuffetStats]:
+    pass
+
+
+@overload
+def summarize_total_and_per_unit_actions(
+    reuse_analysis_result: SymbolicAnalysisOutput,
+) -> dict[tuple, BuffetStats]:
+    pass
+
+
+def summarize_total_and_per_unit_actions(
+    reuse_analysis_result,
+) -> dict[tuple, BuffetStats]:
+    result = {}
+    if isinstance(reuse_analysis_result, IslReuseAnalysisOutput):
+        reads_to_parent = reuse_analysis_result.reads_to_parent
+        reads_to_peer = reuse_analysis_result.reads_to_peer
+        for key, (tags, fill) in reuse_analysis_result.fills.items():
+            read_to_parent = reads_to_parent[key][1]
+            read_to_peer = reads_to_peer[key][1]
+
+            total_fill = get_sum_of_pw_qpolynomial(fill)
+            total_read_to_parent = get_sum_of_pw_qpolynomial(read_to_parent)
+            total_read_to_peer = get_sum_of_pw_qpolynomial(read_to_peer)
+
+            max_per_unit_fill = _sum_over_temporal_max_over_spatial(tags, fill)
+
+            n_read_to_parent_dim = read_to_parent.dim(isl.dim_type.in_)
+            max_per_unit_read_to_parent = _sum_over_temporal_max_over_spatial(
+                tags[:n_read_to_parent_dim], read_to_parent
+            )
+
+            max_per_unit_read_to_peer = _sum_over_temporal_max_over_spatial(
+                tags, read_to_peer
+            )
+
+            buffet_stats = BuffetStats(
+                total_fills=total_fill,
+                total_reads_to_parent=total_read_to_parent,
+                total_reads_to_peer=total_read_to_peer,
+                max_per_unit_fills=max_per_unit_fill,
+                max_per_parent_reads_to_parent=max_per_unit_read_to_parent,
+                max_per_unit_reads_to_peer=max_per_unit_read_to_peer,
+            )
+
+    elif isinstance(reuse_analysis_result, SymbolicAnalysisOutput):
+        for buffet, buffet_stats in reuse_analysis_result.buffet_stats.items():
+            level = buffet.level
+            einsum = buffet.einsum
+            key = (level, buffet.tensor, einsum)
+            result[key] = buffet_stats
+
+    return result
+
+
+@overload
+def isl_buffer_accesses_from_buffet_actions(
+    reuse_analysis_result: IslReuseAnalysisOutput, mapping, workload, is_path=False
+) -> BufferAccesses:
+    pass
+
+
+@overload
+def isl_buffer_accesses_from_buffet_actions(
+    reuse_analysis_result: SymbolicAnalysisOutput, mapping, workload, is_path=False
+) -> BufferAccesses:
+    pass
+
+
+# TODO: is_path should be removed and we should accept only regular mappings
+def isl_buffer_accesses_from_buffet_actions(
+    reuse_analysis_result, mapping, workload: Workload, is_path=False
+) -> BufferAccesses:
+    mapping = mapping.nodes
+
+    parent_buffers = get_parent_buffers(mapping, workload, is_path)
+
+    compute_targets = set()
+    for compute_node in get_leaves(mapping, is_path):
+        assert isinstance(compute_node, Compute)
+        compute_targets.add(compute_node.component)
+
+    summarized_actions = summarize_total_and_per_unit_actions(reuse_analysis_result)
+
+    accesses_results = BufferAccesses()
+    for (buffer_id, tensor, einsum), stats in summarized_actions.items():
+        fill = stats.total_fills  # Writes
+        read_to_parent = stats.total_reads_to_parent  # Reads to parent
+        read_to_peer = stats.total_reads_to_peer  # Reads to peer
+        max_per_unit_fill = stats.max_per_unit_fills
+        max_per_parent_read_to_parent = stats.max_per_parent_reads_to_parent
+        max_per_unit_read_to_peer = stats.max_per_unit_reads_to_peer
+
+        parent_buffer = parent_buffers[(buffer_id, tensor, einsum)]
+        if parent_buffer is not None:
+            parent_is_backing = parent_buffers[(parent_buffer, tensor, einsum)] is None
+
+            accesses = accesses_results.get_accesses(parent_buffer, tensor, einsum)
+            if tensor in workload.einsums[einsum].output_tensor_names:
+                accesses.total_writes += read_to_parent
+                accesses.total_reads += read_to_parent
+
+                accesses.max_per_unit_reads += max_per_parent_read_to_parent
+                accesses.max_per_unit_writes += max_per_parent_read_to_parent
+
+                if read_to_parent == 0 and max_per_parent_read_to_parent != 0:
+                    raise ValueError(
+                        f"read_to_parent is 0 but max_per_parent_read_to_parent is {max_per_parent_read_to_parent}"
+                    )
+                if read_to_parent == 0 and max_per_unit_read_to_peer == 0:
+                    per_unit_to_total = 0
+                else:
+                    per_unit_to_total = max_per_parent_read_to_parent / read_to_parent
+
+                # TODO: Do this per unit properly by recursing on first iteration in symbolic.py
+                # and passing a flag that says whether this is first iteration
+                if parent_is_backing:
+                    elidable_reads = reuse_analysis_result.elidable_reads.get(tensor, 0)
+                    accesses.total_reads -= elidable_reads
+                    accesses.max_per_unit_reads -= per_unit_to_total * elidable_reads
+
+            elif tensor in workload.einsums[einsum].input_tensor_names:
+                accesses.total_reads += read_to_parent
+
+                accesses.max_per_unit_reads += max_per_parent_read_to_parent
+
+        # Fills will write into current buffer except for compute (which does
+        # not have write action) and top-level buffer
+        if buffer_id not in compute_targets and parent_buffer is not None:
+            accesses = accesses_results.get_accesses(buffer_id, tensor, einsum)
+            if tensor in workload.einsums[einsum].output_tensor_names:
+                accesses.total_writes += fill
+                accesses.max_per_unit_writes += max_per_unit_fill
+
+                # # TODO: figure out how to do this per unit
+                # total_elided_writes = get_tensor_size(workload, tensor)
+                # accesses.total_writes -= total_elided_writes
+            else:
+                accesses.total_writes += fill
+                accesses.max_per_unit_writes += max_per_unit_fill
+
+            accesses.total_reads += read_to_peer
+            accesses.max_per_unit_reads += max_per_unit_read_to_peer
+
+    return accesses_results
+
+
+def get_parent_buffers(mapping: Mapping, workload: Workload, is_path):
+    parent_buffers = {}
+    if is_path:
+        paths = [mapping]
+    else:
+        paths = get_paths(mapping)
+
+    for path in paths:
+        leaf = path[-1]
+        einsum = leaf.einsum
+
+        tensor_to_top_buffer = {}
+        for node in path:
+            if isinstance(node, TensorHolder):
+                for tensor in node.tensors:
+                    key = (node.component, tensor, einsum)
+                    if tensor in tensor_to_top_buffer:
+                        parent_buffers[key] = tensor_to_top_buffer[tensor]
+                    else:
+                        parent_buffers[key] = None
+                    tensor_to_top_buffer[tensor] = node.component
+            elif isinstance(node, Compute):
+                for tensor in workload.einsums[einsum].input_tensor_names:
+                    key = (node.component, tensor, einsum)
+                    if tensor in tensor_to_top_buffer:
+                        parent_buffers[key] = tensor_to_top_buffer[tensor]
+                for tensor in workload.einsums[einsum].output_tensor_names:
+                    key = (node.component, tensor, einsum)
+                    if tensor in tensor_to_top_buffer:
+                        parent_buffers[key] = tensor_to_top_buffer[tensor]
+
+    return parent_buffers
+
+
+def _sum_over_temporal_max_over_spatial(tags, actions):
+    return (
+        sum_with_mask(
+            [
+                (
+                    isinstance(t, TemporalTag)
+                    or isinstance(t, PipelineTemporalTag)
+                    or isinstance(t, SequentialTag)
+                )
+                for t in tags
+            ],
+            actions,
+        )
+        .max()
+        .to_python()
+    )
+
+
+def isl_buffer_accesses_from_buffet_actions(
+    reuse_analysis_result, mapping, workload: Workload, is_path=False
+) -> BufferAccesses:
+    mapping = mapping.nodes
+
+    parent_buffers = get_parent_buffers(mapping, workload, is_path)
+
+    compute_targets = set()
+    for compute_node in get_leaves(mapping, is_path):
+        assert isinstance(compute_node, Compute)
+        compute_targets.add(compute_node.component)
+
+    summarized_actions = summarize_total_and_per_unit_actions(reuse_analysis_result)
+
+    accesses_results = BufferAccesses()
+    for (buffer_id, tensor, einsum), stats in summarized_actions.items():
+        fill = stats.total_fills  # Writes
+        read_to_parent = stats.total_reads_to_parent  # Reads to parent
+        read_to_peer = stats.total_reads_to_peer  # Reads to peer
+        max_per_unit_fill = stats.max_per_unit_fills
+        max_per_parent_read_to_parent = stats.max_per_parent_reads_to_parent
+        max_per_unit_read_to_peer = stats.max_per_unit_reads_to_peer
+
+        parent_buffer = parent_buffers[(buffer_id, tensor, einsum)]
+        if parent_buffer is not None:
+            parent_is_backing = parent_buffers[(parent_buffer, tensor, einsum)] is None
+
+            accesses = accesses_results.get_accesses(parent_buffer, tensor, einsum)
+            if tensor in workload.einsums[einsum].output_tensor_names:
+                accesses.total_writes += read_to_parent
+                accesses.total_reads += read_to_parent
+
+                accesses.max_per_unit_reads += max_per_parent_read_to_parent
+                accesses.max_per_unit_writes += max_per_parent_read_to_parent
+
+                if read_to_parent == 0 and max_per_parent_read_to_parent != 0:
+                    raise ValueError(
+                        f"read_to_parent is 0 but max_per_parent_read_to_parent is {max_per_parent_read_to_parent}"
+                    )
+                if read_to_parent == 0 and max_per_unit_read_to_peer == 0:
+                    per_unit_to_total = 0
+                else:
+                    per_unit_to_total = max_per_parent_read_to_parent / read_to_parent
+
+                # TODO: Do this per unit properly by recursing on first iteration in symbolic.py
+                # and passing a flag that says whether this is first iteration
+                if parent_is_backing:
+                    elidable_reads = reuse_analysis_result.elidable_reads.get(tensor, 0)
+                    accesses.total_reads -= elidable_reads
+                    accesses.max_per_unit_reads -= per_unit_to_total * elidable_reads
+
+            elif tensor in workload.einsums[einsum].input_tensor_names:
+                accesses.total_reads += read_to_parent
+
+                accesses.max_per_unit_reads += max_per_parent_read_to_parent
+
+        # Fills will write into current buffer except for compute (which does
+        # not have write action) and top-level buffer
+        if buffer_id not in compute_targets and parent_buffer is not None:
+            accesses = accesses_results.get_accesses(buffer_id, tensor, einsum)
+            if tensor in workload.einsums[einsum].output_tensor_names:
+                accesses.total_writes += fill
+                accesses.max_per_unit_writes += max_per_unit_fill
+
+                # # TODO: figure out how to do this per unit
+                # total_elided_writes = get_tensor_size(workload, tensor)
+                # accesses.total_writes -= total_elided_writes
+            else:
+                accesses.total_writes += fill
+                accesses.max_per_unit_writes += max_per_unit_fill
+
+            accesses.total_reads += read_to_peer
+            accesses.max_per_unit_reads += max_per_unit_read_to_peer
+
+    return accesses_results

accelforge/model/_looptree/capacity/__init__.py

@@ -0,0 +1 @@
+from .capacity import compute_capacity_usage

accelforge/model/_looptree/capacity/aggregators.py

@@ -0,0 +1,36 @@
+def compute_max(child_caps, caps):
+    buf_to_max = {}
+    for child_cap in child_caps:
+        buf_to_child_usage = {}
+
+        for buf, v in child_cap.items():
+            if buf not in buf_to_child_usage:
+                buf_to_child_usage[buf] = 0
+            if buf not in buf_to_max:
+                buf_to_max[buf] = 0
+
+            buf_to_child_usage[buf] += v
+
+        for buf in buf_to_child_usage:
+            buf_to_max[buf] = max(buf_to_max[buf], buf_to_child_usage[buf])
+
+    for buf, c in buf_to_max.items():
+        if buf not in caps:
+            caps[buf] = 0
+        caps[buf] += c
+
+
+def compute_total(child_caps, caps):
+    for child_cap in child_caps:
+        for buf, v in child_cap.items():
+            if buf in caps:
+                caps[buf] += v
+            else:
+                caps[buf] = v
+
+
+CAPACITY_AGGREGATORS = {
+    "sequential": compute_max,
+    "pipeline": compute_total,
+    "parallel": compute_total,
+}