bugfix mapping

stream/inputs/examples/mapping/tpu_like_quad_core.yaml

@@ -1,14 +1,14 @@
 - name: default
   core_allocation: [0, 1, 2, 3]
   inter_core_tiling:
-    - K, 4
+    - K, *
   intra_core_tiling:
     - D, 4
 
 - name: Conv
   core_allocation: [0, 1, 2, 3]
   inter_core_tiling:
-    - K, 4
+    - K, *
   intra_core_tiling:
     - D, 4
 

stream/parser/mapping_factory.py

@@ -71,6 +71,7 @@ def create_intra_core_tiling(self, mapping_data: dict[str, Any]):
     def __convert_layer_dim_int_pair(self, pair: str):
         """Convert strings such as `D, 4` into a LayerDim and int"""
         layer_dim_str = pair.split(",")[0]
-        unrolling = int(pair.split(",")[-1])
+        unrolling_str = pair.split(",")[-1]
+        unrolling = int(unrolling_str) if "*" not in unrolling_str else "*"
         layer_dim = LayerDim(layer_dim_str)
         return layer_dim, unrolling

stream/parser/mapping_validator.py

@@ -20,18 +20,12 @@ class MappingValidator:
         "core_allocation_is_fixed": {"type": "boolean", "default": False},
         "inter_core_tiling": {
             "type": "list",
-            "schema": {
-                "type": "list",
-                "schema": {"type": "string", "regex": r"^[A-Z]+, \d+|\*+$"},
-            },
+            "schema": {"type": "string", "regex": r"^[A-Z]+, [\d+|\*]$"},
             "default": [],
         },
         "intra_core_tiling": {
             "type": "list",
-            "schema": {
-                "type": "list",
-                "schema": {"type": "string", "regex": r"^[A-Z]+, \d+|\*+$"},
-            },
+            "schema": {"type": "string", "regex": r"^[A-Z]+, [\d+|\*]$"},
             "default": [],
         },
         "spatial_mapping": {

stream/parser/onnx/conv.py

@@ -131,18 +131,15 @@ def generate_node(self):
             oa_dimension_shape,
         )
 
-        node_factory = LayerNodeFactory(node_data, self.mapping_data)
+        node_factory = LayerNodeFactory(node_data, mapping_data=None)
         node_attrs = node_factory.create_node_attr()
-
-        # Override spatial mapping by the one defined in the core's dataflows
-        core_allocation = node_attrs.core_allocation
-        spatial_mapping = self.accelerator.get_spatial_mapping_from_core(core_allocation)
-        node_attrs.spatial_mapping = spatial_mapping
+        mapping = self.get_mapping_this_node()
 
         return ComputationNode(
             node_id=self.node_id,
             node_name=self.node.name,
             node_attr=node_attrs,
+            mapping_attr=mapping,
             op_type=ConvParser.OP_TYPE,
             operand_tensor_reshape=None,
         )

stream/parser/onnx/operator_parser.py

@@ -58,7 +58,7 @@ def run(self) -> Generator[ComputationNode, None, None]:
     @abstractmethod
     def get_layer_node_user_format(self, input_shape: list[int], output_shape: list[int]) -> dict[str, Any]: ...
 
-    def get_operand_precision_input_format(self):
+    def get_operand_precision_input_format(self) -> dict[str, int]:
         act_precision = self.get_activation_precision()
         weight_precision = self.get_weight_precision()
         intermediate_output_precision = self.get_intermediate_output_precision()
@@ -83,6 +83,23 @@ def get_operand_precision_input_format(self):
             case _:
                 raise ValueError("No more than 2 layer predecessors expected")
 
+    def get_mapping_this_node(self):
+        """Get the mapping that corresponds to this node's operator. Replace the spatial mapping with the corresponding
+        core's dataflows.
+        NOTE The core's dataflow always precedes the mapping's spatial mapping
+        TODO Mapping based on node name instead of note operator is not yet supported
+        """
+        mapping = self.all_mappings.get(self.node.op_type, self.all_mappings["default"])
+
+        # Override spatial mapping by the one defined in the core's dataflows
+        try:
+            core_dataflow = self.accelerator.get_spatial_mapping_from_core(mapping.core_allocation)
+            mapping.spatial_mapping = core_dataflow
+        except ValueError:
+            pass
+
+        return mapping
+
     def generate_node(self):
         # Get the input and output activation shapes
         input_shape, output_shape = get_node_input_output_dimension_shapes(self.node, self.onnx_model)
