# -*- coding: utf-8 -*-
# Copyright (c) Alibaba, Inc. and its affiliates.
# All rights reserved.
# This file contains code that is adapted from
# timm: https://github.com/huggingface/pytorch-image-models
# pixart: https://github.com/PixArt-alpha/PixArt-alpha
import math
import time
import warnings

import torch
import torch.nn as nn
from torch.cuda import amp
from torch.nn import functional as F
from torch.nn.utils.rnn import pad_sequence
from tqdm import tqdm

from scepter.modules.model.backbone.transformer.pos_embed import apply_2d_rope

try:
    import xformers
    import xformers.ops
    XFORMERS_IS_AVAILABLE = True
except Exception as e:
    XFORMERS_IS_AVAILABLE = False
    warnings.warn(f'{e}')
try:
    from flash_attn import (flash_attn_varlen_func)
    FLASHATTN_IS_AVAILABLE = True
except ImportError:
    FLASHATTN_IS_AVAILABLE = False
    flash_attn_varlen_func = None


def drop_path(x, drop_prob: float = 0., training: bool = False):
    """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).
    This is the same as the DropConnect impl I created for EfficientNet, etc networks, however,
    the original name is misleading as 'Drop Connect' is a different form of dropout in a separate paper...
    See discussion: https://github.com/tensorflow/tpu/issues/494#issuecomment-532968956 ... I've opted for
    changing the layer and argument names to 'drop path' rather than mix DropConnect as a layer name and use
    'survival rate' as the argument.
    """
    if drop_prob == 0. or not training:
        return x
    keep_prob = 1 - drop_prob
    shape = (x.shape[0], ) + (1, ) * (
        x.ndim - 1)  # work with diff dim tensors, not just 2D ConvNets
    random_tensor = keep_prob + torch.rand(
        shape, dtype=x.dtype, device=x.device)
    random_tensor.floor_()  # binarize
    output = x.div(keep_prob) * random_tensor
    return output


