onnxruntime-directml 1.24.1__cp314-cp314-win_amd64.whl

onnxruntime/transformers/fusion_shape.py

@@ -0,0 +1,109 @@
+# -------------------------------------------------------------------------
+# 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 FusionUtils
+from numpy import ndarray
+from onnx import NodeProto, TensorProto
+from onnx_model import OnnxModel
+
+logger = getLogger(__name__)
+
+
+class FusionShape(Fusion):
+    def __init__(self, model: OnnxModel):
+        super().__init__(model, "Shape", "Concat")
+        self.utils = FusionUtils(model)
+        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:
+        shape = self.model.get_shape(input_name)
+        if shape is not None:
+            return len(shape)
+
+        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,
+        concat_node: NodeProto,
+        input_name_to_nodes: dict[str, list[NodeProto]],
+        output_name_to_node: dict[str, NodeProto],
+    ):
+        #
+        # Simplify subgraph like
+        #
+        #          (2d_input)
+        #           /       \
+        #       Shape       shape
+        #       /             \
+        #   Gather(indices=0)  Gather(indices=1)
+        #       |                |
+        #   Unsqueeze(axes=0)   Unsqueeze(axes=0)
+        #          \           /
+        #             Concat
+        #               |
+        #
+        # into  (2d_input) --> Shape -->
+        #
+        opset_version = self.model.get_opset_version()
+
+        inputs = len(concat_node.input)
+        root = None
+        shape_output = None
+        for i in range(inputs):
+            path = self.model.match_parent_path(
+                concat_node,
+                ["Unsqueeze", "Gather", "Shape"],
+                [i, 0, 0],
+                output_name_to_node,
+            )
+            if path is None:
+                return
+
+            unsqueeze, gather, shape = path
+            if i == 0:
+                shape_output = shape.output[0]
+            if root is None:
+                root = shape.input[0]
+                if self.get_dimensions(root) != inputs:
+                    return
+            elif shape.input[0] != root:
+                return
+
+            if not FusionUtils.check_node_attribute(unsqueeze, "axis", 0, default_value=0):
+                return
+
+            if opset_version < 13:
+                if not FusionUtils.check_node_attribute(unsqueeze, "axes", [0]):
+                    return
+            else:
+                if not self.utils.check_node_input_value(unsqueeze, 1, [0]):
+                    return
+
+            value = self.model.get_constant_value(gather.input[1])
+
+            if not (isinstance(value, ndarray) and value.size == 1 and value.item() == i):
+                return
+
+        if self.model.find_graph_output(concat_node.output[0]) is None:
+            self.model.replace_input_of_all_nodes(concat_node.output[0], shape_output)
+            self.increase_counter("Reshape")
+            self.prune_graph = True

onnxruntime/transformers/fusion_simplified_layernorm.py

