accelforge 0.0.1__py3-none-any.whl

accelforge/model/_looptree/capacity/capacity.py

@@ -0,0 +1,47 @@
+from .aggregators import CAPACITY_AGGREGATORS
+
+
+def get_value_from_singular_qpolynomial(qp):
+    return qp.eval(qp.domain().sample_point())
+
+
+def compute_capacity_usage(mapping, occupancy, workload):
+    caps = {}
+    tensor_name_to_id = workload.data_space_name_to_id()
+    einsum_name_to_id = workload.einsum_name_to_id()
+
+    for node in mapping:
+        einsums = get_einsums(mapping)
+        if node["type"] == "storage":
+            buf = node["target"]
+            if buf not in caps:
+                caps[buf] = 0
+
+            for tensor in node["dspace"]:
+                tensor_id = tensor_name_to_id[tensor]
+                max_cap = 0
+                for einsum in einsums:
+                    einsum_id = einsum_name_to_id[einsum]
+                    key = (buf, tensor_id, einsum_id)
+                    if key in occupancy:
+                        max_cap = max(
+                            max_cap,
+                            get_value_from_singular_qpolynomial(occupancy[key][1]),
+                        )
+                caps[buf] += max_cap
+
+        elif node["type"] in ["sequential", "parallel", "pipeline"]:
+            aggregate_capacity = CAPACITY_AGGREGATORS[node["type"]]
+            child_caps = [
+                compute_capacity_usage(b, occupancy, workload) for b in node["branches"]
+            ]
+            aggregate_capacity(child_caps, caps)
+    return caps
+
+
+def get_einsums(mapping):
+    for node in mapping:
+        if node["type"] in ["sequential", "parallel", "pipeline"]:
+            return sum((get_einsums(b) for b in node["branches"]), start=[])
+        elif node["type"] == "compute":
+            return [node["einsum"]]

accelforge/model/_looptree/energy.py

