accelforge 0.0.1__py3-none-any.whl

accelforge/mapper/FFM/_make_pmappings/make_pmappings_from_templates/run_model.py

@@ -0,0 +1,170 @@
+from sympy import Symbol
+import accelforge.frontend.arch as arch
+from accelforge.frontend.mapping import TensorHolder
+from accelforge.mapper.FFM._make_pmappings.pmapper_job import Job
+from accelforge.model._looptree.reuse import symbolic
+from accelforge.model._looptree.reuse.symbolic import (
+    analyze_reuse_and_add_reservations_to_mapping,
+)
+from accelforge.model._looptree.energy import (
+    compute_energy_from_actions,
+    gather_actions,
+)
+from accelforge.model._looptree.latency.memory import component_latency
+from accelforge.mapper.FFM._join_pmappings.pmapping_dataframe import (
+    nameloop2col,
+    tensor2col,
+    firstlatency2col,
+    action2col,
+    energy2col,
+)
+from accelforge.frontend.mapper.metrics import Metrics
+import sympy
+from numbers import Number
+from accelforge.util._sympy.broadcast_max import Max
+
+
+def run_model(
+    job: Job,
+) -> tuple[list[Symbol], dict[str, float], dict[str, float], dict[str, float]]:
+    pmapping = job.mapping
+    spec = job.spec
+    metrics = job.metrics
+    is_copy_op = job.is_copy_operation
+    workload = spec.workload
+
+    df = {}
+
+    reuse = analyze_reuse_and_add_reservations_to_mapping(job)
+
+    latency = component_latency(reuse, job.flattened_arch, pmapping, spec)
+    try:
+        overall_latency = Max(*latency.values()) if latency else 0
+    except Exception as e:
+        for k, v in latency.items():
+            if not isinstance(v, (Number, sympy.Symbol, sympy.Expr)):
+                raise ValueError(
+                    f"Invalid type for latency: {k}: {type(v)} {str(v).strip()}"
+                )
+
+        raise ValueError(
+            f"Error calculating latency for {job.einsum_name}. Could not calculate "
+            f"a symbolic max of the following latencies:\n\t"
+            + "\n\t".join(
+                [f"{k}: {type(v)} {str(v).strip()}" for k, v in latency.items()]
+            )
+        )
+    memory_to_size = {}
+    component_to_non_power_gated_porp = {}
+    usage_df = {}
+
+    non_power_gated_instances = 1
+    for node in job.flattened_arch:
+        if isinstance(node, arch.TensorHolder):
+            if isinstance(node, arch.Memory):
+                memory_to_size[node.name] = node.size
+
+        # If there's no loops that use this spatial fanout, then the model won't output
+        # any usage. We still want to reserve at least one spatial instance in this
+        # case.
+        used_fanout = reuse.fanout.get((node.name, job.einsum_name), {})
+        for s in node.spatial:
+            usage = used_fanout.get(s.name, 1) / s.fanout
+            scaled_usage = usage * s.usage_scale
+            usage_df[f"usage<SEP>spatial<SEP>{node.name}<SEP>{s.name}"] = scaled_usage
+            non_power_gated_instances *= usage
+        component_to_non_power_gated_porp[node.name] = non_power_gated_instances
+
+    actions = gather_actions(reuse, None, use_name=True)
+    energy = compute_energy_from_actions(
+        spec, actions, overall_latency, component_to_non_power_gated_porp
+    )
+    if symbolic.PRINT_FORMULAS:
+        for k, v in energy.items():
+            print(f"{k}: {v}")
+
+    fusable_tensors = workload.tensor_names_used_in_multiple_einsums
+    tensor_to_backing = {}
+    for node in pmapping.nodes:
+        if isinstance(node, TensorHolder):
+            for tensor in node.tensors:
+                if tensor not in tensor_to_backing and tensor in fusable_tensors:
+                    tensor_to_backing[tensor] = node.component
+
+    total_occupancy = {}
+    compute_unit = pmapping.nodes[-1].component
+
+    n_instances = workload.n_instances * workload.einsums[job.einsum_name].n_instances
+
+    n_loop_options = set()
+    for buffet, stats in reuse.buffet_stats.items():
+        if buffet.level == compute_unit:
+            continue
+
+        occupancy = stats.max_occupancy
+
+        if occupancy == 0:
+            continue
+        if stats.persistent:
+            occupancy *= n_instances
+
+        for tensor, backing in tensor_to_backing.items():
+            if (is_copy_op or buffet.tensor == tensor) and buffet.level == backing:
+                df[tensor2col(tensor)] = occupancy / memory_to_size[buffet.level]
+
+        total_occupancy.setdefault(buffet.level, {}).setdefault(stats.n_loops_above, 0)
+        total_occupancy[buffet.level][stats.n_loops_above] += occupancy
+        n_loop_options.add(stats.n_loops_above)
+
+    for memory, occupancies in total_occupancy.items():
+        if memory not in job.memories_track_all:
+            continue
+        running_total = 0
+        for n_loop in n_loop_options:
+            if n_loop in occupancies:
+                running_total += occupancies[n_loop]
+                df[nameloop2col(memory, n_loop)] = (
+                    running_total / memory_to_size[memory]
+                )
+
+    if metrics & Metrics.ACTIONS:
+        detailed_actions = gather_actions(reuse, None, verbose=True, use_name=True)
+        for key, count in detailed_actions.items():
+            df[action2col(key)] = count.total * n_instances
+        detailed_energy = compute_energy_from_actions(
+            spec, detailed_actions, overall_latency, component_to_non_power_gated_porp
+        )
+        for key, energy_val in detailed_energy.items():
+            df[energy2col(key)] = energy_val * n_instances
+        for component, cur_latency in latency.items():
+            df[f"latency<SEP>{component}"] = cur_latency * n_instances
+
+    if metrics & Metrics.LATENCY:
+        df["Total<SEP>latency"] = overall_latency * n_instances
+        # df[f"latency<SEP>compute"] = comp_latency * n_instances
+        # For first latency, we'll follow the convention of treating compute
+        # as a component, similarly to memory (see below).
+        for compute_level, stats in reuse.compute_stats.items():  # FIRST LATENCY
+            for idx, max_first_latency in stats.max_first_latency.items():
+                df[firstlatency2col(compute_level.level, idx)] = (
+                    max_first_latency * n_instances
+                )
+
+    if metrics & Metrics.ENERGY:
+        df["Total<SEP>energy"] = sum(energy.values()) * n_instances
+
+    per_memory_usage_df = {}
+    for memory, occupancies in total_occupancy.items():
+        if job.ignored_resources is not None and memory not in job.ignored_resources:
+            key = f"usage<SEP>memory<SEP>{memory}"
+            per_memory_usage_df[key] = (
+                sum(occupancies.values()) / memory_to_size[memory]
+            )
+
+    return (
+        reuse.symbols,
+        df,
+        per_memory_usage_df,
+        usage_df,
+        reuse.tensor2mapping,
+    )

