accelforge 0.0.1__py3-none-any.whl

accelforge/model/_looptree/run.py

@@ -0,0 +1,122 @@
+from dataclasses import dataclass
+from pathlib import Path
+
+
+@dataclass
+class LoopTreeStatistics:
+    latency: float
+    energy: float
+    actions: dict
+    memory_latency: dict
+    capacity_usage: dict
+
+
+def run_symbolic_model(mapping, workload, architecture):
+    from pytimeloop.looptree.reuse import analyze_reuse_and_add_reservations_to_mapping
+    from pytimeloop.looptree.energy import gather_actions
+
+    job = Job.make_job(mapping=mapping, workload=workload, architecture=architecture)
+    result = analyze_reuse_and_add_reservations_to_mapping(job)
+    actions = gather_actions(result, bindings, use_name=True)
+    pass
+
+
+def run_looptree(config_dir, paths, tmp_path, bindings, call_accelergy):
+    import islpy as isl
+    from bindings.config import Config
+    from bindings.looptree import LooptreeModelApp, LooptreeWorkload
+    from pytimeloop.file import gather_yaml_configs
+    from pytimeloop.looptree.capacity import compute_capacity_usage
+    from pytimeloop.looptree.reuse._isl.des import deserialize_looptree_output
+    from pytimeloop.looptree.energy import gather_actions, compute_energy_from_actions
+    from pytimeloop.looptree.latency import get_latency
+    from pytimeloop.timeloopfe.v4fused import Spec
+    from pytimeloop.timeloopfe.common.backend_calls import call_accelergy_verbose
+
+    yaml_str = gather_yaml_configs(config_dir, paths)
+    config = Config(yaml_str, "yaml")
+    model = LooptreeModelApp(config)
+
+    workload = LooptreeWorkload.parse_cfg(config.root["problem"])
+
+    spec = Spec.from_yaml_files([str(config_dir / p) for p in paths])
+
+    if call_accelergy:
+        if isinstance(tmp_path, Path):
+            tmp_path = str(tmp_path)
+        call_accelergy_verbose(spec, tmp_path)
+        spec = Spec.from_yaml_files(
+            [str(config_dir / p) for p in paths] + [str(Path(tmp_path) / "ERT.yaml")]
+        )
+
+    result = deserialize_looptree_output(model.run(), isl.DEFAULT_CONTEXT)
+
+    actions = gather_actions(result, bindings)
+    energy = compute_energy_from_actions(actions, spec.ERT)
+
+    latency, comp_latency, mem_latency = get_latency(
+        result, spec.mapping, workload, spec.arch, bindings
+    )
+
+    capacity_usage = compute_capacity_usage(
+        spec.mapping.nodes, result.occupancy, workload
+    )
+    component_capacity_usage = {}
+    for level, component in bindings.items():
+        if level in capacity_usage:
+            component_capacity_usage[component] = capacity_usage[level]
+
+    return LoopTreeStatistics(
+        latency, energy, actions, mem_latency, capacity_usage=component_capacity_usage
+    )
+
+
+def run_looptree_symbolic(config_dir, paths, tmp_path, bindings, call_accelergy):
+    from bindings.config import Config
+    from bindings.looptree import LooptreeWorkload, LooptreeWorkloadDependencyAnalyzer
+    from pytimeloop.file import gather_yaml_configs
+    from pytimeloop.looptree.capacity import compute_capacity_usage
+    from pytimeloop.looptree.reuse import analyze_reuse_and_add_reservations_to_mapping
+    from pytimeloop.looptree.energy import gather_actions, compute_energy_from_actions
+    from pytimeloop.looptree.latency import get_latency
+    from pytimeloop.timeloopfe.v4fused import Spec
+    from pytimeloop.timeloopfe.common.backend_calls import call_accelergy_verbose
+    from accelforge.mapper.FFM._make_pmappings.pmapper_job import Job
+
+    yaml_str = gather_yaml_configs(config_dir, paths)
+
+    config = Config(yaml_str, "yaml")
+    workload = LooptreeWorkload.parse_cfg(config.root["problem"])
+    analyzer = LooptreeWorkloadDependencyAnalyzer(workload)
+
+    spec = Spec.from_yaml_files([str(config_dir / p) for p in paths])
+
+    if call_accelergy:
+        if isinstance(tmp_path, Path):
+            tmp_path = str(tmp_path)
+        call_accelergy_verbose(spec, tmp_path)
+        spec = Spec.from_yaml_files(
+            [str(config_dir / p) for p in paths] + [str(Path(tmp_path) / "ERT.yaml")]
+        )
+
+    job = Job.make_job(mapping=spec.mapping, workload=workload, architecture=spec.arch)
+    tile_shapes, result = analyze_reuse_and_add_reservations_to_mapping(job)
+
+    actions = gather_actions(result, bindings, use_name=True)
+    energy = compute_energy_from_actions(actions, spec.ERT)
+
+    latency, comp_latency, mem_latency = get_latency(
+        result, spec.mapping, workload, spec.arch, bindings
+    )
+
+    capacity_usage = compute_capacity_usage(
+        spec.mapping.nodes, result.occupancy, workload
+    )
+    component_capacity_usage = {}
+    for level, component in bindings.items():
+        if level in capacity_usage:
+            component_capacity_usage[component] = capacity_usage[level]
+
+    return LoopTreeStatistics(
+        latency, energy, actions, mem_latency, capacity_usage=component_capacity_usage
+    )