@@ -0,0 +1,150 @@
+from collections import defaultdict
+from collections.abc import Mapping as MappingABC
+from dataclasses import dataclass
+import logging
+from numbers import Number
+from numbers import Real
+
+from accelforge.frontend import arch
+from accelforge.frontend.mapping.mapping import MappingNode
+from accelforge.frontend.spec import Spec
+from accelforge.model._looptree.reuse.symbolic import SymbolicAnalysisOutput
+from accelforge.util._base_analysis_types import (
+    ActionCount,
+    ActionKey,
+    VerboseActionKey,
+)
+from accelforge.frontend.workload import Workload
+from accelforge.frontend.mapping import Mapping
+
+
+def gather_actions(
+    looptree_results: SymbolicAnalysisOutput,
+    bindings: dict[str, str],
+    verbose: bool = False,
+    use_name=False,
+):
+    actions: dict[tuple[str, str], ActionCount] = {}
+    compute_levels = set(c.level for c in looptree_results.compute_stats)
+
+    buffet_keyer = _get_buffet_keyer(verbose, use_name, bindings)
+    compute_keyer = _get_compute_keyer(verbose, use_name, bindings)
+
+    for buffet, accesses in looptree_results.buffet_stats.items():
+        if buffet.level in compute_levels:
+            continue
+
+        level = buffet.level
+
+        if use_name:
+            level = level
+        else:
+            level = bindings[level]
+
+        key = buffet_keyer(buffet, "read")
+        if key not in actions:
+            actions[key] = ActionCount.default()
+        actions[key].total += accesses.net_total_read_actions()
+        actions[key].max_per_unit += accesses.net_max_per_unit_read_actions()
+
+        key = buffet_keyer(buffet, "write")
+        if key not in actions:
+            actions[key] = ActionCount.default()
+        actions[key].total += accesses.net_total_write_actions()
+        actions[key].max_per_unit += accesses.net_max_per_unit_write_actions()
+
+    for compute, ops in looptree_results.compute_stats.items():
+        key = compute_keyer(compute, "compute")
+        if key not in actions:
+            actions[key] = ActionCount.default()
+        actions[key].total += ops.total_ops
+        actions[key].max_per_unit += ops.max_per_unit_ops
+
+    return actions
+
+
+def _get_buffet_keyer(verbose, use_name, bindings):
+    if not verbose:
+
+        def get_buffet_key(buffet, action_name) -> ActionKey:
+            level = buffet.level
+            if use_name:
+                level = level
+            else:
+                level = bindings[level]
+            return ActionKey(level, action_name)
+
+    else:
+
+        def get_buffet_key(buffet, action_name) -> VerboseActionKey:
+            level = buffet.level
+            if use_name:
+                level = level
+            else:
+                level = bindings[level]
+            return VerboseActionKey(level, action_name, buffet.tensor, buffet.einsum)
+
+    return get_buffet_key
+
+
+def _get_compute_keyer(verbose, use_name, bindings):
+    if not verbose:
+
+        def compute_keyer(compute, action_name):
+            level = compute.level
+            if use_name:
+                level = level
+            else:
+                level = bindings[level]
+            return ActionKey(level, action_name)
+
+    else:
+
+        def compute_keyer(compute, action_name):
+            level = compute.level
+            if use_name:
+                level = level
+            else:
+                level = bindings[level]
+            return VerboseActionKey(level, action_name, None, compute.einsum)
+
+    return compute_keyer
+
+
+def compute_energy_from_actions(
+    spec: Spec,
+    action_counts: MappingABC[ActionKey, Real],
+    overall_latency: float,
+    component_to_non_power_gated_porp: dict[str, int] = None,
+) -> dict[ActionKey | VerboseActionKey, Number]:
+    if component_to_non_power_gated_porp is None:
+        logging.warning(
+            "No component_to_non_power_gated_porp provided, will not account for power gating."
+        )
+        component_to_non_power_gated_porp = {}
+
+    energy_result = {}
+    components = {}
+    for key, counts in action_counts.items():
+        if counts.total == 0:
+            continue
+        if key.level not in components:
+            components[key.level] = spec.arch.find(key.level)
+        component_obj = components[key.level]
+        try:
+            energy_per_ac = component_obj.actions[key.action].energy
+        except KeyError as e:
+            raise KeyError(
+                f"Action {key.action} not found in component {key.component}. Action occurred "
+                f"{counts.total} times."
+            ) from None
+        energy_result[key] = counts.total * energy_per_ac
+
+    for component_obj in spec.arch.get_nodes_of_type(arch.Component):
+        energy_result[ActionKey(component_obj.name, "leak")] = (
+            component_obj.total_leak_power
+            * overall_latency
+            * component_to_non_power_gated_porp.get(component_obj.name, 1)
+        )
+
+    return energy_result

accelforge/model/_looptree/equivalent_ranks.py

@@ -0,0 +1,29 @@
+# from bindings.looptree import LooptreeWorkload, LooptreeWorkloadDependencyAnalyzer
+
+
+class EquivalentGroups:
+    def __init__(self):
+        self.group_id_to_ranks = {}
+        self.rank_to_group_id = {}
+
+    @staticmethod
+    def from_workload(
+        workload: "LooptreeWorkload", analyzer: "LooptreeWorkloadDependencyAnalyzer"
+    ):
+        einsum_id_to_name = workload.einsum_id_to_name()
+
+        groups = EquivalentGroups()
+
+        seen_ranks = set()
+        for einsum_id in einsum_id_to_name:
+            for rank_id in workload.einsum_ospace_dimensions(einsum_id):
+                equiv_ranks = analyzer.equivalent_dimensions(einsum_id, rank_id)
+                equiv_ranks = frozenset(equiv_ranks)
+                if equiv_ranks not in seen_ranks:
+                    seen_ranks.add(equiv_ranks)
+                    group_id = len(groups.group_id_to_ranks)
+                    groups.group_id_to_ranks[group_id] = equiv_ranks
+                    for r in equiv_ranks:
+                        groups.rank_to_group_id[r] = group_id
+
+        return groups

accelforge/model/_looptree/latency/__init__.py

@@ -0,0 +1 @@
+from .latency import get_latency

accelforge/model/_looptree/latency/latency.py

