accelforge 0.0.1__py3-none-any.whl

accelforge/_deprecate/tags.py

@@ -0,0 +1,69 @@
+"""
+These classes may be useful when constraining joining pmappings. For example,
+if pmapping A is compatible with any other with the same dataflow (say,
+dataflow DA), and pmapping B is conly compatible with pmappings with dataflow
+DA and untiled weights in the global buffer, we can tag pmapping A with {DA}
+and pmapping B with {DA, untiled_weights}.
+
+It is often useful to compare tags that are "compatible" and tags that "match".
+Moreover, we may want to use either "compatibility" or "matching" as comparison
+for creating a dictionary. For convenience, one can use TagMatch and
+TagCompatibility as classes for the keys.
+"""
+
+from fastfusion.util._frozenset import fzs
+
+
+class TagClass(fzs):
+    pass
+
+
+class Tags(fzs):
+    def __repr__(self):
+        return f"Tags(({super().__repr__()}))"
+
+    def __str__(self):
+        return f"Tags({super().__repr__()})"
+
+    def is_member_of(self, tag_class: TagClass):
+        return all(class_string in self for class_string in tag_class)
+
+    def are_compatible_with(self, tag2):
+        return all(tag2_string in self for tag2_string in tag2) or all(
+            tag1_string in tag2 for tag1_string in self
+        )
+
+    # def filter_membership(tags: set["Tag"], tag_class: TagClass) -> set["Tag"]:
+    #     return {tag for tag in tags if are_compatible_with(tag, tag_class)}
+
+    def matches(self, tag2):
+        return self == tag2
+
+    @staticmethod
+    def from_tuple(t: tuple):
+        return Tags(t)
+
+
+class TagMatch:
+    def __init__(self, tags: Tags):
+        self.tags = tags
+
+    def __str__(self):
+        return f"TagMatch({repr(self.tags)})"
+
+    def __hash__(self):
+        return hash(self.tags)
+
+    def __eq__(self, other: "TagMatch"):
+        return self.tags.matches(other.tags)
+
+
+class TagCompatibility:
+    def __init__(self, tags: Tags):
+        self.tags = tags
+
+    def __hash__(self):
+        return 0  # See note below
+
+    def __eq__(self, other: "TagCompatibility"):
+        return self.tags.are_compatible_with(other.tag)

accelforge/_deprecate/viz/interactive.py

