onnxruntime-directml 1.24.1__cp314-cp314-win_amd64.whl

onnxruntime/transformers/fusion_gelu.py

@@ -0,0 +1,258 @@
+# -------------------------------------------------------------------------
+# 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 FusionGelu(Fusion):
+    def __init__(self, model: OnnxModel):
+        super().__init__(model, "Gelu", "Erf")
+
+    def fuse(self, erf_node, input_name_to_nodes: dict, output_name_to_node: dict):
+        if self.fuse_1(erf_node, input_name_to_nodes, output_name_to_node):
+            return
+        if self.fuse_2(erf_node, input_name_to_nodes, output_name_to_node):
+            return
+        self.fuse_3(erf_node, input_name_to_nodes, output_name_to_node)
+
+    def fuse_1(self, erf_node, input_name_to_nodes: dict, output_name_to_node: dict) -> bool | None:
+        """
+        This pattern is from PyTorch model
+        Fuse Gelu with Erf into one node:
+        Pattern 1:
+                       +-------Mul(0.5)---------------------+
+                       |                                    |
+                       |                                    v
+                    [root] --> Div -----> Erf  --> Add --> Mul -->
+                              (B=1.4142...)       (1)
+
+        Pattern 2:
+                       +------------------------------------+
+                       |                                    |
+                       |                                    v
+                    [root] --> Div -----> Erf  --> Add --> Mul -->Mul -->
+                              (B=1.4142...)       (1)            (0.5)
+
+        Note that constant input for Add and Mul could be first or second input: like either A=0.5 or B=0.5 is fine.
+        """
+        if erf_node.output[0] not in input_name_to_nodes:
+            return
+        children = input_name_to_nodes[erf_node.output[0]]
+        if len(children) != 1 or children[0].op_type != "Add":
+            return
+        add_after_erf = children[0]
+
+        if not self.model.has_constant_input(add_after_erf, 1):
+            return
+
+        if add_after_erf.output[0] not in input_name_to_nodes:
+            return
+        children = input_name_to_nodes[add_after_erf.output[0]]
+        if len(children) != 1 or children[0].op_type != "Mul":
+            return
+        mul_after_erf = children[0]
+
+        div = self.model.match_parent(erf_node, "Div", 0, output_name_to_node)
+        if div is None:
+            return
+
+        if self.model.find_constant_input(div, 1.4142, delta=0.001) != 1:
+            return
+
+        subgraph_input = div.input[0]
+
+        another = 1 if mul_after_erf.input[0] == add_after_erf.output[0] else 0
+        if subgraph_input == mul_after_erf.input[another]:  # pattern 2
+            children = input_name_to_nodes[mul_after_erf.output[0]]
+            if len(children) != 1 or children[0].op_type != "Mul":
+                return
+            mul_half = children[0]
+            if not self.model.has_constant_input(mul_half, 0.5):
+                return
+            subgraph_output = mul_half.output[0]
+        else:  # pattern 1
+            mul_half = self.model.match_parent(mul_after_erf, "Mul", another, output_name_to_node)
+            if mul_half is None:
+                return
+
+            if not self.model.has_constant_input(mul_half, 0.5):
+                return
+
+            if subgraph_input not in mul_half.input:
+                return
+
+            subgraph_output = mul_after_erf.output[0]
+
+        subgraph_nodes = [div, erf_node, add_after_erf, mul_after_erf, mul_half]
+        if not self.model.is_safe_to_fuse_nodes(
+            subgraph_nodes, [subgraph_output], input_name_to_nodes, output_name_to_node
+        ):
+            return
+
+        self.nodes_to_remove.extend(subgraph_nodes)
+        fused_node = helper.make_node(
+            "Gelu", inputs=[subgraph_input], outputs=[subgraph_output], name=self.model.create_node_name("Gelu")
+        )
+        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
+        self.increase_counter("Gelu")
+        return True
+
+    def fuse_2(self, erf_node, input_name_to_nodes: dict, output_name_to_node: dict) -> bool | None:
+        """
+        This pattern is from Keras model
+        Fuse Gelu with Erf into one node:
+                       +------------------------------------------+
+                       |                                          |
+                       |                                          v
+                    [root] --> Div -----> Erf  --> Add --> Mul -->Mul
+                              (B=1.4142...)       (A=1)   (A=0.5)
+
+        Note that constant input for Add and Mul could be first or second input: like either A=0.5 or B=0.5 is fine.
+        """
+        if erf_node.output[0] not in input_name_to_nodes:
+            return
+        children = input_name_to_nodes[erf_node.output[0]]
+        if len(children) != 1 or children[0].op_type != "Add":
+            return
+        add_after_erf = children[0]
+
+        if not self.model.has_constant_input(add_after_erf, 1):
+            return
+
+        if add_after_erf.output[0] not in input_name_to_nodes:
+            return
+        children = input_name_to_nodes[add_after_erf.output[0]]
+        if len(children) != 1 or children[0].op_type != "Mul":
+            return
+        mul_after_erf = children[0]
+
+        if not self.model.has_constant_input(mul_after_erf, 0.5):
+            return
+
+        if mul_after_erf.output[0] not in input_name_to_nodes:
+            return
+        children = input_name_to_nodes[mul_after_erf.output[0]]
+        if len(children) != 1 or children[0].op_type != "Mul":
+            return
+        mul = children[0]
+
+        div = self.model.match_parent(erf_node, "Div", 0, output_name_to_node)
+        if div is None:
+            return
+
+        sqrt_node = None
+        if self.model.find_constant_input(div, 1.4142, delta=0.001) != 1:
+            sqrt_node = self.model.match_parent(div, "Sqrt", 1, output_name_to_node)
+            if sqrt_node is None:
+                return
+            if not self.model.has_constant_input(sqrt_node, 2.0):
+                return
+
+        root_node = self.model.get_parent(div, 0, output_name_to_node)
+        if root_node is None:
+            return
+
+        if root_node.output[0] not in mul.input:
+            return
+
+        subgraph_nodes = [div, erf_node, add_after_erf, mul_after_erf, mul]
+        if sqrt_node:
+            subgraph_nodes.append(sqrt_node)
+
+        if not self.model.is_safe_to_fuse_nodes(
+            subgraph_nodes, [mul.output[0]], input_name_to_nodes, output_name_to_node
+        ):
+            return
+
+        self.nodes_to_remove.extend(subgraph_nodes)
+        fused_node = helper.make_node(
+            "Gelu", inputs=[root_node.output[0]], outputs=[mul.output[0]], name=self.model.create_node_name("Gelu")
+        )
+        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
+        self.increase_counter("Gelu")
+        return True
+
+    def fuse_3(self, erf_node, input_name_to_nodes: dict, output_name_to_node: dict) -> bool | None:
+        """
+        This pattern is from TensorFlow model
+        Fuse Gelu with Erf into one node:
+                       +----------------------------------------------+
+                       |                                              |
+                       |                                              v
+                    [root] --> Mul -----> Erf    -->   Add --> Mul -->Mul
+                               (A=0.7071067690849304)  (B=1)  (B=0.5)
+
+        Note that constant input for Add and Mul could be first or second input: like either A=0.5 or B=0.5 is fine.
+        """
+
+        if erf_node.output[0] not in input_name_to_nodes:
+            return
+        children = input_name_to_nodes[erf_node.output[0]]
+        if len(children) != 1 or children[0].op_type != "Add":
+            return
+        add_after_erf = children[0]
+
+        if not self.model.has_constant_input(add_after_erf, 1):
+            return
+
+        if add_after_erf.output[0] not in input_name_to_nodes:
+            return
+        children = input_name_to_nodes[add_after_erf.output[0]]
+        if len(children) != 1 or children[0].op_type != "Mul":
+            return
+        mul_half = children[0]
+
+        if not self.model.has_constant_input(mul_half, 0.5):
+            return
+
+        first_mul = self.model.match_parent(erf_node, "Mul", 0, output_name_to_node)
+        if first_mul is None:
+            return
+
+        i = self.model.find_constant_input(first_mul, 0.7071067690849304, delta=0.001)
+        if i < 0:
+            return
+
+        root_node = self.model.get_parent(first_mul, 0 if i == 1 else 1, output_name_to_node)
+        if root_node is None:
+            return
+
+        if mul_half.output[0] not in input_name_to_nodes:
+            return
+        children = input_name_to_nodes[mul_half.output[0]]
+        if len(children) != 1 or children[0].op_type != "Mul":
+            return
+        last_mul = children[0]
+
+        if not (last_mul.input[0] == root_node.output[0] or last_mul.input[1] == root_node.output[0]):
+            return
+
+        subgraph_nodes = [first_mul, erf_node, add_after_erf, mul_half, last_mul]
+        if not self.model.is_safe_to_fuse_nodes(
+            subgraph_nodes,
+            [last_mul.output[0]],
+            input_name_to_nodes,
+            output_name_to_node,
+        ):
+            return
+
+        self.nodes_to_remove.extend(subgraph_nodes)
+        fused_node = helper.make_node(
+            "Gelu", inputs=[root_node.output[0]], outputs=[last_mul.output[0]], name=self.model.create_node_name("Gelu")
+        )
+        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
+        self.increase_counter("Gelu")
+        return True

