onnxruntime-directml 1.24.1__cp314-cp314-win_amd64.whl

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,110 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation.  All rights reserved.
+# Licensed under the MIT License.
+# --------------------------------------------------------------------------
+from logging import getLogger
+
+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,222 @@
+# -------------------------------------------------------------------------
+# 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, None, 0, 0, 0],
+        )
+        if qkv_nodes is None:
+            logger.debug("fuse_conformer_attention: failed to match qkv path")
+            return
+
+        reshape_qkv, transpose_qkv, matmul_qkv = qkv_nodes[-3], qkv_nodes[-2], qkv_nodes[-1]
+
+        past_v, present_v = "", ""
+        v_nodes = self.model.match_parent_path(
+            matmul_qkv,
+            ["Concat", "Transpose", "Reshape", "Add", "MatMul"],
+            [1, 1, 0, 0, 1],
+        )
+        if v_nodes is None:
+            v_nodes = self.model.match_parent_path(
+                matmul_qkv,
+                ["Transpose", "Reshape", "Add", "MatMul"],
+                [1, 0, 0, 0],
+            )
+            if v_nodes is None:
+                logger.debug("fuse_conformer_attention: failed to match v path")
+                return
+        else:
+            concat_v = v_nodes[0]
+            concat_parent = self.model.get_parent(concat_v, 0, None)
+            present_v = concat_v.output[0]
+            past_v = concat_parent.output[0]
+
+        add_v, matmul_v = v_nodes[-2], v_nodes[-1]
+
+        attn_mask = ""
+        qk_nodes = self.model.match_parent_path(
+            matmul_qkv,
+            ["Softmax", "Add", "MatMul"],
+            [0, 0, 0],
+        )
+        if qk_nodes is None:
+            qk_nodes = self.model.match_parent_path(
+                matmul_qkv,
+                ["Where", "Softmax", "Where", "Add", "MatMul"],
+                [0, 2, 0, 2, 0],
+            )
+            if qk_nodes is None:
+                logger.debug("fuse_conformer_attention: failed to match qk path")
+                return
+
+            where_qk = qk_nodes[2]
+            mask_nodes = self.model.match_parent_path(
+                where_qk,
+                ["Equal", "Unsqueeze", "Cast"],
+                [0, 0, 0],
+            )
+            if mask_nodes is not None:
+                attn_mask = mask_nodes[-1].output[0]
+
+        add_qk, matmul_qk = qk_nodes[-2], qk_nodes[-1]
+
+        q_nodes = self.model.match_parent_path(
+            matmul_qk,
+            ["Div", "Transpose", "Reshape", "Add", "MatMul"],
+            [0, 0, 0, 0, 1],
+        )
+        if q_nodes is None:
+            q_nodes = self.model.match_parent_path(
+                matmul_qk,
+                ["Mul", "Transpose", "Reshape", "Add", "MatMul"],
+                [0, 0, 0, 0, 0],
+            )
+            if q_nodes is None:
+                logger.debug("fuse_conformer_attention: failed to match q path")
+                return
+
+        reshape_q, add_q, matmul_q = q_nodes[-3], q_nodes[-2], q_nodes[-1]
+
+        extra_q_nodes = self.model.match_parent_path(
+            add_qk,
+            ["Reshape", "Transpose", "MatMul", "Transpose", "Reshape", "Div"],
+            [1, 0, 0, 0, 0, 0],
+        )
+        if extra_q_nodes is not None and q_nodes[0] != extra_q_nodes[-1]:
+            logger.debug("fuse_conformer_attention: failed to match extra q path")
+            return
+
+        past_k, present_k = "", ""
+        k_nodes = self.model.match_parent_path(
+            matmul_qk,
+            ["Transpose", "Concat", "Transpose", "Reshape", "Add", "MatMul"],
+            [1, 0, 1, 0, 0, 1],
+        )
+        if k_nodes is None:
+            k_nodes = self.model.match_parent_path(
+                matmul_qk,
+                ["Transpose", "Transpose", "Reshape", "Add", "MatMul"],
+                [1, 0, 0, 0, 0],
+            )
+            if k_nodes is None:
+                k_nodes = self.model.match_parent_path(
+                    matmul_qk,
+                    ["Transpose", "Reshape", "Add", "MatMul"],
+                    [1, 0, 0, 0],
+                )
+                if k_nodes is None:
+                    logger.debug("fuse_conformer_attention: failed to match k path")
+                    return
+        else:
+            concat_k = k_nodes[1]
+            concat_parent = self.model.get_parent(concat_k, 0, None)
+            past_k = concat_parent.output[0]
+            present_k = concat_k.output[0]
+
+        add_k, matmul_k = k_nodes[-2], k_nodes[-1]
+
+        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 = None
+        use_packed_attention_op = (
+            matmul_q.input[0] == matmul_k.input[0] and matmul_k.input[0] == matmul_v.input[0] and extra_q_nodes is None
+        )
+        if use_packed_attention_op:
+            # Self-attention, use Attention op
+            new_node = self.create_attention_node(
+                mask_index=attn_mask,
+                q_matmul=matmul_q,
+                k_matmul=matmul_k,
+                v_matmul=matmul_v,
+                q_add=add_q,
+                k_add=add_k,
+                v_add=add_v,
+                num_heads=num_heads,
+                hidden_size=hidden_size,
+                first_input=matmul_q.input[0],
+                output=reshape_qkv.output[0],
+                add_qk_str=add_qk.input[1],
+                past_k=past_k,
+                past_v=past_v,
+                present_k=present_k,
+                present_v=present_v,
+            )
+        else:
+            new_node = self.create_multihead_attention_node(
+                q_matmul=matmul_q,
+                k_matmul=matmul_k,
+                v_matmul=matmul_v,
+                q_add=add_q,
+                k_add=add_k,
+                v_add=add_v,
+                num_heads=num_heads,
+                hidden_size=hidden_size,
+                output=reshape_qkv.output[0],
+                key_padding_mask=attn_mask,
+                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([reshape_qkv, transpose_qkv, matmul_qkv])
+        self.nodes_to_remove.extend(qk_nodes)
+
+        # When using MultiHeadAttention, keep MatMul nodes unfused in original graph
+        if not use_packed_attention_op:
+            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()
+
+        if extra_q_nodes is None:
+            # Don't remove Q nodes for conformer-transducer (CT) model since it has
+            # an extra set of nodes attached to the output of the Q path that are not
+            # part of the attention computation
+            self.nodes_to_remove.extend(q_nodes)
+
+        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

onnxruntime/transformers/fusion_constant_fold.py

@@ -0,0 +1,144 @@
+# -------------------------------------------------------------------------
+# 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_model import OnnxModel
+
+logger = getLogger(__name__)
+
+
+class FusionConstantFold(Fusion):
+    def __init__(self, model: OnnxModel):
+        super().__init__(model, "", ["Transpose"])
+        self.count = 0
+
+    def apply(self):
+        super().apply()
+        if self.count > 0:
+            logger.info(f"Constant Folded: {self.count}")
+
+    def fuse(self, node, input_name_to_nodes, output_name_to_node):
+        """
+        Apply multiple fusions on Transpose nodes that can be constant folded.
+        """
+        self.fuse_1(node, input_name_to_nodes, output_name_to_node)
+        self.fuse_2(node, input_name_to_nodes, output_name_to_node)
+
+    def fuse_1(self, node, input_name_to_nodes, output_name_to_node):
+        """
+        Constant fold any initializer data representing a MatMul's
+        weights that are stored in a Transpose op
+
+        Ex: Transpose --> Gemm or Transpose --> MatMul
+        """
+        # Check if Transpose node only has one input and one output
+        if len(node.input) != 1 or len(node.output) != 1:
+            logger.debug("fuse_constant_fold: node has more than one input or output")
+            return
+
+        # Check if input is initializer data
+        proto = self.model.get_initializer(node.input[0])
+        if proto is None:
+            logger.debug("fuse_constant_fold: failed to identify initializer input")
+            return
+
+        # Check that all nodes using input are Transpose ops that also only use the initializer data as input
+        skip = False
+        for child_node in input_name_to_nodes[node.input[0]]:
+            if not (child_node.op_type == "Transpose" and len(node.input) == 1):
+                skip = True
+                break
+        if skip:
+            logger.debug("fuse_constant_fold: other non-Transpose nodes use the initializer")
+            return
+
+        # Check that all nodes using output are Gemm or MatMul ops
+        for child_node in input_name_to_nodes[node.output[0]]:
+            if not (child_node.op_type == "Gemm" or child_node.op_type == "MatMul"):
+                skip = True
+                break
+        if skip:
+            logger.debug("fuse_constant_fold: other non-Gemm and non-MatMul nodes use the transposed data")
+            return
+
+        # Check if initializer data is 2D
+        weight = NumpyHelper.to_array(proto)
+        if len(weight.shape) != 2:
+            logger.debug("fuse_constant_fold: shape of initializer data is not 2D")
+            return
+
+        # Remove old TensorProto and add new TensorProto while re-using same name
+        name = proto.name
+        dtype = proto.data_type
+        self.remove_initializer(proto)
+        self.add_initializer(
+            name=name,
+            data_type=dtype,
+            dims=[weight.shape[1], weight.shape[0]],
+            vals=weight.T,
+        )
+
+        # Update weights input to be the initializer name and not
+        # the output of the Transpose op
+        for child_node in input_name_to_nodes[node.output[0]]:
+            for i in range(len(child_node.input)):
+                if child_node.input[i] == node.output[0]:
+                    child_node.input[i] = node.input[0]
+
+                    if child_node.op_type == "Gemm" and (i == 0 or i == 1):
+                        # Ensure that transA/transB is set to 0 in Gemm
+                        key = "transA" if i == 0 else "transB"
+                        for j, attr_key in enumerate(child_node.attribute):
+                            if attr_key.name == key:
+                                child_node.attribute[j].i = 0
+
+        # Add node to list of nodes to remove
+        self.nodes_to_remove.append(node)
+        self.count += 1
+
+    def fuse_2(self, node, input_name_to_nodes, output_name_to_node):
+        """
+        Constant fold any Transpose --> Transpose ops since the root input
+        is the final result
+
+        Ex: root_input --> Transpose --> Transpose --> next_node to root_input --> next_node
+        """
+        # Check if Transpose node only has one input and one output
+        if len(node.input) != 1 or len(node.output) != 1:
+            logger.debug("fuse_constant_fold: node has more than one input or output")
+            return
+
+        # Check if parent node is Transpose node with only one input and one output
+        parent_node = self.model.match_parent(node, "Transpose", 0)
+        if parent_node is None:
+            logger.debug("fuse_constant_fold: failed to identify parent Transpose node")
+            return
+        if len(parent_node.input) != 1 or len(parent_node.output) != 1:
+            logger.debug("fuse_constant_fold: parent node has more than one input or output")
+            return
+
+        node_perm = node.attribute[0].ints
+        parent_node_perm = parent_node.attribute[0].ints
+
+        if node_perm != parent_node_perm:
+            logger.debug("fuse_constant_fold: Transpose node permutations aren't identical")
+            return
+
+        # For nodes that use output of child Transpose node as an input,
+        # replace that input with root_input
+        root_input = parent_node.input[0]
+        output_nodes = input_name_to_nodes[node.output[0]]
+        for output_node in output_nodes:
+            for i, input_ in enumerate(output_node.input):
+                if input_ == node.output[0]:
+                    output_node.input[i] = root_input
+
+        # Add node to list of nodes to remove
+        self.nodes_to_remove.append(node)
+        self.nodes_to_remove.append(parent_node)
+        self.count += 1