@@ -0,0 +1,159 @@
+import os
+from typing import Optional
+import plotly
+from IPython.display import SVG, display
+import plotly.graph_objs as go
+from ipywidgets import Output, VBox, HBox
+from fastfusion._accelerated_imports import pd
+from fastfusion.mapper.FFM._join_pmappings.pmapping_dataframe import row2pmappings
+from fastfusion.frontend.mapping import Mapping
+
+
+def make_mapping(row, einsum_names, rank_variable_bounds):
+    pmappings = row2pmappings(row, einsum_names, rank_variable_bounds)
+    newmapping = Mapping._from_pmappings(
+        pmappings, rank_variable_bounds=rank_variable_bounds
+    )
+    return newmapping
+
+
+def mapping2svg(
+    mapping: pd.Series,
+    einsum_names: list[str],
+    rank_variable_bounds: Optional[dict[str, dict[str, int]]] = None,
+):
+    mapping: Mapping = make_mapping(mapping, einsum_names, rank_variable_bounds)
+    render = mapping.render()
+    return SVG(render)
+
+
+def diplay_mappings_on_fig(
+    fig: plotly.graph_objs.FigureWidget,
+    data: dict[str, pd.DataFrame],
+    mapping_svg: bool,
+    einsum_names: list[str],
+    rank_variable_bounds: Optional[dict[str, dict[str, int]]] = None,
+):
+    # fig = go.FigureWidget(fig)
+    out = Output()
+    DUAL_OUT = True
+    if mapping_svg:
+        assert DUAL_OUT
+    DUAL_OUT = False
+
+    counter = 0
+
+    @out.capture()
+    def display_mapping(trace, points, selector):
+        if not points.point_inds:
+            return
+        out.clear_output()
+        d = data[trace.name]
+        index = points.point_inds[0]
+        display(mapping2svg(d.iloc[index], einsum_names, rank_variable_bounds))
+        # backing_tensors = set(
+        #     t for tn in d.iloc[index][MAPPING_COLUMN].values() for t in tn.tensors
+        # )
+        # backing_tensors = TensorReservation.get_backing_tensors(backing_tensors)
+        # for t in sorted(backing_tensors):
+        #     print(f"{t.__repr__()},")
+        # for v in d.iloc[index][MAPPING_COLUMN].values():
+        #     print(v)
+
+    out2 = Output()
+
+    @out2.capture()
+    def display_mapping_2(trace, points, selector):
+        if not points.point_inds:
+            return
+        out2.clear_output()
+        d = data[trace.name]
+        index = points.point_inds[0]
+        display(mapping2svg(d.iloc[index]))
+        if mapping_svg:
+            os.makedirs("plots", exist_ok=True)
+            svg = mapping2svg(d.iloc[index])
+            with open(f"plots/{trace.name}.svg", "w") as f:
+                f.write(svg.data)
+        # backing_tensors = set(
+        #     t for tn in d.iloc[index][MAPPING_COLUMN].values() for t in tn.tensors
+        # )
+        # backing_tensors = TensorReservation.get_backing_tensors(backing_tensors)
+        # for t in sorted(backing_tensors):
+        #     print(f"{t.__repr__()},")
+        # for v in d.iloc[index][MAPPING_COLUMN].values():
+        #     print(v)
+
+    for i in fig.data:
+        i.on_hover(display_mapping)
+        i.on_click(display_mapping_2)
+    if not DUAL_OUT:
+        return VBox([fig, out])
+    out.layout.width = "50%"
+    out2.layout.width = "50%"
+    return VBox([fig, HBox([out, out2])])
+
+
+def plotly_show(
+    data: pd.DataFrame | dict[str, pd.DataFrame],
+    x: str,
+    y: str,
+    category: Optional[str] = None,
+    title: Optional[str] = None,
+    show_mapping: Optional[bool] = True,
+    logscales: bool = False,
+    mapping_svg: bool = False,
+    einsum_names: Optional[list[str]] = None,
+    rank_variable_bounds: Optional[dict[str, dict[str, int]]] = None,
+):
+    fig = go.FigureWidget()
+    markers = [
+        "circle",
+        "square",
+        "diamond",
+        "cross",
+        "x",
+        "triangle-up",
+        "triangle-down",
+        "triangle-left",
+        "triangle-right",
+    ]
+    if isinstance(data, dict):
+        for i, (k, v) in enumerate(data.items()):
+            v.sort_values(by=[x, y], inplace=True)
+            fig.add_scatter(
+                x=v[x],
+                y=v[y],
+                name=k,
+                line={"shape": "hv"},
+                mode="markers+lines",
+                marker={"symbol": markers[i % len(markers)]},
+            )
+    else:
+        data.sort_values(by=[x, y], inplace=True)
+        fig.add_scatter(
+            x=data[x],
+            y=data[y],
+            name="",
+            line={"shape": "hv"},
+            mode="markers+lines",
+            marker={"symbol": markers[0]},
+        )
+        data = {"": data}
+    if title is not None:
+        fig.update_layout(title=title)
+    if logscales:
+        fig.update_xaxes(type="log")
+        fig.update_yaxes(type="log")
+    fig.update_xaxes(title_text=x)
+    fig.update_yaxes(title_text=y)
+    fig.update_layout(showlegend=True)
+    # fig = px.scatter(data, x=x, y=y, color=category, title=title, log_x=logscales, log_y=logscales)
+    if show_mapping:
+        assert (
+            einsum_names is not None
+        ), f"einsum_names must be provided if show_mapping is True"
+        return diplay_mappings_on_fig(
+            fig, data, mapping_svg, einsum_names, rank_variable_bounds
+        )
+    return fig

accelforge/_deprecate/viz/reservationtree.py