accelforge/model/_looptree/types.py

@@ -0,0 +1,26 @@
+"""
+Contains shared classes of analysis.
+"""
+
+from dataclasses import dataclass
+from typing import TypeAlias
+
+from accelforge.frontend.mapping import TensorName
+from accelforge.frontend.workload import EinsumName
+
+ComponentName: TypeAlias = str
+
+
+@dataclass(eq=True, frozen=True)
+class Buffet:
+    """
+    A logical buffer that stores a tensor, an einsum operating on it, and the
+    level the buffer exists on in hardware.
+    """
+
+    tensor: TensorName
+    "The tensor held by the buffet."
+    einsum: EinsumName
+    "An einsum operating on the tensor."
+    level: ComponentName
+    "The abstract hardware level the buffet resides in."

accelforge/model/_looptree/visualization/occupancy.py

@@ -0,0 +1,11 @@
+import matplotlib.pyplot as plt
+
+
+def plot_occupancy_graph(output: "IslReuseAnalysisOutput", workload):
+    einsum_rank_to_shape = {
+        einsum: {
+            rank: workload.get_rank_shape(rank)
+            for rank in workload.einsum_ospace_dimensions(einsum)
+        }
+        for einsum in workload.einsum_id_to_name()
+    }

accelforge/model/main.py