@@ -0,0 +1,98 @@
+from typing import overload
+from sympy import Piecewise
+
+# from accelforge.model._looptree._isl.singular import get_value_from_singular_qpolynomial
+from accelforge.frontend.arch import Compute
+from accelforge.model._looptree.latency.processors import LATENCY_PROCESSORS
+from accelforge.model._looptree.reuse.isl import IslReuseAnalysisOutput
+from accelforge.model._looptree.reuse import SymbolicAnalysisOutput
+
+from accelforge.util._sympy.broadcast_max import Max
+
+# from bindings.looptree import SpatialTag
+
+
+def get_latency(looptree_results, mapping, workload, flattened_arch):
+    comp_latency = calculate_compute_latency(looptree_results, mapping, workload)
+    mem_latency = memory_latency(looptree_results, flattened_arch, mapping, workload)
+
+    overall_latency = Max(comp_latency, *mem_latency.values())
+    return overall_latency, comp_latency, mem_latency
+
+
+@overload
+def calculate_compute_latency(
+    reuse_analysis_results: IslReuseAnalysisOutput, mapping, workload
+):
+    pass
+
+
+@overload
+def calculate_compute_latency(
+    reuse_analysis_results: SymbolicAnalysisOutput, mapping, workload
+):
+    pass
+
+
+def calculate_compute_latency(reuse_analysis_results, mapping, workload):
+    if isinstance(reuse_analysis_results, IslReuseAnalysisOutput):
+        return compute_isl_latency(
+            reuse_analysis_results.temporal_steps, mapping, workload
+        )
+    elif isinstance(reuse_analysis_results, SymbolicAnalysisOutput):
+        return compute_summarized_latency(
+            reuse_analysis_results.compute_stats, mapping, workload
+        )
+
+
+def compute_summarized_latency(compute_stats, mapping, workload):
+    # TODO: this is only for single-Einsum!!!
+    longest_compute_latency = 0
+    for stats in compute_stats.values():
+        if longest_compute_latency == 0:
+            longest_compute_latency = stats.max_latency
+        else:
+            longest_compute_latency = Max(longest_compute_latency, stats.max_latency)
+    return longest_compute_latency
+
+
+def compute_isl_latency(temporal_steps, mapping, workload):
+    raise NotImplementedError()
+    return get_value_from_singular_qpolynomial(
+        _compute_latency(mapping.nodes, 0, temporal_steps, workload)[1]
+    ).to_python()
+
+
+def _compute_latency(mapping, top_idx: int, temporal_steps, workload):
+    raise NotImplementedError()
+    einsum_name_to_id = workload.einsum_name_to_id()
+
+    next_top_idx = top_idx
+    for node in mapping:
+        next_top_idx += 1
+
+        if node["type"] in LATENCY_PROCESSORS.keys():
+            children_latencies = [
+                _compute_latency(branch, next_top_idx, temporal_steps, workload)
+                for branch in node["branches"]
+            ]
+
+            return LATENCY_PROCESSORS[node["type"]](top_idx, children_latencies)
+        elif node["type"] == "compute":
+            einsum = node["einsum"]
+            if "incomplete" in node and node["incomplete"]:
+                return ([], 0)
+            einsum_id = einsum_name_to_id[einsum]
+            return temporal_steps[einsum_id]
+
+
+def ops_to_latency(dims, map):
+    raise NotImplementedError()
+    mask = [False] * len(dims)
+    new_dims = []
+    for i, d in enumerate(dims):
+        if d == SpatialTag:
+            mask[i] = True
+        else:
+            new_dims.append(d)
+    return map.domain().identity().card()

accelforge/model/_looptree/latency/memory.py