@@ -0,0 +1,307 @@
+from collections import defaultdict
+import pydot
+from typing import Any, Iterable
+from fastfusion.mapper.FFM._join_pmappings.pmapping_group import (
+    Compatibility,
+    TensorReservation,
+    Loop,
+)
+from fastfusion.util import _expfmt
+from fastfusion.mapper.FFM._join_pmappings.pmapping_dataframe import col2nameloop
+
+PYDOT_NODE_DEFAULTS = {
+    "shape": "box",
+    "fontsize": "10",
+    # "margin": "0.0",
+    "margin": "0.1,0.0",
+}
+
+
+def get_einsum_key(name):
+    return f"Einsum {name}"
+
+
+class Node:
+    def __init__(self, this_level: Iterable[Any] = (), children: Iterable["Node"] = ()):
+        self.this_level = list(this_level)
+        self.children = list(children)
+
+    def access_level(self, index: int) -> "Node":
+        if index == -1:
+            return self
+        if index == None:
+            return self.children[-1].access_level(None) if self.children else self
+        return self.children[-1].access_level(index - 1)
+
+    def _to_yaml(self):
+        this_level = [t.to_yaml() for t in self.this_level]
+        if self.children:
+            this_level.append(
+                {
+                    "type": "sequential",
+                    "children": [c._to_yaml() for c in self.children],
+                }
+            )
+        return this_level
+
+    def to_yaml(self):
+        return {"mapping": "fused", "nodes": self._to_yaml()}
+
+    def to_pydot(self, graph, parent=None, invisible_root: bool = False):
+        label_lines = []
+        for t in self.this_level:
+            label_lines.append(t.pydot_str() if hasattr(t, "pydot_str") else str(t))
+        node_label = "\n".join(sorted(label_lines))
+        if invisible_root:
+            node = None
+        else:
+            node = pydot.Node(id(self), label=node_label, **PYDOT_NODE_DEFAULTS)
+            graph.add_node(node)
+        if parent:
+            reservations = "\n".join(
+                sorted(
+                    f"[{k}] {_expfmt(v)}" for k, v in self.get_reservations().items()
+                )
+            )
+            graph.add_edge(pydot.Edge(parent, node, label=reservations))
+        for child in self.children:
+            child.to_pydot(graph, node, invisible_root=False)
+
+    def add_stats(self, stats: dict[str, Any]):
+        if self.children:
+            self.children[-1].add_stats(stats)
+        else:
+            for k, v in stats.items():
+                self.this_level.append(f"{k}: {_expfmt(v)}")
+
+    def get_reservations(self) -> dict[str, int]:
+        reservations = defaultdict(lambda: 0)
+        for c in self.children:
+            for k, v in c.get_reservations().items():
+                reservations[k] = max(reservations[k], v)
+        for t in self.this_level:
+            if isinstance(t, TensorReservation):
+                reservations[str(t.resource_name)] += t.size
+        return dict(reservations)
+
+    def get_all_tensors(
+        self, _entry: bool = True, start_at: int = 0
+    ) -> list[TensorReservation]:
+        if start_at <= 0:
+            tensors = set(
+                t for t in self.this_level if isinstance(t, TensorReservation)
+            )
+        else:
+            tensors = set()
+        for c in self.children:
+            tensors.update(c.get_all_tensors(_entry=False, start_at=start_at - 1))
+        return sorted(tensors) if _entry else tensors
+
+    def get_backing_tensors(self):
+        return TensorReservation.get_backing_tensors(self.get_all_tensors())
+
+    def merge(self, other: "Node") -> "Node":
+        new_this_level = self.this_level + other.this_level
+        loops = [l for l in new_this_level if isinstance(l, Loop)]
+        new_this_level = [t for t in new_this_level if not isinstance(t, Loop)]
+        assert (
+            len(loops) <= 2
+        ), "Can only merge two nodes with two loops. Loops: " + str(loops)
+        assert len(loops) > 0, "Can only merge two nodes with loops. Loops: " + str(
+            loops
+        )
+        while len(loops) > 1:
+            loops[0] = loops[0].merge_next(loops.pop(1))
+        new_this_level.append(loops[0])
+        return Node(new_this_level, self.children + other.children)
+
+    def get_shared_tensors(
+        self, other: "Node", start_at: int = 0
+    ) -> set[TensorReservation]:
+        return set(self.get_all_tensors(start_at=start_at)) & set(
+            other.get_all_tensors(start_at=start_at)
+        )
+
+    def find_node_with(self, key: Any) -> "Node":
+        if key in self.this_level:
+            return self
+        for c in self.children:
+            found = c.find_node_with(key)
+            if found:
+                return found
+        return None
+
+    def validate_loops(self, einsum2ranks):
+        """
+        Validate the loops in the tree. This is a placeholder function and should be
+        implemented with actual validation logic.
+        """
+        for c in self.children:
+            c.validate_loops(einsum2ranks)
+        for l in self.this_level:
+            if not isinstance(l, Loop):
+                continue
+            for einsum, ranks in einsum2ranks.items():
+                n_ranks_in_einsum = sum(1 for r in l.rank_variable_names if r in ranks)
+                if n_ranks_in_einsum > 1:
+                    return False
+        return True
+
+    def deduplicate_reservations(self, to_deduplicate):
+        for c in self.children:
+            c.deduplicate_reservations(to_deduplicate)
+        i = 0
+        while i < len(self.this_level):
+            if (
+                self.this_level[i] in self.this_level[:i]
+                and self.this_level[i] in to_deduplicate
+            ):
+                self.this_level.pop(i)
+            else:
+                i += 1
+
+
+def mappings2reservationtree(
+    mappings: dict[str, Compatibility],
+    stats: dict[str, Any] = None,
+    skip_backing_tensors_in_right_branch: Iterable[str] = (),
+    still_live_tensors: set[str] = (),
+    skip_merge: bool = False,
+    per_component_energy: dict[dict[str, float]] = None,
+    add_reservations: dict[str, Any] = None,
+) -> Node:
+    prev_mappings = []
+    root = Node()
+    einsum_ids = list(mappings.keys())
+    if stats is not None:
+        assert set(einsum_ids) == set(stats.keys())
+
+    # If a tensor appears in non-back-to-back Einsums, then we need to store it for
+    # all Einsums in between
+    tensors_lifetimes = {e: [] for e in einsum_ids}
+    all_tensors = set().union(*[set(t.tensors) for t in mappings.values()])
+    backers = TensorReservation.get_backing_tensors(all_tensors)
+    for t in backers:
+        first_appearance = min(
+            i for i, ts in enumerate(mappings.values()) if t in ts.tensors
+        )
+        last_appearance = max(
+            i for i, ts in enumerate(mappings.values()) if t in ts.tensors
+        )
+        if t.name in still_live_tensors or still_live_tensors == "all":
+            last_appearance = len(einsum_ids) - 1
+        for i, l in enumerate(tensors_lifetimes.values()):
+            if first_appearance <= i <= last_appearance and t not in l:
+                l.append(t)
+
+    # Add each Einsum to the tree
+    for i, einsum_id in enumerate(einsum_ids):
+        mapping = mappings[einsum_id]
+        index = -1
+        n = root.access_level(index)
+        loops = mapping.loops[index:] if index != -1 else mapping.loops
+        for l in loops:
+            n.children.append(Node())
+            n = n.children[-1]
+            n.this_level.append(l)
+
+        # Add the tensors
+        n.children.append(Node())  # Leaf node
+        n.children[-1].this_level.append(get_einsum_key(einsum_id))
+        if per_component_energy is not None:
+            for k, v in per_component_energy[einsum_id].items():
+                n.children[-1].this_level.append(f"{k} energy: {_expfmt(v)}")
+        id2tensor = defaultdict(set)
+        for t in sorted(mapping.tensors) + tensors_lifetimes[einsum_id]:
+            id2tensor[t.name].add(t)
+        id2tensor = {k: sorted(v) for k, v in id2tensor.items()}
+        for tensor_name, tensors in id2tensor.items():
+            for tensor in tensors:
+                n = root.access_level(tensor.above_loop_index)
+                # TODO if tensor not in n.this_level or tensor not in backers:
+                if tensor not in n.this_level or tensor not in backers:
+                    n.this_level.append(tensor)
+        if stats is not None:
+            root.add_stats(stats[einsum_id])
+        prev_mappings.append(mapping)
+
+    # Start at the root. Iterate through each leaf. Recursively:
+    # - If two leaves have the same tensor and this tensor is a backing tensor,
+    #   merge the two leaves. Remove the duplicate tensor from the right leaf.
+    def merge_nodes(n: Node, level: int = 0):
+        i = 0
+        children = n.children
+        while i < len(children) - 1:
+            for j in range(len(children) - 1, i, -1):
+                shared_tensors = children[i].get_shared_tensors(children[j], start_at=1)
+                shared_tensors |= set(
+                    c
+                    for c in children[i].get_all_tensors(start_at=1)
+                    if c.name in still_live_tensors or still_live_tensors == "all"
+                )
+                if shared_tensors & set(backers):
+                    while j != i:
+                        # print(f'Level {level} merging {shared_tensors} between {i} and {j}')
+                        children[i] = children[i].merge(children.pop(i + 1))
+                        j -= 1
+                    break
+            else:
+                i += 1
+
+            this_level = children[i].this_level
+            for t0 in range(len(this_level)):
+                for t1 in range(len(this_level) - 1, t0, -1):
+                    if this_level[t0] == this_level[t1] and this_level[t0] in backers:
+                        this_level.pop(t1)
+
+        for c in children:
+            merge_nodes(c, level + 1)
+
+        n.children = children
+
+    if not skip_merge:
+        merge_nodes(root)
+
+    n = root
+    skip_backing_tensors_in_right_branch = set(skip_backing_tensors_in_right_branch)
+    while n is not None:
+        i = 0
+        while i < len(n.this_level):
+            t = n.this_level[i]
+            if t in backers and t.name in skip_backing_tensors_in_right_branch:
+                n.this_level.pop(i)
+            else:
+                i += 1
+        n = n.children[-1] if n.children else None
+
+    if add_reservations is not None:
+        for einsum, reservations in add_reservations.items():
+            for resource, size in dict(reservations).items():
+                nameloop = col2nameloop(resource)
+                if nameloop is None:
+                    continue
+                name, _ = nameloop
+                node = root.find_node_with(get_einsum_key(einsum))
+                node.this_level.append(TensorReservation("Reservation", -1, name, size))
+
+    root.deduplicate_reservations(backers)
+
+    return root
+
+
+def mappings2svg(
+    mappings: dict[str, Compatibility],
+    stats: dict[str, Any],
+    per_component_energy: dict[dict[str, float]] = None,
+):
+    root = mappings2reservationtree(
+        mappings, stats, per_component_energy=per_component_energy
+    )
+    graph = pydot.Dot(graph_type="digraph", ranksep="0.2", nodesep="0.2")
+    root.to_pydot(graph)
+    return graph.create_svg().decode("utf-8")
+
+
+def mappings2yaml(mappings: dict[str, Compatibility], stats: dict[str, Any]):
+    root = mappings2reservationtree(mappings, stats)
+    return root.to_yaml()