@@ -0,0 +1,165 @@
+import logging
+
+from fusion_base import Fusion
+from fusion_skiplayernorm import FusionSkipLayerNormalization
+from onnx import helper
+from onnx_model import OnnxModel
+
+logger = logging.getLogger(__name__)
+
+
+class FusionSimplifiedLayerNormalization(Fusion):
+    def __init__(self, model: OnnxModel):
+        super().__init__(model, "SimplifiedLayerNormalization", "Mul")
+
+    def fuse(self, node, input_name_to_nodes: dict, output_name_to_node: dict):
+        if node.op_type != "Mul":
+            return
+
+        sim_ln_nodes = None
+        # RMSNorm formula:
+        #   S = Pow(X, 2) or S = Mul(X, X)
+        #   MS = ReduceMean(S)
+        #   MSEps = Add(MS, epsilon)
+        #   RMS = Sqrt(MSEps)
+        #   InvRMS = Div(1, RMS) or InvRMS = Reciprocal(RMS)
+        #   Normalized = Mul(D, InvRMS)
+        #   Y = Mul(Normalized, Scale)
+        #
+        #  (root_input) ----------------------------------------+
+        #       |                                               |
+        #       v                                               v
+        #      Pow --> ReduceMean --> Add ---> Sqrt --> Div --> Mul --> Mul (node)
+        #      (B=2)                  (A/B=eps)         (A=1)           (A/B=scale)
+        #
+        #  (root_input) ----------------------------------------+
+        #      | |                                              |
+        #      v v                                              v
+        #      Mul --> ReduceMean --> Add ---> Sqrt --> Div --> Mul --> Mul (node)
+        #      (B=2)                  (A/B=eps)         (A=1)           (A/B=scale)
+        #
+        return_indice = []
+        sim_ln_nodes = self.model.match_parent_path(
+            node,
+            ["Mul", "Div", "Sqrt", "Add", "ReduceMean"],
+            [None, 1, 1, 0, None],
+            output_name_to_node=output_name_to_node,
+            return_indice=return_indice,
+        )
+
+        if sim_ln_nodes:
+            mul_node, div_node, _sqrt_node, add_node, reduce_mean_node = sim_ln_nodes
+            if not self.model.has_constant_input(div_node, 1.0):
+                return
+            node_parent = mul_node
+        else:
+            # Div(1, RMS) can also be represented as Reciprocal(RMS) like
+            #
+            #  (root_input) -----------------------------------------------+
+            #       |                                                      |
+            #       v                                                      v
+            #      Pow --> ReduceMean --> Add ---> Sqrt --> Reciprocal --> Mul --> Mul (node)
+            #      (B=2)                  (A/B=eps)                                (A/B=scale)
+            #
+            #  (root_input) -----------------------------------------------+
+            #      | |                                                     |
+            #      v v                                                     v
+            #      Mul --> ReduceMean --> Add ---> Sqrt --> Reciprocal --> Mul --> Mul (node)
+            #      (B=2)                  (A/B=eps)                                (A/B=scale)
+            #
+            return_indice = []
+            sim_ln_nodes = self.model.match_parent_path(
+                node,
+                ["Mul", "Reciprocal", "Sqrt", "Add", "ReduceMean"],
+                [None, 1, 0, 0, None],
+                output_name_to_node=output_name_to_node,
+                return_indice=return_indice,
+            )
+            if sim_ln_nodes is not None:
+                mul_node, _reciprocal_node, _sqrt_node, add_node, reduce_mean_node = sim_ln_nodes
+                node_parent = mul_node
+            else:
+                #  (root_input) --------------------------------+
+                #       |                                       |
+                #       v                                       v
+                #      Pow --> ReduceMean --> Add ---> Sqrt --> Div --> Mul (node)
+                #      (B=2)                  (A/B=eps)                 (A/B=scale)
+                #
+                #  (root_input) --------------------------------+
+                #      | |                                      |
+                #      v v                                      v
+                #      Mul --> ReduceMean --> Add ---> Sqrt --> Div --> Mul (node)
+                #      (B=2)                  (A/B=eps)                 (A/B=scale)
+                #
+                return_indice = []
+                sim_ln_nodes = self.model.match_parent_path(
+                    node,
+                    ["Div", "Sqrt", "Add", "ReduceMean"],
+                    [None, 1, 0, None],
+                    output_name_to_node=output_name_to_node,
+                    return_indice=return_indice,
+                )
+                if sim_ln_nodes is not None:
+                    div_node, _sqrt_node, add_node, reduce_mean_node = sim_ln_nodes
+                    node_parent = div_node
+                else:
+                    return
+
+        reduce_mean_parent = self.model.get_parent(reduce_mean_node, 0, output_name_to_node)
+        if reduce_mean_parent is None or reduce_mean_parent.op_type not in ["Pow", "Mul"]:
+            return
+
+        if reduce_mean_parent.op_type == "Pow":
+            if self.model.find_constant_input(reduce_mean_parent, 2.0) != 1:
+                return
+        else:
+            assert reduce_mean_parent.op_type == "Mul"
+            if reduce_mean_parent[0] != reduce_mean_parent[1]:
+                return
+
+        root_input = reduce_mean_parent.input[0]
+        if root_input not in node_parent.input:
+            return
+
+        _i, epsilon = self.model.get_constant_input(add_node)
+        if epsilon is None or epsilon <= 0 or epsilon > 1.0e-4:
+            logger.warning(f"epsilon value is not expected: {epsilon}")
+            return
+
+        # ReduceMean must have keepdims == 1
+        keepdims = self.model.get_node_attribute(reduce_mean_node, "keepdims")
+        if not keepdims:
+            return
+
+        # ReduceMean axes must refer only to the last dimension.
+        # Axes became an input in opset 18. Before then, axes was an attribute.
+        axes = self.model.get_node_attribute(reduce_mean_node, "axes")
+        if (not axes) and len(reduce_mean_node.input) > 1:
+            axes = self.model.get_constant_value(reduce_mean_node.input[1])
+        # Make sure only one axis as required by SimplifiedLayerNormalization spec.
+        if not axes or len(axes) != 1:
+            return
+
+        self.nodes_to_remove.extend(sim_ln_nodes)
+        self.nodes_to_remove.append(reduce_mean_parent)
+        self.nodes_to_remove.append(node)
+
+        normalize_node = helper.make_node(
+            "SimplifiedLayerNormalization",
+            inputs=[root_input, node.input[1 - return_indice[0]]],
+            outputs=[node.output[0]],
+            name=self.model.create_node_name("SimplifiedLayerNormalization", name_prefix="RMSNorm"),
+        )
+        normalize_node.attribute.extend([helper.make_attribute("epsilon", float(epsilon))])
+        normalize_node.attribute.extend([helper.make_attribute("axis", axes[0])])
+        normalize_node.attribute.extend([helper.make_attribute("stash_type", 1)])
+        self.nodes_to_add.append(normalize_node)
+        self.node_name_to_graph_name[normalize_node.name] = self.this_graph_name
+
+
+class FusionSkipSimplifiedLayerNormalization(FusionSkipLayerNormalization):
+    def __init__(self, model: OnnxModel):
+        super().__init__(model, "SkipSimplifiedLayerNormalization", "SimplifiedLayerNormalization")
+
+    def fuse(self, node, input_name_to_nodes, output_name_to_node):
+        super().fuse(node, input_name_to_nodes, output_name_to_node)

