diff --git a/docs/source/hardware.md b/docs/source/hardware.md index 37f05d549..cd9af7da8 100644 --- a/docs/source/hardware.md +++ b/docs/source/hardware.md @@ -59,13 +59,21 @@ core_connectivity: | `cores` | yes | Map of integer **core id → core file**. Paths may be relative (`./cores/foo.yaml`); a bare filename is resolved against `/cores/` then `/`. | | `offchip_core_id` | yes | The id of the core that fronts external memory (DRAM). Must appear in `cores`. No computation is ever placed here. | | `kernel_library` | no | Path, relative to this file, to the [kernel library](kernel_library.md) the target's AIE kernels are described by. | -| `bandwidth` | no | Map of core id to the measured bandwidth every transfer through that core shares, in bits per cycle: `ceiling` for reads and writes together, `contiguous` one way, and `strided.read` / `strided.write` per contiguous span in bytes. The allocation then runs no transfer through the core faster than its access pattern allows, and bounds every step by the time the core spends on one iteration's transfers. | +| `bandwidth` | no | Map of core id to the measured bandwidth every transfer through that core shares, in bits per cycle: `ceiling` for reads and writes together, `contiguous` one way, and `strided.read` / `strided.write` per contiguous span in bytes. The allocation then runs no transfer through the core faster than its access pattern allows, and bounds every step by the time the core spends on one iteration's transfers. A key can also name one memory port, see [Memory ports](#memory-ports). | | `reconfiguration` | no | What reconfiguring the array between the designs of one dispatch costs, read by the cores' namespace constraints; for `aie2`, `cycles_per_column` per design and `reset_cycles` once. | | `core_connectivity` | yes | List of links and buses connecting the cores (see below). | | `unit_energy_cost` | no | Default energy per transferred word for every connection, unless overridden per-connection. Defaults to `0`. | | `core_coordinates` | no | Map of core id → `[col, row]`. Used for placement-aware models; required for the AIE namespace, optional otherwise. | | `core_memory_sharing` | no | Groups of core ids that share their top-level memory, e.g. `["0, 1", "2, 3"]`. Its capacity bounds everything the group's cores hold in it together: split tiles add up, a tensor each core holds whole counts once. | +### Memory ports + +A `bandwidth` key `..` (e.g. `2.dram.rw_port_1`) gives the measured model to one port of a top-level memory of a ZigZag core. It replaces that port's `bandwidth_max`, in bits per cycle. A core with a core-id key has no per-port model, so it takes no port key. + +A port belongs to a physical memory: cores in one `core_memory_sharing` group share it, and memories listed in one `memory_aliases` group share the ports of the group's first memory. + +The `memory_ports` constraint family (`SolveOptions(families=["memory_ports"])`) uses these rates. Its `interval` option (default true) fits each port's bits per iteration in its rate times the initiation interval; `burst` (default true) fits each step's bits in its rate times that step's latency. + ### Connections: `link` vs `bus` Each entry in `core_connectivity` is one connection: diff --git a/docs/source/variable_tile_co_formulation.md b/docs/source/variable_tile_co_formulation.md index 38f67f0d0..e1b199606 100644 --- a/docs/source/variable_tile_co_formulation.md +++ b/docs/source/variable_tile_co_formulation.md @@ -9,8 +9,13 @@ alongside memory allocation, transfer scheduling, and FIFO sizing. The CO allocates tensors to cores, selects transfer paths, and schedules timeslots for an AIE array. Each workload dimension (e.g. seq_len, embedding, hidden) has a **tile size** that determines how much data each core processes -per invocation. Previously tile sizes were fixed scalars chosen before the CO -ran. Now they are decision variables *inside* the MILP. +per invocation. This note describes tile sizes as decision variables *inside* +the MILP. + +The current allocator (`TransferAndTensorAllocator`) solves with fixed tile +sizes; `TileSearchStage` compares candidate tilings by solving each. Sections +1, 2, 6a, 6b and 7 have no counterpart in the code; the `**Code:**` pointers +name the methods that implement the fixed-tile form of the other sections. ## Variables at a glance @@ -44,7 +49,7 @@ sum_k w[z0, k] == 1 (one-hot) tile_var[z0] == sum_k(tile[k] * w[z0, k]) ``` -**Code:** `__create_tile_selection_vars` (~L730). +**Code:** none; tile sizes are fixed per solve (`TileSearchStage` picks among candidates). ## 2. Joint candidate enumeration @@ -79,8 +84,7 @@ tensor_size_expr = sum_combo(size[combo] * jw[combo]) Since the `w` variables are one-hot per dimension, exactly one `jw` is 1, so the expression evaluates to the correct pre-computed size. -**Code:** `_joint_candidates_for_tensor` (~L1850), -`_joint_binary_for_combo` (~L1920). +**Code:** none; `_add_binary_product` is the binary AND helper. ## 3. Memory capacity constraints @@ -110,7 +114,7 @@ When a tensor depends on multiple tiled dimensions, the (SSIS candidate, tensor candidate) pairs are enumerated, producing pre-computed `ceil(sf * size)` coefficients for each pair. -**Code:** `_memory_capacity_constraints` (~L960). +**Code:** `_memory_capacity_constraints`. ## 4. SSIS loop sizes, reuse, and fire rates @@ -144,9 +148,8 @@ fires[tr] = sum_s( z_stop[t,s] * sum_k(fires_coeff[k,s] * jw[k]) ) The inner sum is a linear expression; multiplying by the `z_stop` binary uses the same big-M `lc` auxiliary pattern. -**Code:** `_ssis_coefficients_for_transfer` (~L300), -`_transfer_fire_rate_constraints` (~L780), -`_reuse_factor_rate_constraints` (~L830). +**Code:** `_init_transfer_fire_helpers` (per stop level reuse, tiles and +buffer descriptor coefficients), `_reuse_factor_rate_constraints`. ## 5. Slot latency constraints @@ -163,7 +166,7 @@ for each tile combo k: slot_latency[s] >= sum_k(raw_lat[k] * jw[k]) ``` -**Code:** `_slot_latency_constraints` (~L1260). +**Code:** `_slot_latency_constraints`. ### 5b. Transfer nodes @@ -191,18 +194,22 @@ lat_sum = sum_s(lc_s) active_latency = y * lat_sum # gated by path choice ``` -**Code:** `_active_transfer_latency` (~L1960). +**Code:** `_active_transfer_latency`. ## 6. Objective - variable iteration count The total latency across all steady-state iterations is: ``` -total = iterations * sum(slot_lat) - (iterations - 1) * overlap - = iterations * (sum(slot_lat) - overlap) + overlap - = iterations * per_iter_net + overlap +total_latency = iterations * sum(slot_lat) - (iterations - 1) * overlap + = iterations * (sum(slot_lat) - overlap) + overlap + = iterations * per_iter_net + overlap ``` +With fixed tiles `iterations` is a constant, and the code states the first +line directly as one linear constraint on `total_latency`, plus `fill`: the +cycles a run waits for the off-chip tensors it holds in one buffer. + When tile sizes are variable, the iteration count changes because `T = workload_size / (K * S)` and `iterations = prod(T)` across all dimensions. Using a fixed iteration count with per-slot scaling (the naive @@ -253,7 +260,7 @@ With 4 dimensions * 4 candidates each, there are 256 global combos. Each adds 1 continuous auxiliary variable and 3 constraints - a negligible ~0.01% increase in model size. -**Code:** `_set_total_latency_and_objective` (~L1585). +**Code:** `_set_total_latency_and_objective` (fixed `iterations`, no `aux_k`). ### 6c. Why not scale slot latencies instead? @@ -287,7 +294,8 @@ The `check_total` field in the objective section recomputes `true_iterations * per_iter - (true_iterations - 1) * overlap` from the extracted variable values, serving as an independent sanity check. -**Code:** `build_solution_report`, `save_solution_report` (~L1780). +**Code:** none; `save_slot_latency_breakdown` writes the closest report, +`slot_latency_breakdown.yaml` (per-slot contributors, totals, slack and reuse). ## The big-M linearisation pattern diff --git a/pyproject.toml b/pyproject.toml index 537788d3d..68809fc53 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -84,6 +84,11 @@ torch_export = "stream.frontends.torch_export:TorchExportFrontend" [project.entry-points."stream.constraints"] aie2 = "stream.opt.allocation.constraint_optimization.context:AIE2Constraints" +# Optional allocation-model constraint families, selected per solve by SolveOptions.families. The +# entry-point name is the family's `name`, its object a ConstraintFamily factory. +[project.entry-points."stream.constraint_families"] +memory_ports = "stream.opt.allocation.constraint_optimization.families.memory_ports:MemoryPorts" + # Core-cost estimator backends: the entry-point name is the backend's `name` (= metadata["backend"]), # its object a CoreCostBackend that `claims` the cores it models and `make`s the estimator. The engine's # own AIE and ZigZag backends register here -- the public path is the plugin path -- so an out-of-tree diff --git a/stream/api.py b/stream/api.py index 22ca780ef..1f57ba421 100644 --- a/stream/api.py +++ b/stream/api.py @@ -12,7 +12,7 @@ import logging from collections.abc import Mapping, Sequence -from dataclasses import dataclass, field +from dataclasses import dataclass, field, replace from typing import Any from zigzag.mapping.temporal_mapping import TemporalMappingType @@ -22,6 +22,7 @@ from stream.hardware.architecture.accelerator import Accelerator from stream.instrumentation import build_instrumentation, fail_instrumentation, finish_instrumentation, instrument from stream.opt.allocation.constraint_optimization.context import build_transfer_context +from stream.opt.allocation.constraint_optimization.families import load_families from stream.opt.solver import ConstraintSelection, GurobiBackend, SolverBackend from stream.stages.allocation.constraint_optimization_allocation import ConstraintOptimizationAllocationStage from stream.stages.codegen.backends import codegen_backend_for @@ -72,6 +73,19 @@ class SolveOptions: tile_search: bool = False instrumentation: Mapping[str, Any] | None = None stage_options: Mapping[str, Any] = field(default_factory=dict) + families: Sequence[str | Mapping[str, Any]] | None = None + """Constraint families to add to the allocation model; None adds none. ``memory_ports`` bounds each top-level + memory port's bits per iteration by its rate (bits per cycle) times the initiation interval (option + ``interval``) and each slot's bits by its rate times the slot latency (option ``burst``), both on by default.""" + + def resolved_constraint_selection(self) -> ConstraintSelection | None: + """``constraint_selection`` with ``families`` folded in; an unknown family or option raises here.""" + if not self.families: + return self.constraint_selection + if self.constraint_selection is not None and self.constraint_selection.families: + raise ValueError("Set constraint families through SolveOptions.families or constraint_selection, not both") + load_families(self.families) + return replace(self.constraint_selection or ConstraintSelection(), families=tuple(self.families)) @dataclass(frozen=True) @@ -169,7 +183,7 @@ def _solve( temporal_mapping_type=TemporalMappingType[options.temporal_mapping_type.upper()], nb_cols_to_use=options.nb_cols_to_use, backend=backend.value, - constraint_selection=options.constraint_selection, + constraint_selection=options.resolved_constraint_selection(), kernel_library=options.kernel_library, tile_search=options.tile_search, **options.stage_options, diff --git a/stream/compiler/dialects/stream.py b/stream/compiler/dialects/stream.py index 17e372106..6a67cc051 100644 --- a/stream/compiler/dialects/stream.py +++ b/stream/compiler/dialects/stream.py @@ -202,6 +202,10 @@ def parse_parameter(cls, parser: AttrParser) -> StrensorVar: return var +# A window a loop moves on is prefetched into a second buffer while the first is read. +DOUBLE_BUFFERED = 2 + + @irdl_attr_definition class StrensorType( ParametrizedAttribute, @@ -217,6 +221,8 @@ class StrensorType( core_allocation: ParameterDef[ArrayAttr[StringAttr]] # how many vars are kept local? reuse_index: ParameterDef[IntAttr] + # the buffers a tile cycles through as a loop outside the reuse window moves it on + buffers: ParameterDef[IntAttr] def __init__( self, @@ -224,6 +230,7 @@ def __init__( ssis: StrensorSpace | StrensorSpaceAttr, core_allocation: Sequence[StringAttr] | ArrayAttr[StringAttr] = tuple(), reuse_index: int | IntAttr = 0, + buffers: int | IntAttr = DOUBLE_BUFFERED, ): if not isinstance(ssis, StrensorSpaceAttr): ssis = StrensorSpaceAttr(ssis) @@ -231,7 +238,9 @@ def __init__( core_allocation = ArrayAttr(core_allocation) if not isinstance(reuse_index, IntAttr): reuse_index = IntAttr(reuse_index) - super().__init__([element_type, ssis, ArrayAttr(core_allocation), reuse_index]) + if not isinstance(buffers, IntAttr): + buffers = IntAttr(buffers) + super().__init__([element_type, ssis, ArrayAttr(core_allocation), reuse_index, buffers]) @property def num_reuse_vars(self) -> int: @@ -246,7 +255,8 @@ def get_relevant_reuse_vars(self) -> Iterable[StrensorVar]: # only get relevant temporal vars relevant_dims = {var.dim for var in self.ssis.data.get_kernel_variables()} for var in vars: - if var.type == StrensorVarType.TEMPORAL and var.dim in relevant_dims: + # A loop of one keeps nothing more local than its body does. + if var.type == StrensorVarType.TEMPORAL and var.dim in relevant_dims and var.size > 1: yield var def get_local_shape(self) -> tuple[int, ...]: @@ -277,6 +287,8 @@ def print_parameters(self, printer: Printer) -> None: if isinstance(self.core_allocation, ArrayAttr): printer.print_string(", ") printer.print_attribute(self.core_allocation) + if self.buffers.data != DOUBLE_BUFFERED: + printer.print_string(f", buffers = {self.buffers.data}") @classmethod def parse_parameters(cls, parser: AttrParser) -> Sequence[Attribute]: @@ -295,9 +307,14 @@ def parse_parameters(cls, parser: AttrParser) -> Sequence[Attribute]: cores = parser.parse_attribute() else: cores = NoneAttr() + buffers = DOUBLE_BUFFERED + if parser.parse_optional_characters(",") is not None: + parser.parse_keyword("buffers") + parser.parse_punctuation("=") + buffers = parser.parse_integer() parser.parse_characters(">") reuse_idx = IntAttr(len(vars) - non_reuse_idx if non_reuse_idx else 0) - return element_type, ss, cores, reuse_idx + return element_type, ss, cores, reuse_idx, IntAttr(buffers) @irdl_op_definition diff --git a/stream/compiler/fifo_depths.py b/stream/compiler/fifo_depths.py index 233ad1209..8f38d4e0e 100644 --- a/stream/compiler/fifo_depths.py +++ b/stream/compiler/fifo_depths.py @@ -36,38 +36,42 @@ def _coords(tile: SSAValue) -> tuple[int, int] | None: def deepen( self, depths: tuple[int, ...], tiles: tuple[SSAValue, ...], object_bytes: int, feed: bool = True ) -> tuple[int, ...]: - """Per-endpoint depths, raised where the endpoint's tile has slack for the extra objects.""" - if not feed: - return depths - if len(depths) != len(tiles): + """Per-endpoint depths, raised as far as each endpoint's tile has slack for the extra objects. Readers on + compute tiles deepen together, since a broadcast advances only once every reader has room.""" + if not feed or len(depths) != len(tiles): return depths out = list(depths) + readers = [] for i, tile in enumerate(tiles): - if out[i] != DEFAULT_DEPTH: - continue coords = self._coords(tile) - if coords is None: + if out[i] != DEFAULT_DEPTH or coords is None or coords not in self.budgets: continue if coords[1] == MEM_ROW: - margin = BYTE_MARGIN + out[i] = self._spend(coords, self._affordable(coords, BYTE_MARGIN, object_bytes), object_bytes) elif coords[1] >= COMPUTE_ROW and i > 0: - margin = COMPUTE_BYTE_MARGIN - else: - continue - budget = self.budgets.get(coords) - if budget is None: - continue - extra = self.max_depth - out[i] - if extra <= 0: - continue - cost = extra * object_bytes - if budget.bytes_free * margin < cost or budget.bds_free * BD_MARGIN < extra: - continue - budget.bytes_free -= cost - budget.bds_free -= extra - out[i] = self.max_depth + readers.append((i, coords)) + if readers: + depth = min(self._affordable(coords, COMPUTE_BYTE_MARGIN, object_bytes) for _, coords in readers) + for i, coords in readers: + out[i] = self._spend(coords, depth, object_bytes) return tuple(out) + def _affordable(self, coords: tuple[int, int], margin: float, object_bytes: int) -> int: + """The deepest depth up to ``max_depth`` whose extra objects fit the tile's slack.""" + budget = self.budgets[coords] + for depth in range(self.max_depth, DEFAULT_DEPTH, -1): + extra = depth - DEFAULT_DEPTH + if budget.bytes_free * margin >= extra * object_bytes and budget.bds_free * BD_MARGIN >= extra: + return depth + return DEFAULT_DEPTH + + def _spend(self, coords: tuple[int, int], depth: int, object_bytes: int) -> int: + budget = self.budgets[coords] + extra = depth - DEFAULT_DEPTH + budget.bytes_free -= extra * object_bytes + budget.bds_free -= extra + return depth + def object_bytes(elem_bits: int, shape: tuple[int, ...]) -> int: return prod(shape) * elem_bits // 8 diff --git a/stream/compiler/kernels/aie_kernel.py b/stream/compiler/kernels/aie_kernel.py index 38ccab3a5..91ed6583a 100644 --- a/stream/compiler/kernels/aie_kernel.py +++ b/stream/compiler/kernels/aie_kernel.py @@ -1,22 +1,26 @@ from abc import ABC, abstractmethod from collections.abc import Mapping, Sequence from dataclasses import dataclass, field -from typing import ClassVar +from typing import TYPE_CHECKING, ClassVar, cast from snaxc.dialects.snax import LayoutCast from snaxc.ir.tsl import Stride, TiledStride, TiledStridedLayout -from xdsl.dialects.builtin import AnyDenseElement, FunctionType, StringAttr, bf16 +from xdsl.dialects.builtin import AnyDenseElement, FunctionType, MemRefType, StringAttr, bf16 from xdsl.dialects.func import CallOp, FuncOp from xdsl.dialects.scf import ForOp, IndexSwitchOp, YieldOp from xdsl.ir import Operation, Region, SSAValue +from xdsl.ir.affine import AffineDimExpr from xdsl.pattern_rewriter import PatternRewriter from xdsl.rewriter import InsertPoint from xdsl.traits import SymbolTable from xdsl_aie.dialects.aie import CoreOp, DeviceOp, ObjectFifoAcquireOp -from stream.compiler.dialects.stream import ComputationNodeOp, StrensorVar, StrensorVarAttr +from stream.compiler.dialects.stream import ComputationNodeOp, StrensorType, StrensorVar, StrensorVarAttr from stream.compiler.kernels.library import CallDim, KernelLibrary, KernelSpec +if TYPE_CHECKING: + from stream.workload.node import ComputationNode + MAC_TILED = "default" CONTIGUOUS = "contiguous" @@ -75,19 +79,14 @@ def induction_variable(op: Operation, var: StrensorVar, occurrence: int = 0) -> @dataclass(frozen=True) class StateOperand: - """A buffer a kernel keeps on its core from one step of a loop to the next. - - Read at ``carried_over - 1`` and written at ``carried_over``, which is the recurrence - :func:`~stream.workload.iterator_type.is_state_operand` recognises and what makes that - dimension SEQUENTIAL. ``rows`` is the extent kept per step; the node supplies the extent - of ``indexed_by``, so splitting that dimension divides the state with it. ``handover`` is - how deep a copy the next step of the computation reads, zero for a state kept private. - """ + """A buffer a kernel keeps on its core from one step of a loop to the next: read at ``carried_over - 1`` and + written at ``carried_over``, ``rows`` per step times the extent of ``indexed_by``. ``handover`` is how deep a + copy the next step of the computation reads, zero for a state kept private.""" name: str rows: int - carried_over: int - indexed_by: int + carried_over: str + indexed_by: str handover: int = 0 @@ -95,7 +94,11 @@ class StateOperand: class AIEKernel(ABC): element_type: AnyDenseElement = bf16 library: KernelLibrary | None = field(default=None, compare=False, repr=False) - DIMS: ClassVar[Mapping[str, int]] = {} + OPERAND_AXES: ClassVar[Mapping[str, tuple[int, int]]] = {} + """Each call dimension's operand and axis: the operand indexes the node's operands, inputs + then output (so ``-1`` is the output), and the axis counts from that operand's last. A + kernel addresses only the trailing axes, so the node's leading ones are batch axes it is + called once per index of.""" @property def unique_name(self) -> str: @@ -123,17 +126,42 @@ def mac(self) -> Mapping[str, int]: def call_shape(self) -> dict[str, int]: return {d.name: int(getattr(self, d.name)) for d in self.spec.dims} - def call_tile(self) -> list[tuple[int, int, CallDim]]: + def dim_positions(self, node: "ComputationNode") -> dict[str, int]: + """Where each of the kernel's dimensions sits in ``node``'s iteration space.""" + positions = {} + for name, (operand, axis) in self.OPERAND_AXES.items(): + expr = node.operand_mapping[operand].results[axis] + if not isinstance(expr, AffineDimExpr): + raise ValueError(f"{node.name}'s {name} axis is not one iteration dimension") + positions[name] = expr.position + return positions + + def output_axes(self, op: ComputationNodeOp) -> list: + """The output dimensions a call's rows and columns run along, where ``OPERAND_AXES`` places ``m`` and ``n``.""" + kernel = [var.dim for var in cast(StrensorType, op.output.type).ssis.data.get_kernel_variables()] + return [kernel[self.OPERAND_AXES[name][1]] for name in ("m", "n")] + + def call_tile(self, node: "ComputationNode") -> list[tuple[int, int, CallDim]]: """Each call dimension the library declares, as (node dimension position, size, declaration).""" - return [(self.DIMS[d.name], int(getattr(self, d.name)), d) for d in self.spec.dims] + positions = self.dim_positions(node) + return [(positions[d.name], int(getattr(self, d.name)), d) for d in self.spec.dims] def validate(self) -> None: self.spec.validate(self.call_shape()) @property def linkwith_name(self) -> str: + """The object of the declared call block.""" return self.spec.object.format(**self.call_shape()) + def call_object(self, op: ComputationNodeOp) -> str: + """The object one call links, at the extents its operands have, inputs then output as ``OPERAND_AXES`` + counts them. A runtime dimension's call takes the tile it is handed rather than the declared + block, and an object compiled for its element count must be compiled for that tile.""" + shapes = [cast(MemRefType[AnyDenseElement], operand.type).get_shape() for operand in op.inputs] + extents = {name: shapes[operand][axis] for name, (operand, axis) in self.OPERAND_AXES.items()} + return self.spec.object.format(**(self.call_shape() | extents)) + @property @abstractmethod def function_name(self) -> str: ... @@ -173,7 +201,7 @@ def rewrite(self, op: ComputationNodeOp, rewriter: PatternRewriter) -> None: while not isinstance(core_op, CoreOp): assert core_op.parent core_op = core_op.parent - core_op.link_with = StringAttr(self.linkwith_name) + core_op.link_with = StringAttr(self.call_object(op)) # replace computation node with func call op rewriter.insert_op(self.function_call(op), InsertPoint.after(op)) diff --git a/stream/compiler/kernels/eltwise_mul.py b/stream/compiler/kernels/eltwise_mul.py index 7419277ec..0da6dade3 100644 --- a/stream/compiler/kernels/eltwise_mul.py +++ b/stream/compiler/kernels/eltwise_mul.py @@ -23,7 +23,7 @@ @dataclass(kw_only=True) class EltwiseMulKernel(AIEKernel): - DIMS: ClassVar[Mapping[str, int]] = {"m": 0, "n": 1} + OPERAND_AXES: ClassVar[Mapping[str, tuple[int, int]]] = {"m": (-1, -2), "n": (-1, -1)} m: int = 32 n: int = 64 layout: str diff --git a/stream/compiler/kernels/flash.py b/stream/compiler/kernels/flash.py index 9a597a22d..6f33fb49a 100644 --- a/stream/compiler/kernels/flash.py +++ b/stream/compiler/kernels/flash.py @@ -77,10 +77,8 @@ SCALE_ROWS = 4 -# The running scale, carried over the key block and indexed by the query, which is the pair -# of dimensions every flash node's iteration space is written in. -QUERY_DIM, KEY_DIM = 0, 1 -STATE_SCALE = StateOperand("flash_state", SCALE_ROWS, carried_over=KEY_DIM, indexed_by=QUERY_DIM, handover=2) +# The running scale, carried over the key block ``n`` and indexed by the query block ``m``. +STATE_SCALE = StateOperand("flash_state", SCALE_ROWS, carried_over="n", indexed_by="m", handover=2) """Rows of ``B_q`` the scale buffer holds: m_{i-1}, m_i, l_i and exp2(m_{i-1} - m_i).""" SNAPSHOT, SNAPSHOT_OBJECT = "passThroughLine", "mha_passThrough.o" @@ -166,8 +164,9 @@ def _stamp_points(device: DeviceOp) -> None: if node is None or node.spatial_index is None: continue dim = _split_dim(node) - point = _spatial_point(node, dim) - if dim is None or point is None: + # A step no split hands out holds every query block on its one core. + point = 0 if dim is None else _spatial_point(node, dim) + if point is None: continue core.attributes[FLASH_POINT] = IntegerAttr.from_int_and_width(point, 32) core.attributes[FLASH_EXTENT] = IntegerAttr.from_int_and_width(_spatial_extent(node, dim), 32) @@ -390,10 +389,6 @@ def _block_index(op: ComputationNodeOp, dim) -> tuple[list[Operation], SSAValue, return ops, result, stride -def _kernel_dims(op: ComputationNodeOp) -> list: - return [var.dim for var in cast(StrensorType, op.output.type).ssis.data.get_kernel_variables()] - - def _store_index(buffer: SSAValue, key: SSAValue, query: SSAValue) -> list[Operation]: return [ first := ConstantOp.from_int_and_width(0, IndexType()), @@ -425,7 +420,7 @@ def work_share(self, index: int, width: int, steps: int) -> float: return attended / (width * steps * key_blocks) def function_call(self, op: ComputationNodeOp) -> Sequence[Operation]: - query, key = _kernel_dims(op) + query, key = self.output_axes(op) key_ops, key_block, _ = _block_index(op, key) query_ops, query_block, _ = _block_index(op, query) return [ @@ -532,7 +527,7 @@ def state_operands(self) -> Sequence[StateOperand]: def function_call(self, op: ComputationNodeOp) -> Sequence[Operation]: device, tile = _device(op), _tile(op) - query, key = _kernel_dims(op) + query, key = self.output_axes(op) key_ops, key_block, key_blocks = _block_index(op, key) query_ops, query_block, query_blocks = _block_index(op, query) state = self._state_buffer(device, tile) @@ -660,7 +655,7 @@ def rewrite(self, op: ComputationNodeOp, rewriter: PatternRewriter) -> None: def function_call(self, op: ComputationNodeOp) -> Sequence[Operation]: device, tile = _device(op), _tile(op) - query, key = _kernel_dims(op) + query, key = self.output_axes(op) key_ops, key_block, key_blocks = _block_index(op, key) query_ops, query_block, query_blocks = _block_index(op, query) state = self._state_buffer(device, tile) @@ -749,6 +744,7 @@ def function_type(self, op: ComputationNodeOp) -> FunctionType: i32, i32, MemRefType(i32, (2,)), + i32, ], outputs=[], ) @@ -771,9 +767,10 @@ def function_call(self, op: ComputationNodeOp) -> Sequence[Operation]: self.check_operands(op) device, tile = _device(op), _tile(op) sources = _partners(device, op, SCORE_SIDE) - kernel_dims = _kernel_dims(op) + kernel_dims = self.output_axes(op) space = cast(StrensorType, op.output.type).ssis.data - reduced = {var.dim for var in space.vars if var.type is not StrensorVarType.KERNEL} - set(kernel_dims) + indexed = {var.dim for var in space.get_kernel_variables()} + reduced = {var.dim for var in space.vars if var.type is not StrensorVarType.KERNEL} - indexed if len(reduced) > 1: raise ValueError(f"kernel {self.function_name} accumulates over one dimension, not {sorted(reduced)}") key_ops, key_block, key_blocks = _block_index(op, next(iter(reduced), None)) @@ -821,9 +818,11 @@ def _step(self, op: ComputationNodeOp, source: TileOp, key_block, key_blocks, in # as the peeled first iteration the kernel was written for. opened := CmpiOp(key_block, zero, "ne"), carried := ExtUIOp(opened, i32), + # The keys that exist: the kernel zeroes the value rows of a block past them. + keys := ConstantOp.from_int_and_width(key_blocks * self.k, i32), CallOp( self.function_name, - [op.inputs[0], op.inputs[1], op.inputs[2], scale, rows.result, carried.result, index], + [op.inputs[0], op.inputs[1], op.inputs[2], scale, rows.result, carried.result, index, keys.result], [], ), last := ConstantOp.from_int_and_width(key_blocks - 1, i32), diff --git a/stream/compiler/kernels/gemm.py b/stream/compiler/kernels/gemm.py index ce880c092..164edeed1 100644 --- a/stream/compiler/kernels/gemm.py +++ b/stream/compiler/kernels/gemm.py @@ -17,7 +17,7 @@ @dataclass(kw_only=True) class GemmKernel(AIEKernelWithZeroing): - DIMS: ClassVar[Mapping[str, int]] = {"m": 0, "k": 1, "n": 2} + OPERAND_AXES: ClassVar[Mapping[str, tuple[int, int]]] = {"m": (-1, -2), "k": (0, -1), "n": (-1, -1)} m: int k: int n: int @@ -61,6 +61,9 @@ def check_operands(self, op: ComputationNodeOp) -> None: assert op.output is not None for operand, expected in zip(op.inputs, ((self.m, self.k), (self.k, self.n), (self.m, self.n)), strict=True): shape = tuple(cast(MemRefType[AnyDenseElement], operand.type).get_shape()) + # Leading axes of extent one are the batch index the call is made at. + while len(shape) > len(expected) and shape[0] == 1: + shape = shape[1:] if shape != expected: raise ValueError( f"kernel {self.function_name} takes a {expected[0]}x{expected[1]} operand " diff --git a/stream/compiler/kernels/matvec.py b/stream/compiler/kernels/matvec.py index a2be31d49..b5f5cd1a5 100644 --- a/stream/compiler/kernels/matvec.py +++ b/stream/compiler/kernels/matvec.py @@ -1,7 +1,7 @@ -from collections.abc import Sequence +from collections.abc import Mapping, Sequence from dataclasses import dataclass from math import prod -from typing import cast +from typing import ClassVar, cast from xdsl.dialects.arith import ConstantOp from xdsl.dialects.builtin import ( @@ -19,6 +19,9 @@ @dataclass(kw_only=True) class MatVecKernel(AIEKernelWithZeroing): + # The output's rows, and the vector's length that a call takes whole. + OPERAND_AXES: ClassVar[Mapping[str, tuple[int, int]]] = {"m": (-1, -1), "n": (0, -1)} + @property def zero_name(self) -> str: return f"zero_vectorized_{self.element_type}" diff --git a/stream/compiler/kernels/silu.py b/stream/compiler/kernels/silu.py index d23487b95..48ba18341 100644 --- a/stream/compiler/kernels/silu.py +++ b/stream/compiler/kernels/silu.py @@ -23,7 +23,7 @@ @dataclass(kw_only=True) class SiluKernel(AIEKernel): - DIMS: ClassVar[Mapping[str, int]] = {"m": 0, "n": 1} + OPERAND_AXES: ClassVar[Mapping[str, tuple[int, int]]] = {"m": (-1, -2), "n": (-1, -1)} m: int = 32 n: int = 64 layout: str diff --git a/stream/compiler/kernels/softmax.py b/stream/compiler/kernels/softmax.py index ffe81337a..58f9d4038 100644 --- a/stream/compiler/kernels/softmax.py +++ b/stream/compiler/kernels/softmax.py @@ -41,7 +41,7 @@ class SoftmaxKernel(AIEKernel): normalizing it, so the tile also has to say where its first row sits globally. """ - DIMS: ClassVar[Mapping[str, int]] = {"m": 0, "n": 1} + OPERAND_AXES: ClassVar[Mapping[str, tuple[int, int]]] = {"m": (-1, -2), "n": (-1, -1)} m: int = 1 n: int layout: str @@ -93,7 +93,7 @@ def row_offset(self, op: ComputationNodeOp) -> tuple[Sequence[Operation], SSAVal drives, named by the ``layer_dim`` its rewrite left on it. """ space = cast(StrensorType, op.output.type).ssis.data - row_dim = next(iter(space.get_kernel_variables())).dim + row_dim = self.output_axes(op)[0] # A point variable carries its position where a sized one carries its size. position = {var.dim: var.size for var in (op.spatial_index.data.vars if op.spatial_index else ())} ops: list[Operation] = [] @@ -122,8 +122,8 @@ def row_offset(self, op: ComputationNodeOp) -> tuple[Sequence[Operation], SSAVal def function_call(self, op: ComputationNodeOp) -> Sequence[Operation]: shape = tuple(cast(MemRefType[AnyDenseElement], op.inputs[0].type).get_shape()) ops: list[Operation] = [ - rows := ConstantOp.from_int_and_width(shape[0], i32), - row_len := ConstantOp.from_int_and_width(shape[1], i32), + rows := ConstantOp.from_int_and_width(shape[-2], i32), + row_len := ConstantOp.from_int_and_width(shape[-1], i32), ] arguments: list[SSAValue] = [op.inputs[0], op.inputs[1], rows.result, row_len.result] if self.causal: diff --git a/stream/compiler/transforms/aie_convert_ofs.py b/stream/compiler/transforms/aie_convert_ofs.py index 280b07efd..0239e18cc 100644 --- a/stream/compiler/transforms/aie_convert_ofs.py +++ b/stream/compiler/transforms/aie_convert_ofs.py @@ -1,7 +1,7 @@ from collections.abc import Callable, Iterable, Sequence from dataclasses import dataclass from functools import reduce -from itertools import product +from itertools import pairwise, product from math import isqrt, prod from typing import Self, cast @@ -36,6 +36,7 @@ DeviceOp, DMABDOp, EndOp, + NextBDOp, ObjectFIFO, ObjectFifoAcquireOp, ObjectFifoLinkOp, @@ -294,6 +295,8 @@ def legalize(self) -> Self: # noqa: PLR0912 NB_COLUMNS = 8 +# The buffer descriptors a fifo's runtime transfers ping pong between. +PING_PONG = 2 def column_of(tile: str) -> int: @@ -916,15 +919,11 @@ def covered(strensor: StrensorType) -> set[tuple[LayerDim, int]]: @staticmethod def held_count(strensor: StrensorType) -> int: - """How many buffers of one fifo element a tile needs to keep its consumer fed. - - Two overlaps the next transfer with the current compute. One is enough when nothing - varies outside the tensor's reuse window: the fifo then hands over the same data - every iteration, which is also the single copy the allocator costed it at. - """ + """Buffers of one fifo element a tile holds: one when nothing varies outside the tensor's reuse window, so + the fifo hands over the same data every iteration, otherwise the buffers the allocator chose.""" variables = strensor.ssis.data.vars outer = variables[: len(variables) - strensor.reuse_index.data] - return 2 if any(var.type == StrensorVarType.TEMPORAL for var in outer) else 1 + return strensor.buffers.data if any(var.type == StrensorVarType.TEMPORAL for var in outer) else 1 @classmethod def held_shape(cls, strensor: StrensorType) -> tuple[int, ...]: @@ -1558,7 +1557,6 @@ def generate_reuse_pattern( # noqa: PLR0912, PLR0915 add_val := ConstantOp.from_int_and_width(0, IndexType()), ] for_op = op.parent_op() - assert isinstance(for_op, ForOp) innermost = None # innermost to outermost: for iter_var in reversed(relevant_reuse_vars): @@ -1691,6 +1689,99 @@ def match_and_rewrite(self, op: RuntimeSequenceOp, rewriter: PatternRewriter): rewriter.replace_matched_op(new_op) +@dataclass +class ChainRereads(RewritePattern): + """Turn a loop of re-reading shim tasks that differ by one constant offset step into one task: a chain of a + descriptor per re-read, each stepping the loop on its iteration dimension. A fifo keeps its tasks where the + chain needs more descriptors than its shim tile has, or the loop more steps than a descriptor iterates.""" + + descriptors: int + iterations: int + + @op_type_rewrite_pattern + def match_and_rewrite(self, op: RuntimeSequenceOp, rewriter: PatternRewriter) -> None: + device = op.parent_op() + assert isinstance(device, DeviceOp) + tiles = { + tile.result: (tile.col.value.data, tile.row.value.data) + for tile in device.walk() + if isinstance(tile, TileOp) + } + shim_of = { + fifo.sym_name.data: col + for fifo in device.walk() + if isinstance(fifo, ObjectFifoOp) + for col, row in (tiles[t] for t in (fifo.producerTile, *fifo.consumerTiles)) + if row == 0 + } + tasks: dict[str, list[DmaConfigureTaskForOp]] = {} + for task in op.body.block.ops: + if isinstance(task, DmaConfigureTaskForOp): + tasks.setdefault(task.alloc.root_reference.data, []).append(task) + # A fifo holds the descriptors SyncDMAs ping pongs between, or one where it has one task. + taken: dict[int, int] = {} + for name, fifo_tasks in tasks.items(): + taken[shim_of[name]] = taken.get(shim_of[name], 0) + min(PING_PONG, len(fifo_tasks)) + for name, fifo_tasks in tasks.items(): + loop = _reread_loop(fifo_tasks) + if loop is None: + continue + rereads, step = loop + if len(fifo_tasks) > self.iterations: + continue + column = shim_of[name] + after = taken[column] - min(PING_PONG, len(fifo_tasks)) + rereads + if after > self.descriptors: + continue + taken[column] = after + first, *rest = fifo_tasks + bd = cast(DMABDOp, first.body.block.first_op) + sizes = list(cast(DenseArrayBase, bd.static_sizes).get_values()) + strides = list(cast(DenseArrayBase, bd.static_strides).get_values()) + sizes[0], strides[0] = len(fifo_tasks), step + bd.static_sizes = DenseArrayBase.from_list(IntegerType(64), sizes) + bd.static_strides = DenseArrayBase.from_list(IntegerType(64), strides) + rewriter.erase_op(cast(Operation, first.body.block.last_op)) + blocks = [first.body.block] + for _ in range(rereads - 1): + blocks.append(Block([bd.clone()])) + first.body.add_block(blocks[-1]) + for block, following in pairwise(blocks): + block.add_op(NextBDOp(following)) + blocks[-1].add_op(EndOp()) + first.repeat_count = IntegerAttr.from_int_and_width(len(fifo_tasks) - 1, 32) + for task in rest: + rewriter.erase_op(task) + + +def _reread_loop(tasks: list[DmaConfigureTaskForOp]) -> tuple[int, int] | None: + """(re-reads a step, offset step) for tasks one loop unrolled around a re-reading transfer, else None.""" + if len(tasks) <= 1: + return None + bds = [] + for task in tasks: + bd = task.body.block.first_op + if len(task.body.blocks) != 1 or not isinstance(bd, DMABDOp) or not isinstance(bd.next_op, EndOp): + return None + bds.append(bd) + sizes = cast(DenseArrayBase, bds[0].static_sizes).get_values() + strides = cast(DenseArrayBase, bds[0].static_strides).get_values() + repeat = tasks[0].repeat_count + if strides[0] != 0 or sizes[0] <= 1 or repeat is None or repeat.value.data != sizes[0] - 1: + return None + offsets = [cast(IntegerAttr, bd.static_offset).value.data for bd in bds] + same = all( + bd.static_sizes == bds[0].static_sizes + and bd.static_strides == bds[0].static_strides + and bd.buffer == bds[0].buffer + for bd in bds + ) + steps = {b - a for a, b in pairwise(offsets)} + if not same or len(steps) != 1 or 0 in steps: + return None + return sizes[0], steps.pop() + + @dataclass class OrderDMAs(RewritePattern): @op_type_rewrite_pattern @@ -1734,7 +1825,7 @@ def match_and_rewrite(self, op: RuntimeSequenceOp, rewriter: PatternRewriter): active_tasks: dict[Attribute, list[DmaConfigureTaskForOp]] = {} # ping ponging between two bds per object fifo, so we can have at most one active task per object fifo at a time - nb_bds_per_of = 2 + nb_bds_per_of = PING_PONG for dma in op.walk(): if not isinstance(dma, DmaConfigureTaskForOp): @@ -1790,6 +1881,9 @@ class AIEConvertOfs(ModulePass): name = "aie-convert-ofs" depths: FifoDepths | None = None + # The buffer descriptors a shim tile has and the steps one iterates; without both every loop stays unrolled. + shim_descriptors: int = 0 + shim_iterations: int = 0 def apply(self, ctx: Context, op: ModuleOp) -> None: # create new shim tile @@ -1803,6 +1897,10 @@ def apply(self, ctx: Context, op: ModuleOp) -> None: ) PatternRewriteWalker(TransferToRuntimeSequence(), apply_recursively=False).rewrite_module(op) PatternRewriteWalker(StrensorToMemref()).rewrite_module(op) + if self.shim_descriptors and self.shim_iterations: + PatternRewriteWalker( + ChainRereads(self.shim_descriptors, self.shim_iterations), apply_recursively=False + ).rewrite_module(op) PatternRewriteWalker(OrderDMAs(), apply_recursively=False).rewrite_module(op) PatternRewriteWalker(SyncDMAs(), apply_recursively=False).rewrite_module(op) PatternRewriteWalker(StartDMAs(), apply_recursively=False).rewrite_module(op) diff --git a/stream/compiler/transforms/convert_stream_to_aie.py b/stream/compiler/transforms/convert_stream_to_aie.py index c2d379e42..53bf4017c 100644 --- a/stream/compiler/transforms/convert_stream_to_aie.py +++ b/stream/compiler/transforms/convert_stream_to_aie.py @@ -1,9 +1,10 @@ from collections.abc import Iterable, Sequence from dataclasses import dataclass +from math import prod from snaxc.dialects.snax import LayoutCast from snaxc.dialects.tsl import TiledStridedLayoutAttr -from snaxc.ir.tsl import Stride, TiledStridedLayout +from snaxc.ir.tsl import Stride, TiledStride, TiledStridedLayout from xdsl.context import Context from xdsl.dialects import scf from xdsl.dialects.arith import ConstantOp @@ -97,8 +98,14 @@ def match_and_rewrite(self, op: ComputationNodeOp, rewriter: PatternRewriter) -> if not layouts: return shaped_operands = [operand for operand in op.operands if isinstance(operand.type, ShapedType)] - for layout, operand in zip(layouts, shaped_operands, strict=True): + for declared, operand in zip(layouts, shaped_operands, strict=True): assert isa(old_type := operand.type, MemRefType[FixedBitwidthType]) + # A batch axis reaches the kernel as one index: a leading axis of extent one. + lead = old_type.get_shape()[: max(0, old_type.get_num_dims() - declared.dimension())] + layout = declared + if lead and all(size == 1 for size in lead): + unit = TiledStride([Stride(prod(old_type.get_shape()), 1)]) + layout = TiledStridedLayout([unit] * len(lead) + list(declared.tstrides), declared.offset) if layout.dimension() != old_type.get_num_dims(): raise ValueError( f"Kernel {op.kernel.data} declares a {layout.dimension()}-D operand layout for a " diff --git a/stream/compiler/transforms/iteration_space_to_for.py b/stream/compiler/transforms/iteration_space_to_for.py index 4bc909fe6..64b78567e 100644 --- a/stream/compiler/transforms/iteration_space_to_for.py +++ b/stream/compiler/transforms/iteration_space_to_for.py @@ -51,6 +51,10 @@ def iteration_space_to_for(block: Block, rewriter: PatternRewriter): # noqa: PL else: raise RuntimeError("non-steady state op encountered") assert ssis is not None + loops = [var for var in ssis.vars if var.type == StrensorVarType.TEMPORAL and var.dim in relevant_dims] + if not loops: + # A steady state of a single iteration runs once, without a loop. + return # some common for loop variables lb = ConstantOp.from_int_and_width(0, IndexType()) @@ -60,18 +64,13 @@ def iteration_space_to_for(block: Block, rewriter: PatternRewriter): # noqa: PL # create for loop nest: for_ops = [] innermost = None - for var in reversed(ssis.vars): - if var.type != StrensorVarType.TEMPORAL: - continue - if var.dim not in relevant_dims: - continue + for var in reversed(loops): ub = ConstantOp.from_int_and_width(var.size, IndexType()) for_op = ForOp(lb, ub, step, [], Block([*for_ops, YieldOp()], arg_types=[IndexType()])) if innermost is None: innermost = for_op for_op.attributes["layer_dim"] = StrensorVarAttr(var) for_ops = [ub, for_op] - assert innermost is not None # insert for loop nest rewriter.insert_op(for_ops, InsertPoint.after(step)) diff --git a/stream/compiler/transforms/unroll.py b/stream/compiler/transforms/unroll.py index 1e4cf3b3f..ca3ba365f 100644 --- a/stream/compiler/transforms/unroll.py +++ b/stream/compiler/transforms/unroll.py @@ -66,6 +66,7 @@ def match_and_rewrite(self, op: ComputationNodeOp, rewriter: PatternRewriter) -> output_tensor.ssis, (core,), output_tensor.reuse_index, + output_tensor.buffers, ) new_ops.append( ComputationNodeOp( @@ -116,6 +117,7 @@ def match_and_rewrite(self, op: TransferOp, rewriter: PatternRewriter) -> None: output.type.ssis, (core,), output.type.reuse_index, + output.type.buffers, ) ssis = StrensorSpaceAttr(StrensorSpace(tuple(var))) pull_ops.append(PullOp(result_type, channel, ssis)) diff --git a/stream/cost_model/bandwidth.py b/stream/cost_model/bandwidth.py index 4efca45de..86f2efb5e 100644 --- a/stream/cost_model/bandwidth.py +++ b/stream/cost_model/bandwidth.py @@ -25,10 +25,15 @@ def from_description(cls, data: Mapping[str, Any]) -> BandwidthModel: strided = {side: {int(k): float(v) for k, v in data["strided"][side].items()} for side in DIRECTIONS} return cls(ceiling=float(data["ceiling"]), contiguous=float(data["contiguous"]), strided=strided) + @classmethod + def flat(cls, bits_per_cycle: float) -> BandwidthModel: + """A rate that every access pattern reaches, as a hardware description's port width declares.""" + return cls(ceiling=bits_per_cycle, contiguous=bits_per_cycle, strided={side: {} for side in DIRECTIONS}) + def efficiency(self, span_bytes: float, direction: str) -> float: """Fraction of the contiguous rate a transfer gets whose contiguous span is ``span_bytes``.""" table = sorted(self.strided[direction].items()) - if span_bytes >= 2 * table[-1][0]: + if not table or span_bytes >= 2 * table[-1][0]: return 1.0 points = [(math.log(s), bw / self.contiguous) for s, bw in table] + [(math.log(2 * table[-1][0]), 1.0)] x = math.log(max(span_bytes, table[0][0])) diff --git a/stream/cost_model/steady_state_scheduler.py b/stream/cost_model/steady_state_scheduler.py index 9f2122198..9306f3b67 100644 --- a/stream/cost_model/steady_state_scheduler.py +++ b/stream/cost_model/steady_state_scheduler.py @@ -115,8 +115,10 @@ def __init__( # noqa: PLR0913 self.latency_total = -1 self.latency_per_iteration = -1 self.overlap_between_iterations = -1 + self.latency_fill = 0 self.performance_stats: dict | None = None self.tensor_depths: TensorDepths = {} + self.single_buffered: set[Tensor] = set() self.nb_cols_to_use = nb_cols_to_use self.transfer_context = build_transfer_context(accelerator, nb_cols_to_use=nb_cols_to_use) @@ -185,6 +187,7 @@ def get_ir(self) -> dict: "total": self.latency_total, "per_iteration": self.latency_per_iteration, "overlap_between_iterations": self.overlap_between_iterations, + "fill": self.latency_fill, }, "backend": self.backend, "solve": solve_ir, @@ -370,6 +373,8 @@ def run(self) -> Workload: latency_per_iteration, overlap, ) + self.latency_fill = round(tta.fill.X) + self.single_buffered = tta.get_single_buffered() # End-to-end MAC utilization: useful MACs vs the whole chip's peak over the full runtime # (so it folds in spatial fill, temporal stalls, idle cores AND transfer overhead). Purely # observational; never let it break the solve. diff --git a/stream/hardware/architecture/accelerator.py b/stream/hardware/architecture/accelerator.py index 39a513f02..0751feee5 100644 --- a/stream/hardware/architecture/accelerator.py +++ b/stream/hardware/architecture/accelerator.py @@ -1,3 +1,4 @@ +from functools import cached_property from typing import Any from zigzag.mapping.spatial_mapping import SpatialMapping @@ -7,6 +8,7 @@ from stream.cost_model.bandwidth import BandwidthModel from stream.cost_model.communication_manager import CommunicationManager from stream.hardware.architecture.core import Core +from stream.hardware.ports import MemoryRef, PortRef, PortRegistry class CoreGraph(DiGraphWrapper[Core]): @@ -20,7 +22,7 @@ class Accelerator: In this Stream version, the cores are actually a graph with directed edges representing communication links. """ - def __init__( + def __init__( # noqa: PLR0913 self, name: str, cores: CoreGraph, @@ -28,13 +30,20 @@ def __init__( offchip_core_id: int | None = None, kernel_library: KernelLibrary | None = None, bandwidth: dict[int, BandwidthModel] | None = None, + *, reconfiguration: dict | None = None, + port_bandwidth: dict[PortRef, BandwidthModel] | None = None, + memory_aliases: dict[MemoryRef, MemoryRef] | None = None, ): self.name = name self.kernel_library = kernel_library self.cores = cores self.offchip_core_id = offchip_core_id self.bandwidth = bandwidth or {} + # (core id, memory instance, port name) -> measured model replacing the port's declared width + self.port_bandwidth = port_bandwidth or {} + # (core id, memory instance) -> the first memory of its `memory_aliases` group, the physical one + self.memory_aliases = memory_aliases or {} self.reconfiguration = reconfiguration or {} # core id -> the first core of its `core_memory_sharing` group, whose top-level memory it uses self.shared_mem_group_ids = shared_mem_group_ids @@ -71,6 +80,11 @@ def get_spatial_mapping_from_core(self, core_allocation: list[int]) -> SpatialMa raise ValueError("Unclear which dataflow to return or no valid dataflow found.") + @cached_property + def ports(self) -> PortRegistry: + """The top-level memory ports of every core.""" + return PortRegistry.from_accelerator(self) + @property def core_list(self) -> list[Core]: return list(self.cores.node_list) diff --git a/stream/hardware/architecture/backends/__init__.py b/stream/hardware/architecture/backends/__init__.py index 7124b9179..77f9e31ff 100644 --- a/stream/hardware/architecture/backends/__init__.py +++ b/stream/hardware/architecture/backends/__init__.py @@ -1,26 +1,6 @@ -"""Pluggable core backends for Stream. - -Each backend supplies the hardware-specific details that a :class:`~stream.hardware.architecture.core.Core` -delegates to. The base :class:`Core` is backend-agnostic; it only holds identity and -stream-level scheduling attributes. - -Both backends implement the same protocol: -- ``get_memory_capacity() -> int`` -- ``get_max_memory_bandwidth(type) -> int`` -- ``get_ir() -> dict`` - -Available backends ------------------- - -``ZigZagCoreBackend`` - Thin wrapper around :mod:`zigzag.hardware.architecture.accelerator.Accelerator` - — full operational-array + memory-hierarchy model used by the ZigZag cost - estimator. - -``AIE2CoreBackend`` - Lightweight description for AIE2 tiles carrying memory capacity and - bandwidth for the simplified cost model. -""" +"""Core backends: the hardware details a :class:`~stream.hardware.architecture.core.Core` delegates to. Each +implements ``get_memory_capacity``, ``get_max_memory_bandwidth``, ``get_ir`` and ``memory_ports``, the top-level +memory ports as :class:`~stream.hardware.ports.PortSpec` (empty when the backend models none).""" from stream.hardware.architecture.backends.aie2 import AIE2CoreBackend from stream.hardware.architecture.backends.zigzag import ZigZagCoreBackend diff --git a/stream/hardware/architecture/backends/aie2.py b/stream/hardware/architecture/backends/aie2.py index 0e5216c82..4e4832add 100644 --- a/stream/hardware/architecture/backends/aie2.py +++ b/stream/hardware/architecture/backends/aie2.py @@ -13,9 +13,11 @@ from __future__ import annotations -from dataclasses import dataclass +from dataclasses import dataclass, field from typing import Literal +from stream.hardware.ports import ANY_OPERAND, READ, WRITE, PortSpec + @dataclass(frozen=True) class AIE2CoreBackend: @@ -34,6 +36,12 @@ class AIE2CoreBackend: memory_capacity_bits: int bandwidth_min: int = 0 bandwidth_max: int = 0 + dma_mm2s: int = 0 + dma_s2mm: int = 0 + dma_channel_bits: int = 0 + dma_buffer_descriptors: int = 0 + dma_iterations: int = 0 + share_group: int = field(default=-1, compare=False) # ------------------------------------------------------------------ # Backend protocol — same interface as ZigZagCoreBackend @@ -47,6 +55,21 @@ def get_max_memory_bandwidth(self, type: Literal["read"] | Literal["write"]) -> """Memory bandwidth in bits/cycle.""" return self.bandwidth_max + def port_specs(self) -> tuple[PortSpec, ...]: + """The tile DMA: MM2S channels read the tile memory onto streams, S2MM channels write into it.""" + if not self.dma_channel_bits: + return () + return tuple( + PortSpec( + self.share_group, + "dma", + name, + channels * self.dma_channel_bits, + frozenset({(ANY_OPERAND, direction)}), + ) + for name, channels, direction in (("mm2s", self.dma_mm2s, READ), ("s2mm", self.dma_s2mm, WRITE)) + ) + def get_ir(self) -> dict: """Serialize backend-specific fields for the IR dict.""" return { @@ -55,4 +78,11 @@ def get_ir(self) -> dict: "bandwidth_min": self.bandwidth_min, "bandwidth_max": self.bandwidth_max, }, + "dma": { + "mm2s": self.dma_mm2s, + "s2mm": self.dma_s2mm, + "channel_bits": self.dma_channel_bits, + "buffer_descriptors": self.dma_buffer_descriptors, + "iterations": self.dma_iterations, + }, } diff --git a/stream/hardware/architecture/backends/zigzag.py b/stream/hardware/architecture/backends/zigzag.py index bbfa9ceda..a5d929ce2 100644 --- a/stream/hardware/architecture/backends/zigzag.py +++ b/stream/hardware/architecture/backends/zigzag.py @@ -11,7 +11,36 @@ from zigzag.datatypes import MemoryOperand from zigzag.hardware.architecture.accelerator import Accelerator as _ZigZagAccelerator -from zigzag.hardware.architecture.memory_port import MemoryPortType +from zigzag.hardware.architecture.memory_level import MemoryLevel +from zigzag.hardware.architecture.memory_port import DataDirection, MemoryPort, MemoryPortType + +from stream.hardware.ports import ( + OUTPUT, + READ, + READ_BY_DATAPATH, + WRITE, + WRITE_BY_DATAPATH, + PortKey, + PortSpec, + Service, + input_role, +) + +ZIGZAG_DIRECTION_NAMES = { + DataDirection.RD_OUT_TO_HIGH: READ, + DataDirection.WR_IN_BY_HIGH: WRITE, + DataDirection.RD_OUT_TO_LOW: READ_BY_DATAPATH, + DataDirection.WR_IN_BY_LOW: WRITE_BY_DATAPATH, +} + + +def operand_role(memory_operand: str) -> str: + """Neutral role of a ZigZag memory operand: ``O`` is the output, ``I`` input k.""" + if memory_operand == "O": + return OUTPUT + if memory_operand.startswith("I") and memory_operand[1:].isdigit(): + return input_role(int(memory_operand[1:])) + raise ValueError(f"memory operand {memory_operand!r} has no operand role.") class ZigZagCoreBackend(_ZigZagAccelerator): @@ -41,6 +70,33 @@ def get_max_memory_bandwidth(self, type: Literal["read"] | Literal["write"]) -> assert first_port is not None, f"{self} does not have a top level memory {type} port." return first_port.bw_max + def port_specs(self) -> tuple[PortSpec, ...]: + """Ports of each operand's top memory level, with the (role, direction) pairs they serve at that level.""" + hierarchy = self.memory_hierarchy + tops = {op: len(hierarchy.get_memory_levels(op)) - 1 for op in hierarchy.get_operands()} + found: dict[PortKey, tuple[MemoryLevel, MemoryPort]] = {} + serves: dict[PortKey, set[Service]] = {} + for op in tops: + level = hierarchy.get_memory_levels(op)[-1] + for port in level.ports: + key = PortKey(level.memory_instance.shared_memory_group_id, level.name, port.name) + found.setdefault(key, (level, port)) + serves.setdefault(key, set()).update( + (operand_role(str(served)), ZIGZAG_DIRECTION_NAMES[direction]) + for served, lv, direction in port.served_op_lv_dir + if tops.get(served) == lv + ) + return tuple( + PortSpec( + share_group=key[0], + memory=level.name, + name=port.name, + bits_per_cycle=float(port.bw_max), + serves=frozenset(serves[key]), + ) + for key, (level, port) in found.items() + ) + def get_ir(self) -> dict: """Serialize ZigZag backend-specific fields (operational array, memory hierarchy, dataflows) into an IR dict.""" diff --git a/stream/hardware/architecture/core.py b/stream/hardware/architecture/core.py index e9b6ba7b9..3667d63ec 100644 --- a/stream/hardware/architecture/core.py +++ b/stream/hardware/architecture/core.py @@ -1,22 +1,17 @@ from __future__ import annotations -from typing import Any +from typing import TYPE_CHECKING, Any from stream.hardware.architecture.backends import AnyBackend, ZigZagCoreBackend +if TYPE_CHECKING: + from stream.hardware.ports import PortSpec -class Core: - """A single hardware core in the Stream accelerator model. - - ``Core`` is a **thin identity object** with pluggable backend. All - hardware-specific details live inside a *backend* object that implements - the backend protocol (``get_memory_capacity``, ``get_max_memory_bandwidth``, - ``get_ir``). - Access to backend attributes is transparent: ``core.operational_array`` - or ``core.mem_hierarchy_dict`` are resolved through ``__getattr__`` - delegation to the backend. - """ +class Core: + """A single hardware core: identity and scheduling attributes, with hardware details in a pluggable backend + implementing the protocol of :mod:`stream.hardware.architecture.backends`. Backend attributes such as + ``core.operational_array`` resolve through ``__getattr__`` delegation.""" def __init__( self, @@ -157,6 +152,11 @@ def get_max_memory_bandwidth(self, type: str) -> int: assert self._backend is not None, f"{self} has no backend" return self._backend.get_max_memory_bandwidth(type) # type: ignore[arg-type] + def port_specs(self) -> tuple[PortSpec, ...]: + """The ports of the core's top-level memories its backend models; empty when it models none.""" + assert self._backend is not None, f"{self} has no backend" + return self._backend.port_specs() + # ------------------------------------------------------------------ # # Serialization # # ------------------------------------------------------------------ # diff --git a/stream/hardware/architecture/noc/communication_link.py b/stream/hardware/architecture/noc/communication_link.py index a0f8ba1d9..cbdfc6dd2 100644 --- a/stream/hardware/architecture/noc/communication_link.py +++ b/stream/hardware/architecture/noc/communication_link.py @@ -3,8 +3,6 @@ if TYPE_CHECKING: from stream.hardware.architecture.core import Core -ENABLE_BROADCASTING = False - def get_bidirectional_edges( core_a: "Core", @@ -19,10 +17,6 @@ def get_bidirectional_edges( link_a_to_b = CommunicationLink(core_a, core_b, bandwidth, unit_energy_cost) link_b_to_a = CommunicationLink(core_b, core_a, bandwidth, unit_energy_cost) - # if have_shared_memory(core_a, core_b): - # # No edge if the cores have a shared memory - # return [] - return [ # A -> B ( diff --git a/stream/hardware/ports.py b/stream/hardware/ports.py new file mode 100644 index 000000000..55d96b463 --- /dev/null +++ b/stream/hardware/ports.py @@ -0,0 +1,149 @@ +"""The memory ports each core's backend declares, keyed by the physical memory they belong to.""" + +from __future__ import annotations + +from collections import defaultdict +from collections.abc import Iterator +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, NamedTuple + +from stream.cost_model.bandwidth import BandwidthModel + +if TYPE_CHECKING: + from stream.hardware.architecture.accelerator import Accelerator + from stream.hardware.architecture.core import Core + +Service = tuple[str, str] +MemoryRef = tuple[int, str] + +OUTPUT = "output" +READ = "read" +WRITE = "write" +READ_BY_DATAPATH = "read_by_datapath" +WRITE_BY_DATAPATH = "write_by_datapath" + + +class PortKey(NamedTuple): + """A physical port: the share group and name of the memory that owns it, and the port name.""" + + share_group: int + memory: str + port: str + + +class PortRef(NamedTuple): + """A port as one core names it, as in a ``bandwidth:`` port key.""" + + core_id: int + memory: str + port: str + + +ANY_OPERAND = "any" +"""Operand role of a port that serves every operand, such as a tile DMA.""" + + +def input_role(k: int) -> str: + """Operand role of the k-th input of a node, counted from 1.""" + return f"input{k}" + + +@dataclass(frozen=True) +class PortSpec: + """One port of a core's top-level memory as its backend declares it. ``serves`` holds (operand role, + direction) and ``bits_per_cycle`` is the declared peak rate.""" + + share_group: int + memory: str + name: str + bits_per_cycle: float + serves: frozenset[Service] + + +@dataclass(frozen=True) +class Port: + """One port of a top-level memory, keyed by (share group, memory, port name), with the cores that reach it. + ``serves`` holds (operand role, direction).""" + + key: PortKey + core_ids: tuple[int, ...] + bandwidth: BandwidthModel = field(hash=False) + serves: frozenset[Service] + + @property + def bits_per_cycle(self) -> float: + """Peak rate of all reads and writes through the port together.""" + return self.bandwidth.ceiling + + +class PortRegistry: + """Every top-level port of the accelerator, and which of them a core reaches.""" + + def __init__(self, ports: dict[PortKey, Port]): + self._ports = ports + self._ports_of: dict[int, tuple[Port, ...]] = defaultdict(tuple) + for port in ports.values(): + for core_id in port.core_ids: + self._ports_of[core_id] += (port,) + + @classmethod + def from_accelerator(cls, accelerator: Accelerator) -> PortRegistry: + """Ports of every core, one per port of each physical memory; a core with a measured bandwidth has none, + as that already covers its traffic. An aliased memory takes the ports of its group's first memory.""" + specs = {core.id: core.port_specs() for core in accelerator.cores.node_list} + share_group = {(core_id, spec.memory): spec.share_group for core_id in specs for spec in specs[core_id]} + found: dict[PortKey, PortSpec] = {} + core_ids: dict[PortKey, set[int]] = defaultdict(set) + serves: dict[PortKey, set[Service]] = defaultdict(set) + key_of: dict[PortRef, PortKey] = {} + for core_id, core_specs in specs.items(): + if core_id in accelerator.bandwidth: + continue + for spec in core_specs: + owner = accelerator.memory_aliases.get((core_id, spec.memory), (core_id, spec.memory)) + key = PortKey(share_group.get(owner, spec.share_group), owner[1], spec.name) + if key not in found or owner == (core_id, spec.memory): + found[key] = spec + core_ids[key].add(core_id) + serves[key].update(spec.serves) + key_of[PortRef(core_id, spec.memory, spec.name)] = key + overrides = _overrides_by_key(accelerator.port_bandwidth, key_of) + ports = { + key: Port( + key=key, + core_ids=tuple(sorted(core_ids[key])), + bandwidth=overrides.get(key, BandwidthModel.flat(spec.bits_per_cycle)), + serves=frozenset(serves[key]), + ) + for key, spec in sorted(found.items()) + } + return cls(ports) + + def __iter__(self) -> Iterator[Port]: + return iter(self._ports.values()) + + def ports_of(self, core: Core) -> tuple[Port, ...]: + """The top-level ports ``core`` reaches; none for a core without a port model.""" + return self._ports_of.get(core.id, ()) + + def port_for(self, core: Core, direction: str, operand: str) -> Port | None: + """The port ``operand``'s top level uses for ``direction`` on ``core``. + When no port serves that pair, any port of ``core`` serving ``direction``, as the bits still pass one.""" + ports = self.ports_of(core) + own = next((p for p in ports if (operand, direction) in p.serves), None) + return own or next((p for p in ports if any(d == direction for _, d in p.serves)), None) + + +def _overrides_by_key( + port_bandwidth: dict[PortRef, BandwidthModel], key_of: dict[PortRef, PortKey] +) -> dict[PortKey, BandwidthModel]: + overrides: dict[PortKey, BandwidthModel] = {} + for ref, model in port_bandwidth.items(): + if ref not in key_of: + raise ValueError( + f"`bandwidth.{ref.core_id}.{ref.memory}.{ref.port}`: core {ref.core_id} declares no port " + f"`{ref.port}` on a top-level memory `{ref.memory}`, or its backend models no memory ports." + ) + if overrides.setdefault(key_of[ref], model) != model: + raise ValueError(f"bandwidth entries disagree on shared port {key_of[ref]}.") + return overrides diff --git a/stream/inputs/aie/hardware/cores/aie_tile.yaml b/stream/inputs/aie/hardware/cores/aie_tile.yaml index 6bd4c71f1..fef85f03c 100644 --- a/stream/inputs/aie/hardware/cores/aie_tile.yaml +++ b/stream/inputs/aie/hardware/cores/aie_tile.yaml @@ -6,3 +6,8 @@ memory: capacity: 524288 # 64 KB in bits (top-level L1 tile memory) bandwidth_min: 512 # bits/cc bandwidth_max: 512 # bits/cc + +dma: # tile DMA channels, each moving 64 bits a cycle as traced on NPU2, the width of the links + mm2s: 2 + s2mm: 2 + channel_bits: 64 diff --git a/stream/inputs/aie/hardware/cores/mem_tile_256KB.yaml b/stream/inputs/aie/hardware/cores/mem_tile_256KB.yaml index 3d91c11a2..96d5a7cf6 100644 --- a/stream/inputs/aie/hardware/cores/mem_tile_256KB.yaml +++ b/stream/inputs/aie/hardware/cores/mem_tile_256KB.yaml @@ -6,3 +6,8 @@ memory: capacity: 2097152 # 256 KB in bits bandwidth_min: 128 # bits/cc bandwidth_max: 768 # bits/cc + +dma: # tile DMA channels, each moving 64 bits a cycle as traced on NPU2, the width of the links + mm2s: 6 + s2mm: 6 + channel_bits: 64 diff --git a/stream/inputs/aie/hardware/cores/mem_tile_512KB.yaml b/stream/inputs/aie/hardware/cores/mem_tile_512KB.yaml index 3c2b94913..8ffbad24b 100644 --- a/stream/inputs/aie/hardware/cores/mem_tile_512KB.yaml +++ b/stream/inputs/aie/hardware/cores/mem_tile_512KB.yaml @@ -6,3 +6,8 @@ memory: capacity: 4194304 # 512 KB in bits, AIE2TargetModel::getMemTileSize bandwidth_min: 128 # bits/cc bandwidth_max: 768 # bits/cc + +dma: # tile DMA channels, each moving 64 bits a cycle as traced on NPU2, the width of the links + mm2s: 6 + s2mm: 6 + channel_bits: 64 diff --git a/stream/inputs/aie/hardware/cores/shim_dma.yaml b/stream/inputs/aie/hardware/cores/shim_dma.yaml index f3dc0e301..22d0c8afb 100644 --- a/stream/inputs/aie/hardware/cores/shim_dma.yaml +++ b/stream/inputs/aie/hardware/cores/shim_dma.yaml @@ -6,3 +6,10 @@ memory: capacity: 10000000000 # effectively unlimited (off-chip backing) bandwidth_min: 128 # bits/cc bandwidth_max: 128 # bits/cc + +dma: # tile DMA channels, each moving 64 bits a cycle as traced on NPU2, the width of the links + mm2s: 2 + s2mm: 2 + channel_bits: 64 + buffer_descriptors: 16 + iterations: 64 diff --git a/stream/inputs/aie/hardware/whole_array_strix.yaml b/stream/inputs/aie/hardware/whole_array_strix.yaml index 32fcd1871..50ba8bd64 100644 --- a/stream/inputs/aie/hardware/whole_array_strix.yaml +++ b/stream/inputs/aie/hardware/whole_array_strix.yaml @@ -100,8 +100,8 @@ bandwidth: strided: read: {32: 19.0, 64: 35.2, 128: 53.5, 256: 57.2, 512: 93.5} write: {32: 17.1, 64: 28.3, 128: 53.2, 256: 54.2, 512: 93.5} -reconfiguration: - cycles_per_column: 69000 +reconfiguration: # a dispatch of SwiGLU's five designs from IRON costs 1.0-1.3 ms over its groups + cycles_per_column: 40000 reset_cycles: 63000 core_connectivity: # ───── Connect shim_dma to mem_tile (Column 0) ───── diff --git a/stream/ir/allocation.py b/stream/ir/allocation.py index 7ce7e2955..dd01be0ca 100644 --- a/stream/ir/allocation.py +++ b/stream/ir/allocation.py @@ -24,6 +24,10 @@ class LatencyInfo(BaseModel): overlap_between_iterations: int = Field( description="Overlap cycles between consecutive iterations (pipeline depth)" ) + fill: int = Field( + default=0, + description="Cycles each run waits before its first iteration for the tensors it holds in one buffer", + ) class CostModelsIR(BaseModel): @@ -357,6 +361,24 @@ class MemoryOccupancyIR(BaseModel): ) +class ResourceActivityIR(BaseModel): + """One resource under the memory_ports family, in ZigZag's port-activity terms.""" + + kind: str = Field(default="memory_port", description="'memory_port', 'shared_bandwidth' or 'link'") + resource: str = Field(description="'.', 'measured', or '->' for a link") + core_ids: list[int] = Field(description="Cores sharing the resource") + core_types: list[str] = Field(default_factory=list, description="Core type of each of core_ids") + bw_bits_per_cycle: float = Field(description="Port bandwidth") + bits_per_iteration: float = Field(description="Bits the solved schedule moves through the port per iteration") + req_bw_aver: float | None = Field(description="Average bandwidth required over the initiation interval") + real_cycle: float = Field(description="Cycles the port needs for one iteration's bits") + allowed_cycle: float = Field(description="Initiation interval: the cycles one iteration gives the port") + stall_or_slack: float = Field(description="real_cycle - allowed_cycle; negative is slack") + utilization: float | None = Field(description="real_cycle / allowed_cycle; 1 when the port sets the interval") + burst_utilization: float | None = Field(description="Highest per-slot bits / (bandwidth * slot latency)") + burst_slot: int | None = Field(description="Slot where burst_utilization occurs") + + class AllocationPerformanceView(BaseModel): """Performance-persona projection of AllocationIR. @@ -383,6 +405,11 @@ class AllocationPerformanceView(BaseModel): default_factory=list, description="Per-core solved residency vs declared capacity: the floor on any capacity reduction", ) + memory_ports: list[ResourceActivityIR] = Field( + default_factory=list, + description="Activity of each memory port, shared-bandwidth core and link, busiest first; empty unless " + "memory_ports is on", + ) class AllocationIR(BaseModel): @@ -490,6 +517,7 @@ def from_internal(cls, scheduler: SteadyStateScheduler) -> AllocationIR: overlap=OverlapIR(**perf_raw["overlap"]) if perf_raw.get("overlap") else None, tensor_reuse=[TensorReuseIR(**d) for d in perf_raw.get("tensor_reuse") or []], memory_occupancy=[MemoryOccupancyIR(**d) for d in perf_raw.get("memory_occupancy") or []], + memory_ports=[ResourceActivityIR(**d) for d in perf_raw.get("memory_ports") or []], ) if perf_raw else None diff --git a/stream/mapping/blocks.py b/stream/mapping/blocks.py index cbc33f109..f0a6536e6 100644 --- a/stream/mapping/blocks.py +++ b/stream/mapping/blocks.py @@ -21,7 +21,7 @@ def block_options(workload: Workload, mapping: Mapping, group: FusedGroup) -> di fixed: set[LayerDim] = set() for node, entry in _group_kernels(workload, mapping, group): dims = workload.get_dims(node) - for position, _, call_dim in entry.kernel.call_tile(): + for position, _, call_dim in entry.kernel.call_tile(node): dim = dims[position] if call_dim.blocks: options[dim] = options.get(dim, set(call_dim.blocks)) & set(call_dim.blocks) @@ -45,7 +45,7 @@ def with_block(workload: Workload, mapping: Mapping, gi: int, group: FusedGroup, dims = workload.get_dims(node) fields = { call_dim.name: size - for position, _, call_dim in entry.kernel.call_tile() + for position, _, call_dim in entry.kernel.call_tile(node) if dims[position] == dim and (call_dim.blocks or call_dim.divisor) } if fields: diff --git a/stream/mapping/chain_placement.py b/stream/mapping/chain_placement.py index 8c22cf989..bf1d75688 100644 --- a/stream/mapping/chain_placement.py +++ b/stream/mapping/chain_placement.py @@ -8,9 +8,9 @@ MIN_CALL_OPERANDS = 2 -def call_sizes(kernel) -> dict[int, int]: - """The size of one call along each node dimension the kernel library declares.""" - return {position: size for position, size, _ in kernel.call_tile()} +def call_sizes(kernel, node) -> dict[int, int]: + """The size of one call along each dimension of ``node`` the kernel library declares.""" + return {position: size for position, size, _ in kernel.call_tile(node)} def is_matmul(kernel) -> bool: @@ -32,7 +32,7 @@ def bandwidth_bound(kernel) -> bool: if is_matmul(kernel): return False operands = max(MIN_CALL_OPERANDS, len(kernel.operand_layouts() or ())) - move_cycles = operands * prod(call_sizes(kernel).values()) / DMA_ELEMENTS_PER_CYCLE + move_cycles = operands * prod(kernel.call_shape().values()) / DMA_ELEMENTS_PER_CYCLE return layer_cost(kernel) <= move_cycles diff --git a/stream/opt/allocation/constraint_optimization/allocation.py b/stream/opt/allocation/constraint_optimization/allocation.py deleted file mode 100644 index 4707b079d..000000000 --- a/stream/opt/allocation/constraint_optimization/allocation.py +++ /dev/null @@ -1,646 +0,0 @@ -""" -ComputeAllocator — a fully-modular MILP wrapper for node-to-core allocation -with *all* solver variables and constants exposed for inspection. - -Highlights ----------- -* Constants collected once and stored under ``self.const`` (see - :class:`ComputeAllocatorConstants`). -* Every solver variable handle is a public attribute—type-annotated precisely - (plain ``dict[tuple, SolverVar]``). -* Helper methods assert the variable dicts are non-``None`` **before** - subscripting, satisfying static type checkers. -* Public API remains: :py:meth:`ComputeAllocator.get_optimal_allocations`. -""" - -from __future__ import annotations - -import logging -from dataclasses import dataclass -from typing import TypeAlias, cast - -from zigzag.datatypes import LayerOperand, MemoryOperand - -from stream.cost_model.core_cost_lut import CoreCostLUT -from stream.hardware.architecture.accelerator import Accelerator -from stream.hardware.architecture.core import Core -from stream.opt.allocation.constraint_optimization.config import ConstraintOptStageConfig -from stream.opt.allocation.constraint_optimization.context import ConstraintContext, build_constraint_context -from stream.opt.allocation.constraint_optimization.utils import ( - convert_ids, - get_energies, - get_latencies, - invert_ids_list, -) -from stream.opt.solver import ( - SolverBackend, - SolverModel, - SolverParams, - SolverVar, - SolverVarType, - create_solver, -) -from stream.workload.workload import Workload - -logger = logging.getLogger(__name__) - -# --------------------------------------------------------------------------- # -# Type aliases # -# --------------------------------------------------------------------------- # -ALLOCATION_T: TypeAlias = list[tuple[int, int, tuple[int, int]]] -ALLOCATION_INTERNAL_T: TypeAlias = list[tuple[int, Core, int]] - -LatDict = dict[tuple[int, Core, int], int] -EnergyDict = dict[tuple[int, Core], float] -WeightDict = dict[int, int] -DepDict = dict[tuple[int, int], int] -CapDict = dict[Core, float] -GroupDict = dict[tuple[int, int], list[int]] # key = (node.id, node.group) -SplitDict = dict[int, dict[Core, dict[int, int]]] - - -# --------------------------------------------------------------------------- # -# Constants container # -# --------------------------------------------------------------------------- # -@dataclass(slots=True) -class ComputeAllocatorConstants: - """Immutable data needed to build the MILP model.""" - - lat: LatDict - energies: EnergyDict - weights: WeightDict - deps: DepDict - capacities: CapDict - groups: GroupDict - splits: SplitDict - node_count: int - cores: list[Core] - - -# --------------------------------------------------------------------------- # -# Allocator # -# --------------------------------------------------------------------------- # -class ComputeAllocator: - """Constraint-optimization wrapper exposing solver internals.""" - - # ------------------------------------------------------------------ # - # Construction # - # ------------------------------------------------------------------ # - def __init__( - self, - workload: Workload, - accelerator: Accelerator, - cost_lut: CoreCostLUT, - context: ConstraintContext, - *, - iterations: int = 1, - backend: str = "ORTOOLS_GSCIP", - ) -> None: - self.workload = workload - self.accelerator = accelerator - self.cost_lut = cost_lut - self.context = context - self.iterations = iterations - self.backend_str = backend - - # Solver model and main assignment tensor - self.model: SolverModel | None = None - self.asgn: dict[tuple[Core, int, int], SolverVar] | None = None - - # ---------------- variable handles (precise typing) ------------- # - self.k_vec: dict[tuple[int, int], SolverVar] | None = None - self.k_splits: dict[int, SolverVar] | None = None - self.core_asgn: dict[tuple[Core, int], SolverVar] | None = None - self.lat_id_core: dict[tuple[int, Core], SolverVar] | None = None - self.slot_asgn: dict[tuple[int, int], SolverVar] | None = None - self.slot_idx: dict[int, SolverVar] | None = None - self.w_split: dict[int, SolverVar] | None = None - self.w_core: dict[Core, SolverVar] | None = None - self.lat_core_slot: dict[tuple[Core, int], SolverVar] | None = None - self.lat_slot: dict[int, SolverVar] | None = None - self.lat_iter: SolverVar | None = None - self.idle_start: dict[tuple[Core, int], SolverVar] | None = None - self.idle_end: dict[tuple[Core, int], SolverVar] | None = None - self.idle_sum: dict[Core, SolverVar] | None = None - self.idle_min: SolverVar | None = None - self.total_lat: SolverVar | None = None - - # ---------------- sets populated during model build ------------- # - self.node_ids: list[int] = [] - self.cores: list[Core] = [] - self.slots: list[int] = [] - self.p_vals: list[int] = [] - - # Constants container (populated in _prepare_constants) - self.const: ComputeAllocatorConstants | None = None - - # ------------------------------------------------------------------ # - # Public API # - # ------------------------------------------------------------------ # - def get_optimal_allocations(self) -> ALLOCATION_T: - """Return sorted ``(slot, core_name, node_id)`` tuples.""" - self.const = self._prepare_constants() - self._build_constraint_model(self.const) - alloc = self._solve_model(self.const.node_count) - return invert_ids_list(alloc, self.const.node_count) # type: ignore - - # ------------------------------------------------------------------ # - # Prepare constants # - # ------------------------------------------------------------------ # - def _prepare_constants(self) -> ComputeAllocatorConstants: - cores, caps = self._collect_core_data() - nodes, ids = self._collect_node_data() - lat, splits = self._calculate_latencies(nodes, cores, ids) - en = self._calculate_energies(nodes, cores, ids) - deps = self._build_dependency_map(nodes, ids) - groups = self._group_nodes(nodes, ids) - weights = self._calculate_weights(groups, ids) - caps = self._relax_caps_if_single_node(nodes, caps) - - return ComputeAllocatorConstants( - lat=lat, - energies=en, - weights=weights, - deps=deps, - capacities=caps, - groups=groups, - splits=splits, - node_count=len(nodes), - cores=cores, - ) - - # ----------------- low-level helpers -------------------------------- # - def _collect_core_data(self) -> tuple[list[Core], CapDict]: - cores = sorted(self.context.compute_cores, key=lambda c: c.id) - if not cores: - raise ValueError("No eligible compute cores found for ComputeAllocator.") - caps = {core: self.context.capacities[core] for core in cores} - return cores, caps - - def _collect_node_data(self) -> tuple[list, dict]: - nodes = sorted(self.workload.node_list, key=lambda n: n.id) - return nodes, convert_ids(nodes) - - def _calculate_latencies( - self, - nodes, - cores: list[Core], - ids: dict, - ) -> tuple[LatDict, SplitDict]: - raw_lat, raw_split = get_latencies( - nodes, - [c.id for c in cores], - self.accelerator, - self.cost_lut, - impossible_lat=0, - ) - lat: LatDict = {(ids[n], c, k): v for (n, c, k), v in raw_lat.items()} - split: SplitDict = {ids[n]: {c: {k: raw_split[n][c][k] for k in raw_split[n][c]} for c in cores} for n in nodes} - return lat, split - - def _calculate_energies( - self, - nodes, - cores: list[Core], - ids: dict, - ) -> EnergyDict: - return get_energies(nodes, [c.id for c in cores], self.accelerator, self.cost_lut, 0, ids) - - def _build_dependency_map(self, nodes, ids: dict) -> DepDict: - out_op = LayerOperand("O") - return { - (ids[p], ids[c]): cast(int, p.operand_size_bit[out_op]) - for p, c in self.workload.edges() - if p in nodes and c in nodes - } - - @staticmethod - def _group_nodes(nodes, ids: dict) -> GroupDict: - groups: GroupDict = {} - for n in nodes: - key = (n.id, n.group) - groups.setdefault(key, []).append(ids[n]) - return groups - - @staticmethod - def _calculate_weights(groups: GroupDict, ids: dict) -> WeightDict: - weights: WeightDict = {} - for grp in groups.values(): - ref_sz = None - for idx, nid in enumerate(grp): - node = next(k for k, v in ids.items() if v == nid) - i2_ops = [op for op in node.constant_operands if node.memory_operand_links[op] == MemoryOperand("I2")] - bits = sum(node.operand_size_bit[op] for op in i2_ops) - if idx == 0: - ref_sz = bits - weights[nid] = bits - else: - assert bits == ref_sz, "Mismatch in grouped weight sizes" - weights[nid] = 0 - return weights - - @staticmethod - def _relax_caps_if_single_node( - nodes, - caps: CapDict, - ) -> CapDict: - if len(nodes) == 1: - return {c: 1e11 for c in caps} - return caps - - # ------------------------------------------------------------------ # - # Model construction # - # ------------------------------------------------------------------ # - def _build_constraint_model(self, const: ComputeAllocatorConstants) -> None: - self._create_basic_sets(const.lat, const.capacities) - self._create_variables() - self._add_k_split_constraints() - self._add_core_assignment_constraints(const.splits) - self._add_slot_assignment_constraints() - self._add_group_constraints(const.groups) - self._add_dependency_constraints(const.deps) - self._add_weight_constraints(const.weights, const.capacities) - self._add_latency_constraints(const.lat) - self._add_overlap_constraints() - self._set_objective() - - # -------------------- basic sets ----------------------------------- # - def _create_basic_sets(self, lat: LatDict, caps: CapDict) -> None: - node_core_k = list(lat.keys()) - self.node_ids = sorted({n for n, _, _ in node_core_k}) - self.cores = sorted(caps, key=lambda c: c.id) - self.p_vals = list(range(1, len(self.cores) + 1)) - self.slots = list(range(len(self.node_ids))) - - self.model = create_solver(SolverBackend[self.backend_str], "compute_alloc") - self.model.set_param(SolverParams.VERBOSITY, 1) - self.model.set_param(SolverParams.TIME_LIMIT, self.compute_cfg.time_limit) - self.model.set_param(SolverParams.THREADS, 0) - self.model.set_param(SolverParams.POOL_GAP, self.compute_cfg.gap) - - # -------------------- variable creation ---------------------------- # - def _create_variables(self) -> None: - assert self.model is not None - m = self.model - - # Replace tupledict creation with plain dict loops - self.k_vec = { - (n, k): m.add_var(vtype=SolverVarType.BINARY, name=f"k_vec[{n},{k}]") - for n in self.node_ids - for k in self.p_vals - } - self.k_splits = {n: m.add_var(vtype=SolverVarType.INTEGER, name=f"k_splits[{n}]") for n in self.node_ids} - self.core_asgn = { - (c, n): m.add_var(vtype=SolverVarType.BINARY, name=f"core_asgn[{c.id},{n}]") - for c in self.cores - for n in self.node_ids - } - self.lat_id_core = { - (n, c): m.add_var(vtype=SolverVarType.INTEGER, name=f"lat_id_core[{n},{c.id}]") - for n in self.node_ids - for c in self.cores - } - self.slot_asgn = { - (s, n): m.add_var(vtype=SolverVarType.BINARY, name=f"slot_asgn[{s},{n}]") - for s in self.slots - for n in self.node_ids - } - self.asgn = { - (c, s, n): m.add_var(vtype=SolverVarType.BINARY, name=f"asgn[{c.id},{s},{n}]") - for c in self.cores - for s in self.slots - for n in self.node_ids - } - self.slot_idx = {n: m.add_var(vtype=SolverVarType.INTEGER, name=f"slot_idx[{n}]") for n in self.node_ids} - self.w_split = {n: m.add_var(vtype=SolverVarType.INTEGER, name=f"w_split[{n}]") for n in self.node_ids} - self.w_core = {c: m.add_var(vtype=SolverVarType.INTEGER, name=f"w_core[{c.id}]") for c in self.cores} - self.lat_core_slot = { - (c, s): m.add_var(vtype=SolverVarType.INTEGER, name=f"lat_core_slot[{c.id},{s}]") - for c in self.cores - for s in self.slots - } - self.lat_slot = {s: m.add_var(vtype=SolverVarType.INTEGER, name=f"lat_slot[{s}]") for s in self.slots} - self.idle_start = { - (c, s): m.add_var(vtype=SolverVarType.BINARY, name=f"idle_start[{c.id},{s}]") - for c in self.cores - for s in self.slots - } - self.idle_end = { - (c, s): m.add_var(vtype=SolverVarType.BINARY, name=f"idle_end[{c.id},{s}]") - for c in self.cores - for s in self.slots - } - self.idle_sum = {c: m.add_var(vtype=SolverVarType.INTEGER, name=f"idle_sum[{c.id}]") for c in self.cores} - - # Scalars - self.lat_iter = m.add_var(vtype=SolverVarType.INTEGER, name="lat_iter") - self.idle_min = m.add_var(vtype=SolverVarType.INTEGER, name="idle_min") - self.total_lat = m.add_var(vtype=SolverVarType.INTEGER, name="total_lat") - - # -------------------- constraint layers ---------------------------- # - def _add_k_split_constraints(self) -> None: - assert self.model and self.k_vec and self.k_splits - m = self.model - for n in self.node_ids: - m.add_constr( - m.quicksum(self.k_vec[(n, k)]._raw for k in self.p_vals) == 1, - name=f"k_vec_one_hot_{n}", - ) - for n in self.node_ids: - m.add_constr( - self.k_splits[n]._raw == m.quicksum(self.k_vec[(n, k)]._raw * k for k in self.p_vals)._raw, - name=f"k_splits_def_{n}", - ) - - def _add_core_assignment_constraints(self, splits: SplitDict) -> None: - assert self.model and self.core_asgn and self.k_splits and self.k_vec - m = self.model - for n in self.node_ids: - m.add_constr( - m.quicksum(self.core_asgn[(c, n)]._raw for c in self.cores) == self.k_splits[n]._raw, - name=f"core_asgn_sum_{n}", - ) - for n in self.node_ids: - for c in self.cores: - if any(splits[n][c][p] > 0 for p in self.p_vals): - m.add_constr( - self.core_asgn[(c, n)]._raw - <= m.quicksum(splits[n][c][k] * self.k_vec[(n, k)]._raw for k in self.p_vals)._raw, - name=f"core_asgn_split_{c.id}_{n}", - ) - else: - # Explicitly zero the assignment - m.add_constr(self.core_asgn[(c, n)]._raw == 0, name=f"core_asgn_zero_{c.id}_{n}") - - def _add_slot_assignment_constraints(self) -> None: - assert self.model and self.slot_asgn and self.asgn and self.core_asgn and self.k_splits - m = self.model - for n in self.node_ids: - m.add_constr( - m.quicksum(self.slot_asgn[(s, n)]._raw for s in self.slots) == 1, - name=f"slot_asgn_one_hot_{n}", - ) - for c in self.cores: - for n in self.node_ids: - m.add_constr( - self.core_asgn[(c, n)]._raw == m.quicksum(self.asgn[(c, s, n)]._raw for s in self.slots)._raw, - name=f"core_asgn_asgn_{c.id}_{n}", - ) - for n in self.node_ids: - for s in self.slots: - m.add_constr( - self.slot_asgn[(s, n)]._raw * m.quicksum(self.asgn[(c, s, n)]._raw for c in self.cores)._raw - == self.slot_asgn[(s, n)]._raw * self.k_splits[n]._raw, - name=f"slot_asgn_ksplits_{s}_{n}", - ) - for c in self.cores: - for s in self.slots: - m.add_constr( - m.quicksum(self.asgn[(c, s, n)]._raw for n in self.node_ids) <= 1, - name=f"asgn_at_most_one_{c.id}_{s}", - ) - - def _add_group_constraints(self, groups: GroupDict) -> None: - assert self.model and self.asgn - m = self.model - for grp in groups.values(): - for n_i, n_j in zip(grp, grp[1:], strict=False): - for c in self.cores: - m.add_constr( - m.quicksum(self.asgn[(c, s, n_i)]._raw for s in self.slots) - == m.quicksum(self.asgn[(c, s, n_j)]._raw for s in self.slots), - name=f"group_{n_i}_{n_j}_{c.id}", - ) - - def _add_dependency_constraints(self, deps: DepDict) -> None: - assert self.model and self.slot_idx and self.slot_asgn - m = self.model - for n in self.node_ids: - m.add_constr( - self.slot_idx[n]._raw == m.quicksum(s * self.slot_asgn[(s, n)]._raw for s in self.slots)._raw, - name=f"slot_idx_def_{n}", - ) - for p, c in deps: - m.add_constr(self.slot_idx[c]._raw >= self.slot_idx[p]._raw + 1, name=f"dep_{p}_{c}") - - def _add_weight_constraints(self, weights: WeightDict, caps: CapDict) -> None: - assert self.model and self.w_split and self.k_splits and self.w_core and self.asgn - m = self.model - for n in self.node_ids: - m.add_constr( - self.w_split[n]._raw * self.k_splits[n]._raw >= weights[n], - name=f"w_split_{n}", - ) - for c in self.cores: - m.add_constr( - self.w_core[c]._raw - >= m.quicksum( - self.w_split[n]._raw * m.quicksum(self.asgn[(c, s, n)]._raw for s in self.slots)._raw - for n in self.node_ids - )._raw, - name=f"w_core_{c.id}", - ) - m.add_constr(self.w_core[c]._raw <= caps[c], name=f"w_cap_{c.id}") - - def _add_latency_constraints(self, lat: LatDict) -> None: - assert ( - self.model - and self.lat_id_core - and self.k_vec - and self.lat_core_slot - and self.asgn - and self.lat_slot - and self.lat_iter - ) - m = self.model - # Remove gp.multidict: lat is already a dict with (n, c, k) tuple keys - for n in self.node_ids: - for c in self.cores: - m.add_constr( - self.lat_id_core[(n, c)]._raw - == m.quicksum(self.k_vec[(n, k)]._raw * lat[(n, c, k)] for k in self.p_vals)._raw, - name=f"lat_id_{n}_{c.id}", - ) - for c in self.cores: - for s in self.slots: - m.add_constr( - self.lat_core_slot[(c, s)]._raw - == m.quicksum( - self.lat_id_core[(n, c)]._raw * self.asgn[(c, s, n)]._raw for n in self.node_ids - )._raw, - name=f"lat_core_slot_{c.id}_{s}", - ) - # lat_slot[s] = max(lat_core_slot[c, s] for c in cores) - # Since we MINIMIZE lat_iter (which sums lat_slot), the objective forces lat_slot down. - # Therefore: lat_slot[s] >= lat_core_slot[c, s] for all c is sufficient. - for s in self.slots: - for c in self.cores: - m.add_constr( - self.lat_slot[s]._raw >= self.lat_core_slot[(c, s)]._raw, - name=f"lat_slot_max_{s}_{c.id}", - ) - m.add_constr( - self.lat_iter._raw == m.quicksum(self.lat_slot[s]._raw for s in self.slots)._raw, - name="lat_iter_def", - ) - - def _add_overlap_constraints(self) -> None: - assert ( - self.model - and self.idle_start - and self.idle_end - and self.idle_sum - and self.idle_min - and self.asgn - and self.lat_slot - and self.total_lat - ), "Model and required variables must be initialized before adding overlap constraints" - m = self.model - incl: dict[tuple[Core, int], SolverVar] = { - (c, s): m.add_var(vtype=SolverVarType.INTEGER, name=f"incl[{c.id},{s}]") - for c in self.cores - for s in self.slots - } - excl: dict[tuple[Core, int], SolverVar] = { - (c, s): m.add_var(vtype=SolverVarType.INTEGER, name=f"excl[{c.id},{s}]") - for c in self.cores - for s in self.slots - } - - for c in self.cores: - for s in self.slots: - m.add_constr( - incl[(c, s)]._raw - == m.quicksum( - m.quicksum(self.asgn[(c, t, n)]._raw for n in self.node_ids)._raw for t in range(s + 1) - )._raw, - name=f"incl_{c.id}_{s}", - ) - for c in self.cores: - for s in self.slots: - m.add_constr( - excl[(c, s)]._raw - == m.quicksum( - m.quicksum(self.asgn[(c, t, n)]._raw for n in self.node_ids)._raw for t in range(s) - )._raw, - name=f"excl_{c.id}_{s}", - ) - - eps = 1e-4 - big_m = len(self.node_ids) + eps - tot_nodes_core: dict[Core, SolverVar] = { - c: m.add_var(vtype=SolverVarType.INTEGER, name=f"tot_nodes_core[{c.id}]") for c in self.cores - } - for c in self.cores: - m.add_constr( - tot_nodes_core[c]._raw - == m.quicksum(self.asgn[(c, s, n)]._raw for s in self.slots for n in self.node_ids)._raw, - name=f"tot_nodes_core_{c.id}", - ) - - for c in self.cores: - for s in self.slots: - m.add_constr( - 1 >= incl[(c, s)]._raw + eps - big_m * (1 - self.idle_start[(c, s)]._raw), - name=f"idle_start_ub_{c.id}_{s}", - ) - m.add_constr( - 1 <= incl[(c, s)]._raw + big_m * self.idle_start[(c, s)]._raw, - name=f"idle_start_lb_{c.id}_{s}", - ) - m.add_constr( - excl[(c, s)]._raw >= tot_nodes_core[c]._raw - 1 + eps - big_m * (1 - self.idle_end[(c, s)]._raw), - name=f"idle_end_ub_{c.id}_{s}", - ) - m.add_constr( - excl[(c, s)]._raw <= tot_nodes_core[c]._raw - 1 + big_m * self.idle_end[(c, s)]._raw, - name=f"idle_end_lb_{c.id}_{s}", - ) - for c in self.cores: - m.add_constr( - self.idle_sum[c]._raw - == m.quicksum( - (self.idle_start[(c, s)]._raw + self.idle_end[(c, s)]._raw) * self.lat_slot[s]._raw - for s in self.slots - )._raw, - name=f"idle_sum_{c.id}", - ) - - # idle_min = min(idle_sum[c] for c in cores) - # Use big-M binary selector pattern: - # idle_min <= idle_sum[c] for all c (upper-bound from below) - # idle_min >= idle_sum[c] - M*(1-b_c) for all c, with sum(b_c)==1 - big_m_min = len(self.node_ids) * max(1, len(self.slots)) # safe upper bound on any idle_sum - b_min: dict[Core, SolverVar] = { - c: m.add_var(vtype=SolverVarType.BINARY, name=f"b_min[{c.id}]") for c in self.cores - } - for c in self.cores: - m.add_constr(self.idle_min._raw <= self.idle_sum[c]._raw, name=f"idle_min_ub_{c.id}") - m.add_constr( - self.idle_min._raw >= self.idle_sum[c]._raw - big_m_min * (1 - b_min[c]._raw), - name=f"idle_min_lb_{c.id}", - ) - m.add_constr( - m.quicksum(b_min[c]._raw for c in self.cores) == 1, - name="idle_min_selector_one_hot", - ) - assert self.lat_iter is not None, "lat_iter must be set before this constraint" - m.add_constr( - self.total_lat._raw == self.iterations * self.lat_iter._raw - (self.iterations - 1) * self.idle_min._raw, - name="total_lat_def", - ) - - def _set_objective(self) -> None: - assert self.model and self.total_lat - self.model.set_objective(self.total_lat._raw, sense="minimize") - - # ------------------------------------------------------------------ # - # Solve & extract # - # ------------------------------------------------------------------ # - def _solve_model(self, node_cnt: int) -> ALLOCATION_INTERNAL_T: - assert self.model and self.asgn - self.model.optimize() - - if self.model.get_sol_count() == 0: - self.model.compute_iis() - self.model.write("infeasible.ilp") - raise ValueError("No feasible solution; see infeasible.ilp") - if self.model.get_status() not in {"OPTIMAL", "TIME_LIMIT"}: - raise RuntimeError(f"Solver status {self.model.get_status()}") - - alloc: ALLOCATION_INTERNAL_T = [ - (slot, core, nid) for (core, slot, nid), var in self.asgn.items() if round(var.X) == 1 - ] - return sorted(alloc, key=lambda t: (t[0], getattr(t[1], "id", t[1]), t[2])) - - -# --------------------------------------------------------------------------- # -# Functional façade # -# --------------------------------------------------------------------------- # -def get_optimal_allocations( - workload: Workload, - accelerator: Accelerator, - cost_lut: CoreCostLUT, - *, - context: ConstraintContext | None = None, - stage_config: ConstraintOptStageConfig | None = None, - iterations: int = 1, - backend: str = "ORTOOLS_GSCIP", -) -> ALLOCATION_T: - """Backwards-compatible helper preserving the original functional API.""" - if context is None: - logger.warning( - "get_optimal_allocations called without explicit context. " - "Building defaults; please pass ConstraintOptStageConfig explicitly." - ) - stage_cfg = stage_config or ConstraintOptStageConfig() - context = context or build_constraint_context(accelerator, stage_cfg) - - return ComputeAllocator( - workload, - accelerator, - cost_lut, - context, - iterations=iterations, - backend=backend, - ).get_optimal_allocations() diff --git a/stream/opt/allocation/constraint_optimization/config.py b/stream/opt/allocation/constraint_optimization/config.py index b6d5d7adf..57a6b3f07 100644 --- a/stream/opt/allocation/constraint_optimization/config.py +++ b/stream/opt/allocation/constraint_optimization/config.py @@ -1,71 +1,11 @@ from __future__ import annotations from dataclasses import dataclass, field -from enum import Enum, auto - -from zigzag.datatypes import MemoryOperand - - -class CoreRole(Enum): - COMPUTE = auto() - CACHE = auto() - IO = auto() - SHIM = auto() - - -@dataclass(frozen=True) -class DmaBudget: - mem_to_dram: int = 0 - dram_to_mem: int = 0 - shim_s2mm: int = 0 - shim_mm2s: int = 0 - - -@dataclass(frozen=True) -class BufferingConfig: - double_buffer_factor: int = 1 - - def __post_init__(self) -> None: - if self.double_buffer_factor < 1: - raise ValueError("double_buffer_factor must be >= 1") - - -@dataclass(frozen=True) -class IoRoutingRules: - force_via_cache: bool = False - max_cols_to_use: int | None = None # e.g., cap AIE column usage - - -@dataclass(frozen=True) -class ObjectiveWeights: - idle_weight: float = 0.0 - - -@dataclass(frozen=True) -class CoreConstraintProfile: - """ - Per-core-type constraint knobs. Mapped by core.type. - """ - - type_name: str - roles: set[CoreRole] - capacity_operand: MemoryOperand | None = None - dma_limits: DmaBudget = field(default_factory=DmaBudget) - buffering: BufferingConfig = field(default_factory=BufferingConfig) - io_routing: IoRoutingRules = field(default_factory=IoRoutingRules) - - def __post_init__(self) -> None: - if not self.roles: - raise ValueError(f"CoreConstraintProfile {self.type_name} must declare at least one role") @dataclass(frozen=True) class TransferMilpConfig: - # Currently unused; defaults mirror existing behavior. nb_cols_to_use: int = 4 - force_io_via_cache: bool = True - mem_dma_channels: int = 6 - shim_dma_channels: int = 2 def __post_init__(self) -> None: if self.nb_cols_to_use <= 0: @@ -75,50 +15,11 @@ def __post_init__(self) -> None: @dataclass(frozen=True) class ConstraintOptStageConfig: transfer: TransferMilpConfig = field(default_factory=TransferMilpConfig) - profiles: dict[str, CoreConstraintProfile] = field(default_factory=dict) - - def __post_init__(self) -> None: - if not self.profiles: - object.__setattr__(self, "profiles", default_core_profiles()) - elif len({p.type_name for p in self.profiles.values()}) != len(self.profiles): - raise ValueError("Profile type_names must be unique") @classmethod - def from_legacy_kwargs(cls, **kwargs) -> ConstraintOptStageConfig: - """ - Temporary adapter to convert legacy kwargs into a typed config. - """ + def from_kwargs(cls, **kwargs) -> ConstraintOptStageConfig: + """Build the config from the stage context, reading `nb_cols_to_use`.""" transfer_cfg = TransferMilpConfig( nb_cols_to_use=kwargs.get("nb_cols_to_use", 4), ) - profiles = kwargs.get("core_profiles", default_core_profiles()) - return cls(transfer=transfer_cfg, profiles=profiles) - - -def default_core_profiles() -> dict[str, CoreConstraintProfile]: - return { - "compute": CoreConstraintProfile( - type_name="compute", - roles={CoreRole.COMPUTE}, - capacity_operand=MemoryOperand("I2"), - ), - "memory": CoreConstraintProfile( - type_name="memory", - roles={CoreRole.CACHE}, - dma_limits=DmaBudget(mem_to_dram=6, dram_to_mem=6, shim_s2mm=2, shim_mm2s=2), - ), - "shim": CoreConstraintProfile( - type_name="shim", - roles={CoreRole.SHIM, CoreRole.IO}, - ), - } - - -def pick_profile(profiles: dict[str, CoreConstraintProfile], core_type: str) -> CoreConstraintProfile: - return profiles.get(core_type) or profiles.get("compute") or next(iter(profiles.values())) - - -def ensure_profile_registry(profiles: dict[str, CoreConstraintProfile] | None) -> dict[str, CoreConstraintProfile]: - if profiles is None or not profiles: - return default_core_profiles() - return profiles + return cls(transfer=transfer_cfg) diff --git a/stream/opt/allocation/constraint_optimization/context.py b/stream/opt/allocation/constraint_optimization/context.py index 6d3600dd5..5e5a5eeb0 100644 --- a/stream/opt/allocation/constraint_optimization/context.py +++ b/stream/opt/allocation/constraint_optimization/context.py @@ -7,13 +7,6 @@ from stream.hardware.architecture.accelerator import Accelerator from stream.hardware.architecture.core import Core -from stream.opt.allocation.constraint_optimization.config import ( - ConstraintOptStageConfig, - CoreConstraintProfile, - CoreRole, - ensure_profile_registry, - pick_profile, -) from stream.opt.allocation.constraint_optimization.timeslot_allocation import _resource_key from stream.plugins import load_group @@ -34,58 +27,6 @@ class MemoryReuseEntry: unexpressible: tuple[tuple[SolverVar, SolverVar], ...] -# ============================================================================ -# ConstraintContext – used by the *timeslot* allocation stage -# ============================================================================ - - -@dataclass(frozen=True) -class ConstraintContext: - core_profiles: dict[Core, CoreConstraintProfile] - compute_cores: list[Core] - cache_cores: list[Core] - io_cores: list[Core] - capacities: dict[Core, float] - offchip_core_id: int | None - - def core_ids(self) -> list[int]: - return sorted(c.id for c in self.compute_cores) - - -def build_constraint_context(accelerator: Accelerator, cfg: ConstraintOptStageConfig) -> ConstraintContext: - profiles_registry = ensure_profile_registry(cfg.profiles) - - core_profiles: dict[Core, CoreConstraintProfile] = {} - for core in accelerator.core_list: - core_profiles[core] = pick_profile(profiles_registry, core.type) - - def _cores_with(role: CoreRole) -> list[Core]: - return [c for c, p in core_profiles.items() if role in p.roles] - - def _eligible_for_compute(core: Core) -> bool: - profile = core_profiles[core] - return bool(profile.roles) - - compute_cores = [ - c - for c in core_profiles - if _eligible_for_compute(c) and (accelerator.offchip_core_id is None or c.id != accelerator.offchip_core_id) - ] - cache_cores = _cores_with(CoreRole.CACHE) - io_cores = _cores_with(CoreRole.IO) + _cores_with(CoreRole.SHIM) - - capacities: dict[Core, float] = {c: c.get_memory_capacity() for c in compute_cores} - - return ConstraintContext( - core_profiles=core_profiles, - compute_cores=compute_cores, - cache_cores=cache_cores, - io_cores=io_cores, - capacities=capacities, - offchip_core_id=accelerator.offchip_core_id, - ) - - # ============================================================================ # Namespace-specific MILP constraint strategies # ---------------------------------------------------------------------------- diff --git a/stream/opt/allocation/constraint_optimization/families/__init__.py b/stream/opt/allocation/constraint_optimization/families/__init__.py new file mode 100644 index 000000000..92f288f02 --- /dev/null +++ b/stream/opt/allocation/constraint_optimization/families/__init__.py @@ -0,0 +1,75 @@ +"""Constraint families: optional groups of constraints and quantities an allocator builds on request.""" + +from __future__ import annotations + +from collections.abc import Callable, Mapping, Sequence +from typing import TYPE_CHECKING, Any, ClassVar, Protocol, runtime_checkable + +from stream.plugins import load_group + +if TYPE_CHECKING: + from stream.opt.allocation.constraint_optimization.quantities import QuantityRegistry + from stream.opt.allocation.constraint_optimization.transfer_and_tensor_allocation import ( + TransferAndTensorAllocator, + ) + +FAMILY_GROUP = "stream.constraint_families" +SLOT_PRESSURE = "slot_pressure" +FamilySpec = str | Mapping[str, Any] + + +class ConstraintFamily(Protocol): + """``declare`` runs before the overlap is defined and ``constrain`` after it, both with the registry.""" + + name: ClassVar[str] + + def declare(self, alloc: TransferAndTensorAllocator, q: QuantityRegistry) -> None: ... + + def constrain(self, alloc: TransferAndTensorAllocator, q: QuantityRegistry) -> None: ... + + +@runtime_checkable +class ReportingFamily(Protocol): + """A family that adds sections to the solved schedule's performance report.""" + + def report(self, alloc: TransferAndTensorAllocator, q: QuantityRegistry) -> dict[str, Any]: ... + + +def available_families() -> dict[str, Callable[..., ConstraintFamily]]: + """Family factories by name from the ``stream.constraint_families`` entry points.""" + factories: dict[str, Callable[..., ConstraintFamily]] = {} + for plugin in load_group(FAMILY_GROUP): + if plugin.name in factories and factories[plugin.name] is not plugin.obj: + raise ValueError(f"Constraint family {plugin.name!r} is registered twice") + factories[plugin.name] = plugin.obj + return factories + + +def parse_spec(spec: FamilySpec) -> tuple[str, dict[str, Any]]: + """``"memory_ports"`` or ``{"memory_ports": {"burst": True}}`` to a name and its options.""" + if isinstance(spec, str): + return spec, {} + if len(spec) != 1: + raise ValueError(f"A family entry names exactly one family, got {sorted(spec)}") + ((name, options),) = spec.items() + return name, dict(options or {}) + + +def load_families(specs: Sequence[FamilySpec]) -> tuple[ConstraintFamily, ...]: + """Instantiate the families ``specs`` name, in order; an unknown name raises with the known ones.""" + if not specs: + return () + if isinstance(specs, str): + raise TypeError(f"Constraint families are a list of names, got the string {specs!r}") + known = available_families() + families: list[ConstraintFamily] = [] + for name, options in map(parse_spec, specs): + if name not in known: + raise KeyError(f"Unknown constraint family {name!r}; available: {', '.join(sorted(known)) or 'none'}") + if any(f.name == name for f in families): + raise ValueError(f"Constraint family {name!r} is selected twice") + try: + families.append(known[name](**options)) + except TypeError as exc: + raise TypeError(f"Constraint family {name!r}: {exc}") from None + return tuple(families) diff --git a/stream/opt/allocation/constraint_optimization/families/memory_ports.py b/stream/opt/allocation/constraint_optimization/families/memory_ports.py new file mode 100644 index 000000000..8e8bf93cc --- /dev/null +++ b/stream/opt/allocation/constraint_optimization/families/memory_ports.py @@ -0,0 +1,185 @@ +"""Memory-port bandwidth: what each top-level port moves must fit in the time the schedule gives it.""" + +from __future__ import annotations + +from collections import defaultdict +from math import ceil +from typing import TYPE_CHECKING, Any, ClassVar + +from stream.hardware.ports import PortKey +from stream.opt.allocation.constraint_optimization.families import SLOT_PRESSURE +from stream.opt.allocation.constraint_optimization.families.traffic import DmaStream, dma_streams, node_traffic + +if TYPE_CHECKING: + from stream.opt.allocation.constraint_optimization.quantities import QuantityRegistry + from stream.opt.allocation.constraint_optimization.transfer_and_tensor_allocation import ( + TransferAndTensorAllocator, + ) + +# (coefficient, expression, upper bound of the expression): one stream's or node's bits on a port +Terms = list[tuple[float, Any, float]] + + +def _port_name(key: PortKey) -> str: + return "_".join(str(part) for part in key) + + +class MemoryPorts: + """With ``interval``, the interval bound: a port's bits per iteration fit in its rate times the initiation + interval, for every port. With ``burst``, the burst bound: each slot's bits on a port fit in its rate times + that slot's latency. With neither, the model is unchanged and the family only reports. Rates in bits/cycle.""" + + name: ClassVar[str] = "memory_ports" + + def __init__(self, interval: bool = True, burst: bool = True) -> None: + self.interval = interval + self.burst = burst + + def declare(self, alloc: TransferAndTensorAllocator, q: QuantityRegistry) -> None: + for port in alloc.accelerator.ports: + q.add("port_rate", port.bits_per_cycle, index=port.key) + per_iteration: dict[PortKey, Terms] = defaultdict(list) + per_slot: dict[tuple[PortKey, int], Terms] = defaultdict(list) + for stream in dma_streams(alloc): + for side in stream.sides: + term = (stream.bits_per_cycle * side.share / side.efficiency, stream.active_cycles, stream.latency_ub) + per_iteration[side.port.key].append(term) + per_slot[(side.port.key, stream.slot)].append(term) + for traffic in node_traffic(alloc): + per_iteration[traffic.port.key].append((traffic.bits, 1, 1)) + per_slot[(traffic.port.key, traffic.slot)].append((traffic.bits, 1, 1)) + + for key, terms in per_iteration.items(): + q.add("port_demand", self._sum(alloc, terms), index=key) + for key, terms in per_slot.items(): + q.add("port_demand_slot", self._sum(alloc, terms), index=key) + if not (self.interval or self.burst): + return + for key, iteration_terms in per_iteration.items(): + rate = q.get("port_rate", key).expr + demands = [iteration_terms, *(terms for (k, _), terms in per_slot.items() if k == key)] + bound = max(sum(c * ub for c, _, ub in terms) for terms in demands) / rate + cycles = q.get("port_demand", key).expr / rate + # One cycle of margin against float rounding between this bound and the solved demand. + q.add(SLOT_PRESSURE, cycles, index=("memory_ports", key), upper_bound=ceil(bound) + 1) + + def constrain(self, alloc: TransferAndTensorAllocator, q: QuantityRegistry) -> None: + if "port_demand" not in q: + return + if self.interval: + interval = q.get("iteration").expr - q.get("overlap").expr + for key, demand in q.indexed("port_demand").items(): + rate = q.get("port_rate", key).expr + alloc.model.add_constr(rate * interval >= demand.expr, name=f"port_interval_{_port_name(key)}") + if self.burst: + for (key, slot), demand in q.indexed("port_demand_slot").items(): + alloc.model.add_constr( + q.get("port_rate", key).expr * q.get("slot_latency", slot).expr >= demand.expr, + name=f"port_burst_{_port_name(key)}_{slot}", + ) + + def report(self, alloc: TransferAndTensorAllocator, q: QuantityRegistry) -> dict[str, Any]: + """Per resource in ZigZag's port-activity terms, busiest first: ``real_cycle`` it needs for one iteration's + bits, ``allowed_cycle`` the initiation interval, and ``stall_or_slack`` their difference. Resources are + memory ports, measured shared bandwidth and links; the last two are read from the solved transfers.""" + interval = alloc.model.value(q.get("iteration").expr) - alloc.model.value(q.get("overlap").expr) + rows = [*self._port_rows(alloc, q, interval), *_shared_rows(alloc, q, interval), *_link_rows(alloc, interval)] + rows.sort(key=lambda row: -(row["utilization"] or 0.0)) + return {"memory_ports": rows} + + @staticmethod + def _port_rows(alloc: TransferAndTensorAllocator, q: QuantityRegistry, interval: float) -> list[dict[str, Any]]: + if "port_demand" not in q: + return [] + value = alloc.model.value + cores = {port.key: port.core_ids for port in alloc.accelerator.ports} + rows = [] + for key, demand in q.indexed("port_demand").items(): + rate = q.get("port_rate", key).expr + bursts = { + slot: value(d.expr) / (rate * alloc.slot_latency[slot].X) + for (k, slot), d in q.indexed("port_demand_slot").items() + if k == key and alloc.slot_latency[slot].X > 0 + } + busiest = max(bursts, key=bursts.__getitem__, default=None) + row = _activity_row( + alloc, "memory_port", f"{key.memory}.{key.port}", cores[key], rate, value(demand.expr), interval + ) + rows.append( + row | {"burst_utilization": bursts.get(busiest) if busiest is not None else None, "burst_slot": busiest} + ) + return rows + + @staticmethod + def _sum(alloc: TransferAndTensorAllocator, terms: Terms) -> Any: + return alloc.model.quicksum(coefficient * value for coefficient, value, _ in terms)._raw + + +def _activity_row( + alloc: TransferAndTensorAllocator, + kind: str, + name: str, + core_ids: tuple[int, ...], + rate: float, + bits: float, + interval: float, +) -> dict[str, Any]: + core_types = [alloc.accelerator.get_core(core_id).core_type for core_id in core_ids] + real_cycle = bits / rate + return { + "kind": kind, + "resource": name, + "core_ids": list(core_ids), + "core_types": core_types, + "bw_bits_per_cycle": rate, + "bits_per_iteration": bits, + "req_bw_aver": bits / interval if interval > 0 else None, + "real_cycle": real_cycle, + "allowed_cycle": interval, + "stall_or_slack": real_cycle - interval, + "utilization": real_cycle / interval if interval > 0 else None, + "burst_utilization": None, + "burst_slot": None, + } + + +def _stream_bits(alloc: TransferAndTensorAllocator) -> list[tuple[DmaStream, float]]: + return [(stream, alloc.model.value(stream.bits_per_cycle * stream.active_cycles)) for stream in dma_streams(alloc)] + + +def _shared_rows(alloc: TransferAndTensorAllocator, q: QuantityRegistry, interval: float) -> list[dict[str, Any]]: + """A core with a measured bandwidth: its solved busy time, which counts the access pattern's slow down.""" + rows = [] + streams = _stream_bits(alloc) + for core_id, model in alloc.shared_bandwidth.items(): + bits = sum(b for s, b in streams if any(c.id == core_id for c in (*s.choice.sources, *s.choice.targets))) + row = _activity_row(alloc, "shared_bandwidth", "measured", (core_id,), model.ceiling, bits, interval) + if ("shared_busy" in q) and core_id in q.indexed("shared_busy"): + busy = alloc.model.value(q.get("shared_busy", core_id).expr) + row |= { + "real_cycle": busy, + "stall_or_slack": busy - interval, + "utilization": busy / interval if interval > 0 else None, + } + rows.append(row) + return rows + + +def _link_rows(alloc: TransferAndTensorAllocator, interval: float) -> list[dict[str, Any]]: + bits: dict[Any, float] = defaultdict(float) + for stream, stream_bits in _stream_bits(alloc): + for link in stream.choice.links_used: + bits[link] += stream_bits + return [ + _activity_row(alloc, "link", _link_name(link), _link_cores(link), link.bandwidth, b, interval) + for link, b in bits.items() + if b > 0 + ] + + +def _link_name(link: Any) -> str: + return "->".join(str(getattr(end, "id", end)) for end in (link.sender, link.receiver)) + + +def _link_cores(link: Any) -> tuple[int, ...]: + return tuple(end.id for end in (link.sender, link.receiver) if hasattr(end, "id")) diff --git a/stream/opt/allocation/constraint_optimization/families/traffic.py b/stream/opt/allocation/constraint_optimization/families/traffic.py new file mode 100644 index 000000000..cfa678bdf --- /dev/null +++ b/stream/opt/allocation/constraint_optimization/families/traffic.py @@ -0,0 +1,141 @@ +"""The bits each DMA stream and each node moves through the top-level memory ports, per iteration.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any + +from stream.cost_model.bandwidth import contiguous_span_bytes +from stream.hardware.ports import OUTPUT, READ, WRITE, Port, input_role +from stream.opt.allocation.constraint_optimization.utils import active_fraction, get_active_latency +from stream.stages.estimation.core_cost_backends import CoreCostBackend, port_traffic, select_backend +from stream.workload.workload import ComputationNode, TransferNode + +if TYPE_CHECKING: + from stream.cost_model.communication_manager import MulticastPathPlan + from stream.hardware.architecture.core import Core + from stream.opt.allocation.constraint_optimization.transfer_and_tensor_allocation import ( + TransferAndTensorAllocator, + ) + + +@dataclass(frozen=True) +class PortShare: + """``share`` of a stream's bits pass ``port`` in ``direction`` (1/|sources| on a read, 1 on a write) at + ``efficiency`` of the port's rate for the stream's access pattern.""" + + port: Port + direction: str + share: float + efficiency: float + + +@dataclass(frozen=True) +class DmaStream: + """``bits_per_cycle * active_cycles`` is the stream's bits per iteration; ``active_cycles``, at most + ``latency_ub``, is the transfer's latency on its chosen path, or that path's choice when the latency rounds to 0.""" + + choice: MulticastPathPlan + active_cycles: Any + bits_per_cycle: float + latency_ub: float + slot: int + sides: tuple[PortShare, ...] + + +@dataclass(frozen=True) +class NodeTraffic: + """Bits per iteration a node moves through ``port`` of the core it runs on.""" + + port: Port + bits: float + slot: int + + +def transfer_operand_role(alloc: TransferAndTensorAllocator, tr: TransferNode, read_side: bool) -> str: + """Operand role the tensor has on one side: the output for a producer, input k for a consumer's input k.""" + producer = next(iter(alloc.workload.predecessors(tr)), None) + if read_side and isinstance(producer, ComputationNode): + return OUTPUT + for consumer in alloc.workload.successors(tr): + if isinstance(consumer, ComputationNode): + for tensor in tr.outputs: + if tensor in consumer.inputs: + return input_role(consumer.inputs.index(tensor) + 1) + return OUTPUT + + +def _span_bytes(tr: TransferNode) -> float: + tensor = tr.inputs[0] + full = tuple(tensor.subview.source.type.get_shape()) + return contiguous_span_bytes(tuple(tensor.shape), full, tensor.operand_type.bitwidth) + + +def _target_share(alloc: TransferAndTensorAllocator, tr: TransferNode) -> float: + """Share of the transferred tensor one target receives: its tile under the consumer's inter-core tiling, + as memory access estimation counts it. A broadcast gives every target the whole tensor.""" + tensor = tr.outputs[0] + tile = alloc.workload.get_tensor_of_transfer_to_single_core(tensor, tr, alloc.mapping) + return tile.size_bits() / tensor.size_bits() + + +def _sides(alloc: TransferAndTensorAllocator, tr: TransferNode, choice: Any) -> list[PortShare]: + ports = alloc.accelerator.ports + sides: list[PortShare] = [] + span = _span_bytes(tr) + write_share = _target_share(alloc, tr) + for cores, direction, share in ( + (choice.sources, READ, 1.0 / len(choice.sources)), + (choice.targets, WRITE, write_share), + ): + operand = transfer_operand_role(alloc, tr, read_side=direction == READ) + for core in cores: + port = ports.port_for(core, direction, operand) + if port is None: + continue + # Targets in one core_memory_sharing group receive a broadcast once, into their shared memory. + broadcast_again = any(s.port.key == port.key and s.direction == WRITE for s in sides) + if direction == WRITE and write_share == 1.0 and broadcast_again: + continue + sides.append(PortShare(port, direction, share, port.bandwidth.efficiency(span, direction))) + return sides + + +def dma_streams(alloc: TransferAndTensorAllocator) -> list[DmaStream]: + """Every transfer latency on a path choice that moves bits, with the ports its sources read and targets write.""" + streams: list[DmaStream] = [] + if "transfer_latency" not in alloc.quantities: + return [] + for (tr, choice), quantity in alloc.quantities.indexed("transfer_latency").items(): + link_latency = float(alloc.transfer_latency_for_path(tr, choice)) + if link_latency <= 0: + continue + latency = get_active_latency(tr, link_latency, alloc.ssis) + bits = tr.inputs[0].size_bits() * active_fraction(tr, alloc.ssis) + sides = tuple(_sides(alloc, tr, choice)) + if latency > 0: + stream = DmaStream(choice, quantity.expr, bits / latency, latency, alloc.slot_of[tr], sides) + else: + # Charged at a reuse factor of 1, an upper bound, since y / R needs a variable of its own. + gate = alloc.y_path_choice[(tr, choice)]._raw + stream = DmaStream(choice, gate, bits, 1.0, alloc.slot_of[tr], sides) + streams.append(stream) + return streams + + +def node_traffic(alloc: TransferAndTensorAllocator) -> list[NodeTraffic]: + """The top-level port traffic of each node from its cost backend, on the physical core it runs on.""" + traffic: list[NodeTraffic] = [] + backends: dict[Core, CoreCostBackend] = {} + for node in alloc.ssc_nodes: + placed = {core for group in alloc.mapping.get(node).resource_allocation for core in group} + fraction = active_fraction(node, alloc.ssis) + for core in alloc.cost_lut.get_cores(node): + if core not in placed: + continue + backend = backends.setdefault(core, select_backend(core)) + for operand, direction, bits in port_traffic(backend, alloc.cost_lut.get_cost(node, core)): + port = alloc.accelerator.ports.port_for(core, direction, operand) + if port is not None: + traffic.append(NodeTraffic(port, bits * fraction, alloc.slot_of[node])) + return traffic diff --git a/stream/opt/allocation/constraint_optimization/quantities.py b/stream/opt/allocation/constraint_optimization/quantities.py new file mode 100644 index 000000000..bdd2f6f2c --- /dev/null +++ b/stream/opt/allocation/constraint_optimization/quantities.py @@ -0,0 +1,56 @@ +"""Named model expressions an allocator exposes to objectives, bounds and constraint families.""" + +from __future__ import annotations + +from collections.abc import Hashable +from dataclasses import dataclass +from typing import Any + + +@dataclass +class Quantity: + """One model expression; ``upper_bound`` is a value the expression provably never exceeds.""" + + expr: Any + upper_bound: float | None = None + + +class QuantityRegistry: + """Quantities by name, either scalar or indexed by a hashable key such as a slot or a core.""" + + def __init__(self) -> None: + self._scalars: dict[str, Quantity] = {} + self._indexed: dict[str, dict[Hashable, Quantity]] = {} + + def add(self, name: str, expr: Any, *, index: Hashable | None = None, upper_bound: float | None = None) -> None: + quantity = Quantity(expr, upper_bound) + if index is None: + if name in self._scalars or name in self._indexed: + raise ValueError(f"Quantity {name!r} is already registered") + self._scalars[name] = quantity + return + if name in self._scalars: + raise ValueError(f"Quantity {name!r} is already registered as a scalar") + entries = self._indexed.setdefault(name, {}) + if index in entries: + raise ValueError(f"Quantity {name!r}[{index!r}] is already registered") + entries[index] = quantity + + def get(self, name: str, index: Hashable | None = None) -> Quantity: + if index is None and name in self._scalars: + return self._scalars[name] + if index is not None and index in self._indexed.get(name, {}): + return self._indexed[name][index] + label = name if index is None else f"{name}[{index!r}]" + raise KeyError(f"Unknown quantity {label}; available: {', '.join(self.names())}") + + def indexed(self, name: str) -> dict[Hashable, Quantity]: + if name not in self._indexed: + raise KeyError(f"Unknown indexed quantity {name!r}; available: {', '.join(self.names())}") + return dict(self._indexed[name]) + + def names(self) -> list[str]: + return sorted({*self._scalars, *self._indexed}) + + def __contains__(self, name: str) -> bool: + return name in self._scalars or name in self._indexed diff --git a/stream/opt/allocation/constraint_optimization/timeslot_allocation.py b/stream/opt/allocation/constraint_optimization/timeslot_allocation.py index 895fc8d28..62680e20f 100644 --- a/stream/opt/allocation/constraint_optimization/timeslot_allocation.py +++ b/stream/opt/allocation/constraint_optimization/timeslot_allocation.py @@ -112,43 +112,6 @@ def _add_alloc( self._node_types[node] = node_type # ................................................. public slot helpers - def add_node_to_next_slot( - self, - node: Node, - res: Resource, - *, - min_slot: int = 0, - node_type: NodeType = NodeType.STEADY_STATE, - ) -> int: - """Place *node* on *res* at the first free slot ≥ *min_slot*.""" - next_slot = max(self._res_max_slot.get(res, -1) + 1, min_slot) - while next_slot in self._slot_res_to_node and res in self._slot_res_to_node[next_slot]: - next_slot += 1 - self._add_alloc(next_slot, res, node, node_type) - return next_slot - - def get_allocations_in_slot(self, slot: int) -> dict[Resource, Node]: - return dict(self._slot_res_to_node.get(slot, {})) - - def get_resources_for_node(self, node: Node) -> set[Resource]: - return set(self._node_to_res.get(node, set())) - - def get_resources_for_node_id(self, node_id: int) -> set[Resource]: - nodes = [n for n in self._node_to_res if n.id == node_id] - return set(res for n in nodes for res in self._node_to_res[n]) - - def get_timeslot_of_node_on_resource(self, node: Node, res: Resource) -> int: - for slot, res_map in self._slot_res_to_node.items(): - if res_map.get(res) is node: - return slot - raise ValueError(f"{node.name} not on {_resource_key(res)}.") - - def get_timeslot_of_node(self, node: Node) -> int: - slots = [slot for slot, res_map in self._slot_res_to_node.items() if node in res_map.values()] - if not slots: - raise ValueError(f"{node.name} not scheduled.") - return max(slots) - def get_timeslot_runtime(self, slot: int) -> int: if slot < self.slot_min or slot > self.slot_max: raise ValueError(f"Slot {slot} is out of range ({self.slot_min} - {self.slot_max}).") diff --git a/stream/opt/allocation/constraint_optimization/transfer_and_tensor_allocation.py b/stream/opt/allocation/constraint_optimization/transfer_and_tensor_allocation.py index 143db71dc..be6bb8b4f 100644 --- a/stream/opt/allocation/constraint_optimization/transfer_and_tensor_allocation.py +++ b/stream/opt/allocation/constraint_optimization/transfer_and_tensor_allocation.py @@ -3,6 +3,7 @@ import os import re from collections import defaultdict +from functools import cached_property from math import ceil, prod from typing import Any, TypeAlias @@ -10,8 +11,8 @@ import yaml from xdsl.ir.affine import AffineDimExpr -# GRB.Callback constants — used only by _mip_progress_callback (Gurobi-specific). -# gurobipy is optional (Gurobi backend only); GRB is touched solely on the Gurobi solve path. +# GRB supplies the callback codes of _mip_progress_callback and the Gurobi status names of the solve +# summary; gurobipy is optional, so GRB is None without it and only the Gurobi solve path touches it. try: from gurobipy import GRB except ModuleNotFoundError: @@ -39,10 +40,13 @@ TransferAndTensorContext, build_transfer_context, ) +from stream.opt.allocation.constraint_optimization.families import SLOT_PRESSURE, ReportingFamily, load_families +from stream.opt.allocation.constraint_optimization.quantities import QuantityRegistry from stream.opt.allocation.constraint_optimization.timeslot_allocation import ( _resource_key, ) from stream.opt.allocation.constraint_optimization.utils import ( + active_fraction, get_active_latency, get_transfer_latency_for_path, ) @@ -61,7 +65,6 @@ from stream.workload.node import HasOutputs, TransferType from stream.workload.steady_state.iteration_space import ( IterationVariableType, - LoopEffect, Reuse, SteadyStateIterationSpace, ) @@ -205,7 +208,6 @@ def __init__( # noqa: PLR0913, PLR0915 # auxiliary indicators self._handover_bits: dict[int, int] = {} self.transfer_side_indicator: dict[tuple[TransferNode, Core, bool], SolverVar] = {} - self.transfer_core_indicator: dict[tuple[TransferNode, Core], SolverVar] = {} self.tensor_core_indicator: dict[tuple[Tensor, Core], SolverVar] = {} self.same_core_indicator: dict[tuple[Tensor, Tensor, Core], SolverVar] = {} @@ -222,8 +224,11 @@ def __init__( # noqa: PLR0913, PLR0915 self._shared_latency_cache: dict[tuple[int, TransferNode, MulticastPathPlan], SolverVar] = {} self.slot_latency: dict[int, SolverVar] = {} self.overlap: SolverVar | None = None + self.fill: SolverVar | None = None self.total_latency: SolverVar | None = None self.recurrence_bound: int = 0 + self.quantities = QuantityRegistry() + self.families = load_families(self.constraint_selection.families) # transfer fire helpers init self._ensure_same_ssis_for_all_transfers() @@ -331,11 +336,8 @@ def _normalize_path_choices(self, raw: Any) -> tuple[MulticastPathPlan, ...]: # ------------------------------------------------------------ # # internal helpers # # ------------------------------------------------------------ # - @staticmethod - def _mem_factor(t: Tensor, core: Core) -> int: - return 1 - - def _transfer_latency_for_path(self, tr: TransferNode, path: MulticastPathPlan) -> int: + def transfer_latency_for_path(self, tr: TransferNode, path: MulticastPathPlan) -> int: + """Cycles one firing of ``tr`` takes on ``path``, before its active fraction and reuse.""" # A transfer served out of memory the two cores share reads in place: no bytes cross a link, # so it adds no time to the slot, the same reason it spends no DMA channel. if self._choice_shares_memory(tr, path): @@ -719,27 +721,6 @@ def _tensor_on_core_expr(self, t: Tensor, core: Core): self.x_tensor_choice[(t, choice)]._raw for choice in self._tensor_choices(t) if core in choice ) - def _transfer_uses_core_var(self, tr: TransferNode, core: Core) -> SolverVar: - key = (tr, core) - if key in self.transfer_core_indicator: - return self.transfer_core_indicator[key] - - u = self.model.add_var(vtype=SolverVarType.BINARY, name=f"u_{tr.name}_{_resource_key(core)}") - self.transfer_core_indicator[key] = u - - occ_exprs = [self._tensor_on_core_expr(t, core) for t in tr.tensors if isinstance(t, Tensor)] - if not occ_exprs: - self.model.add_constr(u == 0, name=f"u_zero_{tr.name}_{_resource_key(core)}") - return u - - for i, occ in enumerate(occ_exprs): - self.model.add_constr(u >= occ, name=f"u_lb_{tr.name}_{_resource_key(core)}_{i}") - self.model.add_constr( - u <= self.model.quicksum(occ_exprs), - name=f"u_ub_{tr.name}_{_resource_key(core)}", - ) - return u - def _transfer_side_uses_core_var(self, tr: TransferNode, core: Core, incoming: bool) -> SolverVar: """Whether any tensor this transfer brings to (or takes from) this core sits on it. @@ -830,13 +811,6 @@ def _same_core_var(self, src_tensor: Tensor, dst_tensor: Tensor, core: Core) -> ) return v - def _all_candidate_cores_for_transfer(self, tr: TransferNode) -> set[Core]: - out: set[Core] = set() - for t in tr.tensors: - if isinstance(t, Tensor): - out.update(self._candidate_cores_for_tensor(t)) - return out - # ------------------------------------------------------------ # # model construction # # ------------------------------------------------------------ # @@ -845,21 +819,29 @@ def _build_model(self): self._create_vars() self._index_choice_metadata() self._create_constraints() - self._overlap_and_objective() + for family in self.families: + family.declare(self, self.quantities) + self._overlap() + for family in self.families: + family.constrain(self, self.quantities) + self._dma_channels_and_objective() # ...................... VARIABLES ................... # def _create_vars(self): self.__create_tensor_placement_vars() self.__create_transfer_path_vars() self.__create_reuse_vars() + self.__create_single_buffer_vars() self.__create_slot_latency_vars() def __create_slot_latency_vars(self): for s in range(self.max_slot + 1): self.slot_latency[s] = self.model.add_var(vtype=SolverVarType.INTEGER, name=f"L_{s}") + self.quantities.add("slot_latency", self.slot_latency[s]._raw, index=s) def __create_reuse_vars(self): self.z_stop: dict[tuple[Tensor, int], SolverVar] = {} + self.z_single: dict[tuple[Tensor, int], SolverVar] = {} for t in self.workload.tensors: sizes = self.ssis[t].get_applicable_temporal_sizes() for stop in range(-1, len(sizes)): @@ -885,6 +867,31 @@ def __create_reuse_vars(self): name=f"zStop_AtLeast_{t.name}_L{stop}", ) + @cached_property + def _sole_outputs(self) -> set[Tensor]: + """The tensors a transfer moves to a single reader.""" + return {tr.outputs[0] for tr in self.transfer_nodes if len(tr.outputs) == 1} + + def _may_single_buffer(self, t, stop: int) -> bool: + """Whether a tensor a transfer moves to one reader may hold its one-tile window at ``stop`` in + a single buffer where an outer loop moves it on: the next window then waits for the last + read of this one.""" + return ( + stop >= 0 + and t in self.tensors_to_optimize_reuse_for + and t in self._sole_outputs + and self.rotation_levels[(t, stop)] + and self.tiles_needed_levels[(t, stop)] == 1 + ) + + def __create_single_buffer_vars(self): + for t in self.tensors_to_optimize_reuse_for: + for stop in range(len(self.ssis[t].get_applicable_temporal_variables())): + if self._may_single_buffer(t, stop): + v = self.model.add_var(vtype=SolverVarType.BINARY, name=f"zSingle_{t.name}_L{stop}") + self.model.add_constr(v <= self.z_stop[(t, stop)], name=f"zSingle_AtStop_{t.name}_L{stop}") + self.z_single[(t, stop)] = v + def __create_transfer_path_vars(self): for tr in self.transfer_nodes: for i, choice in enumerate(self.possible_transfer_allocations[tr]): @@ -998,6 +1005,21 @@ def _add_destination_tensor_coherence_constraints( y <= self._tensor_on_core_expr(dst_tensor, dst_core), name=f"path_dst_match_{tr.name}_{dst_tensor.name}_{_resource_key(dst_core)}_choice_{i}", ) + # And the tensor sits only where the chosen path delivers it: a core holding a copy + # nothing wrote cannot be where a later transfer reads it from. An empty path moves + # nothing, so the colocation constraints place its tensor. + if dst_tensor in self.tensor_fixed: + continue + for core in self._candidate_cores_for_tensor(dst_tensor): + delivering = [ + self.y_path_choice[(tr, choice)]._raw + for choice in choices + if core in self.choice_dst_cores[(tr, choice)] or self.choice_has_empty_path[(tr, choice)] + ] + self.model.add_constr( + self._tensor_on_core_expr(dst_tensor, core) <= self.model.quicksum(delivering), + name=f"dst_delivered_{tr.name}_{dst_tensor.name}_{_resource_key(core)}", + ) def _add_empty_path_coherence_constraints(self, tr: TransferNode, choices: tuple[MulticastPathPlan, ...]) -> None: """ @@ -1046,15 +1068,18 @@ def _link_contention_constraints(self): ) # ...................... memory capacity .................... # - def _resident_tiles(self, t, stop: int) -> int: - """Tiles of this tensor codegen keeps resident when it stops reuse at ``stop``.""" + def _resident_tiles(self, t, stop: int, single: bool = False) -> int: + """Tiles of this tensor codegen keeps resident when it stops reuse at ``stop``, held in one + buffer where ``single``.""" tiles = self.tiles_needed_levels[(t, stop)] - return max(tiles, 2) if self.rotation_levels.get((t, stop)) else tiles + return max(tiles, 2) if self.rotation_levels.get((t, stop)) and not single else tiles def _min_resident_bits(self, t, tensor_size: int) -> int: - """The least this tensor can keep resident under any reuse-stop choice.""" + """The least this tensor can keep resident under any reuse-stop and buffering choice.""" stops = range(-1, len(self.ssis[t].get_applicable_temporal_variables())) - return min(ceil(self._resident_tiles(t, stop) * tensor_size) for stop in stops) + return min( + ceil(self._resident_tiles(t, stop, self._may_single_buffer(t, stop)) * tensor_size) for stop in stops + ) def _capacity_screen(self): """Fail before building the model when a memory cannot fit its pinned tensors under any reuse choice.""" @@ -1125,13 +1150,23 @@ def _memory_capacity_constraints(self): min_req: int | None = None for stop in range(-1, len(self.ssis[t].get_applicable_temporal_variables())): req_size = ceil(self._resident_tiles(t, stop) * tensor_size) - min_req = req_size if min_req is None else min(min_req, req_size) + single = (t, stop) in self.z_single + least = ceil(self._resident_tiles(t, stop, single) * tensor_size) + min_req = least if min_req is None else min(min_req, least) uz = self._add_binary_product( a=u, b=self.z_stop[(t, stop)], base_name=f"memload_{t.name}_{_resource_key(holder)}_L{stop}", ) held[(t, memory)].append((uz, req_size, t.name)) + if single: + # The second buffer the window would rotate through is not held. + us = self._add_binary_product( + a=u, + b=self.z_single[(t, stop)], + base_name=f"memsingle_{t.name}_{_resource_key(holder)}_L{stop}", + ) + held[(t, memory)].append((us, least - req_size, t.name)) if min_req is not None: # bytes this tensor's tile adds if resident on holder least_bytes[memory] += min_req / 8 for memory, value in least_bytes.items(): @@ -1422,13 +1457,55 @@ def _force_final_output_reuse_levels(self): ) # ...................... overlap + objective ................. # - def _overlap_and_objective(self) -> None: - max_s = self.max_slot - big_m = self.big_m - - self._init_idle_indicators(max_s, big_m) - self._create_idle_latency_vars(max_s) + def _overlap(self) -> None: + self._init_idle_indicators(self.max_slot, self.big_m) + self._create_idle_latency_vars(self.max_slot) self._define_overlap_var() + self._resident_fill() + + def _resident_fill(self) -> None: + """Cycles each run waits for the off-chip windows it holds in one buffer to fill: such a window fills + before the iteration that reads it, overlapping none. The fills share each path and shared-bandwidth core + as one iteration's transfers do.""" + self.fill = fill = self.model.add_var(vtype=SolverVarType.CONTINUOUS, lb=0.0, name="resident_fill") + shared: dict[int, list[Any]] = defaultdict(list) + singles: dict[Tensor, list[tuple[int, Any]]] = defaultdict(list) + for (t, s), single in self.z_single.items(): + singles[t].append((s, single)) + for (tr, choice), y in self.y_path_choice.items(): + t = tr.outputs[0] + if t not in self.tensors_to_optimize_reuse_for: + continue + sizes = self.ssis[t].get_applicable_temporal_sizes() + # (tiles a run waits for, the choice that makes it wait, its name) + whole = ( + [ + (self.tiles_needed_levels[(t, s)], self.z_stop[(t, s)], f"L{s}") + for s in range(len(sizes)) + if not self.rotation_levels[(t, s)] + ] + if self._is_const_i(tr) + else [] + ) + choices = whole + [(prod(sizes[s + 1 :]), single, f"single_L{s}") for s, single in singles[t]] + held = [ + (tiles, self._add_binary_product(a=y, b=z, base_name=f"fill_{tr.name}_{hash(choice)}_{name}")) + for tiles, z, name in choices + ] + if not held: + continue + cycles = self.transfer_latency_for_path(tr, choice) + self.model.add_constr( + fill >= self.model.quicksum(tiles * cycles * w._raw for tiles, w in held), + name=f"fill_{tr.name}_{hash(choice)}", + ) + for core_id, model in self.shared_bandwidth.items(): + share = self._shared_cycles(core_id, tr, choice, model.ceiling) + shared[core_id] += [tiles * share * w._raw for tiles, w in held] + for core_id, terms in shared.items(): + self.model.add_constr(fill >= self.model.quicksum(terms), name=f"fill_shared_{core_id}") + + def _dma_channels_and_objective(self) -> None: if self.constraint_selection.dma_channels: self._add_dma_usage_constraints() else: @@ -1553,8 +1630,9 @@ def _create_idle_latency_vars(self, max_s: int) -> None: for tr in self.transfer_nodes: for choice in self._path_choices(tr): - lat = ceil(self._transfer_latency_for_path(tr, choice)) + lat = ceil(self.transfer_latency_for_path(tr, choice)) slot_latency_ub = max(slot_latency_ub, lat) + slot_latency_ub = max(slot_latency_ub, self._family_slot_pressure_bound()) for res in {r for r, _ in self.idle_ind}: terms = [] @@ -1575,10 +1653,21 @@ def _create_idle_latency_vars(self, max_s: int) -> None: ) self.idle_lat[res] = v + def _family_slot_pressure_bound(self) -> int: + """The largest slot latency a family's constraints can force, 0 when no family declares one.""" + if SLOT_PRESSURE not in self.quantities: + return 0 + pressures = self.quantities.indexed(SLOT_PRESSURE) + if unbounded := [index for index, q in pressures.items() if q.upper_bound is None]: + raise ValueError(f"{SLOT_PRESSURE} quantities need an upper_bound: {unbounded}") + return ceil(max(q.upper_bound for q in pressures.values())) + def _define_overlap_var(self) -> None: overlap = self.model.add_var(vtype=SolverVarType.INTEGER, name="overlap") self.overlap = overlap iteration = self.model.quicksum(v._raw for v in self.slot_latency.values()) + self.quantities.add("overlap", overlap._raw) + self.quantities.add("iteration", iteration) for res, v in self.idle_lat.items(): if not self._bounds_overlap(res): continue @@ -1665,9 +1754,6 @@ def _recurrence_bound(self) -> int: for n in carriers ) - def _transfer_dma_usage_expr(self, tr: TransferNode, core: Core): - return self.model.quicksum(self._tensor_on_core_expr(t, core) for t in tr.tensors if isinstance(t, Tensor)) - def _add_dma_usage_constraints(self) -> None: """ Directional DMA accounting for all on-chip cores. @@ -1743,42 +1829,43 @@ def _add_dma_usage_constraints(self) -> None: ) def _set_total_latency_and_objective(self) -> None: - self.total_lat = self.model.add_var(vtype=SolverVarType.INTEGER, name="total_latency") - self.total_latency = self.total_lat + total_latency = self.model.add_var(vtype=SolverVarType.INTEGER, name="total_latency") + self.total_latency = total_latency assert self.overlap is not None, "Overlap variable must be initialized before objective." self.model.add_constr( - self.total_lat + total_latency == self.iterations * self.model.quicksum(v._raw for v in self.slot_latency.values()) - (self.iterations - 1) * self.overlap + + self.fill ) # Primary objective: minimize total latency (+ DMA balancing if enabled) if self.constraint_selection.dma_channels: - primary_expr = self.total_lat._raw + self.max_core_dma_in._raw + self.max_core_dma_out._raw + primary_expr = total_latency._raw + self.max_core_dma_in._raw + self.max_core_dma_out._raw else: - primary_expr = self.total_lat._raw + primary_expr = total_latency._raw # Slot latency only sees a transfer when it is the longest thing in its slot, # so a transfer hiding behind compute is free to the primary objective however # often it fires. Offchip bandwidth is shared by every slot, so it is not free # on hardware. - self._traffic_terms = [ + traffic_terms = [ (t.size_bits() / self.reuse_levels[(t, s)], self.z_stop[(t, s)]) for t in self.tensors_to_optimize_reuse_for for s in range(-1, len(self.ssis[t].get_applicable_temporal_variables())) ] - traffic_expr = self.model.quicksum(bits * z._raw for bits, z in self._traffic_terms) + traffic_expr = self.model.quicksum(bits * z._raw for bits, z in traffic_terms) - self._traffic_weight = 0.0 + traffic_weight = 0.0 if ( not self.shared_bandwidth and self.constraint_selection.offchip_traffic_cost and (bw := self._offchip_bandwidth()) ): - self._traffic_weight = self.iterations / bw - primary_expr = primary_expr + self._traffic_weight * traffic_expr + traffic_weight = self.iterations / bw + primary_expr = primary_expr + traffic_weight * traffic_expr - # Third objective (tiebreaker): minimize total buffering depth + # Buffering depth, minimised once latency and offchip traffic are settled. buffering_expr = self.model.quicksum( self.tiles_needed_levels[(t, s)] * self.z_stop[(t, s)]._raw for t in self.tensors_to_optimize_reuse_for @@ -1791,13 +1878,15 @@ def _set_total_latency_and_objective(self) -> None: for choice in self.possible_transfer_allocations[tr] ) + self.quantities.add("primary", primary_expr) + levels = [ + ("latency", primary_expr), + ("offchip_traffic", traffic_expr), + ("buffering", buffering_expr), + ("route_hops", hops_expr), + ] self.model.set_lexicographic_objectives( - [ - ObjectiveLevel(expr=primary_expr, priority=4, name="latency"), - ObjectiveLevel(expr=traffic_expr, priority=3, name="offchip_traffic"), - ObjectiveLevel(expr=buffering_expr, priority=2, name="buffering"), - ObjectiveLevel(expr=hops_expr, priority=1, name="route_hops"), - ], + [ObjectiveLevel(expr=expr, priority=len(levels) - i, name=name) for i, (name, expr) in enumerate(levels)], sense="minimize", ) @@ -2331,6 +2420,10 @@ def get_tensor_reuse_levels( reuse_levels[t] = stop return reuse_levels + def get_single_buffered(self) -> set[Tensor]: + """The tensors whose moving window the solve holds in one buffer.""" + return {t for (t, _), z in self.z_single.items() if z.X > self.VAR_THRESHOLD} + def get_tensor_depths( self, ) -> TensorDepths: @@ -2368,11 +2461,6 @@ def get_chosen_memory_cores(self) -> MemoryAlloc: chosen_memory_cores[tr] = self._fixed_tensor_choice(tensor) return chosen_memory_cores - def update_transfer_memory_core_allocation(self) -> None: - chosen_memory_cores = self.get_chosen_memory_cores() - for tr, cores in chosen_memory_cores.items(): - self.mapping.update_memory_allocation_for_node(tr, (cores,)) - def get_tensor_allocations(self) -> TensorAlloc: tensor_alloc: TensorAlloc = {} for t in self.tensor_fixed: @@ -2388,13 +2476,6 @@ def get_tensor_allocations(self) -> TensorAlloc: tensor_alloc[t] = chosen[0] return tensor_alloc - def _check_io_transfers_firing_levels(self) -> None: - for tr in self.transfer_nodes: - for t in tr.tensors: - stop_max = len(self.ssis[t].get_applicable_temporal_variables()) - stop = next(s for s in range(-1, stop_max) if self.z_stop[(t, s)].X > self.VAR_THRESHOLD) - assert stop >= 0 - def _retrieve_core_allocation(self, node: Node) -> tuple[tuple[Core, ...], ...]: if isinstance(node, InEdge): assert self.accelerator.offchip_core_id is not None @@ -2588,7 +2669,7 @@ def _active_transfer_latency( if (tr, choice) in self._transfer_latency_cache: active_latency_absent_loops_and_reuse_factor = self._transfer_latency_cache[(tr, choice)] else: - latency_constant = float(self._transfer_latency_for_path(tr, choice)) + latency_constant = float(self.transfer_latency_for_path(tr, choice)) active_latency_absent_loops = get_active_latency(tr, latency_constant, self.ssis) reuse_factor_expr = self.reuse_factors[tr]._raw @@ -2608,6 +2689,9 @@ def _active_transfer_latency( selectors=selectors, ) self._transfer_latency_cache[(tr, choice)] = active_latency_absent_loops_and_reuse_factor + self.quantities.add( + "transfer_latency", active_latency_absent_loops_and_reuse_factor._raw, index=(tr, choice) + ) return active_latency_absent_loops_and_reuse_factor @@ -2647,20 +2731,7 @@ def _shared_bandwidth_bounds(self, overlap, iteration) -> None: self.model.add_constr(busy == self.model.quicksum(terms), name=f"shared_busy_{core_id}_def") self.model.add_constr(overlap <= iteration - busy, name=f"shared_bound_{core_id}") self.shared_busy[core_id] = busy - - def _active_compute_latency( - self, - n: ComputationNode, - runtime_constant: float, - ) -> int: - # Get the temporal steady state fraction of 'ABSENT' loops - ssis_t = self.ssis.get(n).get_temporal_variables() - total_product = prod([ssis_var.size for ssis_var in ssis_t]) - product_without_absent = prod([ssis_var.size for ssis_var in ssis_t if ssis_var.effect != LoopEffect.ABSENT]) - fraction = product_without_absent / total_product if total_product > 0 else 1.0 - # Scale the runtime constant by the fraction to get the effective latency - active_latency = int(round(runtime_constant * fraction)) - return active_latency + self.quantities.add("shared_busy", busy._raw, index=core_id) def _mip_progress_callback(self, model, where): if where not in (GRB.Callback.MIP, GRB.Callback.MIPSOL, GRB.Callback.PRESOLVE): @@ -2958,14 +3029,18 @@ def _node_active(n: Node) -> tuple[int, list[Core]]: "overlap": self._overlap_section(), "tensor_reuse": self._tensor_reuse_breakdown(), "memory_occupancy": self._memory_occupancy(), - } + } | self._family_reports() + + def _family_reports(self) -> dict[str, Any]: + reports: dict[str, Any] = {} + for family in self.families: + if isinstance(family, ReportingFamily): + reports |= family.report(self, self.quantities) + return reports def primary_cost(self) -> float: """The solved value of the latency objective, whichever lexicographic level the backend ended on.""" - cost = float(self.total_lat.X) - if self.constraint_selection.dma_channels: - cost += float(self.max_core_dma_in.X) + float(self.max_core_dma_out.X) - return cost + self._traffic_weight * sum(bits * float(z.X) for bits, z in self._traffic_terms) + return self.model.value(self.quantities.get("primary").expr) def throughput_bound(self) -> float: """The pipelined compute bound of the steady state, from the solved allocation.""" @@ -2982,7 +3057,7 @@ def throughput_bound(self) -> float: for busy in self.shared_busy.values(): per_iteration = max(per_iteration, float(busy.X)) chain = sum(float(v.X) for v in self.slot_latency.values()) - return self.iterations * per_iteration + max(0.0, chain - per_iteration) + return self.iterations * per_iteration + max(0.0, chain - per_iteration) + float(self.fill.X) def capacity_slack(self) -> dict[int, dict[str, float]]: """Unused capacity per core: memory in bytes, fifo depth and buffer descriptors in slots.""" @@ -3172,12 +3247,8 @@ def _scalar(v: Any) -> Any: latencies = [self.cost_lut.get_cost(n, c).latency_total for c in cores] runtime = ceil(max(latencies)) if latencies else 0 active = get_active_latency(n, float(runtime), self.ssis) - fraction: float | None = None try: - ssis_t = self.ssis.get(n).get_temporal_variables() - total = prod([v.size for v in ssis_t]) - present = prod([v.size for v in ssis_t if v.effect != LoopEffect.ABSENT]) - fraction = (present / total) if total > 0 else 1.0 + fraction: float | None = active_fraction(n, self.ssis) except Exception: fraction = None breakdown.setdefault( @@ -3202,15 +3273,15 @@ def _scalar(v: Any) -> Any: if float(y.X) < 0.5: # noqa: PLR2004 continue s = int(self.slot_of[tr]) - raw = int(self._transfer_latency_for_path(tr, choice)) + raw = int(self.transfer_latency_for_path(tr, choice)) active_abs = int(get_active_latency(tr, float(raw), self.ssis)) try: - reuse_factor = float(self.reuse_factors[tr].getValue()) + reuse_factor = self.model.value(self.reuse_factors[tr]) except Exception: reuse_factor = None cached_var = self._transfer_latency_cache.get((tr, choice)) try: - contribution = float(cached_var.getValue()) if cached_var is not None else None + contribution = self.model.value(cached_var) if cached_var is not None else None except Exception: contribution = None tensor_bits: int | None = None diff --git a/stream/opt/allocation/constraint_optimization/utils.py b/stream/opt/allocation/constraint_optimization/utils.py index 5415d9cc7..68cd9c5f4 100644 --- a/stream/opt/allocation/constraint_optimization/utils.py +++ b/stream/opt/allocation/constraint_optimization/utils.py @@ -1,309 +1,10 @@ from math import ceil, prod -from typing import TYPE_CHECKING - -from zigzag.datatypes import LayerDim, UnrollFactor from stream.cost_model.communication_manager import MulticastPathPlan -from stream.cost_model.core_cost_lut import CoreCostLUT -from stream.hardware.architecture.accelerator import Accelerator -from stream.hardware.architecture.core import Core from stream.workload.node import Node, TransferNode -from stream.workload.steady_state.computation import SteadyStateComputation from stream.workload.steady_state.iteration_space import LoopEffect, SteadyStateIterationSpace from stream.workload.workload import ComputationNode -if TYPE_CHECKING: - from stream.opt.allocation.constraint_optimization.timeslot_allocation import TimeSlotAllocation - -MODULATION_NUMBER = 1 << 20 # Must be higher than any node's sub id - - -def convert_id(i: int, j: int) -> int: - assert 0 <= j < MODULATION_NUMBER - return MODULATION_NUMBER * i + j - - -def invert_id(k: int) -> tuple[int, int]: - i, j = divmod(k, MODULATION_NUMBER) - return i, j - - -def convert_ids(nodes: list[ComputationNode]): - return {} - - -def invert_ids_list( - ids_list: list[tuple[int, Core | int, int]], nb_nodes: int -) -> list[tuple[int, int, tuple[int, int]]]: - return [] - - -def get_loop_size(loops: list[tuple[LayerDim, UnrollFactor]]) -> int: - return int(prod(unroll_factor for _, unroll_factor in loops)) - - -def get_latencies( - nodes: list[ComputationNode], - core_ids: list[int], - accelerator: Accelerator, - cost_lut: CoreCostLUT, - impossible_lat: float = 1e11, -) -> tuple[dict[tuple[ComputationNode, Core, int], int], dict]: - cores = [accelerator.get_core(core_id) for core_id in core_ids] - latencies = {(node, core): impossible_lat for node in nodes for core in cores} - possible_allocations: dict[ComputationNode, list[Core]] = {} - inter_core_tiling_sizes = {} - - for node in nodes: - possible_allocations[node] = [] - for core in cores: - if node.chosen_core_allocation is not None and core.id != node.chosen_core_allocation: - continue - try: - equal_node = cost_lut.get_equal_node(node) - assert equal_node, f"No equal node for {node} found in CoreCostLUT" - cce = cost_lut.get_cost(equal_node, core) - inter_core_tiling_size = get_loop_size(node.inter_core_tiling) - inter_core_tiling_sizes[(node, core)] = inter_core_tiling_size - lat = cce.latency_total - possible_allocations[node].append(core) - except ValueError: - lat = impossible_lat - latencies[(node, core)] = lat - - latencies_with_split = {} - possible_allocation_splits = {} - p_max = len(cores) # maximum parallelization factor - - for node in nodes: - possible_allocation_splits[node] = {} - for core in cores: - possible_allocation_splits[node][core] = {} - if core in possible_allocations[node]: - p_t = int(inter_core_tiling_sizes[node, core]) - for p in range(1, p_max + 1): - if p <= len(possible_allocations[node]): - lat = int(latencies[(node, core)] / min(p_t, p)) - possible_allocation_splits[node][core][p] = 1 - else: - lat = impossible_lat - possible_allocation_splits[node][core][p] = 0 - latencies_with_split[(node, core, p)] = lat - else: - for p in range(1, p_max + 1): - latencies_with_split[(node, core, p)] = impossible_lat - possible_allocation_splits[node][core][p] = 0 - - return latencies_with_split, possible_allocation_splits - - -def get_energies( - nodes: list[ComputationNode], - core_ids: list[int], - accelerator: Accelerator, - cost_lut: CoreCostLUT, - impossible_energy: float = 1e11, - ids: dict[ComputationNode, int] | None = None, -) -> dict[tuple[int, Core], float]: - if ids is None: - ids = {node: node.id for node in nodes} - cores = [accelerator.get_core(core_id) for core_id in core_ids] - energies = {(ids[node], core): impossible_energy for node in nodes for core in cores} - - for node in nodes: - for core in cores: - try: - eq_node = cost_lut.get_equal_node(node) - assert eq_node, f"No equal node for {node} found in CoreCostLUT" - cme = cost_lut.get_cost(eq_node, core) - en = cme.energy_total - except ValueError: - en = impossible_energy - energies[(ids[node], core)] = en - - return energies - - -def get_node_latencies(allocation: "TimeSlotAllocation", cost_lut, accelerator, latency_attr): - node_latencies = {} - for node in allocation.get_computation_nodes(): - cores = allocation.get_resources_for_node(node) - assert all(isinstance(core, Core) for core in cores), f"Node {node} has non-core resources: {cores}" - core_ids = [core.id for core in cores if isinstance(core, Core)] - latencies, _ = get_latencies([node], core_ids, accelerator, cost_lut, latency_attr=latency_attr) - p = len(cores) - for core in cores: - assert isinstance(core, Core), f"Resource {core} for node {node} is not a Core." - combined_id = (node, core, p) - lat = latencies[combined_id] - node_latencies[node, core] = lat - return node_latencies - - -def get_timestep_latencies(allocation: "TimeSlotAllocation", timesteps): - timestep_latencies = {t: 0 for t in range(max(timesteps) + 1)} - for timestep in timesteps: - max_latency = 0 - for _, node in allocation.get_allocations_in_slot(timestep).items(): - assert node.runtime is not None, f"Node {node} has no runtime set." - max_latency = max(max_latency, int(node.runtime)) - timestep_latencies[timestep] = max_latency - return timestep_latencies - - -def get_node_start_timesteps(allocation: "TimeSlotAllocation", timestep_latencies): - starts = {} - for timestep, core, node in allocation.allocations: - start = get_start_time_of_timestep(timestep, timestep_latencies) - starts[node, core] = start - return starts - - -def get_start_time_of_timestep(timestep, timestep_latencies, t_start=0): - for t in range(timestep): - t_end = t_start + timestep_latencies[t] - t_start = t_end - return t_start - - -def calculate_total_latency(allocation: "TimeSlotAllocation", iterations) -> tuple[int, str]: - timesteps = allocation.slots - timestep_latencies = get_timestep_latencies(allocation, timesteps) - starts = get_node_start_timesteps(allocation, timestep_latencies) - total_timestep_latency = sum(timestep_latencies.values()) - overlap = compute_iterations_overlap(allocation, timestep_latencies, starts, total_timestep_latency) - total_lat = iterations * total_timestep_latency - (iterations - 1) * overlap - total_lat_str = ( - "total_lat = N * T - (N - 1) * overlap --> " - f"{total_lat} = {iterations} * {total_timestep_latency} - {iterations - 1} * {overlap}" - ) - return total_lat, total_lat_str - - -def compute_iterations_overlap(allocation: "TimeSlotAllocation", timestep_latencies, starts, iteration_latency): - slacks = {} - for core in allocation.resources: - relevant_starts = [v for k, v in starts.items() if k[1] == core] - earliest_start = min(relevant_starts) - latest_start = max(relevant_starts) - latest_node_on_this_core = next((k for k, v in starts.items() if v == latest_start and k[1] == core))[0] - latest_timestep = allocation.get_timeslot_of_node_on_resource(latest_node_on_this_core, core) - timestep_latency = timestep_latencies[latest_timestep] - latest_end = latest_start + timestep_latency - slack = iteration_latency - latest_end + earliest_start - assert slack >= 0 - slacks[core] = slack - overlap = min(slacks.values()) - return overlap - - -def get_partitioned_nodes( - node: ComputationNode, - core_allocations: list[Core], - accelerator: Accelerator, - cost_lut: CoreCostLUT, - multiplicity: int, -) -> list[SteadyStateComputation]: - """ - Get the partitioned SteadyStateComputation nodes for a given ComputationNode based on the core allocations. - """ - # If the node is not partitioned, return it as is - if len(core_allocations) == 1: - core = next(iter(core_allocations)) - latencies, _ = get_latencies( - [ - node, - ], - [ - core.id, - ], - accelerator, - cost_lut, - ) - runtime = latencies[(node, core, 1)] - mapping_attr = node.extract_inter_core_mapping_attr() - possible_resource_allocation = [ - core, - ] - new_node = SteadyStateComputation( - id=node.id, - sub_id=node.sub_id, - node_name=f"{node.name}.{node.sub_id}", - node_attr=node.extract_node_attr(), - mapping_attr=mapping_attr, - operand_tensor_reshape=node.operand_tensor_reshape, - produces_final_output=node.produces_final_output, - group_id=node.group, - input_names=node.input_names, - partially_constant_operands=node.partially_constant_operands, - possible_resource_allocation=possible_resource_allocation, - ssis_multiplicity=multiplicity, - ) - new_node.set_runtime(runtime) - new_node.update_loop_ranges(node.loop_ranges) - return [ - new_node, - ] - # Get the inter-core-tiling - inter_core_tiling = node.inter_core_tiling - if len(inter_core_tiling) > 1: - raise NotImplementedError( - f"Partitioning of nodes with inter-core tiling {inter_core_tiling} is not supported yet." - ) - tiling_dim = inter_core_tiling[0][0] - original_size = node.layer_dim_sizes[tiling_dim] - nb_tiles = len(core_allocations) - # Get the original loop range of the node for the tiling dimension - original_loop_range = node.loop_ranges[tiling_dim] - # Calculate the new loop range for each partitioned node - partitioned_loop_ranges = [ - ( - original_loop_range[0] + i * (original_loop_range[1] // nb_tiles), - original_loop_range[0] + (i + 1) * (original_loop_range[1] // nb_tiles), - ) - for i in range(nb_tiles) - ] - size_per_tile = ceil(original_size / nb_tiles) - partitioned_nodes: list[SteadyStateComputation] = [] - latencies, _ = get_latencies( - [ - node, - ], - [core.id for core in core_allocations], - accelerator, - cost_lut, - ) - runtimes = {core: latencies[(node, core, len(core_allocations))] for core in core_allocations} - for i, core in enumerate(core_allocations): - # Update the layer_dim_sizes for the smaller partitioned tile - node_attr = node.extract_node_attr() - node_attr.layer_dim_sizes[tiling_dim] = size_per_tile - inter_core_mapping_attr = node.extract_inter_core_mapping_attr() - inter_core_mapping_attr.inter_core_tiling = [(tiling_dim, nb_tiles)] - possible_resource_allocation = [ - core, - ] - # Create a new ComputationNode for each core allocation - partitioned_node = SteadyStateComputation( - id=node.id, - sub_id=node.sub_id, - node_name=f"{node.name}.{node.sub_id}.part{i}", - node_attr=node_attr, - mapping_attr=inter_core_mapping_attr, - operand_tensor_reshape=node.operand_tensor_reshape, - produces_final_output=node.produces_final_output, - group_id=node.group, - input_names=node.input_names, - partially_constant_operands=node.partially_constant_operands, - possible_resource_allocation=possible_resource_allocation, - ssis_multiplicity=multiplicity, - ) - partitioned_node.loop_ranges[tiling_dim] = partitioned_loop_ranges[i] - partitioned_node.set_chosen_core_allocation(core.id) - partitioned_node.set_runtime(runtimes[core]) - partitioned_nodes.append(partitioned_node) - return partitioned_nodes - def get_transfer_latency_for_path(tr: TransferNode, path: MulticastPathPlan) -> int: """Cycles one firing of this transfer costs on this path. @@ -328,12 +29,12 @@ def get_active_transfer_latency_for_path(tr: TransferNode, choice: MulticastPath return active_latency +def active_fraction(n: Node, ssis: dict[ComputationNode, SteadyStateIterationSpace]) -> float: + """Fraction of the steady-state iterations in which ``n`` is not idle on an absent loop.""" + temporal = ssis.get(n).get_temporal_variables() + total = prod(v.size for v in temporal) + return prod(v.size for v in temporal if v.effect != LoopEffect.ABSENT) / total if total else 1.0 + + def get_active_latency(n: Node, runtime_constant: float, ssis: dict[ComputationNode, SteadyStateIterationSpace]) -> int: - # Get the temporal steady state fraction of 'ABSENT' loops - ssis_t = ssis.get(n).get_temporal_variables() - total_product = prod([ssis_var.size for ssis_var in ssis_t]) - product_without_absent = prod([ssis_var.size for ssis_var in ssis_t if ssis_var.effect != LoopEffect.ABSENT]) - fraction = product_without_absent / total_product if total_product > 0 else 1.0 - # Scale the runtime constant by the fraction to get the effective latency - active_latency = int(round(runtime_constant * fraction)) - return active_latency + return int(round(runtime_constant * active_fraction(n, ssis))) diff --git a/stream/opt/solver/solver.py b/stream/opt/solver/solver.py index 29c2184ff..bec3e8b06 100644 --- a/stream/opt/solver/solver.py +++ b/stream/opt/solver/solver.py @@ -6,7 +6,7 @@ import logging import math from abc import ABC, abstractmethod -from collections.abc import Sequence +from collections.abc import Mapping, Sequence from dataclasses import dataclass from enum import Enum, auto from typing import Any @@ -108,6 +108,8 @@ class ConstraintSelection: offchip_contention: bool = True offchip_traffic_cost: bool = True pipelining: PipeliningModel = PipeliningModel.OCCUPANCY + families: tuple[str | Mapping[str, Any], ...] = () + """Constraint families to build, by entry-point name, or ``{name: options}``; none by default.""" def __post_init__(self) -> None: if not self.memory_capacity and self.object_fifo_depth: @@ -375,6 +377,10 @@ def get_status(self) -> str: def get_sol_count(self) -> int: """Return the number of solutions found.""" + @abstractmethod + def value(self, expr: Any) -> float: + """Solved value of a variable, a linear expression or a constant.""" + @abstractmethod def solve_stats(self) -> SolveStats: """Return structured solve statistics. @@ -461,10 +467,6 @@ def set_lexicographic_objectives( """ raise NotImplementedError(f"{type(self).__name__} does not support lexicographic objectives") - def infinity(self) -> float: - """Convenience accessor for INFINITY class constant.""" - return self.INFINITY - # --------------------------------------------------------------------------- # Gurobi private implementations @@ -701,6 +703,12 @@ def get_status(self) -> str: def get_sol_count(self) -> int: return self._model.SolCount + def value(self, expr: Any) -> float: + raw = _unwrap(expr) + if isinstance(raw, int | float): + return float(raw) + return float(raw.getValue() if hasattr(raw, "getValue") else raw.X) + def _mip_gap(self) -> float | None: """Relative optimality gap; falls back to the primal/dual bounds when ``MIPGap`` is absent.""" if self._model.SolCount <= 0: @@ -1102,6 +1110,14 @@ def get_sol_count(self) -> int: return 0 return 1 if self._result.has_primal_feasible_solution() else 0 + def value(self, expr: Any) -> float: + raw = _unwrap_ort(expr) + if isinstance(raw, int | float): + return float(raw) + if self._result is None or not self._result.has_primal_feasible_solution(): + raise ValueError("No solution available") + return float(mathopt.evaluate_expression(raw, self._result.variable_values())) + def _mip_gap(self) -> float | None: """Relative optimality gap, derived from MathOpt's primal/dual bounds.""" if self._result is None or not self._result.has_primal_feasible_solution(): diff --git a/stream/parser/accelerator_factory.py b/stream/parser/accelerator_factory.py index 9db859b4f..edec372b3 100644 --- a/stream/parser/accelerator_factory.py +++ b/stream/parser/accelerator_factory.py @@ -8,6 +8,7 @@ from stream.hardware.architecture.backends import AIE2CoreBackend, ZigZagCoreBackend from stream.hardware.architecture.core import Core from stream.hardware.architecture.noc.communication_link import CommunicationLink, get_bidirectional_edges +from stream.parser.accelerator_validator import parse_memory_ref, parse_port_ref from stream.parser.core_validator import ALLOWED_KINDS, ALLOWED_NAMESPACES, CoreValidatorRegistry @@ -44,18 +45,33 @@ def create(self) -> Accelerator: cores_graph = self.create_core_graph(cores) # Take next available core id - return Accelerator( + accelerator = Accelerator( name=self.data["name"], cores=cores_graph, offchip_core_id=offchip_core_id, shared_mem_group_ids=shared_mem_group_ids, kernel_library=KernelLibrary.load(self.data.get("kernel_library")), bandwidth={ - int(core): BandwidthModel.from_description(model) + core: BandwidthModel.from_description(model) for core, model in self.data.get("bandwidth", {}).items() + if isinstance(core, int) }, reconfiguration=self.data.get("reconfiguration"), + port_bandwidth={ + parse_port_ref(ref): BandwidthModel.from_description(model) + for ref, model in self.data.get("bandwidth", {}).items() + if isinstance(ref, str) + }, + memory_aliases={ + parse_memory_ref(alias): parse_memory_ref(group[0]) + for group in self.data.get("memory_aliases") or [] + for alias in group[1:] + }, ) + if accelerator.port_bandwidth: + # Building the registry rejects a port key that names no port a core's backend models. + _ = accelerator.ports + return accelerator def create_core( self, @@ -83,10 +99,17 @@ def create_core( if namespace == "aie2": # ---- AIE2 native path: lightweight backend ---- mem = core_data["memory"] + dma = core_data.get("dma", {}) backend = AIE2CoreBackend( memory_capacity_bits=mem["capacity"], bandwidth_min=mem.get("bandwidth_min", 0), bandwidth_max=mem.get("bandwidth_max", 0), + dma_mm2s=dma.get("mm2s", 0), + dma_s2mm=dma.get("s2mm", 0), + dma_channel_bits=dma.get("channel_bits", 0), + dma_buffer_descriptors=dma.get("buffer_descriptors", 0), + dma_iterations=dma.get("iterations", 0), + share_group=core_id if shared_mem_group_id is None else shared_mem_group_id, ) core = Core( backend=backend, diff --git a/stream/parser/accelerator_validator.py b/stream/parser/accelerator_validator.py index 8b414e039..33c6f26d9 100644 --- a/stream/parser/accelerator_validator.py +++ b/stream/parser/accelerator_validator.py @@ -1,6 +1,7 @@ import copy import logging import os +import re from functools import reduce from itertools import combinations from typing import Any @@ -8,6 +9,7 @@ from cerberus import Validator from zigzag.utils import open_yaml +from stream.hardware.ports import MemoryRef, PortRef from stream.parser.core_validator import ALLOWED_KINDS, ALLOWED_NAMESPACES, CoreValidatorRegistry logger = logging.getLogger(__name__) @@ -21,6 +23,21 @@ r"(?:\.ya?ml)?$" # optional ".yaml" or ".yml" ) +# "..", e.g. "2.dram.rw_port_1" +PORT_REF_REGEX = r"^\d+\.[A-Za-z0-9_\-]+\.[A-Za-z0-9_\-]+$" + + +def parse_port_ref(ref: str) -> PortRef: + """Split a ``bandwidth:`` port key into (core id, memory instance name, port name).""" + core, memory, port = ref.split(".") + return PortRef(int(core), memory, port) + + +def parse_memory_ref(ref: str) -> MemoryRef: + """Split a ``memory_aliases`` entry ``.``.""" + core, _, memory = ref.partition(".") + return int(core), memory + def resolve_core_path(core_file_name: str, accelerator_dirname: str) -> str | None: """Resolve a ``cores:`` reference to a path on disk, or None when nothing matches.""" @@ -79,8 +96,13 @@ class AcceleratorValidator: }, # Id of the core that acts as the off-chip memory controller "offchip_core_id": {"type": "integer", "min": 0, "required": True}, - # Core id -> measured bandwidth every transfer through that core shares - "bandwidth": {"type": "dict", "required": False, "keysrules": {"type": "integer"}}, + # Core id -> measured bandwidth every transfer through that core shares, or a port key -> + # measured bandwidth replacing that top-level port's declared width + "bandwidth": { + "type": "dict", + "required": False, + "keysrules": {"anyof": [{"type": "integer"}, {"type": "string", "regex": PORT_REF_REGEX}]}, + }, # What reconfiguring the array between designs costs, read by the namespace constraints "reconfiguration": {"type": "dict", "required": False}, # Optional unit_energy_cost used for connections that don't specify their own @@ -177,6 +199,7 @@ def validate(self) -> bool: self.validate_core_connectivity() self.validate_core_mem_sharing() self.validate_memory_aliases() + self.validate_bandwidth() if not self.is_valid and self.errors: logger.critical("Accelerator validation failed with %d issue(s).", len(self.errors)) @@ -369,6 +392,20 @@ def validate_core_connectivity(self): conn.setdefault("type", "link") conn.setdefault("unit_energy_cost", self.data.get("unit_energy_cost", 0)) + def validate_bandwidth(self) -> None: + """Every ``bandwidth:`` key names an existing core; the accelerator factory checks that a port key + names a modelled port.""" + keys = self.data.get("bandwidth", {}) + measured_cores = {key for key in keys if isinstance(key, int)} + for core_id in measured_cores - set(self.data["cores"]): + self.invalidate(f"`bandwidth` names unknown core {core_id}.") + for ref in (key for key in keys if isinstance(key, str) and re.match(PORT_REF_REGEX, key)): + core_id, _, _ = parse_port_ref(ref) + if core_id not in self.data["cores"]: + self.invalidate(f"`bandwidth.{ref}`: unknown core {core_id}.") + elif core_id in measured_cores: + self.invalidate(f"`bandwidth.{ref}`: core {core_id} already has a measured core bandwidth.") + def validate_core_mem_sharing(self): # Replace string of core ids with tuple of ints mem_sharing_data = self.data["core_memory_sharing"] @@ -402,8 +439,7 @@ def validate_memory_aliases(self): """Every ``.`` in `memory_aliases` must name a memory that exists.""" for group in self.data.get("memory_aliases", []): for ref in group: - core_id_str, _, mem_name = ref.partition(".") - core_id = int(core_id_str) + core_id, mem_name = parse_memory_ref(ref) core_data = self.data["cores"].get(core_id) if core_data is None: self.invalidate(f"`memory_aliases` entry '{ref}' names unknown core id {core_id}.") diff --git a/stream/parser/core_validator.py b/stream/parser/core_validator.py index c971bcb0e..d09113f26 100644 --- a/stream/parser/core_validator.py +++ b/stream/parser/core_validator.py @@ -239,6 +239,17 @@ class ZigZagOffchipCoreValidator(ZigZagBaseCoreValidator): "bandwidth_max": {"type": "integer", "required": True, "min": 0}, }, }, + "dma": { + "type": "dict", + "required": False, + "schema": { + "mm2s": {"type": "integer", "required": True, "min": 1}, + "s2mm": {"type": "integer", "required": True, "min": 1}, + "channel_bits": {"type": "integer", "required": True, "min": 1}, + "buffer_descriptors": {"type": "integer", "required": False, "min": 1}, + "iterations": {"type": "integer", "required": False, "min": 1}, + }, + }, } diff --git a/stream/parser/mapping_factory.py b/stream/parser/mapping_factory.py index 334d25b57..a9e518310 100644 --- a/stream/parser/mapping_factory.py +++ b/stream/parser/mapping_factory.py @@ -139,13 +139,16 @@ def _group_kernels(self, layers: tuple[str, ...]): yield node, kernel def _call_tile_tiling(self, layers: tuple[str, ...]) -> tuple[tuple[LayerDim, int], ...]: - """A group's default tiling: its kernels' call tiles, with the dimensions its layers share innermost.""" + """A group's default tiling, innermost first: its kernels' call tiles, with the dimensions its layers + share innermost, then one at a time each dimension no kernel addresses, such as a batch axis.""" tiling: dict[LayerDim, int] = {} + addressed: set[LayerDim] = set() kernels = list(self._group_kernels(layers)) for node, kernel in kernels: node_dims = self.workload.get_dims(node) - for position, size, dim in kernel.call_tile(): + for position, size, dim in kernel.call_tile(node): layer_dim = node_dims[position] + addressed.add(layer_dim) extent = self.workload.get_dimension_size(layer_dim) if size < extent or dim.keep_whole: tiling.setdefault(layer_dim, min(size, extent)) @@ -154,7 +157,13 @@ def _call_tile_tiling(self, layers: tuple[str, ...]) -> tuple[tuple[LayerDim, in for dim in a & b: if dim in shared: shared[dim] += 1 - return tuple(sorted(tiling.items(), key=lambda item: shared[item[0]])) + batch = dict.fromkeys( + dim + for node, _ in kernels + for dim in self.workload.get_dims(node) + if dim not in addressed and self.workload.get_dimension_size(dim) > 1 + ) + return (*sorted(tiling.items(), key=lambda item: shared[item[0]]), *((dim, 1) for dim in batch)) def _runtime_dims(self, layers: tuple[str, ...], tiling) -> tuple[LayerDim, ...]: """Tiling dimensions every kernel of the group takes as a runtime size.""" @@ -162,7 +171,7 @@ def _runtime_dims(self, layers: tuple[str, ...], tiling) -> tuple[LayerDim, ...] runtime: set[LayerDim] = set() for node, kernel in self._group_kernels(layers): node_dims = self.workload.get_dims(node) - for position, _, dim in kernel.call_tile(): + for position, _, dim in kernel.call_tile(node): (runtime if dim.runtime else compiled).add(node_dims[position]) return tuple(dim for dim, _ in tiling if dim in runtime and dim not in compiled) diff --git a/stream/parser/onnx/matmul.py b/stream/parser/onnx/matmul.py index 7ab3a2825..064a1eb50 100644 --- a/stream/parser/onnx/matmul.py +++ b/stream/parser/onnx/matmul.py @@ -5,25 +5,22 @@ class MatMulParser(OnnxOperatorParser): - """Parses an ONNX MatMul into an affine ``ComputationNode``. - - Semantics follow ``numpy.matmul``: the last two axes are the matrix axes ``(m, k) @ (k, n)`` and - leading axes are numpy-broadcast batch axes (right-aligned). Iteration space ``(batch..., m, n, k)`` - with ``k`` the single contraction (REDUCTION); a broadcast batch axis (size 1 or absent) indexes - that operand at constant 0, keeping the map affine. - """ + """Parses an ONNX MatMul, as ``numpy.matmul``, into a ``ComputationNode`` over ``(batch..., m, n, k)`` with + ``k`` the contraction. A broadcast batch axis (size 1 against a larger output axis, or absent) indexes its + operand at constant 0, keeping the map affine.""" EXPECTED_NB_OF_INPUTS = 2 # A and B def _batch_exprs(self, operand_batch: tuple[int, ...], out_batch: tuple[int, ...]) -> list[AffineExpr]: - """Batch-axis index expressions: the iteration dim when the axis matches the output extent, - constant 0 when the operand broadcasts over it.""" + """Batch-axis index expressions: the iteration dim when the axis matches the output extent, even + an extent of one, so every operand's axis stays the node's; constant 0 when the operand broadcasts + over it.""" num_batch = len(out_batch) offset = num_batch - len(operand_batch) # right-align operand batch axes to the output's exprs: list[AffineExpr] = [] for j, size in enumerate(operand_batch): out_pos = offset + j - if size == out_batch[out_pos] and size != 1: + if size == out_batch[out_pos]: exprs.append(AffineExpr.dimension(out_pos)) else: exprs.append(AffineExpr.constant(0)) diff --git a/stream/stages/allocation/constraint_optimization_allocation.py b/stream/stages/allocation/constraint_optimization_allocation.py index 02670c83f..03bd69c17 100644 --- a/stream/stages/allocation/constraint_optimization_allocation.py +++ b/stream/stages/allocation/constraint_optimization_allocation.py @@ -45,13 +45,7 @@ def __init__( self.cost_lut = self.ctx.get("cost_lut") config = self.ctx.get("constraint_opt_config") - if config is None: - logger.warning( - "ConstraintOptimizationAllocationStage: legacy kwargs configuration path is deprecated. " - "Please pass a ConstraintOptStageConfig. Building config from kwargs for now." - ) - config = ConstraintOptStageConfig.from_legacy_kwargs(**self.ctx.data) - self.config = config + self.config = config or ConstraintOptStageConfig.from_kwargs(**self.ctx.data) self.output_path = self.ctx.get("output_path") self.backend: str = self.ctx.get("backend", "ORTOOLS_GSCIP") diff --git a/stream/stages/codegen/aie_code_generation.py b/stream/stages/codegen/aie_code_generation.py index 4d9d608d1..1b9e528c8 100644 --- a/stream/stages/codegen/aie_code_generation.py +++ b/stream/stages/codegen/aie_code_generation.py @@ -12,6 +12,7 @@ from stream.compiler.context.aie_context import AIEContext from stream.compiler.dialects.stream import ( + DOUBLE_BUFFERED, ComputationNodeOp, FusionGroupOp, InEdgeOp, @@ -152,17 +153,20 @@ def create_transfer_op( # noqa: PLR0912, PLR0915 result_types = [result_type] * len(node.outputs) elif len(node.outputs) > 1: result_types = [] - for (index, space), reader in zip(spaces, self._readers(node, workload), strict=True): + for (index, space), reader, output in zip(spaces, self._readers(node, workload), node.outputs, strict=True): allocation = full_mapping.get(reader).resource_allocation[0] reader_cores = allocation.sources if isinstance(allocation, MulticastPathPlan) else allocation cores_attr = [StringAttr(f"tile_{c.col_id}_{c.row_id}") for c in reader_cores] - result_types.append(StrensorType(input_type.element_type, space, cores_attr, index)) + result_types.append( + StrensorType(input_type.element_type, space, cores_attr, index, self._buffers(output)) + ) else: result_type = StrensorType( input_type.element_type, ss, cores, reuse_index, + self._buffers(node.outputs[0]), ) result_types = [result_type] op = TransferOp( @@ -179,6 +183,10 @@ def create_transfer_op( # noqa: PLR0912, PLR0915 return op + def _buffers(self, tensor: Tensor) -> int: + """The buffers the solve holds a moving window of ``tensor`` in.""" + return 1 if tensor in self.ctx.get("scheduler").single_buffered else DOUBLE_BUFFERED + @staticmethod def _readers(node: TransferNode, workload: Workload) -> list[HasInputs]: """Per output of a transfer, the node that reads it.""" @@ -402,6 +410,15 @@ def inputs(node: HasInputs) -> Iterable[SSAValue]: return module + def _shim_limits(self) -> tuple[int, int]: + """The buffer descriptors a shim tile has for the runtime sequence and the steps one iterates, + as the accelerator declares them; 0 where it does not. A trace writes its buffer out through + one descriptor.""" + shims = [core for core in self.ctx.get("accelerator").core_list if core.type == "shim"] + descriptors = min((core.dma_buffer_descriptors for core in shims), default=0) + iterations = min((core.dma_iterations for core in shims), default=0) + return descriptors - 1 if descriptors and self.trace_size else descriptors, iterations + def _fifo_depths(self) -> FifoDepths | None: """A depth policy funded by the capacity the solved allocation left unused.""" scheduler = self.ctx.get("scheduler") @@ -458,7 +475,10 @@ def codegen_main(self) -> None: IterationSpaceToFor().apply(self.context, module) with open(output_path + "/with_for.mlir", "w") as f: f.write(str(module)) - AIEConvertOfs(depths=self._fifo_depths()).apply(self.context, module) + descriptors, iterations = self._shim_limits() + AIEConvertOfs(depths=self._fifo_depths(), shim_descriptors=descriptors, shim_iterations=iterations).apply( + self.context, module + ) with open(output_path + "/convert_of.mlir", "w") as f: f.write(str(module)) ConvertStreamToAIEPass().apply(self.context, module) diff --git a/stream/stages/estimation/core_cost_backends.py b/stream/stages/estimation/core_cost_backends.py index 3ef6cf442..8910579a0 100644 --- a/stream/stages/estimation/core_cost_backends.py +++ b/stream/stages/estimation/core_cost_backends.py @@ -3,8 +3,11 @@ from __future__ import annotations import logging -from typing import TYPE_CHECKING, Protocol +from typing import TYPE_CHECKING, Protocol, runtime_checkable +from zigzag.hardware.architecture.memory_port import DataDirection + +from stream.hardware.architecture.backends.zigzag import ZIGZAG_DIRECTION_NAMES, operand_role from stream.plugins import load_group from stream.stages.estimation.zigzag_cost_estimator import ZigZagCostEstimator @@ -22,6 +25,8 @@ CORE_COST_BACKENDS_GROUP = "stream.core_cost_backends" CONTRACT_VERSION = 1 +PortTraffic = tuple[tuple[str, str, float], ...] + class CoreEstimator(Protocol): """What a backend produces: the object the stage calls once per node-core pair.""" @@ -57,6 +62,20 @@ def make(self, context: CoreCostContext) -> CoreEstimator: ... +@runtime_checkable +class PortTrafficSource(Protocol): + """Optional backend capability: the bits a costed node moves through its core's top-level memory ports.""" + + def port_traffic(self, entry: CoreCostEntry) -> PortTraffic: + """(operand role, direction, bits) per evaluation of ``entry``'s node, before its active fraction.""" + ... + + +def port_traffic(backend: object, entry: CoreCostEntry) -> PortTraffic: + """``backend``'s port traffic for ``entry``; none from a backend without the capability.""" + return backend.port_traffic(entry) if isinstance(backend, PortTrafficSource) else () + + class AIEBackend: """AIE compute tiles: the kernel-library-priced estimator.""" @@ -91,6 +110,24 @@ def make(self, context: CoreCostContext) -> CoreEstimator: nb_spatial_mappings_generated=context.nb_spatial_mappings_generated, ) + def port_traffic(self, entry: CoreCostEntry) -> PortTraffic: + """Words each operand's top level moves to and from the datapath, times the evaluated port's width.""" + cme = entry.cme + if cme is None: + return () + traffic: list[tuple[str, str, float]] = [] + for layer_op in cme.layer.layer_operands: + mem_op = cme.memory_operand_links.layer_to_mem_op(layer_op) + top = cme.mapping.mem_level[layer_op] - 1 + level = cme.accelerator.get_memory_level(mem_op, top) + accesses = cme.memory_word_access[layer_op][top] + for direction in (DataDirection.RD_OUT_TO_LOW, DataDirection.WR_IN_BY_LOW): + words = accesses.get(direction) + port = next((p for p in level.ports if (mem_op, top, direction) in p.served_op_lv_dir), None) + if words and port is not None: + traffic.append((operand_role(str(mem_op)), ZIGZAG_DIRECTION_NAMES[direction], words * port.bw_max)) + return tuple(traffic) + # Entry-point targets registered under the public distribution (see pyproject.toml). AIE_BACKEND = AIEBackend() diff --git a/stream/stages/generation/kernel_state.py b/stream/stages/generation/kernel_state.py index 5e2b0dd0c..e9aaae0d2 100644 --- a/stream/stages/generation/kernel_state.py +++ b/stream/stages/generation/kernel_state.py @@ -29,21 +29,18 @@ def _extent(node: ComputationNode, position: int) -> int | None: return None -def state_tensor(node: ComputationNode, state) -> tuple[Tensor, AffineMap] | None: - """The operand and the map that make ``state`` a recurrence on ``node``. - - ``rows`` of state per step of the carried dimension, indexed by the dimension the node - splits over, so a split of that dimension divides the state with it. Reading the carried - dimension at ``-1`` is what marks it: the extent that reaches the tensor is clipped to - its own shape, so the carry costs ``rows`` however the carried dimension is tiled. - """ - indexed = _extent(node, state.indexed_by) - if indexed is None or _extent(node, state.carried_over) is None: +def state_tensor(node: ComputationNode, state, positions: dict[str, int]) -> tuple[Tensor, AffineMap] | None: + """The operand and map that make ``state`` a recurrence on ``node``: ``rows`` per step of the carried + dimension, read at ``-1``, indexed by the dimension the node splits over. The carry costs ``rows`` however + the carried dimension is tiled, since the extent reaching the tensor is clipped to its shape.""" + carried, index = positions[state.carried_over], positions[state.indexed_by] + indexed = _extent(node, index) + if indexed is None or _extent(node, carried) is None: return None rank = len(node.operand_mapping[0].results) if node.operand_mapping else 0 - dims = max(rank, state.carried_over + 1, state.indexed_by + 1) - carry: AffineExpr = AffineDimExpr(state.carried_over) - 1 - mapping = AffineMap(dims, 0, (carry, AffineDimExpr(state.indexed_by))) + dims = max(rank, carried + 1, index + 1) + carry: AffineExpr = AffineDimExpr(carried) - 1 + mapping = AffineMap(dims, 0, (carry, AffineDimExpr(index))) tensor = Tensor.create(f"{state.name}_{node.name}", node.outputs[0].operand_type, (state.rows, indexed)) return tensor, mapping @@ -61,13 +58,15 @@ def __init__(self, list_of_callables: list[StageCallable], ctx: StageContext): def run(self): rebuilt, changed = {}, False for node in self.workload.nodes: - declared = self._declared(node) + kernel = self._kernel(node) + declared = list(kernel.state_operands()) if kernel is not None else [] if not declared: rebuilt[node.name] = node continue inputs, maps = list(node.inputs), list(node.operand_mapping) + positions = kernel.dim_positions(node) for state in declared: - built = state_tensor(node, state) + built = state_tensor(node, state, positions) if built is None: continue tensor, mapping = built @@ -85,8 +84,7 @@ def run(self): sub_stage = self.list_of_callables[0](self.list_of_callables[1:], self.ctx) yield from sub_stage.run() - def _declared(self, node): + def _kernel(self, node): if not isinstance(node, ComputationNode): - return [] - kernel = getattr(self.mapping.get(node), "kernel", None) - return list(kernel.state_operands()) if kernel is not None else [] + return None + return getattr(self.mapping.get(node), "kernel", None) diff --git a/stream/stages/generation/placement_generation.py b/stream/stages/generation/placement_generation.py index 6d453f7cf..dadab0c4c 100644 --- a/stream/stages/generation/placement_generation.py +++ b/stream/stages/generation/placement_generation.py @@ -121,29 +121,31 @@ def _linear(self, nodes: list[ComputationNode]) -> bool: def _place_alone(self, node: ComputationNode, grid, columns, rows) -> None: kernel = self.mapping.get(node).kernel - granule = call_sizes(kernel) + granule = call_sizes(kernel, node) + at = kernel.dim_positions(node) if bandwidth_bound(kernel): self._place_one_row(node, kernel, grid, columns, rows, self.mapping) return if not is_matmul(kernel): self._place_one_row(node, kernel, grid, columns, rows, self.mapping) - depth = self._fitting_split(node, 0, len(rows), granule.get(0, 1)) + depth = self._fitting_split(node, at["m"], len(rows), granule.get(at["m"], 1)) narrow = tuple(grid[(columns[0], row)] for row in rows[:depth]) - self._pending_reserves.append(lambda m, n=node, c=narrow, sp=depth: self._assign(n, c, ((0, sp),), m)) + self._pending_reserves.append(lambda m, n=node, c=narrow, sp=((at["m"], depth),): self._assign(n, c, sp, m)) return - split = [(0, self._fitting_split(node, 0, len(rows), granule.get(0, 1)))] + split = [(at["m"], self._fitting_split(node, at["m"], len(rows), granule.get(at["m"], 1)))] cols_used = 1 - if is_matmul(kernel) and self._fitting_split(node, 2, len(columns), granule.get(2, 1)) == len(columns): + if self._fitting_split(node, at["n"], len(columns), granule.get(at["n"], 1)) == len(columns): cols_used = len(columns) - split.append((2, cols_used)) + split.append((at["n"], cols_used)) cores = tuple(grid[(col, row)] for col in columns[:cols_used] for row in rows[: split[0][1]]) self._assign(node, cores, tuple(split)) def _place_one_row(self, node, kernel, grid, columns, rows, mapping) -> None: row = rows[0] - width = self._fitting_split(node, 0, len(columns)) + at = kernel.dim_positions(node) + width = self._fitting_split(node, at["m"], len(columns)) cores = tuple(grid[(col, row)] for col in columns[:width]) - self._assign(node, cores, ((0, width),), mapping) + self._assign(node, cores, ((at["m"], width),), mapping) width = self._row_width(node, kernel, next(iter(grid.values()))) self._retile(node, kernel, mapping, m=1, n=width, layout="contiguous") @@ -155,15 +157,18 @@ def _place_stacked(self, nodes: list[ComputationNode], grid, columns, rows) -> N if kernel is not None and any(s.handover for s in kernel.state_operands()) ) counts = row_counts([layer_cost(k) for k in kernels], len(rows), state_consumers=state_consumers) - granule = max(call_sizes(kernel).get(0, 1) for kernel in kernels) - extent = self.workload.get_dimension_size(self.workload.get_dims(nodes[0])[0]) + rows_at = [kernel.dim_positions(node)["m"] for node, kernel in zip(nodes, kernels, strict=True)] + granule = max( + call_sizes(kernel, node).get(at, 1) for node, kernel, at in zip(nodes, kernels, rows_at, strict=True) + ) + extent = self.workload.get_dimension_size(self.workload.get_dims(nodes[0])[rows_at[0]]) width = widest_columns(extent, granule, max(counts), len(columns)) offset = 0 - for node, r in zip(nodes, counts, strict=True): + for node, r, at in zip(nodes, counts, rows_at, strict=True): layer_rows = rows[offset : offset + r] offset += r cores = tuple(grid[(col, row)] for row in layer_rows for col in columns[:width]) - self._assign(node, cores, ((0, width * r),)) + self._assign(node, cores, ((at, width * r),)) for node, kernel, r, r_next in zip(nodes, kernels, counts, (*counts[1:], counts[-1]), strict=True): if r > r_next: self._retile(node, kernel, tiled_out=True) @@ -183,13 +188,14 @@ def _assign_tenancy(self, nodes, kernels, budgets, grid, columns, rows, mapping) for node, kernel, budget in zip(nodes, kernels, budgets, strict=True): tenant = columns[first : first + budget] first += budget - granule = call_sizes(kernel) + granule = call_sizes(kernel, node) + at = kernel.dim_positions(node) width = 1 - split = [(0, self._fitting_split(node, 0, len(rows), granule.get(0, 1)))] + split = [(at["m"], self._fitting_split(node, at["m"], len(rows), granule.get(at["m"], 1)))] if is_matmul(kernel) and budget > 1: - width = self._fitting_split(node, 2, budget, granule.get(2, 1)) + width = self._fitting_split(node, at["n"], budget, granule.get(at["n"], 1)) if width > 1: - split.append((2, width)) + split.append((at["n"], width)) cores = tuple(grid[(col, row)] for col in tenant[:width] for row in rows[: split[0][1]]) self._assign(node, cores, tuple(split), mapping) @@ -202,7 +208,7 @@ def _fitting_split(self, node: ComputationNode, position: int, target: int, gran return 1 def _row_width(self, node: ComputationNode, kernel, core: Core) -> int: - extent = self.workload.get_dimension_size(self.workload.get_dims(node)[1]) + extent = self.workload.get_dimension_size(self.workload.get_dims(node)[kernel.dim_positions(node)["n"]]) operands = max(MIN_CALL_OPERANDS, len(kernel.operand_layouts() or ())) budget = (core.get_memory_capacity() // 8) // (operands * BUFFERS * ELEMENT_BYTES) return max(w for w in range(1, min(extent, budget) + 1) if extent % w == 0) diff --git a/stream/visualization/constraint_optimization.py b/stream/visualization/constraint_optimization.py deleted file mode 100644 index 55b791b1a..000000000 --- a/stream/visualization/constraint_optimization.py +++ /dev/null @@ -1,79 +0,0 @@ -import json -import logging -from typing import TYPE_CHECKING - -from stream.cost_model.core_cost_lut import CoreCostLUT -from stream.hardware.architecture.accelerator import Accelerator -from stream.opt.allocation.constraint_optimization.timeslot_allocation import TimeSlotAllocation -from stream.opt.allocation.constraint_optimization.utils import ( - get_node_latencies, - get_node_start_timesteps, - get_timestep_latencies, -) - -if TYPE_CHECKING: - pass - -logger = logging.getLogger(__name__) - - -def to_perfetto_json( - allocation: TimeSlotAllocation, - cost_lut: CoreCostLUT, - accelerator: Accelerator, - latency_attr: str, - json_path: str, -): - """ - Allocation is a list of tuples, with each tuple being of form (timestep, allocation, node_id). Allocation is a core. - cost_lut is a CoreCostLUT storing for each node and each core the hardware performance. - """ - timesteps = allocation.slots - resources = allocation.resources - node_latencies = get_node_latencies(allocation, cost_lut, accelerator, latency_attr) - timestep_latencies = get_timestep_latencies(allocation, timesteps) - starts = get_node_start_timesteps(allocation, timestep_latencies) - # total_timestep_latency = sum(timestep_latencies.values()) - # overlap = compute_iterations_overlap(allocation, timestep_latencies, starts, total_timestep_latency) - # offset = total_timestep_latency - overlap - - # Prepare JSON data for Perfetto - perfetto_data = [] - - # Add thread names (cores) - for resource in resources: - thread_name_event = { - "name": "thread_name", - "ph": "M", - "pid": "waco", - "tid": resource, - "args": {"name": f"Resource {resource}"}, - "cname": "blue", - } - perfetto_data.append(thread_name_event) - - # Add events for each iteration - # for iteration in range(iterations): - # iteration_offset = iteration * offset - for slot in range(allocation.slot_min, allocation.slot_max + 1): - for resource, node in allocation.get_allocations_in_slot(slot).items(): - start = starts[node, resource] # + iteration_offset - runtime = node_latencies[node, resource] - event = { - "name": f"{node}", - "cat": "compute", - "ph": "X", - "ts": start, - "dur": runtime, - "pid": "waco", - "tid": resource, - "cname": "blue", - "args": {"Runtime": runtime, "Id": node.id, "Name": node.node_name}, - } - perfetto_data.append(event) - - # Write JSON data to file - with open(json_path, "w") as f: - json.dump(perfetto_data, f, indent=2) - - logger.info(f"Plotted WACO result to {json_path}") diff --git a/stream/visualization/memory_usage.py b/stream/visualization/memory_usage.py deleted file mode 100644 index c8575975c..000000000 --- a/stream/visualization/memory_usage.py +++ /dev/null @@ -1,197 +0,0 @@ -import logging -from typing import TYPE_CHECKING - -import numpy as np -from brokenaxes import brokenaxes -from matplotlib import pyplot as plt -from matplotlib.gridspec import GridSpec - -if TYPE_CHECKING: - from stream.cost_model.cost_model import StreamCostModelEvaluation - -logger = logging.getLogger(__name__) - -SMALL_SIZE = 8 -MEDIUM_SIZE = 10 -BIGGER_SIZE = 12 - -plt.rc("font", size=SMALL_SIZE) # controls default text sizes -plt.rc("axes", titlesize=SMALL_SIZE) # fontsize of the axes title -plt.rc("axes", labelsize=SMALL_SIZE) # fontsize of the x and y labels -plt.rc("xtick", labelsize=SMALL_SIZE) # fontsize of the tick labels -plt.rc("ytick", labelsize=SMALL_SIZE) # fontsize of the tick labels -plt.rc("legend", fontsize=SMALL_SIZE) # legend fontsize -plt.rc("figure", titlesize=BIGGER_SIZE) # fontsize of the figure title - - -def humanbytes(bytes): - """Return the given bytes as a human friendly KB, MB, GB, or TB string.""" - bytes = float(bytes) - KB = float(1024) - MB = float(KB**2) # 1,048,576 - GB = float(KB**3) # 1,073,741,824 - TB = float(KB**4) # 1,099,511,627,776 - - if bytes < KB: - return f"{bytes} {'Bytes' if bytes == 0 or bytes > 1 else 'Byte'}" - elif KB <= bytes < MB: - return f"{bytes / KB:.2f} kB" - elif MB <= bytes < GB: - return f"{bytes / MB:.2f} MB" - elif GB <= bytes < TB: - return f"{bytes / GB:.2f} GB" - elif TB <= bytes: - return f"{bytes / TB:.2f} TB" - else: - raise ValueError(f"Cannot convert {bytes} bytes to human readable format.") - - -def plot_memory_usage( - scme: "StreamCostModelEvaluation", - section_start_percent: tuple[int, ...] = (0,), - percent_shown: tuple[int, ...] = (100,), - show_dram: bool = False, - show_human_bytes: bool = True, - fig_path: str = "outputs/memory_usage.png", - fig_size: tuple[int, int] = (16, 9), -): - memory_manager = scme.accelerator.memory_manager - cpti, mopti = _filter_memory_instances(memory_manager, show_dram) - cpti, mopti = _remove_zero_usage_instances(memory_manager, cpti, mopti) - total_nb_of_top_instances = len(cpti) - latency = _get_latency(scme) - xlims = _calculate_xlims(latency, section_start_percent, percent_shown) - fig, baxs = _create_figure_and_axes(total_nb_of_top_instances, xlims, fig_size) - fig.suptitle("Memory usage through time (Bytes)") - _plot_all_instances(baxs, cpti, mopti, memory_manager, xlims, show_human_bytes) - fig.savefig(fig_path) - logger.info(f"Saved memory usage fig to {fig_path}") - - -def _filter_memory_instances(memory_manager, show_dram): - cpti = memory_manager.cores_per_top_instance.copy() - mopti = memory_manager.memory_operands_per_top_instance.copy() - if not show_dram: - offchip_core_id = memory_manager.offchip_core_id - to_remove = [ - top_instance - for top_instance, cores in memory_manager.cores_per_top_instance.items() - if all(core.id == offchip_core_id for core in cores) - ] - for top_instance in to_remove: - cpti.pop(top_instance, None) - mopti.pop(top_instance, None) - return cpti, mopti - - -def _remove_zero_usage_instances(memory_manager, cpti, mopti): - to_remove = [] - for ti, stored_cumsum in memory_manager.top_instance_stored_cumsum.items(): - _, stored_bits = zip(*stored_cumsum, strict=False) - if max(stored_bits) == 0: - to_remove.append(ti) - for ti in to_remove: - cpti.pop(ti, None) - mopti.pop(ti, None) - return cpti, mopti - - -def _get_latency(scme): - assert scme.latency is not None, "SCME latency is None, cannot plot memory usage." - return int(scme.latency) - - -def _calculate_xlims(latency, section_start_percent, percent_shown): - x_starts = [int((start / 100) * latency) for start in section_start_percent] - x_ends = [ - int(((start + percent) / 100) * latency) - for (start, percent) in zip(section_start_percent, percent_shown, strict=False) - ] - return tuple(zip(x_starts, x_ends, strict=False)) - - -def _create_figure_and_axes(total_nb_of_top_instances, xlims, fig_size): - fig = plt.figure(figsize=fig_size) - gridspecs = GridSpec(total_nb_of_top_instances, 1) - baxs = [brokenaxes(xlims=xlims, subplot_spec=gridspec, wspace=0.05, d=0.005) for gridspec in gridspecs] - return fig, baxs - - -def _plot_all_instances(baxs, cpti, mopti, memory_manager, xlims, show_human_bytes): - for ( - ax, - (ti_cores, cores_for_this_ti), - (ti_memory_operands, memory_operands_for_this_ti), - ) in zip(baxs, cpti.items(), mopti.items(), strict=False): - assert ti_cores is ti_memory_operands, "Sanity check for same ordering of memory manager dicts failed." - _plot_single_instance( - ax, - ti_cores, - cores_for_this_ti, - memory_operands_for_this_ti, - memory_manager, - xlims, - show_human_bytes, - is_last=(ax is baxs[-1]), - ) - - -def _plot_single_instance( - ax, - ti, - cores_for_this_ti, - memory_operands_for_this_ti, - memory_manager, - xlims, - show_human_bytes, - is_last, -): - ti_cumsum = memory_manager.top_instance_stored_cumsum[ti] - ti_cumsum_bytes = ti_cumsum.astype(float) - ti_cumsum_bytes[:, 1] /= 8 - timesteps = ti_cumsum_bytes[:, 0] - stored_bytes = ti_cumsum_bytes[:, 1] - peak_usage_bytes = max(stored_bytes) - if not peak_usage_bytes > 0: - return - if min(stored_bytes) < 0: - logger.warn(f"We used negative amount of memory on top instance {ti}.") - ax.plot(timesteps, stored_bytes, drawstyle="steps-post") - ax.axhline( - y=peak_usage_bytes, - xmin=min(timesteps), - xmax=max(timesteps), - color="r", - linestyle="dashed", - ) - memory_capacity_bytes = memory_manager.top_instance_capacities[ti] / 8 - mem_text = ( - humanbytes(peak_usage_bytes) + " / " + humanbytes(np.array(memory_capacity_bytes)) - if show_human_bytes - else f"{peak_usage_bytes} B / {np.array(memory_capacity_bytes)} B" - ) - ax.text( - xlims[-1][1] - 1, - peak_usage_bytes, - mem_text, - color="r", - va="bottom", - ha="right", - fontsize=BIGGER_SIZE, - ) - core_memory_operand_zipped = zip(cores_for_this_ti, memory_operands_for_this_ti, strict=False) - formatted_str = "\n".join([f"{core}: {memory_operands}" for core, memory_operands in core_memory_operand_zipped]) - y_label = f"{ti.name}\n{formatted_str}" - ax.set_ylim(bottom=0, top=1.05 * peak_usage_bytes) - ax.ticklabel_format(axis="x", style="sci", scilimits=(0, 0)) - ax.ticklabel_format(axis="y", style="sci", scilimits=(0, 0)) - if not is_last: - ax.set_xticklabels([]) - ax.set_ylabel( - y_label, - rotation=0, - fontsize=BIGGER_SIZE, - labelpad=70, - loc="center", - va="center", - ) diff --git a/stream/visualization/steady_state_trace.py b/stream/visualization/steady_state_trace.py index 249140f0e..f7fc4d26a 100644 --- a/stream/visualization/steady_state_trace.py +++ b/stream/visualization/steady_state_trace.py @@ -337,7 +337,7 @@ def export_steady_state_trace( # noqa: PLR0912, PLR0915 if chosen_path is None: continue - one_transfer_lat = float(tta._transfer_latency_for_path(node, chosen_path)) + one_transfer_lat = float(tta.transfer_latency_for_path(node, chosen_path)) reuse_factor = tta.reuse_factors[node].X ssis = tta.ssis[node] active_transfer_lat = get_active_transfer_latency_for_path(node, chosen_path, reuse_factor, tta.ssis) @@ -476,7 +476,7 @@ def export_steady_state_trace( # noqa: PLR0912, PLR0915 if chosen_path is None: continue - one_transfer_lat = float(tta._transfer_latency_for_path(node, chosen_path)) + one_transfer_lat = float(tta.transfer_latency_for_path(node, chosen_path)) ssis = tta.ssis[node] reuse_factor = tta.reuse_factors[node].X active_transfer_lat = get_active_transfer_latency_for_path(node, chosen_path, reuse_factor, tta.ssis) diff --git a/stream/workload/utils.py b/stream/workload/utils.py index 38b329cae..15ef8f3fb 100644 --- a/stream/workload/utils.py +++ b/stream/workload/utils.py @@ -256,7 +256,8 @@ def collect_spatial_unrollings(workload: "Workload", mapping: "Mapping"): for node in workload.get_iteration_space_nodes(): node_mapping = mapping.get(node) assert node_mapping is not None, f"No mapping found for node {node.name}" - unrollings = workload.get_unique_dims_inter_core_tiling(node, mapping) + # A factor of one spreads nothing: the node then has no spatial variable of that dimension. + unrollings = tuple(u for u in workload.get_unique_dims_inter_core_tiling(node, mapping) if u[1] > 1) _reject_illegal_spatial_unroll(workload, node, unrollings, node_mapping.kernel is not None) spatial_unrollings[node] = unrollings diff --git a/tests/conftest.py b/tests/conftest.py index f933eac6c..86055b329 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,8 +1,27 @@ import json import pathlib +from collections.abc import Callable, Mapping +from typing import Any import pytest +from stream.api import SolveOptions, evaluate_mapping +from stream.inputs.testing.workload.make_2_conv import TwoConvWorkloadConfig +from stream.opt.allocation.constraint_optimization import families +from stream.opt.allocation.constraint_optimization import transfer_and_tensor_allocation as tta + +TWO_CONV = TwoConvWorkloadConfig( + batch_size=1, + in_channels=8, + height=32, + width=32, + out_channels_1=16, + out_channels_2=32, + kernel_size=3, + in_dtype="bf16", + weight_dtype="bf16", +) + def pytest_configure(config): config.addinivalue_line("markers", "slow: marks tests as slow (deselect with '-m \"not slow\"')") @@ -55,3 +74,62 @@ def pytest_terminal_summary(terminalreporter, exitstatus, config): # noqa: ARG0 return out = pathlib.Path(__file__).parent.parent / "metrics_current.json" out.write_text(json.dumps(dict(sorted(_metrics_store.items())), indent=2, sort_keys=True)) + + +def _solve_allocator( + hardware: str, + workload: Any, + output_path: str, + mapping: Any = None, + options: SolveOptions | None = None, + *, + families_available: Mapping[str, Callable[..., Any]] | None = None, + hook: str = "solve", +) -> tta.TransferAndTensorAllocator: + """Run ``evaluate_mapping`` and return the first allocator that reached its ``hook`` method (``solve`` or + ``_build_model``); ``families_available`` replaces the entry-point families.""" + captured: list[tta.TransferAndTensorAllocator] = [] + original = getattr(tta.TransferAndTensorAllocator, hook) + + def capture(self: tta.TransferAndTensorAllocator, *args: Any, **kwargs: Any) -> Any: + captured.append(self) + return original(self, *args, **kwargs) + + with pytest.MonkeyPatch.context() as patch: + if families_available is not None: + patch.setattr(families, "available_families", lambda: dict(families_available)) + patch.setattr(tta.TransferAndTensorAllocator, hook, capture) + evaluate_mapping(hardware, workload, output_path, mapping, options) + return captured[0] + + +def _model_size(alloc: tta.TransferAndTensorAllocator) -> tuple[int, int]: + """(variables, linear constraints) of the allocator's built model.""" + raw = alloc.model._model # type: ignore[attr-defined] + return sum(1 for _ in raw.variables()), sum(1 for _ in raw.linear_constraints()) + + +def _interval(alloc: tta.TransferAndTensorAllocator) -> float: + """The solved initiation interval: the iteration minus the overlap with the next one.""" + return alloc.model.value(alloc.quantities.get("iteration").expr) - alloc.overlap.X + + +@pytest.fixture(scope="session") +def two_conv() -> TwoConvWorkloadConfig: + return TWO_CONV + + +@pytest.fixture(scope="session") +def solved_allocator() -> Callable[..., tta.TransferAndTensorAllocator]: + """:func:`_solve_allocator`; session scoped so module-scoped fixtures can solve once.""" + return _solve_allocator + + +@pytest.fixture(scope="session") +def model_size() -> Callable[[tta.TransferAndTensorAllocator], tuple[int, int]]: + return _model_size + + +@pytest.fixture(scope="session") +def interval() -> Callable[[tta.TransferAndTensorAllocator], float]: + return _interval diff --git a/tests/integration/test_aie_codegen.py b/tests/integration/test_aie_codegen.py index 3798d3bc2..c27116f94 100644 --- a/tests/integration/test_aie_codegen.py +++ b/tests/integration/test_aie_codegen.py @@ -1,4 +1,4 @@ -"""SwiGLU lowers to an NPU2 design end to end, through the code generation backend that claims the array.""" +"""Workloads lower to an NPU2 design end to end, through the code generation backend that claims the array.""" from __future__ import annotations @@ -9,7 +9,9 @@ pytest.importorskip("snaxc", reason="the AIE dialects are a separate install, via stream-setup-aie") from stream.api import SolveOptions, generate_code # noqa: E402 +from stream.inputs.aie.mapping.make_gemm_mapping import make_gemm_mapping # noqa: E402 from stream.inputs.aie.mapping.make_swiglu_mapping import make_swiglu_mapping # noqa: E402 +from stream.inputs.aie.workload.make_onnx_gemm import make_gemm_workload # noqa: E402 from stream.inputs.aie.workload.make_onnx_swiglu import make_swiglu_workload # noqa: E402 ACCELERATOR = "stream/inputs/aie/hardware/whole_array_strix.yaml" @@ -25,3 +27,13 @@ def test_swiglu_generates_an_npu2_design(last_gemm_down: bool, tmp_path): module = str(generate_code(ACCELERATOR, workload, str(tmp_path), mapping, options).context.get("module")) assert "aie.device(npu2)" in module assert len(re.findall(r"aiex?\.", module)) > 100 + + +@pytest.mark.slow +def test_a_single_tile_gemm_generates_an_npu2_design(tmp_path): + """One kernel tile on one core: no split to unroll and no loop to iterate, so it runs once.""" + workload = make_gemm_workload(64, 64, 64, "bf16", "bf16") + mapping = make_gemm_mapping(64, 64, 64, 64, 64, 64, nb_rows_to_use=1, nb_cols_to_use=1) + options = SolveOptions(nb_cols_to_use=1, stage_options={"npu": "npu2"}) + module = str(generate_code(ACCELERATOR, workload, str(tmp_path), mapping, options).context.get("module")) + assert module.count("func.call @matmul_bf16_bf16_64_64_64") == 1 diff --git a/tests/integration/test_constraint_toggles.py b/tests/integration/test_constraint_toggles.py index 01786a131..7dfabc939 100644 --- a/tests/integration/test_constraint_toggles.py +++ b/tests/integration/test_constraint_toggles.py @@ -28,7 +28,6 @@ ) REL_TOL = 0.01 -_ALLOC_CREATE_SOLVER = "stream.opt.allocation.constraint_optimization.allocation.create_solver" _TTA_CREATE_SOLVER = "stream.opt.allocation.constraint_optimization.transfer_and_tensor_allocation.create_solver" _LICENSE_CHECK = "stream.api.GurobiBackend.check_license" _BUILD_TRANSFER_CONTEXT = ( @@ -223,14 +222,8 @@ def _tight_fifo_init(self, *args, **kwargs): @pytest.mark.slow def test_dma_channels_flip(): - """dma_channels guard: tight DMA limit causes infeasibility; disabling restores feasibility. - - The guard in _overlap_and_objective() is structurally wired. Proof: - - DMA channels=1 (all tiles) + dma_channels=True -> RuntimeError - - DMA channels=1 (all tiles) + dma_channels=False -> success - - Patch target: build_transfer_context in TTA's own namespace (imported at line 22). - """ + """A tight DMA limit is infeasible with dma_channels on and feasible with it off; build_transfer_context is + patched in the allocator's own namespace.""" # Enabled + tight limit -> infeasible with tempfile.TemporaryDirectory() as tmpdir: with patch(_BUILD_TRANSFER_CONTEXT, side_effect=_build_transfer_context_tight_dma): @@ -321,7 +314,6 @@ def test_cross_backend_parity(cs: ConstraintSelection): ort_factory = _make_ortools_factory() with tempfile.TemporaryDirectory() as tmpdir: with ( - patch(_ALLOC_CREATE_SOLVER, side_effect=ort_factory), patch(_TTA_CREATE_SOLVER, side_effect=ort_factory), patch(_LICENSE_CHECK), ): diff --git a/tests/integration/test_cross_backend.py b/tests/integration/test_cross_backend.py index fbb3353ed..a206bb490 100644 --- a/tests/integration/test_cross_backend.py +++ b/tests/integration/test_cross_backend.py @@ -46,8 +46,7 @@ # Relative tolerance for cross-backend comparison (1%) REL_TOL = 0.01 -# Patch targets — both allocator modules import create_solver into their own namespace -_ALLOC_CREATE_SOLVER = "stream.opt.allocation.constraint_optimization.allocation.create_solver" +# Patch target: the allocator module imports create_solver into its own namespace _TTA_CREATE_SOLVER = "stream.opt.allocation.constraint_optimization.transfer_and_tensor_allocation.create_solver" _LICENSE_CHECK = "stream.api.GurobiBackend.check_license" @@ -190,18 +189,13 @@ def test_gemm_gurobi_baseline(): @pytest.mark.slow def test_gemm_cross_backend(): - """ORToolsBackend (GSCIP) produces OPTIMAL on TETRA gemm instance, matching Gurobi within 1%. - - Covers objective quality on the gemm workload. - The pipeline's ``create_solver`` calls are patched in both allocator modules - so the ORToolsBackend is used transparently in place of GurobiBackend. - """ + """ORToolsBackend (GSCIP) solves the TETRA gemm instance to OPTIMAL, within 1% of Gurobi; the allocator + module's ``create_solver`` is patched to return an ORToolsBackend.""" results = [] ort_factory = _make_ortools_factory(mathopt.SolverType.GSCIP) with tempfile.TemporaryDirectory() as tmpdir: with ( - patch(_ALLOC_CREATE_SOLVER, side_effect=ort_factory), patch(_TTA_CREATE_SOLVER, side_effect=ort_factory), patch(_LICENSE_CHECK), ): @@ -232,7 +226,6 @@ def test_gemm_highs(): with tempfile.TemporaryDirectory() as tmpdir: with ( - patch(_ALLOC_CREATE_SOLVER, side_effect=highs_factory), patch(_TTA_CREATE_SOLVER, side_effect=highs_factory), patch(_LICENSE_CHECK), ): @@ -291,7 +284,6 @@ def test_swiglu_cross_backend(): with tempfile.TemporaryDirectory() as tmpdir: with ( - patch(_ALLOC_CREATE_SOLVER, side_effect=ort_factory), patch(_TTA_CREATE_SOLVER, side_effect=ort_factory), patch(_LICENSE_CHECK), ): @@ -318,7 +310,6 @@ def test_swiglu_highs(): with tempfile.TemporaryDirectory() as tmpdir: with ( - patch(_ALLOC_CREATE_SOLVER, side_effect=highs_factory), patch(_TTA_CREATE_SOLVER, side_effect=highs_factory), patch(_LICENSE_CHECK), ): diff --git a/tests/unit/test_block_options.py b/tests/unit/test_block_options.py index b13067c83..2581267c2 100644 --- a/tests/unit/test_block_options.py +++ b/tests/unit/test_block_options.py @@ -7,6 +7,8 @@ from dataclasses import dataclass # noqa: E402 from typing import Any # noqa: E402 +from xdsl.ir.affine import AffineMap # noqa: E402 + from stream.compiler.kernels.flash import PartialSoftmaxKernel # noqa: E402 from stream.compiler.kernels.gemm import GemmKernel # noqa: E402 from stream.compiler.kernels.library import KernelLibrary # noqa: E402 @@ -75,12 +77,19 @@ class _Group: intra_core_tiling: tuple -def _node(name): - return ComputationNode(type=name, name=name, inputs=(), outputs=(), operand_mapping=()) +# The iteration spaces the Gemm and elementwise parsers give: (m, k, n) and (m, n). +GEMM_MAPS = tuple( + AffineMap.from_callable(f) for f in (lambda m, k, n: (m, k), lambda m, k, n: (k, n), lambda m, k, n: (m, n)) +) +ELEMENTWISE_MAPS = (AffineMap.identity(2), AffineMap.identity(2)) + + +def _node(name, maps): + return ComputationNode(type=name, name=name, inputs=(), outputs=(), operand_mapping=maps) def _attention(lib): - gemm, softmax = _node("Attn_QK"), _node("Attn_Softmax") + gemm, softmax = _node("Attn_QK", GEMM_MAPS), _node("Attn_Softmax", ELEMENTWISE_MAPS) workload = _Workload({gemm: (M, K, N), softmax: (M, S)}) mapping = _Mapping( { @@ -92,7 +101,7 @@ def _attention(lib): def test_a_group_whose_only_block_rule_is_a_divisor_offers_no_candidates(): - gemm = _node("Gemm_Left") + gemm = _node("Gemm_Left", GEMM_MAPS) workload = _Workload({gemm: (M, K, N)}) mapping = _Mapping({gemm: _Entry(GemmKernel(m=64, k=64, n=64, layout="default", library=library()))}) assert block_options(workload, mapping, _Group(layers=("Gemm_Left",), intra_core_tiling=TILED)) == {} diff --git a/tests/unit/test_call_tile_tiling.py b/tests/unit/test_call_tile_tiling.py index ea65ccef4..8fa564ff2 100644 --- a/tests/unit/test_call_tile_tiling.py +++ b/tests/unit/test_call_tile_tiling.py @@ -8,7 +8,7 @@ class _Kernel: def __init__(self, *tile): self.tile = list(tile) - def call_tile(self): + def call_tile(self, node): return self.tile @@ -43,3 +43,13 @@ def test_only_dimensions_every_kernel_sizes_at_runtime_may_grow(): silu = _node(("s", "h"), _Kernel((1, 32, CallDim("n", runtime=True)), (0, 16, CallDim("m", runtime=True)))) factory = _Factory({"gemm": gemm, "silu": silu}, {"s": 256, "h": 512}) assert factory._runtime_dims(("gemm", "silu"), (("s", 16), ("h", 32))) == ("s",) + + +def test_a_dimension_no_kernel_addresses_is_tiled_one_at_a_time_outermost(): + """Attention's heads lead every operand of a batched matmul, which its kernel is called + once per head of; a unit axis needs no loop.""" + gemm = _node( + ("u", "b", "s", "e", "h"), _Kernel((3, 64, CallDim("k")), (4, 32, CallDim("n")), (2, 16, CallDim("m"))) + ) + factory = _Factory({"gemm": gemm}, {"u": 1, "b": 4, "s": 256, "e": 128, "h": 512}) + assert factory._call_tile_tiling(("gemm",)) == (("e", 64), ("h", 32), ("s", 16), ("b", 1)) diff --git a/tests/unit/test_chain_rereads.py b/tests/unit/test_chain_rereads.py new file mode 100644 index 000000000..07b223294 --- /dev/null +++ b/tests/unit/test_chain_rereads.py @@ -0,0 +1,78 @@ +import pytest + +pytest.importorskip("xdsl_aie", reason="the AIE dialects are a separate install, via stream-setup-aie") + +from xdsl.dialects.builtin import MemRefType, ModuleOp, bf16 # noqa: E402 +from xdsl.ir import Block, Region # noqa: E402 +from xdsl.pattern_rewriter import PatternRewriteWalker # noqa: E402 +from xdsl_aie.dialects.aie import ( # noqa: E402 + DeviceOp, + DMABDOp, + EndOp, + NextBDOp, + ObjectFifoOp, + RuntimeSequenceOp, + TileOp, +) +from xdsl_aie.dialects.aiex import DmaConfigureTaskForOp # noqa: E402 + +from stream.compiler.transforms.aie_convert_ofs import ChainRereads # noqa: E402 + +HEAD = 64 * 512 +# A head's key, re-read once per query block: the iteration dimension at stride zero. +REREAD = ([4, 8, 64, 64], [0, 64, 512, 1]) + + +def _sequence(offsets, sizes_strides=REREAD): + """A runtime sequence feeding one fifo a task per offset, as a loop unrolls them.""" + shim, mem = TileOp(0, 0), TileOp(0, 1) + fifo = ObjectFifoOp.from_referenced_type(shim, [mem], "key", 2, bf16, [64, 64]) + block = Block(arg_types=[MemRefType(bf16, [2 * len(offsets), 64, 512])]) + sizes, strides = sizes_strides + tasks = [ + DmaConfigureTaskForOp( + "key", + Region(Block([DMABDOp(block.args[0], offset, HEAD, sizes, strides), EndOp()])), + repeat_count=sizes[0] - 1, + ) + for offset in offsets + ] + block.add_ops(tasks) + return ModuleOp([DeviceOp(0, [shim, mem, fifo, RuntimeSequenceOp(Region(block))])]) + + +def _chain(module, descriptors, iterations=64): + PatternRewriteWalker(ChainRereads(descriptors, iterations), apply_recursively=False).rewrite_module(module) + return [op for op in module.walk() if isinstance(op, DmaConfigureTaskForOp)] + + +def test_a_loop_around_a_rereading_transfer_becomes_one_chained_task(): + """Each head's key re-read four times: four descriptors, each stepping through the heads.""" + (task,) = _chain(_sequence([h * HEAD for h in range(3)]), descriptors=16) + assert task.repeat_count.value.data == 2 + blocks = list(task.body.blocks) + assert len(blocks) == 4 + for block in blocks: + bd = block.first_op + assert list(bd.static_sizes.get_values()) == [3, 8, 64, 64] + assert list(bd.static_strides.get_values()) == [HEAD, 64, 512, 1] + assert isinstance(bd.next_op, NextBDOp if block is not blocks[-1] else EndOp) + + +def test_a_chain_longer_than_the_shim_descriptors_keeps_its_tasks(): + assert len(_chain(_sequence([h * HEAD for h in range(3)]), descriptors=3)) == 3 + + +def test_a_loop_longer_than_a_descriptor_iterates_keeps_its_tasks(): + assert len(_chain(_sequence([h * HEAD for h in range(3)]), descriptors=16, iterations=2)) == 3 + + +@pytest.mark.parametrize( + "offsets, sizes_strides", + [ + ([0, HEAD, 3 * HEAD], REREAD), # no one step between the heads + ([0, HEAD, 2 * HEAD], ([1, 8, 64, 64], [0, 64, 512, 1])), # read once a head + ], +) +def test_only_an_evenly_stepped_rereading_loop_chains(offsets, sizes_strides): + assert len(_chain(_sequence(offsets, sizes_strides), descriptors=16)) == 3 diff --git a/tests/unit/test_constraint_families.py b/tests/unit/test_constraint_families.py new file mode 100644 index 000000000..e1253ce1e --- /dev/null +++ b/tests/unit/test_constraint_families.py @@ -0,0 +1,108 @@ +from collections.abc import Callable +from typing import Any, ClassVar + +import pytest + +from stream.api import SolveOptions +from stream.inputs.testing.mapping.make_2_conv_mapping import make_2_conv_mapping +from stream.inputs.testing.workload.make_2_conv import TwoConvWorkloadConfig, make_2_conv_workload +from stream.opt.allocation.constraint_optimization import families +from stream.opt.allocation.constraint_optimization import transfer_and_tensor_allocation as tta +from stream.opt.allocation.constraint_optimization.families import SLOT_PRESSURE, load_families, parse_spec +from stream.opt.allocation.constraint_optimization.quantities import QuantityRegistry +from stream.opt.solver import ConstraintSelection + +ACCELERATOR = "stream/inputs/examples/hardware/tpu_like_quad_core.yaml" +PRESSURE_BOUND = 10**9 + + +class CapIteration: + """Test family: registers the iteration as a quantity and caps it at ``cap`` cycles.""" + + name: ClassVar[str] = "cap_iteration" + + def __init__(self, cap: float = 1e12) -> None: + self.cap = cap + + def declare(self, alloc: tta.TransferAndTensorAllocator, q: QuantityRegistry) -> None: + q.add(SLOT_PRESSURE, 0, index="test", upper_bound=PRESSURE_BOUND) + + def constrain(self, alloc: tta.TransferAndTensorAllocator, q: QuantityRegistry) -> None: + iteration = q.get("iteration").expr + q.add("capped_iteration", iteration) + alloc.model.add_constr(iteration <= self.cap, name="cap_iteration") + + +def solve( + solved_allocator: Callable[..., Any], two_conv: TwoConvWorkloadConfig, out: Any, options: SolveOptions +) -> Any: + workload, mapping = make_2_conv_workload(two_conv), make_2_conv_mapping(two_conv) + return solved_allocator( + ACCELERATOR, + workload, + str(out), + mapping, + options, + families_available={CapIteration.name: CapIteration}, + hook="_build_model", + ) + + +def test_spec_forms_parse_to_name_and_options() -> None: + assert parse_spec("memory_ports") == ("memory_ports", {}) + assert parse_spec({"memory_ports": {"burst": False}}) == ("memory_ports", {"burst": False}) + assert parse_spec({"memory_ports": None}) == ("memory_ports", {}) + with pytest.raises(ValueError, match="exactly one family"): + parse_spec({"a": {}, "b": {}}) + + +@pytest.fixture +def cap_only(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(families, "available_families", lambda: {CapIteration.name: CapIteration}) + + +def test_unknown_family_names_the_available_ones(cap_only: None) -> None: + with pytest.raises(KeyError, match="available: cap_iteration"): + load_families(["cap_iterations"]) + + +def test_family_selected_twice_is_rejected(cap_only: None) -> None: + with pytest.raises(ValueError, match="selected twice"): + load_families(["cap_iteration", {"cap_iteration": {"cap": 1}}]) + + +def test_no_families_loads_nothing() -> None: + assert load_families([]) == () + + +def test_solve_options_fold_families_into_the_constraint_selection() -> None: + assert SolveOptions().resolved_constraint_selection() is None + selection = SolveOptions(families=["memory_ports"]).resolved_constraint_selection() + assert selection == ConstraintSelection(families=("memory_ports",)) + + +@pytest.mark.slow +def test_family_adds_its_constraint_quantity_and_slot_bound( + solved_allocator: Callable, two_conv: TwoConvWorkloadConfig, model_size: Callable, tmp_path: Any +) -> None: + base = solve(solved_allocator, two_conv, tmp_path / "base", SolveOptions()) + options = SolveOptions(families=[{"cap_iteration": {"cap": 1e9}}]) + with_family = solve(solved_allocator, two_conv, tmp_path / "family", options) + base_vars, base_cons = model_size(base) + family_vars, family_cons = model_size(with_family) + assert family_vars == base_vars + assert family_cons == base_cons + 1 + assert "capped_iteration" in with_family.quantities + assert "capped_iteration" not in base.quantities + assert with_family._family_slot_pressure_bound() == PRESSURE_BOUND + assert base._family_slot_pressure_bound() == 0 + + +def test_a_bare_family_name_is_rejected() -> None: + with pytest.raises(TypeError, match="list of names"): + load_families("cap_iteration") + + +def test_solve_options_reject_a_bare_family_name() -> None: + with pytest.raises(TypeError, match="list of names"): + SolveOptions(families="memory_ports").resolved_constraint_selection() diff --git a/tests/unit/test_constraint_selection.py b/tests/unit/test_constraint_selection.py index a99caf263..f6e041bda 100644 --- a/tests/unit/test_constraint_selection.py +++ b/tests/unit/test_constraint_selection.py @@ -106,25 +106,21 @@ def test_importable(): def _make_tta_stub(constraint_selection, *, bind_objective=False): - """Create a minimal TTA-like object with constraint_selection set. - - We cannot instantiate a real TTA without a Workload, so we create - a mock that has the constraint_selection attribute and delegates - to the real _create_constraints / _overlap_and_objective methods. - - bind_objective: if True, also bind the real _set_total_latency_and_objective - so its DMA conditional logic executes. Only needed for the objective test. - """ + """A mock allocator that runs the real _create_constraints and _dma_channels_and_objective; bind_objective + also binds the real _set_total_latency_and_objective so its DMA branch runs.""" + from stream.opt.allocation.constraint_optimization.quantities import QuantityRegistry from stream.opt.allocation.constraint_optimization.transfer_and_tensor_allocation import ( TransferAndTensorAllocator, ) tta = MagicMock(spec=TransferAndTensorAllocator) tta.constraint_selection = constraint_selection + tta.quantities = QuantityRegistry() tta.shared_bandwidth = {} + tta.fill = 0 # Bind the real dispatch methods so if-guards execute tta._create_constraints = TransferAndTensorAllocator._create_constraints.__get__(tta) - tta._overlap_and_objective = TransferAndTensorAllocator._overlap_and_objective.__get__(tta) + tta._dma_channels_and_objective = TransferAndTensorAllocator._dma_channels_and_objective.__get__(tta) if bind_objective: tta._set_total_latency_and_objective = TransferAndTensorAllocator._set_total_latency_and_objective.__get__(tta) return tta @@ -168,9 +164,7 @@ def test_dma_guard(): """TTA with dma_channels=False does not call _add_dma_usage_constraints.""" cs = ConstraintSelection(dma_channels=False) tta = _make_tta_stub(cs) - tta.max_slot = 0 - tta.big_m = 10 - tta._overlap_and_objective() + tta._dma_channels_and_objective() tta._add_dma_usage_constraints.assert_not_called() @@ -178,14 +172,12 @@ def test_dma_enabled(): """TTA with dma_channels=True calls _add_dma_usage_constraints.""" cs = ConstraintSelection(dma_channels=True) tta = _make_tta_stub(cs) - tta.max_slot = 0 - tta.big_m = 10 - tta._overlap_and_objective() + tta._dma_channels_and_objective() tta._add_dma_usage_constraints.assert_called_once() def test_dma_objective_no_dma_terms(): - """When dma_channels=False, primary objective = total_lat only (no DMA vars).""" + """When dma_channels=False, primary objective = total_latency only (no DMA vars).""" cs = ConstraintSelection(dma_channels=False, offchip_traffic_cost=False) tta = _make_tta_stub(cs, bind_objective=True) # Set up minimal mocks for _set_total_latency_and_objective @@ -202,7 +194,7 @@ def test_dma_objective_no_dma_terms(): tta.transfer_nodes = [] tta.possible_transfer_allocations = {} tta._set_total_latency_and_objective() - # Verify lexicographic objectives were set with primary = total_lat only + # Verify lexicographic objectives were set with primary = total_latency only mock_model.set_lexicographic_objectives.assert_called_once() objectives = mock_model.set_lexicographic_objectives.call_args[0][0] primary = next(o for o in objectives if o.name == "latency") @@ -303,9 +295,7 @@ def test_skip_warnings(caplog): tta = _make_tta_stub(cs) with caplog.at_level(logging.WARNING): tta._create_constraints() - tta.max_slot = 0 - tta.big_m = 10 - tta._overlap_and_objective() + tta._dma_channels_and_objective() assert "memory_capacity" in caplog.text.lower() assert "object_fifo_depth" in caplog.text.lower() assert "buffer_descriptors" in caplog.text.lower() @@ -320,9 +310,7 @@ def test_all_enabled_calls_all(): tta._memory_capacity_constraints.assert_called_once() tta._object_fifo_depth_constraints.assert_called_once() tta._buffer_descriptor_constraints.assert_called_once() - tta.max_slot = 0 - tta.big_m = 10 - tta._overlap_and_objective() + tta._dma_channels_and_objective() tta._add_dma_usage_constraints.assert_called_once() diff --git a/tests/unit/test_delivered_placement.py b/tests/unit/test_delivered_placement.py new file mode 100644 index 000000000..e701d344e --- /dev/null +++ b/tests/unit/test_delivered_placement.py @@ -0,0 +1,54 @@ +"""A tensor a transfer produces is held only on the cores the transfer's chosen path writes to.""" + +from __future__ import annotations + +from xdsl.dialects.builtin import bf16 + +from stream.opt.allocation.constraint_optimization.transfer_and_tensor_allocation import TransferAndTensorAllocator +from stream.opt.solver import SolverBackend, SolverParams, SolverVarType, create_solver +from stream.workload.workload import Tensor + + +class _Node: + def __init__(self, name: str, outputs=()): + self.name, self.outputs = name, list(outputs) + + +class _Tile: + def __init__(self, core_id: int, name: str): + self.id, self.name = core_id, name + + +def _placed(one_tile: bool) -> tuple[str, ...]: + """Where the key lands when the shim's path writes one memory tile, or two.""" + alloc = object.__new__(TransferAndTensorAllocator) + alloc.model = model = create_solver(SolverBackend.ORTOOLS_GSCIP, "delivered") + model.set_param(SolverParams.VERBOSITY, 0) + alloc._name_counter = {} + key = Tensor.create("key", bf16, (64, 64)) + transfer = _Node("Transfer(key)", [key]) + tile_a, tile_b = _Tile(1, "a"), _Tile(7, "b") + placements = ((tile_a,), (tile_a, tile_b)) + alloc.possible_tensor_allocations = {key: placements} + alloc.tensor_fixed = [] + alloc.x_tensor_choice = {(key, p): model.add_var(vtype=SolverVarType.BINARY, name=f"x{len(p)}") for p in placements} + model.add_constr(model.quicksum(v._raw for v in alloc.x_tensor_choice.values()) == 1) + path = "one" if one_tile else "two" + alloc.y_path_choice = {(transfer, path): model.add_var(vtype=SolverVarType.BINARY, name="y")} + model.add_constr(alloc.y_path_choice[(transfer, path)] == 1) + alloc.choice_dst_cores = {(transfer, path): {tile_a} if one_tile else {tile_a, tile_b}} + alloc.choice_has_empty_path = {(transfer, path): False} + alloc._add_destination_tensor_coherence_constraints(transfer, (path,)) + # Spreading the key would let the next transfer read it from more tiles at once. + model.set_objective(-alloc.x_tensor_choice[(key, placements[1])]._raw) + model.optimize() + chosen = next(p for p in placements if alloc.x_tensor_choice[(key, p)].X > 0.5) + return tuple(t.name for t in chosen) + + +def test_a_tensor_is_not_held_where_its_transfer_never_writes(): + assert _placed(one_tile=True) == ("a",) + + +def test_a_tensor_may_be_held_on_every_tile_its_transfer_writes(): + assert _placed(one_tile=False) == ("a", "b") diff --git a/tests/unit/test_dma_channels.py b/tests/unit/test_dma_channels.py index 44a2e95c4..0aa2c857a 100644 --- a/tests/unit/test_dma_channels.py +++ b/tests/unit/test_dma_channels.py @@ -138,12 +138,12 @@ def _allocator(context, broadcast: bool = False): def test_a_neighbour_transfer_reads_in_place_so_its_path_latency_is_zero(aie2): """A compute-to-compute transfer between neighbours is served out of shared memory: it spends no - channel (above) and moves no bytes over a link, so ``_transfer_latency_for_path`` returns 0.""" + channel (above) and moves no bytes over a link, so ``transfer_latency_for_path`` returns 0.""" alloc = _allocator(aie2) tr = _Transfer(TransferType.COMPUTE_TO_COMPUTE) plan = _Plan([_core(0, "compute", 0, 2)], [_core(1, "compute", 0, 3)]) # north neighbour assert alloc._choice_shares_memory(tr, plan) is True - assert alloc._transfer_latency_for_path(tr, plan) == 0 + assert alloc.transfer_latency_for_path(tr, plan) == 0 def test_a_transfer_across_the_array_is_not_shared(aie2): diff --git a/tests/unit/test_eltwise_mul_kernel.py b/tests/unit/test_eltwise_mul_kernel.py index 26063f82a..b27ca10c2 100644 --- a/tests/unit/test_eltwise_mul_kernel.py +++ b/tests/unit/test_eltwise_mul_kernel.py @@ -21,3 +21,22 @@ def kernel(m: int, n: int, layout: str = "default"): def test_every_tile_calls_the_runtime_sized_multiply(m: int, n: int): assert kernel(m, n).function_name == "eltwise_mul_bf16_vector_size" assert kernel(m, n).linkwith_name == "mul.o" + + +def test_a_call_links_the_object_built_for_the_tile_it_takes(tmp_path): + """A call handed half its declared row links the half-row object: an object compiled for a + fixed element count would otherwise run past the tile it was given.""" + from types import SimpleNamespace + + from xdsl.dialects.builtin import MemRefType, bf16 + + from stream.compiler.kernels.eltwise_mul import EltwiseMulKernel + from stream.compiler.kernels.library import KernelLibrary + + toml = (Path(__file__).parents[2] / "stream/inputs/aie/kernels/aie2p.toml").read_text() + sized = tmp_path / "sized.toml" + sized.write_text(toml.replace('object = "mul.o"', 'object = "mul_{m}x{n}.o"')) + multiply = EltwiseMulKernel(m=1, n=4096, layout="contiguous", library=KernelLibrary.load(sized)) + tile = SimpleNamespace(type=MemRefType(bf16, [1, 2048])) + assert multiply.linkwith_name == "mul_1x4096.o" + assert multiply.call_object(SimpleNamespace(inputs=[tile, tile, tile])) == "mul_1x2048.o" diff --git a/tests/unit/test_fifo_depths.py b/tests/unit/test_fifo_depths.py index 4649ad6d3..33915fa45 100644 --- a/tests/unit/test_fifo_depths.py +++ b/tests/unit/test_fifo_depths.py @@ -38,6 +38,7 @@ def test_no_descriptors_keeps_defaults(): def test_budget_is_debited_across_fifos(): policy = FifoDepths(budgets={(0, 1): TileBudget(bytes_free=70000, bds_free=40)}) assert policy.deepen((2, 2), tiles((0, 1), (0, 2)), 16384) == (4, 2) + assert policy.deepen((2, 2), tiles((0, 1), (0, 2)), 16384) == (3, 2) assert policy.deepen((2, 2), tiles((0, 1), (0, 2)), 16384) == (2, 2) @@ -59,3 +60,21 @@ def test_object_bytes(): from xdsl.dialects.builtin import BFloat16Type assert object_bytes(elem_bits(BFloat16Type()), (64, 64)) == 8192 + + +def test_a_reader_deepens_as_far_as_its_slack_allows(): + """Two more 8 KB objects do not fit half of 31 KB, one does.""" + policy = FifoDepths(budgets={(0, 4): TileBudget(bytes_free=31744, bds_free=7)}) + assert policy.deepen((4, 2), tiles((0, 1), (0, 4)), 8192) == (4, 3) + + +def test_a_broadcast_deepens_its_readers_together(): + """Every reader takes the depth the tightest one affords.""" + policy = FifoDepths( + budgets={ + (0, 2): TileBudget(bytes_free=1 << 20, bds_free=40), + (1, 2): TileBudget(bytes_free=20000, bds_free=40), + } + ) + assert policy.deepen((4, 2, 2), tiles((0, 1), (0, 2), (1, 2)), 8192) == (4, 3, 3) + assert policy.budgets[(0, 2)].bytes_free == (1 << 20) - 8192 diff --git a/tests/unit/test_flash_kernels.py b/tests/unit/test_flash_kernels.py index 6053e1d8a..b5a1116a9 100644 --- a/tests/unit/test_flash_kernels.py +++ b/tests/unit/test_flash_kernels.py @@ -30,6 +30,16 @@ def column(): Score GEMM, online softmax and value GEMM on three neighbouring tiles, each in the query loop over the key loop the blocked design nests them in. """ + return _column(split=True) + + +@pytest.fixture(scope="module") +def unsplit(): + """The same column where one query block is all there is, so no core splits it.""" + return _column(split=False) + + +def _column(split): from xdsl.dialects.arith import ConstantOp from xdsl.dialects.builtin import IndexType, IntegerAttr, ModuleOp, bf16 from xdsl.dialects.scf import ForOp, YieldOp @@ -61,9 +71,10 @@ def column(): query, key, head = LayerDim(0), LayerDim(2), LayerDim(3) query_loop = StrensorVar(StrensorVarType.TEMPORAL, BLOCKS // COLUMNS, query) key_loop = StrensorVar(StrensorVarType.TEMPORAL, BLOCKS, key) - outer = (query_loop, StrensorVar(StrensorVarType.SPATIAL, COLUMNS, query), key_loop) + outer = (query_loop, StrensorVar(StrensorVarType.SPATIAL, COLUMNS, query), key_loop) if split else (key_loop,) # A point variable carries its position, so this is the first column of the four. - position = StrensorSpaceAttr(StrensorSpace((StrensorVar(StrensorVarType.POINT, 0, query),))) + points = (StrensorVar(StrensorVarType.POINT, 0, query),) if split else () + position = StrensorSpaceAttr(StrensorSpace(points)) def node(kernel_name, operands, kernel_dims): """One computation node, its operands taken from fifos the way codegen leaves them.""" @@ -85,7 +96,7 @@ def node(kernel_name, operands, kernel_dims): def core(row, body): """The node on its own tile, wrapped in the key loop and then the query loop.""" ops = body - for var in (key_loop, query_loop): + for var in (key_loop, query_loop) if split else (key_loop,): bounds = [ConstantOp.from_int_and_width(bound, IndexType()) for bound in (0, var.size, 1)] loop = ForOp(*bounds, [], Region(Block([*ops, YieldOp()], arg_types=[IndexType()]))) loop.attributes["layer_dim"] = StrensorVarAttr(var) @@ -239,4 +250,11 @@ def test_the_declared_state_is_the_size_the_core_buffer_holds(): # _core_buffer allocates SCALE_ROWS * m elements; the declaration says rows per step, and # the node supplies the query extent that a split then divides. assert state.rows * kernel.m == SCALE_ROWS * kernel.m - assert (state.carried_over, state.indexed_by) == (1, 0) + assert (state.carried_over, state.indexed_by) == ("n", "m") + + +def test_an_unsplit_step_still_hands_its_scale_on(unsplit): + """With no split to say which cores pair up, the two halves of a step share one column.""" + from xdsl_aie.dialects.aie import ObjectFifoOp + + assert [op.sym_name.data for op in unsplit.walk() if isinstance(op, ObjectFifoOp)] == ["flash_scale_0_1_0_2"] diff --git a/tests/unit/test_get_ir.py b/tests/unit/test_get_ir.py index 22dc17bec..6b08f6849 100644 --- a/tests/unit/test_get_ir.py +++ b/tests/unit/test_get_ir.py @@ -223,11 +223,13 @@ def test_post_solve_latency_values(self): scheduler.latency_total = 1000 scheduler.latency_per_iteration = 250 scheduler.overlap_between_iterations = 50 + scheduler.latency_fill = 120 result = scheduler.get_ir() assert result["latency"]["total"] == 1000 assert result["latency"]["per_iteration"] == 250 assert result["latency"]["overlap_between_iterations"] == 50 + assert result["latency"]["fill"] == 120 def test_backend_and_constraint_selection_in_ir(self): """Test 3: get_ir() includes backend (str) and constraint_selection (dict or None).""" diff --git a/tests/unit/test_hardware_described_costs.py b/tests/unit/test_hardware_described_costs.py index f1ff13da1..903c99804 100644 --- a/tests/unit/test_hardware_described_costs.py +++ b/tests/unit/test_hardware_described_costs.py @@ -20,7 +20,7 @@ def test_the_strix_array_declares_its_off_chip_bandwidth_and_reconfiguration(): accelerator = _parse("stream/inputs/aie/hardware/whole_array_strix.yaml") assert set(accelerator.bandwidth) == {accelerator.offchip_core_id} assert isinstance(accelerator.bandwidth[accelerator.offchip_core_id], BandwidthModel) - assert accelerator.reconfiguration == {"cycles_per_column": 69000, "reset_cycles": 63000} + assert accelerator.reconfiguration == {"cycles_per_column": 40000, "reset_cycles": 63000} def test_hardware_without_either_declares_neither(): diff --git a/tests/unit/test_held_buffers.py b/tests/unit/test_held_buffers.py new file mode 100644 index 000000000..7fb70a806 --- /dev/null +++ b/tests/unit/test_held_buffers.py @@ -0,0 +1,68 @@ +"""How many buffers a tile holds of a tensor whose window a loop outside its reuse moves on.""" + +from __future__ import annotations + +import pytest + +pytest.importorskip("snaxc", reason="the AIE dialects are a separate install, via stream-setup-aie") + +from xdsl.context import Context # noqa: E402 +from xdsl.dialects.builtin import StringAttr, bf16 # noqa: E402 +from xdsl.parser import Parser # noqa: E402 + +from stream.compiler.dialects.stream import ( # noqa: E402 + Stream, + StrensorSpace, + StrensorType, + StrensorVar, + StrensorVarType, +) +from stream.compiler.transforms.aie_convert_ofs import ChannelToObjectFifoPass # noqa: E402 +from stream.datatypes import LayerDim # noqa: E402 + +HEAD, QUERY, KEY = LayerDim(0), LayerDim(1), LayerDim(2) + + +def key(buffers: int = 2, heads: int = 4) -> StrensorType: + """A head's whole key on a core, reused over the query loop and moved on by the heads loop.""" + space = StrensorSpace( + ( + StrensorVar(StrensorVarType.TEMPORAL, heads, HEAD), + StrensorVar(StrensorVarType.TEMPORAL, 8, QUERY), + StrensorVar(StrensorVarType.KERNEL, 64, KEY), + ) + ) + return StrensorType(bf16, space, [StringAttr("tile_0_2")], reuse_index=2, buffers=buffers) + + +@pytest.mark.parametrize("buffers", [1, 2]) +def test_a_moving_window_is_held_in_the_buffers_the_solve_chose(buffers): + assert ChannelToObjectFifoPass.held_count(key(buffers)) == buffers + + +def test_a_window_nothing_moves_on_is_held_once(): + space = StrensorSpace( + (StrensorVar(StrensorVarType.TEMPORAL, 8, QUERY), StrensorVar(StrensorVarType.KERNEL, 64, KEY)) + ) + assert ChannelToObjectFifoPass.held_count(StrensorType(bf16, space, [StringAttr("tile_0_2")], reuse_index=2)) == 1 + + +@pytest.mark.parametrize("buffers", [1, 2]) +def test_the_buffer_count_survives_printing(buffers): + context = Context() + context.load_dialect(Stream) + assert Parser(context, str(key(buffers))).parse_attribute() == key(buffers) + + +def test_a_loop_of_one_keeps_nothing_more_local(): + """A query block that is the whole of a core's rows leaves its loop one step long.""" + space = StrensorSpace( + ( + StrensorVar(StrensorVarType.TEMPORAL, 4, HEAD), + StrensorVar(StrensorVarType.TEMPORAL, 1, QUERY), + StrensorVar(StrensorVarType.TEMPORAL, 2, KEY), + StrensorVar(StrensorVarType.KERNEL, 32, QUERY), + StrensorVar(StrensorVarType.KERNEL, 64, KEY), + ) + ) + assert StrensorType(bf16, space, [StringAttr("tile_0_2")], reuse_index=4).get_local_shape() == (2,) diff --git a/tests/unit/test_ir_models.py b/tests/unit/test_ir_models.py index bfba3327f..6f5bcb550 100644 --- a/tests/unit/test_ir_models.py +++ b/tests/unit/test_ir_models.py @@ -664,6 +664,21 @@ def test_json_round_trip(self): "iteration_count": 4200, }, "performance": { + "memory_ports": [ + { + "resource": "dram.rw_port_1", + "core_ids": [6], + "bw_bits_per_cycle": 64.0, + "bits_per_iteration": 196864.0, + "req_bw_aver": 64.0, + "real_cycle": 3076.0, + "allowed_cycle": 3076.0, + "stall_or_slack": 0.0, + "utilization": 1.0, + "burst_utilization": 1.0, + "burst_slot": 2, + } + ], "per_node": { "MatMul": { "kind": "compute", @@ -738,6 +753,10 @@ def test_optimality_gap_and_status_survive(self): assert ir.solve.status == "TIME_LIMIT" assert ir.algorithmic_view().solve.mip_gap == 0.04 + def test_port_activity_survives(self): + (port,) = self._ir().performance.memory_ports + assert (port.resource, port.utilization, port.stall_or_slack) == ("dram.rw_port_1", 1.0, 0.0) + def test_aggregate_extras_survive(self): aggregate = self._ir().performance.aggregate diff --git a/tests/unit/test_kernel_binding.py b/tests/unit/test_kernel_binding.py index 873bdc421..e025fbde8 100644 --- a/tests/unit/test_kernel_binding.py +++ b/tests/unit/test_kernel_binding.py @@ -2,8 +2,13 @@ pytest.importorskip("snaxc", reason="the AIE kernels are a separate install, via stream-setup-aie") +from xdsl.dialects.builtin import bf16 # noqa: E402 +from xdsl.ir.affine import AffineMap # noqa: E402 + from stream.compiler.kernels.library import KernelLibrary # noqa: E402 from stream.compiler.kernels.registry import AIE_KERNELS # noqa: E402 +from stream.workload.node import ComputationNode # noqa: E402 +from stream.workload.tensor import Tensor # noqa: E402 LIBRARY = KernelLibrary.from_dict( { @@ -19,13 +24,46 @@ ) -def test_a_kernel_binds_the_library_dimensions_to_its_node_dimensions(): +def _gemm_node(heads: int, maps): + lead = (heads,) if heads else () + a, b = Tensor.create("a", bf16, (*lead, 32, 64)), Tensor.create("b", bf16, (*lead, 64, 16)) + out = Tensor.create("c", bf16, (*lead, 32, 16)) + return ComputationNode(type="MatMul", name="mm", inputs=(a, b), outputs=(out,), operand_mapping=maps) + + +@pytest.mark.parametrize( + "node, expected", + [ + # Gemm iterates (m, k, n). + ( + _gemm_node( + 0, + tuple( + AffineMap.from_callable(f) + for f in (lambda m, k, n: (m, k), lambda m, k, n: (k, n), lambda m, k, n: (m, n)) + ), + ), + [(1, 64, "k"), (2, 16, "n"), (0, 32, "m")], + ), + # A batched MatMul iterates (h, m, n, k): the heads lead and the kernel never sees them. + ( + _gemm_node( + 4, + tuple( + AffineMap.from_callable(f) + for f in (lambda h, m, n, k: (h, m, k), lambda h, m, n, k: (h, k, n), lambda h, m, n, k: (h, m, n)) + ), + ), + [(3, 64, "k"), (2, 16, "n"), (1, 32, "m")], + ), + ], + ids=["gemm", "batched_matmul"], +) +def test_a_kernel_binds_the_library_dimensions_to_its_node_dimensions(node, expected): + """By what each dimension does, the output's rows and columns and the contraction, not by + where a parser happened to put it.""" gemm = AIE_KERNELS["gemm"](m=32, k=64, n=16, layout="default", library=LIBRARY) - assert [(position, size, d.name) for position, size, d in gemm.call_tile()] == [ - (1, 64, "k"), - (2, 16, "n"), - (0, 32, "m"), - ] + assert [(position, size, d.name) for position, size, d in gemm.call_tile(node)] == expected assert gemm.spec.symbol == "matmul_bf16_bf16" diff --git a/tests/unit/test_kernel_state.py b/tests/unit/test_kernel_state.py index b08f429b7..c779411a8 100644 --- a/tests/unit/test_kernel_state.py +++ b/tests/unit/test_kernel_state.py @@ -42,7 +42,9 @@ def _softmax_with_state(): outputs=(probs,), operand_mapping=(IDENTITY, IDENTITY), ) - state, mapping = state_tensor(node, StateOperand("flash_state", ROWS, carried_over=1, indexed_by=0)) + state, mapping = state_tensor( + node, StateOperand("flash_state", ROWS, carried_over="n", indexed_by="m"), {"m": 0, "n": 1} + ) node = replace(node, inputs=(scores, state), operand_mapping=(IDENTITY, mapping, IDENTITY)) workload = Workload([InEdge(name="scores_in", outputs=(scores,)), node, OutEdge(name="probs_out", inputs=(probs,))]) return workload, node, state, probs diff --git a/tests/unit/test_matmul_parser.py b/tests/unit/test_matmul_parser.py index ec1ca7ce3..468b598f0 100644 --- a/tests/unit/test_matmul_parser.py +++ b/tests/unit/test_matmul_parser.py @@ -65,6 +65,16 @@ def test_matmul_broadcast_batch_maps_to_constant(): assert a_map.results[0].value == 0 +def test_a_batch_axis_of_one_on_every_operand_is_the_nodes_batch_dimension(): + """Attention over one head: the head axis is no broadcast, so every operand indexes it by the node's + batch dimension, as it would for two heads.""" + from xdsl.ir.affine import AffineDimExpr + + node = _parse_single(_matmul_model((1, 8, 4), (1, 4, 6), (1, 8, 6))) + for operand_map in node.operand_mapping: + assert operand_map.results[0] == AffineDimExpr(0) + + def test_matmul_footprint_matches_numpy_contraction_shape(): """The affine maps agree with numpy.matmul on the produced/consumed index extents.""" a, b = np.random.default_rng(0).random((2, 5, 3)), np.random.default_rng(1).random((2, 3, 7)) diff --git a/tests/unit/test_memory_ports.py b/tests/unit/test_memory_ports.py new file mode 100644 index 000000000..d8ee557c7 --- /dev/null +++ b/tests/unit/test_memory_ports.py @@ -0,0 +1,178 @@ +from collections.abc import Callable +from types import SimpleNamespace +from typing import Any +from unittest.mock import MagicMock + +import pytest +from zigzag.cost_model.cost_model import CostModelEvaluation +from zigzag.cost_model.port_activity import PortActivity +from zigzag.datatypes import LayerOperand +from zigzag.hardware.architecture.memory_port import DataDirection + +from stream.api import SolveOptions +from stream.inputs.testing.mapping.make_2_conv_mapping import make_2_conv_mapping +from stream.inputs.testing.workload.make_2_conv import TwoConvWorkloadConfig, make_2_conv_workload +from stream.opt.allocation.constraint_optimization import transfer_and_tensor_allocation as tta +from stream.opt.allocation.constraint_optimization.families import traffic +from stream.opt.allocation.constraint_optimization.families.memory_ports import MemoryPorts +from stream.opt.allocation.constraint_optimization.quantities import QuantityRegistry +from stream.workload.steady_state.iteration_space import LoopEffect + +ACCELERATOR = "stream/inputs/examples/hardware/tpu_like_quad_core.yaml" +DRAM = (6, "dram", "rw_port_1") +Solve = Callable[[list[Any]], tta.TransferAndTensorAllocator] + + +@pytest.fixture(scope="module") +def solve( + tmp_path_factory: pytest.TempPathFactory, solved_allocator: Callable[..., Any], two_conv: TwoConvWorkloadConfig +) -> Solve: + def run(family_specs: list[Any]) -> tta.TransferAndTensorAllocator: + return solved_allocator( + ACCELERATOR, + make_2_conv_workload(two_conv), + str(tmp_path_factory.mktemp("two_conv")), + make_2_conv_mapping(two_conv), + SolveOptions(families=family_specs), + families_available={MemoryPorts.name: MemoryPorts}, + ) + + return run + + +@pytest.fixture(scope="module") +def base(solve: Solve) -> tta.TransferAndTensorAllocator: + return solve([]) + + +@pytest.fixture(scope="module") +def ports(solve: Solve) -> tta.TransferAndTensorAllocator: + return solve(["memory_ports"]) + + +def rate(alloc: tta.TransferAndTensorAllocator, key: Any) -> float: + return alloc.quantities.get("port_rate", key).expr + + +@pytest.mark.slow +def test_memory_ports_add_constraints_but_no_variables(base: Any, ports: Any, model_size: Callable) -> None: + (base_vars, base_cons), (port_vars, port_cons) = model_size(base), model_size(ports) + assert port_vars == base_vars + assert port_cons > base_cons + + +@pytest.mark.slow +def test_every_port_moves_its_demand_within_the_interval(ports: Any, interval: Callable) -> None: + for key, demand in ports.quantities.indexed("port_demand").items(): + assert ports.model.value(demand.expr) <= rate(ports, key) * interval(ports) + 1e-6 + + +@pytest.mark.slow +def test_every_port_moves_each_slot_demand_within_that_slot(ports: Any) -> None: + for (key, slot), demand in ports.quantities.indexed("port_demand_slot").items(): + slot_latency = ports.model.value(ports.quantities.get("slot_latency", slot).expr) + assert ports.model.value(demand.expr) <= rate(ports, key) * slot_latency + 1e-6 + + +@pytest.mark.slow +def test_offchip_port_binds_and_stretches_the_schedule(base: Any, ports: Any, interval: Callable) -> None: + assert ports.total_latency.X > base.total_latency.X + utilisation = { + key: ports.model.value(d.expr) / (rate(ports, key) * interval(ports)) + for key, d in ports.quantities.indexed("port_demand").items() + } + assert max(utilisation, key=lambda key: utilisation[key]) == DRAM + + +@pytest.mark.slow +def test_no_slot_outgrows_the_slot_pressure_a_port_declares(ports: Any) -> None: + worst = max(q.upper_bound for q in ports.quantities.indexed("slot_pressure").values()) + assert all(ports.model.value(q.expr) <= worst for q in ports.quantities.indexed("slot_latency").values()) + + +@pytest.mark.slow +def test_burst_off_builds_one_constraint_per_port(solve: Solve, base: Any, model_size: Callable) -> None: + interval_only = solve([{"memory_ports": {"burst": False}}]) + n_ports = len(interval_only.quantities.indexed("port_demand")) + assert model_size(interval_only)[1] == model_size(base)[1] + n_ports + + +@pytest.mark.slow +def test_without_bounds_the_family_only_reports(solve: Solve, base: Any, model_size: Callable) -> None: + report_only = solve([{"memory_ports": {"interval": False, "burst": False}}]) + assert model_size(report_only) == model_size(base) + assert report_only.total_latency.X == base.total_latency.X + rows = report_only.compute_performance_stats()["memory_ports"] + assert {row["kind"] for row in rows} == {"memory_port", "link"} + assert rows[0]["utilization"] == max(row["utilization"] for row in rows) + + +def test_an_accelerator_without_port_models_gets_no_port_constraint() -> None: + alloc = MagicMock() + registry = QuantityRegistry() + MemoryPorts().constrain(alloc, registry) + alloc.model.add_constr.assert_not_called() + + +def port_registry() -> tuple[MagicMock, QuantityRegistry]: + alloc = MagicMock() + registry = QuantityRegistry() + for name, value in (("iteration", 10), ("overlap", 2)): + registry.add(name, value) + registry.add("port_rate", 64.0, index=DRAM) + registry.add("port_demand", 256.0, index=DRAM) + registry.add("port_demand_slot", 256.0, index=(DRAM, 0)) + registry.add("slot_latency", 4, index=0) + return alloc, registry + + +def constraint_names(alloc: MagicMock) -> list[str]: + return [call.kwargs["name"] for call in alloc.model.add_constr.call_args_list] + + +def test_burst_off_keeps_only_the_interval_bound() -> None: + alloc, registry = port_registry() + MemoryPorts(burst=False).constrain(alloc, registry) + assert constraint_names(alloc) == ["port_interval_6_dram_rw_port_1"] + + +def test_interval_off_keeps_only_the_burst_bound() -> None: + alloc, registry = port_registry() + MemoryPorts(interval=False).constrain(alloc, registry) + assert constraint_names(alloc) == ["port_burst_6_dram_rw_port_1_0"] + + +def test_a_stream_whose_active_latency_rounds_to_zero_keeps_its_bits(monkeypatch: pytest.MonkeyPatch) -> None: + tr, choice, alloc = MagicMock(), MagicMock(), MagicMock() + tr.inputs[0].size_bits.return_value = 4096 + alloc.quantities = QuantityRegistry() + alloc.quantities.add("transfer_latency", "gated", index=(tr, choice)) + alloc.transfer_latency_for_path.return_value = 1 + loops = [SimpleNamespace(size=2, effect=LoopEffect.INVARIANT), SimpleNamespace(size=8, effect=LoopEffect.ABSENT)] + alloc.ssis.get.return_value.get_temporal_variables.return_value = loops + alloc.y_path_choice = {(tr, choice): SimpleNamespace(_raw="y")} + alloc.slot_of = {tr: 0} + monkeypatch.setattr(traffic, "_sides", lambda *_: []) + (stream,) = traffic.dma_streams(alloc) + # Active for 1 of 8 iterations: 1 cycle rounds to 0, while 4096 / 8 bits still move. + assert (stream.active_cycles, stream.bits_per_cycle * stream.latency_ub) == ("y", 512) + + +def zigzag_stall(real_cycles: list[int], window: int, periods: int) -> float: + """ZigZag's combined stall of double-buffered streams sharing one port (CostModelEvaluation, step 2).""" + activities = [ + PortActivity(r, window, window, periods, LayerOperand("I"), 0, DataDirection.WR_IN_BY_HIGH) for r in real_cycles + ] + union = CostModelEvaluation._CostModelEvaluation__calc_mem_updating_window_union(None, activities) # type: ignore[attr-defined] + positive = sum(a.stall_or_slack for a in activities if a.stall_or_slack > 0) + negative = sum(a.stall_or_slack for a in activities if a.stall_or_slack <= 0) + return positive + max(0, negative + sum(a.mem_updating_window for a in activities) - union) + + +@pytest.mark.parametrize("bits", [[4096], [4096, 2048], [1024, 1024, 3072, 512]]) +def test_interval_bound_is_zigzag_stall_free_window_of_double_buffered_streams(bits: list[int]) -> None: + rate, periods = 64, 10_000 + real = [b // rate for b in bits] + interval = sum(bits) / rate # the interval bound at equality: rate * interval = the streams' bits + assert zigzag_stall(real, round(interval), periods) == 0 + assert zigzag_stall(real, round(interval) - 1, periods) > 0 diff --git a/tests/unit/test_no_gurobipy_leakage.py b/tests/unit/test_no_gurobipy_leakage.py index e74d6872f..767bee915 100644 --- a/tests/unit/test_no_gurobipy_leakage.py +++ b/tests/unit/test_no_gurobipy_leakage.py @@ -1,9 +1,5 @@ -"""Verify no gurobipy imports leaked outside the solver backend. - -The only permitted gurobipy imports are: -1. stream/opt/solver/solver.py (the backend implementation) -2. transfer_and_tensor_allocation.py: `from gurobipy import GRB` (callback constants) -""" +"""No gurobipy import outside the solver backend (stream/opt/solver/solver.py), except +`from gurobipy import GRB` in transfer_and_tensor_allocation.py for callback codes and status names.""" import ast from pathlib import Path @@ -13,7 +9,6 @@ # Files that MUST NOT have any gurobipy imports CLEAN_FILES = [ - CO_DIR / "allocation.py", CO_DIR / "context.py", CO_DIR / "utils.py", REPO_ROOT / "stream" / "api.py", diff --git a/tests/unit/test_ports.py b/tests/unit/test_ports.py new file mode 100644 index 000000000..81374db9f --- /dev/null +++ b/tests/unit/test_ports.py @@ -0,0 +1,147 @@ +from typing import Any + +import pytest +from zigzag.utils import open_yaml + +from stream.cost_model.bandwidth import BandwidthModel +from stream.cost_model.core_cost import CoreCostEntry +from stream.hardware.architecture.accelerator import Accelerator +from stream.hardware.ports import OUTPUT, READ_BY_DATAPATH, WRITE, WRITE_BY_DATAPATH, input_role +from stream.parser.accelerator_factory import AcceleratorFactory +from stream.parser.accelerator_validator import AcceleratorValidator +from stream.stages.estimation.core_cost_backends import ZIGZAG_BACKEND, port_traffic +from stream.stages.parsing.accelerator_parser import parse_accelerator + +FUSEMAX = "stream/inputs/examples/hardware/fusemax.yaml" +AIE2 = "stream/inputs/aie/hardware/whole_array.yaml" +ALIASED = "tests/fixtures/hardware/tpu_v7_aliased.yaml" +MEASURED = { + "ceiling": 296.6, + "contiguous": 148.3, + "strided": {"read": {32: 19.0, 512: 93.5}, "write": {32: 17.1, 512: 93.5}}, +} + + +def fusemax_with_bandwidth(bandwidth: dict[Any, Any]) -> tuple[bool, dict[str, Any]]: + """Validate fusemax with a ``bandwidth:`` section; returns (valid, normalized data).""" + data = open_yaml(FUSEMAX) + data["bandwidth"] = bandwidth + validator = AcceleratorValidator(data, FUSEMAX) + return validator.validate(), validator.normalized_data + + +def fusemax_accelerator(bandwidth: dict[Any, Any]) -> Accelerator: + valid, data = fusemax_with_bandwidth(bandwidth) + assert valid + return AcceleratorFactory(data).create() + + +@pytest.fixture(scope="module") +def fusemax() -> Accelerator: + return parse_accelerator(FUSEMAX) + + +def test_cores_sharing_a_memory_share_its_ports(fusemax: Accelerator): + registry = fusemax.ports + sram = [port for port in registry if port.key.memory == "sram"] + assert {port.key for port in sram} == {(0, "sram", "r_port_1"), (0, "sram", "w_port_1")} + assert all(port.core_ids == (0, 1) for port in sram) + assert registry.ports_of(fusemax.get_core(0)) == registry.ports_of(fusemax.get_core(1)) + + +def test_each_operand_gets_its_own_dram_port(fusemax: Accelerator): + registry = fusemax.ports + dram = fusemax.get_core(2) + assert len(registry.ports_of(dram)) == 3 + roles = (input_role(1), input_role(2), OUTPUT) + names = {role: registry.port_for(dram, READ_BY_DATAPATH, role) for role in roles} + assert {role: port and port.key.port for role, port in names.items()} == { + "input1": "rw_port_1", + "input2": "rw_port_2", + "output": "rw_port_3", + } + + +def test_a_direction_the_operand_does_not_use_falls_back_to_a_port_serving_it(fusemax: Accelerator): + registry = fusemax.ports + port = registry.port_for(fusemax.get_core(2), WRITE_BY_DATAPATH, input_role(1)) + assert port is not None and port.key.port == "rw_port_3" + + +def test_declared_port_width_becomes_the_port_model(fusemax: Accelerator): + # fusemax_dram.yaml: rw ports of 3200 b/cc + port = next(p for p in fusemax.ports if p.key == (2, "dram", "rw_port_1")) + assert port.bits_per_cycle == 3200 + assert port.bandwidth.efficiency(4, "read") == 1.0 + + +def test_aie2_tiles_expose_their_dma_channels_as_ports(): + accelerator = parse_accelerator(AIE2) + compute, mem_tile = accelerator.get_core(2), accelerator.get_core(1) + rates = {(p.key.port, p.bits_per_cycle) for p in accelerator.ports.ports_of(compute)} + assert rates == {("mm2s", 2 * 64), ("s2mm", 2 * 64)} + assert accelerator.ports.port_for(mem_tile, WRITE, input_role(1)).bits_per_cycle == 6 * 64 + assert accelerator.ports.port_for(compute, READ_BY_DATAPATH, input_role(1)) is None + + +def test_a_core_with_measured_bandwidth_has_no_ports(): + accelerator = fusemax_accelerator({2: MEASURED}) + registry = accelerator.ports + assert registry.ports_of(accelerator.get_core(2)) == () + assert {port.key.memory for port in registry} == {"sram"} + assert set(accelerator.bandwidth) == {2} + + +def test_a_port_key_replaces_that_ports_declared_width(): + accelerator = fusemax_accelerator({"2.dram.rw_port_1": MEASURED, "1.sram.r_port_1": MEASURED}) + ports = {port.key: port for port in accelerator.ports} + measured = BandwidthModel.from_description(MEASURED) + assert ports[(2, "dram", "rw_port_1")].bandwidth == measured + assert ports[(0, "sram", "r_port_1")].bandwidth == measured + assert ports[(2, "dram", "rw_port_2")].bits_per_cycle == 3200 + assert accelerator.bandwidth == {} + + +@pytest.mark.parametrize("key", [9, "9.dram.rw_port_1", "2.dram"]) +def test_bandwidth_keys_naming_no_core_are_rejected(key: Any): + assert not fusemax_with_bandwidth({key: MEASURED})[0] + + +@pytest.mark.parametrize("key", ["2.sram.rw_port_1", "2.dram.rw_port_9", "0.rf_I.r_port_1"]) +def test_port_keys_naming_no_top_level_port_are_rejected_at_parse_time(key: Any): + with pytest.raises(ValueError, match="declares no port"): + fusemax_accelerator({key: MEASURED}) + + +def test_aie2_port_keys_are_rejected_at_parse_time(): + data = open_yaml(AIE2) + data["bandwidth"] = {"0.mem.port": MEASURED} + validator = AcceleratorValidator(data, AIE2) + assert validator.validate() + with pytest.raises(ValueError, match="models no memory ports"): + AcceleratorFactory(validator.normalized_data).create() + + +def test_a_port_key_on_a_measured_core_is_rejected(): + assert not fusemax_with_bandwidth({2: MEASURED, "2.dram.rw_port_1": MEASURED})[0] + + +class _BackendWithoutPortTraffic: + name = "bare" + priority = 0 + + +def test_a_cost_backend_without_port_traffic_reports_none(): + entry = CoreCostEntry(energy_total=0.0, latency_total=0.0, ideal_cycle=0.0, ideal_temporal_cycle=0.0) + assert port_traffic(_BackendWithoutPortTraffic(), entry) == () + assert port_traffic(ZIGZAG_BACKEND, entry) == () + + +def test_aliased_memories_share_one_physical_port(): + # tpu_v7_aliased.yaml: `8.vmem`, `0.operand_buffer` and `1.operand_buffer` are one scratchpad. + accelerator = parse_accelerator(ALIASED) + vmem = [port for port in accelerator.ports if 8 in port.core_ids] + assert {port.key for port in vmem} == {(8, "vmem", "r_port_1"), (8, "vmem", "w_port_1")} + assert all(port.core_ids == (0, 1, 8) for port in vmem) + assert all(port.bits_per_cycle == 262144 for port in vmem) + assert accelerator.ports.ports_of(accelerator.get_core(0)) == tuple(vmem) diff --git a/tests/unit/test_quantities.py b/tests/unit/test_quantities.py new file mode 100644 index 000000000..26d8db725 --- /dev/null +++ b/tests/unit/test_quantities.py @@ -0,0 +1,48 @@ +from collections.abc import Callable +from typing import Any + +import pytest + +from stream.inputs.testing.mapping.make_2_conv_mapping import make_2_conv_mapping +from stream.inputs.testing.workload.make_2_conv import TwoConvWorkloadConfig, make_2_conv_workload +from stream.opt.allocation.constraint_optimization import transfer_and_tensor_allocation as tta +from stream.opt.allocation.constraint_optimization.quantities import QuantityRegistry + +ACCELERATOR = "stream/inputs/examples/hardware/tpu_like_quad_core.yaml" + + +@pytest.fixture(scope="module") +def alloc( + tmp_path_factory: pytest.TempPathFactory, solved_allocator: Callable[..., Any], two_conv: TwoConvWorkloadConfig +) -> tta.TransferAndTensorAllocator: + out = str(tmp_path_factory.mktemp("two_conv")) + return solved_allocator(ACCELERATOR, make_2_conv_workload(two_conv), out, make_2_conv_mapping(two_conv)) + + +def test_registry_returns_scalar_and_indexed_quantities() -> None: + registry = QuantityRegistry() + registry.add("overlap", 3, upper_bound=10) + registry.add("slot_latency", 1, index=0) + registry.add("slot_latency", 2, index=1) + assert registry.get("overlap").upper_bound == 10 + assert registry.get("slot_latency", 1).expr == 2 + assert set(registry.indexed("slot_latency")) == {0, 1} + assert registry.names() == ["overlap", "slot_latency"] + + +def test_registry_rejects_duplicates_and_names_the_available_quantities() -> None: + registry = QuantityRegistry() + registry.add("overlap", 3) + with pytest.raises(ValueError, match="already registered"): + registry.add("overlap", 4) + with pytest.raises(KeyError, match="available: overlap"): + registry.get("buffering") + + +@pytest.mark.slow +def test_allocator_registers_what_it_builds(alloc: tta.TransferAndTensorAllocator) -> None: + expected = {"slot_latency", "overlap", "iteration", "transfer_latency", "primary"} + assert expected <= set(alloc.quantities.names()) + slots = alloc.quantities.indexed("slot_latency") + iteration = alloc.model.value(alloc.quantities.get("iteration").expr) + assert iteration == sum(alloc.model.value(q.expr) for q in slots.values()) diff --git a/tests/unit/test_resident_fill.py b/tests/unit/test_resident_fill.py new file mode 100644 index 000000000..d888ad4f6 --- /dev/null +++ b/tests/unit/test_resident_fill.py @@ -0,0 +1,88 @@ +"""The cycles a run waits for the off-chip tensors it holds in a single buffer, which fill before the first +iteration reads them and overlap no iteration.""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest + +from stream.opt.allocation.constraint_optimization.quantities import QuantityRegistry +from stream.opt.allocation.constraint_optimization.transfer_and_tensor_allocation import ( + TransferAndTensorAllocator, +) +from stream.opt.solver import SolverBackend, SolverParams, SolverVarType, create_solver + +TILES = 32 +MOVES = 3 +# Per tile: (cycles on its own path, cycles of the shared off-chip bandwidth), a strided key +# and a contiguous value as the attention trace moved them. +KEY, VALUE = (1225, 1225), (1024, 442) + + +class _Node: + """A tensor or transfer: a name, its outputs, and identity for a key.""" + + def __init__(self, name: str, outputs=()): + self.name, self.outputs = name, list(outputs) + + +def _fill(*, rotating: bool, single: bool = False) -> float: + """The solved fill for a key and a value held at reuse level 0, rotated or not, and if rotated + in one buffer or two. An outer loop of ``MOVES`` moves a rotated window on.""" + allocator = object.__new__(TransferAndTensorAllocator) + allocator.model = model = create_solver(SolverBackend.ORTOOLS_GSCIP, "fill") + model.set_param(SolverParams.VERBOSITY, 0) + allocator.quantities = QuantityRegistry() + allocator._name_counter = {} + allocator.shared_bandwidth = {0: SimpleNamespace(ceiling=1.0)} + cycles = {} + allocator.y_path_choice, allocator.z_stop, allocator.z_single = {}, {}, {} + allocator.tensors_to_optimize_reuse_for, allocator.ssis = [], {} + allocator.rotation_levels, allocator.tiles_needed_levels = {}, {} + for name, cost in (("key", KEY), ("value", VALUE)): + tensor = _Node(name) + transfer = _Node(f"Transfer({name})", [tensor]) + cycles[transfer] = cost + y = model.add_var(vtype=SolverVarType.BINARY, name=f"y_{name}") + z = model.add_var(vtype=SolverVarType.BINARY, name=f"z_{name}") + model.add_constr(y == 1) + model.add_constr(z == 1) + allocator.y_path_choice[(transfer, "path")] = y + allocator.z_stop[(tensor, 0)] = z + allocator.z_stop[(tensor, 1)] = outer = model.add_var(vtype=SolverVarType.BINARY, name=f"z1_{name}") + model.add_constr(outer == 0) + allocator.tensors_to_optimize_reuse_for.append(tensor) + allocator.ssis[tensor] = SimpleNamespace( + get_applicable_temporal_variables=lambda: [None, None], + get_applicable_temporal_sizes=lambda: [TILES, MOVES], + ) + allocator.rotation_levels[(tensor, 0)] = rotating + allocator.rotation_levels[(tensor, 1)] = False + allocator.tiles_needed_levels[(tensor, 0)] = 1 if single else TILES + allocator.tiles_needed_levels[(tensor, 1)] = MOVES * allocator.tiles_needed_levels[(tensor, 0)] + if single: + allocator.z_single[(tensor, 0)] = z + allocator._is_const_i = lambda transfer: True + allocator.transfer_latency_for_path = lambda transfer, path: cycles[transfer][0] + allocator._shared_cycles = lambda core, transfer, path, rate: cycles[transfer][1] + allocator._resident_fill() + model.set_objective(allocator.fill._raw) + model.optimize() + return allocator.fill.X + + +def test_a_whole_window_held_in_one_buffer_fills_before_the_run(): + """Each fill takes its own path's time and all share the off-chip bandwidth, so the run + waits for the bandwidth: both tensors' shares end to end, beyond either fill alone.""" + assert _fill(rotating=False) == pytest.approx(TILES * (KEY[1] + VALUE[1])) + + +def test_a_rotating_window_is_prefetched_and_costs_no_wait(): + assert _fill(rotating=True) == 0 + + +def test_a_rotating_window_held_in_one_buffer_waits_each_time_it_moves_on(): + """A head's key held once on its core: every next head waits for its fill, on its own path + and through the shared bandwidth, so the run waits for the bandwidth once a head.""" + assert _fill(rotating=True, single=True) == pytest.approx(MOVES * (KEY[1] + VALUE[1])) diff --git a/tests/unit/test_shared_memory.py b/tests/unit/test_shared_memory.py index b8d18d02e..4ba343928 100644 --- a/tests/unit/test_shared_memory.py +++ b/tests/unit/test_shared_memory.py @@ -82,7 +82,7 @@ def test_a_handover_between_cores_sharing_a_memory_stays_in_it(): assert handovers for tr in handovers: for choice in alloc.possible_transfer_allocations[tr]: - assert alloc._transfer_latency_for_path(tr, choice) == 0 + assert alloc.transfer_latency_for_path(tr, choice) == 0 assert not alloc.links_in_choice[(tr, choice)] # Each copy is its source's buffer, and these hold nothing beyond their sources. copies = {t.name for tr in handovers for t in tr.outputs} diff --git a/tests/unit/test_slot_latency_breakdown.py b/tests/unit/test_slot_latency_breakdown.py new file mode 100644 index 000000000..bc603891c --- /dev/null +++ b/tests/unit/test_slot_latency_breakdown.py @@ -0,0 +1,22 @@ +from pathlib import Path + +import pytest +import yaml + +from stream.api import evaluate_mapping +from stream.inputs.testing.mapping.make_2_conv_mapping import make_2_conv_mapping +from stream.inputs.testing.workload.make_2_conv import TwoConvWorkloadConfig, make_2_conv_workload + +ACCELERATOR = "stream/inputs/examples/hardware/tpu_like_quad_core.yaml" + + +@pytest.mark.slow +def test_slot_latency_breakdown_reports_solved_reuse_factor_and_contribution( + tmp_path: Path, two_conv: TwoConvWorkloadConfig +) -> None: + evaluate_mapping(ACCELERATOR, make_2_conv_workload(two_conv), str(tmp_path), make_2_conv_mapping(two_conv)) + (path,) = tmp_path.rglob("slot_latency_breakdown.yaml") + slots = yaml.safe_load(path.read_text())["slots"] + transfers = [tr for slot in slots for tr in slot["transfer_contributors"]] + assert transfers + assert all(tr["reuse_factor"] is not None and tr["contribution"] is not None for tr in transfers)