Name tensor axis dims in the graph view, memory sharing and operator types in the hardware IR

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,
         }

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)

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"],
         )

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
 

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"]]

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]