Mamba support

.gitignore

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

main_stream_co copy.py

@@ -0,0 +1,69 @@
+import logging as _logging
+import re
+
+from stream.api import optimize_allocation_co
+from stream.utils import CostModelEvaluationLUT
+from stream.visualization.memory_usage import plot_memory_usage
+from stream.visualization.schedule import (
+    visualize_timeline_plotly,
+)
+
+_logging_level = _logging.INFO
+_logging_format = "%(asctime)s - %(name)s.%(funcName)s +%(lineno)s - %(levelname)s - %(message)s"
+_logging.basicConfig(level=_logging_level, format=_logging_format)
+
+############################################INPUTS############################################
+accelerator = "stream/inputs/examples/hardware/tpu_like_quad_core.yaml"
+workload_path = "outputs/custom_ssm.onnx"
+mapping_path = "stream/inputs/examples/mapping/tpu_like_quad_core copy.yaml"
+mode = "fused"
+layer_stacks = [tuple(range(0, 11)), tuple(range(11, 22))] + list((i,) for i in range(22, 49))
+##############################################################################################
+
+################################PARSING###############################
+hw_name = accelerator.split("/")[-1].split(".")[0]
+wl_name = re.split(r"/|\.", workload_path)[-1]
+if wl_name == "onnx":
+    wl_name = re.split(r"/|\.", workload_path)[-2]
+experiment_id = f"{hw_name}-{wl_name}-{mode}-constraint_optimization"
+######################################################################
+
+scme = optimize_allocation_co(
+    hardware=accelerator,
+    workload=workload_path,
+    mapping=mapping_path,
+    mode=mode,
+    layer_stacks=layer_stacks,
+    experiment_id=experiment_id,
+    output_path="outputs",
+    skip_if_exists=False,
+)
+
+############PLOTTING#############
+plot_full_schedule = True
+draw_dependencies = True
+plot_data_transfer = True
+section_start_percent = (0,)
+percent_shown = (100,)
+#################################
+
+#########################PLOTTING PATHS##############################
+timeline_fig_path_plotly = f"outputs/{experiment_id}/schedule.html"
+memory_fig_path = f"outputs/{experiment_id}/memory.png"
+#####################################################################
+
+#####################CostModelEvaluationLUT LOAD#############################
+cost_lut_path = f"outputs/{experiment_id}/cost_lut_post_co.pickle"
+cost_lut = CostModelEvaluationLUT(cost_lut_path)
+#############################################################################
+
+# Plotting schedule timeline of best SCME
+visualize_timeline_plotly(
+    scme,
+    draw_dependencies=draw_dependencies,
+    draw_communication=plot_data_transfer,
+    fig_path=timeline_fig_path_plotly,
+    cost_lut=cost_lut,
+)
+# Plotting memory usage of best SCME
+plot_memory_usage(scme, section_start_percent, percent_shown, fig_path=memory_fig_path)

stream/inputs/examples/mapping/tpu_like_quad_core copy.yaml

@@ -0,0 +1,55 @@
+- name: default
+  core_allocation: [0, 1, 2, 3]
+  intra_core_tiling:
+    - D, 1
+  inter_core_tiling:
+    - K, *
+
+- name: Conv
+  core_allocation: [0, 1, 2, 3]
+  intra_core_tiling:
+    - OY, 1
+  inter_core_tiling:
+    - K, *
+
+- name: Gemm
+  core_allocation: [0, 1, 2, 3]
+  intra_core_tiling:
+    - D, 1
+  inter_core_tiling:
+    - H, *
+
+- name: Pool
+  core_allocation: [4]
+  intra_core_tiling:
+    - OY, 1
+  inter_core_tiling:
+    - K, *
+
+- name: MaxPool
+  core_allocation: [4]
+  intra_core_tiling:
+    - OY, 1
+  inter_core_tiling:
+    - K, *
+
+- name: AveragePool
+  core_allocation: [4]
+  intra_core_tiling:
+    - OY, 1
+  inter_core_tiling:
+    - K, *
+
+- name: GlobalAveragePool
+  core_allocation: [4]
+  intra_core_tiling:
+    - OY, 1
+  inter_core_tiling:
+    - K, *
+
+- name: Add
+  core_allocation: [5]
+  intra_core_tiling:
+    - D, 1
+  inter_core_tiling:
+    - H, *

stream/node_tensor.py

@@ -47,7 +47,7 @@ def _get_and_increment_pointer(self):
     @property
     def shape(self) -> None:  # type: ignore
         """Protect the original shape attribute to prevent errors"""
-        raise ValueError("The numpy shape of NodeTensor is hidden in an abstraction layer")
+        raise ValueError("The numpy shape of NodeTensor is hidden in an abstraction layer. Call `tensor_shape` instead")
 
     @property
     def full_shape(self):
@@ -125,6 +125,25 @@ def gather(self, gather_indices: int | list[int], axis: int) -> "NodeTensor":
         axis = axis - 1 if axis < 0 else axis
         return (np.take(self.as_ndarray(), gather_indices, axis=axis)).view(NodeTensor)
 