@@ -0,0 +1,120 @@
+from collections import defaultdict
+
+from accelforge.frontend import arch
+from accelforge.frontend.arch import Leaf, Memory, TensorHolder, Component
+from accelforge.frontend.mapping import Compute, Mapping
+from accelforge.frontend.spec import Spec
+
+from accelforge.model._looptree.accesses import isl_buffer_accesses_from_buffet_actions
+from accelforge.model._looptree.mapping_utilities import get_leaves
+from accelforge.model._looptree.reuse.isl import IslReuseAnalysisOutput
+from accelforge.model._looptree.reuse import SymbolicAnalysisOutput
+from accelforge.model._looptree.types import Buffet
+
+from accelforge.model._looptree.reuse.symbolic import BuffetStats
+from accelforge.util._parse_expressions import MATH_FUNCS, parse_expression
+from accelforge.util._sympy.broadcast_max import Max, Min
+import sympy as sp
+
+
+def isl_to_summarized(
+    looptree_results: IslReuseAnalysisOutput, mapping, workload
+) -> SymbolicAnalysisOutput:
+    accesses_stats = isl_buffer_accesses_from_buffet_actions(
+        looptree_results, mapping, workload, is_path=False
+    )
+    buffet_stats = {
+        Buffet(level=component, tensor=tensor, einsum=einsum): BuffetStats(
+            max_per_unit_read_actions=accesses.max_per_unit_reads,
+            max_per_unit_write_actions=accesses.max_per_unit_writes,
+        )
+        for (component, tensor, einsum), accesses in accesses_stats.items()
+    }
+    return SymbolicAnalysisOutput(buffet_stats=buffet_stats)
+
+
+def component_latency(
+    looptree_results: SymbolicAnalysisOutput,
+    flattened_arch: list[Leaf],
+    mapping: Mapping,
+    spec: Spec,
+):
+    component_to_actions: dict[str, dict[str, float]] = defaultdict(
+        lambda: defaultdict(lambda: 0)
+    )
+    name2component: dict[str, Component] = {node.name: node for node in flattened_arch}
+
+    compute_obj = flattened_arch[-1]
+    if not isinstance(compute_obj, arch.Compute):
+        raise ValueError("Last node in flattened_arch must be a Compute")
+
+    for buffet, buffet_stats in looptree_results.buffet_stats.items():
+        component = buffet.level
+        actions = component_to_actions[component]
+        if component not in name2component:
+            raise ValueError(f"Component {component} found in mapping but not arch")
+
+        for action in name2component[component].actions:
+            actions[f"{action.name}_actions"] += 0
+
+        if isinstance(name2component[component], TensorHolder):
+            actions["read_actions"] += (
+                buffet_stats.max_per_unit_read_actions
+                - buffet_stats.min_per_unit_skipped_first_read_actions
+            )
+            if not isinstance(name2component[component], arch.ProcessingStage):
+                actions["write_actions"] += (
+                    buffet_stats.max_per_unit_write_actions
+                    - buffet_stats.min_per_unit_skipped_first_write_actions
+                )
+        elif isinstance(name2component[component], arch.Compute):
+            pass
+        else:
+            raise NotImplementedError(
+                f"Component {component} is not a TensorHolder or Compute"
+            )
+
+    longest_compute_latency = Max(
+        0, *[s.max_latency for s in looptree_results.compute_stats.values()]
+    )
+    component_to_actions[compute_obj.name]["compute_actions"] = longest_compute_latency
+
+    # TODO: Unhardcode "compute" name"
+    component_to_action_latency = defaultdict(dict)
+    for component, actions in component_to_actions.items():
+        component_obj = name2component[component]
+        for action, count in actions.items():
+            action_name = action.rsplit("_", 1)[0]
+            latency = component_obj.actions[action_name].latency
+            component_to_action_latency[component][f"{action_name}_latency"] = (
+                latency * count
+            )
+
+    component_latency = {}
+
+    symbol_table_base = {
+        **dict(spec.variables),
+        "variables": spec.variables,
+        "max": Max,
+        "min": Min,
+        "sum": sp.Add,
+    }
+
+    for component, actions in component_to_actions.items():
+        component_obj = name2component[component]
+        symbol_table = {
+            "action2latency": component_to_action_latency[component],
+            **symbol_table_base,
+            **dict(name2component[component]),
+            **actions,
+            **component_to_action_latency[component],
+        }
+        if name2component[component].total_latency is not None:
+            component_latency[component] = parse_expression(
+                name2component[component].total_latency,
+                symbol_table,
+                attr_name="latency",
+                location=component,
+            )
+
+    return component_latency

accelforge/model/_looptree/latency/processors.py

