Parsing update (#83)

* ignore .sh files

* rename all files to lowercase and restructure stages

* add results stages and fix imports

* fix precommit line length req

* allow layers with eq of type O[a]=I[a]*W[] (single operand) for max, relu, etc

* add accelerator and mapping arguments to ONNXOperatorParsr

* make W/I layer op a constant

* complete rebase dev-rename-files

* rename module import Accelerator -> accelerator

---------

Co-authored-by: Arne Symons <[email protected]>

zigzag/__init__.py

@@ -1,3 +1 @@
 __version__ = "3.6.1"
-
-# from . import visualization as visualization

zigzag/datatypes.py

@@ -93,13 +93,16 @@ def __init__(self, name: str):
 class Constants:
     """! Store constant objects used throughout ZigZag (instead of hardcoding them)"""
 
-    OUTPUT_LAYER_OP = LayerOperand(AcceleratorValidator.OUTPUT_OPERAND_STR)
-    FINAL_OUTPUT_LAYER_OP = LayerOperand(AcceleratorValidator.FINAL_OUTPUT_OPERAND_STR)
-    OUTPUT_MEM_OP = MemoryOperand(AcceleratorValidator.OUTPUT_OPERAND_STR)
-    FINAL_OUTPUT_MEM_OP = MemoryOperand(AcceleratorValidator.FINAL_OUTPUT_OPERAND_STR)
+    LAYER_OP_I = LayerOperand("I")
+    LAYER_OP_W = LayerOperand("W")
+    OUTPUT_LAYER_OP = LayerOperand("O")
+    FINAL_OUTPUT_LAYER_OP = LayerOperand("O_final")
+
+    MEM_OP_1 = MemoryOperand("I1")
+    MEM_OP_2 = MemoryOperand("I2")
+    OUTPUT_MEM_OP = MemoryOperand("O")
+    FINAL_OUTPUT_MEM_OP = MemoryOperand("O_final")
 
-    MEM_OP_1 = MemoryOperand(AcceleratorValidator.MEM_OP_1_STR)
-    MEM_OP_2 = MemoryOperand(AcceleratorValidator.MEM_OP_2_STR)
     UNKNOWN_DIM_OPERATOR = LayerDim("*")
 
 

zigzag/parser/accelerator_validator.py

@@ -15,13 +15,6 @@ class AcceleratorValidator:
     DIMENSION_REGEX = r"^D\d$"
     PORT_REGEX = r"^[r]?[w]?_port_\d+$"
 
-    # Intermediate output operand. Hard coded, and must be specified by the user as such
-    OUTPUT_OPERAND_STR = "O"
-    # Final output operand after scaling. Hard coded, and must be specified by the user as such
-    FINAL_OUTPUT_OPERAND_STR = "O_final"
-    MEM_OP_1_STR = "I1"
-    MEM_OP_2_STR = "I2"
-
     SCHEMA = {
         "name": {"type": "string", "required": True},
         "memories": {

zigzag/parser/onnx/gemm_parser.py

@@ -1,7 +1,5 @@
 from typing import Any
 
-from onnx import ModelProto, NodeProto
-
 from zigzag.parser.onnx.onnx_operator_parser import ONNXOperatorParser
 from zigzag.parser.onnx.utils import (
     get_attribute_ints_with_name,
@@ -14,17 +12,6 @@
 class GemmParser(ONNXOperatorParser):
     """! Parses an ONNX Gemm operator into a LayerNode"""
 
-    def __init__(
-        self,
-        node_id: int,
-        node: NodeProto,
-        nodes_outputs: dict[int, Any],
-        mapping_data: list[dict[str, Any]],
-        onnx_model: ModelProto,
-    ) -> None:
-        super().__init__(node_id, node, nodes_outputs, onnx_model)
-        self.mapping_data = mapping_data
-
     def run(self) -> LayerNode:
         """! Run the parser"""
         return self.generate_layer_node()

zigzag/parser/onnx/onnx_model_parser.py

@@ -1,12 +1,13 @@
 import logging
-from typing import Any
+from typing import Any, Type
 
-from onnx import ModelProto
+from onnx import ModelProto, NodeProto
 
 from zigzag.parser.onnx.conv_parser import ConvParser
 from zigzag.parser.onnx.default_node_parser import DefaultNodeParser
 from zigzag.parser.onnx.gemm_parser import GemmParser
 from zigzag.parser.onnx.matmul_parser import MatMulParser
+from zigzag.parser.onnx.onnx_operator_parser import ONNXOperatorParser
 from zigzag.parser.onnx.utils import (
     parse_dynamic_onnx_model,
     parse_onnx_model_from_path,
@@ -20,6 +21,14 @@
 class ONNXModelParser:
     """! Parses the ONNX model into a workload."""
 
+    # Map the node's op_type to the corresponding Parser class
+    PARSER_MAPPING: dict[str, Type[ONNXOperatorParser]] = {
+        "QLinearConv": ConvParser,
+        "Conv": ConvParser,
+        "MatMul": MatMulParser,
+        "Gemm": GemmParser,
+    }
+
     def __init__(self, onnx_model: str | ModelProto, mapping_yaml_path: str) -> None:
         assert isinstance(onnx_model, (str, ModelProto)), f"Given onnx_model is of type {type(onnx_model)}."
         assert isinstance(mapping_yaml_path, str) and mapping_yaml_path.split(".")[-1] == "yaml"
@@ -41,6 +50,12 @@ def run(self) -> ONNXWorkload:
 
         return self.parse_workload_from_onnx_model_and_mapping()
 
+    def get_parser_class(self, node: NodeProto):
+        parser_class = ONNXModelParser.PARSER_MAPPING.get(node.op_type)
+        if not parser_class:
+            return DefaultNodeParser
+        return parser_class
+
     def parse_workload_from_onnx_model_and_mapping(self):
         """! Converts an onnx model into a workload object.
         We scan the model for all convolutional layers, and setup a Layer object for each of those using the mapping.
@@ -72,15 +87,14 @@ def parse_workload_from_onnx_model_and_mapping(self):
             nodes_inputs[node_id] = node.input
             nodes_outputs[node_id] = node.output
 
-            if node.op_type in ["QLinearConv", "Conv"]:
-                parser = ConvParser(node_id, node, nodes_outputs, self.mapping_data, self.onnx_model)
-            elif node.op_type in ["MatMul"]:
-                parser = MatMulParser(node_id, node, nodes_outputs, self.mapping_data, self.onnx_model)
-            elif node.op_type in ["Gemm"]:
-                parser = GemmParser(node_id, node, nodes_outputs, self.mapping_data, self.onnx_model)
-            # it is not a convolutional node, so create a DummyNode
-            else:
-                parser = DefaultNodeParser(node_id, node, nodes_outputs, self.onnx_model)
+            parser_class = self.get_parser_class(node)
+            parser = parser_class(
+                node_id=node_id,
+                node=node,
+                nodes_outputs=nodes_outputs,
+                onnx_model=self.onnx_model,
+                mapping_data=self.mapping_data,
+            )
 
             node_obj = parser.run()
             # Add the node_obj to the ONNXWorkload

zigzag/parser/onnx/onnx_operator_parser.py

@@ -3,6 +3,7 @@
 
 from onnx import ModelProto, NodeProto
 
+from zigzag.hardware.architecture.accelerator import Accelerator
 from zigzag.parser.onnx.utils import get_attribute_ints_with_name, get_onnx_tensor_type
 from zigzag.workload.layer_node_abc import LayerNodeABC
 
@@ -21,11 +22,16 @@ def __init__(
         node: NodeProto,
         nodes_outputs: dict[int, Any],
         onnx_model: ModelProto,
+        *,
+        mapping_data: list[dict[str, Any]] | None = None,
+        accelerator: Accelerator | None = None,
     ) -> None:
         self.node_id = node_id
         self.node = node
         self.nodes_outputs = nodes_outputs
         self.onnx_model = onnx_model
+        self.mapping_data = mapping_data
+        self.accelerator = accelerator
 
     @abstractmethod
     def run(self) -> LayerNodeABC:
@@ -47,13 +53,15 @@ def get_weight_name(self, node: NodeProto):
         """! Return the name of the weight input of this node depending on its operator type
         @param node (NodeProto): The node
         """
-        op_type = node.op_type  # 'Conv', 'QLinearConv', ...
-        if op_type == "Conv":
-            return node.input[1]
-        elif op_type == "QLinearConv":
-            return node.input[3]
-        else:
-            raise NotImplementedError(f"Retrieving weight name for onnx node of type {op_type} is not supported.")
+        match node.op_type:
+            case "Conv":
+                return node.input[1]
+            case "QLinearConv":
+                return node.input[3]
+            case _:
+                raise NotImplementedError(
+                    f"Retrieving weight name for onnx node of type {node.op_type} is not supported."
+                )
 
     def get_node_predecessors(self) -> list[int]:
         """Compute node input sources"""

zigzag/parser/workload_validator.py

@@ -10,9 +10,10 @@
 class WorkloadValidator:
     """Class to validate user-defined workloads from yaml files according to the given schema and rules"""
 
-    EQUATION_REGEX = r"^O(\[\w+\])+\+?=\w(\[\w+\])+[*+]\w(\[\w+\])+$"
+    EQUATION_REGEX = r"^O(\[\w+\])+\+?=\w(\[\w+\])+[*+]\w(\[(?:\w+)?\])+$"
     LAYER_DIM_RELATION_REGEX = r"^(\w+)\s*=\s*(?:(\w+)\s*\*\s*)?(\w+)\s*\+\s*(?:(\w+)\s*\*\s*)?(\w+)$"
     ALLOWED_OPERATORS: list[str] = [
+        "Max",
         "Conv",
         "Pooling",
         "Add",

zigzag/stages/mapping/spatial_mapping_generation.py

@@ -180,8 +180,8 @@ def generate_spatial_mappings(self) -> Generator[SpatialMapping, None, None]:
             yield candidate
 
     def limit_unrolling_to_mem_bandwidth(
-        self, mapping: dict[OADimension, dict[LayerDim, UnrollFactorInt]]
-    ) -> dict[OADimension, dict[LayerDim, UnrollFactorInt]]:
+        self, mapping: dict[OADimension, dict[LayerDim, int]]
+    ) -> dict[OADimension, dict[LayerDim, int]]:
         """! Scale the given unroll factors such that they do not exceed the bandwidths of the memory structure"""
 
         def conditional_log(layer_dim: LayerDim, oa_dim: OADimension, value: int, mem_name: str):
@@ -296,7 +296,7 @@ def adjust_unrolling_factors(factors: list[UnrollFactor], max_unrolling: float)
                         mapping[oa_dim].update(unrollings)
         return mapping
 
-    def get_max_unrolling(self) -> dict[OADimension, dict[LayerDim, UnrollFactorInt]]:
+    def get_max_unrolling(self) -> dict[OADimension, dict[LayerDim, int]]:
         """! Generate a SpatialMapping that contains the maximal unroll factor for every Operational
         Array OADimension and every Layer Dimensions. Note that this is NOT a valid mapping as each
         OADimension contains ALL Layer Dimensions, maximally unrolled."""
@@ -316,7 +316,7 @@ def get_max_unrolling(self) -> dict[OADimension, dict[LayerDim, UnrollFactorInt]
     def generate_spatial_mapping_single_oa_dim(
         self,
         unroll_hints: set[LayerDim],
-        max_unrollings: dict[LayerDim, UnrollFactorInt],
+        max_unrollings: dict[LayerDim, int],
         oa_dim_size: int,
     ) -> Generator[MappingSingleOADim, None, None]:
         """! Generate a list of possible mappings for the given OADimension. Possible mappings include

zigzag/workload/dummy_node.py

@@ -1,4 +1,4 @@
-from zigzag.datatypes import LayerOperand
+from zigzag.datatypes import Constants
 from zigzag.workload.layer_attributes import InputOperandSource
 from zigzag.workload.layer_node_abc import LayerNodeABC
 
@@ -18,16 +18,17 @@ def __init__(self, node_id: int, predecessors: list[int], node_type: str, node_n
         """
         LayerNodeABC.__init__(self, node_id, node_name)
 
-        if len(predecessors) == 0:
-            self.input_operand_source: InputOperandSource = {}
-        elif len(predecessors) == 1:
-            self.input_operand_source = {LayerOperand("I"): predecessors[0]}
-        else:
-            # We currently don't support more than 2 sources so we can also use `I` and `W` for the layer operands
-            self.input_operand_source = {
-                LayerOperand("I"): predecessors[0],
-                LayerOperand("W"): predecessors[1],
-            }
+        match len(predecessors):
+            case 0:
+                self.input_operand_source: InputOperandSource = {}
+            case 1:
+                self.input_operand_source = {Constants.LAYER_OP_I: predecessors[0]}
+            case _:
+                # We currently don't support more than 2 sources so we can also use `I` and `W` for the layer operands
+                self.input_operand_source = {
+                    Constants.LAYER_OP_I: predecessors[0],
+                    Constants.LAYER_OP_W: predecessors[1],
+                }
 
         self.type = node_type
         # We assume these nodes are mapped on a core with id -1

zigzag/workload/onnx_workload.py

@@ -25,7 +25,6 @@ def add(self, node_id: int, node_obj: LayerNodeABC):
         self.add_workload_node(node_obj)
         edges: list[tuple[LayerNodeABC, LayerNodeABC]] = []
         for parent_id in node_obj.input_operand_source.values():
-            # for parent_id in parents:
             parent_node_obj = self.node_id_to_obj[parent_id]
             edges.append((parent_node_obj, node_obj))
             self.add_workload_edges_from(edges)