onnxruntime-directml 1.20.0__cp313-cp313-win_amd64.whl

onnxruntime/transformers/fusion_base.py

@@ -0,0 +1,137 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation.  All rights reserved.
+# Licensed under the MIT License.
+# --------------------------------------------------------------------------
+from collections import defaultdict
+from logging import getLogger
+from typing import Any, Dict, List, Optional, Sequence, Union
+
+import numpy as np
+from onnx import NodeProto, helper
+from onnx_model import OnnxModel
+
+logger = getLogger(__name__)
+
+
+class Fusion:
+    """
+    Base class for Graph Fusion
+    """
+
+    def __init__(
+        self,
+        model: OnnxModel,
+        fused_op_type: str,
+        search_op_types: Union[str, List[str]],
+        description: str = "",
+    ):
+        self.search_op_types: List[str] = [search_op_types] if isinstance(search_op_types, str) else search_op_types
+        self.fused_op_type: str = fused_op_type
+        self.description: str = f"{fused_op_type}({description})" if description else fused_op_type
+        self.model: OnnxModel = model
+        self.nodes_to_remove: List = []
+        self.nodes_to_add: List = []
+        self.prune_graph: bool = False
+        self.node_name_to_graph_name: dict = {}
+        self.this_graph_name: Optional[str] = None
+        # It is optional that subclass updates fused_count since we will also check nodes_to_add to get counter.
+        self.fused_count: defaultdict = defaultdict(int)
+
+    def increase_counter(self, fused_op_name: str):
+        """
+        Increase counter of a fused operator.
+        """
+        self.fused_count[fused_op_name] += 1
+
+    def fuse(
+        self,
+        node: NodeProto,
+        input_name_to_nodes: Dict[str, List[NodeProto]],
+        output_name_to_node: Dict[str, NodeProto],
+    ):
+        """Interface for fusion that starts from a node"""
+        raise NotImplementedError
+
+    def apply(self):
+        """
+        Apply graph fusion on the whole model graph.
+        It searched nodes of given operators, and start fusion on each of those nodes.
+        """
+        logger.debug(f"start {self.description} fusion...")
+        input_name_to_nodes = self.model.input_name_to_nodes()
+        output_name_to_node = self.model.output_name_to_node()
+
+        # This assumes that two search ops will not be fused at same time!
+        for search_op_type in self.search_op_types:
+            for node in self.model.get_nodes_by_op_type(search_op_type):
+                graph = self.model.get_graph_by_node(node)
+                if graph is None:
+                    raise Exception("Can not find node in any graph")
+                self.this_graph_name = graph.name
+                self.fuse(node, input_name_to_nodes, output_name_to_node)
+
+        op_list = [node.op_type for node in self.nodes_to_add]
+        if self.fused_count:
+            for key, value in self.fused_count.items():
+                if value:
+                    logger.info(f"Fused {key}: {value}")
+        else:
+            count = op_list.count(self.fused_op_type)
+            if count > 0:
+                logger.info(f"Fused {self.description}: {count}")
+
+        self.model.remove_nodes(self.nodes_to_remove)
+        self.model.add_nodes(self.nodes_to_add, self.node_name_to_graph_name)
+
+        if self.prune_graph:
+            self.model.prune_graph()
+        elif self.nodes_to_remove or self.nodes_to_add:
+            self.model.update_graph()
+
+    def add_initializer(self, name: str, data_type: int, dims: Sequence[int], vals: Any, raw: bool = True):
+        if raw:
+            np_type = helper.tensor_dtype_to_np_dtype(data_type)
+            if not isinstance(vals, np.ndarray):
+                bytes = np.array(vals, dtype=np_type).tobytes()
+            else:
+                bytes = vals.astype(np_type).tobytes()
+            tensor = helper.make_tensor(
+                name=name,
+                data_type=data_type,
+                dims=dims,
+                vals=bytes,
+                raw=True,
+            )
+        else:
+            tensor = helper.make_tensor(
+                name=name,
+                data_type=data_type,
+                dims=dims,
+                vals=vals,
+                raw=False,
+            )
+
+        self.model.add_initializer(tensor, self.this_graph_name)
+        return tensor
+
+    def add_nodes_to_remove(self, nodes: List[NodeProto]):
+        # Some nodes are shared between paths (e.g. rotary embedding nodes in the Q and K paths).
+        # When path A is fused, its shared nodes are added to `self.nodes_to_remove`. But when path B
+        # is fused, its shared nodes are also added to `self.nodes_to_remove`. When the nodes are
+        # iteratively removed from `self.nodes_to_remove`, path A's shared nodes are removed first.
+        # Since path A's shared nodes are removed, path B's shared nodes are not removed because they
+        # were previously removed for path A. This causes an error to print in remove_node that a node
+        # has failed to be removed.
+        #
+        # To avoid this error, we pre-emptively check if the shared nodes are already in `self.nodes_to_remove`.
+        # We could alternatively convert `self.nodes_to_remove` to a set to avoid this issue, but there could
+        # be scenarios where the nodes need to be removed in a specific order and converting to a set would
+        # lose this order.
+        for node in nodes:
+            if node not in self.nodes_to_remove:
+                self.nodes_to_remove.append(node)
+
+    def add_nodes_to_remove_with_nodes_to_keep(self, nodes: List[NodeProto], nodes_to_keep: List[NodeProto]):
+        for node in nodes:
+            if node not in self.nodes_to_remove and node not in nodes_to_keep:
+                self.nodes_to_remove.append(node)

