import torch
from torch._inductor.select_algorithm import extern_kernels
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
from torch._inductor.runtime import triton_helpers
from torch._inductor.runtime.triton_helpers import math as tl_math
from torch import nn
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_clone_0(in_ptr0, out_ptr0, xnumel, XBLOCK: tl.constexpr):
    xnumel = 512
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x0 = xindex % 4
    x1 = xindex // 4 % 4
    x2 = xindex // 16 % 8
    x3 = xindex // 128
    x4 = xindex
    tmp0 = tl.load(in_ptr0 + (x0 + 4 * x2 + 32 * x1 + 128 * x3), xmask)
    tl.store(out_ptr0 + x4, tmp0, xmask)


@triton.jit
def triton_poi_fused__softmax_1(in_ptr0, out_ptr0, xnumel, XBLOCK: tl.constexpr
    ):
    xnumel = 512
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x2 = xindex
    x1 = xindex // 4
    tmp0 = tl.load(in_ptr0 + x2, xmask)
    tmp3 = tl.load(in_ptr0 + 4 * x1, xmask, eviction_policy='evict_last')
    tmp5 = tl.load(in_ptr0 + (1 + 4 * x1), xmask, eviction_policy='evict_last')
    tmp8 = tl.load(in_ptr0 + (2 + 4 * x1), xmask, eviction_policy='evict_last')
    tmp11 = tl.load(in_ptr0 + (3 + 4 * x1), xmask, eviction_policy='evict_last'
        )
    tmp1 = 1.0
    tmp2 = tmp0 * tmp1
    tmp4 = tmp3 * tmp1
    tmp6 = tmp5 * tmp1
    tmp7 = triton_helpers.maximum(tmp4, tmp6)
    tmp9 = tmp8 * tmp1
    tmp10 = triton_helpers.maximum(tmp7, tmp9)
    tmp12 = tmp11 * tmp1
    tmp13 = triton_helpers.maximum(tmp10, tmp12)
    tmp14 = tmp2 - tmp13
    tmp15 = 0.5
    tmp16 = tmp14 * tmp15
    tmp17 = tl_math.exp(tmp16)
    tl.store(out_ptr0 + x2, tmp17, xmask)


@triton.jit
def triton_poi_fused__softmax_2(in_ptr0, out_ptr0, xnumel, XBLOCK: tl.constexpr
    ):
    xnumel = 512
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x2 = xindex
    x1 = xindex // 4
    tmp0 = tl.load(in_ptr0 + x2, xmask)
    tmp1 = tl.load(in_ptr0 + 4 * x1, xmask, eviction_policy='evict_last')
    tmp2 = tl.load(in_ptr0 + (1 + 4 * x1), xmask, eviction_policy='evict_last')
    tmp4 = tl.load(in_ptr0 + (2 + 4 * x1), xmask, eviction_policy='evict_last')
    tmp6 = tl.load(in_ptr0 + (3 + 4 * x1), xmask, eviction_policy='evict_last')
    tmp3 = tmp1 + tmp2
    tmp5 = tmp3 + tmp4
    tmp7 = tmp5 + tmp6
    tmp8 = tmp0 / tmp7
    tl.store(out_ptr0 + x2, tmp8, xmask)


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


