Stages¶

Stream's mapping flow is a pipeline of stages. Each stage does one job - parse an input, generate tilings, estimate cost, run the MILP allocation - and passes shared state to the next through a StageContext. This makes the flow easy to read, configure, and extend.

The framework lives in stream/stages/.

Execution model¶

• Stage - the base unit of work. Each stage runs the stages after it in its list_of_callables and yields what they yield; a LeafStage ends the list. A MainStage takes the ordered list and a context and runs the pipeline.
• StageContext (stream/stages/context.py) - the fields the stages hand each other. Inputs (hardware, workload, mapping path, backend, output path) go in; results (group_cycles, group_latencies, allocation, workload, accelerator, ...) come out. You read results with ctx.get("...").

The public API functions in stream/api.py assemble the right stage list for you - you normally don't build a MainStage by hand.

The CO pipeline¶

evaluate_mapping, select_mapping and generate_code run one pipeline:

1. AcceleratorParserStage - the accelerator model (cores, memories, interconnect), already parsed from the hardware YAML so the plugins below can be chosen for it. The workload comes from the frontend that loads it (see Workload).
2. Mapping - either:
3. FixedMappingGenerationStage (a mapping was given) - split the workload at the mapping's fused groups and build each group's mapping; or
4. the stages of the mapping generator that claims the accelerator (no mapping given) - for any accelerator, ExpandNormalizationStage + GenericMappingGenerationStage propose one.
5. FusionGroupIterationStage - run the stages below once per fused group, on that group's workload and mapping.
6. The stage of the code generation backend that claims the accelerator, for generate_code only - for AIE2 arrays, AIECodeGenerationStage.
7. PlacementGenerationStage - place a group whose kernels are known but whose cores are not, and offer one variant per other compiled block.
8. KernelStateStage - the state a kernel carries between iterations.
9. TileSearchStage - with tile_search, price the tile candidates around the mapping's seed, each in its own [p<placement>_]tile_<index>/ folder, and keep the fastest, whose files it copies into the group's folder.
10. TilingGenerationStage - generate the intra-/inter-core tilings for each node.
11. CoreCostEstimationStage - estimate per-(node, core) cost through the core-cost backend that claims each core.
12. SteadyStateLoweringStage - lower the group to its steady state (stream.allocation.lowering): make the transfers explicit with the placements and routes each may take, and fix the iteration spaces and timeslots, as an AllocationProblem.
13. AllocationStage - build the MILP (AllocationModel) for that problem from its constraint families and solve it: decide tensor placement and transfer paths, producing the Allocation the context carries as allocation.
14. MemoryAccessesEstimationStage - estimate memory traffic for the chosen allocation.

The mapping generators (stream.mapping_generators) and code generation backends (stream.codegen_backends) are entry-point groups: an object with a name, a priority, a claims(accelerator) predicate and stages() or stage() extends the pipeline for a new kind of hardware, the highest priority among those that claim it winning.

Windowed operators¶

A convolution or a pool reads each spatial axis of its input through a sliding window, s*o + d*f - p of its output dim o and kernel dim f; the same per-axis strides, dilations and padding serve every windowed parser and frontend. The pipeline handles such a read like any other affine access:

• Couplings. A reader that indexes a producer's axis as s*o + windows + c merges the producer's dim into its own o, with a remainder dim of size s when s > 1, so fused convolutions and pools share their row and column axes; the window and the constant stay in the access maps. Only where the producer's extent is s times the reader's: an unpadded reader leaves its producer's axis its own, as an unfused one. A mapping naming the producer's dim s*z + r tiles or splits z, so its tile is a multiple of s, and entries naming one fused axis cut it alike.
• Halos. A tile's footprint is its interior window, not clipped at the tensor's origin, and a transfer's copy holds the window of the node it reaches. Each loop that slides that window carries its halo, the window less its step, on its IterationVariable: the rows a fused loop shares between iterations, the columns cores share.
• Line buffer. A sliding loop keeps its halo resident: the copy holds the window (times its buffering) and the transfer moves only the step. A target core is fed by every source core whose tile its window overlaps, and the halo it takes from a core it shares no memory with counts on the links, ports, DMA channels and off-chip traffic like any transfer. ZigZag costs the interior tile, which needs no border padding.
• Boundary iterations. A producer's first tile also covers what its reader's first window reaches ahead (conv1 computes 9, 8, 8, 7 rows under a 3x3 conv tiled to 8 rows), and its last is shorter (Workload.get_sliding_work). The fill adds, at the interior rate per element, the most a node's work runs ahead of its steady pace over the nested fused loops: its longer first and shorter last tiles cancel on its own cores, so what counts is who waits on them, a transfer's reader and a node's readers behind a moving transfer or on other cores.

Constraint families¶

The allocation model is built from constraint families, each a group of constraints and the quantities they define. Stream's own families are built in; another package registers its own in the stream.constraint_families entry-point group. SolveOptions(families=...) lists the families of a solve, by name or as {name: options}, the options being the family's constructor arguments. Without it a solve builds the default set, stream.api.default_families(hardware): Stream's own families and those of each core namespace the accelerator has, such as aie2's.

A family declares the quantities it requires and those it provides, and the builder runs each one after every family that provides what it requires: Stream's own families keep the order of the table, and any other family runs as soon as its requirements allow, so a namespace's limit lands right after the quantity it bounds. A family that defines quantities others read before its own constraints can be built, as memory_ports declares the slot latencies its ports can force before the overlap is built, does that in a separate declare step. A selection that leaves out what a family requires is rejected, so switching a constraint group off means leaving out its family and those that need it, which default_families(hardware, without=[...]) does:

from stream.api import SolveOptions, default_families

hardware = "stream/inputs/aie/hardware/whole_array_strix.yaml"
no_dma = SolveOptions(families=default_families(hardware, without=["dma_channels"]))  # also drops aie2_dma_channels
span = SolveOptions(families=[*default_families(hardware, without=["overlap"]), {"overlap": {"model": "span"}}])
ports_report_only = SolveOptions(families=default_families(hardware, options={"memory_ports": {"interval": False, "burst": False}}))

A family's build(ctx) (and declare) receives a FormulationContext: ctx.space, the read-only problem and the choices derived from it (each tensor's placements, each transfer's routes and the links they use, each tensor's reuse stops and the tiles they hold); ctx.vars, the core decision variables (x places a tensor, y routes a transfer, z_stop stops a tensor's reuse, z_single holds its window in one buffer, slot_latency); ctx.model, the SolverModel; ctx.quantities, the QuantityRegistry of the quantities the families provide; and the modelling helpers families share, such as binary_product and tensor_uses_core_var. A family that creates a constraint with ctx.add_constr(expr, name=..., resource=..., kind=..., subject=..., rule=..., bound=...) states what it stands for: the core or link it binds, the hardware limit (ResourceKind) it is part of and that limit's bound, the tensor whose demand it carries, or the StructuralRule it enforces. When a model has no solution, the diagnosis maps the solver's IIS back to cores, links and causes through these tags alone. A family can also contribute to the objective: objective(ctx) runs once every family has built and returns ObjectiveLevels; the levels of one name are summed and the solve minimizes them lexicographically, highest priority first. Stream's levels are latency (priority 4: the run's latency from overlap, the DMA peaks from dma_channels and the weighted off-chip traffic from offchip_traffic), offchip_traffic (3), buffering (2, from object_fifo_depth) and route_hops (1, from path_choice). A family's report(ctx) adds sections to the solved allocation's performance report; one that fails is logged and leaves its section None. Before any family builds, the model checks that every memory fits the tensors pinned to it under some reuse choice, whichever families a solve selects, and fails one that cannot with an InfeasibleAllocationError.

A core namespace (a HardwareNamespace) contributes its families by naming them in its families, next to the facts the model reads of it: which cores share memory, what the toolchain reserves, and what a dispatch of several designs costs. Another package plugs into the allocation model through two entry-point groups:

• stream.constraint_families - a family factory under the family's name, whose keyword arguments are its options.
• stream.namespaces - a HardwareNamespace subclass under the core namespace it describes, built with from_config whenever the accelerator has a core of that namespace.

Stream's own families and its aie2 namespace are built in, so an install whose entry points are stale still has them.

Migrating from Stream 1.x¶

Stream 2.0 replaces ConstraintSelection with the family list; each of its toggles has an exact equivalent, which builds the same model:

The rest of the 2.0 changes an extension meets: the stream.constraints entry-point group is stream.namespaces (an entry point left in the old group is ignored, with a warning), and NamespaceConstraints is HardwareNamespace (AIE2Constraints, AIE2Namespace), whose add_*_constraints hooks are families now (a namespace that still defines one is rejected with the family that replaces it); the context carries the solved allocation (an Allocation, read by AllocationIR.from_internal) instead of the scheduler; AllocationIR lists its families instead of a constraint_selection; and a stage declares its contract (below) instead of REQUIRED_FIELDS.

Stage contracts¶

A stage declares the context fields it touches, as tuples of field names on the class:

• reads - the fields it needs when it starts; each must be in the context it is given or written by a stage before it, and not None.
• optional_reads - the fields it reads when they are there, such as the options a plugin's stages take through SolveOptions(stage_options=...).
• writes - the fields it sets before the stages after it run.
• result_reads, result_writes - the fields it reads and sets on the context the stages after it yield, as FusionGroupIterationStage does with each group's allocation.

MainStage checks a pipeline's contracts before it runs anything: following the list, and through the stages a stage runs after it (the inner pipelines of FusionGroupIterationStage and TileSearchStage are the rest of the list), each read must be in the initial context or in the writes of a stage before it, and each result read in what the stages after it write. A pipeline that breaks the rule fails with a StageContractError naming the stage and the field. While a stage runs, the context lets it read only the fields its contract names and write only its writes, raising a StageContractError otherwise; outside a stage, as when a caller reads the result, every field is open. ctx.data is the raw store, for a stage that snapshots and restores the whole context, as TileSearchStage does between candidates.

A test checks this table against the stages' declarations. Constraint families follow the same rule one level down: a family declares the quantities it requires and provides, and the model builder orders and checks them (see Constraint families).

Writing a custom stage¶

Subclass Stage, declare its contract, and run the stages after it:

from stream.stages.stage import Stage

class LayerCountStage(Stage):
    reads = ("workload",)
    writes = ("layer_count",)

    def run(self):
        self.ctx.set(layer_count=len(self.ctx.get("workload").get_computation_nodes()))
        yield from self.list_of_callables[0](self.list_of_callables[1:], self.ctx).run()

Insert it at the right position in the list passed to MainStage. A stage that reduces (keeps only the best of several results) yields once after its loop rather than inside it. Out-of-tree stages, such as those a mapping generator's stages() or a code generation backend's stage() returns, declare their contracts the same way. A stage that declares none, such as one written against Stream 1.x, and a stage callable that is not a Stage class, such as a factory or a functools.partial, run unchecked: the pipeline logs once that it cannot check them, and checks the fields of the stages after them only as they run.