accelforge/mapper/FFM/_make_pmappings/make_pmappings_from_templates/symbol_relations.py

@@ -0,0 +1,174 @@
+from collections import defaultdict
+from numbers import Number
+
+from sympy import Symbol
+
+from accelforge.frontend.workload import Workload
+from accelforge.frontend._workload_isl._symbolic import get_stride_and_halo
+from accelforge.frontend.mapping import (
+    Loop,
+    Mapping,
+)
+
+
+class SymbolRelations:
+    def __init__(self):
+        self.what_tiles_symbol: list[tuple[Symbol | int, Symbol | int]] = []
+        self.tile_shape_and_initial: list[tuple[Symbol | int, Symbol | int]] = []
+        self.delta_choices: list[tuple[Symbol, frozenset[int]]] = []
+        self.bounds: tuple[tuple[Symbol, int, int], ...] = ()
+
+    def make_bounds(self):
+        all_symbols = set(
+            s for w in self.what_tiles_symbol for s in w if isinstance(s, Symbol)
+        )
+        self.bounds = tuple((s, 1, self.get_max_size(s)) for s in all_symbols)
+
+    def is_stride(self, symbol: Symbol) -> bool:
+        """Check if `symbol` is a stride."""
+        for tile_shape, initial in self.tile_shape_and_initial:
+            if tile_shape == symbol:
+                return True
+            if initial == symbol:
+                return False
+        return True
+
+    def is_initial_tile_shape(self, symbol: Symbol) -> bool:
+        """Check if `symbol` is a initial tile shape."""
+        for tile_shape, initial in self.tile_shape_and_initial:
+            if tile_shape == symbol:
+                return False
+            if initial == symbol:
+                return True
+        return False
+
+    def get_tile_shape(self, symbol: Symbol) -> Symbol | int:
+        """Get the stride corresponding to the initial tile shape `symbol`."""
+        for tile_shape, initial in self.tile_shape_and_initial:
+            if initial == symbol:
+                return tile_shape
+        raise ValueError(f"Symbol {symbol} not found as initial in {self}")
+
+    def get_initial(self, symbol: Symbol, none_if_fail: bool = False) -> Symbol | int:
+        for tile_shape, initial in self.tile_shape_and_initial:
+            if tile_shape == symbol:
+                return initial
+        if not none_if_fail:
+            raise ValueError(f"Symbol {symbol} not found as tile_shape in {self}")
+        else:
+            return None
+
+    def get_delta_choices(self, symbol: Symbol) -> frozenset[int]:
+        """Get the possible initial deltas for the rank variable represented by `symbol`."""
+        for initial, choices in self.delta_choices:
+            if initial == symbol:
+                return choices
+        raise ValueError(f"Symbol {symbol} not found in {self}")
+
+    def get_inner_tiles(
+        self, symbol: Symbol, none_if_fail: bool = False
+    ) -> Symbol | int | None:
+        """Get tiles within the tile represented by `symbol`."""
+        for tiled_by, what_tiles in self.what_tiles_symbol:
+            if tiled_by == symbol:
+                return what_tiles
+        if none_if_fail:
+            return None
+        raise ValueError(f"Symbol {symbol} not found in {self}")
+
+    def get_outer_tiles(
+        self, symbol: Symbol, none_if_fail: bool = False
+    ) -> Symbol | int | None:
+        """Get the tile that contain the tile represented by `symbol`."""
+        for tiled_by, what_tiles in self.what_tiles_symbol:
+            if what_tiles == symbol:
+                return tiled_by
+        if none_if_fail:
+            return None
+        raise ValueError(f"Symbol {symbol} not found in {self}")
+
+    def get_max_size(self, symbol: Symbol) -> Number:
+        while not isinstance(symbol, Number):
+            symbol = self.get_outer_tiles(symbol)
+        return symbol
+
+    @staticmethod
+    def from_pmapping_and_shape(
+        pmapping: Mapping, shape: dict[str, int], workload: Workload
+    ) -> "SymbolRelations":
+        initial_delta_choices = get_initial_delta_choices(
+            pmapping.nodes[-1].einsum, workload
+        )
+
+        relation = SymbolRelations()
+        last_seen_loop_per_rank_var: dict[str, Symbol | int] = dict(shape)
+        for node in pmapping.nodes:
+            if not isinstance(node, Loop):
+                continue
+            prev = last_seen_loop_per_rank_var.get(node.rank_variable, None)
+            # If we're a symbol and we've seen an outer loop with the same rank variable,
+            # then we tile that one.
+            if prev is not None:
+                relation.what_tiles_symbol.append((prev, node.tile_shape))
+            last_seen_loop_per_rank_var[node.rank_variable] = node.tile_shape
+
+            if (
+                isinstance(node.initial_tile_shape, Symbol)
+                and node.initial_tile_shape != node.tile_shape
+            ):
+                relation.tile_shape_and_initial.append(
+                    (node.tile_shape, node.initial_tile_shape)
+                )
+                relation.delta_choices.append(
+                    (
+                        node.initial_tile_shape,
+                        frozenset(initial_delta_choices[node.rank_variable]),
+                    )
+                )
+
+        for r, s in last_seen_loop_per_rank_var.items():
+            if isinstance(s, Symbol):
+                relation.what_tiles_symbol.append((s, 1))
+
+        relation.make_bounds()
+        return relation
+
+
+def get_initial_delta_choices(einsum_name: str, workload: Workload):
+    stride_and_halo = get_stride_and_halo(workload)
+    einsum = workload.einsums[einsum_name]
+
+    choices = defaultdict(lambda: set([0]))
+    consumer_chains = []
+    stack = [[(None, einsum)]]
+    while stack:
+        cur_chain = stack.pop()
+        last_tensor, last_einsum = cur_chain[-1]
+        for tensor in last_einsum.output_tensor_names:
+            einsums_with_tensor_as_input = workload.einsums_with_tensor_as_input(tensor)
+
+            if len(einsums_with_tensor_as_input) == 0:
+                consumer_chains.append(cur_chain)
+
+            for next_einsum in einsums_with_tensor_as_input:
+                stack.append(cur_chain + [(tensor, next_einsum)])
+
+    for chain in consumer_chains:
+        for (_, producer), (tensor, consumer) in zip(
+            list(reversed(chain))[1:], reversed(chain)
+        ):
+            rank_stride_and_halo = stride_and_halo[(consumer.name, tensor)]
+            if tensor is None:
+                break  # done
+
+            for cons_rank_var in consumer.rank_variables:
+                for prod_rank_var in producer.rank_variables:
+                    prod_rank = prod_rank_var.upper()
+                    for cons_choice in choices[cons_rank_var]:
+                        key = (prod_rank, cons_rank_var)
+                        if key not in rank_stride_and_halo:
+                            continue
+                        stride, halo = rank_stride_and_halo[key]
+                        choices[prod_rank_var].add(int(cons_choice * stride + halo))
+
+    return choices