@@ -92,14 +109,7 @@ def generate_node(self):
         node_factory = LayerNodeFactory(node_data, mapping_data=[])
         node_attrs = node_factory.create_node_attr()
 
-        mapping = self.all_mappings.get(self.node.op_type, self.all_mappings["default"])
-
-        # Override spatial mapping by the one defined in the core's dataflows
-        try:
-            core_dataflow = self.accelerator.get_spatial_mapping_from_core(mapping.core_allocation)
-            mapping.spatial_mapping = core_dataflow
-        except ValueError:
-            pass
+        mapping = self.get_mapping_this_node()
 
         return ComputationNode(
             node_id=self.node_id,

stream/parser/onnx/pooling.py

@@ -114,16 +114,13 @@ def generate_node(self):
             ia_dimension_shape,
             oa_dimension_shape,
         )
-        node_factory = LayerNodeFactory(node_data, self.mapping_data)
+        node_factory = LayerNodeFactory(node_data, None)
         node_attrs = node_factory.create_node_attr()
-
-        # Override spatial mapping by the one defined in the core's dataflows
-        core_allocation = node_attrs.core_allocation
-        spatial_mapping = self.accelerator.get_spatial_mapping_from_core(core_allocation)
-        node_attrs.spatial_mapping = spatial_mapping
+        mapping = self.get_mapping_this_node()
 
         return PoolingNode(
             node_id=self.node_id,
             node_name=self.node.name,
             node_attr=node_attrs,
+            mapping_attr=mapping,
         )

stream/stages/allocation/constraint_optimization_allocation.py

