AndroidRemoteController/python/py/Lib/site-packages/onnxruntime/transformers/onnx_model_tnlr.py

# -------------------------------------------------------------------------
# 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()