class MultiHeadAttention(nn.Module):
    def __init__(self,
                 dim,
                 context_dim=None,
                 num_heads=None,
                 head_dim=None,
                 attn_drop=0.0,
                 qkv_bias=False,
                 dropout=0.0,
                 backend=None,
                 **block_kwargs):
        super().__init__()
        # consider head_dim first, then num_heads
        num_heads = dim // head_dim if head_dim else num_heads
        head_dim = dim // num_heads
        assert num_heads * head_dim == dim
        context_dim = context_dim or dim
        self.dim = dim
        self.context_dim = context_dim
        self.num_heads = num_heads
        self.head_dim = head_dim
        self.scale = math.pow(head_dim, -0.25)
        # layers
        self.q = nn.Linear(dim, dim, bias=qkv_bias)
        self.k = nn.Linear(context_dim, dim, bias=qkv_bias)
        self.v = nn.Linear(context_dim, dim, bias=qkv_bias)
        self.o = nn.Linear(dim, dim)
        self.dropout = nn.Dropout(dropout)
        self.attention_op = None
        self.attn_drop = nn.Dropout(attn_drop)
        self.backend = backend
        assert self.backend in ('flash_attn', 'xformer_attn', 'pytorch_attn',
                                None)
        if FLASHATTN_IS_AVAILABLE and self.backend in ('flash_attn', None):
            self.backend = 'flash_attn'
            self.softmax_scale = block_kwargs.get('softmax_scale', None)
            self.causal = block_kwargs.get('causal', False)
            self.window_size = block_kwargs.get('window_size', (-1, -1))
            self.deterministic = block_kwargs.get('deterministic', False)
        elif XFORMERS_IS_AVAILABLE and self.backend in ('xformer_attn', None):
            self.backend = 'xformer_attn'
        else:
            self.backend = 'pytorch_attn'

    def xformer_attn(self, x, context=None, mask=None, **kwargs):
        context = x if context is None else context
        b, n, d = x.size(0), self.num_heads, self.head_dim
        # compute query, key, value
        q = self.q(x).view(b, -1, n, d)
        k = self.k(context).view(b, -1, n, d)
        v = self.v(context).view(b, -1, n, d)

        attn_bias = None
        if mask is not None:
            assert mask.ndim in [2, 3]
            mask = mask.view(b, 1, 1,
                             -1) if mask.ndim == 2 else mask.unsqueeze(1)
            # To use an `attn_bias` with a sequence length that is not a multiple of 8,
            # you need to ensure memory is aligned by slicing a bigger tensor.
            # Example: use `attn_bias = torch.zeros([1, 1, 5, 8])[:,:,:,:5]`
            # instead of `torch.zeros([1, 1, 5, 5])
            q_size = math.ceil(q.size(1) / 8) * 8
            k_size = math.ceil(k.size(1) / 8) * 8
            attn_bias = x.new_zeros(b, n, q_size,
                                    k_size)[:, :, :q.size(1), :k.size(1)]
            attn_bias = attn_bias.masked_fill_(mask == 0,
                                               torch.finfo(x.dtype).min).to(
                                                   q.dtype)
        x = xformers.ops.memory_efficient_attention(q,
                                                    k,
                                                    v,
                                                    p=self.attn_drop.p,
                                                    attn_bias=attn_bias)
        x = x.reshape(b, -1, n * d)
        x = self.o(x)
        x = self.dropout(x)
        return x

    def flash_attn(self, x, context=None, mask=None, **kwargs):
        '''
         The implementation will be very slow when mask is not None,
         because we need rearange the x/context features according to mask.
        Args:
            x:
            context:
            mask:
            **kwargs:
        Returns: x
        '''
        context = x if context is None else context
        dtype = kwargs.get('dtype', torch.float16)
        q_lens = kwargs.get('q_lens', None)

        # if mask is not None or q_lens is not None:
        #     warnings.warn("Detected mask or q_lens is not None, "
        #                   "which will be very slow because of the x/context features' rearrangement,"
        #                   "please use FlashMultiHeadAttention instead.")
        def half(x):
            return x if x.dtype in [torch.float16, torch.bfloat16
                                    ] else x.to(dtype)

        b, n, d = x.size(0), self.num_heads, self.head_dim

        q = self.q(x).view(b, -1, n, d)  # [B, Lq, Nq, C1].
        k = self.k(context).view(b, -1, n, d)  # [B, Lk, Nk, C1]
        v = self.v(context).view(
            b, -1, n, d)  # [B, Lk, Nk, C2] Nq must be divisible by Nk.

        assert q.device.type == 'cuda' and q.size(-1) <= 256
        lq, lk, out_dtype = int(q.size(1)), int(k.size(1)), q.dtype
        # preprocess query
        if q_lens is None:
            q_lens = torch.tensor([lq] * b,
                                  dtype=torch.int32).to(q.device,
                                                        non_blocking=True)
            # q_lens = (q.flatten(2, ).bool() + 1).sum(dim=-1).bool().sum(dim=-1)
            q = half(q.flatten(0, 1))
        else:
            q = half(torch.cat([q_v[:q_l] for q_v, q_l in zip(q, q_lens)]))

        # preprocess key, value
        if mask is None:
            k_lens = torch.tensor([lk] * b,
                                  dtype=torch.int32).to(k.device,
                                                        non_blocking=True)
            # k_lens = (k.flatten(2, ).bool() + 1).sum(dim=-1).bool().sum(dim=-1)
            k = half(k.flatten(0, 1))
            v = half(v.flatten(0, 1))
        else:
            assert mask.ndim in [1, 2, 3]
            k_lens = mask if mask.ndim == 1 else mask.flatten(start_dim=1).sum(
                dim=-1)
            k = half(torch.cat([k_v[:k_l] for k_v, k_l in zip(k, k_lens)]))
            v = half(torch.cat([v_v[:v_l] for v_v, v_l in zip(v, k_lens)]))

        x = flash_attn_varlen_func(
            q,
            k,
            v,
            cu_seqlens_q=torch.cat([q_lens.new_zeros([1]),
                                    q_lens]).cumsum(0, dtype=torch.int32),
            cu_seqlens_k=torch.cat([k_lens.new_zeros([1]),
                                    k_lens]).cumsum(0, dtype=torch.int32),
            max_seqlen_q=int(torch.max(q_lens).cpu().numpy()),
            max_seqlen_k=int(torch.max(k_lens).cpu().numpy()),
            dropout_p=self.attn_drop.p,
            softmax_scale=self.softmax_scale,
            causal=self.causal,
            window_size=self.window_size,  # -1 means infinite context window
            deterministic=self.deterministic).unflatten(0, (b, lq))
        x = x.type(out_dtype)
        x = x.flatten(2)
        # output
        x = self.o(x)
        x = self.dropout(x)
        return x

    def pytorch_attn(self, x, context=None, mask=None, **kwargs):
        """x:       [B, L, C].
           context: [B, L', C'] or None.
        """
        context = x if context is None else context
        b, n, d = x.size(0), self.num_heads, self.head_dim

        # compute query, key, value
        q = self.q(x).view(b, -1, n, d)
        k = self.k(context).view(b, -1, n, d)
        v = self.v(context).view(b, -1, n, d)
        # attention bias
        attn_bias = x.new_zeros(b, n, q.size(1), k.size(1))
        if mask is not None:
            assert mask.ndim in [2, 3]
            mask = mask.view(b, 1, 1,
                             -1) if mask.ndim == 2 else mask.unsqueeze(1)
            attn_bias = attn_bias.masked_fill_(mask == 0,
                                               torch.finfo(x.dtype).min).to(
                                                   q.dtype)

        # compute attention (T5 does not use scaling)
        attn = torch.einsum('binc,bjnc->bnij', q * self.scale,
                            k * self.scale) + attn_bias
        attn = F.softmax(attn.float(), dim=-1).type_as(attn)
        x = torch.einsum('bnij,bjnc->binc', attn, v.float())
        # output
        x = x.reshape(b, -1, n * d)
        x = self.o(x)
        x = self.dropout(x)
        return x

    def forward(self, x, context=None, mask=None, **kwargs):
        """x: [B, L, C].
           context: [B, L', C'] or None.
        """
        x = getattr(self, self.backend)(x,
                                        context=context,
                                        mask=mask,
                                        **kwargs)
        return x


