stream-dse 2.3.0__tar.gz → 2.4.0__tar.gz

{stream_dse-2.3.0 → stream_dse-2.4.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: stream-dse
-Version: 2.3.0
+Version: 2.4.0
 Summary: Stream - Multi-core accelerator design space exploration with layer-fused scheduling
 Author-email: Arne Symons <arne.symons@kuleuven.be>, Linyan Mei <linyan.mei@kuleuven.be>
 License: MIT License

{stream_dse-2.3.0 → stream_dse-2.4.0}/pyproject.toml

@@ -21,7 +21,7 @@ namespaces = true
 
 [project]
 name = "stream-dse"
-version = "2.3.0"
+version = "2.4.0"
 description = "Stream - Multi-core accelerator design space exploration with layer-fused scheduling"
 readme = "README.md"
 authors = [{ name = "Arne Symons", email = "arne.symons@kuleuven.be" }, { name = "Linyan Mei", email = "linyan.mei@kuleuven.be" }]
@@ -148,7 +148,7 @@ generic = "stream.stages.generation.mapping_generators:GENERIC_MAPPING"
 Homepage = "https://github.com/KULeuven-MICAS/stream"
 
 [tool.bumpver]
-current_version = "2.3.0"
+current_version = "2.4.0"
 version_pattern = "MAJOR.MINOR.PATCH"
 commit_message = "bump version {old_version} -> {new_version}"
 commit = true

stream_dse-2.4.0/stream/__init__.py

@@ -0,0 +1 @@
+__version__ = "2.4.0"

{stream_dse-2.3.0 → stream_dse-2.4.0}/stream/aie_setup.py

@@ -31,11 +31,11 @@ from pathlib import Path
 _XDSL_AIE = "git+https://github.com/xdslproject/xdsl-aie.git@f2540442628aa782847ababd1dc76c0e323ba459"
 _SNAX_MLIR = "git+https://github.com/kuleuven-micas/snax-mlir.git@1c01c5d100df128c9fa01d3336ebea98e19b20cf"
 
-_MLIR_AIE_PIN = "mlir_aie==1.4.4"
-_LLVM_AIE_PIN = "llvm-aie==22.0.0.2026092101+0006955e"
+_MLIR_AIE_PIN = "mlir_aie==1.4.6.dev15+ge5a3fa1"
+_LLVM_AIE_PIN = "llvm-aie==22.0.0.2026102001+4f5df13d"
 _AIE_WHEEL_INDICES = [
     "--extra-index-url",
-    "https://github.com/Xilinx/mlir-aie/releases/expanded_assets/v1.4.4",
+    "https://github.com/Xilinx/mlir-aie/releases/expanded_assets/latest-wheels-4",
     "--extra-index-url",
     "https://github.com/Xilinx/llvm-aie/releases/expanded_assets/nightly",
     "--extra-index-url",

{stream_dse-2.3.0 → stream_dse-2.4.0}/stream/allocation/allocation.py

@@ -13,7 +13,6 @@ from stream.opt.allocation.constraint_optimization.report import guarded
 from stream.workload.node import InEdge, OutEdge, Tensor, declared_of, edge_tensor
 from stream.workload.steady_state.iteration_space import (
     IterationVariableType,
-    LoopEffect,
     Reuse,
     SteadyStateIterationSpace,
 )
@@ -164,6 +163,22 @@ class Allocation:
                 add(str(layer_dim), size, "core_temporal")
         return loops
 
+    def _group_loops(self) -> list[dict]:
+        """The group's shared loops, outermost first; a dim split over cores is one loop at the largest split."""
+        loops: list[dict] = []
+        splits: dict[str, int] = {}
+        for cn in self.problem.workload.get_computation_nodes():
+            parts: dict[str, int] = {}
+            for iv in reversed(self.ssis[cn].variables) if cn in self.ssis else ():
+                loop = {"dim": str(iv.dimension), "size": iv.size, "type": iv.type.name.lower()}
+                if iv.type in (IterationVariableType.SPATIAL, IterationVariableType.SPATIOTEMPORAL):
+                    parts[loop["dim"]] = parts.get(loop["dim"], 1) * iv.size
+                elif iv.size > 1 and loop not in loops:
+                    loops.append(loop)
+            for dim, n in parts.items():
+                splits[dim] = max(splits.get(dim, 1), n)
+        return loops + [{"dim": dim, "size": n, "type": "spatial"} for dim, n in splits.items() if n > 1]
+
     def _steady_state_ir(self) -> dict:
         """Serialise the tiled/steady-state inspection view: operators, loop nest and transfer graph, and with the
         decision space the per-core detail of every tiled node, axis and tensor."""
@@ -175,16 +190,7 @@ class Allocation:
             }
             for cn in self.problem.source_workload.get_computation_nodes()
         ]
