Land the allocator and AIE codegen changes IRON's stream-dse operators need

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 `<accelerator_dir>/cores/` then `<accelerator_dir>/`. |
 | `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 `<core id>.<memory>.<port>` (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:

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
 

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

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,

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,21 +221,26 @@ 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,
         element_type: FixedBitwidthType,
         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)
         if not isinstance(core_allocation, ArrayAttr):
             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

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