refactor get_mesh

stream/hardware/architecture/accelerator.py

@@ -48,14 +48,11 @@ def __jsonrepr__(self):
         return {"name": self.name, "cores": self.cores}
 
     def get_core(self, core_id: int) -> Core:
-        """
+        """s
         Return the core with id 'core_id'.
         Raises ValueError() when a core_id is not found in the available cores.
         """
-        core = next((core for core in self.core_list if core.id == core_id), None)
-        if core is None:
-            raise ValueError(f"Requested core with id {core_id} is not present in accelerator.")
-        return core
+        return self.cores.get_node_with_id(core_id)
 
     @property
     def core_list(self) -> list[Core]:

stream/hardware/architecture/noc/bus.py

@@ -3,25 +3,7 @@
 
 from stream.hardware.architecture.accelerator import CoreGraph
 from stream.hardware.architecture.noc.communication_link import CommunicationLink
-
-
-def have_shared_memory(a: Core, b: Core):
-    """Returns True if core a and core b have a shared top level memory
-
-    Args:
-        a (Core): First core
-        b (Core): Second core
-    """
-    top_level_memory_instances_a = set(
-        [level.memory_instance for level, out_degree in a.memory_hierarchy.out_degree() if out_degree == 0]
-    )
-    top_level_memory_instances_b = set(
-        [level.memory_instance for level, out_degree in b.memory_hierarchy.out_degree() if out_degree == 0]
-    )
-    for memory_instance_a in top_level_memory_instances_a:
-        if memory_instance_a in top_level_memory_instances_b:
-            return True
-    return False
+from stream.hardware.architecture.utils import get_bidirectional_edges
 
 
 def get_bus(
@@ -35,85 +17,41 @@ def get_bus(
     """Return a graph of the cores where each core is connected to a single bus.
 
     Args:
-        cores (list): list of core objects
-        bandwidth (int): bandwidth of the communication bus
-        unit_energy_cost (float): The unit energy cost of having a communication-link active. This does not include the
+        cores: list of core objects
+        bandwidth: bandwidth of the communication bus
+        unit_energy_cost: The unit energy cost of having a communication-link active. This does not include the
         involved memory read/writes.