accelforge/_deprecate/viz/ski_slope.py

@@ -0,0 +1,88 @@
+import matplotlib.axes as mpax
+import matplotlib.pyplot as plt
+from fastfusion._accelerated_imports import np
+from fastfusion._accelerated_imports import pd
+
+from fastfusion.mapper.FFM._join_pmappings.pmapping_dataframe import PmappingDataframe
+
+
+DATAFLOW_COLUMN = "dataflow"
+
+
+def plot_ski_slope(
+    data: pd.DataFrame,
+    categorize_by_dataflow: bool = False,
+    split_by_dataflow: bool = False,
+    ax: mpax.Axes = None,
+    **kwargs
+):
+    if ax is None:
+        fig, ax = plt.subplots()
+    else:
+        fig = ax.get_figure()
+
+    if categorize_by_dataflow or split_by_dataflow:
+        _add_dataflow_to_data(data)
+
+    if not split_by_dataflow:
+        data = PmappingDataframe(data).data
+
+    separated_datas = []
+    labels = []
+    if categorize_by_dataflow or split_by_dataflow:
+        for dataflow, sub_df in data.groupby(by=DATAFLOW_COLUMN):
+            separated_datas.append(PmappingDataframe(sub_df).data)
+            labels.append(dataflow)
+    else:
+        separated_datas.append(data)
+        labels.append(None)
+
+    for label, sub_df in zip(labels, separated_datas):
+        ax.plot(
+            *make_staircase(
+                sub_df["Occupancy"].to_numpy(), sub_df["Offchip Accesses"].to_numpy()
+            ),
+            label=label,
+            **kwargs
+        )
+
+    ax.set_xlabel("Capacity")
+    ax.set_ylabel("Off-chip Accesses")
+    ax.set_xscale("log")
+    ax.set_yscale("log")
+
+    return fig, ax
+
+
+def make_staircase(x: np.array, y: np.array):
+    """Make staircase from x, y coordinates."""
+    sort_idx = np.argsort(x)
+    x, y = x[sort_idx], y[sort_idx]
+
+    shifted_x = x[1:]
+    shifted_y = y[:-1]
+    x = np.concat([x, shifted_x])
+    y = np.concat([y, shifted_y])
+
+    sort_idx = np.lexsort([x, -y])
+    x, y = x[sort_idx], y[sort_idx]
+
+    return x, y
+
+
+def _add_dataflow_to_data(data: pd.DataFrame):
+    data[DATAFLOW_COLUMN] = data["__Mappings"].apply(_dataflow_from_fullmapping)
+
+
+def _dataflow_from_fullmapping(fullmapping: str):
+    fullmapping = fullmapping.strip("[")
+    fullmapping = fullmapping.strip("]")
+    dataflow = []
+    for term in fullmapping.split(","):
+        if term[0] == "T":
+            dataflow.append(int(term[1:].split(" size ")))
+    return tuple(dataflow)
+
+
+if __name__ == "__main__":
+    print(make_staircase(np.array([1, 2, 3, 4]), np.array([4, 3, 2, 1])))