@@ -0,0 +1,222 @@
+from copy import copy, deepcopy
+from uuid import uuid4
+
+import pandas as pd
+
+from accelforge.frontend import arch
+from accelforge.frontend.arch import Memory
+from accelforge.frontend.renames import EinsumName
+from accelforge.frontend.spec import Mapping, Spec
+from accelforge.frontend.mapping import Compute, Split, Nested, NodeList, TensorHolder
+from accelforge.frontend.workload import Workload
+from accelforge.frontend._workload_isl._symbolic import (
+    get_stride_and_halo_of_einsum,
+    get_rank_variable_relevancy,
+)
+
+
+def evaluate_mapping(
+    spec: Spec,
+    flattened_arches: dict[(EinsumName, str), list[arch.Leaf]] | None = None,
+    parsed_specs: dict[EinsumName, Spec] | None = None,
+):
+    """
+    Evaluate a mapping.
+
+    Parameters
+    ----------
+    spec:
+        The specification of architecture, workload, and mapping.
+    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. If provided, then these will be used instead of
+        re-parsing the architecture.
+    parsed_specs:
+        A dictionary of Einsum names to parsed specifications. These contain the parsed
+        specification for that particular Einsum. If provided, then these will be used
+        instead of re-parsing the specification.
+    """
+    from accelforge.mapper.FFM._join_pmappings.compatibility import Compatibility
+    from accelforge.mapper.FFM._join_pmappings.pmapping_dataframe import (
+        PmappingDataframe,
+    )
+    from accelforge.mapper.FFM._join_pmappings.pmapping_group import PmappingGroup
+    from accelforge.mapper.FFM._join_pmappings.join_pmappings import (
+        clean_compress_and_join_pmappings,
+    )
+    from accelforge.mapper.FFM.pmappings import MultiEinsumPmappings
+    from accelforge.mapper.FFM._make_pmappings.make_pmappings import (
+        get_rank_variable_bounds_for_all_einsums,
+    )
+    from accelforge.mapper.FFM._make_pmappings.make_pmappings_from_templates.run_model import (
+        run_model,
+    )
+    from accelforge.mapper.FFM._make_pmappings.pmapper_job import Job
+
+    assert (parsed_specs is not None) == (
+        flattened_arches is not None
+    ), f"Provide either flattened_arches or parsed_specs, not both."
+
+    original_job = Job(
+        metrics=spec.model.metrics,
+        rank_variable_bounds=get_rank_variable_bounds_for_all_einsums(spec),
+        spec=spec,
+    )
+
+    einsum2pmappings = {}
+    pmapping_objects = {}
+    einsum2jobs = {}
+    assert not getattr(
+        spec, "_parsed", False
+    ), "Spec must not be parsed before evaluating a mapping"
+    for pmapping in _split_mapping_to_pmappings(spec.mapping, spec.workload):
+        einsum_name = pmapping.nodes[-1].einsum
+        compute_name = pmapping.nodes[-1].component
+        pmapping_id = uuid4()
+        job = copy(original_job)
+
+        if flattened_arches is not None:
+            flattened_arch = flattened_arches[(einsum_name, compute_name)]
+            cur_spec = parsed_specs[einsum_name]
+
+        else:
+            cur_spec = spec.calculate_component_area_energy_latency_leak(
+                einsum_name=einsum_name,
+                area=False,
+            )
+            flattened_arch = cur_spec._get_flattened_architecture(
+                compute_node=pmapping.nodes[-1].component
+            )
+
+        job.spec = cur_spec
+        pmapping.remove_reservations()
+        pmapping.split_loop_with_multiple_rank_variables()
+        pmapping.split_tensor_holders_with_multiple_tensors()
+        _add_backing_to_tensor_holders(pmapping)
+
+        job.mapping = pmapping
+        job.einsum_name = pmapping.nodes[-1].einsum
+        job.tensor_to_relevancy = {
+            tensor: get_rank_variable_relevancy(
+                job.spec.workload.einsums[job.einsum_name], tensor
+            )
+            for tensor in job.spec.workload.einsums[job.einsum_name].tensor_names
+        }
+        einsum2jobs[job.einsum_name] = job
+
+        job.flattened_arch = flattened_arch
+        job.memories_track_all = [
+            m.name for m in flattened_arch if isinstance(m, Memory)
+        ]
+
+        job.stride_and_halo = get_stride_and_halo_of_einsum(
+            job.einsum_name, cur_spec.workload
+        )
+        job.fusable_tensors = set(
+            cur_spec.workload.tensor_names_used_in_multiple_einsums
+            & set(job.tensor_to_relevancy)
+        )
+
+        _, df, _, _, tensor2mapping = run_model(job)
+        new_df = {}
+        for key, value in df.items():
+            if "Total" in key:
+                new_df[key] = value
+            else:
+                new_df[f"{job.einsum_name}<SEP>{key}"] = value
+        df = new_df
+        df[f"{job.einsum_name}<SEP>mapping"] = pmapping_id
+
+        einsum = cur_spec.workload.einsums[job.einsum_name]
+        rank_variable_to_ranks = {
+            t.name: t.rank_variable2ranks for t in einsum.tensor_accesses
+        }
+        compatibility = Compatibility.from_mapping(
+            job.mapping, einsum.tensor_names, rank_variable_to_ranks
+        )
+
+        einsum2pmappings[job.einsum_name] = [
+            PmappingGroup(
+                compatibility,
+                PmappingDataframe(pd.DataFrame(df, columns=df.keys(), index=[0]), 1, 1),
+            )
+        ]
+        pmapping_objects[job.einsum_name] = {pmapping_id: job.mapping}
+
+    m = MultiEinsumPmappings(
+        einsum2pmappings,
+        pmapping_objects,
+        einsum2jobs,
+        can_combine_multiple_runs=True,
+        einsums_with_pmappings_generated=spec.workload.einsum_names,
+        flattened_arches=flattened_arches,
+        parsed_specs=parsed_specs,
+    )
+
+    return clean_compress_and_join_pmappings(spec, m)
+
+
+def _add_backing_to_tensor_holders(pmapping: Mapping):
+    seen_tensors = set()
+    for node in pmapping.nodes:
+        if isinstance(node, TensorHolder):
+            new_tensors = set(node.tensors) - seen_tensors
+            node._backing = new_tensors
+            seen_tensors.update(new_tensors)
+
+
+def _split_mapping_to_pmappings(mapping: Mapping, workload: Workload):
+    """
+    A DFS-like algorithm to split a mapping into pmappings at Split nodes.
+
+    DFS has to be modified because the tree has list of nodes for nested nodes
+    instead of links to children.
+    """
+    dfs_stack: list[NodeList] = [mapping.nodes]
+    cur_pmapping = []
+
+    while dfs_stack:
+        # nodes_segment is a list of nested nodes with a Split or Compute at the end.
+        nodes_segment = dfs_stack.pop()
+        assert isinstance(nodes_segment[-1], (Split, Compute))
+
+        cur_pmapping.append(nodes_segment[:-1])
+
+        last_node = nodes_segment[-1]
+        if isinstance(last_node, Split):
+            for segment in last_node.nodes:
+                assert isinstance(segment, Nested)
+                dfs_stack.append(segment.nodes)
+        else:
+            assert isinstance(last_node, Compute)
+
+            mapping = Mapping()
+            mapping.nodes = deepcopy(
+                [n for ns in cur_pmapping for n in ns] + [last_node]
+            )
+            _remove_storage_of_unrelevant_tensors(mapping, workload)
+            yield mapping
+
+            cur_pmapping.pop()  # Remove the last segment
+
+
+def _remove_storage_of_unrelevant_tensors(pmapping: Mapping, workload: Workload):
+    """
+    Remove tensors from Storage nodes that are not relevant to the Einsum being
+    mapped.
+    """
+    einsum_name = pmapping.nodes[-1].einsum
+    einsum = workload.einsums[einsum_name]
+    relevant_tensors = set(t.name for t in einsum.tensor_accesses)
+
+    new_nodes = []
+    for node in pmapping.nodes:
+        if isinstance(node, TensorHolder):
+            node.tensors = [t for t in node.tensors if t in relevant_tensors]
+            if node.tensors:
+                new_nodes.append(node)
+        else:
+            new_nodes.append(node)
+
+    pmapping.nodes = new_nodes

