Skip to content

Commit

Permalink
Add onnx node name propagation; fix LayerTemporalOrdering attribute
Browse files Browse the repository at this point in the history
  • Loading branch information
asyms committed May 16, 2024
1 parent f890e92 commit 3a6c2ac
Show file tree
Hide file tree
Showing 7 changed files with 23 additions and 40 deletions.
12 changes: 12 additions & 0 deletions zigzag/parser/MappingValidator.py
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,18 @@ class MappingValidator:
},
"required": False,
},
"temporal_ordering": {
"type": "list",
"schema": {
"type": "list",
"items": [
{"type": "string"},
{"type": "integer"}
],
"minlength": 2,
"maxlength": 2
}
}
}

def __init__(self, data: Any):
Expand Down
26 changes: 1 addition & 25 deletions zigzag/parser/onnx/ConvParser.py
Original file line number Diff line number Diff line change
Expand Up @@ -32,30 +32,6 @@ def run(self) -> LayerNode:
"""! Run the parser and return the created LayerNode object"""
return self.generate_layer_node_for_conv()

# 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.")

# def get_input_output_weight_data_type(self):
# """! Return the data type of the input, output and weight tensors of this node."""
# input_name = self.node.input[0]
# output_name = self.node.output[0]
# weight_name = self.get_weight_name(self.node)

# input_elem_type = get_onnx_tensor_type(input_name, self.onnx_model).elem_type
# output_elem_type = get_onnx_tensor_type(output_name, self.onnx_model).elem_type
# weight_elem_type = get_onnx_tensor_type(weight_name, self.onnx_model).elem_type

# return input_elem_type, output_elem_type, weight_elem_type

def get_layer_node_input_format(
self,
kernel_shape: list[int],
Expand All @@ -73,7 +49,7 @@ def get_layer_node_input_format(

data: dict[str, Any] = {}
data["id"] = self.node_id
data["name"] = f"Layer{self.node_id}"
data["name"] = self.node_name
data["operator_type"] = self.node.op_type
# IMPORTANT: If any of the input loops require padding, they should be defined as the rightmost dimensions
# in the equation. This is because we construct the dimensionality order and then add the padding to those last
Expand Down
2 changes: 1 addition & 1 deletion zigzag/parser/onnx/GemmParser.py
Original file line number Diff line number Diff line change
Expand Up @@ -40,7 +40,7 @@ def get_layer_node_input_format(

data: dict[str, Any] = {}
data["id"] = self.node_id
data["name"] = f"Layer{self.node_id}"
data["name"] = self.node_name
data["operator_type"] = self.node.op_type
data["equation"] = "O[b][k]+=W[k][c]*I[b][c]"
data["loop_dims"] = ["B", "C", "K"]
Expand Down
2 changes: 1 addition & 1 deletion zigzag/parser/onnx/MatMulParser.py
Original file line number Diff line number Diff line change
Expand Up @@ -47,7 +47,7 @@ def get_layer_node_input_format(

data: dict[str, Any] = {}
data["id"] = self.node_id
data["name"] = f"Layer{self.node_id}"
data["name"] = self.node_name
data["operator_type"] = self.node.op_type
data["equation"] = "O[b][k]+=W[k][c]*I[b][c]"
data["loop_dims"] = ["B", "C", "K"]
Expand Down
4 changes: 2 additions & 2 deletions zigzag/parser/workload_factory.py
Original file line number Diff line number Diff line change
Expand Up @@ -245,7 +245,7 @@ def create_temporal_ordering(self) -> LayerTemporalOrdering:
"""! This attribute lacks support within the MappingValidator. Returns an empty instance in case it is not
provided (to be compatible with older code) or raises an error if it is present in the user-provided data.
"""
if "temporal_ordering" not in self.mapping_data or self.mapping_data["temporal_ordering"] is None:
if "temporal_ordering" not in self.mapping_data or not self.mapping_data["temporal_ordering"]:
return LayerTemporalOrdering.empty()

raise NotImplementedError()
return LayerTemporalOrdering(self.mapping_data["temporal_ordering"])
4 changes: 2 additions & 2 deletions zigzag/stages/TemporalOrderingConversionStage.py
Original file line number Diff line number Diff line change
Expand Up @@ -55,8 +55,8 @@ def convert_user_temporal_mapping(self, user_temporal_mapping: LayerTemporalOrde
spatial_mapping = self.spatial_mapping
layer = self.layer
layer_dim_sizes = layer.layer_dim_sizes
for i, utm in list(enumerate(user_temporal_mapping))[::-1]:
if utm[0] not in layer_dim_sizes:
for i, utm in list(enumerate(user_temporal_mapping.data))[::-1]:
if utm[0] not in layer_dim_sizes.layer_dims:
logger.warning(
f"Supplied temporal ordering {utm} for layer {layer} thrown out because loop not present in the layer"
)
Expand Down
13 changes: 4 additions & 9 deletions zigzag/workload/layer_attributes.py
Original file line number Diff line number Diff line change
Expand Up @@ -170,19 +170,14 @@ def extract_pr_loop_info(relations: list["LayerDimRelation"]) -> tuple[PrLoop, L


class LayerTemporalOrdering(LayerAttribute):
"""
# TODO is this ever used?
"""
DEFAULT = []

def __init__(self, data: dict[LayerOperand, UnrollFactorInt]):
self.data = data
def __init__(self, data: list[list[str, int]]):
self.data = [(LayerDim(loop[0]), loop[1]) for loop in data]

@staticmethod
def empty():
return LayerTemporalOrdering({})

def __delitem__(self, x: LayerOperand):
del self.data[x]
return LayerTemporalOrdering(LayerTemporalOrdering.DEFAULT)


class LayerPadding(LayerAttribute):
Expand Down

0 comments on commit 3a6c2ac

Please sign in to comment.