def flash_preprocess(x, context=None, q_mask=None, mask=None):
    context = x if context is None else context
    b, x_l, x_hidden_size = x.shape
    x = x.flatten(0, 1)
    if q_mask is None:
        q_lens = torch.tensor([x_l] * b,
                              dtype=torch.int32).to(x.device,
                                                    non_blocking=True)
    else:
        assert q_mask.ndim in [1, 2, 3]
        q_lens = q_mask if q_mask.ndim == 1 else q_mask.flatten(
            start_dim=1).sum(dim=-1)

    mask_b, mask_l, mask_hidden_size = context.shape

    if mask is None:
        mask_lens = torch.tensor([mask_l] * mask_b,
                                 dtype=torch.int32).to(context.device,
                                                       non_blocking=True)
    else:
        assert mask.ndim in [1, 2, 3]
        mask_lens = mask if mask.ndim == 1 else mask.flatten(start_dim=1).sum(
            dim=-1)

    return_data = {
        'x':
        x,
        'context':
        torch.cat([u[:v] for u, v in zip(context, mask_lens)]),
        'cu_seqlens_q':
        torch.cat([q_lens.new_zeros([1]), q_lens]).cumsum(0,
                                                          dtype=torch.int32),
        'max_seqlen_q':
        int(torch.max(q_lens).cpu().numpy()),
        'cu_seqlens_k':
        torch.cat([mask_lens.new_zeros([1]),
                   mask_lens]).cumsum(0, dtype=torch.int32),
        'max_seqlen_k':
        int(torch.max(mask_lens).cpu().numpy())
    }
    return return_data