onnxruntime/transformers/fusion_bias_add.py

@@ -0,0 +1,58 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation.  All rights reserved.
+# Licensed under the MIT License.
+# --------------------------------------------------------------------------
+from logging import getLogger
+from typing import Dict
+
+from fusion_base import Fusion
+from numpy import ndarray
+from onnx import helper
+from onnx_model import OnnxModel
+
+logger = getLogger(__name__)
+
+
+class FusionBiasAdd(Fusion):
+    def __init__(self, model: OnnxModel):
+        super().__init__(model, "BiasAdd", "Add")
+
+    def fuse(self, add_node, input_name_to_nodes: Dict, output_name_to_node: Dict):
+        """
+        Fuse Add bias and Add skip connection into BiasAdd
+        """
+
+        nodes = self.model.match_parent_path(
+            add_node,
+            ["Add", "MatMul", "BiasSplitGelu", "MatMul", "SkipLayerNormalization"],
+            [0, None, 0, 0, 0],
+            output_name_to_node,
+        )
+
+        if nodes is None:
+            return
+
+        bias_node = nodes[0]
+        skip_layer_norm = nodes[-1]
+
+        # Check skip connection is from SkipLayerNormalization output
+        if add_node.input[1] not in skip_layer_norm.output:
+            return
+
+        bias_index, bias_value = self.model.get_constant_input(bias_node)
+        if not (isinstance(bias_index, int) and (bias_value is not None) and isinstance(bias_value, ndarray)):
+            return
+        if bias_value.ndim != 1:
+            return
+
+        self.nodes_to_remove.extend([add_node, bias_node])
+        node_name = self.model.create_node_name("BiasAdd")
+        fused_node = helper.make_node(
+            "BiasAdd",
+            inputs=[bias_node.input[1 - bias_index], bias_node.input[bias_index], add_node.input[1]],
+            outputs=[add_node.output[0]],
+            name=node_name,
+        )
+        fused_node.domain = "com.microsoft"
+        self.nodes_to_add.append(fused_node)
+        self.node_name_to_graph_name[node_name] = self.this_graph_name

onnxruntime/transformers/fusion_biasgelu.py