@@ -0,0 +1,92 @@
+import pprint
+
+import islpy as isl
+
+# from bindings.looptree import PipelineSpatialTag
+# from pytimeloop._isl.sum import sum_until_idx, make_reduction_map
+# from pytimeloop._isl.qpolynomial import from_pw_qpolynomial_fold
+
+
+def process_sequential_latency(top_idx: int, latencies):
+    common_dim_tags = latencies[0][0][:top_idx]
+    try:
+        total_sequential_latency = sum(
+            sum_until_idx(top_idx, latency) for dim_tags, latency in latencies
+        )
+    except:
+        print("Bad input:")
+        pprint.pp(latencies)
+        raise
+    return common_dim_tags, total_sequential_latency
+
+
+def process_pipeline_latency(top_idx: int, latencies):
+    sequential_latency = process_sequential_latency(top_idx, latencies)[1]
+
+    all_dim_tags = latencies[0][0]
+    dim_tags = all_dim_tags[:]
+    for pipeline_idx in range(len(dim_tags)):
+        if isinstance(dim_tags[pipeline_idx], PipelineSpatialTag):
+            break
+
+    try:
+        dim_tags = dim_tags[: pipeline_idx + 1]
+        summed_latency = sum(
+            sum_until_idx(pipeline_idx + 1, latency) for tags, latency in latencies
+        )
+    except:
+        print("Bad input:")
+        pprint.pp(latencies)
+        raise
+
+    space = summed_latency.get_domain_space()
+    hidden_latency_map = make_hidden_latency_map(dim_tags, space, len(latencies))
+    hidden_latencies = hidden_latency_map.apply_pw_qpolynomial(summed_latency)
+
+    reduction_map = make_reduction_map(space, len(dim_tags) - 1, 1)
+    reduction_map = reduction_map.intersect_range(summed_latency.domain()).coalesce()
+    hidden_latencies, is_tight = reduction_map.apply_pw_qpolynomial_fold(
+        isl.PwQPolynomialFold.from_pw_qpolynomial(isl.fold.min, hidden_latencies)
+    )
+    hidden_latencies = from_pw_qpolynomial_fold(hidden_latencies)
+
+    # Remove last one
+    domain = hidden_latencies.domain()
+    hidden_latencies = hidden_latencies.subtract_domain(domain.lexmax())
+
+    hidden_latency = sum_until_idx(top_idx, hidden_latencies)
+
+    return all_dim_tags[:top_idx], sequential_latency - hidden_latency
+
+
+LATENCY_PROCESSORS = {
+    "sequential": process_sequential_latency,
+    "pipeline": process_pipeline_latency,
+}
+
+
+def make_hidden_latency_map(dim_tags, space, n_stages):
+    """
+    space: [..., t, ps]
+    returns: [..., t, ps] -> [..., t', ps'] : P*t+ps+1 <= P*t'+ps' < P*t+ps+P
+    """
+    assert len(dim_tags) >= 2
+
+    t_idx = len(dim_tags) - 2
+    ps_idx = len(dim_tags) - 1
+
+    tprime = isl.Aff.var_on_domain(space, isl.dim_type.set, t_idx)
+    ps_prime = isl.Aff.var_on_domain(space, isl.dim_type.set, ps_idx)
+    inner = n_stages * tprime + ps_prime
+
+    lower = n_stages * tprime + ps_prime + 1
+    upper = n_stages * tprime + ps_prime + n_stages
+
+    hidden_latency_map = lower.le_map(inner).intersect(upper.gt_map(inner))
+
+    # Make other dimensions equal
+    for i in range(t_idx):
+        var = isl.Aff.var_on_domain(space, isl.dim_type.set, i)
+        hidden_latency_map = hidden_latency_map.intersect(var.eq_map(var))
+
+    return hidden_latency_map

accelforge/model/_looptree/mapping_utilities.py

