Incorporate zigzag filename changes

.gitignore

@@ -1,4 +1,5 @@
 *.out
+typings
 
 # Byte-compiled / optimized / DLL files
 __pycache__/

.pre-commit-config.yaml

@@ -10,7 +10,7 @@ repos:
         types_or: [python, pyi]
         args: [--extend-select, I, --fix, --exit-non-zero-on-fix, --line-length, "120"]
   - repo: https://github.com/psf/black
-    rev: 23.3.0
+    rev: 24.8.0
     hooks:
       - id: black
         args: [--line-length, "120"]

main_stream.py

@@ -2,7 +2,7 @@
 import pickle
 import re
 
-from zigzag.stages.MainStage import MainStage
+from zigzag.stages.main import MainStage
 
 from stream.classes.stages.AcceleratorParserStage import (
     AcceleratorParserStage as AcceleratorParserStage_,

main_stream_layer_splitting.py

@@ -4,7 +4,7 @@
 import pickle
 import re
 
-from zigzag.stages.MainStage import MainStage
+from zigzag.stages.main import MainStage
 
 from stream.classes.stages.AcceleratorParserStage import AcceleratorParserStage as AcceleratorParserStage_
 from stream.classes.stages.GenerateCNWorkloadHybridStage import GenerateCNWorkloadHybridStage

main_stream_mode_3.py

@@ -2,7 +2,7 @@
 import pickle
 import re
 
-from zigzag.stages.MainStage import MainStage
+from zigzag.stages.main import MainStage
 
 from stream.classes.stages.AcceleratorParserStage import AcceleratorParserStage as AcceleratorParserStage_
 from stream.classes.stages.GenerateCNWorkloadHybridStage import GenerateCNWorkloadHybridStage

main_stream_mode_4.py

@@ -2,7 +2,7 @@
 import pickle
 import re
 
-from zigzag.stages.MainStage import MainStage
+from zigzag.stages.main import MainStage
 
 from stream.classes.stages.AcceleratorParserStage import AcceleratorParserStage as AcceleratorParserStage_
 from stream.classes.stages.GenerateCNWorkloadHybridStage import GenerateCNWorkloadHybridStage

main_testing_1_core_with_testing_workload.py

@@ -1,7 +1,7 @@
 import logging as _logging
 import re
 
-from zigzag.stages.MainStage import MainStage
+from zigzag.stages.main import MainStage
 
 from stream.classes.cost_model.cost_model import StreamCostModelEvaluation
 from stream.classes.stages.AcceleratorParserStage import AcceleratorParserStage as AcceleratorParserStage_

main_testing_2_cores_shared.py

@@ -1,7 +1,7 @@
 import logging as _logging
 import re
 
-from zigzag.stages.MainStage import MainStage
+from zigzag.stages.main import MainStage
 
 from stream.classes.stages.AcceleratorParserStage import AcceleratorParserStage as AcceleratorParserStage_
 from stream.classes.stages.GenerateCNWorkloadHybridStage import GenerateCNWorkloadHybridStage

main_testing_2_cores_with_testing_workload.py

@@ -1,7 +1,7 @@
 import logging as _logging
 import re
 
-from zigzag.stages.MainStage import MainStage
+from zigzag.stages.main import MainStage
 
 from stream.classes.cost_model.cost_model import StreamCostModelEvaluation
 from stream.classes.stages.AcceleratorParserStage import AcceleratorParserStage as AcceleratorParserStage_

main_testing_2_cores_with_testing_workload_3_layers.py

@@ -1,7 +1,7 @@
 import logging as _logging
 import re
 
-from zigzag.stages.MainStage import MainStage
+from zigzag.stages.main import MainStage
 
 from stream.classes.cost_model.cost_model import StreamCostModelEvaluation
 from stream.classes.stages.AcceleratorParserStage import AcceleratorParserStage as AcceleratorParserStage_

main_testing_4_cores_with_testing_workload.py

@@ -1,7 +1,7 @@
 import logging as _logging
 import re
 
-from zigzag.stages.MainStage import MainStage
+from zigzag.stages.main import MainStage
 
 from stream.classes.cost_model.cost_model import StreamCostModelEvaluation
 from stream.classes.stages.AcceleratorParserStage import AcceleratorParserStage as AcceleratorParserStage_

stream/api.py

@@ -1,4 +1,4 @@
-from zigzag.stages.MainStage import MainStage
+from zigzag.stages.main import MainStage
 
 from stream.classes.stages.AcceleratorParserStage import AcceleratorParserStage
 from stream.classes.stages.GenerateCNWorkloadHybridStage import GenerateCNWorkloadHybridStage

stream/classes/cost_model/communication_manager.py

@@ -3,7 +3,7 @@
 from typing import TYPE_CHECKING, Any
 
 from zigzag.datatypes import Constants, MemoryOperand
-from zigzag.hardware.architecture.Core import Core
+from zigzag.hardware.architecture.core import Core
 
 from stream.classes.hardware.architecture.utils import intersections
 from stream.classes.workload.computation_node import ComputationNode

stream/classes/cost_model/memory_manager.py

@@ -4,9 +4,9 @@
 
 import numpy as np
 from zigzag.datatypes import MemoryOperand
-from zigzag.hardware.architecture.Core import Core
+from zigzag.hardware.architecture.core import Core
+from zigzag.hardware.architecture.memory_instance import MemoryInstance
 from zigzag.hardware.architecture.memory_level import MemoryLevel
-from zigzag.hardware.architecture.MemoryInstance import MemoryInstance
 
 from stream.classes.workload.tensor import Tensor
 

stream/classes/cost_model/scheduler.py

@@ -3,7 +3,7 @@
 from typing import TYPE_CHECKING
 
 from zigzag.datatypes import Constants, LayerOperand, MemoryOperand
-from zigzag.hardware.architecture.Core import Core
+from zigzag.hardware.architecture.core import Core
 
 from stream.classes.workload.computation_node import ComputationNode
 from stream.classes.workload.onnx_workload import ComputationNodeWorkload

stream/classes/hardware/architecture/accelerator.py

@@ -1,8 +1,8 @@
 from math import ceil
 
 from zigzag.datatypes import MemoryOperand
-from zigzag.hardware.architecture.Core import Core
-from zigzag.hardware.architecture.MemoryInstance import MemoryInstance
+from zigzag.hardware.architecture.core import Core
+from zigzag.hardware.architecture.memory_instance import MemoryInstance
 from zigzag.mapping.spatial_mapping import SpatialMapping
 
 from stream.classes.cost_model.communication_manager import CommunicationManager

stream/classes/hardware/architecture/noc/bus.py

@@ -1,11 +1,11 @@
 from zigzag.datatypes import Constants
-from zigzag.hardware.architecture.Core import Core
+from zigzag.hardware.architecture.core import Core
 
 from stream.classes.hardware.architecture.accelerator import CoreGraph
 from stream.classes.hardware.architecture.noc.communication_link import CommunicationLink
 
 
-def have_shared_memory(a, b):
+def have_shared_memory(a: Core, b: Core):
     """Returns True if core a and core b have a shared top level memory
 
     Args:
@@ -25,12 +25,12 @@ def have_shared_memory(a, b):
 
 
 def get_bus(
-    cores,
-    bandwidth,
-    unit_energy_cost,
-    pooling_core=None,
-    simd_core=None,
-    offchip_core=None,
+    cores: list[Core],
+    bandwidth: int,
+    unit_energy_cost: float,
+    pooling_core: Core | None = None,
+    simd_core: Core | None = None,
+    offchip_core: Core | None = None,
 ):
     """Return a graph of the cores where each core is connected to a single bus.
 
@@ -46,7 +46,7 @@ def get_bus(
     """
     bus = CommunicationLink("Any", "Any", bandwidth, unit_energy_cost)
 
-    edges = []
+    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

stream/classes/hardware/architecture/noc/communication_link.py

@@ -5,7 +5,7 @@
 from stream.classes.cost_model.communication_manager import CommunicationLinkEvent
 
 if TYPE_CHECKING:
-    from zigzag.hardware.architecture.Core import Core
+    from zigzag.hardware.architecture.core import Core
 
     from stream.classes.workload.tensor import Tensor
 

stream/classes/hardware/architecture/noc/mesh_2d.py

@@ -1,6 +1,6 @@
 import numpy as np
 from zigzag.datatypes import Constants
-from zigzag.hardware.architecture.Core import Core
+from zigzag.hardware.architecture.core import Core
 
 from stream.classes.hardware.architecture.accelerator import CoreGraph
 from stream.classes.hardware.architecture.noc.communication_link import CommunicationLink
@@ -35,7 +35,7 @@ def get_2d_mesh(
     simd_core: Core | None = None,
     offchip_core: Core | None = None,
 ):
-    """Return a 2D mesh graph of the cores where each core is connected to its N, E, S, W neighbour.
+    """Return a 2D mesh graph of the cores where each core is connected to its N, E, S, W neighbor.
     We build the mesh by iterating through the row and then moving to the next column.
     Each connection between two cores includes two links, one in each direction, each with specified bandwidth.
     Thus there are a total of ((nb_cols-1)*2*nb_rows + (nb_rows-1)*2*nb_cols) links in the noc.

stream/classes/io/accelerator_factory.py

@@ -1,6 +1,6 @@
 from typing import Any
 
-from zigzag.hardware.architecture.Core import Core
+from zigzag.hardware.architecture.core import Core
 from zigzag.parser.accelerator_factory import CoreFactory
 
 from stream.classes.hardware.architecture.accelerator import Accelerator
@@ -23,12 +23,13 @@ def create(self) -> Accelerator:
             core = core_factory.create(core_id)
             cores.append(core)
 
-        if self.data["graph"]["type"] == "2d_mesh":
-            cores_graph = self.create_2d_mesh(cores)
-        elif self.data["graph"]["type"] == "bus":
-            cores_graph = self.create_bus(cores)
-        else:
-            raise ValueError(f"Invalid graph type {self.data['graph']['type']}.")
+        match self.data["graph"]["type"]:
+            case "2d_mesh":
+                cores_graph = self.create_2d_mesh(cores)
+            case "bus":
+                cores_graph = self.create_bus(cores)
+            case _:
+                raise ValueError(f"Invalid graph type {self.data['graph']['type']}.")
 
         offchip_core_id: int | None = self.data["graph"]["offchip_core_id"]
         return Accelerator(name=self.data["name"], cores=cores_graph, offchip_core_id=offchip_core_id)

stream/classes/io/accelerator_validator.py

@@ -3,14 +3,15 @@
 from typing import Any
 
 from cerberus import Validator
-from zigzag.parser.AcceleratorValidator import AcceleratorValidator as CoreValidator
+from zigzag.parser.accelerator_validator import AcceleratorValidator as CoreValidator
 from zigzag.utils import open_yaml
 
 logger = logging.getLogger(__name__)
 
 
 class AcceleratorValidator:
     INPUT_DIR_LOCATION = "stream/inputs/"
+    GRAPH_TYPES = ["2d_mesh", "bus"]
 
     SCHEMA = {
         "name": {"type": "string", "required": True},
@@ -26,7 +27,7 @@ class AcceleratorValidator:
             "type": "dict",
             "required": True,
             "schema": {
-                "type": {"type": "string", "required": True},
+                "type": {"type": "string", "required": True, "allowed": GRAPH_TYPES},
                 "nb_rows": {"type": "integer", "required": False},
                 "nb_cols": {"type": "integer", "required": False},
                 "bandwidth": {"type": "integer", "required": True},

stream/classes/io/onnx/concat.py

@@ -1,20 +1,16 @@
-from zigzag.parser.onnx.ONNXOperatorParser import ONNXOperatorParser
 from zigzag.parser.onnx.utils import OnnxTensorCategory, get_onnx_tensor_type
 
+from stream.classes.io.onnx.operator_parser import OnnxOperatorParser
 from stream.classes.workload.concat_node import ConcatNode
 
 
-class ConcatParser(ONNXOperatorParser):
+class ConcatParser(OnnxOperatorParser):
     """Parses an onnx gather operator into a ConcatNode."""
 
-    def run(self):
-        return self.generate_node()
-
     def generate_node(self):
         predecessors = self.get_node_predecessors()
 
         axis = self.get_axis_value()
-        output_names = [self.node.output[0]]
 
         input_1, input_2 = self.node.input[0], self.node.input[1]
 
@@ -25,15 +21,13 @@ def generate_node(self):
 
             constant_shape = tuple(constant_tensor.shape)
             variable_input_first = True
-            input_names = [input_2]
         except ValueError:  # Try second one as constant input
             constant_tensor = get_onnx_tensor_type(input_2, self.onnx_model)
             if constant_tensor.category != OnnxTensorCategory.HIDDEN or "constant" not in input_2.lower():
                 raise ValueError
 
             constant_shape = tuple(constant_tensor.shape)
             variable_input_first = True
-            input_names = [input_1]
 
         return ConcatNode(
             node_id=self.node_id,
@@ -42,8 +36,6 @@ def generate_node(self):
             axis=axis,
             constant_shape=constant_shape,
             variable_input_first=variable_input_first,
-            input_names=input_names,
-            output_names=output_names,
         )
 
     def get_axis_value(self):