onnxruntime-directml 1.24.1__cp314-cp314-win_amd64.whl

onnxruntime/transformers/onnx_model_tnlr.py

@@ -0,0 +1,226 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation.  All rights reserved.
+# Licensed under the MIT License.
+# --------------------------------------------------------------------------
+import logging
+
+from fusion_attention import AttentionMask, FusionAttention
+from fusion_utils import NumpyHelper
+from onnx import NodeProto, helper
+from onnx_model import OnnxModel
+from onnx_model_bert import BertOnnxModel
+
+logger = logging.getLogger(__name__)
+
+
+class FusionTnlrAttention(FusionAttention):
+    """
+    Fuse TNLR Attention subgraph into one Attention node.
+    TNLR Attention has extra addition after qk nodes and adopts [S, B, NH] as I/O shape.
+    """
+
+    def __init__(
+        self,
+        model: OnnxModel,
+        hidden_size: int,
+        num_heads: int,
+        attention_mask: AttentionMask,
+    ):
+        super().__init__(model, hidden_size, num_heads, attention_mask)
+
+    def create_attention_node(
+        self,
+        mask_index: str,
+        matmul: NodeProto,
+        add: NodeProto,
+        num_heads: int,
+        hidden_size: int,
+        input: str,
+        output: str,
+        add_qk_str: str,
+    ) -> NodeProto | None:
+        assert num_heads > 0
+        if hidden_size > 0 and (hidden_size % num_heads) != 0:
+            logger.debug(f"input hidden size {hidden_size} is not a multiple of num of heads {num_heads}")
+            return None
+
+        weight = self.model.get_initializer(matmul.input[1])
+        bias = self.model.get_initializer(add.input[1]) or self.model.get_initializer(add.input[0])
+
+        if weight is None or bias is None:
+            return None
+
+        qkv_weight = NumpyHelper.to_array(weight)
+        qkv_bias = NumpyHelper.to_array(bias)
+
+        attention_node_name = self.model.create_node_name("Attention")
+
+        tensor_dtype = weight.data_type
+        np_type = helper.tensor_dtype_to_np_dtype(tensor_dtype)
+        weight = helper.make_tensor(
+            name=attention_node_name + "_qkv_weight",
+            data_type=tensor_dtype,
+            dims=[hidden_size, 3 * hidden_size],
+            vals=qkv_weight.astype(np_type).tobytes(),
+            raw=True,
+        )
+        self.model.add_initializer(weight, self.this_graph_name)
+
+        bias = helper.make_tensor(
+            name=attention_node_name + "_qkv_bias",
+            data_type=tensor_dtype,
+            dims=[3 * hidden_size],
+            vals=qkv_bias.astype(np_type).tobytes(),
+            raw=True,
+        )
+        self.model.add_initializer(bias, self.this_graph_name)
+
+        attention_inputs = [
+            input,
+            attention_node_name + "_qkv_weight",
+            attention_node_name + "_qkv_bias",
+        ]
+        if mask_index is not None:
+            attention_inputs.append(mask_index)
+        else:
+            attention_inputs.append("")
+
+        if add_qk_str is not None:
+            attention_inputs.append("")
+            attention_inputs.append(add_qk_str)
+
+        attention_node = helper.make_node(
+            "Attention",
+            inputs=attention_inputs,
+            outputs=[output],
+            name=attention_node_name,
+        )
+        attention_node.domain = "com.microsoft"
+        attention_node.attribute.extend([helper.make_attribute("num_heads", num_heads)])
+
+        return attention_node
+
+    def fuse(self, normalize_node, input_name_to_nodes, output_name_to_node):
+        # Sometimes we can not fuse skiplayernormalization since the add before layernorm has an output that used by nodes outside skiplayernorm
+        # Conceptually we treat add before layernorm as skiplayernorm node since they share the same pattern
+        start_node = normalize_node
+        if normalize_node.op_type != "SkipLayerNormalization":
+            return
+
+        # SkipLayerNormalization has two inputs, and one of them is the root input for attention.
+        qkv_nodes = self.model.match_parent_path(
+            start_node,
+            ["Where", "Add", "MatMul", "Reshape", "Transpose", "MatMul"],
+            [1, 1, 1, 0, 0, 0],
+        )
+        if qkv_nodes is not None:
+            (_, _, matmul_below, reshape_qkv, transpose_qkv, matmul_qkv) = qkv_nodes
+        else:
+            return
+
+        other_inputs = []
+        for _i, input in enumerate(start_node.input):
+            if input not in output_name_to_node:
+                continue
+
+            if input == qkv_nodes[0].output[0]:
+                continue
+            other_inputs.append(input)
+        if len(other_inputs) != 1:
+            return
+
+        root_input = other_inputs[0]
+
+        v_nodes = self.model.match_parent_path(
+            matmul_qkv,
+            ["Transpose", "Reshape", "Slice", "Add", "MatMul"],
+            [1, 0, 0, 0, 1],
+        )
+        if v_nodes is None:
+            return
+        (_, _, _, add, matmul) = v_nodes
+
+        upper_nodes = self.model.match_parent_path(matmul, ["Transpose"], [0])
+        transpose = upper_nodes[0]
+
+        qk_nodes = self.model.match_parent_path(matmul_qkv, ["Softmax", "Add", "MatMul"], [0, 0, 0])
+        if qk_nodes is None:
+            return
+        (_, add_qk, matmul_qk) = qk_nodes
+
+        q_nodes = self.model.match_parent_path(
+            matmul_qk,
+            ["Mul", "Transpose", "Reshape", "Slice", "Add", "MatMul"],
+            [0, 0, 0, 0, 0, 1],
+        )
+        if q_nodes is None:
+            return
+        add = q_nodes[-2]
+        matmul = q_nodes[-1]
+
+        k_nodes = self.model.match_parent_path(
+            matmul_qk,
+            ["Transpose", "Reshape", "Slice", "Add", "MatMul"],
+            [1, 0, 0, 0, 1],
+        )
+        if k_nodes is None:
+            return
+        add = k_nodes[-2]
+        matmul = k_nodes[-1]
+
+        relative_position_bias_nodes = self.model.match_parent_path(add_qk, ["Reshape", "Where"], [1, 0])
+        if relative_position_bias_nodes is None:
+            return
+
+        if matmul.input[0] == root_input:
+            mask_index = None
+            attention_last_node = reshape_qkv
+            # number of heads are same for all the paths, hence to create attention node, we pass the q_num_heads
+            # the input_hidden_size represents the input hidden size, this is used as needed but hidden sizes for Q, K are extracted appropriately
+            new_node = self.create_attention_node(
+                mask_index,
+                matmul,
+                add,
+                self.num_heads,
+                self.hidden_size,
+                root_input,
+                attention_last_node.output[0],
+                relative_position_bias_nodes[0].input[0],
+            )
+            if new_node is None:
+                return
+
+            self.nodes_to_add.append(new_node)
+            self.node_name_to_graph_name[new_node.name] = self.this_graph_name
+
+            # Add a transpose node after the attention node
+            back_transpose = helper.make_node(
+                "Transpose",
+                ["back_transpose_in_" + new_node.name],
+                [new_node.output[0]],
+                "back_transpose_" + new_node.name,
+                perm=[1, 0, 2],
+            )
+            self.model.add_node(back_transpose, self.this_graph_name)
+            new_node.input[0] = transpose.input[0]
+            new_node.output[0] = "back_transpose_in_" + new_node.name
+
+            self.nodes_to_remove.extend([attention_last_node, transpose_qkv, matmul_qkv])
+            self.nodes_to_remove.extend(qk_nodes)
+            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.nodes_to_remove.extend(mask_nodes)
+            self.prune_graph = True
+
+
+class TnlrOnnxModel(BertOnnxModel):
+    def __init__(self, model, num_heads, hidden_size):
+        super().__init__(model, num_heads, hidden_size)
+        self.attention_mask = AttentionMask(self)
+        self.attention_fusion = FusionTnlrAttention(self, self.hidden_size, self.num_heads, self.attention_mask)
+
+    def fuse_attention(self):
+        self.attention_fusion.apply()