class FlashMultiHeadAttention(nn.Module):
    def __init__(self,
                 dim,
                 context_dim=None,
                 num_heads=None,
                 head_dim=None,
                 attn_drop=0.0,
                 qkv_bias=False,
                 dropout=0.0,
                 softmax_scale=None,
                 causal=False,
                 window_size=(-1, -1),
                 deterministic=False,
                 **block_kwargs):
        super().__init__()
        # consider head_dim first, then num_heads
        num_heads = dim // head_dim if head_dim else num_heads
        head_dim = dim // num_heads
        assert num_heads * head_dim == dim
        context_dim = context_dim or dim
        self.dim = dim
        self.context_dim = context_dim
        self.num_heads = num_heads
        self.head_dim = head_dim
        self.scale = math.pow(head_dim, -0.25)
        # layers
        self.q = nn.Linear(dim, dim, bias=qkv_bias)
        self.k = nn.Linear(context_dim, dim, bias=qkv_bias)
        self.v = nn.Linear(context_dim, dim, bias=qkv_bias)
        self.o = nn.Linear(dim, dim)
        self.dropout = nn.Dropout(dropout)
        self.attention_op = None
        self.attn_drop = nn.Dropout(attn_drop)
        self.softmax_scale = softmax_scale
        self.causal = causal
        self.window_size = window_size
        self.deterministic = deterministic

    def forward(self,
                x,
                context=None,
                cu_seqlens_q=None,
                max_seqlen_q=None,
                cu_seqlens_k=None,
                max_seqlen_k=None,
                dtype=torch.float16,
                **kwargs):
        '''
            The implementation used the rearanaged x/context according to q_lens or k_lens.
            Args:
                x: [batch_size * max_seq_len or sum(q_lens) , heads, hidden_size].
                context:  [batch_size * max_seq_len or sum(q_lens) , heads, hidden_size].
                cu_seqlens_q: cumsum of seq_q to index the postion of query in the batch.
                max_seqlen_q: max length of query.
                cu_seqlens_k: cumsum of seq_k to index the postion of key/value in the batch.
                max_seqlen_k: max length of key/value.
                dtype: the dtype for attention.
                **kwargs:
            Returns: x
            '''
        context = x if context is None else context

        def half(x):
            return x if x.dtype in [torch.float16, torch.bfloat16
                                    ] else x.to(dtype)

        n, d, out_dtype = self.num_heads, self.head_dim, x.dtype
        q = self.q(x).view(-1, n, d)  # [B * Lq, Nq, C1].
        k = self.k(context).view(-1, n, d)  # [B * Lk, Nk, C1]
        v = self.v(context).view(
            -1, n, d)  # [B * Lk, Nk, C2] Nq must be divisible by Nk.
        q, k, v = half(q), half(k), half(v)
        assert q.device.type == 'cuda' and d <= 256
        x = flash_attn_varlen_func(
            q,
            k,
            v,
            cu_seqlens_q=cu_seqlens_q,
            cu_seqlens_k=cu_seqlens_k,
            max_seqlen_q=max_seqlen_q,
            max_seqlen_k=max_seqlen_k,
            dropout_p=self.attn_drop.p,
            softmax_scale=self.softmax_scale,
            causal=self.causal,
            window_size=self.window_size,  # -1 means infinite context window
            deterministic=self.deterministic).unflatten(0, (x.shape[0], ))
        x = x.flatten(1).type(out_dtype)
        # output
        x = self.o(x)
        x = self.dropout(x)
        return x


def multi_head_varlen_attention(q_img,
                                k_img,
                                v_img,
                                q_txt,
                                k_txt,
                                v_txt,
                                n,
                                d,
                                img_lens,
                                txt_lens,
                                dropout_p=0.0,
                                flash_dtype=torch.bfloat16):
    '''
        q/k/v: b, s, n*d
        q_lens/k_lens: b,
    '''
    from flash_attn import flash_attn_varlen_func
    q_lens = k_lens = img_lens + txt_lens

    cu_seqlens_q = torch.cat([q_lens.new_zeros([1]),
                              q_lens]).cumsum(0, dtype=torch.int32)
    cu_seqlens_k = torch.cat([k_lens.new_zeros([1]),
                              k_lens]).cumsum(0, dtype=torch.int32)
    max_seqlen_q = q_lens.max()
    max_seqlen_k = k_lens.max()

    # concat img & txt for joint attention
    q = torch.cat([
        torch.cat([i[:i_len], t[:t_len]], dim=0)
        for i, i_len, t, t_len in zip(q_img, img_lens, q_txt, txt_lens)
    ],
                  dim=0).view(-1, n, d)

    k = torch.cat([
        torch.cat([i[:i_len], t[:t_len]], dim=0)
        for i, i_len, t, t_len in zip(k_img, img_lens, k_txt, txt_lens)
    ],
                  dim=0).view(-1, n, d)

    v = torch.cat([
        torch.cat([i[:i_len], t[:t_len]], dim=0)
        for i, i_len, t, t_len in zip(v_img, img_lens, v_txt, txt_lens)
    ],
                  dim=0).view(-1, n, d)

    # attention
    dtype = q.dtype
    if dtype != flash_dtype:
        q = q.type(flash_dtype)
        k = k.type(flash_dtype)
        v = v.type(flash_dtype)

    with amp.autocast():
        x = flash_attn_varlen_func(q=q,
                                   k=k,
                                   v=v,
                                   cu_seqlens_q=cu_seqlens_q,
                                   cu_seqlens_k=cu_seqlens_k,
                                   max_seqlen_q=max_seqlen_q,
                                   max_seqlen_k=max_seqlen_k,
                                   dropout_p=dropout_p).type(dtype)

    return x, cu_seqlens_q