accelforge/mapper/FFM/_make_pmappings/pmapper_job.py

@@ -0,0 +1,282 @@
+from collections.abc import Sequence
+from dataclasses import dataclass, field
+import logging
+from numbers import Number
+from typing import Any, Callable
+from uuid import UUID, uuid4
+
+import accelforge.frontend.arch as arch
+from accelforge.frontend.mapping import (
+    Mapping,
+)
+from accelforge.frontend.spec import Spec
+from accelforge.frontend._workload_isl._symbolic import Relevant, PartiallyRelevant
+from accelforge.frontend.workload import (
+    Einsum,
+    EinsumName,
+    RankVariable,
+    SymbolTable,
+    TensorName,
+    Workload,
+    Rank,
+)
+
+from accelforge.frontend.mapper import Metrics
+from accelforge.mapper.FFM._join_pmappings.compatibility import (
+    Compatibility,
+)
+from accelforge.mapper.FFM._make_pmappings.contraints.constraints import (
+    MappingConstraints,
+    _ConstraintLambda,
+)
+from accelforge.util.parallel import _expfmt
+from accelforge.util._itertools import first
+from accelforge.frontend.mapping import Reservation as ReservationNode
+
+
+def make_compatibility(
+    mapping: Mapping,
+    fusable_tensors: set[TensorName],
+    workload: Workload,
+    rank_variable_bounds: dict[RankVariable, int],
+    stride_and_halo,
+) -> Compatibility:
+
+    einsum = workload.einsums[mapping.nodes[-1].einsum]
+    rank_variable_to_ranks = {
+        t.name: t.rank_variable2ranks for t in einsum.tensor_accesses
+    }
+    return Compatibility.from_mapping(mapping, fusable_tensors, rank_variable_to_ranks)
+
+
+@dataclass
+class Job:
+    spec: Spec | None
+    metrics: Metrics
+    rank_variable_bounds: dict[RankVariable, int]
+
+    job_id: UUID = field(default_factory=uuid4)
+
+    stride_and_halo: (
+        dict[TensorName, dict[tuple[Rank, RankVariable], tuple[int, int]]] | None
+    ) = None
+    mapping: Mapping | None = None
+    constraints: MappingConstraints | None = None
+    fusable_tensors: set[TensorName] | None = None
+    flattened_arch: list[arch.Leaf] | None = None
+
+    einsum_name: EinsumName | None = None
+    """If the Job is for a single einsum, this is the einsum name."""
+
+    _compatibility: Compatibility | None = None
+    memories_track_all: list[str] | None = None
+    memories_track_pmappings_only: list[str] | None = None
+    ignored_resources: set[str] | None = None
+    time_limit: float | int = float("inf")
+    memory_limit: float | int = float("inf")
+    messages: list[str] = field(default_factory=list)
+    pmapping_keep_rates: dict[str, float] = field(default_factory=dict)
+    tensor_to_relevancy: (
+        dict[TensorName, dict[RankVariable, Relevant | PartiallyRelevant]] | None
+    ) = None
+
+    n_total_pmappings: int = 1
+    n_valid_pmappings: int = 1
+    n_evaluated_pmappings: int = 0
+
+    _update_compatibility_with_tile_shapes_args: dict[str, Any] | None = None
+
+    symbol_table: SymbolTable | None = None
+
+    @property
+    def einsum(self) -> Einsum:
+        return self.spec.workload.einsums[self.einsum_name]
+
+    @property
+    def compatibility(self) -> Compatibility:
+        if self._compatibility is None:
+            self._make_compatibility_and_updater()
+        return self._compatibility
+
+    @compatibility.setter
+    def compatibility(self, compatibility: Compatibility):
+        self._compatibility = compatibility
+
+    def update_compatibility_with_tile_shapes(
+        self, tile_shapes: Sequence[Number], tensor2size: dict
+    ) -> Callable[[Sequence[Number], dict], Compatibility]:
+        if self._update_compatibility_with_tile_shapes_args is None:
+            self._make_compatibility_and_updater()
+        return update_compatibility_with_tile_shapes(
+            self._compatibility,
+            tile_shapes=tile_shapes,
+            tensor2size=tensor2size,
+            **self._update_compatibility_with_tile_shapes_args,
+        )
+
+    def _make_compatibility_and_updater(self):
+        from accelforge.model._looptree.reuse.symbolic import (
+            quick_insert_reservation_nodes,
+        )
+
+        with_reservations = quick_insert_reservation_nodes(self)
+        self._compatibility = make_compatibility(
+            with_reservations,
+            self.fusable_tensors,
+            self.spec.workload,
+            self.rank_variable_bounds,
+            self.stride_and_halo,
+        )
+
+    @property
+    def is_copy_operation(self) -> bool:
+        return self.spec.workload.einsums[self.einsum_name].is_copy_operation
+
+    @classmethod
+    def make_job(
+        cls,
+        **kwargs,
+    ) -> "Job":
+        defaults = {
+            "spec": None,
+            "mapping": None,
+            "workload": None,
+            "architecture": None,
+        }
+        kwargs = {**defaults, **kwargs}
+        return cls(**kwargs)
+
+    def pretty_str(self) -> str:
+        constraints = self.constraints.get_all_constraints()
+        node2constraints: dict[int, list[_ConstraintLambda]] = {}
+        for constraint in constraints:
+            for target_index in constraint._target_node_indices:
+                l = node2constraints.setdefault(target_index, [])
+                l.append(constraint)
+
+        # Reservations are added after mapping generation so it messes up the indexing
+        mapping = [n for n in self.mapping.nodes if not isinstance(n, ReservationNode)]
+
+        s = ""
+        s += "=" * 80 + "\n"
+        s += f"Mapper job with ID {self.job_id}\n"
+        s += f"Einsum name: {self.einsum_name}\n"
+        s += f"Rank variable bounds: {self.rank_variable_bounds}\n"
+        s += f"Compute node name: {self.flattened_arch[-1].name}\n"
+        s += f"Mapping:\n"
+        for i, node in enumerate(mapping):
+            cur_constraints = sorted(
+                constraints.index(c) for c in node2constraints.get(i, [])
+            )
+            s += f"\t{i} {node.compact_str()} constrained by {cur_constraints}\n"
+        s += self.constraints.pretty_str()
+        s += f"Messages:\n"
+        for m in self.messages:
+            s += f"\t{m}\n"
+
+        s += f"Total pmappings: {self.n_total_pmappings}\n"
+        s += f"Valid pmappings: {self.n_valid_pmappings}\n"
+        s += f"One in {_expfmt(self.n_total_pmappings / self.n_valid_pmappings)} pmappings is valid\n"
+        s += f"Number of pmappings evaluated: {self.n_evaluated_pmappings}\n"
+        s += f"One in {_expfmt(self.n_evaluated_pmappings / self.n_total_pmappings)} pmappings was evaluated\n"
+        s += f"Pmapping elimination reasons:\n"
+        for cause, keep_rate in self.pmapping_keep_rates.items():
+            s += f"\t{cause} kept one in {_expfmt(1/keep_rate)} pmappings\n"
+        s += "=" * 80 + "\n"
+        return s
+
+    def set_total_pmappings(self, n_pmappings: int):
+        self.n_total_pmappings = n_pmappings
+
+    def log_porp_pmappings_kept(
+        self,
+        cause: str,
+        porp_kept: float,
+        out_of: int = None,
+    ):
+        if out_of is not None:
+            n_kept = porp_kept * out_of + (self.n_total_pmappings - out_of)
+            porp_kept = n_kept / self.n_total_pmappings
+
+        if any(x == 0 for x in self.pmapping_keep_rates.values()):
+            return
+
+        self.pmapping_keep_rates.setdefault(cause, 1)
+        self.pmapping_keep_rates[cause] *= porp_kept
+
+    def log_message(self, message: str):
+        self.messages.append(message)
+        logging.info(message)
+
+    def __copy__(self) -> "Job":
+        new = self.__class__(**self.__dict__)
+        new.messages = self.messages.copy()
+        new.pmapping_keep_rates = self.pmapping_keep_rates.copy()
+        return new
+
+
+class SameSpecJobs(list[Job]):
+    @property
+    def spec(self) -> Spec:
+        return first(self).spec
+
+    @property
+    def rank_variable_bounds(self) -> dict[RankVariable, int]:
+        return first(self).rank_variable_bounds
+
+    @property
+    def metrics(self) -> Metrics:
+        return first(self).metrics
+
+
+class SameEinsumJobs(SameSpecJobs):
+    def check_invariance(self):
+        all_einsums = set(job.einsum_name for job in self)
+        if len(all_einsums) > 1:
+            raise RuntimeError("broken invariance: not all Einsums are equal.")
+
+    @property
+    def fusable_tensors(self) -> set[TensorName]:
+        return first(self).fusable_tensors
+
+    @property
+    def einsum_name(self) -> set[EinsumName]:
+        return first(self).einsum_name
+
+    @property
+    def rank_variable_bounds(self) -> dict[RankVariable, int]:
+        return first(self).rank_variable_bounds
+
+    @property
+    def stride_and_halo(
+        self,
+    ) -> dict[tuple[str, str], dict[tuple[str, str], tuple[int, int]]]:
+        return first(self).stride_and_halo
+
+    @property
+    def is_copy_op(self) -> bool:
+        return first(self).is_copy_operation
+
+
+class SameCompatibilityJobs(SameEinsumJobs):
+    """Jobs with the same compatibility before tile shape exploration."""
+
+    def check_invariance(self):
+        all_compatibilities = set(job.compatibility for job in self)
+        if len(all_compatibilities) > 1:
+            raise RuntimeError(
+                "broken invariance: " "not all compatibilities are equal."
+            )
+
+    @property
+    def compatibility(self) -> Compatibility:
+        return first(self).compatibility
+
+    @property
+    def update_compatibility_with_tile_shapes(
+        self,
+    ) -> Callable[[Sequence[Number], dict], Compatibility]:
+        return first(self).update_compatibility_with_tile_shapes
+
+    def split(self) -> list["SameCompatibilityJobs"]:
+        return [SameCompatibilityJobs([j]) for j in self]