onnxruntime/transformers/onnx_model_unet.py

@@ -0,0 +1,258 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation.  All rights reserved.
+# Licensed under the MIT License.
+# --------------------------------------------------------------------------
+
+import logging
+
+from fusion_attention_unet import FusionAttentionUnet
+from fusion_bias_add import FusionBiasAdd
+from fusion_biassplitgelu import FusionBiasSplitGelu
+from fusion_group_norm import FusionGroupNorm
+from fusion_nhwc_conv import FusionNhwcConv
+from fusion_options import FusionOptions
+from fusion_skip_group_norm import FusionSkipGroupNorm
+from fusion_transpose import FusionInsertTranspose, FusionTranspose
+from import_utils import is_installed
+from onnx import ModelProto
+from onnx_model import OnnxModel
+from onnx_model_bert import BertOnnxModel
+
+logger = logging.getLogger(__name__)
+
+
+class UnetOnnxModel(BertOnnxModel):
+    def __init__(self, model: ModelProto, num_heads: int = 0, hidden_size: int = 0):
+        """Initialize UNet ONNX Model.
+
+        Args:
+            model (ModelProto): the ONNX model
+            num_heads (int, optional): number of attention heads. Defaults to 0 (detect the parameter automatically).
+            hidden_size (int, optional): hidden dimension. Defaults to 0 (detect the parameter automatically).
+        """
+        assert (num_heads == 0 and hidden_size == 0) or (num_heads > 0 and hidden_size % num_heads == 0)
+
+        super().__init__(model, num_heads=num_heads, hidden_size=hidden_size)
+
+    def preprocess(self):
+        self.remove_useless_div()
+
+    def postprocess(self):
+        self.prune_graph()
+        self.remove_unused_constant()
+
+    def remove_useless_div(self):
+        """Remove Div by 1"""
+        div_nodes = [node for node in self.nodes() if node.op_type == "Div"]
+
+        nodes_to_remove = []
+        for div in div_nodes:
+            if self.find_constant_input(div, 1.0) == 1:
+                nodes_to_remove.append(div)
+
+        for node in nodes_to_remove:
+            self.replace_input_of_all_nodes(node.output[0], node.input[0])
+
+        if nodes_to_remove:
+            self.remove_nodes(nodes_to_remove)
+            logger.info("Removed %d Div nodes", len(nodes_to_remove))
+
+    def convert_conv_to_nhwc(self):
+        # Transpose weights in offline might help since ORT does not apply constant-folding on Transpose nodes.
+        conv_to_nhwc_conv = FusionNhwcConv(self, update_weight=True)
+        conv_to_nhwc_conv.apply()
+
+    def merge_adjacent_transpose(self):
+        fusion_transpose = FusionTranspose(self)
+        fusion_transpose.apply()
+
+        remove_count = 0
+        nodes = self.get_nodes_by_op_type("Transpose")
+        for node in nodes:
+            permutation = OnnxModel.get_node_attribute(node, "perm")
+            assert isinstance(permutation, list)
+            if permutation != list(range(len(permutation))):
+                continue
+            assert not (
+                self.find_graph_output(node.output[0])
+                or self.find_graph_input(node.input[0])
+                or self.find_graph_output(node.input[0])
+            )
+
+            # Let all children nodes skip current Transpose node and link to its parent
+            # Note that we cannot update parent node output since parent node might have more than one children.
+            self.replace_input_of_all_nodes(node.output[0], node.input[0])
+
+            self.remove_node(node)
+            remove_count += 1
+
+        total = len(fusion_transpose.nodes_to_remove) + remove_count
+        if total:
+            logger.info("Removed %d Transpose nodes", total)
+
+    def fuse_multi_head_attention(self, options: FusionOptions | None = None):
+        # Self Attention
+        enable_packed_qkv = (options is None) or options.enable_packed_qkv
+        self_attention_fusion = FusionAttentionUnet(
+            self,
+            self.hidden_size,
+            self.num_heads,
+            is_cross_attention=False,
+            enable_packed_qkv=enable_packed_qkv,
+            enable_packed_kv=False,
+        )
+        self_attention_fusion.apply()
+
+        # Cross Attention
+        enable_packed_kv = (options is None) or options.enable_packed_kv
+        cross_attention_fusion = FusionAttentionUnet(
+            self,
+            self.hidden_size,
+            self.num_heads,
+            is_cross_attention=True,
+            enable_packed_qkv=False,
+            enable_packed_kv=enable_packed_kv,
+        )
+        cross_attention_fusion.apply()
+
+    def fuse_bias_add(self):
+        fusion = FusionBiasAdd(self)
+        fusion.apply()
+
+    def optimize(self, options: FusionOptions | None = None):
+        if is_installed("tqdm"):
+            import tqdm  # noqa: PLC0415
+            from tqdm.contrib.logging import logging_redirect_tqdm  # noqa: PLC0415
+
+            with logging_redirect_tqdm():
+                steps = 18
+                progress_bar = tqdm.tqdm(range(steps), initial=0, desc="fusion")
+                self._optimize(options, progress_bar)
+        else:
+            logger.info("tqdm is not installed. Run optimization without progress bar")
+            self._optimize(options, None)
+
+    def _optimize(self, options: FusionOptions | None = None, progress_bar=None):
+        if (options is not None) and not options.enable_shape_inference:
+            self.disable_shape_inference()
+
+        self.utils.remove_identity_nodes()
+        if progress_bar:
+            progress_bar.update(1)
+
+        # Remove cast nodes that having same data type of input and output based on symbolic shape inference.
+        self.utils.remove_useless_cast_nodes()
+        if progress_bar:
+            progress_bar.update(1)
+
+        if (options is None) or options.enable_layer_norm:
+            self.fuse_layer_norm()
+        if progress_bar:
+            progress_bar.update(1)
+
+        if (options is None) or options.enable_gelu:
+            self.fuse_gelu()
+        if progress_bar:
+            progress_bar.update(1)
+
+        self.preprocess()
+        if progress_bar:
+            progress_bar.update(1)
+
+        self.fuse_reshape()
+        if progress_bar:
+            progress_bar.update(1)
+
+        if (options is None) or options.enable_group_norm:
+            channels_last = (options is None) or options.group_norm_channels_last
+            group_norm_fusion = FusionGroupNorm(self, channels_last)
+            group_norm_fusion.apply()
+
+            insert_transpose_fusion = FusionInsertTranspose(self)
+            insert_transpose_fusion.apply()
+        if progress_bar:
+            progress_bar.update(1)
+
+        if (options is None) or options.enable_bias_splitgelu:
+            bias_split_gelu_fusion = FusionBiasSplitGelu(self)
+            bias_split_gelu_fusion.apply()
+        if progress_bar:
+            progress_bar.update(1)
+
+        if (options is None) or options.enable_attention:
+            # self.save_model_to_file("before_mha.onnx")
+            self.fuse_multi_head_attention(options)
+        if progress_bar:
+            progress_bar.update(1)
+
+        if (options is None) or options.enable_skip_layer_norm:
+            self.fuse_skip_layer_norm()
+        if progress_bar:
+            progress_bar.update(1)
+
+        self.fuse_shape()
+        if progress_bar:
+            progress_bar.update(1)
+
+        # Remove reshape nodes that having same shape of input and output based on symbolic shape inference.
+        self.utils.remove_useless_reshape_nodes()
+        if progress_bar:
+            progress_bar.update(1)
+
+        if (options is None) or options.enable_skip_group_norm:
+            skip_group_norm_fusion = FusionSkipGroupNorm(self)
+            skip_group_norm_fusion.apply()
+        if progress_bar:
+            progress_bar.update(1)
+
+        if (options is None) or options.enable_bias_skip_layer_norm:
+            # Fuse SkipLayerNormalization and Add Bias before it.
+            self.fuse_add_bias_skip_layer_norm()
+        if progress_bar:
+            progress_bar.update(1)
+
+        if options is not None and options.enable_gelu_approximation:
+            self.gelu_approximation()
+        if progress_bar:
+            progress_bar.update(1)
+
+        if options is None or options.enable_nhwc_conv:
+            self.convert_conv_to_nhwc()
+            self.merge_adjacent_transpose()
+        if progress_bar:
+            progress_bar.update(1)
+
+        if options is not None and options.enable_bias_add:
+            self.fuse_bias_add()
+        if progress_bar:
+            progress_bar.update(1)
+
+        self.postprocess()
+        if progress_bar:
+            progress_bar.update(1)
+
+        logger.info(f"opset version: {self.get_opset_version()}")
+
+    def get_fused_operator_statistics(self):
+        """
+        Returns node count of fused operators.
+        """
+        op_count = {}
+        ops = [
+            "Attention",
+            "MultiHeadAttention",
+            "LayerNormalization",
+            "SkipLayerNormalization",
+            "BiasSplitGelu",
+            "GroupNorm",
+            "SkipGroupNorm",
+            "NhwcConv",
+            "BiasAdd",
+        ]
+
+        for op in ops:
+            nodes = self.get_nodes_by_op_type(op)
+            op_count[op] = len(nodes)
+
+        logger.info(f"Optimized operators:{op_count}")
+        return op_count