+    def split(self, split_indices: list[int], axis: int) -> "list[NodeTensor]":
+        axis = axis - 1 if axis < 0 else axis
+        return [t.view(NodeTensor) for t in np.split(self.as_ndarray(), split_indices, axis=axis)]
+
+    def slice(self, starts: int, ends: int, axis: int, steps: int) -> "NodeTensor":
+        assert starts != 1 and ends != -1
+        axis = len(self.tensor_shape) - 1 if axis < 0 else axis
+        match axis:
+            case 0:
+                return self.as_ndarray()[starts:ends:steps, ...].view(NodeTensor)
+            case 1:
+                return self.as_ndarray()[:, starts:ends:steps, ...].view(NodeTensor)
+            case 2:
+                return self.as_ndarray()[:, :, starts:ends:steps, ...].view(NodeTensor)
+            case 3:
+                return self.as_ndarray()[:, :, :, starts:ends:steps, ...].view(NodeTensor)
+            case _:
+                raise NotImplementedError
+
     def concat_with_empty(self, shape: tuple[int, ...], axis: int, variable_input_first: bool):
         empty_shape = self.convert_to_full_shape(shape)
         empty_tensor = np.zeros(empty_shape, dtype=object)

stream/onnx_utils.py

@@ -0,0 +1,107 @@
+import numpy as np
+from onnx import AttributeProto, ModelProto, NodeProto, numpy_helper
+from zigzag.parser.onnx.utils import get_onnx_tensor_type
+
+
+def get_attribute_as_ints(
+    node: NodeProto, attribute_name: str, default: list[int] | int | None = None
+) -> list[int] | int:
+    """! Return the value of an attribute of given name from the given attributes
+    If name does not exist in attrs, the default provided by the caller is used.
+    If the caller doesn't supply a default, an error is thrown.
+
+    """
+    attrs = node.attribute
+    attrs_names = [attr.name for attr in attrs]
+    try:
+        name_idx = attrs_names.index(attribute_name)
+        value = attrs[name_idx]
+        attr_type = value.type
+        if attr_type == AttributeProto.AttributeType.INT:  # type: ignore
+            return int(value.i)
+        elif attr_type == AttributeProto.AttributeType.INTS:  # type: ignore
+            return list(value.ints)
+        elif attr_type == AttributeProto.AttributeType.TENSOR:  # type: ignore
+            return list(numpy_helper.to_array(value.t).tolist())  # type: ignore
+        else:
+            raise NotImplementedError(f"Attribute extraction of type {attr_type} not supported.")
+    except ValueError as exc:
+        if default is not None:
+            return default
+        else:
+            raise ValueError(
+                f"Node {node.name} has no attribute called {attribute_name} and no default was given. Attributes = {attrs_names}."
+            ) from exc
+
+
+def get_onnx_input_shapes(node: NodeProto, onnx_model: ModelProto) -> list[list[int]]:
+    """Return the shape of each input operand"""
+    input_names = node.input
+    input_shapes = [get_onnx_tensor_type(name, onnx_model).shape for name in input_names]
+    return input_shapes
+
+
+def get_onnx_output_shapes(node: NodeProto, onnx_model: ModelProto) -> list[list[int]]:
+    """Return the shape of each output operand"""
+
+    output_names = node.output
+    output_shapes = [get_onnx_tensor_type(name, onnx_model).shape for name in output_names]
+    return output_shapes
+
+
+def has_asymmetric_input_data(node: NodeProto, onnx_model: ModelProto):
+    """Return true iff the node has two inputs and the input nodes have a different shape"""
+    if len(node.input) != 2:
+        return False
+
+    input_shape1, input_shape2 = get_onnx_input_shapes(node, onnx_model)
+    return input_shape1 != input_shape2
+
+
+def get_constant_tensor_int(onnx_model: ModelProto, constant_output_name: str):
+    """In some cases, the constants to a node (e.g. slice and split indices) are saved as tensors within a constant
+    node. The output name of the constant nodes corresponds to the input name of the node that uses this constant
+    tensor."""
+
+    for node in onnx_model.graph.node:
+        if node.op_type == "Constant" and node.output[0] == constant_output_name:
+            for attr in node.attribute:
+                if attr.name == "value":
+                    tensor = attr.t  # This is an ONNX TensorProto
+                    # Decode tensor to a numpy array
+                    array = np.frombuffer(tensor.raw_data, dtype=int)
+                    array = array.reshape([dim for dim in tensor.dims])
+
+                    return [int(i) for i in array]
+
+    raise ValueError(f"Cannot find {constant_output_name}")
+
+
+def get_axis_attribute(node: NodeProto):
+    """Find the value of the axis associated with this ONNX node"""
+    ATTR_NAME = "axis"
+
+    value = get_attribute_as_ints(node, ATTR_NAME)
+    if not isinstance(value, int):
+        raise ValueError(f"{ATTR_NAME} attribute as list of ints not supported")
+    return value
+
+
+def get_split_attribute(node: NodeProto, onnx_model: ModelProto):
+    output_name = next(n for n in node.input if "split" in n.lower())
+    return get_constant_tensor_int(onnx_model, output_name)
+
+
+def get_slice_attributes(node: NodeProto, onnx_model: ModelProto):
+    """Get the `starts`, `ends`, `axes` and `steps` tensors for a slice node.
+    NOTE: this assumes that the attributes are given as inputs in this order"""
+    if len(node.input) != 5:
+        raise NotImplementedError("Unsure how to get slice attributes from Node")
+
+    starts_output_name, ends_output_name, axes_output_name, steps_output_name = node.input[1:5]
+
+    starts_value = get_constant_tensor_int(onnx_model, starts_output_name)
+    ends_value = get_constant_tensor_int(onnx_model, ends_output_name)
+    axes_value = get_constant_tensor_int(onnx_model, axes_output_name)
+    steps_value = get_constant_tensor_int(onnx_model, steps_output_name)
+    return starts_value, ends_value, axes_value, steps_value

