onnxruntime-directml 1.20.0__cp313-cp313-win_amd64.whl

onnxruntime/transformers/fusion_attention_clip.py

@@ -0,0 +1,257 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation.  All rights reserved.
+# Licensed under the MIT License.
+# --------------------------------------------------------------------------
+from logging import getLogger
+from typing import Tuple
+
+from fusion_attention import AttentionMask, FusionAttention
+from fusion_options import AttentionMaskFormat
+from onnx import NodeProto
+from onnx_model import OnnxModel
+
+logger = getLogger(__name__)
+
+
+class FusionAttentionClip(FusionAttention):
+    """
+    Fuse Attention subgraph of Clip into one Attention node.
+    """
+
+    def __init__(
+        self,
+        model: OnnxModel,
+        hidden_size: int,
+        num_heads: int,
+    ):
+        attention_mask = AttentionMask(model)
+        attention_mask.mask_format = AttentionMaskFormat.NoMask
+
+        super().__init__(
+            model,
+            hidden_size,
+            num_heads,
+            attention_mask,
+            use_multi_head_attention=False,
+            search_op_types=["SkipLayerNormalization"],
+        )
+
+    def get_num_heads_and_hidden_size(self, reshape_q: NodeProto) -> Tuple[int, int]:
+        """Detect num_heads and hidden_size for ONNX model from MiDaS
+        Args:
+            reshape_q (NodeProto): reshape node for q
+        Returns:
+            Tuple[int, int]: num_heads and hidden_size
+        """
+        concat = self.model.match_parent(reshape_q, "Concat", 1)
+        if concat is None or len(concat.input) != 4:
+            return self.num_heads, self.hidden_size
+
+        # The shape is a tensor like [?, ?, num_heads, head_size]
+        num_head_value = self.model.get_constant_value(concat.input[2])
+        if num_head_value is None:
+            return self.num_heads, self.hidden_size  # Fall back to user specified value
+
+        if len(num_head_value) != 1 or num_head_value[0] <= 0:
+            return self.num_heads, self.hidden_size  # Fall back to user specified value
+
+        num_heads = num_head_value[0]
+
+        head_size_value = self.model.get_constant_value(concat.input[3])
+        if head_size_value is None:
+            return self.num_heads, self.hidden_size  # Fall back to user specified value
+
+        if len(head_size_value) != 1 or head_size_value[0] <= 0:
+            return self.num_heads, self.hidden_size  # Fall back to user specified value
+
+        head_size = head_size_value[0]
+
+        hidden_size = num_heads * head_size
+
+        if self.num_heads > 0 and num_heads != self.num_heads:
+            if self.num_heads_warning:
+                logger.warning(f"--num_heads is {self.num_heads}. Detected value is {num_heads}. Using detected value.")
+                self.num_heads_warning = False  # Do not show the warning more than once
+
+        if self.hidden_size > 0 and hidden_size != self.hidden_size:
+            if self.hidden_size_warning:
+                logger.warning(
+                    f"--hidden_size is {self.hidden_size}. Detected value is {hidden_size}. Using detected value."
+                )
+                self.hidden_size_warning = False  # Do not show the warning more than once
+
+        return num_heads, hidden_size
+
+    def fuse(self, normalize_node, input_name_to_nodes, output_name_to_node):
+        skip_input_index = None
+        node_before_layer_norm = None
+        for i in [1, 0]:
+            parent = self.model.match_parent(normalize_node, "SkipLayerNormalization", i)
+            if parent is not None:
+                skip_input_index = i
+                node_before_layer_norm = parent
+
+        root_input = None
+        if node_before_layer_norm is not None:
+            root_input = node_before_layer_norm.output[0]
+        else:
+            # Deal with the first attention after the embedding layer.
+            for i in [0, 1]:
+                node_before_layer_norm = None
+
+                node_before_layer_norm_1 = self.model.match_parent(normalize_node, "Add", i)
+                node_before_layer_norm_2 = self.model.match_parent(normalize_node, "LayerNormalization", i)
+                if node_before_layer_norm_1 is not None:
+                    #           Add -----------+
+                    #            |             |
+                    #        LayerNorm         |
+                    #            |             |
+                    #        LayerNorm         |
+                    #            |             |
+                    #   Attention subgraph     |
+                    #            |             |
+                    #      SkipLayerNorm ------+
+                    node_before_layer_norm = node_before_layer_norm_1
+                elif node_before_layer_norm_2 is not None:
+                    #           Add
+                    #            |
+                    #        LayerNorm --------+
+                    #            |             |
+                    #        LayerNorm         |
+                    #            |             |
+                    #   Attention subgraph     |
+                    #            |             |
+                    #      SkipLayerNorm ------+
+                    node_before_layer_norm = node_before_layer_norm_2
+
+                if node_before_layer_norm is None:
+                    continue
+                child = self.model.find_first_child_by_type(
+                    node_before_layer_norm, "LayerNormalization", input_name_to_nodes, False
+                )
+                if child is None:
+                    continue
+                root_input = child.output[0]
+                skip_input_index = i
+                break
+
+            if skip_input_index is None:
+                return
+
+        qkv_nodes = self.model.match_parent_path(
+            normalize_node,
+            ["Add", "MatMul", "Reshape", "Transpose", "Reshape", "MatMul"],
+            [1 - skip_input_index, None, None, 0, 0, 0],
+        )
+        if qkv_nodes is None:
+            return
+
+        (_, _, reshape_qkv, transpose_qkv, _, matmul_qkv) = qkv_nodes
+
+        v_nodes = self.model.match_parent_path(
+            matmul_qkv, ["Reshape", "Transpose", "Reshape", "Add", "MatMul"], [1, 0, 0, 0, None]
+        )
+        if v_nodes is None:
+            logger.debug("fuse_attention: failed to match v path")
+            return
+        (_, _, reshape_v, add_v, matmul_v) = v_nodes
+
+        add_mask = None
+        add_mask_indices = []
+        qk_nodes = None
+        qk_nodes_1 = self.model.match_parent_path(
+            matmul_qkv,
+            ["Softmax", "Reshape", "Add", "Reshape", "MatMul"],
+            [0, 0, 0, None, 0],
+            return_indice=add_mask_indices,
+        )
+        qk_nodes_2 = self.model.match_parent_path(
+            matmul_qkv,
+            ["Softmax", "MatMul"],
+            [0, 0],
+        )
+        if qk_nodes_1 is not None:
+            qk_nodes = qk_nodes_1
+            assert len(add_mask_indices) == 1
+            causal_mask_input_index = 1 - add_mask_indices[0]
+
+            (_softmax_qk, _, add_mask, _, matmul_qk) = qk_nodes
+        elif qk_nodes_2 is not None:
+            qk_nodes = qk_nodes_2
+            (_softmax_qk, matmul_qk) = qk_nodes
+        else:
+            logger.debug("fuse_attention: failed to match qk path")
+            return
+
+        q_nodes = self.model.match_parent_path(
+            matmul_qk, ["Reshape", "Transpose", "Reshape", "Mul", "Add", "MatMul"], [0, 0, 0, 0, None, None]
+        )
+        if q_nodes is None:
+            logger.debug("fuse_attention: failed to match q path")
+            return
+        (_, _transpose_q, reshape_q, mul_q, add_q, matmul_q) = q_nodes
+
+        k_nodes = self.model.match_parent_path(
+            matmul_qk, ["Transpose", "Reshape", "Transpose", "Reshape", "Add", "MatMul"], [1, 0, 0, 0, 0, None]
+        )
+        if k_nodes is None:
+            logger.debug("fuse_attention: failed to match k path")
+            return
+
+        (_transpose_k, _reshape_k, _, _, add_k, matmul_k) = k_nodes
+        if matmul_q.input[0] != root_input or matmul_k.input[0] != root_input or matmul_v.input[0] != root_input:
+            logger.debug("fuse_attention: expect to have same input to q, k and v matmul")
+            return
+
+        num_heads, hidden_size = self.get_num_heads_and_hidden_size(reshape_q)
+        if num_heads <= 0 or hidden_size <= 0:
+            logger.debug("fuse_attention: failed to detect num_heads or hidden_size")
+            return
+
+        attention_last_node = reshape_qkv
+
+        if add_mask is not None:
+            # Here we do not match the whole subgraph since it is very complex. Instead, we just check whether a key path
+            # of computing causal mask.
+            causal_mask_nodes = self.model.match_parent_path(
+                add_mask,
+                ["Concat", "Expand", "Unsqueeze", "Unsqueeze", "Where", "Less"],
+                [causal_mask_input_index, 0, 0, 0, 0, 0],
+            )
+            if causal_mask_nodes is None:
+                # If the model is exported with batch_size == 1, there is no Concat node
+                causal_mask_nodes = self.model.match_parent_path(
+                    add_mask,
+                    ["Expand", "Unsqueeze", "Unsqueeze", "Where", "Less"],
+                    [causal_mask_input_index, 0, 0, 0, 0],
+                )
+                if causal_mask_nodes is None:
+                    logger.debug("fuse_attention: failed to match causal mask subgraph")
+                    return
+
+        new_node = self.create_attention_node(
+            mask_index=None,
+            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,
+            input=root_input,
+            output=attention_last_node.output[0],
+            add_qk_str=None,
+            scale=None,
+            causal=(add_mask is not None),
+        )
+        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
+
+        self.nodes_to_remove.extend([attention_last_node, transpose_qkv])
+
+        # Use prune graph to remove nodes since they are shared by all attention nodes.
+        self.prune_graph = True