From df410a89631de283ee8c8df144cccd51d3293e37 Mon Sep 17 00:00:00 2001 From: Arne Symons Date: Tue, 6 Oct 2026 17:23:52 +0200 Subject: [PATCH] Name each graph tensor axis's unique dims, and the cores sharing a memory and a core's operators in the hardware IR --- stream/hardware/architecture/accelerator.py | 9 +++++++ stream/hardware/architecture/core.py | 1 + stream/ir/accelerator.py | 8 ++++-- stream/ir/graph_view.py | 30 ++++++++++++++++----- tests/unit/test_graph_view.py | 12 +++++++++ tests/unit/test_ir_models.py | 10 +++++++ 6 files changed, 61 insertions(+), 9 deletions(-) diff --git a/stream/hardware/architecture/accelerator.py b/stream/hardware/architecture/accelerator.py index 0751feee5..477056f01 100644 --- a/stream/hardware/architecture/accelerator.py +++ b/stream/hardware/architecture/accelerator.py @@ -60,6 +60,14 @@ def get_core(self, core_id: int) -> Core: def nb_shared_mem_groups(self) -> int: return len(set(self.shared_mem_group_ids.values())) + @property + def core_memory_sharing(self) -> list[list[int]]: + """The ids of each group of cores that share one top-level memory.""" + groups: dict[int, list[int]] = {} + for core_id, owner in sorted(self.shared_mem_group_ids.items()): + groups.setdefault(owner, []).append(core_id) + return [ids for ids in groups.values() if len(ids) > 1] + def memory_of(self, core: Core) -> Core: """The core owning the top-level memory ``core`` uses: the first of its `core_memory_sharing` group.""" return self.get_core(self.shared_mem_group_ids.get(core.id, core.id)) @@ -147,6 +155,7 @@ def get_ir(self) -> dict: "num_cores": len(list(self.cores.nodes)), "offchip_core_id": self.offchip_core_id, "nb_shared_mem_groups": self.nb_shared_mem_groups, + "core_memory_sharing": self.core_memory_sharing, "cores": cores_ir, "core_connectivity": links_ir, } diff --git a/stream/hardware/architecture/core.py b/stream/hardware/architecture/core.py index 3667d63ec..7c1a52d53 100644 --- a/stream/hardware/architecture/core.py +++ b/stream/hardware/architecture/core.py @@ -177,6 +177,7 @@ def get_ir(self) -> dict: "row_id": self.row_id, "col_id": self.col_id, "utilization": self.utilization, + "operator_types": getattr(self, "operator_types", None), } # Merge backend-specific fields (uniform protocol) diff --git a/stream/ir/accelerator.py b/stream/ir/accelerator.py index 558a2ffe8..a423aa159 100644 --- a/stream/ir/accelerator.py +++ b/stream/ir/accelerator.py @@ -28,8 +28,8 @@ class CoreIR(BaseModel): name: str = Field(description="Core human-readable name") core_type: str = Field(description="Core type namespace, e.g. 'aie2.compute' or 'zigzag.compute'") type: str = Field(description="Core role: 'compute', 'offchip', or 'dma'") - row_id: int = Field(description="Row position in the 2-D core grid (-1 for offchip)") - col_id: int = Field(description="Column position in the 2-D core grid (-1 for offchip)") + row_id: int | None = Field(description="Row position in the 2-D core grid (-1 for offchip, None when ungridded)") + col_id: int | None = Field(description="Column position in the 2-D core grid (-1 for offchip, None when ungridded)") utilization: float = Field(description="Core utilization ratio in [0, 1]") extra_fields: dict[str, Any] = Field( default_factory=dict, @@ -85,6 +85,9 @@ class AcceleratorIR(BaseModel): description="ID of the offchip memory core, or None if the accelerator has no offchip core" ) nb_shared_mem_groups: int = Field(description="Number of shared memory groups in the accelerator") + core_memory_sharing: list[list[int]] = Field( + default_factory=list, description="The ids of each group of cores that share one top-level memory" + ) cores: list[CoreIR] = Field(description="All cores with common and type-specific fields") core_connectivity: list[dict[str, Any]] = Field( description="Connectivity entries: bus (bidirectional, multiple cores) or link (directed pair)" @@ -120,6 +123,7 @@ def from_internal(cls, accelerator: Accelerator) -> AcceleratorIR: num_cores=raw["num_cores"], offchip_core_id=raw["offchip_core_id"], nb_shared_mem_groups=raw["nb_shared_mem_groups"], + core_memory_sharing=raw["core_memory_sharing"], cores=cores, core_connectivity=raw["core_connectivity"], ) diff --git a/stream/ir/graph_view.py b/stream/ir/graph_view.py index 3a8a842a5..9e7ec402d 100644 --- a/stream/ir/graph_view.py +++ b/stream/ir/graph_view.py @@ -6,6 +6,7 @@ from typing import TYPE_CHECKING, Literal from pydantic import BaseModel, ConfigDict, Field +from xdsl.ir.affine import AffineExpr from stream.datatypes import LayerDim from stream.workload.affine_access import footprint, map_dim_positions @@ -50,6 +51,9 @@ class TensorRefIR(BaseModel): name: str shape: list[int] = Field(default_factory=list, description="Dimension sizes, in tensor-axis order") + dims: list[str] = Field( + default_factory=list, description="Per tensor axis, the unique dims it spans, as the nodes' dims write them" + ) class OperandReuseIR(BaseModel): @@ -173,7 +177,7 @@ def from_workload(cls, workload: Workload, fusion_capacity: int | None = None) - nodes = [_node_ir(by_name[name], dims, block_of, region_of, proposed_of) for name in order] edges = [] for s, t in workload.edges: - shared = _shared(s, t) + shared = _shared(s, t, dims) edges.append( GraphEdgeIR(source=s.name, target=t.name, shared_tensors=[x.name for x in shared], tensors=shared) ) @@ -197,6 +201,10 @@ def __init__(self, workload: Workload): self._global_idxs = workload.global_idxs _, self._expressions = workload.unique_dimensions() self._size = workload.get_dimension_size + self._workload = workload + + def neighbours(self, node) -> list: + return [*self._workload.predecessors(node), *self._workload.successors(node)] def dims(self, node: HasIterationSpace) -> list[LayerDim]: span = self._global_idxs[node] @@ -217,17 +225,25 @@ def _topo_names(workload: Workload) -> list[str]: return [n.name for n in workload.dataflow_sort()] -def _tensor_ref(t: Tensor) -> TensorRefIR: +def _tensor_ref(t: Tensor, resolver: _DimResolver, *nodes) -> TensorRefIR: + """``t`` with its shape and, read through the first of ``nodes`` that iterates over it, each axis' unique dims.""" try: shape = [int(s) for s in t.shape] except Exception: # noqa: BLE001 -- a symbolic/unknown shape renders as no dims, not a crash shape = [] - return TensorRefIR(name=t.name, shape=shape) + node = next((n for n in nodes if isinstance(n, HasIterationSpace) and t in n.tensors), None) + dims = [] if node is None else [_axis(r, resolver.dims(node)) for r in node.get_mapping(t).results] + return TensorRefIR(name=t.name, shape=shape, dims=dims) + + +def _axis(expr: AffineExpr, dims: list[LayerDim]) -> str: + text = str(expr.replace_dims_and_symbols(dims, [])).replace("+ -", "- ") + return text[1:-1] if text.startswith("(") and text.endswith(")") else text -def _shared(src, dst) -> list[TensorRefIR]: +def _shared(src, dst, resolver: _DimResolver) -> list[TensorRefIR]: if isinstance(src, HasOutputs) and isinstance(dst, HasInputs): - return [_tensor_ref(t) for t in src.outputs if t in dst.inputs] + return [_tensor_ref(t, resolver, src, dst) for t in src.outputs if t in dst.inputs] return [] @@ -333,9 +349,9 @@ def _node_ir( ) # Boundary nodes (graph inputs/outputs) carry a single tensor -- keep its shape. if isinstance(node, InEdge) and node.outputs: - ir.tensor = _tensor_ref(node.outputs[0]) + ir.tensor = _tensor_ref(node.outputs[0], resolver, *resolver.neighbours(node)) elif isinstance(node, OutEdge) and node.inputs: - ir.tensor = _tensor_ref(node.inputs[0]) + ir.tensor = _tensor_ref(node.inputs[0], resolver, *resolver.neighbours(node)) if not isinstance(node, HasIterationSpace): return ir diff --git a/tests/unit/test_graph_view.py b/tests/unit/test_graph_view.py index fe53cc88e..37313dfdb 100644 --- a/tests/unit/test_graph_view.py +++ b/tests/unit/test_graph_view.py @@ -6,6 +6,8 @@ from __future__ import annotations +from stream.frontends import load_workload +from stream.inputs.testing.workload.make_2_conv import TwoConvWorkloadConfig, make_2_conv_workload from stream.inputs.testing.workload.make_scan import make_scan_workload from stream.ir import WorkloadGraphView from stream.parser.onnx.model import ONNXModelParser @@ -161,3 +163,13 @@ def test_a_loaded_workload_surfaces_proposed_regions_under_a_capacity(): view = _view_of_file("stream/inputs/testing/workload/attention_head.onnx", fusion_capacity=UNBOUNDED) assert view["proposed_regions"] assert any(n["proposed_region"] is not None for n in view["nodes"]) + + +def test_edges_name_the_unique_dims_of_each_tensor_axis(): + """Between two convs a tensor's axes are the first conv's output dims; the graph input slides each kernel + window over the output rows and columns.""" + config = TwoConvWorkloadConfig(1, 32, 32, 8, 16, 32, 3, "bf16", "bf16") + view = _view(load_workload(make_2_conv_workload(config))) + tensors = {(e.source, e.target): e.tensors for e in view.edges} + assert [t.dims for t in tensors["Conv1", "Conv2"]] == [["z3", "z8", "z5", "z4"]] + assert [t.dims for t in tensors["input", "Conv1"]] == [["z3", "z2", "(z5 + z1) - 1", "(z4 + z0) - 1"]] diff --git a/tests/unit/test_ir_models.py b/tests/unit/test_ir_models.py index 03e448a9a..5fc7e63fc 100644 --- a/tests/unit/test_ir_models.py +++ b/tests/unit/test_ir_models.py @@ -33,6 +33,7 @@ WorkloadAlgorithmicView, WorkloadCompilerView, ) +from stream.stages.parsing.accelerator_parser import parse_accelerator # --------------------------------------------------------------------------- # Fixtures: synthetic dicts matching get_ir() shapes @@ -101,6 +102,7 @@ "num_cores": 3, "offchip_core_id": 2, "nb_shared_mem_groups": 1, + "core_memory_sharing": [], "cores": [ { "id": 0, @@ -770,3 +772,11 @@ def test_allocation_ir_records_which_overlays_were_loaded(monkeypatch): ir = AllocationIR.from_internal(allocation) assert ir.overlays == ["vendor-overlay", "vendor-overlay-acme"] assert ir.model_dump()["overlays"] == ["vendor-overlay", "vendor-overlay-acme"] + + +def test_the_hardware_ir_names_the_cores_sharing_a_memory_and_what_a_specialized_core_runs(): + """FuseMax's array and vector core share one memory; only the vector core restricts its operators.""" + ir = AcceleratorIR.from_internal(parse_accelerator("stream/inputs/examples/hardware/fusemax.yaml")) + assert ir.core_memory_sharing == [[0, 1]] + served = {core.id: core.extra_fields["operator_types"] for core in ir.cores} + assert served[0] is None and "MaxPool" in served[1]