stream/parser/onnx/asymmetric_simd.py

@@ -1,12 +1,9 @@
 from typing import Any
 
-from zigzag.parser.onnx.utils import (
-    get_node_input_output_dimension_shapes,
-)
 from zigzag.parser.workload_factory import LayerNodeFactory
 
+from stream.onnx_utils import get_onnx_input_shapes, get_onnx_output_shapes
 from stream.parser.onnx.operator_parser import OnnxComputeOperatorParser
-from stream.utils import get_onnx_input_shapes
 from stream.workload.computation.computation_node import ComputationNode
 
 
@@ -30,7 +27,7 @@ def get_layer_node_user_format(self, input_shape: list[int], output_shape: list[
         data["name"] = self.node.name
         data["operator_type"] = self.node.op_type
         data["operand_source"] = self.get_operand_source_input_format()
-        data["operand_precision"] = self.get_operand_precision_input_format()
+        data["operand_precision"] = self.get_operand_precision_user_format()
         data["dimension_relations"] = []
         data["loop_sizes"] = output_shape
 
@@ -41,8 +38,15 @@ def get_layer_node_user_format(self, input_shape: list[int], output_shape: list[
 
     def generate_node(self):
         # Get the input and output activation shapes
-        input_shape1, input_shape2 = get_onnx_input_shapes(self.node, self.onnx_model)
-        _, output_shape = get_node_input_output_dimension_shapes(self.node, self.onnx_model)
+        input_shapes = get_onnx_input_shapes(self.node, self.onnx_model)
+        if len(input_shapes) != 2:
+            raise NotImplementedError("Only SIMD nodes with input length 2 are supported")
+        input_shape1, input_shape2 = input_shapes
+
+        output_shapes = get_onnx_output_shapes(self.node, self.onnx_model)
+        if len(output_shapes) != 1:
+            raise NotImplementedError("Only SIMD nodes with input length 2 are supported")
+        output_shape = output_shapes.pop()
 
         if input_shape1 == output_shape:
             non_batched_input_shape = input_shape2
@@ -57,11 +61,13 @@ def generate_node(self):
         node_factory = LayerNodeFactory(node_data, mapping_data=None)
         node_attrs = node_factory.create_node_attr()
         mapping = self.get_mapping_this_node()
+        input_names = list(self.node.input)
 
         return ComputationNode(
             node_id=self.node_id,
             node_name=self.node.name,
             node_attr=node_attrs,
             mapping_attr=mapping,
             op_type=self.node.op_type,
+            input_names=input_names,
         )

stream/parser/onnx/concat.py

@@ -7,10 +7,21 @@
 class ConcatParser(OnnxOperatorParser):
     """Parses an onnx gather operator into a ConcatNode."""
 
+    def get_axis_value(self):
+        AXIS_ATTR = "axis"
+
+        """Find the value of the axis associated with this concat node in ONNX"""
+        # `axis` is an attribute of the node
+        try:
+            axis_attr = next(filter(lambda x: x.name == AXIS_ATTR, self.node.attribute))
+            return axis_attr.i
+        except StopIteration:
+            raise ValueError("Axis attribute not found in ONNX node")
+
     def generate_node(self):
         predecessors = self.get_node_predecessors()
-
         axis = self.get_axis_value()
+        input_names = list(self.node.input)
 
         input_1, input_2 = self.node.input[0], self.node.input[1]
 
@@ -36,13 +47,5 @@ def generate_node(self):
             axis=axis,
             constant_shape=constant_shape,
             variable_input_first=variable_input_first,
+            input_names=input_names,
         )
-
-    def get_axis_value(self):
-        """Find the value of the axis associated with this concat node in ONNX"""
-        # `axis` is an attribute of the node
-        try:
-            axis_attr = next(filter(lambda x: x.name == "axis", self.node.attribute))
-            return axis_attr.i
-        except StopIteration:
-            raise ValueError("Axis attribute not found in ONNX node")