import torch
import triton
import triton.language as tl
from torch._inductor.runtime.triton_heuristics import grid
from torch._C import _cuda_getCurrentRawStream as get_raw_stream
import torch.nn.functional as F
import torch.nn as nn
from torch.nn.utils import weight_norm
from torch.optim.lr_scheduler import *
assert_size_stride = torch._C._dynamo.guards.assert_size_stride
empty_strided_cuda = torch._C._dynamo.guards._empty_strided_cuda
reinterpret_tensor = torch._C._dynamo.guards._reinterpret_tensor


@triton.jit
def triton_poi_fused_0(in_ptr0, out_ptr0, xnumel, XBLOCK: tl.constexpr):
    xnumel = 64
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x3 = xindex
    x0 = xindex % 4
    x1 = xindex // 4 % 4
    x2 = xindex // 16
    tmp0 = tl.load(in_ptr0 + x3, xmask)
    tl.store(out_ptr0 + (x0 + 4 * x2 + 16 * x1), tmp0, xmask)


def call(args):
    arg0_1, arg1_1, arg2_1 = args
    args.clear()
    assert_size_stride(arg0_1, (4, 4, 4), (16, 4, 1))
    assert_size_stride(arg1_1, (4, 4, 4), (16, 4, 1))
    assert_size_stride(arg2_1, (4, 4, 4), (16, 4, 1))
    with torch.cuda._DeviceGuard(0):
        torch.cuda.set_device(0)
        buf0 = empty_strided_cuda((1, 4, 4, 4), (64, 4, 16, 1), torch.float32)
        get_raw_stream(0)
        triton_poi_fused_0[grid(64)](arg0_1, buf0, 64, XBLOCK=64, num_warps
            =1, num_stages=1)
        del arg0_1
        buf1 = empty_strided_cuda((1, 4, 4, 4), (64, 4, 16, 1), torch.float32)
        triton_poi_fused_0[grid(64)](arg1_1, buf1, 64, XBLOCK=64, num_warps
            =1, num_stages=1)
        del arg1_1
        buf2 = empty_strided_cuda((1, 4, 4, 4), (64, 4, 16, 1), torch.float32)
        triton_poi_fused_0[grid(64)](arg2_1, buf2, 64, XBLOCK=64, num_warps
            =1, num_stages=1)
        del arg2_1
        buf3 = torch.ops.aten._scaled_dot_product_efficient_attention.default(
            buf0, buf1, buf2, None, False, scale=1.0)
        del buf0
        del buf1
        del buf2
        buf4 = buf3[0]
        del buf3
    return reinterpret_tensor(buf4, (4, 4, 4), (4, 16, 1), 0),


def linear(x):
    return x


def activation(func_a):
    """Activation function wrapper
    """
    try:
        f = eval(func_a)
    except:
        f = linear
    return f


class DropoutWrapper(nn.Module):
    """
    This is a dropout wrapper which supports the fix mask dropout
    """

    def __init__(self, dropout_p=0, enable_vbp=True):
        super(DropoutWrapper, self).__init__()
        """variational dropout means fix dropout mask
        ref: https://discuss.pytorch.org/t/dropout-for-rnns/633/11
        """
        self.enable_variational_dropout = enable_vbp
        self.dropout_p = dropout_p

    def forward(self, x):
        """
            :param x: batch * len * input_size
        """
        if self.training is False or self.dropout_p == 0:
            return x
        if len(x.size()) == 3:
            mask = 1.0 / (1 - self.dropout_p) * torch.bernoulli((1 - self.
                dropout_p) * (x.data.new(x.size(0), x.size(2)).zero_() + 1))
            mask.requires_grad = False
            return mask.unsqueeze(1).expand_as(x) * x
        else:
            return F.dropout(x, p=self.dropout_p, training=self.training)


class MultiheadAttentionWrapperNew(nn.Module):
    """Multi-headed attention.
    See "Attention Is All You Need" for more details.
    """

    def __init__(self, query_dim, key_dim, value_dim, prefix='attention',
        opt={}, dropout=None):
        super().__init__()
        self.prefix = prefix
        self.num_heads = opt.get('{}_head'.format(self.prefix), 1)
        self.dropout = DropoutWrapper(opt.get('{}_dropout'.format(self.
            prefix), 0)) if dropout is None else dropout
        self.qkv_dim = [query_dim, key_dim, value_dim]
        assert query_dim == key_dim, 'query dim must equal with key dim'
        self.hidden_size = opt.get('{}_hidden_size'.format(self.prefix), 64)
        self.proj_on = opt.get('{}_proj_on'.format(prefix), False)
        self.share = opt.get('{}_share'.format(self.prefix), False)
        self.layer_norm_on = opt.get('{}_norm_on'.format(self.prefix), False)
        self.scale_on = opt.get('{}_scale_on'.format(self.prefix), False)
        if self.proj_on:
            self.proj_modules = nn.ModuleList([nn.Linear(dim, self.
                hidden_size) for dim in self.qkv_dim[0:2]])
            if self.layer_norm_on:
                for proj in self.proj_modules:
                    proj = weight_norm(proj)
            if self.share and self.qkv_dim[0] == self.qkv_dim[1]:
                self.proj_modules[1] = self.proj_modules[0]
            self.f = activation(opt.get('{}_activation'.format(self.prefix),
                'relu'))
            self.qkv_head_dim = [self.hidden_size // self.num_heads] * 3
            self.qkv_head_dim[2] = value_dim // self.num_heads
            assert self.qkv_head_dim[0
                ] * self.num_heads == self.hidden_size, 'hidden size must be divisible by num_heads'
            assert self.qkv_head_dim[2
                ] * self.num_heads == value_dim, 'value size must be divisible by num_heads'
        else:
            self.qkv_head_dim = [(emb // self.num_heads) for emb in self.
                qkv_dim]
            assert self.qkv_head_dim[0] * self.num_heads == self.qkv_dim[0
                ], 'query size must be divisible by num_heads'
            assert self.qkv_head_dim[1] * self.num_heads == self.qkv_dim[1
                ], 'key size must be divisible by num_heads'
            assert self.qkv_head_dim[2] * self.num_heads == self.qkv_dim[2
                ], 'value size must be divisible by num_heads'
        if self.scale_on:
            self.scaling = self.qkv_head_dim[0] ** -0.5
        self.drop_diagonal = opt.get('{}_drop_diagonal'.format(self.prefix),
            False)
        self.output_size = self.qkv_dim[2]

    def forward(self, input_0, input_1, input_2):
        arg0_1 = input_0
        arg1_1 = input_1
        arg2_1 = input_2
        output = call([arg0_1, arg1_1, arg2_1])
        return output[0]


Measured latency_ms: 0.044483