import torch
import torch.nn as nn

class Model(nn.Module):
    def __init__(self, input_size: int, hidden_size: int, output_size: int):
        """
        Initialize the Vanilla RNN model.
        
        :param input_size: The number of input features (int).
        :param hidden_size: The size of the hidden state (int).
        :param output_size: The number of output features (int).
        """
        super(Model, self).__init__()
        self.input_size = input_size
        self.hidden_size = hidden_size
        self.output_size = output_size
        self.hidden = torch.randn((batch_size, hidden_size))
        
        # Define the RNN cell components (input to hidden, hidden to hidden, and hidden to output)
        self.i2h = nn.Linear(input_size + hidden_size, hidden_size)  # Input to hidden
        self.h2o = nn.Linear(hidden_size, output_size)  # Hidden to output
        self.tanh = nn.Tanh()  # Activation function for hidden state
    
    def forward(self, x: torch.Tensor, initial_hidden=None) -> torch.Tensor:
        """
        Forward pass of the Vanilla RNN.
        
        :param x: Input tensor of shape (batch_size, input_size).
        :param hidden: Hidden state tensor of shape (batch_size, hidden_size).
        :return: Output tensor of shape (batch_size, output_size), and the new hidden state.
        """
        if initial_hidden is not None:
            self.hidden.copy_(initial_hidden)
        self.hidden = self.hidden.to(x.device)
        combined = torch.cat((x, self.hidden), dim=1)  # Concatenate input and hidden state
        self.hidden = self.tanh(self.i2h(combined))  # Update hidden state
        output = self.h2o(self.hidden)  # Compute output
        return output

batch_size = 256
input_size = 16384
hidden_size = 16384
output_size = 8192
sequence_length = 256

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

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

import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline

torch.backends.cuda.matmul.allow_tf32 = True

_CPP = r'''
#include <torch/extension.h>
#include <vector>

void tanh_inplace_cuda(torch::Tensor a);

std::vector<torch::Tensor> rnn_forward(torch::Tensor x, torch::Tensor hidden_in,
                                       torch::Tensor i2h_weight, torch::Tensor i2h_bias,
                                       torch::Tensor h2o_weight, torch::Tensor h2o_bias) {
    auto combined = torch::cat({x, hidden_in}, 1);
    auto hidden = torch::addmm(i2h_bias, combined, i2h_weight.t());
    tanh_inplace_cuda(hidden);
    auto out = torch::addmm(h2o_bias, hidden, h2o_weight.t());
    return {out, hidden};
}
'''

_CUDA = r'''
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>

__global__ void tanh_kernel(float* p, int n) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < n) p[idx] = tanhf(p[idx]);
}

void tanh_inplace_cuda(torch::Tensor a) {
    int n = (int)a.numel();
    tanh_kernel<<<(n + 255) / 256, 256, 0, at::cuda::getCurrentCUDAStream()>>>(a.data_ptr<float>(), n);
}
'''

_ext = load_inline(
    name='rnn_ext_submit_tf32_fused_tanh',
    cpp_sources=_CPP,
    cuda_sources=_CUDA,
    functions=['rnn_forward'],
    extra_cflags=['-O3'],
    extra_cuda_cflags=['-O3'],
    verbose=False,
)

batch_size = 256

class ModelNew(nn.Module):
    def __init__(self, input_size: int, hidden_size: int, output_size: int):
        super(ModelNew, self).__init__()
        self.input_size = input_size
        self.hidden_size = hidden_size
        self.output_size = output_size
        self.hidden = torch.randn((batch_size, hidden_size))
        self.i2h = nn.Linear(input_size + hidden_size, hidden_size)
        self.h2o = nn.Linear(hidden_size, output_size)
        self.tanh = nn.Tanh()

    def forward(self, x: torch.Tensor, initial_hidden=None) -> torch.Tensor:
        h = initial_hidden if initial_hidden is not None else self.hidden.to(x.device)
        out, new_hidden = _ext.rnn_forward(x, h, self.i2h.weight, self.i2h.bias, self.h2o.weight, self.h2o.bias)
        self.hidden = new_hidden
        return out