onnxruntime/transformers/fusion_gelu_approximation.py

@@ -0,0 +1,25 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation.  All rights reserved.
+# Licensed under the MIT License.
+# --------------------------------------------------------------------------
+
+from fusion_base import Fusion
+from onnx import helper
+from onnx_model import OnnxModel
+
+
+class FusionGeluApproximation(Fusion):
+    def __init__(self, model: OnnxModel):
+        super().__init__(model, "FastGelu", ["Gelu", "BiasGelu"], "GeluApproximation")
+
+    def fuse(self, node, input_name_to_nodes, output_name_to_node):
+        new_node = helper.make_node(
+            "FastGelu",
+            inputs=node.input,
+            outputs=node.output,
+            name=self.model.create_node_name("FastGelu", node.op_type + "_Approximation"),
+        )
+        new_node.domain = "com.microsoft"
+        self.nodes_to_remove.append(node)
+        self.nodes_to_add.append(new_node)
+        self.node_name_to_graph_name[new_node.name] = self.this_graph_name

onnxruntime/transformers/fusion_gemmfastgelu.py

@@ -0,0 +1,121 @@
+# -------------------------------------------------------------------------
+# 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 NodeProto, TensorProto, helper
+from onnx_model import OnnxModel
+
+logger = getLogger(__name__)
+
+
+class FusionGemmFastGelu(Fusion):
+    def __init__(self, model: OnnxModel):
+        super().__init__(model, "GemmFastGelu", "FastGelu", "GemmFastGelu")
+        self.shape_infer = None
+        self.shape_infer_done = False
+
+    def get_dimensions_from_tensor_proto(self, tensor_proto: TensorProto) -> int | None:
+        if tensor_proto.type.tensor_type.HasField("shape"):
+            return len(tensor_proto.type.tensor_type.shape.dim)
+        else:
+            return None
+
+    def get_dimensions(self, input_name: str) -> int | None:
+        graph_input = self.model.find_graph_input(input_name)
+        if graph_input:
+            return self.get_dimensions_from_tensor_proto(graph_input)
+
+        if not self.shape_infer_done:
+            self.shape_infer = self.model.infer_runtime_shape(update=True)
+            self.shape_infer_done = True
+
+        if self.shape_infer is not None:
+            return self.get_dimensions_from_tensor_proto(self.shape_infer.known_vi_[input_name])
+
+        return None
+
+    def fuse(
+        self,
+        node: NodeProto,
+        input_name_to_nodes: dict[str, list[NodeProto]],
+        output_name_to_node: dict[str, NodeProto],
+    ):
+        """
+        This pattern is from PyTorch bert model
+        Fuse MatMul with FastGelu into one node:
+
+            [root] --> MatMul --> FastGelu -->
+
+        """
+        has_bias = False
+        if len(node.input) == 2:
+            has_bias = True
+
+        match_nodes = self.model.match_parent_path(node, ["MatMul"], [0])
+        if match_nodes is None:
+            return
+        matmul = match_nodes[0]
+
+        # matmul input X should >= two dimension, input weight should be two dimension
+        weight_index = -1
+        x_dims = 0
+        weight = None
+
+        for i, input in enumerate(matmul.input):
+            initializer = self.model.get_initializer(input)
+            if initializer is None:
+                x_dims = self.get_dimensions(matmul.input[i])
+            else:
+                weight_index = i
+                weight = NumpyHelper.to_array(initializer)
+        if weight is None:
+            return
+        if len(weight.shape) != 2:
+            return
+        if x_dims < len(weight.shape):
+            return
+
+        # bias weight should be one dimension
+        bias_index = -1
+        if has_bias:
+            bias_weight = None
+            for i, input in enumerate(node.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, matmul]
+        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)
+
+        inputs = (
+            [matmul.input[1 - weight_index], matmul.input[weight_index], node.input[bias_index]]
+            if has_bias
+            else [matmul.input[1 - weight_index], matmul.input[weight_index]]
+        )
+
+        fused_node = helper.make_node(
+            "GemmFastGelu",
+            inputs=inputs,
+            outputs=node.output,
+            name=self.model.create_node_name("GemmFastGelu"),
+        )
+        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