-        loops: list[dict] = []
-        seen: set = set()
-        for ssis in self.ssis.values():
-            for iv in ssis.variables:
-                if iv.effect is LoopEffect.ABSENT:
-                    continue
-                key = (str(iv.dimension), int(iv.size))
-                if int(iv.size) > 1 and key not in seen:
-                    seen.add(key)
-                    loops.append({"dim": str(iv.dimension), "size": int(iv.size), "type": iv.type.name.lower()})
+        loops = self._group_loops()
         expanded = False
         for cn in self.problem.source_workload.get_computation_nodes():
             core_loops = self._core_loops(cn)

{stream_dse-2.3.0 → stream_dse-2.4.0}/stream/allocation/lowering.py

@@ -135,6 +135,8 @@ class _Lowering:
         with span("iteration_spaces"):
             self.ssis = self.generate_ssis()
             self.iterations = self.calculate_iterations()
+        with span("routes"):
+            self.route_transfers()
         with span("timeslots"):
             timeslots = self.ssw.get_timeslots(self.mapping)
         return AllocationProblem(
@@ -325,7 +327,7 @@ class _Lowering:
         for node in self.ssw.get_computation_nodes():
             unique_dims_tiling = (self.ssw.get_unique_dims_inter_core_tiling(node, self.mapping),)
             self.mapping.update_inter_core_tiling(node, unique_dims_tiling)
-        # Add transfer node mappings
+        # Add transfer node mappings; their routes follow once the iteration spaces give each copy's halos.
         for node in self.ssw.get_transfer_nodes():
             assert len(node.inputs) == 1, "Transfer node must have exactly one input tensor."
             src = cast(HasOutputs, list(self.ssw.predecessors(node))[0])
@@ -333,6 +335,19 @@ class _Lowering:
             self.update_mapping_for_transfer(node, src, dsts)
         return self.mapping.with_updated_workload(self.ssw, self.workload)  # updates FusedGroups
 
+    def route_transfers(self) -> None:
+        """Give each transfer the routes of each of its placements, moving what each copy reads per firing."""
+        for node in self.ssw.get_transfer_nodes():
+            src = cast(HasOutputs, next(iter(self.ssw.predecessors(node))))
+            node_mapping = self.mapping.get(node)
+            dst_allocs, tilings = node_mapping.memory_allocation, node_mapping.inter_core_tiling
+            self.mapping.set_for_node(
+                node,
+                resource_allocation=self.determine_possible_transfer_plans(node, src, dst_allocs, tilings),
+                inter_core_tiling=tilings,
+                memory_allocation=dst_allocs,
+            )
+
     def _fresh(self, name: str) -> str:
         """``name``, suffixed until no tensor of the workload has it, so a copy cannot be taken for another tensor."""
         taken = {t.name for t in self.workload.tensors}
@@ -414,11 +429,13 @@ class _Lowering:
 
     def add_halos(self, tensor: Tensor, transfer: TransferNode, ssis: SteadyStateIterationSpace) -> None:
         """Give each loop that slides the window of a transfer's copy the halo its consecutive tiles share."""
-        sliding = (IterationVariableType.TEMPORAL, IterationVariableType.SPATIAL)
-        loops = [v for v in ssis if v.relevant and v.type in sliding]
-        windows = self.ssw.get_windows(tensor, transfer, self.mapping, [v.dimension for v in loops])
-        for loop in loops:
-            loop.halo = windows[loop.dimension][1] if loop.dimension in windows else 0
+        for kind in (IterationVariableType.TEMPORAL, IterationVariableType.SPATIAL):
+            loops = [v for v in ssis if v.relevant and v.type is kind]
+            dims = [v.dimension for v in loops]
+            temporal = kind is IterationVariableType.TEMPORAL
+            windows = self.ssw.get_windows(tensor, transfer, self.mapping, dims, temporal=temporal)
+            for loop in loops:
+                loop.halo = windows[loop.dimension][1] if loop.dimension in windows else 0
 
     def update_mapping_for_transfer(self, node: TransferNode, src: HasOutputs, dsts: tuple[HasInputs, ...]) -> None:
         possible_dst_allocs = self.determine_possible_memory_allocations(node, src, dsts)
@@ -429,15 +446,6 @@ class _Lowering:
             inter_core_tiling=possible_inter_core_tiling,
             memory_allocation=possible_dst_allocs,
         )
