Upload 379 files

610d9c2 verified about 3 hours ago

11.8 kB

# -------------------------------------------------------------------------
# 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

	# -------------------------------------------------------------------------
	# 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