import torch
import torch.nn as nn

class Model(nn.Module):
    def __init__(self, input_size, hidden_size, num_layers, output_size, dropout=0.0):
        """
        Initialize the LSTM model.

        :param input_size: The number of expected features in the input `x`
        :param hidden_size: The number of features in the hidden state `h`
        :param num_layers: Number of recurrent layers
        :param output_size: The number of output features
        :param dropout: If non-zero, introduces a Dropout layer on the outputs of each LSTM layer except the last layer, with dropout probability equal to `dropout`
        """
        super(Model, self).__init__()
        # Initialize hidden state with random values
        self.lstm = nn.LSTM(input_size, hidden_size, num_layers, batch_first=True, dropout=dropout, bidirectional=False)
        self.fc = nn.Linear(hidden_size, output_size)
    
    def forward(self, x,h0,c0):
        """
        Forward pass through the LSTM model.

        :param x: The input tensor, shape (batch_size, sequence_length, input_size)
        :return: The output tensor, shape (batch_size, sequence_length, output_size)
        """
        
        # Forward propagate LSTM
        out, state = self.lstm(x, (h0, c0))  # out: tensor of shape (batch_size, seq_length, hidden_size)
        
        # Decode the hidden state of the last time step
        out = self.fc(out[:, -1, :])  # out: tensor of shape (batch_size, output_size)
        
        return state[0]

# Test code
batch_size = 10
sequence_length = 512
input_size = 128
hidden_size = 256
num_layers = 6
output_size = 10
dropout = 0.0

def get_inputs():
    return [torch.rand(batch_size, sequence_length, input_size),torch.rand((num_layers, batch_size, hidden_size)),torch.rand((num_layers, batch_size, hidden_size))]

def get_init_inputs():
    return [input_size, hidden_size, num_layers, output_size, dropout]