accelforge/plotting/__init__.py

@@ -0,0 +1,2 @@
+from . import mappings
+from . import specs

accelforge/plotting/mappings.py

@@ -0,0 +1,219 @@
+from collections.abc import Iterable, Sequence
+
+import matplotlib.pyplot as plt
+import pandas as pd
+
+from accelforge.mapper.FFM import Mappings
+from accelforge.mapper.FFM._pareto_df.df_convention import col2energy, col2action
+from accelforge.util._base_analysis_types import VerboseActionKey
+
+
+def plot_latency_comparison(
+    mappings: Iterable[Mappings] | Mappings,
+    labels=None,
+):
+    """
+    Plot latency comparison of multiple mappings.
+
+    Parameters
+    ----------
+    mappings:
+        A mapping to plot or an iterable of mappings to plot.
+    labels:
+        Labels to use for each Mapping class in `mappings`.
+    """
+    fig, ax = _plot_column_comparison(mappings, labels, "Total<SEP>energy")
+    ax.set_ylabel("Latency (s)")
+    return fig, ax
+
+
+def plot_action_breakdown(
+    mappings: Iterable[Mappings] | Mappings,
+    separate_by: Sequence[str],
+    stack_by: Sequence[str] = None,
+    labels: Iterable[str] = None,
+):
+    """
+    Plot actions breakdown.
+
+    Parameters
+    ----------
+    mappings:
+        A mapping to plot or an iterable of mappings to plot. Each mapping will
+        be plotted in a new subplot.
+    labels:
+        Labels to use for each Mapping class in `mappings`.
+    separate_by:
+        A list that has elements in {"einsum", "tensor", "component", "action"}.
+        Different bars will be created based on `separate_by`.
+        The order from left to right will determine grouping of the breakdown.
+    stack_by:
+        A list that has elements in {"einsum", "tensor", "component", "action"}.
+        Different components in a stacked bar will be created based on `stack_by`.
+        By default, will stack actions.
+    """
+    if stack_by is None:
+        stack_by = ["action"]
+    fig, axes = _plot_breakdown(
+        mappings, labels, separate_by, stack_by, "action", col2action
+    )
+    axes[0].set_ylabel("Actions")
+    return fig, axes
+
+
+def plot_energy_breakdown(
+    mappings: Iterable[Mappings] | Mappings,
+    separate_by: Sequence[str],
+    stack_by: Sequence[str] = None,
+    labels: Iterable[str] = None,
+):
+    """
+    Plot energy breakdown.
+
+    Parameters
+    ----------
+    mappings:
+        A mapping to plot or an iterable of mappings to plot. Each mapping will
+        be plotted in a new subplot.
+    labels:
+        Labels to use for each Mapping class in `mappings`.
+    separate_by:
+        A list that has elements in {"einsum", "tensor", "component", "action"}.
+        Different bars will be created based on `separate_by`.
+        The order from left to right will determine grouping of the breakdown.
+    stack_by:
+        A list that has elements in {"einsum", "tensor", "component", "action"}.
+        Different components in a stacked bar will be created based on `stack_by`.
+    """
+    fig, axes = _plot_breakdown(
+        mappings, labels, separate_by, stack_by, "energy", col2energy
+    )
+    axes[0].set_ylabel("Energy (pJ)")
+    return fig, axes
+
+
+def _plot_breakdown(mappings, labels, separate_by, stack_by, col_keyword: str, keyer):
+    mappings = [mappings] if isinstance(mappings, Mappings) else list(mappings)
+    n_axes = sum(map(len, (m.data for m in mappings)))
+
+    fig, axes = plt.subplots(1, n_axes, sharey=True)
+    if n_axes == 1:
+        axes = [axes]
+
+    labels = (
+        labels + "-" if labels is not None else [f"{i}-" for i in range(len(mappings))]
+    )
+    assert len(labels) == len(mappings)
+
+    if len(separate_by) == 0:
+        raise ValueError("Missing categories by which to breakdown energy")
+
+    idx = 0
+    for label, df in zip(labels, (m.data for m in mappings)):
+        colnames = [c for c in df.columns if col_keyword in c and "Total" not in c]
+        bar_components = list(
+            _get_bar_components(colnames, keyer, separate_by, stack_by)
+        )
+
+        for j, (_key, row) in enumerate(df.iterrows()):
+            ax = axes[idx]
+            idx += 1
+
+            ax.set_title(f"{label}mapping{j}")
+
+            bars = []
+            label2heights = {}
+            for name, constituents in bar_components:
+                bars.append(name)
+                for stack_name, subconstituents in constituents:
+                    if not stack_name in label2heights:
+                        label2heights[stack_name] = []
+            for label in label2heights:
+                label2heights[label] = [0] * len(bars)
+
+            for name, constituents in bar_components:
+                bar_i = bars.index(name)
+                for stack_name, subconstituents in constituents:
+                    heights = label2heights[stack_name]
+
+                    height = 0
+                    for colname in subconstituents:
+                        col = df[colname].iloc[0]
+                        height += col
+                    heights[bar_i] = height
+                    assert len(heights) == len(bars)
+
+            for label, heights in label2heights.items():
+                ax.bar(bars, height=heights, label=label)
+                ax.set_xticklabels(bars, rotation=90)
+
+    for ax in axes:
+        ax.legend()
+    return fig, axes
+
+
+def plot_energy_comparison(mappings: Iterable[Mappings] | Mappings, labels=None):
+    """
+    Plot energy comparison of multiple mappings.
+
+    Parameters
+    ----------
+    mappings:
+        A mapping to plot or an iterable of mappings to plot.
+    labels:
+        Labels to use for each Mapping class in `mappings`.
+    """
+    fig, ax = _plot_column_comparison(mappings, labels, "Total<SEP>energy")
+    ax.set_ylabel("Energy (pJ)")
+    return fig, ax
+
+
+def _plot_column_comparison(mappings, labels, colname):
+    fig, ax = plt.subplots()
+
+    mappings = [mappings] if isinstance(mappings, Mappings) else list(mappings)
+    labels = labels + "-" if labels is not None else [""] * len(mappings)
+    assert len(labels) == len(mappings)
+
+    for label, df in zip(labels, (m.data for m in mappings)):
+        bars = [f"{label}mapping{i}" for i in range(len(df))]
+        heights = df[colname]
+        ax.bar(bars, heights)
+
+    return fig, ax
+
+
+def _get_bar_components(colnames, keyer, separate_by, stack_by=None):
+    if not stack_by:
+        stack_by = []
+
+    split_colnames = []
+    for c in colnames:
+        key = keyer(c)
+        if not isinstance(key, VerboseActionKey):
+            continue
+        split_colnames.append([key.einsum, key.level, key.tensor, key.action, c])
+    transposed_colnames = zip(*split_colnames)
+    df = pd.DataFrame(
+        {
+            k: v
+            for k, v in zip(
+                ["einsum", "component", "tensor", "action", "colname"],
+                transposed_colnames,
+            )
+        }
+    )
+
+    result = []
+    for group, subdf in df.groupby(by=separate_by):
+        group = ", ".join(group)
+        if not stack_by:
+            result.append((group, [(None, subdf["colname"])]))
+        else:
+            finer_separation = []
+            for subgroup, stack_df in subdf.groupby(by=stack_by):
+                stack_df = stack_df.sort_values(by="colname")
+                subgroup = ", ".join(subgroup)
+                finer_separation.append((subgroup, stack_df["colname"]))
+            result.append((group, finer_separation))
+    return result