-        possible_allocations = self.determine_possible_transfer_plans(
-            node, src, possible_dst_allocs, possible_inter_core_tiling
-        )
-        self.mapping.set_for_node(
-            node,
-            resource_allocation=possible_allocations,
-            inter_core_tiling=possible_inter_core_tiling,
-            memory_allocation=possible_dst_allocs,
-        )
 
     def determine_possible_memory_allocations(
         self, node: TransferNode, src: HasOutputs, dsts: tuple[HasInputs, ...]
@@ -561,31 +569,48 @@ class _Lowering:
         possible_tilings: tuple[InterCoreTiling, ...],
     ) -> tuple[MulticastPathPlan, ...]:
         """The route of each placement of ``node``, moving what each target reads from the sources that hold it, of
-        each placement the allocation may choose for the input it moves (:func:`placement_options`)."""
-        all_possible_resource_plans = []
+        each placement the allocation may choose for the input it moves (:func:`placement_options`), with and without
+        a local copy from a larger copy in a target's memory."""
+        all_possible_resource_plans: dict[MulticastPathPlan, None] = {}
         bitwidth = node.inputs[0].operand_type.bitwidth
         for placement, src_allocs in self._held_options(src):
             mapping = self.mapping if placement is None else self.mapping.with_edge_placements({src.name: placement})
             for dst_allocs, tiling in zip(possible_dst_allocs, possible_tilings, strict=True):
-                demand = self.ssw.get_transfer_demand(
-                    node,
-                    mapping,
-                    tiling,
-                    len(src_allocs),
-                    len(dst_allocs),
-                    self.hardware.pairs_by_spatial_index((*src_allocs, *dst_allocs)),
-                    dst_allocs,
-                    self._reads_in,
-                    partial(self._iteration_cost, src_allocs, dst_allocs, bitwidth),
-                )
-                possible_resource_plans = self.accelerator.communication_manager.get_possible_transfer_plan(
-                    src_allocs, dst_allocs, demand, bitwidth, self.hardware.shares_memory, self.route_alternatives
-                )
-                if placement is not None:
-                    possible_resource_plans = [replace(plan, placement=placement) for plan in possible_resource_plans]
-                all_possible_resource_plans.extend(possible_resource_plans)
+                by_spatial_index = self.hardware.pairs_by_spatial_index((*src_allocs, *dst_allocs))
+                demands = {
+                    self.ssw.get_transfer_demand(
+                        node,
+                        mapping,
+                        tiling,
+                        len(src_allocs),
+                        len(dst_allocs),
+                        by_spatial_index,
+                        dst_allocs,
+                        self._reads_in,
+                        partial(self._iteration_cost, src_allocs, dst_allocs, bitwidth),
+                        ssis=self.ssis.get,
+                        fusion_splits=self.fusion_splits,
+                        allow_local_copy=allow_local_copy,
+                        shares_memory=partial(self._shares_memory, src_allocs, dst_allocs),
+                    ): None
+                    # Code generation pairing by spatial index moves every copy by its own transfer.
+                    for allow_local_copy in ((False,) if by_spatial_index else (False, True))
+                }
+                for demand in demands:
+                    possible_resource_plans = self.accelerator.communication_manager.get_possible_transfer_plan(
+                        src_allocs, dst_allocs, demand, bitwidth, self.hardware.shares_memory, self.route_alternatives
+                    )
+                    if placement is not None:
+                        possible_resource_plans = [
+                            replace(plan, placement=placement) for plan in possible_resource_plans
+                        ]
+                    all_possible_resource_plans.update(dict.fromkeys(possible_resource_plans))
         return tuple(all_possible_resource_plans)
 