accelforge/_version.py

@@ -0,0 +1,15 @@
+from typing import Annotated
+
+__version__ = "0.5"
+
+
+def assert_version(version: str):
+    if str(version) != str(__version__):
+        raise ValueError(
+            f"Version {str(version)} is not supported. "
+            f"Only version {__version__} is supported."
+        )
+    return str(version)
+
+
+version: Annotated[str, assert_version] = __version__

accelforge/examples.py

@@ -0,0 +1,39 @@
+from pathlib import Path
+
+EXAMPLES_DIR = Path(__file__).parent.parent / "examples"
+
+
+class Directory:
+    def __init__(self, path):
+        self.path: Path = path
+
+    def __repr__(self):
+        return f"Directory({self.path})"
+
+    def __getattr__(self, name: Path):
+        target_stem: Path = self.path / name
+        if target_stem.is_dir():
+            return Directory(self.path / name)
+
+        target_yaml = target_stem.with_suffix(".yaml")
+        if target_yaml.is_file():
+            return target_yaml
+
+        raise ValueError("Not found: {target_stem} or {target_yaml}")
+
+    def iter(self):
+        if not self.path.is_dir():
+            return
+        for path in self.path.iterdir():
+            yield path.stem
+
+
+examples = Directory(EXAMPLES_DIR)
+"""
+Convenient variable for getting path to examples directory.
+
+For example:
+```
+path_to_gh100_yaml: pathlib.Path = examples.arches.gh100
+```
+"""

accelforge/frontend/__init__.py

@@ -0,0 +1,10 @@
+"""Timeloop Spec. Each piece below (minus processors) corresponds to a top key in the Timeloop spec."""
+
+from .spec import *
+import accelforge.frontend.arch as arch
+import accelforge.frontend.config as config
+import accelforge.frontend.mapping as mapping
+import accelforge.frontend.renames as renames
+import accelforge.frontend.spec as spec
+import accelforge.frontend.variables as variables
+import accelforge.frontend.workload as workload