@@ -0,0 +1,66 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation.  All rights reserved.
+# Licensed under the MIT License.
+# --------------------------------------------------------------------------
+
+from logging import getLogger
+
+from fusion_base import Fusion
+from fusion_utils import NumpyHelper
+from onnx import helper
+from onnx_model import OnnxModel
+
+logger = getLogger(__name__)
+
+
+class FusionBiasGelu(Fusion):
+    def __init__(self, model: OnnxModel, is_fastgelu):
+        if is_fastgelu:
+            super().__init__(model, "FastGelu", "FastGelu", "add bias")
+        else:
+            super().__init__(model, "BiasGelu", "Gelu")
+
+    def fuse(self, node, input_name_to_nodes, output_name_to_node):
+        gelu_op_type = node.op_type
+        fuse_op_type = "BiasGelu" if gelu_op_type == "Gelu" else "FastGelu"
+
+        if len(node.input) != 1:
+            return
+
+        nodes = self.model.match_parent_path(node, ["Add", "MatMul"], [0, None])
+        if nodes is None:
+            return
+        (add, matmul) = nodes
+
+        bias_weight = None
+        # bias should be one dimension
+        bias_index = -1
+        for i, input in enumerate(add.input):
+            initializer = self.model.get_initializer(input)
+            if initializer is None:
+                continue
+            bias_index = i
+            bias_weight = NumpyHelper.to_array(initializer)
+            break
+        if bias_weight is None:
+            return
+        if len(bias_weight.shape) != 1:
+            return
+
+        subgraph_nodes = [node, add]
+        if not self.model.is_safe_to_fuse_nodes(
+            subgraph_nodes, [node.output[0]], input_name_to_nodes, output_name_to_node
+        ):
+            return
+
+        self.nodes_to_remove.extend(subgraph_nodes)
+
+        fused_node = helper.make_node(
+            fuse_op_type,
+            inputs=[matmul.output[0], add.input[bias_index]],
+            outputs=node.output,
+            name=self.model.create_node_name(fuse_op_type, gelu_op_type + "_AddBias_"),
+        )
+        fused_node.domain = "com.microsoft"
+        self.nodes_to_add.append(fused_node)
+        self.node_name_to_graph_name[fused_node.name] = self.this_graph_name

onnxruntime/transformers/fusion_biassplitgelu.py