+    def _shares_memory(self, sources: Sequence[Memory], targets: Sequence[Memory], source: int, target: int) -> bool:
+        """Whether the source at position ``source`` and the target at position ``target`` share a memory."""
+        return self.hardware.shares_memory(sources[source], targets[target])
+
     def calculate_iterations(self) -> int:
         """Calculate the amount of steady state iterations based on all nodes' SSIS."""
         iterations_per_node = {node: prod(ssis.get_temporal_sizes()) for node, ssis in self.ssis.items()}
@@ -659,7 +684,7 @@ class _Lowering:
         times the links they cross from the nearest source holding them."""
         held: dict[tuple[Memory, str], list[Box]] = defaultdict(list)
         carried = 0.0
-        for item in items:
+        for item in (item for item in items if item.sources):
             target = targets[item.target]
             source = min((sources[s] for s in item.sources), key=lambda memory: self.accelerator.hops(memory, target))
             for memory, direction in ((source, READ), (target, WRITE)):

{stream_dse-2.3.0 → stream_dse-2.4.0}/stream/allocation/solution.py

@@ -34,8 +34,8 @@ class AllocationSolution:
     depths: Mapping[Tensor, int]
     single_buffered: frozenset[Tensor]
     latency: Latency
-    slot_latencies: Mapping[int, float]
-    reuse_factors: Mapping[TransferNode, float]
+    slot_latencies: Mapping[int, int]
+    reuse_factors: Mapping[TransferNode, int]
     route_cycles: Mapping[TransferNode, int]
     primary_cost: float
     solve_stats: SolveStats

{stream_dse-2.3.0 → stream_dse-2.4.0}/stream/allocation/steady_state.py

@@ -7,12 +7,16 @@ from math import prod
 from typing import TYPE_CHECKING, Any
 
 from stream.opt.allocation.constraint_optimization.families.traffic import link_name
-from stream.opt.allocation.constraint_optimization.report import guarded
+from stream.opt.allocation.constraint_optimization.report import guarded, reported_reuse
+from stream.opt.allocation.constraint_optimization.utils import absent_trips
+from stream.workload.demand import Box, DemandItem, covered, firing_count
 from stream.workload.node import ComputationNode, Tensor, TransferNode
 
 if TYPE_CHECKING:
     from stream.allocation.allocation import Allocation
+    from stream.cost_model.communication_manager import MulticastPathPlan
     from stream.datatypes import LayerDim
+    from stream.hardware.architecture.memory import Memory
     from stream.opt.allocation.constraint_optimization.space import DecisionSpace
     from stream.workload.node import HasIterationSpace
     from stream.workload.steady_state.iteration_space import SteadyStateIterationSpace
@@ -35,7 +39,7 @@ def steady_state_detail(allocation: Allocation, space: DecisionSpace) -> dict[st
     return {
         "iterations": allocation.problem.iterations,
         "axes": _axes(allocation, sliding),
-        "tensors": _tensors(allocation),
+        "tensors": _tensors(allocation, space),
         "nodes": {name: found for name, found in nodes.items() if found is not None},
     }
 
@@ -53,13 +57,16 @@ def _axes(allocation: Allocation, sliding: set[LayerDim]) -> list[dict[str, Any]
     ]
 
 
-def _tensors(allocation: Allocation) -> list[dict[str, Any]]:
-    solution = allocation.solution
+def _tensors(allocation: Allocation, space: DecisionSpace) -> list[dict[str, Any]]:
+    """Each tensor as the solution placed and buffered it, a transfer's copy at the window its reader holds."""
+    solution, workload, mapping = allocation.solution, allocation.problem.workload, allocation.problem.mapping
+    writers = {copy: tn for tn in workload.get_transfer_nodes() for copy in tn.outputs}
+    held = {t: workload.get_tensor_single_core(t, tn, mapping) for t, tn in writers.items() if t in allocation.ssis}
     return [
         {
             "name": t.name,
-            "shape": _shape(t),
-            "bits": int(t.size_bits()),
+            "shape": _shape(held.get(t, t)),
+            "bits": space.held_bits(t, held.get(t, t)),
             "loops": _loops(ssis),
             "placement": [memory.name for memory in solution.tensor_placements.get(t, ())],
             "reuse_level": solution.reuse_levels.get(t),
@@ -90,6 +97,11 @@ def _shape(tensor: Tensor) -> list[int]:
     return [int(s) for s in tensor.shape]
 
 
+def _per_firing(points: int, spanned: int) -> int | float:
+    """Points over the firings a demand spans, an integer where they divide."""
+    return points // spanned if points % spanned == 0 else points / spanned
+
+
 class _NodeDetail:
     """What one tiled node does on its cores, read off the solution and the decision space it was solved in."""
 
@@ -104,11 +116,12 @@ class _NodeDetail:
         }
 
     def compute(self, cn: ComputationNode) -> dict[str, Any]:
-        """Its cores, the output tile at each core position and whether a reader waits on it."""
+        """Its cores, how often it runs, the output tile at each core position and whether a reader waits on it."""
         tile = self.workload.get_tensor_single_core
         return {
             "dims": [str(dim) for dim in self.workload.get_dims(cn)],
             "loops": _loops(self.allocation.ssis.get(cn)),
+            "calls": self.allocation.problem.iterations // absent_trips(self.space.ssis[cn]),
             "cores": [core.id for core in self.space.cores(cn)],
             "tiles": [_shape(tile(cn.outputs[0], cn, self.mapping, p)) for p in range(self._positions(cn))],
             "waited_on": self.space.is_waited_on(cn),
@@ -116,9 +129,9 @@ class _NodeDetail:
         }
 
     def transfer(self, tn: TransferNode) -> dict[str, Any]:
-        """Its route, the window of each copy a target holds and what one firing moves to it, the halos its loops
-        keep, what each source hands each target, and the bits it moves. A target holds the copies of the nodes it
-        runs, each at its position among that copy's targets modulo the positions, as sources pair with targets."""
+        """Its loops, route, copies and the copy each target holds (the copy of a node the target runs), the window
+        and moved tile of that copy, each copy's halos, the demand each source serves at each target, the bits it moves
+        and the iterations one firing serves."""
         space, solution = self.space, self.allocation.solution
         route = solution.transfer_routes[tn]
         # Pair each copy with the targets it may land in; a target no copy may land in gets the first copy.
@@ -133,26 +146,93 @@ class _NodeDetail:
             for copy, p in at
         ]
         return {
+            "loops": _loops(self.allocation.ssis.get(tn)),
             "sources": [m.name for m in route.sources],
             "targets": [m.name for _, m in held],
+            "copy_names": [copy.name for copy in tn.outputs],
+            "target_copies": [tn.outputs.index(copy) for copy, _ in held],
             "links": [link_name(link) for link in route.links_used],
             "windows": [_shape(self.workload.get_tensor_single_core(copy, tn, self.mapping, p)) for copy, p in at],
             "moved": [_shape(tile) for tile in moved],
             "halos": [
-                {"dim": str(loop.dimension), "type": loop.type.name.lower(), "halo": loop.halo}
-                for loop in space.ssis.get(tn.outputs[0]) or ()
+                {"dim": str(loop.dimension), "type": loop.type.name.lower(), "halo": loop.halo, "copy": i}
+                for i, copy in enumerate(tn.outputs)
+                for loop in reversed(space.ssis[copy].variables)
                 if loop.halo
             ],
-            "handed": [
-                {"source": route.sources[i].name, "target": route.targets[j].name, "elements": elements}
-                for (i, j), elements in space.overlaps(tn).items()
-            ],
-            "moved_bits": 0 if space.choice_shares_memory(tn, route) else space.moved_bits(tn, route),
+            "demand": self._demand_by_kind(tn, route),
+            **self._copy_bits(tn),
+            "link_bits": self._link_bits(tn, route),
             "route_cycles": solution.route_cycles[tn],
-            "reuse_factor": solution.reuse_factors.get(tn),
             **self._sliding(tn),
         }
 