@@ -0,0 +1,71 @@
+from typing import Generator, List, Tuple
+
+from accelforge.frontend.mapping import (
+    Compute,
+    Mapping,
+    MappingNode,
+    Pipeline,
+    Sequential,
+)
+from accelforge.frontend.workload import Workload
+
+
+def get_paths(root: Mapping) -> Generator[Tuple[MappingNode, Compute], None, None]:
+    """
+    Given a MappingNode, get the paths to all all leaves in post-order.
+
+    :param root:    The root of the child exploration.
+
+    :type root:     MappingNode
+
+    :returns:       A generator of all the MappingNodes to a Compute leaf.
+    :rtype:         Generator[List[MappingNode]]
+    """
+    cur_path: List[MappingNode] = []
+    for node in root.nodes:
+        cur_path.append(node)
+        match node:
+            # Pipelines or sequentials should have their paths expanded.
+            # Mappings naturally get expanded.
+            case Mapping() | Pipeline() | Sequential():
+                for child in node.nodes:
+                    for subpath in get_paths(child):
+                        yield tuple(cur_path) + subpath
+            # Computes are leaves so should get a yield here.
+            case Compute():
+                yield tuple(cur_path)
+            # Not implemented so continue.
+            case _:
+                # TODO: Check this is correct
+                continue
+                raise NotImplementedError(
+                    f"{type(node)} does not have type elucidation.\n"
+                    f"---\n"
+                    f"node={node}"
+                )
+
+
+def get_leaves(mapping: Mapping, is_path):
+    if is_path:
+        yield mapping[-1]
+        return
+    for node in mapping:
+        if isinstance(node, Pipeline) or isinstance(node, Sequential):
+            for child in node.children:
+                yield from get_leaves(child, is_path)
+        elif isinstance(node, Compute):
+            yield node
+
+
+def get_intermediate_tensors(workload: Workload):
+    result = set()
+    for einsum in workload.einsum_id_to_name():
+        written_tensors = workload.einsums[einsum].output_tensor_names
+        for tensor in written_tensors:
+            reader_einsums = workload.reader_einsums(tensor)
+            for reader in reader_einsums:
+                if reader in workload.einsum_id_to_name():
+                    result.add(tensor)
+                    break
+
+    return result

accelforge/model/_looptree/reuse/__init__.py

@@ -0,0 +1,4 @@
+from .symbolic import (
+    SymbolicAnalysisOutput,
+    analyze_reuse_and_add_reservations_to_mapping,
+)

accelforge/model/_looptree/reuse/isl/__init__.py

@@ -0,0 +1 @@
+from .des import IslReuseAnalysisOutput

accelforge/model/_looptree/reuse/isl/des.py

@@ -0,0 +1,59 @@
+"""
+TODO:   Is this file still necessary? It is referenced elsewhere but is no longer
+        the format we are looking for.
+"""
+
+from dataclasses import dataclass, field
+
+import islpy as isl
+
+
+@dataclass
+class IslReuseAnalysisOutput:
+    ops: dict = field(default_factory=dict)
+    fills: dict = field(default_factory=dict)
+    occupancy: dict = field(default_factory=dict)
+    op_occupancy: dict = field(default_factory=dict)
+    reads_to_peer: dict = field(default_factory=dict)
+    reads_to_parent: dict = field(default_factory=dict)
+    temporal_steps: dict = field(default_factory=dict)
+    fanout: dict = field(default_factory=dict)
+    op_intensity: dict = field(default_factory=dict)
+
+
+def deserialize_looptree_output(
+    looptree_output, isl_ctx: isl.Context  #: bindings.looptree.LooptreeResult,
+) -> IslReuseAnalysisOutput:
+    output = IslReuseAnalysisOutput()
+
+    output.ops = {
+        k: (dims, isl.PwQPolynomial.read_from_str(isl_ctx, v))
+        for k, (dims, v) in looptree_output.ops.items()
+    }
+
+    output.fills = {
+        k: (dims, isl.PwQPolynomial.read_from_str(isl_ctx, v))
+        for k, (dims, v) in looptree_output.fills.items()
+    }
+
+    output.occupancy = {
+        k: (dims, isl.PwQPolynomial.read_from_str(isl_ctx, v))
+        for k, (dims, v) in looptree_output.occupancy.items()
+    }
+
+    output.reads_to_peer = {
+        k: (dims, isl.PwQPolynomial.read_from_str(isl_ctx, v))
+        for k, (dims, v) in looptree_output.reads_to_peer.items()
+    }
+
+    output.reads_to_parent = {
+        k: (dims, isl.PwQPolynomial.read_from_str(isl_ctx, v))
+        for k, (dims, v) in looptree_output.reads_to_parent.items()
+    }
+
+    output.temporal_steps = {
+        k: (dims, isl.PwQPolynomial.read_from_str(isl_ctx, v))
+        for k, (dims, v) in looptree_output.temporal_steps.items()
+    }
+
+    return output