stream-dse 2.1.0__tar.gz → 2.2.0__tar.gz

{stream_dse-2.1.0 → stream_dse-2.2.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: stream-dse
-Version: 2.1.0
+Version: 2.2.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
@@ -33,6 +33,7 @@ Requires-Python: >=3.12
 Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: zigzag-dse==3.9.1
+Requires-Dist: six
 Requires-Dist: cerberus
 Requires-Dist: ortools>=9.15
 Requires-Dist: pydantic<2.12,>=2.0
@@ -62,7 +63,7 @@ Dynamic: license-file
 [![Python 3.12+](https://img.shields.io/badge/python-3.12%2B-blue)](https://www.python.org/)
 [![Docs](https://img.shields.io/badge/docs-stream-blue)](https://kuleuven-micas.github.io/stream/)
 
-**Stream** is a design space exploration (DSE) and constraint-optimization framework for **heterogeneous dataflow accelerators**: accelerator systems built by combining cores that each have their own dataflow and performance model (**AIE** and **TPU-like** are two example core types among others). Scheduling is **layer-fused**, and the **constraint optimization** uses MILP (Mixed-Integer Linear Programming) to decide tensor placement and transfer paths across the cores of such a system. Stream builds on top of [ZigZag](https://zigzag-project.github.io/zigzag/) for per-core cost estimation.
+**Stream** is a design space exploration (DSE) and constraint-optimization framework for **heterogeneous dataflow accelerators**: accelerator systems built by combining cores that each have their own dataflow and performance model (**AIE** and **TPU-like** are two example core types among others). Scheduling is **layer-fused**, and the **constraint optimization** uses MILP (Mixed-Integer Linear Programming) to decide tensor placement and transfer paths between memories. Stream builds on top of [ZigZag](https://zigzag-project.github.io/zigzag/) for per-core cost estimation.
 
 ### 📖 [**Explore the Documentation**](https://kuleuven-micas.github.io/stream/)
 ### 🚀 [**Getting Started Guide**](https://kuleuven-micas.github.io/stream/getting-started/)
@@ -138,13 +139,13 @@ Price a small two-Conv workload (a committed test fixture) on a TPU-like quad-co
 
 ## 🧩 Hardware and Core Types
 
-An accelerator in Stream is described as a **system of heterogeneous dataflow cores**. Core roles include compute, memory, shim, and offchip; example dataflow core types include **AIE**, **TPU-like**, and **pooling**.
+An accelerator in Stream is described as a **system of heterogeneous dataflow cores** that compute out of **memories**, joined by links and buses. Several cores can share one memory, and off-chip DRAM is a memory too. Example dataflow core types include **AIE**, **TPU-like**, and **pooling**. See the [Hardware docs](https://kuleuven-micas.github.io/stream/hardware/) for the format.
 
 Hardware and mapping files are organized as follows:
 
 - `stream/inputs/examples/hardware/` - system-level hardware YAMLs (e.g. `tpu_like_quad_core.yaml`, `eyeriss_like_*.yaml`, `simba*.yaml`, `fusemax.yaml`).
-- `stream/inputs/examples/hardware/cores/` - per-core-type YAMLs (e.g. `tpu_like.yaml`, `pooling.yaml`, `simd.yaml`, `offchip.yaml`, `eyeriss_like.yaml`).
-- `stream/inputs/aie/hardware/` and `stream/inputs/aie/hardware/cores/` - AMD AIE example core types (e.g. `aie_tile.yaml`, `mem_tile_256KB.yaml`, `shim_dma.yaml`).
+- `stream/inputs/examples/hardware/cores/` - per-core-type YAMLs (e.g. `tpu_like.yaml`, `pooling.yaml`, `simd.yaml`, `eyeriss_like.yaml`).
+- `stream/inputs/aie/hardware/` and `stream/inputs/aie/hardware/cores/` - AMD AIE example core types (e.g. `aie_tile.yaml`, `mem_tile_256KB.yaml`).
 - `stream/inputs/examples/mapping/`, `stream/inputs/aie/mapping/`, and `stream/inputs/testing/mapping/` - mapping descriptions.
 
 A mapping can be generated (as in Quick Start above) or hand-written and passed as `mapping`.
@@ -210,9 +211,9 @@ with tempfile.TemporaryDirectory() as tmp:
     print("cycles:", estimate.cycles)
 ```
 
-A `MappingEstimate` holds `cycles`, the fused groups' estimates plus the reconfiguration the hardware declares, the per-group `group_cycles`, and the solved `context`, whose useful keys are `allocation`, `workload`, `accelerator` and `group_latencies`. `SolveOptions` sets the solver backend, the number of columns, the constraint families of the allocation model, the kernel library, the tile search, the solve's time limit and solver log, whether each solve writes its reports, traces and figures (`artifacts`), and instrumentation (such as `timing`), and its `stage_options` carries what a plugin's stages read, such as `fusion_cut_points` and `intra_core_tiling` for the generic mapping generator or `npu` and `trace_size` for the AIE code generator.
+A `MappingEstimate` holds `cycles`, the fused groups' estimates plus the reconfiguration the hardware declares, the per-group `group_cycles`, and the solved `context`, whose useful keys are `allocation`, `workload`, `accelerator` and `group_latencies`. `SolveOptions` sets the solver backend, the number of columns, the constraint families of the allocation model, the kernel library, the tile search, the solve's time limit and solver log, whether the run writes its pictures and each solve its reports, traces and figures (`artifacts`), and instrumentation (such as `timing`), and its `stage_options` carries what a plugin's stages read, such as `fusion_cut_points` and `intra_core_tiling` for the generic mapping generator or `npu` and `trace_size` for the AIE code generator.
 
-Whatever depends on the hardware is found through entry-point groups, so a separate package extends Stream without a fork: `stream.frontends` (workload formats), `stream.mapping_generators` (a mapping when none is given), `stream.namespaces` (namespace facts and the constraint families a namespace adds), `stream.constraint_families` (the constraint families the allocation model is built from), `stream.core_cost_backends` (per-core cost) and `stream.codegen_backends` (code generation).
+Whatever depends on the hardware is found through entry-point groups, so a separate package extends Stream without a fork: `stream.frontends` (workload formats), `stream.mapping_generators` (a mapping when none is given), `stream.namespaces` (namespace facts and the constraint families a namespace adds), `stream.constraint_families` (the constraint families the allocation model is built from), `stream.core_cost_backends` (per-core cost), `stream.codegen_backends` (code generation) and `stream.accelerator_keys` (top-level accelerator keys an overlay declares and checks; Stream reads none of them).
 
 ---
 

{stream_dse-2.1.0 → stream_dse-2.2.0}/README.md

@@ -4,7 +4,7 @@
 [![Python 3.12+](https://img.shields.io/badge/python-3.12%2B-blue)](https://www.python.org/)
 [![Docs](https://img.shields.io/badge/docs-stream-blue)](https://kuleuven-micas.github.io/stream/)
 
-**Stream** is a design space exploration (DSE) and constraint-optimization framework for **heterogeneous dataflow accelerators**: accelerator systems built by combining cores that each have their own dataflow and performance model (**AIE** and **TPU-like** are two example core types among others). Scheduling is **layer-fused**, and the **constraint optimization** uses MILP (Mixed-Integer Linear Programming) to decide tensor placement and transfer paths across the cores of such a system. Stream builds on top of [ZigZag](https://zigzag-project.github.io/zigzag/) for per-core cost estimation.
+**Stream** is a design space exploration (DSE) and constraint-optimization framework for **heterogeneous dataflow accelerators**: accelerator systems built by combining cores that each have their own dataflow and performance model (**AIE** and **TPU-like** are two example core types among others). Scheduling is **layer-fused**, and the **constraint optimization** uses MILP (Mixed-Integer Linear Programming) to decide tensor placement and transfer paths between memories. Stream builds on top of [ZigZag](https://zigzag-project.github.io/zigzag/) for per-core cost estimation.
 
 ### 📖 [**Explore the Documentation**](https://kuleuven-micas.github.io/stream/)
 ### 🚀 [**Getting Started Guide**](https://kuleuven-micas.github.io/stream/getting-started/)
@@ -80,13 +80,13 @@ Price a small two-Conv workload (a committed test fixture) on a TPU-like quad-co
 
 ## 🧩 Hardware and Core Types
 
-An accelerator in Stream is described as a **system of heterogeneous dataflow cores**. Core roles include compute, memory, shim, and offchip; example dataflow core types include **AIE**, **TPU-like**, and **pooling**.
+An accelerator in Stream is described as a **system of heterogeneous dataflow cores** that compute out of **memories**, joined by links and buses. Several cores can share one memory, and off-chip DRAM is a memory too. Example dataflow core types include **AIE**, **TPU-like**, and **pooling**. See the [Hardware docs](https://kuleuven-micas.github.io/stream/hardware/) for the format.
 
 Hardware and mapping files are organized as follows:
 
 - `stream/inputs/examples/hardware/` - system-level hardware YAMLs (e.g. `tpu_like_quad_core.yaml`, `eyeriss_like_*.yaml`, `simba*.yaml`, `fusemax.yaml`).
-- `stream/inputs/examples/hardware/cores/` - per-core-type YAMLs (e.g. `tpu_like.yaml`, `pooling.yaml`, `simd.yaml`, `offchip.yaml`, `eyeriss_like.yaml`).
-- `stream/inputs/aie/hardware/` and `stream/inputs/aie/hardware/cores/` - AMD AIE example core types (e.g. `aie_tile.yaml`, `mem_tile_256KB.yaml`, `shim_dma.yaml`).
+- `stream/inputs/examples/hardware/cores/` - per-core-type YAMLs (e.g. `tpu_like.yaml`, `pooling.yaml`, `simd.yaml`, `eyeriss_like.yaml`).
+- `stream/inputs/aie/hardware/` and `stream/inputs/aie/hardware/cores/` - AMD AIE example core types (e.g. `aie_tile.yaml`, `mem_tile_256KB.yaml`).
 - `stream/inputs/examples/mapping/`, `stream/inputs/aie/mapping/`, and `stream/inputs/testing/mapping/` - mapping descriptions.
 
 A mapping can be generated (as in Quick Start above) or hand-written and passed as `mapping`.
@@ -152,9 +152,9 @@ with tempfile.TemporaryDirectory() as tmp:
     print("cycles:", estimate.cycles)
 ```
 
-A `MappingEstimate` holds `cycles`, the fused groups' estimates plus the reconfiguration the hardware declares, the per-group `group_cycles`, and the solved `context`, whose useful keys are `allocation`, `workload`, `accelerator` and `group_latencies`. `SolveOptions` sets the solver backend, the number of columns, the constraint families of the allocation model, the kernel library, the tile search, the solve's time limit and solver log, whether each solve writes its reports, traces and figures (`artifacts`), and instrumentation (such as `timing`), and its `stage_options` carries what a plugin's stages read, such as `fusion_cut_points` and `intra_core_tiling` for the generic mapping generator or `npu` and `trace_size` for the AIE code generator.
+A `MappingEstimate` holds `cycles`, the fused groups' estimates plus the reconfiguration the hardware declares, the per-group `group_cycles`, and the solved `context`, whose useful keys are `allocation`, `workload`, `accelerator` and `group_latencies`. `SolveOptions` sets the solver backend, the number of columns, the constraint families of the allocation model, the kernel library, the tile search, the solve's time limit and solver log, whether the run writes its pictures and each solve its reports, traces and figures (`artifacts`), and instrumentation (such as `timing`), and its `stage_options` carries what a plugin's stages read, such as `fusion_cut_points` and `intra_core_tiling` for the generic mapping generator or `npu` and `trace_size` for the AIE code generator.
 
-Whatever depends on the hardware is found through entry-point groups, so a separate package extends Stream without a fork: `stream.frontends` (workload formats), `stream.mapping_generators` (a mapping when none is given), `stream.namespaces` (namespace facts and the constraint families a namespace adds), `stream.constraint_families` (the constraint families the allocation model is built from), `stream.core_cost_backends` (per-core cost) and `stream.codegen_backends` (code generation).
+Whatever depends on the hardware is found through entry-point groups, so a separate package extends Stream without a fork: `stream.frontends` (workload formats), `stream.mapping_generators` (a mapping when none is given), `stream.namespaces` (namespace facts and the constraint families a namespace adds), `stream.constraint_families` (the constraint families the allocation model is built from), `stream.core_cost_backends` (per-core cost), `stream.codegen_backends` (code generation) and `stream.accelerator_keys` (top-level accelerator keys an overlay declares and checks; Stream reads none of them).
 
 ---
 

{stream_dse-2.1.0 → stream_dse-2.2.0}/pyproject.toml

@@ -8,22 +8,20 @@ include = ["stream*"]  # only ship the stream package, not top-level scripts/out
 namespaces = true
 
 [tool.setuptools.package-data]
-# Ship the static hardware-description YAMLs so a pip-installed package can run
-# codegen without a checkout. (.py-only is the setuptools default; declare the
-# data files explicitly.) cores/ holds the per-tile descriptions and kernels/ the
-# kernel library referenced by the top-level accelerator YAMLs. Test
-# workloads/onnx fixtures are intentionally NOT shipped.
+# Ship the hardware YAMLs and kernels, testing hardware included for its generators, so a pip install needs no checkout.
 "stream" = [
     "inputs/aie/hardware/*.yaml",
     "inputs/aie/hardware/cores/*.yaml",
     "inputs/aie/kernels/*.toml",
     "inputs/examples/hardware/*.yaml",
     "inputs/examples/hardware/cores/*.yaml",
+    "inputs/testing/hardware/*.yaml",
+    "inputs/testing/hardware/cores/*.yaml",
 ]
 
 [project]
 name = "stream-dse"
-version = "2.1.0"
+version = "2.2.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" }]
@@ -35,6 +33,8 @@ classifiers = [
 keywords = ["stream", "multi-core", "accelerator", "layer-fused", "scheduling", "zigzag", "dse", "design-space-exploration", "machine-learning", "deep-learning", "mapping"]
 dependencies = [
     "zigzag-dse==3.9.1",
+    # zigzag's mapping validator imports six without declaring it
+    "six",
     "cerberus",
     "ortools>=9.15",
     "pydantic>=2.0,<2.12",
@@ -100,11 +100,14 @@ aie2 = "stream.stages.codegen.backends:AIE2_CODEGEN"
 [project.entry-points."stream.mapping_generators"]
 generic = "stream.stages.generation.mapping_generators:GENERIC_MAPPING"
 
-# A plugin distribution declares itself here so it can be allowlisted (STREAM_ALLOWED_OVERLAYS) and
+# A plugin distribution declares itself here so it can be allowlisted (STREAM_OVERLAYS) and
 # ordered by an OVERLAY_PRIORITY int on its module -- public entry points are priority 0, so a plugin
 # outranks a built-in of the same name. The value is the distribution's module.
 [project.entry-points."stream.overlays"]
 
+# Overlay accelerator keys: the name is the key, the object a cerberus `rule` and `check(accelerator, value)`.
+[project.entry-points."stream.accelerator_keys"]
+
 [project.entry-points."stream.workload_blocks"]
 # Built-in blocks register in-tree (stream.workload.blocks); overlays add their own entries here.
 
@@ -122,7 +125,7 @@ generic = "stream.stages.generation.mapping_generators:GENERIC_MAPPING"
 Homepage = "https://github.com/KULeuven-MICAS/stream"
 
 [tool.bumpver]
-current_version = "2.1.0"
+current_version = "2.2.0"
 version_pattern = "MAJOR.MINOR.PATCH"
 commit_message = "bump version {old_version} -> {new_version}"
 commit = true

stream_dse-2.2.0/stream/__init__.py

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

{stream_dse-2.1.0 → stream_dse-2.2.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.dev73+g18ca6c1"
+_MLIR_AIE_PIN = "mlir_aie==1.4.4"
 _LLVM_AIE_PIN = "llvm-aie==22.0.0.2026092101+0006955e"
 _AIE_WHEEL_INDICES = [
     "--extra-index-url",
-    "https://github.com/Xilinx/mlir-aie/releases/expanded_assets/latest-wheels-4",
+    "https://github.com/Xilinx/mlir-aie/releases/expanded_assets/v1.4.4",
     "--extra-index-url",
     "https://github.com/Xilinx/llvm-aie/releases/expanded_assets/nightly",
     "--extra-index-url",

{stream_dse-2.1.0 → stream_dse-2.2.0}/stream/allocation/allocation.py

@@ -5,9 +5,12 @@ from dataclasses import dataclass
 from math import prod
 from typing import TYPE_CHECKING, Any
 
+from stream.allocation.steady_state import steady_state_detail
+from stream.cost_model.roofline import group_point, hardware_roofline
 from stream.hardware.architecture.core import Core
+from stream.mapping.mapping import held_bits
 from stream.opt.allocation.constraint_optimization.report import guarded
-from stream.workload.node import Tensor
+from stream.workload.node import InEdge, OutEdge, Tensor, declared_of, edge_tensor
 from stream.workload.steady_state.iteration_space import (
     IterationVariableType,
     LoopEffect,
@@ -19,6 +22,7 @@ if TYPE_CHECKING:
     from stream.allocation.problem import AllocationProblem
     from stream.allocation.solution import AllocationSolution
     from stream.mapping.mapping import Mapping
+    from stream.opt.allocation.constraint_optimization.space import DecisionSpace
     from stream.workload.node import ComputationNode, HasIterationSpace
     from stream.workload.workload import Workload
 
@@ -39,8 +43,8 @@ _LOOP_NEST_DEPTH: dict[str, int] = {
 @dataclass(frozen=True)
 class Allocation:
     """A solved steady state as downstream reads it: the problem it solves, the mapping and iteration
-    spaces the solution decided, the constraint families it was built from with their options, and the
-    solution itself."""
+    spaces the solution decided, the constraint families it was built from with their options, the
+    solution itself and, where kept, the decision space it was solved in."""
 
     problem: AllocationProblem
     mapping: Mapping
@@ -48,6 +52,7 @@ class Allocation:
     backend: str
     families: tuple[tuple[str, dict[str, Any]], ...]
     solution: AllocationSolution
+    space: DecisionSpace | None = None
 
     @property
     def estimated_cycles(self) -> float:
@@ -56,8 +61,9 @@ class Allocation:
         return float(self.solution.latency.total)
 
     def get_ir(self) -> dict:
-        """The allocation as a plain dict: latencies, solve configuration, statistics and families, fusion splits,
-        mapping, performance report and the steady-state inspection view, None when it cannot be built."""
+        """The allocation as a plain dict: latencies, solve configuration, statistics, families, fusion splits, mapping,
+        each input and output edge's tensor and the bits of it each off-chip memory holds, performance report, and the
+        steady-state view and roofline point, each None when it cannot be built."""
         stats = self.solution.solve_stats
         latency = self.solution.latency
         return {
@@ -79,11 +85,55 @@ class Allocation:
             },
             "families": [{"name": name, "options": options} for name, options in self.families],
             "fusion_splits": {str(dim): size for dim, size in self.problem.fusion_splits.items()},
-            "mapping": self.mapping.get_ir(),
+            "mapping": self._mapping_ir(),
+            "edge_tensors": self._edge_tensors(),
+            "offchip_holdings": guarded("off-chip holdings", self._offchip_holdings),
             "performance": self.solution.performance,
             "steady_state": guarded("steady-state IR view", self._steady_state_ir),
+            "roofline": guarded("roofline point", self.roofline_point),
         }
 
+    def _edges(self) -> dict[str, InEdge | OutEdge]:
+        return {edge.name: edge for edge in self.problem.source_workload.nodes if isinstance(edge, InEdge | OutEdge)}
+
+    def _mapping_ir(self) -> dict[str, Any]:
+        """The mapping, each edge's placement split counted in the axes the workload declared its tensor in."""
+        mapping, edges = self.mapping.get_ir(), self._edges()
+        for name, placement in mapping["edge_placements"].items():
+            if placement["split"] is not None:
+                placement["split"] = declared_of(edges[name]).axes[placement["split"]]
+        return mapping
+
+    def _edge_tensors(self) -> dict[str, dict[str, Any]]:
+        """The whole tensor each input and output edge of the group brings in or takes out, as the workload declared
+        it."""
+        return {
+            name: {
+                "tensor": edge_tensor(edge).name,
+                "shape": list((declared := declared_of(edge)).shape),
+                "axes": list(declared.axes),
+                "element_bits": edge_tensor(edge).operand_type.bitwidth,
+                "direction": "in" if isinstance(edge, InEdge) else "out",
+            }
+            for name, edge in self._edges().items()
+        }
+
+    def _offchip_holdings(self) -> dict[str, dict[str, float]]:
+        """The bits of each tensor of the group's edges each off-chip memory holds, by the placements the solve chose
+        (:func:`held_bits`)."""
+        problem = self.problem
+        mapping = problem.source_mapping.with_edge_placements(self.mapping.edge_placements)
+        held = held_bits(problem.accelerator, [problem.source_workload], [mapping])
+        return {
+            memory.name: {tensor.name: bits for tensor, bits in tensors.items()}
+            for memory, tensors in held.items()
+            if memory.offchip
+        }
+
+    def roofline_point(self) -> dict[str, Any]:
+        """Where this group sits on its accelerator's roofline."""
+        return group_point(self, hardware_roofline(self.problem.accelerator))
+
     def _core_loops(self, cn: ComputationNode) -> list[dict]:
         """The loop nest inside one core (ZigZag mapping), as ``core_*`` loops; empty for a non-ZigZag core."""
         lut_node = next((n for n in self.problem.cost_lut.get_nodes() if n.name == cn.name), None)
@@ -115,7 +165,8 @@ class Allocation:
         return loops
 
     def _steady_state_ir(self) -> dict:
-        """Serialise the tiled/steady-state inspection view: operators, loop nest and transfer graph."""
+        """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."""
         operators = [
             {
                 "name": cn.name,
@@ -160,11 +211,18 @@ class Allocation:
                     "kind": "transfer",
                     "transfer_type": getattr(transfer_type, "name", None),
                     "tensor": out.name if out is not None else None,
+                    "copies": len(tn.outputs),
                     "elements": int(prod(out.shape)) if out is not None else 0,
                 }
             )
         edges = [{"source": s.name, "target": t.name} for s, t in self.problem.workload.edges()]
-        return {"operators": operators, "loops": loops, "tiled_graph": {"nodes": tiled_nodes, "edges": edges}}
+        view = {"operators": operators, "loops": loops, "tiled_graph": {"nodes": tiled_nodes, "edges": edges}}
+        if self.space is None:
+            return view
+        detail = steady_state_detail(self, self.space)
+        for node in tiled_nodes:
+            node.update(detail["nodes"].get(node["name"], {}))
+        return view | {key: detail[key] for key in ("iterations", "axes", "tensors")}
 
 
 def solved_iteration_spaces(
@@ -219,14 +277,22 @@ def _propagate_spatial_reuse(ssis: IterationSpaces, spatial_side: Tensor, tempor
 
 
 def solved_mapping(workload: Workload, mapping: Mapping, solution: AllocationSolution) -> Mapping:
-    """A copy of ``mapping`` with each transfer bound to the route and memory cores the solution chose."""
-    solved = mapping.copy()
+    """A copy of ``mapping`` with each transfer bound to the route and memories the solution chose, and each input
+    whose placement the solution chose placed so."""
+    chosen = {
+        edge.name: route.placement
+        for tr, route in solution.transfer_routes.items()
+        if route.placement is not None
+        for edge in workload.predecessors(tr)
+    }
+    solved = mapping.with_edge_placements(chosen).copy()
+    placements = {tr: (placement,) for tr, placement in solution.memory_placements.items()}
     for tr, route in solution.transfer_routes.items():
         solved.set_for_node(
             tr,
             resource_allocation=(route,),
             inter_core_tiling=tuple(),
-            memory_allocation=solution.memory_cores.get(tr, tuple()),
+            memory_allocation=placements.get(tr, ()),
         )
     for tr in workload.get_transfer_nodes():
         assert len(solved.get(tr).resource_allocation) == 1, (