@@ -0,0 +1,111 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation.  All rights reserved.
+# Licensed under the MIT License.
+# --------------------------------------------------------------------------
+from logging import getLogger
+from typing import Dict
+
+from fusion_base import Fusion
+from onnx import helper
+from onnx_model import OnnxModel
+
+logger = getLogger(__name__)
+
+
+class FusionBiasSplitGelu(Fusion):
+    def __init__(self, model: OnnxModel):
+        super().__init__(model, "BiasSplitGelu", "Gelu")
+
+    def fuse(self, gelu_node, input_name_to_nodes: Dict, output_name_to_node: Dict):
+        """
+        [root] --->Add -------------------->  Slice ---------------> Mul -->
+                   |                            ^                    ^
+                   |                            |                    |
+                   +----------------------------+---Slice --> Gelu---+
+                   |                            |     ^
+                   |                            |-----|
+                   |                            |     |
+                   |                           Mul   Mul
+                   |                            ^     ^
+                   v                            |     |
+                  Shape ---> Gather --> Add --> Div --+
+        """
+        if gelu_node.output[0] not in input_name_to_nodes:
+            return
+        children = input_name_to_nodes[gelu_node.output[0]]
+        if len(children) != 1 or children[0].op_type != "Mul":
+            return
+        mul_after_gelu = children[0]
+
+        slice_before_gelu = self.model.match_parent(gelu_node, "Slice", 0, output_name_to_node)
+        if slice_before_gelu is None:
+            return
+
+        if self.model.find_constant_input(slice_before_gelu, -1, delta=0.001) != 3:
+            return
+
+        add_output = slice_before_gelu.input[0]
+
+        start_index_nodes = self.model.match_parent_path(
+            slice_before_gelu,
+            ["Div", "Add", "Gather", "Shape", "Add"],
+            [1, 0, 0, 0, 0],
+            output_name_to_node,  # Mul(1) is optional
+        )
+        if start_index_nodes is None:
+            start_index_nodes = self.model.match_parent_path(
+                slice_before_gelu,
+                ["Mul", "Div", "Add", "Gather", "Shape", "Add"],
+                [1, 0, 0, 0, 0, 0],
+                output_name_to_node,
+            )
+
+        if start_index_nodes is None or start_index_nodes[-2].input[0] != add_output:
+            return
+
+        end_index_nodes = self.model.match_parent_path(slice_before_gelu, ["Mul", "Div"], [2, 0], output_name_to_node)
+
+        if (
+            end_index_nodes is None or end_index_nodes[1] not in start_index_nodes
+        ):  # the Div is parent of both two Mul nodes
+            return
+
+        slice_before_mul = self.model.match_parent(mul_after_gelu, "Slice", 0, output_name_to_node)
+        if slice_before_mul is None:
+            return
+
+        if (
+            slice_before_mul.input[2] != slice_before_gelu.input[1]
+        ):  # end index of slice_before_mul is start index of slice_before_gelu
+            return
+
+        subgraph_nodes = [
+            *start_index_nodes,
+            end_index_nodes[0],
+            mul_after_gelu,
+            gelu_node,
+            slice_before_mul,
+            slice_before_gelu,
+        ]
+        subgraph_output = mul_after_gelu.output[0]
+        if not self.model.is_safe_to_fuse_nodes(
+            subgraph_nodes, [subgraph_output], input_name_to_nodes, output_name_to_node
+        ):
+            logger.info("Skip fuse BiasSplitGelu since it is not safe to fuse the subgraph.")
+            return
+
+        add_node = start_index_nodes[-1]
+        bias_index, _value = self.model.get_constant_input(add_node)
+        if not isinstance(bias_index, int):
+            return
+        self.nodes_to_remove.extend(subgraph_nodes)
+        node_name = self.model.create_node_name("BiasSplitGelu", name_prefix="BiasSplitGelu")
+        fused_node = helper.make_node(
+            "BiasSplitGelu",
+            inputs=[add_node.input[1 - bias_index], add_node.input[bias_index]],
+            outputs=[subgraph_output],
+            name=node_name,
+        )
+        fused_node.domain = "com.microsoft"
+        self.nodes_to_add.append(fused_node)
+        self.node_name_to_graph_name[node_name] = self.this_graph_name

onnxruntime/transformers/fusion_conformer_attention.py