def call(args):
    primals_1, primals_2, primals_3, primals_4, primals_5, primals_6 = args
    args.clear()
    assert_size_stride(primals_1, (4, 4, 4), (16, 4, 1))
    assert_size_stride(primals_2, (32, 4), (4, 1))
    assert_size_stride(primals_3, (32, 4), (4, 1))
    assert_size_stride(primals_4, (32, 4), (4, 1))
    assert_size_stride(primals_5, (4, 32), (32, 1))
    assert_size_stride(primals_6, (4,), (1,))
    with torch.cuda._DeviceGuard(0):
        torch.cuda.set_device(0)
        buf0 = empty_strided_cuda((16, 32), (32, 1), torch.float32)
        extern_kernels.mm(reinterpret_tensor(primals_1, (16, 4), (4, 1), 0),
            reinterpret_tensor(primals_2, (4, 32), (1, 4), 0), out=buf0)
        del primals_2
        buf1 = empty_strided_cuda((16, 32), (32, 1), torch.float32)
        extern_kernels.mm(reinterpret_tensor(primals_1, (16, 4), (4, 1), 0),
            reinterpret_tensor(primals_3, (4, 32), (1, 4), 0), out=buf1)
        del primals_3
        buf2 = empty_strided_cuda((16, 32), (32, 1), torch.float32)
        extern_kernels.mm(reinterpret_tensor(primals_1, (16, 4), (4, 1), 0),
            reinterpret_tensor(primals_4, (4, 32), (1, 4), 0), out=buf2)
        del primals_4
        buf3 = empty_strided_cuda((4, 8, 4, 4), (128, 16, 4, 1), torch.float32)
        get_raw_stream(0)
        triton_poi_fused_clone_0[grid(512)](buf1, buf3, 512, XBLOCK=128,
            num_warps=4, num_stages=1)
        buf4 = reinterpret_tensor(buf1, (4, 8, 4, 4), (128, 16, 4, 1), 0)
        del buf1
        triton_poi_fused_clone_0[grid(512)](buf0, buf4, 512, XBLOCK=128,
            num_warps=4, num_stages=1)
        buf5 = reinterpret_tensor(buf0, (32, 4, 4), (16, 4, 1), 0)
        del buf0
        extern_kernels.bmm(reinterpret_tensor(buf3, (32, 4, 4), (16, 4, 1),
            0), reinterpret_tensor(buf4, (32, 4, 4), (16, 1, 4), 0), out=buf5)
        buf6 = empty_strided_cuda((32, 4, 4), (16, 4, 1), torch.float32)
        triton_poi_fused__softmax_1[grid(512)](buf5, buf6, 512, XBLOCK=128,
            num_warps=4, num_stages=1)
        buf7 = buf5
        del buf5
        triton_poi_fused__softmax_2[grid(512)](buf6, buf7, 512, XBLOCK=256,
            num_warps=4, num_stages=1)
        buf8 = reinterpret_tensor(buf6, (4, 8, 4, 4), (128, 16, 4, 1), 0)
        del buf6
        triton_poi_fused_clone_0[grid(512)](buf2, buf8, 512, XBLOCK=128,
            num_warps=4, num_stages=1)
        buf9 = reinterpret_tensor(buf2, (32, 4, 4), (16, 4, 1), 0)
        del buf2
        extern_kernels.bmm(buf7, reinterpret_tensor(buf8, (32, 4, 4), (16, 
            4, 1), 0), out=buf9)
        buf10 = empty_strided_cuda((4, 4, 8, 4), (128, 32, 4, 1), torch.float32
            )
        triton_poi_fused_clone_3[grid(512)](buf9, buf10, 512, XBLOCK=128,
            num_warps=4, num_stages=1)
        del buf9
        buf11 = empty_strided_cuda((16, 4), (4, 1), torch.float32)
        extern_kernels.addmm(primals_6, reinterpret_tensor(buf10, (16, 32),
            (32, 1), 0), reinterpret_tensor(primals_5, (32, 4), (1, 32), 0),
            alpha=1, beta=1, out=buf11)
        del primals_6
    return reinterpret_tensor(buf11, (4, 4, 4), (16, 4, 1), 0
        ), buf7, reinterpret_tensor(primals_1, (16, 4), (4, 1), 0
        ), buf7, reinterpret_tensor(buf10, (16, 32), (32, 1), 0
        ), primals_5, reinterpret_tensor(buf8, (32, 4, 4), (16, 1, 4), 0
        ), reinterpret_tensor(buf3, (32, 4, 4), (16, 1, 4), 0
        ), reinterpret_tensor(buf4, (32, 4, 4), (16, 4, 1), 0)


class MultiHeadAttentionNew(nn.Module):

    def __init__(self, emb_size, n_heads=8, mask=False):
        """
        Arguments:

        emb_size: Size of input Embeddings
        n_heads: Number of heads for MultiHead Attention

        Layers:

        tokeys: For Keys
        toquery: For Query
        tovalue: For Value

        combine_heads: Convert the ( (number of heads) * (emb_size) ) into (emb_size) for output

        Note:
            The input dimension we give to the Linear Layer (nn.Linear) is the last dimension (no. of columns) in size
            (batch_size ,seq_length, input_dims) ==> input_dims : give to nn.Linear
            So in order for the tensor to pass through linear layer having the last dim equal to input_dim of
            Linear layer is necessary.

        """
        super(MultiHeadAttentionNew, self).__init__()
        self.emb_size = emb_size
        self.n_heads = n_heads
        """
        Explanation Comment:
            For a single head attention the input will be of emb_size and output will be of emb_size Linear(emb_size,emb_size) but in multihead attention
            there are number of heads (n_heads) and each head has it's own key(k), query(q) and value(v) so total number of k,q,v
            will be equal to n_heads.
            We can do that in a single linear layer by giving the size of output layer = n_head * emb_size (here emb_size = output size)
            Same for k,q,v. nn.Linear(emb_size,n_heads*emb_size)

        """
        self.tokeys = nn.Linear(emb_size, n_heads * emb_size, bias=False)
        self.toquery = nn.Linear(emb_size, n_heads * emb_size, bias=False)
        self.tovalue = nn.Linear(emb_size, n_heads * emb_size, bias=False)
        self.combine_heads = nn.Linear(n_heads * emb_size, emb_size)
        self.Attnweights = None
        self.mask = mask

    def __MatrixMask(self, matrix, mask_value=float('-inf'), diagonal=False):
        """
        In self attention mask is used when we are trying to predict the next work based on the previous sequence
        but attention inherently contains information about all the words in the sequence.
        In order to accuretly predict the next word without sort of cheating by looking ahead with attention we set all
        the attention weights of the next words, tokens = -inf or zero, so that there is no input from the sequence ahead
        of the current word


        [ 1 2 3                 [ 1 0 0
          4 5 6    =========>     4 5 0
          7 8 9 ]                 7 8 9 ]

        """
        row, column = matrix.size(-2), matrix.size(-1)
        offset = 0 if diagonal else 1
        upper_triangular_idx = torch.triu_indices(row, column, offset)
        matrix[..., upper_triangular_idx[0], upper_triangular_idx[1]
            ] = mask_value
        return matrix

    def forward(self, input_0):
        primals_2 = self.tokeys.weight
        primals_3 = self.toquery.weight
        primals_4 = self.tovalue.weight
        primals_5 = self.combine_heads.weight
        primals_6 = self.combine_heads.bias
        primals_1 = input_0
        output = call([primals_1, primals_2, primals_3, primals_4,
            primals_5, primals_6])
        return output[0]


Measured latency_ms: 0.123553