import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline
_cpp_src = r'''
#include <torch/extension.h>
torch::Tensor lstm_forward(torch::Tensor x, torch::Tensor h0, torch::Tensor c0,
    torch::Tensor wih0, torch::Tensor whh0, torch::Tensor bih0, torch::Tensor bhh0,
    torch::Tensor wih1, torch::Tensor whh1, torch::Tensor bih1, torch::Tensor bhh1,
    torch::Tensor wih2, torch::Tensor whh2, torch::Tensor bih2, torch::Tensor bhh2,
    torch::Tensor wih3, torch::Tensor whh3, torch::Tensor bih3, torch::Tensor bhh3,
    torch::Tensor wih4, torch::Tensor whh4, torch::Tensor bih4, torch::Tensor bhh4,
    torch::Tensor wih5, torch::Tensor whh5, torch::Tensor bih5, torch::Tensor bhh5);
'''
_cuda_src = r'''
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <cooperative_groups.h>
#include <math.h>
namespace cg = cooperative_groups;
#define HSIZE 256
#define BSIZE 10
#define TSIZE 512
#define LAYERS 6
#define THREADS 64
#define HGRP 4
__device__ __forceinline__ float sigm(float x) { return 1.0f/(1.0f+expf(-x)); }
__device__ __forceinline__ float warp_sum(float v){
    v += __shfl_down_sync(0xffffffff, v, 16);
    v += __shfl_down_sync(0xffffffff, v, 8);
    v += __shfl_down_sync(0xffffffff, v, 4);
    v += __shfl_down_sync(0xffffffff, v, 2);
    v += __shfl_down_sync(0xffffffff, v, 1);
    return v;
}
__global__ void lstm_kernel(const float* __restrict__ x, const float* __restrict__ h0, const float* __restrict__ c0,
 const float* __restrict__ wih0,const float* __restrict__ whh0,const float* __restrict__ bih0,const float* __restrict__ bhh0,
 const float* __restrict__ wih1,const float* __restrict__ whh1,const float* __restrict__ bih1,const float* __restrict__ bhh1,
 const float* __restrict__ wih2,const float* __restrict__ whh2,const float* __restrict__ bih2,const float* __restrict__ bhh2,
 const float* __restrict__ wih3,const float* __restrict__ whh3,const float* __restrict__ bih3,const float* __restrict__ bhh3,
 const float* __restrict__ wih4,const float* __restrict__ whh4,const float* __restrict__ bih4,const float* __restrict__ bhh4,
 const float* __restrict__ wih5,const float* __restrict__ whh5,const float* __restrict__ bih5,const float* __restrict__ bhh5,
 float* __restrict__ workspace, float* __restrict__ hbuf, float* __restrict__ hn) {
    cg::grid_group grid = cg::this_grid();
    int groups = HSIZE/HGRP;
    int b = blockIdx.x / groups;
    int hbase = (blockIdx.x - b*groups) * HGRP;
    int tid = threadIdx.x, lane=tid&31, warp=tid>>5;
    __shared__ float sh[4][2];
    float cval[HGRP];
    for(int layer=0; layer<LAYERS; ++layer){
        const float *wih,*whh,*bih,*bhh; int K;
        if(layer==0){wih=wih0;whh=whh0;bih=bih0;bhh=bhh0;K=128;}
        else if(layer==1){wih=wih1;whh=whh1;bih=bih1;bhh=bhh1;K=HSIZE;}
        else if(layer==2){wih=wih2;whh=whh2;bih=bih2;bhh=bhh2;K=HSIZE;}
        else if(layer==3){wih=wih3;whh=whh3;bih=bih3;bhh=bhh3;K=HSIZE;}
        else if(layer==4){wih=wih4;whh=whh4;bih=bih4;bhh=bhh4;K=HSIZE;}
        else {wih=wih5;whh=whh5;bih=bih5;bhh=bhh5;K=HSIZE;}
        int hcur=layer&1, hnxt=hcur^1;
        if(tid==0){
            #pragma unroll
            for(int q=0;q<HGRP;++q){ int h=hbase+q; hbuf[(hcur*BSIZE+b)*HSIZE+h]=h0[(layer*BSIZE+b)*HSIZE+h]; cval[q]=c0[(layer*BSIZE+b)*HSIZE+h]; }
        }
        grid.sync();
        int inbuf=(layer-1)&1, outbuf=layer&1;
        for(int t=0;t<TSIZE;++t){
            #pragma unroll
            for(int q=0;q<HGRP;++q){
                int h=hbase+q;
                float si=0.f,sf=0.f,sg=0.f,so=0.f;
                for(int k=tid;k<K;k+=THREADS){
                    float xv = (layer==0) ? x[(b*TSIZE+t)*128+k] : workspace[((inbuf*TSIZE+t)*BSIZE+b)*HSIZE+k];
                    si += xv*wih[h*K+k]; sf += xv*wih[(HSIZE+h)*K+k];
                    sg += xv*wih[(2*HSIZE+h)*K+k]; so += xv*wih[(3*HSIZE+h)*K+k];
                }
                for(int k=tid;k<HSIZE;k+=THREADS){
                    float hv=hbuf[(hcur*BSIZE+b)*HSIZE+k];
                    si += hv*whh[h*HSIZE+k]; sf += hv*whh[(HSIZE+h)*HSIZE+k];
                    sg += hv*whh[(2*HSIZE+h)*HSIZE+k]; so += hv*whh[(3*HSIZE+h)*HSIZE+k];
                }
                si=warp_sum(si); sf=warp_sum(sf); sg=warp_sum(sg); so=warp_sum(so);
                if(lane==0){ sh[0][warp]=si; sh[1][warp]=sf; sh[2][warp]=sg; sh[3][warp]=so; }
                __syncthreads();
                if(tid==0){
                    float ai=sh[0][0]+sh[0][1], af=sh[1][0]+sh[1][1], ag=sh[2][0]+sh[2][1], ao=sh[3][0]+sh[3][1];
                    ai += bih[h]+bhh[h]; af += bih[HSIZE+h]+bhh[HSIZE+h]; ag += bih[2*HSIZE+h]+bhh[2*HSIZE+h]; ao += bih[3*HSIZE+h]+bhh[3*HSIZE+h];
                    float ig=sigm(ai), fg=sigm(af), gg=tanhf(ag), og=sigm(ao);
                    cval[q]=fg*cval[q]+ig*gg; float hnew=og*tanhf(cval[q]);
                    hbuf[(hnxt*BSIZE+b)*HSIZE+h]=hnew; workspace[((outbuf*TSIZE+t)*BSIZE+b)*HSIZE+h]=hnew;
                }
                __syncthreads();
            }
            grid.sync(); int tmp=hcur; hcur=hnxt; hnxt=tmp;
        }
        if(tid==0){
            #pragma unroll
            for(int q=0;q<HGRP;++q){ int h=hbase+q; hn[(layer*BSIZE+b)*HSIZE+h]=hbuf[(hcur*BSIZE+b)*HSIZE+h]; }
        }
        grid.sync();
    }
}

torch::Tensor lstm_forward(torch::Tensor x, torch::Tensor h0, torch::Tensor c0,
    torch::Tensor wih0, torch::Tensor whh0, torch::Tensor bih0, torch::Tensor bhh0,
    torch::Tensor wih1, torch::Tensor whh1, torch::Tensor bih1, torch::Tensor bhh1,
    torch::Tensor wih2, torch::Tensor whh2, torch::Tensor bih2, torch::Tensor bhh2,
    torch::Tensor wih3, torch::Tensor whh3, torch::Tensor bih3, torch::Tensor bhh3,
    torch::Tensor wih4, torch::Tensor whh4, torch::Tensor bih4, torch::Tensor bhh4,
    torch::Tensor wih5, torch::Tensor whh5, torch::Tensor bih5, torch::Tensor bhh5){
    auto hn=torch::empty({LAYERS,BSIZE,HSIZE}, x.options());
    auto workspace=torch::empty({2,TSIZE,BSIZE,HSIZE}, x.options());
    auto hbuf=torch::empty({2,BSIZE,HSIZE}, x.options());
    const float *xp=x.data_ptr<float>(), *h0p=h0.data_ptr<float>(), *c0p=c0.data_ptr<float>();
    const float *pwih0=wih0.data_ptr<float>(),*pwhh0=whh0.data_ptr<float>(),*pbih0=bih0.data_ptr<float>(),*pbhh0=bhh0.data_ptr<float>();
    const float *pwih1=wih1.data_ptr<float>(),*pwhh1=whh1.data_ptr<float>(),*pbih1=bih1.data_ptr<float>(),*pbhh1=bhh1.data_ptr<float>();
    const float *pwih2=wih2.data_ptr<float>(),*pwhh2=whh2.data_ptr<float>(),*pbih2=bih2.data_ptr<float>(),*pbhh2=bhh2.data_ptr<float>();
    const float *pwih3=wih3.data_ptr<float>(),*pwhh3=whh3.data_ptr<float>(),*pbih3=bih3.data_ptr<float>(),*pbhh3=bhh3.data_ptr<float>();
    const float *pwih4=wih4.data_ptr<float>(),*pwhh4=whh4.data_ptr<float>(),*pbih4=bih4.data_ptr<float>(),*pbhh4=bhh4.data_ptr<float>();
    const float *pwih5=wih5.data_ptr<float>(),*pwhh5=whh5.data_ptr<float>(),*pbih5=bih5.data_ptr<float>(),*pbhh5=bhh5.data_ptr<float>();
    float *wsp=workspace.data_ptr<float>(), *hbp=hbuf.data_ptr<float>(), *hnp=hn.data_ptr<float>();
    void* args[]={&xp,&h0p,&c0p,&pwih0,&pwhh0,&pbih0,&pbhh0,&pwih1,&pwhh1,&pbih1,&pbhh1,&pwih2,&pwhh2,&pbih2,&pbhh2,&pwih3,&pwhh3,&pbih3,&pbhh3,&pwih4,&pwhh4,&pbih4,&pbhh4,&pwih5,&pwhh5,&pbih5,&pbhh5,&wsp,&hbp,&hnp};
    cudaLaunchCooperativeKernel((void*)lstm_kernel, dim3(BSIZE*(HSIZE/HGRP)), dim3(THREADS), args, 0, at::cuda::getCurrentCUDAStream());
    C10_CUDA_KERNEL_LAUNCH_CHECK(); return hn;
}
'''
_ext = load_inline(name='lstm_coop_ext_submit', cpp_sources=[_cpp_src], cuda_sources=[_cuda_src], functions=['lstm_forward'], extra_cuda_cflags=['-O3'], extra_cflags=['-O3'], verbose=False)
class ModelNew(nn.Module):
    def __init__(self,input_size,hidden_size,num_layers,output_size,dropout=0.0):
        super().__init__(); self.lstm=nn.LSTM(input_size,hidden_size,num_layers,batch_first=True,dropout=dropout,bidirectional=False); self.fc=nn.Linear(hidden_size,output_size)
    def forward(self,x,h0,c0):
        return _ext.lstm_forward(x.contiguous(),h0.contiguous(),c0.contiguous(),
            self.lstm.weight_ih_l0.contiguous(),self.lstm.weight_hh_l0.contiguous(),self.lstm.bias_ih_l0.contiguous(),self.lstm.bias_hh_l0.contiguous(),
            self.lstm.weight_ih_l1.contiguous(),self.lstm.weight_hh_l1.contiguous(),self.lstm.bias_ih_l1.contiguous(),self.lstm.bias_hh_l1.contiguous(),
            self.lstm.weight_ih_l2.contiguous(),self.lstm.weight_hh_l2.contiguous(),self.lstm.bias_ih_l2.contiguous(),self.lstm.bias_hh_l2.contiguous(),
            self.lstm.weight_ih_l3.contiguous(),self.lstm.weight_hh_l3.contiguous(),self.lstm.bias_ih_l3.contiguous(),self.lstm.bias_hh_l3.contiguous(),
            self.lstm.weight_ih_l4.contiguous(),self.lstm.weight_hh_l4.contiguous(),self.lstm.bias_ih_l4.contiguous(),self.lstm.bias_hh_l4.contiguous(),
            self.lstm.weight_ih_l5.contiguous(),self.lstm.weight_hh_l5.contiguous(),self.lstm.bias_ih_l5.contiguous(),self.lstm.bias_hh_l5.contiguous())