@@ -365,6 +365,8 @@ def set_fixed_allocations_for_workload(
         layer_ids_not_in_ss = [
             layer_id for stack in self.layer_stacks for layer_id in stack if layer_id not in layer_ids
         ]
+
+        # TODO this piece of code is outdates
         for layer_id_not_in_ss in layer_ids_not_in_ss:
             layer_ids_idx = np.searchsorted(layer_ids, layer_id_not_in_ss)
             for n in filter(lambda n: n.id == layer_id_not_in_ss, workload.node_list):

stream/stages/generation/hint_loops_generation.py

@@ -32,16 +32,17 @@ def __init__(
         self.hint_loops = kwargs.get("hint_loops", None)
 
     def run(self):
-        if self.mode == "fused":
-            if self.hint_loops is None:
-                self.hint_loops = self.get_hint_loops_fused()
-                self.kwargs["cn_define_mode"] = 3
-        elif self.mode == "lbl":
-            if self.hint_loops is None:
-                self.hint_loops = self.get_hint_loops_lbl()
-                self.kwargs["cn_define_mode"] = 3
-        else:
-            raise ValueError("Unsupported mode for hint loops determination.")
+        match self.mode:
+            case "fused":
+                if self.hint_loops is None:
+                    self.hint_loops = self.get_hint_loops_fused()
+                    self.kwargs["cn_define_mode"] = 3
+            case "lbl":
+                if self.hint_loops is None:
+                    self.hint_loops = self.get_hint_loops_lbl()
+                    self.kwargs["cn_define_mode"] = 3
+            case _:
+                raise ValueError("Unsupported mode for hint loops determination.")
 
         self.kwargs["accelerator"] = self.accelerator
         self.kwargs["workload"] = self.workload
@@ -73,9 +74,7 @@ def get_hint_loops_fused(self):
             if len(computation_nodes) > 1:
                 for n in computation_nodes:
                     # NOTE Take first entry of fusion_partition_dims for now
-                    hint_loops[(n.id,)] = [
-                        (n.fusion_partition_dims[0], "all"),
-                    ]
+                    hint_loops[(n.id,)] = [(n.fusion_partition_dims[0], "all")]
             else:
                 hint_loops[stack] = []
         return hint_loops

stream/stages/generation/hint_loops_partitioned_workload_generation.py

@@ -241,6 +241,8 @@ def get_finer_nodes(
 
         # Take away the outer_temporal_loops to create finer CNs for this node
         finer_node_attrs = original_node.extract_node_attr()
+        finer_node_mapping = original_node.extract_inter_core_mapping_attr()
+
         for outer_tl in outer_temporal_loops:
             outer_dim = outer_tl.dimension
             outer_size = outer_tl.size
@@ -305,8 +307,7 @@ def get_finer_nodes(
                 dim_max = dim_min + (finer_span[loop_dim] if loop_dim in finer_span else 1)
                 dim_min_max[loop_dim] = (dim_min, dim_max)
 
-            # Add the loop ranges for this cn to a copy of the finer node attributes
-            finer_node_attrs_copy = deepcopy(finer_node_attrs)
+            # finer_node_mapping_copy = deepcopy(original_node.extract_mapping_attr())
             group_id = group_id_manager.get_group_id(original_node, dim_min_max)
 
             # Create the computation node object with the computed ranges of the loop dimensions
@@ -324,7 +325,8 @@ def get_finer_nodes(
                 node_id=original_node_id,
                 sub_id=n,
                 node_name=node_name,
-                node_attr=finer_node_attrs_copy,
+                node_attr=finer_node_attrs,
+                mapping_attr=finer_node_mapping,
                 op_type=original_node.type,
                 produces_final_output=produces_final_output,
                 group_id=group_id,

stream/workload/computation/computation_node.py

@@ -1,3 +1,4 @@
+from copy import deepcopy
 from math import prod
 from typing import TypeAlias
 
@@ -28,7 +29,7 @@ class ComputationNode(LayerNode, Node):
     producer/consumer of another layer.
     """
 
-    too_large_operands: list[LayerOperand]
+    too_large_operands: list[MemoryOperand]
 
     # Map the node's op_type to the corresponding layer dimension to split on for fusion
     FUSION_DIM_MAPPING: dict[str, list[LayerDim]] = {
@@ -253,11 +254,22 @@ def reshape_operand_tensor(self, tensor: NodeTensor, operand: LayerOperand):
             new_shape = self.operand_tensor_reshape[operand]
             return tensor.reshape(new_shape)
 
-    def set_too_large_operands(self, too_large_operands: list[MemoryOperand]):
+    def set_too_large_operands(self, too_large_operands: list[LayerOperand]):
         self.too_large_operands = too_large_operands
 
     def update_loop_ranges(self, new_ranges: LoopRanges):
         """Override the loop ranges with a new value for each of the given LayerDims. Keep the old range for the
         LayerDims not defined in `new_ranges`"""
         for layer_dim in new_ranges:
             self.loop_ranges[layer_dim] = new_ranges[layer_dim]
+
+    def extract_inter_core_mapping_attr(self):
+        mapping_attr = InterCoreMappingAttributes(
+            op_type=self.type,
+            spatial_mapping=self.spatial_mapping,
+            core_allocation=self.core_allocation,
+            core_allocation_is_fixed=self.core_allocation_is_fixed,
+            intra_core_tiling=self.intra_core_tiling,
+            inter_core_tiling=self.inter_core_tiling,
+        )
+        return deepcopy(mapping_attr)

stream/workload/computation/pooling_node.py

@@ -1,6 +1,7 @@
 from zigzag.workload.layer_node import LayerNodeAttributes
 
 from stream.workload.computation.computation_node import ComputationNode
+from stream.workload.mapping import InterCoreMappingAttributes
 
 
 class PoolingNode(ComputationNode):
@@ -9,10 +10,12 @@ def __init__(
         node_id: int,
         node_name: str,
         node_attr: LayerNodeAttributes,
+        mapping_attr: InterCoreMappingAttributes,
     ):
         super().__init__(
             node_id=node_id,
             node_name=node_name,
             node_attr=node_attr,
+            mapping_attr=mapping_attr,
             op_type="pooling",
         )