class RMSNorm(nn.Module):
    def __init__(self, dim, eps=1e-6):
        super().__init__()
        self.dim = dim
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))

    def forward(self, x):
        return self._norm(x.float()).type_as(x) * self.weight

    def _norm(self, x):
        return x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps)


class FullAttention(nn.Module):
    def __init__(self,
                 dim,
                 num_heads=None,
                 head_dim=None,
                 dropout=0.0,
                 qkv_bias=False,
                 qk_norm=False,
                 eps=1e-6,
                 flash_dtype=torch.bfloat16):

        # consider head_dim first, then num_heads
        num_heads = dim // head_dim if head_dim else num_heads
        head_dim = dim // num_heads
        assert num_heads * head_dim == dim
        assert flash_dtype in (None, torch.float16, torch.bfloat16)
        super().__init__()
        self.dim = dim
        self.num_heads = num_heads
        self.head_dim = head_dim
        self.scale = math.pow(head_dim, -0.25)
        self.flash_dtype = flash_dtype
        # layers
        self.qkv_W = nn.Linear(dim, dim * 3, bias=qkv_bias)

        self.out_proj = nn.Linear(dim, dim)
        self.dropout = nn.Dropout(dropout)

        if qk_norm:
            from apex.normalization import FusedRMSNorm
            self.q_img_norm = FusedRMSNorm(head_dim, eps=eps)
            self.k_img_norm = FusedRMSNorm(head_dim, eps=eps)
            self.q_txt_norm = FusedRMSNorm(head_dim, eps=eps)
            self.k_txt_norm = FusedRMSNorm(head_dim, eps=eps)
        else:
            self.q_img_norm = nn.Identity()
            self.k_img_norm = nn.Identity()
            self.q_txt_norm = nn.Identity()
            self.k_txt_norm = nn.Identity()

    def forward(self,
                img,
                txt,
                img_lens=None,
                txt_lens=None,
                padded_pos_index=None):
        '''
            img: B, L, C
            txt: B, L', C
        '''
        b, img_len, c = img.shape
        txt_len, n, d = txt.shape[1], self.num_heads, self.head_dim

        # compute query, key, value
        img_txt = torch.cat([img, txt], dim=1)
        img_tokens, txt_tokens = torch.split(self.qkv_W(img_txt),
                                             [img_len, txt_len],
                                             dim=1)

        q_img, k_img, v_img = img_tokens.chunk(3, dim=-1)
        q_txt, k_txt, v_txt = txt_tokens.chunk(3, dim=-1)

        # multi-head qk norm
        q_img, k_img = q_img.view(b, -1, n, d), k_img.view(b, -1, n, d)
        q_txt, k_txt = q_txt.view(b, -1, n, d), k_txt.view(b, -1, n, d)
        q_img, q_txt = self.q_img_norm(q_img).view(
            b, -1, n * d), self.q_txt_norm(q_txt).view(b, -1, n * d)
        k_img, k_txt = self.k_img_norm(k_img).view(
            b, -1, n * d), self.k_txt_norm(k_txt).view(b, -1, n * d)

        # add position
        q_img, k_img = apply_2d_rope(q_img, k_img, padded_pos_index, n, d)

        # support varying length
        if img_lens is None:
            img_lens = torch.tensor([img.size(1)] * b,
                                    dtype=torch.int32,
                                    device=img.device)
        if txt_lens is None:
            txt_lens = torch.tensor([txt.size(1)] * b,
                                    dtype=torch.int32,
                                    device=txt.device)

        # attention
        x, cu_seqlens_q = multi_head_varlen_attention(
            q_img,
            k_img,
            v_img,
            q_txt,
            k_txt,
            v_txt,
            n,
            d,
            img_lens,
            txt_lens,
            dropout_p=self.dropout.p if self.training else 0.0,
            flash_dtype=self.flash_dtype)

        # output proj.
        x = x.reshape(-1, n * d)
        x = self.out_proj(x)
        x = self.dropout(x)

        # split img & txt and padding to max_len
        img = pad_sequence(tuple([
            x[s:s + img_len] for s, e, img_len in zip(
                cu_seqlens_q[:-1], cu_seqlens_q[1:], img_lens)
        ]),
                           batch_first=True)
        txt = pad_sequence(tuple([
            x[s + img_len:e] for s, e, img_len in zip(
                cu_seqlens_q[:-1], cu_seqlens_q[1:], img_lens)
        ]),
                           batch_first=True)

        return img, txt