onnxruntime/transformers/onnx_model_vae.py

@@ -0,0 +1,42 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation.  All rights reserved.
+# Licensed under the MIT License.
+# --------------------------------------------------------------------------
+
+from logging import getLogger
+
+from fusion_attention_vae import FusionAttentionVae
+from fusion_options import FusionOptions
+from onnx import ModelProto
+from onnx_model_unet import UnetOnnxModel
+
+logger = getLogger(__name__)
+
+
+class VaeOnnxModel(UnetOnnxModel):
+    def __init__(self, model: ModelProto, num_heads: int = 0, hidden_size: int = 0):
+        assert (num_heads == 0 and hidden_size == 0) or (num_heads > 0 and hidden_size % num_heads == 0)
+        super().__init__(model, num_heads=num_heads, hidden_size=hidden_size)
+
+    def fuse_multi_head_attention(self, options: FusionOptions | None = None):
+        # Self Attention
+        self_attention_fusion = FusionAttentionVae(self, self.hidden_size, self.num_heads)
+        self_attention_fusion.apply()
+
+    def get_fused_operator_statistics(self):
+        """
+        Returns node count of fused operators.
+        """
+        op_count = {}
+        ops = [
+            "Attention",
+            "GroupNorm",
+            "SkipGroupNorm",
+            "NhwcConv",
+        ]
+        for op in ops:
+            nodes = self.get_nodes_by_op_type(op)
+            op_count[op] = len(nodes)
+
+        logger.info(f"Optimized operators:{op_count}")
+        return op_count