@@ -0,0 +1,143 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation.  All rights reserved.
+# Licensed under the MIT License.
+# --------------------------------------------------------------------------
+import logging
+
+from fusion_attention import AttentionMask, FusionAttention
+from onnx_model import OnnxModel
+
+logger = logging.getLogger(__name__)
+
+
+class FusionConformerAttention(FusionAttention):
+    """
+    Fuse Conformer Attention subgraph into one MultiHeadAttention node.
+    """
+
+    def __init__(
+        self,
+        model: OnnxModel,
+        hidden_size: int,
+        num_heads: int,
+        attention_mask: AttentionMask,
+    ):
+        super().__init__(model, hidden_size, num_heads, attention_mask)
+
+    def fuse(self, normalize_node, input_name_to_nodes, output_name_to_node):
+        # SkipLayerNormalization has two inputs, and one of them is the root input for attention.
+        qkv_nodes = self.model.match_parent_path(
+            normalize_node,
+            ["Add", "MatMul", "Reshape", "Transpose", "MatMul"],
+            [1, 1, 0, 0, 0],
+        )
+        if qkv_nodes is not None:
+            (
+                _,
+                _,
+                reshape_qkv,
+                transpose_qkv,
+                matmul_qkv,
+            ) = qkv_nodes
+        else:
+            logger.debug("fuse_conformer_attention: failed to match qkv path")
+            return
+
+        v_nodes = self.model.match_parent_path(
+            matmul_qkv,
+            ["Concat", "Transpose", "Reshape", "Add", "MatMul"],
+            [1, 1, 0, 0, 1],
+        )
+
+        add_v = None
+        if v_nodes is not None:
+            (concat_v, _, _, add_v, matmul_v) = v_nodes
+            concat_parent = self.model.get_parent(concat_v, 0, None)
+            present_v = concat_v.output[0]
+            past_v = concat_parent.output[0]
+        else:
+            logger.debug("fuse_conformer_attention: failed to match v path")
+            return
+
+        qk_nodes = self.model.match_parent_path(matmul_qkv, ["Softmax", "Add", "MatMul"], [0, 0, 0])
+
+        if qk_nodes is not None:
+            _, add_qk, matmul_qk = qk_nodes
+        else:
+            logger.debug("fuse_conformer_attention: failed to match qk path")
+            return
+
+        q_nodes = self.model.match_parent_path(
+            matmul_qk,
+            ["Div", "Transpose", "Reshape", "Add", "MatMul"],
+            [0, 0, 0, 0, 1],
+        )
+        if q_nodes is not None:
+            _, _, reshape_q, add_q, matmul_q = q_nodes
+        else:
+            logger.debug("fuse_conformer_attention: failed to match q path")
+            return
+
+        k_nodes = self.model.match_parent_path(
+            matmul_qk,
+            ["Transpose", "Concat", "Transpose", "Reshape", "Add", "MatMul"],
+            [1, 0, 1, 0, 0, 1],
+        )
+
+        matmul_k = None
+        if k_nodes is not None:
+            _, concat_k, _, _, add_k, matmul_k = k_nodes
+            concat_parent = self.model.get_parent(concat_k, 0, None)
+            past_k = concat_parent.output[0]
+            present_k = concat_k.output[0]
+        else:
+            logger.debug("fuse_conformer_attention: failed to match k path")
+            return
+
+        attention_last_node = reshape_qkv
+        num_heads, hidden_size = self.get_num_heads_and_hidden_size(reshape_q)
+
+        if num_heads <= 0 or hidden_size <= 0 or (hidden_size % num_heads) != 0:
+            logger.debug("fuse_conformer_attention: failed to detect num_heads or hidden_size")
+            return
+
+        new_node = self.create_multihead_attention_node(
+            matmul_q,
+            matmul_k,
+            matmul_v,
+            add_q,
+            add_k,
+            add_v,
+            num_heads,
+            hidden_size,
+            attention_last_node.output[0],
+            add_qk=add_qk.input[1],
+            past_k=past_k,
+            past_v=past_v,
+            present_k=present_k,
+            present_v=present_v,
+        )
+
+        if new_node is None:
+            logger.debug("fuse_conformer_attention: MultiHeadAttention node creation failed")
+            return
+
+        self.nodes_to_add.append(new_node)
+        self.node_name_to_graph_name[new_node.name] = self.this_graph_name
+
+        self.nodes_to_remove.extend([attention_last_node, transpose_qkv, matmul_qkv])
+        self.nodes_to_remove.extend(qk_nodes)
+
+        # When using multihead attention, keep MatMul nodes in original graph
+        if q_nodes[-1].op_type == "MatMul":
+            q_nodes.pop()
+        if k_nodes[-1].op_type == "MatMul":
+            k_nodes.pop()
+        if v_nodes[-1].op_type == "MatMul":
+            v_nodes.pop()
+
+        self.nodes_to_remove.extend(k_nodes)
+        self.nodes_to_remove.extend(v_nodes)
+
+        # Use prune graph to remove mask nodes since they are shared by all attention nodes.
+        self.prune_graph = True