class FFNSwiGLU(nn.Module):
    def __init__(self, in_features, hidden_features):
        super().__init__()
        self.W1 = nn.Linear(in_features, hidden_features, bias=False)
        self.W2 = nn.Linear(in_features, hidden_features, bias=False)
        self.W3 = nn.Linear(hidden_features, in_features, bias=False)
        self.silu = nn.SiLU()

    def forward(self, x):
        return self.W3(self.silu(self.W1(x)) * self.W2(x))


if __name__ == '__main__':
    # Align results for different attention implementation
    torch.manual_seed(2023)
    hidden_dim = 4096
    q_weight = torch.randn((hidden_dim, hidden_dim))
    q_bias = torch.zeros((hidden_dim))
    k_weight = torch.randn((hidden_dim, hidden_dim))
    k_bias = torch.zeros((hidden_dim))
    v_weight = torch.randn((hidden_dim, hidden_dim))
    v_bias = torch.zeros((hidden_dim))
    o_weight = torch.randn((hidden_dim, hidden_dim))
    o_bias = torch.randn((hidden_dim))
    pytorch_attn = MultiHeadAttention(hidden_dim,
                                      context_dim=hidden_dim,
                                      num_heads=32,
                                      head_dim=None,
                                      attn_drop=0.0,
                                      dropout=0.0,
                                      backend='pytorch_attn')

    pytorch_attn.load_state_dict({
        'q.weight': q_weight,
        'k.weight': k_weight,
        'v.weight': v_weight,
        'o.weight': o_weight,
        'o.bias': o_bias
    })
    pytorch_attn.to(0)

    xformer_attn = MultiHeadAttention(hidden_dim,
                                      context_dim=hidden_dim,
                                      num_heads=32,
                                      head_dim=None,
                                      attn_drop=0.0,
                                      dropout=0.0,
                                      backend='xformer_attn')

    xformer_attn.load_state_dict({
        'q.weight': q_weight,
        'k.weight': k_weight,
        'v.weight': v_weight,
        'o.weight': o_weight,
        'o.bias': o_bias
    })
    xformer_attn.to(0)

    flash_attn = MultiHeadAttention(hidden_dim,
                                    context_dim=hidden_dim,
                                    num_heads=32,
                                    head_dim=None,
                                    attn_drop=0.0,
                                    dropout=0.0,
                                    backend='flash_attn',
                                    dtype=torch.float16)
    flash_attn.load_state_dict({
        'q.weight': q_weight,
        'k.weight': k_weight,
        'v.weight': v_weight,
        'o.weight': o_weight,
        'o.bias': o_bias
    })
    flash_attn.to(0)

    improved_flash_attn = FlashMultiHeadAttention(hidden_dim,
                                                  context_dim=hidden_dim,
                                                  num_heads=32,
                                                  head_dim=None,
                                                  attn_drop=0.0,
                                                  dropout=0.0,
                                                  backend='flash_attn',
                                                  dtype=torch.float16)

    improved_flash_attn.load_state_dict({
        'q.weight': q_weight,
        'k.weight': k_weight,
        'v.weight': v_weight,
        'o.weight': o_weight,
        'o.bias': o_bias
    })
    improved_flash_attn.to(0)

    batch_size = 1
    query_length = 1024
    key_length = 1024
    # mask = None
    run_num = 10
    torch.cuda.empty_cache()
    x = torch.randn((batch_size, query_length, hidden_dim)).to(0)
    context = torch.randn((batch_size, key_length, hidden_dim)).to(0)
    # mask = torch.cat([torch.ones((batch_size, 80)), torch.zeros((batch_size, key_length - 80))], dim=1).long().to(0)
    # mask = torch.randint(1, key_length, [batch_size]).to(0)
    mask = None
    st = time.time()
    for i in tqdm(range(run_num)):
        pytorch_res = pytorch_attn(x.clone(), context.clone(),
                                   mask.clone() if mask is not None else mask)
        if i == run_num - 1:
            free_mem, total_mem = torch.cuda.mem_get_info(0)
            free_mem, total_mem = free_mem / (1024**3), total_mem / (1024**3)
            mem_msg = f'GPU {0}: free mem {free_mem:.3f}G, total mem {total_mem:.3f}G \n'
            pytorch_res_data = pytorch_res.clone().detach().cpu()
            print('pytorch attn ', mem_msg,
                  f'Cost time per time {(time.time() - st) / run_num}s')
    #
    torch.cuda.empty_cache()
    st = time.time()
    for i in tqdm(range(run_num)):
        xformer_res = xformer_attn(x.clone(), context.clone(),
                                   mask.clone() if mask is not None else mask)
        if i == run_num - 1:
            free_mem, total_mem = torch.cuda.mem_get_info(0)
            free_mem, total_mem = free_mem / (1024**3), total_mem / (1024**3)
            mem_msg = f'GPU {0}: free mem {free_mem:.3f}G, total mem {total_mem:.3f}G \n'
            xformer_res_data = xformer_res.clone().detach().cpu()
            print('xformer attn ', mem_msg,
                  f'Cost time per time {(time.time() - st) / run_num}s')
    #
    torch.cuda.empty_cache()
    # mask = None
    st = time.time()
    for i in tqdm(range(run_num)):
        flash_res = flash_attn(x.clone(), context.clone(),
                               mask.clone() if mask is not None else mask)
        if i == run_num - 1:
            free_mem, total_mem = torch.cuda.mem_get_info(0)
            free_mem, total_mem = free_mem / (1024**3), total_mem / (1024**3)
            mem_msg = f'GPU {0}: free mem {free_mem:.3f}G, total mem {total_mem:.3f}G \n'
            flash_res_data = flash_res.clone().detach().cpu()
            print('flash attn ', mem_msg,
                  f'Cost time per time {(time.time() - st) / run_num}s')

    # recommend this style for multi blocks to save the preprocess time.

    flash_input = flash_preprocess(
        x.clone(),
        context.clone(),
        mask=mask.clone() if mask is not None else mask)
    st = time.time()
    for i in tqdm(range(run_num)):
        improved_flash_res_v1 = improved_flash_attn(**flash_input).reshape(
            (batch_size, -1, hidden_dim))
        if i == 0:
            improved_flash_res_v1_data = improved_flash_res_v1.clone().detach(
            ).cpu()
        if i == run_num - 1:
            free_mem, total_mem = torch.cuda.mem_get_info(0)
            free_mem, total_mem = free_mem / (1024**3), total_mem / (1024**3)
            mem_msg = f'GPU {0}: free mem {free_mem:.3f}G, total mem {total_mem:.3f}G \n'
            print('improved flash attn ', mem_msg,
                  f'Cost time per time {(time.time() - st) / run_num}s')
    #
    print(pytorch_res_data, xformer_res_data, flash_res_data,
          improved_flash_res_v1_data)
    print(pytorch_res_data.shape, xformer_res_data.shape, flash_res_data.shape,
          improved_flash_res_v1_data.shape)
    print(
        torch.sum(pytorch_res_data) / (batch_size * query_length * hidden_dim),
        torch.sum(xformer_res_data) / (batch_size * query_length * hidden_dim),
        torch.sum(flash_res_data) / (batch_size * query_length * hidden_dim),
        torch.sum(improved_flash_res_v1_data) /
        (batch_size * query_length * hidden_dim))