+    def _copy_bits(self, tn: TransferNode) -> dict[str, Any]:
+        """Per copy, the bits one of its firings moves and its reuse factor, a box several copies read charged to the
+        one firing most often; the transfer's own are its most often firing copy's reuse and bits per such firing."""
+        space, solution = self.space, self.allocation.solution
+        route = solution.transfer_routes[tn]
+        reuse = [
+            reported_reuse(space.reuse_levels[(copy, solution.reuse_levels[copy])], space.ssis[tn])
+            if copy in solution.reuse_levels
+            else reported_reuse(solution.reuse_factors[tn], space.ssis[tn])
+            for copy in tn.outputs
+        ]
+        moved = [0.0] * len(tn.outputs)
+        if not space.choice_shares_memory(tn, route):
+            for copies in space.copy_sets(tn, route):
+                first = min(copies, key=lambda c: reuse[c])
+                moved[first] += space.sent_bits(tn, route, copies)
+        top = min((r for r, bits in zip(reuse, moved, strict=True) if bits), default=reuse[0])
+        return {
+            "moved_bits": sum(bits * top / r for bits, r in zip(moved, reuse, strict=True)),
+            "reuse_factor": top,
+            "copy_moved_bits": moved,
+            "copy_reuse_factors": reuse,
+        }
+
+    def _demand_by_kind(self, tn: TransferNode, route: MulticastPathPlan) -> list[dict[str, Any]]:
+        """Per source, target and kind, the elements of the target's demand the source serves per firing: over links
+        ('moved'), from a source in or sharing the target's memory ('in_place'), from a larger copy in that memory
+        ('local_copy'), or as a source's partial sum of a split reduction ('partial_sum')."""
+        served: dict[tuple[Memory, Memory, str], list[Box]] = {}
+        kind = "in_place" if self.space.choice_shares_memory(tn, route) else "moved"
+        sent = [item for item in route.delivered_items if not item.reduce]
+        delivered = {(item.target, item.box, item.copies) for item in sent}
+        for item in sent:
+            served.setdefault((route.sources[item.sources[0]], route.targets[item.target], kind), []).append(item.box)
+        for item in route.demand:
+            for key in self._served_without_links(route, item, delivered):
+                served.setdefault(key, []).append(item.box)
+        spanned = firing_count(route.demand)
+        return [
+            {"source": one.name, "target": other.name, "kind": kind, "elements": _per_firing(covered(boxes), spanned)}
+            for (one, other, kind), boxes in served.items()
+        ]
+
+    def _served_without_links(
+        self, route: MulticastPathPlan, item: DemandItem, delivered: set[tuple[int, Box, tuple[int, ...]]]
+    ) -> list[tuple[Memory, Memory, str]]:
+        """The source that serves a demand item no link carries, and the kind."""
+        target = route.targets[item.target]
+        if item.reduce:
+            return [(route.sources[s], target, "partial_sum") for s in item.sources]
+        if not item.sources:
+            return [(target, target, "local_copy")]
+        if (item.target, item.box, item.copies) in delivered:
+            return []
+        held = [s for s in item.sources if route.sources[s] is target] or [
+            s for s in item.sources if self.space.hardware.shares_memory(route.sources[s], target)
+        ]
+        return [(route.sources[(held or list(item.sources))[0]], target, "in_place")]
+
+    def _link_bits(self, tn: TransferNode, route: MulticastPathPlan) -> list[float]:
+        """Per link of the route, the bits one firing carries over it, as its route cycles count them."""
+        if self.space.choice_shares_memory(tn, route):
+            return [0.0] * len(route.links_used)
+        moved, shares = self.space.moved_bits(tn, route), dict(route.link_shares)
+        return [moved * shares.get(link, 0.0) if shares else moved for link in route.links_used]
+
     def _positions(self, node: HasIterationSpace) -> int:
         tiling = self.mapping.get(node).inter_core_tiling
         return prod(f for _, f in tiling[0]) if tiling else 1

