From cdfa78e69fd423af55e9449e3449c51aecec5b99 Mon Sep 17 00:00:00 2001
From: Robin Geens <35200290+RobinGeens@users.noreply.github.com>
Date: Mon, 7 Oct 2024 12:39:58 +0200
Subject: [PATCH] Remove "core" concept from ZigZag (#89)

* update pre-commit version

* move DiGraphWrapper from Stream to ZigZag

* fix typeguard issues

* disable pylint issue with super() method

* add mem sharing group to mem instance

* fix bug: shared_memory_group_id not stored in instance

* define class instance types in class body, to ensure they are included in stubs

* memory_sharing_list -> mem_sharing_list

* change open_yaml return type

* Remove Accelerator layer and rename core into accelerator

* fix typeguard issue'

* make has_same_performance method for all classes that make up an accelerator

* **kwargs should be of type Any
---
 example.py                                    |   2 +-
 zigzag/cost_model/cost_model.py               |  39 ++---
 zigzag/cost_model/cost_model_imc.py           |   5 +-
 zigzag/hardware/architecture/accelerator.py   | 146 +++++++++++++++---
 zigzag/hardware/architecture/core.py          | 140 -----------------
 .../hardware/architecture/memory_hierarchy.py |   7 +-
 .../hardware/architecture/memory_instance.py  |  37 ++++-
 zigzag/hardware/architecture/memory_level.py  |   9 ++
 zigzag/mapping/mapping.py                     |   4 +-
 zigzag/opt/loma/engine.py                     |   3 +-
 zigzag/opt/loma/memory_allocator.py           |   3 +-
 zigzag/parser/accelerator_factory.py          |  36 ++---
 .../evaluation/cost_model_evaluation.py       |   4 +-
 zigzag/stages/exploit_data_locality_stages.py |  28 +---
 .../mapping/spatial_mapping_conversion.py     |   3 +-
 .../mapping/spatial_mapping_generation.py     |  23 +--
 zigzag/stages/results/visualization.py        |   2 +-
 zigzag/stages/stage.py                        |   6 +-
 zigzag/stages/workload_iterator.py            |   4 +-
 .../visualization/graph/memory_hierarchy.py   |   2 +-
 zigzag/visualization/results/plot_cme.py      |   4 +-
 zigzag/visualization/results/print_mapping.py |   3 +-
 zigzag/workload/layer_node_abc.py             |   3 +
 23 files changed, 250 insertions(+), 263 deletions(-)
 delete mode 100644 zigzag/hardware/architecture/core.py

diff --git a/example.py b/example.py
index 7e6f8e1de..4b0fae918 100644
--- a/example.py
+++ b/example.py
@@ -8,7 +8,7 @@
 from zigzag.visualization.results.print_mapping import print_mapping
 
 model = "resnet"
-workload_path = "zigzag/inputs/workload/resnet18.yaml"  # or "zigzag/inputs/workload/resnet18.onnx"
+workload_path = "zigzag/inputs/workload/resnet18.onnx"  # or "zigzag/inputs/workload/resnet18.yaml"
 accelerator_path = "zigzag/inputs/hardware/tpu_like.yaml"
 mapping_path = "zigzag/inputs/mapping/tpu_like.yaml"
 pickle_filename = f"outputs/{model}-saved_list_of_cmes.pickle"
diff --git a/zigzag/cost_model/cost_model.py b/zigzag/cost_model/cost_model.py
index 2541d5487..028d03b96 100644
--- a/zigzag/cost_model/cost_model.py
+++ b/zigzag/cost_model/cost_model.py
@@ -31,6 +31,8 @@ class MemoryUtilization(TypedDict):
 class CostModelEvaluationABC(metaclass=ABCMeta):
     """! Superclass for CostModelEvaluation and CumulativeCME"""
 
+    accelerator: Accelerator
+
     @abstractmethod
     def __init__(self) -> None:
         # Attributes that all subclasses should define
@@ -54,6 +56,10 @@ def __init__(self) -> None:
 
         self.accelerator: Accelerator | None
 
+    @property
+    def core(self):
+        return self.accelerator
+
     def __add__(self, other: "CostModelEvaluationABC") -> "CumulativeCME":
         result = CumulativeCME()
 
@@ -130,10 +136,8 @@ def __add__(self, other: "CostModelEvaluationABC") -> "CumulativeCME":
         for cme in (self, other):
             if isinstance(cme, CostModelEvaluation):
                 result.layer_ids.append(cme.layer.id)
-                result.core_ids.append(cme.core_id)
             elif isinstance(cme, CumulativeCME):
                 result.layer_ids += cme.layer_ids
-                result.core_ids += cme.core_ids
 
         # Select accelerator for result
         if isinstance(self, CumulativeCME) and self.accelerator is None:
@@ -252,7 +256,6 @@ def __init__(self):
         self.memory_word_access: dict[LayerOperand, list[MemoryAccesses]] = dict()
 
         self.layer_ids: list[int] = []
-        self.core_ids: list[int] = []
 
         self.mac_energy: float = 0.0
         self.mem_energy: float = 0.0
@@ -307,25 +310,22 @@ def __init__(
         @param layer the layer to run
         @param access_same_data_considered_as_no_access (optional)
         """
-        self.accelerator: Accelerator = accelerator  # type: ignore
+        self.accelerator = accelerator
         self.layer: LayerNode = layer
         self.spatial_mapping = spatial_mapping
         self.spatial_mapping_int = spatial_mapping_int  # the original spatial mapping without decimal
         self.temporal_mapping = temporal_mapping
         self.access_same_data_considered_as_no_access = access_same_data_considered_as_no_access
 
-        self.core_id = layer.core_allocation[0]
-        core = accelerator.get_core(self.core_id)
-        self.mem_level_list = core.memory_hierarchy.mem_level_list
-        self.mem_hierarchy_dict = core.mem_hierarchy_dict
-        self.mem_size_dict = core.mem_size_dict
-        self.mem_r_bw_dict, self.mem_w_bw_dict = core.get_memory_bw_dict()
-        self.mem_r_bw_min_dict, self.mem_w_bw_min_dict = core.get_memory_bw_min_dict()
-        self.mem_sharing_tuple = tuple(tuple(i.items()) for i in core.mem_sharing_list)
+        self.mem_level_list = accelerator.memory_hierarchy.mem_level_list
+        self.mem_hierarchy_dict = accelerator.mem_hierarchy_dict
+        self.mem_size_dict = accelerator.mem_size_dict
+        self.mem_r_bw_dict, self.mem_w_bw_dict = accelerator.get_memory_bw_dict()
+        self.mem_r_bw_min_dict, self.mem_w_bw_min_dict = accelerator.get_memory_bw_min_dict()
+        self.mem_sharing_tuple = tuple(tuple(i.items()) for i in accelerator.mem_sharing_list)
         self.memory_operand_links = layer.memory_operand_links
 
         self.cumulative_layer_ids: list[int] = []  # In case the CME results from adding other CMEs together
-        self.cumulative_core_ids: list[int] = []
 
         # generate the integer spatial mapping from fractional spatial mapping (due to greedy mapping support).
         # Later the fractional one is used for calculating energy, and the integer one is used for calculating latency
@@ -547,8 +547,7 @@ def calc_energy(self) -> None:
 
     def calc_mac_energy_cost(self) -> None:
         """! Calculate the dynamic MAC energy"""
-        core = self.accelerator.get_core(self.core_id)
-        operational_array = core.operational_array
+        operational_array = self.accelerator.operational_array
         assert isinstance(
             operational_array, OperationalArray
         ), "This method expects an OperationalArray instance. Otherwise, the method should be overridden in a subclass."
@@ -563,8 +562,7 @@ def calc_memory_energy_cost(self):
 
         The energy total consumption is saved in self.energy_total.
         """
-        core = self.accelerator.get_core(self.core_id)
-        mem_hierarchy = core.memory_hierarchy
+        mem_hierarchy = self.accelerator.memory_hierarchy
 
         mem_energy_breakdown: dict[LayerOperand, list[float]] = {}
         mem_energy_breakdown_further: dict[LayerOperand, list[AccessEnergy]] = {}
@@ -1112,10 +1110,7 @@ def calc_overall_latency(self, cycles_per_mac: float = 1) -> None:
         """
         # the ideal cycle count assuming the MAC array is 100% utilized
         ideal_cycle = int(
-            ceil(
-                self.layer.total_mac_count / self.accelerator.get_core(self.core_id).operational_array.total_unit_count
-            )
-            * cycles_per_mac
+            ceil(self.layer.total_mac_count / self.accelerator.operational_array.total_unit_count) * cycles_per_mac
         )
 
         # the ideal temporal cycle count given the spatial mapping (the spatial mapping can be non-ideal)
@@ -1167,7 +1162,7 @@ def get_total_inst_bandwidth(self, memory_instance: MemoryInstance) -> MemoryAcc
         return total_inst_bw
 
     def __str__(self):
-        return f"CostModelEvaluation({self.layer}, core {self.core_id})"
+        return f"CostModelEvaluation({self.layer}, {self.accelerator})"
 
     def __repr__(self):
         return str(self)
diff --git a/zigzag/cost_model/cost_model_imc.py b/zigzag/cost_model/cost_model_imc.py
index 5f2931354..87dc4a1ec 100644
--- a/zigzag/cost_model/cost_model_imc.py
+++ b/zigzag/cost_model/cost_model_imc.py
@@ -41,9 +41,8 @@ def __init__(
         access_same_data_considered_as_no_access: bool = True,
     ):
         self.is_imc = True
-        self.core = next(iter(accelerator.cores))
-        assert isinstance(self.core.operational_array, ImcArray)
-        self.operational_array: ImcArray = self.core.operational_array
+        assert isinstance(accelerator.operational_array, ImcArray)
+        self.operational_array: ImcArray = accelerator.operational_array
         super().__init__(
             accelerator=accelerator,
             layer=layer,
diff --git a/zigzag/hardware/architecture/accelerator.py b/zigzag/hardware/architecture/accelerator.py
index 39fecbcec..b5a5734a2 100644
--- a/zigzag/hardware/architecture/accelerator.py
+++ b/zigzag/hardware/architecture/accelerator.py
@@ -1,31 +1,141 @@
-from zigzag.hardware.architecture.core import Core
+from zigzag.datatypes import MemoryOperand
+from zigzag.hardware.architecture.memory_hierarchy import MemoryHierarchy
+from zigzag.hardware.architecture.memory_instance import MemoryInstance
+from zigzag.hardware.architecture.memory_level import MemoryLevel
+from zigzag.hardware.architecture.operational_array import OperationalArrayABC
+from zigzag.mapping.spatial_mapping import SpatialMapping
 from zigzag.utils import json_repr_handler
 
 
 class Accelerator:
-    """!  The Accelerator class houses a set of Cores with an additional Global Buffer.
-    This Global Buffer sits above the cores, and can optionally be disabled.
+    """! This class houses the array of multipliers and the attached memory hierarchy.
+    This class supports a singular multiplier array and memory hierarchy, runtime flexibility should be implemented
+    on top.
     """
 
-    def __init__(self, name: str, core_set: set[Core]):
-        self.name: str = name
-        self.cores = sorted(list(core_set), key=lambda core: core.id)
+    id: int
+    name: str
+    operational_array: OperationalArrayABC
+    memory_hierarchy: MemoryHierarchy
+    dataflows: SpatialMapping | None
+    mem_hierarchy_dict: dict[MemoryOperand, list[MemoryLevel]]
+    mem_hierarchy_dict: dict[MemoryOperand, list[MemoryLevel]]
+    mem_size_dict: dict[MemoryOperand, list[int]]
+    mem_r_bw_dict: dict[MemoryOperand, list[int]]
+    mem_w_bw_dict: dict[MemoryOperand, list[int]]
+    mem_r_bw_min_dict: dict[MemoryOperand, list[int]]
+    mem_w_bw_min_dict: dict[MemoryOperand, list[int]]
+    mem_sharing_list: list[dict[MemoryOperand, int]]
+
+    def __init__(
+        self,
+        core_id: int,
+        name: str,
+        operational_array: OperationalArrayABC,
+        memory_hierarchy: MemoryHierarchy,
+        dataflows: SpatialMapping | None = None,
+    ):
+        self.id = core_id
+        self.name = name
+        self.operational_array = operational_array
+        self.memory_hierarchy = memory_hierarchy
+        self.mem_hierarchy_dict = {}
+
+        self.dataflows = dataflows
+        self.recalculate_memory_hierarchy_information()
+
+    def get_memory_level(self, mem_op: MemoryOperand, mem_lv: int) -> MemoryLevel:
+        """! Returns a specific memory level in the memory hierarchy for the memory operand"""
+        # Sort the nodes topologically and filter out all memories that don't store mem_op
+        memory = [node for node in self.memory_hierarchy.topological_sort() if mem_op in node.operands]
+        return memory[mem_lv]
+
+    def recalculate_memory_hierarchy_information(self) -> None:
+        self.__generate_memory_hierarchy_dict()
+        self.__generate_mem_sharing_list()
+
+    def __generate_memory_hierarchy_dict(self):
+        mem_operands = self.memory_hierarchy.nb_levels.keys()
+        self.mem_hierarchy_dict = {}
+        self.mem_size_dict = {}
+        self.mem_r_bw_dict = {}
+        self.mem_w_bw_dict = {}
+        self.mem_r_bw_min_dict = {}
+        self.mem_w_bw_min_dict = {}
+        for mem_op in mem_operands:
+            self.mem_hierarchy_dict[mem_op] = [
+                node for node in self.memory_hierarchy.topological_sort() if mem_op in node.operands
+            ]
+            self.mem_size_dict[mem_op] = [
+                node.memory_instance.size
+                for node in self.memory_hierarchy.topological_sort()
+                if mem_op in node.operands
+            ]
+            self.mem_r_bw_dict[mem_op] = [
+                node.memory_instance.r_bw
+                for node in self.memory_hierarchy.topological_sort()
+                if mem_op in node.operands
+            ]
+            self.mem_w_bw_dict[mem_op] = [
+                node.memory_instance.w_bw
+                for node in self.memory_hierarchy.topological_sort()
+                if mem_op in node.operands
+            ]
+            self.mem_r_bw_min_dict[mem_op] = [
+                node.memory_instance.r_bw_min
+                for node in self.memory_hierarchy.topological_sort()
+                if mem_op in node.operands
+            ]
+            self.mem_w_bw_min_dict[mem_op] = [
+                node.memory_instance.w_bw_min
+                for node in self.memory_hierarchy.topological_sort()
+                if mem_op in node.operands
+            ]
+
+    def __generate_mem_sharing_list(self):
+        """! Generates a list of dictionary that indicates which operand's which memory levels are sharing the same
+        physical memory"""
+        self.mem_sharing_list = []
+        for mem_lv in self.mem_hierarchy_dict.values():
+            for mem in mem_lv:
+                operand_mem_share = mem.mem_level_of_operands
+                if len(operand_mem_share) > 1 and operand_mem_share not in self.mem_sharing_list:
+                    self.mem_sharing_list.append(operand_mem_share)
+
+    def get_top_memory_instance(self, mem_op: MemoryOperand) -> MemoryInstance:
+        if mem_op not in self.memory_hierarchy.get_operands():
+            raise ValueError(f"Memory operand {mem_op} not in {self}.")
+        mem_level = self.memory_hierarchy.get_operand_top_level(mem_op)
+        mem_instance = mem_level.memory_instance
+        return mem_instance
+
+    def get_memory_bw_dict(self):
+        return self.mem_r_bw_dict, self.mem_w_bw_dict
+
+    def get_memory_bw_min_dict(self):
+        return self.mem_r_bw_min_dict, self.mem_w_bw_min_dict
 
     def __str__(self) -> str:
-        return f"Accelerator({self.name})"
+        return f"Core({self.id})"
 
     def __repr__(self) -> str:
         return str(self)
 
     def __jsonrepr__(self):
-        """! JSON representation used for saving this object to a json file."""
-        return json_repr_handler({"name": self.name, "cores": self.cores})
-
-    def get_core(self, core_id: int) -> Core:
-        """! 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.cores if core.id == core_id), None)
-        if not core:
-            raise ValueError(f"Requested core with id {core_id} is not present in accelerator.")
-        return core
+        return json_repr_handler(self.__dict__)
+
+    def __hash__(self) -> int:
+        return self.id
+
+    def __eq__(self, other: object) -> bool:
+        return (
+            isinstance(other, Accelerator)
+            and self.id == other.id
+            and self.operational_array == other.operational_array
+            and self.memory_hierarchy == other.memory_hierarchy
+        )
+
+    def has_same_performance(self, other: "Accelerator") -> bool:
+        return self.operational_array == other.operational_array and self.memory_hierarchy.has_same_performance(
+            other.memory_hierarchy
+        )
diff --git a/zigzag/hardware/architecture/core.py b/zigzag/hardware/architecture/core.py
deleted file mode 100644
index 7f574afba..000000000
--- a/zigzag/hardware/architecture/core.py
+++ /dev/null
@@ -1,140 +0,0 @@
-from zigzag.datatypes import MemoryOperand
-from zigzag.hardware.architecture.memory_hierarchy import MemoryHierarchy
-from zigzag.hardware.architecture.memory_instance import MemoryInstance
-from zigzag.hardware.architecture.memory_level import MemoryLevel
-from zigzag.hardware.architecture.operational_array import OperationalArrayABC
-from zigzag.mapping.spatial_mapping import SpatialMapping
-from zigzag.utils import json_repr_handler
-
-
-class Core:
-    """! The Core class houses the array of multipliers and the attached memory hierarchy.
-    This class supports a singular multiplier array and memory hierarchy, runtime flexibility should be implemented
-    on top.
-    """
-
-    id: int
-    operational_array: OperationalArrayABC
-    memory_hierarchy: MemoryHierarchy
-    dataflows: SpatialMapping | None
-    mem_hierarchy_dict: dict[MemoryOperand, list[MemoryLevel]]
-    mem_hierarchy_dict: dict[MemoryOperand, list[MemoryLevel]]
-    mem_size_dict: dict[MemoryOperand, list[int]]
-    mem_r_bw_dict: dict[MemoryOperand, list[int]]
-    mem_w_bw_dict: dict[MemoryOperand, list[int]]
-    mem_r_bw_min_dict: dict[MemoryOperand, list[int]]
-    mem_w_bw_min_dict: dict[MemoryOperand, list[int]]
-    mem_sharing_list: list[dict[MemoryOperand, int]]
-
-    def __init__(
-        self,
-        core_id: int,
-        operational_array: OperationalArrayABC,
-        memory_hierarchy: MemoryHierarchy,
-        dataflows: SpatialMapping | None = None,
-    ):
-        self.id = core_id
-        self.operational_array = operational_array
-        self.memory_hierarchy = memory_hierarchy
-        self.mem_hierarchy_dict = {}
-
-        self.dataflows = dataflows
-        self.recalculate_memory_hierarchy_information()
-
-    def get_memory_level(self, mem_op: MemoryOperand, mem_lv: int) -> MemoryLevel:
-        """! Returns a specific memory level in the memory hierarchy for the memory operand"""
-        # Sort the nodes topologically and filter out all memories that don't store mem_op
-        memory = [node for node in self.memory_hierarchy.topological_sort() if mem_op in node.operands]
-        return memory[mem_lv]
-
-    def recalculate_memory_hierarchy_information(self):
-        self.__generate_memory_hierarchy_dict()
-        self.__generate_mem_sharing_list()
-
-    def __generate_memory_hierarchy_dict(self):
-        mem_operands = self.memory_hierarchy.nb_levels.keys()
-        self.mem_hierarchy_dict = {}
-        self.mem_size_dict = {}
-        self.mem_r_bw_dict = {}
-        self.mem_w_bw_dict = {}
-        self.mem_r_bw_min_dict = {}
-        self.mem_w_bw_min_dict = {}
-        for mem_op in mem_operands:
-            self.mem_hierarchy_dict[mem_op] = [
-                node for node in self.memory_hierarchy.topological_sort() if mem_op in node.operands
-            ]
-            self.mem_size_dict[mem_op] = [
-                node.memory_instance.size
-                for node in self.memory_hierarchy.topological_sort()
-                if mem_op in node.operands
-            ]
-            self.mem_r_bw_dict[mem_op] = [
-                node.memory_instance.r_bw
-                for node in self.memory_hierarchy.topological_sort()
-                if mem_op in node.operands
-            ]
-            self.mem_w_bw_dict[mem_op] = [
-                node.memory_instance.w_bw
-                for node in self.memory_hierarchy.topological_sort()
-                if mem_op in node.operands
-            ]
-            self.mem_r_bw_min_dict[mem_op] = [
-                node.memory_instance.r_bw_min
-                for node in self.memory_hierarchy.topological_sort()
-                if mem_op in node.operands
-            ]
-            self.mem_w_bw_min_dict[mem_op] = [
-                node.memory_instance.w_bw_min
-                for node in self.memory_hierarchy.topological_sort()
-                if mem_op in node.operands
-            ]
-
-    def __generate_mem_sharing_list(self):
-        """! Generates a list of dictionary that indicates which operand's which memory levels are sharing the same
-        physical memory"""
-        self.mem_sharing_list = []
-        for mem_lv in self.mem_hierarchy_dict.values():
-            for mem in mem_lv:
-                operand_mem_share = mem.mem_level_of_operands
-                if len(operand_mem_share) > 1 and operand_mem_share not in self.mem_sharing_list:
-                    self.mem_sharing_list.append(operand_mem_share)
-
-    def get_top_memory_instance(self, mem_op: MemoryOperand) -> MemoryInstance:
-        if mem_op not in self.memory_hierarchy.get_operands():
-            raise ValueError(f"Memory operand {mem_op} not in {self}.")
-        mem_level = self.memory_hierarchy.get_operand_top_level(mem_op)
-        mem_instance = mem_level.memory_instance
-        return mem_instance
-
-    def get_memory_bw_dict(self):
-        return self.mem_r_bw_dict, self.mem_w_bw_dict
-
-    def get_memory_bw_min_dict(self):
-        return self.mem_r_bw_min_dict, self.mem_w_bw_min_dict
-
-    def __str__(self) -> str:
-        return f"Core({self.id})"
-
-    def __repr__(self) -> str:
-        return str(self)
-
-    def __jsonrepr__(self):
-        return json_repr_handler(self.__dict__)
-
-    def __hash__(self) -> int:
-        return self.id
-
-    def __eq__(self, other: object) -> bool:
-        return (
-            isinstance(other, Core)
-            and self.id == other.id
-            and self.operational_array == other.operational_array
-            and self.memory_hierarchy == other.memory_hierarchy
-        )
-
-    def has_same_performance(self, other: object) -> bool:
-        return (
-            isinstance(other, Core)
-            and self.operational_array == other.operational_array
-            and self.memory_hierarchy == other.memory_hierarchy
-        )
diff --git a/zigzag/hardware/architecture/memory_hierarchy.py b/zigzag/hardware/architecture/memory_hierarchy.py
index 4885a392a..157e3bc0b 100644
--- a/zigzag/hardware/architecture/memory_hierarchy.py
+++ b/zigzag/hardware/architecture/memory_hierarchy.py
@@ -50,7 +50,7 @@ def add_memory(
         operands: list[MemoryOperand],
         port_alloc: PortAllocation,
         served_dimensions: ServedMemDimensions,
-    ):
+    ) -> None:
         """! Adds a memory to the memory hierarchy graph.
         NOTE: memory level need to be added from bottom level (e.g., Reg) to top level (e.g., DRAM) for each operand !!!
 
@@ -165,3 +165,8 @@ def __eq__(self, other: object) -> bool:
             and self.nb_levels == other.nb_levels
             and all([self_ml == other_ml for self_ml, other_ml in zip(self.node_list, other.node_list)])
         )
+
+    def has_same_performance(self, other: "MemoryHierarchy"):
+        return self.nb_levels == other.nb_levels and all(
+            [self_ml.has_same_performance(other_ml) for self_ml, other_ml in zip(self.node_list, other.node_list)]
+        )
diff --git a/zigzag/hardware/architecture/memory_instance.py b/zigzag/hardware/architecture/memory_instance.py
index 5a9ed55df..447bd3e72 100644
--- a/zigzag/hardware/architecture/memory_instance.py
+++ b/zigzag/hardware/architecture/memory_instance.py
@@ -23,6 +23,24 @@ class MemoryInstance:
     double_buffering_support: bool
     shared_memory_group_id: int
 
+    name: str
+    size: int
+    r_bw: int
+    w_bw: int
+    r_cost: float
+    w_cost: float
+    area: float
+    r_port: int
+    w_port: int
+    rw_port: int
+    latency: int
+    min_r_granularity: int
+    min_w_granularity: int
+    mem_type: str
+    auto_cost_extraction: bool
+    double_buffering_support: bool
+    shared_memory_group_id: int
+
     def __init__(
         self,
         name: str,
@@ -105,8 +123,25 @@ def __jsonrepr__(self):
     def __eq__(self, other: object) -> bool:
         return isinstance(other, MemoryInstance) and self.__dict__ == other.__dict__
 
+    def has_same_performance(self, other: "MemoryInstance"):
+        """Wether using this instance will result in the same estimations as using the other instance. This method
+        differs from __eq__ since it does not consider e.g. the shared_memory_group_id"""
+        return (
+            self.size == other.size
+            and self.r_bw == other.r_bw
+            and self.w_bw == other.w_bw
+            and self.r_cost == other.r_cost
+            and self.w_cost == other.w_cost
+            and self.r_port_nb == other.r_port_nb
+            and self.w_port_nb == other.w_port_nb
+            and self.rw_port_nb == other.rw_port_nb
+            and self.latency == other.latency
+            and self.double_buffering_support == other.double_buffering_support
+            and self.r_bw_min == other.r_bw_min
+            and self.w_bw_min == other.w_bw_min
+        )
+
     def __hash__(self):
-        # id(self)  # unique for every object within its lifetime
         return hash(frozenset(self.__dict__.values()))
 
     def __str__(self):
diff --git a/zigzag/hardware/architecture/memory_level.py b/zigzag/hardware/architecture/memory_level.py
index 8f6f68932..124f156f9 100644
--- a/zigzag/hardware/architecture/memory_level.py
+++ b/zigzag/hardware/architecture/memory_level.py
@@ -177,5 +177,14 @@ def __eq__(self, other: Any) -> bool:
             and self.served_dimensions == other.served_dimensions
         )
 
+    def has_same_performance(self, other: "MemoryLevel"):
+        return (
+            self.memory_instance.has_same_performance(other.memory_instance)
+            and self.operands == other.operands
+            and self.mem_level_of_operands == other.mem_level_of_operands
+            and self.port_list == other.port_list
+            and self.served_dimensions == other.served_dimensions
+        )
+
     def __hash__(self) -> int:
         return hash_sha512(self.id)
diff --git a/zigzag/mapping/mapping.py b/zigzag/mapping/mapping.py
index 269b1d686..53b7b077a 100644
--- a/zigzag/mapping/mapping.py
+++ b/zigzag/mapping/mapping.py
@@ -369,9 +369,7 @@ def calc_data_access(self):
                 if (
                     self.access_same_data_considered_as_no_access
                     and mem_level == 0
-                    and self.accelerator.get_core(self.layer_node.core_allocation[0]).mem_r_bw_dict[
-                        self.layer_node.memory_operand_links[operand]
-                    ][mem_level]
+                    and self.accelerator.mem_r_bw_dict[self.layer_node.memory_operand_links[operand]][mem_level]
                     >= self.data_bit_per_level[operand][mem_level]
                     // self.spatial_mapping.unit_unique[operand][mem_level + 1]
                 ):
diff --git a/zigzag/opt/loma/engine.py b/zigzag/opt/loma/engine.py
index 7939370e2..838b9d85d 100644
--- a/zigzag/opt/loma/engine.py
+++ b/zigzag/opt/loma/engine.py
@@ -74,8 +74,7 @@ def __init__(
         # TODO: Take into account that data might be stored in lower level,
         # TODO: thus adapt the memory hierarchy.
         # TODO: The fact that there is a global buffer above the cores requires attention.
-        core_id = layer.core_allocation[0]
-        self.memory_hierarchy = accelerator.get_core(core_id).memory_hierarchy
+        self.memory_hierarchy = accelerator.memory_hierarchy
 
         self.show_progress_bar = kwargs.get("loma_show_progress_bar", False)
 
diff --git a/zigzag/opt/loma/memory_allocator.py b/zigzag/opt/loma/memory_allocator.py
index 1bbf3af10..5db343893 100644
--- a/zigzag/opt/loma/memory_allocator.py
+++ b/zigzag/opt/loma/memory_allocator.py
@@ -82,8 +82,7 @@ def run(self):
         """
 
         # self.nodes contains the different memory nodes in bottom-up fashion
-        core_id = self.layer.core_allocation[0]
-        memory_hierarchy = self.accelerator.get_core(core_id).memory_hierarchy
+        memory_hierarchy = self.accelerator.memory_hierarchy
         top_levels = {mem_op: memory_hierarchy.get_operand_top_level(mem_op) for mem_op in self.mem_ops}
         for node in memory_hierarchy.topological_sort():
             self.allocate_node(node, top_levels)
diff --git a/zigzag/parser/accelerator_factory.py b/zigzag/parser/accelerator_factory.py
index 656979fd0..ec40aeef7 100644
--- a/zigzag/parser/accelerator_factory.py
+++ b/zigzag/parser/accelerator_factory.py
@@ -8,7 +8,6 @@
     UnrollFactor,
 )
 from zigzag.hardware.architecture.accelerator import Accelerator
-from zigzag.hardware.architecture.core import Core
 from zigzag.hardware.architecture.imc_array import ImcArray
 from zigzag.hardware.architecture.memory_hierarchy import MemoryHierarchy
 from zigzag.hardware.architecture.memory_instance import MemoryInstance
@@ -21,34 +20,34 @@
 from zigzag.hardware.architecture.operational_unit import Multiplier
 from zigzag.mapping.spatial_mapping import MappingSingleOADim, SpatialMapping
 
+# class AcceleratorFactory:
+#     """! Converts valid user-provided accelerator data into an `Accelerator` instance"""
 
-class AcceleratorFactory:
-    """! Converts valid user-provided accelerator data into an `Accelerator` instance"""
+#     def __init__(self, data: dict[str, Any]):
+#         """! Generate an `Accelerator` instance from the validated user-provided data."""
+#         self.data = data
 
-    def __init__(self, data: dict[str, Any]):
-        """! Generate an `Accelerator` instance from the validated user-provided data."""
-        self.data = data
-
-    def create(self) -> Accelerator:
-        """! Create an Accelerator instance from the user-provided data.
-        NOTE the memory instances must be defined from lowest to highest.
-        """
-        core_factory = CoreFactory(self.data)
-        core = core_factory.create()
-        return Accelerator(name=self.data["name"], core_set={core})
+#     def create(self) -> Accelerator:
+#         """! Create an Accelerator instance from the user-provided data.
+#         NOTE the memory instances must be defined from lowest to highest.
+#         """
+#         core_factory = CoreFactory(self.data)
+#         core = core_factory.create()
+#         return Accelerator(name=self.data["name"], core_set={core})
 
 
-class CoreFactory:
-    """! Converts valid user-provided accelerator data into a `Core` instance"""
+class AcceleratorFactory:
+    """! Converts valid user-provided accelerator data into an `Accelerator` instance"""
 
     def __init__(self, data: dict[str, Any]):
         """! Generate an `Core` instance from the validated user-provided data."""
         self.data = data
 
-    def create(self, core_id: int = 1, shared_mem_group_id: int | None = None) -> Core:
+    def create(self, core_id: int = 0, shared_mem_group_id: int | None = None) -> Accelerator:
         """! Create a Core instance from the user-provided data.
         NOTE the memory instances must be defined from lowest to highest.
         """
+        name = self.data["name"]
         operational_array = self.create_operational_array()
 
         mem_graph = MemoryHierarchy(operational_array)
@@ -61,8 +60,9 @@ def create(self, core_id: int = 1, shared_mem_group_id: int | None = None) -> Co
             )
             memory_factory.add_memory_to_graph(mem_graph)
 
-        return Core(
+        return Accelerator(
             core_id=core_id,
+            name=name,
             operational_array=operational_array,
             memory_hierarchy=mem_graph,
             dataflows=dataflows,
diff --git a/zigzag/stages/evaluation/cost_model_evaluation.py b/zigzag/stages/evaluation/cost_model_evaluation.py
index 79b9b0f6a..46df24457 100644
--- a/zigzag/stages/evaluation/cost_model_evaluation.py
+++ b/zigzag/stages/evaluation/cost_model_evaluation.py
@@ -39,9 +39,7 @@ def __init__(
 
     def run(self):
         """! Run the cost model stage by calling the internal zigzag cost model with the correct inputs."""
-        core_id = self.layer.core_allocation[0]
-        core = self.accelerator.get_core(core_id)
-        operational_array = core.operational_array
+        operational_array = self.accelerator.operational_array
         if isinstance(operational_array, ImcArray):
             cme = CostModelEvaluationForIMC(
                 accelerator=self.accelerator,
diff --git a/zigzag/stages/exploit_data_locality_stages.py b/zigzag/stages/exploit_data_locality_stages.py
index 53f1fabe8..2adf4ac9f 100644
--- a/zigzag/stages/exploit_data_locality_stages.py
+++ b/zigzag/stages/exploit_data_locality_stages.py
@@ -9,7 +9,6 @@
 
 from zigzag.datatypes import MemoryOperand
 from zigzag.hardware.architecture.accelerator import Accelerator
-from zigzag.hardware.architecture.core import Core
 from zigzag.hardware.architecture.memory_hierarchy import MemoryHierarchy
 from zigzag.hardware.architecture.memory_instance import MemoryInstance
 from zigzag.hardware.architecture.memory_level import MemoryLevel, ServedMemDimensions
@@ -36,13 +35,11 @@ def __init__(
         super().__init__(list_of_callables, **kwargs)
         self.accelerator = accelerator
         self.workload = workload
-        core_id = accelerator.cores[0].id  # correct only for single-core hardware
-        core = accelerator.get_core(core_id)
         # Remove dummy (non-conv) layers (ReLU, Pooling..) in the layer graph
         self.workload_no_dummy = self.workload.get_copy_no_dummy()
 
         # record of all memory levels
-        self.core_mem_level_list = core.memory_hierarchy.mem_level_list
+        self.core_mem_level_list = accelerator.memory_hierarchy.mem_level_list
         # record of top mem level per layer: dict[layer_idx, mem_level]
         self.mem_update_list: dict[int, dict[MemoryOperand, int]] = {}
         # record of input, output size per layer (unit: bit): dict[layer_idx: [operand: size]]
@@ -228,8 +225,7 @@ def update_top_mem_level(self):
 
     def check_if_mem_serve_all_oa_dims(self, mem: MemoryLevel, accelerator: Accelerator):
         """! Function to check if mem serve all hardware dimensions"""
-        core = accelerator.cores[0]
-        operational_array = core.operational_array
+        operational_array = accelerator.operational_array
         oa_dim_nb = len(operational_array.dimension_sizes)
         mem_served_oa_dim_nb = len(mem.served_dimensions)
         return mem_served_oa_dim_nb == oa_dim_nb
@@ -303,9 +299,8 @@ def run(self):
     def generate_accelerator_with_removing_unused_memory(self) -> Accelerator:
         # Remove no-use memory level according to update_mem_list and mem_update_weight
         curr_id = self.layer.id
-        core = next(iter(self.accelerator.cores))
-        operational_array = core.operational_array
-        memory_hierarchy = core.memory_hierarchy
+        operational_array = self.accelerator.operational_array
+        memory_hierarchy = self.accelerator.memory_hierarchy
 
         act_layer_op = self.layer.get_act_layer_op()
         weight_layer_op = self.layer.get_weight_layer_op()
@@ -373,20 +368,13 @@ def generate_accelerator_with_removing_unused_memory(self) -> Accelerator:
                     served_dimensions=new_served_dimensions,
                 )
 
-        # Create the new core
-        new_core = Core(
-            core_id=core.id,
-            operational_array=operational_array,
-            memory_hierarchy=new_memory_hierarchy,
-        )
-
         # Create the new accelerator
-        name = self.accelerator.name
-        new_name = name + "-removing-nouse-mem"
-        new_cores = {new_core}
+        new_name = self.accelerator.name + "-removing-nouse-mem"
         new_accelerator = Accelerator(
+            core_id=self.accelerator.id,
             name=new_name,
-            core_set=new_cores,
+            operational_array=operational_array,
+            memory_hierarchy=new_memory_hierarchy,
         )
 
         logger.info("Update mem architecture for layer %s...", self.layer)
diff --git a/zigzag/stages/mapping/spatial_mapping_conversion.py b/zigzag/stages/mapping/spatial_mapping_conversion.py
index d5205b7d4..67afcc04e 100644
--- a/zigzag/stages/mapping/spatial_mapping_conversion.py
+++ b/zigzag/stages/mapping/spatial_mapping_conversion.py
@@ -195,8 +195,7 @@ def generate_mapping_per_mem_lvl(self, user_spatial_mapping: SpatialMapping) ->
         # TODO This should be a class
         """
         mapping_per_mem_lvl: SpatialMappingPerMemLvl = {}
-        core_id = self.layer.core_allocation[0]
-        mem_hierarchy = self.accelerator.get_core(core_id).memory_hierarchy
+        mem_hierarchy = self.accelerator.memory_hierarchy
         for layer_op in self.memory_operand_links.layer_operands:
             mem_op = self.memory_operand_links.layer_to_mem_op(layer_op)
             usm_copy = user_spatial_mapping.copy()
diff --git a/zigzag/stages/mapping/spatial_mapping_generation.py b/zigzag/stages/mapping/spatial_mapping_generation.py
index 0c3d0a3c1..fbfd6d6ae 100644
--- a/zigzag/stages/mapping/spatial_mapping_generation.py
+++ b/zigzag/stages/mapping/spatial_mapping_generation.py
@@ -14,7 +14,6 @@
     UnrollFactorInt,
 )
 from zigzag.hardware.architecture.accelerator import Accelerator
-from zigzag.hardware.architecture.core import Core
 from zigzag.hardware.architecture.memory_hierarchy import MemoryHierarchy
 from zigzag.hardware.architecture.memory_instance import MemoryInstance
 from zigzag.hardware.architecture.memory_level import ServedMemDimensions
@@ -73,10 +72,8 @@ def __init__(
         self.nb_mappings_generated = nb_mappings_generated
 
         self.layer_dim_sizes = self.layer.layer_dim_sizes
-        core_id = layer.core_allocation[0]
-        self.core = self.accelerator.get_core(core_id)
-        self.oa_dim_sizes = self.core.operational_array.dimension_sizes
-        self.memory_hierarchy = self.core.memory_hierarchy
+        self.oa_dim_sizes = self.accelerator.operational_array.dimension_sizes
+        self.memory_hierarchy = self.accelerator.memory_hierarchy
 
         # Spatial mapping hint
         self.spatial_mapping_hint = self.layer.spatial_mapping_hint
@@ -701,7 +698,7 @@ def modify_innermost_input_mem_size(self, user_spatial_mapping: SpatialMapping)
         # Initialize the new memory hierarchy
         mh_name = self.memory_hierarchy.name
         new_mh_name = mh_name + "-supporting-diagonal-map"
-        operational_array = self.core.operational_array
+        operational_array = self.accelerator.operational_array
         new_memory_hierarchy = MemoryHierarchy(operational_array, new_mh_name)
         # Add memories to the new memory hierarchy with the correct attributes
         for memory_level in self.memory_hierarchy.mem_level_list:
@@ -729,18 +726,12 @@ def modify_innermost_input_mem_size(self, user_spatial_mapping: SpatialMapping)
                 served_dimensions=new_served_dimensions,
             )
         # Create the new core
-        new_core = Core(
-            core_id=self.core.id,
+        new_name = self.accelerator.name + "-supporting-diagonal-map"
+        new_accelerator = Accelerator(
+            core_id=self.accelerator.id,
+            name=new_name,
             operational_array=operational_array,
             memory_hierarchy=new_memory_hierarchy,
         )
 
-        # Create the new accelerator
-        name = self.accelerator.name
-        new_name = name + "-supporting-diagonal-map"
-        new_cores = {new_core}
-        new_accelerator = Accelerator(
-            name=new_name,
-            core_set=new_cores,
-        )
         return new_accelerator
diff --git a/zigzag/stages/results/visualization.py b/zigzag/stages/results/visualization.py
index cbf1c4777..63f9a35cd 100644
--- a/zigzag/stages/results/visualization.py
+++ b/zigzag/stages/results/visualization.py
@@ -48,7 +48,7 @@ def __save_loop_ordering(self, cme: CostModelEvaluation):
 
     def __save_mem_hierarchy(self, cme: CostModelEvaluation):
         visualize_memory_hierarchy_graph(
-            cme.accelerator.cores[0].memory_hierarchy,
+            cme.accelerator.memory_hierarchy,
             save_path=self.dump_folder + "/mem_hierarchy.png",
         )
         self.figure_is_saved = True
diff --git a/zigzag/stages/stage.py b/zigzag/stages/stage.py
index d04a12329..8d576fad5 100644
--- a/zigzag/stages/stage.py
+++ b/zigzag/stages/stage.py
@@ -7,6 +7,8 @@
 class Stage(metaclass=ABCMeta):
     """! Abstract superclass for Runnables"""
 
+    kwargs: dict[str, Any]
+
     def __init__(
         self,
         list_of_callables: list["StageCallable"],
@@ -24,8 +26,8 @@ def __init__(
 
         if list_of_callables in ([], tuple(), set(), None) and not self.is_leaf():
             raise ValueError(
-                "List of callables empty on a non leaf runnable, so nothing can be generated.\
-                    Final callable in list_of_callables must return Stage instances that have is_leaf() == True"
+                "List of callables empty on a non leaf runnable, so nothing can be generated. "
+                "Final callable in list_of_callables must return Stage instances that have is_leaf() == True"
             )
 
     @abstractmethod
diff --git a/zigzag/stages/workload_iterator.py b/zigzag/stages/workload_iterator.py
index 2dfeb5e96..1ee334dc8 100644
--- a/zigzag/stages/workload_iterator.py
+++ b/zigzag/stages/workload_iterator.py
@@ -35,9 +35,7 @@ def run(self):
                 continue
             # Skip Pooling, Add layers for imc. This happens only when the workload is manually defined.
             # No skipping if the workload is from onnx.
-            core_id = layer.core_allocation[0]
-            core = self.accelerator.get_core(core_id)
-            operational_array = core.operational_array
+            operational_array = self.accelerator.operational_array
             if isinstance(operational_array, ImcArray) and layer.type in [
                 "Pooling",
                 "Add",
diff --git a/zigzag/visualization/graph/memory_hierarchy.py b/zigzag/visualization/graph/memory_hierarchy.py
index 7da5f40c1..bf0c0adbc 100644
--- a/zigzag/visualization/graph/memory_hierarchy.py
+++ b/zigzag/visualization/graph/memory_hierarchy.py
@@ -52,4 +52,4 @@ def visualize_memory_hierarchy_graph(graph: MemoryHierarchy, save_path: str = ""
     with open("../list_of_cmes.pickle", "rb") as handle:
         list_of_cme = pickle.load(handle)
     cme = list_of_cme[0]
-    visualize_memory_hierarchy_graph(cme.accelerator.cores[0].memory_hierarchy)
+    visualize_memory_hierarchy_graph(cme.accelerator.memory_hierarchy)
diff --git a/zigzag/visualization/results/plot_cme.py b/zigzag/visualization/results/plot_cme.py
index aad84cddb..92620e24f 100644
--- a/zigzag/visualization/results/plot_cme.py
+++ b/zigzag/visualization/results/plot_cme.py
@@ -51,7 +51,7 @@ def get_energy_array(
     """Convert the given list of cmes to a numpy array with the energy per layer, memory level, operand and data
     direction. Output shape: (cmes, all_mems, all_ops, data_directions)"""
 
-    mem_hierarchy = cmes[0].accelerator.get_core(cmes[0].layer.core_allocation[0]).memory_hierarchy
+    mem_hierarchy = cmes[0].accelerator.memory_hierarchy
     access_energy: dict[int, defaultdict[LayerOperand, defaultdict[MemoryInstance, AccessEnergy]]] = {
         idx: defaultdict(lambda: defaultdict(lambda: AccessEnergy(0, 0, 0, 0))) for idx in range(len(cmes))
     }
@@ -100,7 +100,7 @@ def get_latency_array(cmes: list[CostModelEvaluation]):
 def bar_plot_cost_model_evaluations_breakdown(cmes: list[CostModelEvaluationABC], save_path: str):
     # Input-specific
     cmes_to_plot: list[CostModelEvaluation] = [cme for cme in cmes if isinstance(cme, CostModelEvaluation)]
-    mem_hierarchy = cmes_to_plot[0].accelerator.get_core(cmes_to_plot[0].layer.core_allocation[0]).memory_hierarchy
+    mem_hierarchy = cmes_to_plot[0].accelerator.memory_hierarchy
     memories = [mem_level.memory_instance for mem_level in mem_hierarchy.mem_level_list]
     all_mems = sorted(memories, key=lambda x: x.size)
     all_ops = list({layer_op for cme in cmes_to_plot for layer_op in cme.layer.layer_operands})
diff --git a/zigzag/visualization/results/print_mapping.py b/zigzag/visualization/results/print_mapping.py
index f3c1971f3..3c29c8baf 100644
--- a/zigzag/visualization/results/print_mapping.py
+++ b/zigzag/visualization/results/print_mapping.py
@@ -17,7 +17,6 @@ def get_temporal_spatial_loops(
     """
     # TODO documentation, split this up into multiple, sensible functions
     """
-    core = cme.accelerator.get_core(cme.layer.core_allocation[0])
     operand_links = cme.layer.memory_operand_links
 
     tm: dict[LayerOperand, list[list[tuple[LayerDim, UnrollFactor]]]] = pickle_deepcopy(
@@ -37,7 +36,7 @@ def get_temporal_spatial_loops(
             idx = tm[layer_op][level].index(tl)
             tm[layer_op][level].pop(idx)
             # Get the name of this memory level
-            mem_name = core.get_memory_level(mem_op, level).memory_instance.name
+            mem_name = cme.accelerator.get_memory_level(mem_op, level).memory_instance.name
             mem_names.append(mem_name)
             all_mem_names.add(mem_name)
         mem_names_tuple = tuple(mem_names)
diff --git a/zigzag/workload/layer_node_abc.py b/zigzag/workload/layer_node_abc.py
index da81d6ff1..e29168fa2 100644
--- a/zigzag/workload/layer_node_abc.py
+++ b/zigzag/workload/layer_node_abc.py
@@ -6,6 +6,9 @@
 class LayerNodeABC(metaclass=ABCMeta):
     """Represents a single layer of a workload in any form."""
 
+    id: int
+    name: str
+
     def __init__(self, node_id: int, node_name: str):
         self.id = node_id
         self.name = node_name