File size: 11,777 Bytes
610d9c2 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 | # -------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License. See License.txt in the project root for
# license information.
# --------------------------------------------------------------------------
from __future__ import annotations
from collections import deque
import onnx
from ..onnx_model import ONNXModel
class Fusion:
"""
Base class for fusions.
"""
def __init__(self, model: ONNXModel, fused_op_type: str, search_op_type: str):
self.search_op_type: str = search_op_type
self.fused_op_type: str = fused_op_type
self.model: ONNXModel = model
self.nodes_to_remove: list = []
self.nodes_to_add: list = []
self._new_node_name_prefix = self.fused_op_type + "_fused_" + self.search_op_type + "_"
self._new_node_name_suffix = None # int|None used to create unique node names for the fused ops.
def fuse(
self,
node: onnx.NodeProto,
input_name_to_nodes: dict[str, list[onnx.NodeProto]],
output_name_to_node: dict[str, onnx.NodeProto],
):
"""
Interface function for derived fusion classes. Tries to fuse a node sequence containing
the specified node.
"""
raise NotImplementedError
def apply(self) -> bool:
"""
Apply graph fusion on the entire model graph.
"""
input_name_to_nodes = self.model.input_name_to_nodes()
output_name_to_node = self.model.output_name_to_node()
for node in self.model.nodes():
if node.op_type == self.search_op_type:
self.fuse(node, input_name_to_nodes, output_name_to_node)
self.model.remove_nodes(self.nodes_to_remove)
self.model.add_nodes(self.nodes_to_add)
graph_updated = bool(self.nodes_to_remove or self.nodes_to_add)
if graph_updated:
self.model.remove_unused_constant()
return graph_updated
def create_unique_node_name(self):
prefix = self._new_node_name_prefix
if self._new_node_name_suffix is None:
largest_suffix: int = self.model.get_largest_node_name_suffix(prefix)
self._new_node_name_suffix = largest_suffix + 1
new_name = f"{prefix}{self._new_node_name_suffix!s}"
self._new_node_name_suffix += 1
return new_name
@staticmethod
def is_safe_to_fuse_nodes(
nodes_to_remove: list[onnx.NodeProto],
keep_outputs: list[str],
input_name_to_nodes: dict[str, list[onnx.NodeProto]],
output_name_to_node: dict[str, onnx.NodeProto],
) -> bool:
for node_to_remove in nodes_to_remove:
for output_to_remove in node_to_remove.output:
if output_to_remove in keep_outputs:
continue
if output_to_remove in input_name_to_nodes:
for impacted_node in input_name_to_nodes[output_to_remove]:
if impacted_node not in nodes_to_remove:
# Not safe to remove nodes since output is used by impacted_node
return False
return True
@staticmethod
def get_node_attribute(node: onnx.NodeProto, attribute_name: str):
for attr in node.attribute:
if attr.name == attribute_name:
value = onnx.helper.get_attribute_value(attr)
return value
return None
@staticmethod
def input_index(node_output: str, child_node: onnx.NodeProto) -> int:
for index, input_name in enumerate(child_node.input):
if input_name == node_output:
return index
return -1
@staticmethod
def tensor_shape_to_list(tensor_type) -> list[int]:
shape_list = []
for d in tensor_type.shape.dim:
if d.HasField("dim_value"):
shape_list.append(d.dim_value) # known dimension
elif d.HasField("dim_param"):
shape_list.append(d.dim_param) # unknown dimension with symbolic name
else:
shape_list.append("?") # shall not happen
return shape_list
def get_constant_input(self, node: onnx.NodeProto):
for i, inp in enumerate(node.input):
value = self.model.get_constant_value(inp)
if value is not None:
return i, value
return None, None
def find_constant_input(self, node: onnx.NodeProto, expected_value: float, delta: float = 0.000001) -> int:
i, value = self.get_constant_input(node)
if value is not None and value.size == 1 and abs(value - expected_value) < delta:
return i
return -1
def has_constant_input(self, node: onnx.NodeProto, expected_value: float, delta: float = 0.000001) -> bool:
return self.find_constant_input(node, expected_value, delta) >= 0
def is_constant_with_specified_rank(self, output_name: str, rank: int) -> bool:
value = self.model.get_constant_value(output_name)
if value is None:
return False # Not an initializer
if len(value.shape) != rank:
return False # Wrong dimensions
return True
def match_first_parent(
self,
node: onnx.NodeProto,
parent_op_type: str,
output_name_to_node: dict[str, onnx.NodeProto] | None = None,
exclude: list[onnx.NodeProto] = [], # noqa: B006
) -> tuple[onnx.NodeProto | None, int | None]:
"""
Find parent node based on constraints on op_type.
Args:
node: current node.
parent_op_type (str): constraint of parent node op_type.
output_name_to_node (dict): dictionary with output name as key, and node as value.
exclude (list): list of nodes that are excluded (not allowed to match as parent).
Returns:
parent: The matched parent node. None if not found.
index: The input index of matched parent node. None if not found.
"""
if output_name_to_node is None:
output_name_to_node = self.model.output_name_to_node()
for i, inp in enumerate(node.input):
if inp in output_name_to_node:
parent = output_name_to_node[inp]
if parent.op_type == parent_op_type and parent not in exclude:
return parent, i
return None, None
def match_parent(
self,
node: onnx.NodeProto,
parent_op_type: str,
input_index: int | None = None,
output_name_to_node: dict[str, onnx.NodeProto] | None = None,
exclude: list[onnx.NodeProto] = [], # noqa: B006
return_indice: list[int] | None = None,
) -> onnx.NodeProto | None:
"""
Find parent node based on constraints on op_type and index.
When input_index is None, we will find the first parent node based on constraints,
and return_indice will be appended the corresponding input index.
Args:
node (str): current node name.
parent_op_type (str): constraint of parent node op_type.
input_index (int or None): only check the parent given input index of current node.
output_name_to_node (dict): dictionary with output name as key, and node as value.
exclude (list): list of nodes that are excluded (not allowed to match as parent).
return_indice (list): a list to append the input index when input_index is None.
Returns:
parent: The matched parent node.
"""
assert node is not None
assert input_index is None or input_index >= 0
if output_name_to_node is None:
output_name_to_node = self.model.output_name_to_node()
if input_index is None:
parent, index = self.match_first_parent(node, parent_op_type, output_name_to_node, exclude)
if return_indice is not None:
return_indice.append(index)
return parent
if input_index >= len(node.input):
# Input index out of bounds.
return None
parent = self.model.get_parent(node, input_index, output_name_to_node)
if parent is not None and parent.op_type == parent_op_type and parent not in exclude:
return parent
return None
def match_parent_path(
self,
node: onnx.NodeProto,
parent_op_types: list[str],
parent_input_index: list[int] | None = None,
output_name_to_node: dict[str, onnx.NodeProto] | None = None,
return_indice: list[int] | None = None,
) -> list[onnx.NodeProto] | None:
"""
Find a sequence of input edges based on constraints on parent op_type and index.
When input_index is None, we will find the first parent node based on constraints,
and return_indice will be appended the corresponding input index.
Args:
node (str): current node name.
parent_op_types (str): constraint of parent node op_type of each input edge.
parent_input_index (list): constraint of input index of each input edge. None means no constraint.
output_name_to_node (dict): dictionary with output name as key, and node as value.
return_indice (list): a list to append the input index
When there is no constraint on input index of an edge.
Returns:
parents: a list of matched parent node.
"""
if parent_input_index is not None:
assert len(parent_input_index) == len(parent_op_types)
if output_name_to_node is None:
output_name_to_node = self.model.output_name_to_node()
current_node = node
matched_parents = []
for i, op_type in enumerate(parent_op_types):
matched_parent = self.match_parent(
current_node,
op_type,
parent_input_index[i] if parent_input_index is not None else None,
output_name_to_node,
exclude=[],
return_indice=return_indice,
)
if matched_parent is None:
return None
matched_parents.append(matched_parent)
current_node = matched_parent
return matched_parents
def match_parent_paths(
self,
node: onnx.NodeProto,
paths: list[tuple[list[str], list[int]]],
output_name_to_node: dict[str, onnx.NodeProto],
) -> tuple[int, list[onnx.NodeProto] | None, list[int] | None]:
"""
Find a matching parent path to the given node.
"""
for i, path in enumerate(paths):
return_indice = []
matched = self.match_parent_path(node, path[0], path[1], output_name_to_node, return_indice)
if matched:
return i, matched, return_indice
return -1, None, None
def find_first_child_by_type(
self,
node: onnx.NodeProto,
child_type: str,
input_name_to_nodes: dict[str, list[onnx.NodeProto]] | None = None,
recursive: bool = True,
) -> onnx.NodeProto | None:
children = self.model.get_children(node, input_name_to_nodes)
dq = deque(children)
while len(dq) > 0:
current_node = dq.pop()
if current_node.op_type == child_type:
return current_node
if recursive:
children = self.model.get_children(current_node, input_name_to_nodes)
for child in children:
dq.appendleft(child)
return None
|