accelforge/plotting/specs.py

@@ -0,0 +1,57 @@
+from collections.abc import Iterable
+
+import matplotlib.axes as axes
+import matplotlib.pyplot as plt
+
+from accelforge.frontend.spec import Spec
+
+
+def plot_area(
+    specs: Iterable[Spec], labels: Iterable[str] = None, ax: axes.Axes = None
+):
+    """
+    Plot area of one or more specs.
+
+    Parameters
+    ----------
+    specs:
+        An iterable of specifications.
+    labels:
+        An iterable of the same length as `specs` to use as labels in the plot.
+    ax:
+        An matplotlib Axes to use. A new one is created by default.
+    """
+    if ax is None:
+        fig, ax = plt.subplots()
+    ax.set_ylabel("Area (m^2)")
+
+    if labels is None:
+        labels = [f"spec-{i}" for i in range(len(specs))]
+    assert len(labels) == len(specs)
+
+    component2color = {}
+    for i, (label, spec) in enumerate(zip(labels, specs)):
+        heights = []
+        colors = []
+        names = []
+        height = 0
+        for component, area in spec.arch.per_component_total_area.items():
+            height += area
+            if component not in component2color:
+                color = plt.cm.tab10(len(component2color))
+                component2color[component] = color
+            else:
+                color = component2color[component]
+            heights.append(height)
+            colors.append(color)
+            names.append(component)
+
+        heights = reversed(heights)
+        colors = reversed(colors)
+        names = reversed(names)
+        for height, color, name in zip(heights, colors, names):
+            ax.bar(i, height=height, label=name, color=color)
+
+    ax.set_xticks(range(len(specs)), labels)
+
+    ax.legend()

accelforge/util/__init__.py

@@ -0,0 +1,4 @@
+from .parallel import _expfmt, _lambdify_type_check
+from .parallel import *
+from ._frozenset import fzs
+from ._parse_expressions import LiteralString

accelforge/util/_base_analysis_types.py

@@ -0,0 +1,24 @@
+from dataclasses import dataclass
+from typing import Any
+
+
+@dataclass(frozen=True)
+class ActionKey:
+    level: str
+    action: str
+
+
+@dataclass(frozen=True)
+class VerboseActionKey(ActionKey):
+    tensor: str | None
+    einsum: str
+
+
+@dataclass
+class ActionCount:
+    total: Any
+    max_per_unit: Any
+
+    @staticmethod
+    def default():
+        return ActionCount(0, 0)