-        pooling_core (Core, optional): If provided, the pooling core that is added.
-        simd_core (Core, optional): If provided, the simd core that is added.
-        offchip_core (Core, optional): If provided, the offchip core that is added.
-        offchip_bandwidth (int, optional): If offchip_core is provided, this is the
+        pooling_core: If provided, the pooling core that is added.
+        simd_core: If provided, the simd core that is added.
+        offchip_core: If provided, the offchip core that is added.
+        offchip_bandwidth: If offchip_core is provided, this is the
     """
-    bus = CommunicationLink("Any", "Any", bandwidth, unit_energy_cost)
+
+    def get_edges_bus(core_a: Core, core_b: Core):
+        return get_bidirectional_edges(core_a, core_b, bandwidth, unit_energy_cost, link_type="bus")
+
+    def get_edges_link(core_a: Core, core_b: Core):
+        return get_bidirectional_edges(core_a, core_b, bandwidth, unit_energy_cost, link_type="link")
 
     edges: list[tuple[Core, Core, dict[str, CommunicationLink]]] = []
     pairs = [(a, b) for idx, a in enumerate(cores) for b in cores[idx + 1 :]]
-    for pair in pairs:
-        (sender, receiver) = pair
-        if not have_shared_memory(sender, receiver):
-            edges.append((sender, receiver, {"cl": bus}))
-            edges.append((receiver, sender, {"cl": bus}))
+
+    for core_a, core_b in pairs:
+        edges += get_edges_bus(core_a, core_b)
 
     # If there is a pooling core, also add two edges from each core to the pooling core: one in each direction
     if pooling_core:
-        if not isinstance(pooling_core, Core):
-            raise ValueError("The given pooling_core is not a Core object.")
         for core in cores:
-            if not have_shared_memory(core, pooling_core):
-                edges.append(
-                    (
-                        core,
-                        pooling_core,
-                        {"cl": CommunicationLink(core, pooling_core, bandwidth, unit_energy_cost)},
-                    )
-                )
-                edges.append(
-                    (
-                        pooling_core,
-                        core,
-                        {"cl": CommunicationLink(pooling_core, core, bandwidth, unit_energy_cost)},
-                    )
-                )
+            edges += get_edges_link(core, pooling_core)
 
     # If there is a simd core, also add two edges from each core to the pooling core: one in each direction
-    # For now, assume the simd operations come for free, so bandwidth is infinite and unit energy cost is 0
-    float("inf")
     if simd_core:
-        if not isinstance(simd_core, Core):
-            raise ValueError("The given simd_core is not a Core object.")
         for core in cores:
-            if not have_shared_memory(core, simd_core):
-                edges.append(
-                    (
-                        core,
-                        simd_core,
-                        {"cl": bus},
-                    )
-                )
-                edges.append(
-                    (
-                        simd_core,
-                        core,
-                        {"cl": bus},
-                    )
-                )
-        # If there is a pooling core, also add two edges from/to the pooling core
-        if pooling_core:
-            if not have_shared_memory(pooling_core, simd_core):
-                edges.append(
-                    (
-                        pooling_core,
-                        simd_core,
-                        {"cl": bus},
-                    )
-                )
-                edges.append(
-                    (
-                        simd_core,
-                        pooling_core,
-                        {"cl": bus},
-                    )
-                )
+            edges += get_edges_bus(core, simd_core)
+
+    # If there is a pooling core, also add two edges from/to the pooling core
+    if pooling_core and simd_core:
+        edges += get_edges_bus(pooling_core, simd_core)
 
     # If there is an offchip core, add a single link for writing to and a single link for reading from the offchip
     if offchip_core:
@@ -134,8 +72,6 @@ def get_bus(
         else:
             to_offchip_link = CommunicationLink("Any", offchip_core, offchip_write_bandwidth, unit_energy_cost)
             from_offchip_link = CommunicationLink(offchip_core, "Any", offchip_read_bandwidth, unit_energy_cost)
-        if not isinstance(offchip_core, Core):
-            raise ValueError("The given offchip_core is not a Core object.")
         for core in cores:
             edges.append((core, offchip_core, {"cl": to_offchip_link}))
             edges.append((offchip_core, core, {"cl": from_offchip_link}))

stream/hardware/architecture/noc/communication_link.py

@@ -25,13 +25,13 @@ def __init__(
         self.receiver = receiver
         self.bandwidth = bandwidth
         self.unit_energy_cost = unit_energy_cost
-        self.bidirectional = bidirectional
+        self.bidirectional = bidirectional  # TODO this property is not in use?
 
-        self.events = []
+        self.events: list[CommunicationLinkEvent] = []
         self.active_periods = [(0, float("inf"), 0)]
         self.active_ts = np.array([0, float("inf")])
         self.active_deltas = np.array([0, 0])
-        self.tensors = {}
+        self.tensors: dict[Tensor, list[CommunicationLinkEvent]] = {}
 
     def __str__(self) -> str:
         return f"CommunicationLink({self.sender}, {self.receiver}, bw={self.bandwidth})"
@@ -40,15 +40,7 @@ def __repr__(self) -> str:
         return str(self)
 
     def __hash__(self) -> int:
-        return hash(
-            (
-                self.sender,
-                self.receiver,
-                self.bandwidth,
-                self.unit_energy_cost,
-                self.bidirectional,
-            )
-        )
+        return hash((self.sender, self.receiver, self.bandwidth, self.unit_energy_cost, self.bidirectional))
 
     def __eq__(self, other: object) -> bool:
         return isinstance(other, CommunicationLink) and (self.sender, self.receiver, self.bandwidth) == (
@@ -63,7 +55,7 @@ def get_name_for_schedule_plot(self) -> str:
         else:
             return f"{self.sender} -> {self.receiver}"
 
-    def transfer(self, cle: CommunicationLinkEvent) -> float:
+    def transfer(self, link_event: CommunicationLinkEvent) -> float:
         """Transfer data on this communication link at timestep.
         The transfer can take longer than necessary for this link if another lower-bandwidth link is involved.
 
@@ -75,8 +67,8 @@ def transfer(self, cle: CommunicationLinkEvent) -> float:
         Returns:
             int: The end time when communication on this link is finished
         """
-        energy_cost = cle.energy
-        self.update_activity(cle)
+        energy_cost = link_event.energy
+        self.update_activity(link_event)
         return energy_cost
 
     def block(
@@ -89,10 +81,10 @@ def block(
         """Block this communication link from start timestep for a given duration.
 
         Args:
-            start (int): The timestep at which the blocking starts.
-            duration (int): The duration of the blocking.
-            tensors (list): A list of tensors for which we are blocking the link.
-            activity (int): The bandwidth activity in bits/cc.
+            start: The timestep at which the blocking starts.
+            duration: The duration of the blocking.
+            tensors: A list of tensors for which we are blocking the link.
+            activity: The percentage of the link bandwidth used
         """
         end = start + duration
         # Create a CLEvent
@@ -137,8 +129,8 @@ def update_activity(self, event: CommunicationLinkEvent):
 
     def get_idle_window(self, activity: float, duration: int, earliest_t: int, tensors: list["Tensor"]):
         """
-        Get the earliest time window of duration 'duration' from 'earliest_t'
-        with atleast 'activity' percent available.
+        Get the earliest time window of duration `duration` from `earliest_t` with at least `activity` percent
+        available.
         """
         valid_windows: list[tuple[int, int]] = []
         ## Check if this tensor has already been transferred on this link before