onnxruntime/transformers/onnx_utils.py

@@ -0,0 +1,55 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation.  All rights reserved.
+# Licensed under the MIT License.
+# --------------------------------------------------------------------------
+from fusion_utils import NumpyHelper
+from onnx import ModelProto, TensorProto
+from onnx.external_data_helper import set_external_data
+from onnx_model import OnnxModel
+
+from onnxruntime import OrtValue
+
+
+def extract_raw_data_from_model(model: ModelProto):
+    """
+    Extract external data from model and return the external data as a list of tuples (name, value).
+    Note this function does not handle external data that is not loaded into the model as raw data.
+
+    Args:
+        model (ModelProto): the model proto to extract external data from.
+    Returns:
+        (external_names, external_values): a tuple of two lists of external data names and values.
+    """
+    external_data = []
+    onnx_model = OnnxModel(model)
+    for graph in onnx_model.graphs():
+        for initializer in graph.initializer:
+            name = initializer.name
+
+            if initializer.HasField("raw_data"):
+                numpy_tensor = NumpyHelper.to_array(initializer)
+                ort_value = OrtValue.ortvalue_from_numpy(numpy_tensor)
+                external_data.append((name, ort_value))
+                # mimic set_external_data
+                set_external_data(initializer, location="foo.bin")
+                initializer.name = name
+                initializer.ClearField("raw_data")
+
+    return zip(*external_data, strict=False)
+
+
+def has_external_data(model: ModelProto):
+    """
+    Check if the model has external data.
+
+    Args:
+        model (ModelProto): the model proto to check for external data.
+    Returns:
+        bool: True if the model has external data, False otherwise.
+    """
+    onnx_model = OnnxModel(model)
+    for graph in onnx_model.graphs():
+        for initializer in graph.initializer:
+            if initializer.HasField("data_location") and initializer.data_location == TensorProto.EXTERNAL:
+                return True
+    return False