{stream_dse-2.3.0 → stream_dse-2.4.0}/stream/api.py

@@ -242,10 +242,10 @@ def _solve(
         raise
     finish_instrumentation(observers)
     groups = sorted(ctx.get("group_cycles"))
-    columns = [len(ctx.get("group_columns")[i]) for i in groups]
+    cores = [ctx.get("group_cores")[i] for i in groups]
     return MappingEstimate(
         mapping=mapping,
         group_cycles=tuple(ctx.get("group_cycles")[i] for i in groups),
-        dispatch_cycles=ctx.get("allocation").problem.hardware.dispatch_overhead_cycles(columns),
+        dispatch_cycles=ctx.get("allocation").problem.hardware.dispatch_overhead_cycles(cores),
         context=ctx,
     )

{stream_dse-2.3.0 → stream_dse-2.4.0}/stream/compiler/kernels/aie_kernel.py

@@ -8,7 +8,7 @@ from snaxc.ir.tsl import Stride, TiledStride, TiledStridedLayout
 from xdsl.dialects.builtin import AnyDenseElement, FunctionType, MemRefType, StringAttr, SymbolRefAttr, 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 import Attribute, Operation, Region, SSAValue
 from xdsl.ir.affine import AffineDimExpr
 from xdsl.pattern_rewriter import PatternRewriter
 from xdsl.printer import Printer
@@ -100,14 +100,16 @@ class KernelDeclaration(FuncOp):
 @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``, ``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."""
+    written at ``carried_over``, ``rows`` per step times the extent of ``indexed_by``, of ``element_type`` (the
+    output's where None). ``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: str
     indexed_by: str
     handover: int = 0
+    element_type: Attribute | None = None
 
 
 @dataclass(kw_only=True)
@@ -122,25 +124,37 @@ class AIEKernel(ABC):
     FALLBACK: ClassVar[Mapping[str, tuple[str, str | None]]] = {}
     """The symbol and object, ``None`` for the library's, stream declares for a call it names otherwise,
     where no provider binds the kernel."""
+    ACCUMULATOR: ClassVar[str | None] = None
+    """The format the kernel keeps a sum in across the calls adding into it, beside its output; None where it
+    adds into the output itself."""
+    SCRATCH_BITS: ClassVar[int] = 0
+    """Bits of a buffer the binding keeps on each core the kernel runs on, beside its operands and its state."""
 
     @property
     def unique_name(self) -> str:
         return self.function_name
 
+    @property
+    def library_name(self) -> str:
+        """The kernel library's entry for this kernel: its function, built as the library describes it."""
+        return self.function_name
+
     @property
     def spec(self) -> KernelSpec:
         if self.library is None:
             raise ValueError(f"{type(self).__name__} needs a kernel library")
-        spec = self.library.spec(self.function_name)
+        spec = self.library.spec(self.library_name)
         if spec is None:
-            raise ValueError(f"the kernel library does not describe {self.function_name}")
+            raise ValueError(f"the kernel library does not describe {self.library_name}")
         return spec
 
     @property
     def mac(self) -> Mapping[str, int]:
+        """The MAC tile the kernel computes in, which is also the tiling it leaves an operand in."""
         if self.library is None:
             raise ValueError(f"{type(self).__name__} needs a kernel library")
-        return self.library.mac
+        spec = self.library.spec(self.library_name)
+        return spec.mac if spec is not None and spec.mac else self.library.mac
 
     def call_shape(self) -> dict[str, int]:
         return {d.name: int(getattr(self, d.name)) for d in self.spec.dims}
@@ -179,13 +193,18 @@ class AIEKernel(ABC):
         """What a call to ``call`` resolves to: the provider's binding, or stream's own declaration."""
         spec = self.spec
         if spec.binding is not None:
-            entry = {"binding": spec.binding, "args": {**dims, "npu": bindings.npu}}
+            args = {**dims, **spec.options, **self.binding_options(), "npu": bindings.npu}
+            entry = {"binding": spec.binding, "args": args}
             return bindings.resolve(entry if call == self.function_name else entry | {"companion": call})
         symbol, object_file = self.FALLBACK.get(call, (call, None))
         if (object_file := object_file or spec.object) is None:
             raise ValueError(f"the kernel library names neither a binding nor an object for {self.function_name}")
         return bindings.resolve({"symbol": symbol, "object": object_file})
 
+    def binding_options(self) -> dict[str, bool]:
+        """What the provider builds the kernel's object with besides the call's dimensions."""
+        return {}
+
     @property
     @abstractmethod
     def function_name(self) -> str: ...