onnxruntime/transformers/fusion_skip_group_norm.py

@@ -0,0 +1,254 @@
+# -------------------------------------------------------------------------
+# 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 FusionSkipGroupNorm(Fusion):
+    """
+    Fuse Add + GroupNorm into one node: SkipGroupNorm.
+    """
+
+    def __init__(self, model: OnnxModel):
+        super().__init__(model, "SkipGroupNorm", "GroupNorm")
+        # Update shape inference is needed since other fusions might add new edge which does not have shape info yet.
+        self.shape_infer_helper = self.model.infer_runtime_shape(update=True)
+
+        if self.shape_infer_helper is None:
+            logger.warning("SkipGroupNorm fusion will be skipped since symbolic shape inference disabled or failed.")
+
+    def create_transpose_node(self, input_name: str, perm: list[int], output_name=None):
+        """Append a Transpose node after an input"""
+        node_name = self.model.create_node_name("Transpose")
+        if output_name is None:
+            output_name = node_name + "_out" + "-" + input_name
+        transpose_node = helper.make_node("Transpose", inputs=[input_name], outputs=[output_name], name=node_name)
+        transpose_node.attribute.extend([helper.make_attribute("perm", perm)])
+        return transpose_node
+
+    def get_skip_index(self, add, is_channel_last: bool):
+        """Add has two inputs. This classifies which input is skip based on shape info (skip allows broadcast)."""
+        skip = -1
+        broadcast = False
+
+        assert self.shape_infer_helper is not None
+        shape_a = self.shape_infer_helper.get_edge_shape(add.input[0])
+        shape_b = self.shape_infer_helper.get_edge_shape(add.input[1])
+        assert shape_a is not None and shape_b is not None
+
+        if len(shape_a) == 4 and len(shape_b) == 4:
+            if shape_a == shape_b:
+                skip = 1
+            else:
+                c = 3 if is_channel_last else 1
+                h = 1 if is_channel_last else 2
+                w = 2 if is_channel_last else 3
+                if shape_a[0] == shape_b[0] and shape_a[c] == shape_b[c]:
+                    if shape_b[h] == 1 and shape_b[w] == 1:
+                        skip = 1
+                        broadcast = True
+                    elif shape_a[h] == 1 and shape_a[w] == 1:
+                        skip = 0
+                        broadcast = True
+
+        if skip < 0:
+            logger.debug(
+                "skip SkipGroupNorm fusion since shape of Add inputs (%s, %s) are not expected",
+                add.input[0],
+                add.input[1],
+            )
+        return skip, broadcast
+
+    def has_multiple_consumers(self, output_name, input_name_to_nodes):
+        """Whether an output has multiple consumers (like graph output or more than one children nodes)"""
+        return self.model.find_graph_output(output_name) is not None or (
+            output_name in input_name_to_nodes and len(input_name_to_nodes[output_name]) > 1
+        )
+
+    def remove_if_safe(self, node, input_name_to_nodes):
+        """Remove a node if it is safe (only one children, and not graph output)"""
+        if not self.has_multiple_consumers(node.output[0], input_name_to_nodes):
+            self.nodes_to_remove.extend([node])
+
+    def is_bias_1d(self, bias_name: str):
+        """Whether bias is an initializer of one dimension"""
+        initializer = self.model.get_initializer(bias_name)
+        if initializer is None:
+            return False
+
+        bias_weight = NumpyHelper.to_array(initializer)
+        if bias_weight is None:
+            logger.debug("Bias weight not found")
+            return False
+
+        if len(bias_weight.shape) != 1:
+            logger.debug("Bias weight is not 1D")
+            return False
+        return True
+
+    def match_bias_path(self, node, input_name_to_nodes, output_name_to_node):
+        """
+        Match the bias graph pattern from an Transpose node after Reshape node like in below example.
+        It checks whether the bias is 1D initializer. If so, remove Add and redirect MatMul output to Reshape.
+        """
+        # Before Fusion:
+        #                        MatMul  (bias)
+        #                            \  /     (shape)
+        #                             Add    /
+        #                               \   /
+        #       (a)                   Reshape
+        #        \                       |
+        # Transpose([0, 3, 1, 2])   Transpose([0, 3, 1, 2])  --- the start node, this func only handles the above nodes.
+        #                        \  /
+        #                         Add
+        #                         / \
+        #                      (c)  Transpose([0,2,3,1])
+        #                              |
+        #                           GroupNorm
+        #                              |
+        #                             (d)
+        #
+        # After Fusion (the nodes below Reshape is handled in the fuse function):
+        #                    MatMul (shape)
+        #                       \   /
+        #                (a)   Reshape
+        #                  \    /
+        #                SkipGroupNorm
+        #                  /    \
+        #                (d)   Transpose([0, 3, 1, 2])
+        #                         \
+        #                         (c)
+
+        add_input_index = []
+        bias_nodes = self.model.match_parent_path(
+            node, ["Reshape", "Add", "MatMul"], [0, 0, None], output_name_to_node, add_input_index
+        )
+        if bias_nodes is None:
+            return None
+
+        (reshape, add_bias, matmul) = bias_nodes
+        bias = bias_nodes[1].input[1 - add_input_index[0]]
+        if not self.is_bias_1d(bias):
+            return None
+
+        reshape.input[0] = matmul.output[0]
+        self.remove_if_safe(add_bias, input_name_to_nodes)
+
+        return bias
+
+    def match_transpose_from_nhwc(self, output_name, input_name_to_nodes, output_name_to_node):
+        """Match whether an output is from a Transpose(perm=[0,3,1,2]) node."""
+        parent = output_name_to_node.get(output_name, None)
+        if parent is not None and parent.op_type == "Transpose":
+            permutation = OnnxModel.get_node_attribute(parent, "perm")
+            if permutation == [0, 3, 1, 2]:
+                self.remove_if_safe(parent, input_name_to_nodes)
+                return parent
+        return None
+
+    def fuse(self, node, input_name_to_nodes, output_name_to_node):
+        # This fusion requires shape information, so skip it if shape is not available.
+        if self.shape_infer_helper is None:
+            return
+
+        # Before Fusion:
+        #     (a)  (b)
+        #       \  /
+        #       Add
+        #       /\
+        #   (c)   Transpose([0,2,3,1])
+        #            \
+        #          GroupNorm
+        #             |
+        #            (d)
+        #
+        # After Fusion:
+        #           (a)              (b)
+        #             \              /
+        #   Transpose([0,2,3,1])   Transpose([0,2,3,1])
+        #                \        /
+        #              SkipGroupNorm
+        #                  /    \
+        #                 /    Transpose([0, 3, 1, 2])
+        #                /        \
+        #               (d)       (c)
+        nodes = self.model.match_parent_path(node, ["Transpose", "Add"], [0, 0], output_name_to_node)
+        if nodes is None:
+            return
+
+        (transpose, add) = nodes
+        if transpose in self.nodes_to_remove or add in self.nodes_to_remove:
+            return
+
+        if self.has_multiple_consumers(transpose.output[0], input_name_to_nodes):
+            return
+
+        permutation = OnnxModel.get_node_attribute(transpose, "perm")
+        if permutation != [0, 2, 3, 1]:
+            return
+
+        inputs = []
+        bias = None
+        for i in range(2):
+            matched_transpose = self.match_transpose_from_nhwc(add.input[i], input_name_to_nodes, output_name_to_node)
+            if matched_transpose:
+                # When there is an Transpose node before Add (see examples in match_bias_path), we do not need to
+                # insert another Transpose node. The existing Transpose node will be removed in prune_graph if it
+                # has only one consumer.
+                inputs.append(matched_transpose.input[0])
+                # See whether it match bias pattern.
+                if bias is None:
+                    bias = self.match_bias_path(matched_transpose, input_name_to_nodes, output_name_to_node)
+            else:
+                # Otherwise, insert a Transpose node before Add.
+                new_transpose = self.create_transpose_node(add.input[i], [0, 2, 3, 1])
+                self.model.add_node(new_transpose, self.this_graph_name)
+                inputs.append(new_transpose.output[0])
+
+        skip, broadcast = self.get_skip_index(add, is_channel_last=False)
+        if skip < 0:
+            return
+
+        inputs = [inputs[1 - skip], node.input[1], node.input[2], inputs[skip]]
+        if bias:
+            inputs = [*inputs, bias]
+
+        outputs = node.output
+
+        new_node_name = self.model.create_node_name(self.fused_op_type, name_prefix="SkipGroupNorm")
+        if self.has_multiple_consumers(add.output[0], input_name_to_nodes):
+            add_out_name = new_node_name + "_add_out"
+            outputs.append(add_out_name)
+
+            # Insert a Transpose node after add output.
+            add_out_transpose = self.create_transpose_node(add_out_name, [0, 3, 1, 2], add.output[0])
+            self.model.add_node(add_out_transpose, self.this_graph_name)
+
+        skip_group_norm = helper.make_node(
+            self.fused_op_type,
+            inputs=inputs,
+            outputs=outputs,
+            name=new_node_name,
+        )
+        skip_group_norm.domain = "com.microsoft"
+
+        self.increase_counter(
+            f"SkipGroupNorm(add_out={int(len(outputs) > 1)} bias={int(bias is not None)} broadcast={int(broadcast)})"
+        )
+
+        # Pass attributes from GroupNorm node to SkipGroupNorm
+        for att in node.attribute:
+            skip_group_norm.attribute.extend([att])
+
+        self.nodes_to_remove.extend([add, transpose, node])
+        self.nodes_to_add.append(skip_group_norm)
+        self.node_name_to_graph_name[skip_group_norm.name] = self.this_graph_name
+        self.prune_graph = True