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

_cpp_src = r"""
#include <torch/extension.h>
void concat_cols(torch::Tensor x, torch::Tensor h, torch::Tensor out);
void tanh_inplace(torch::Tensor y);
"""

_cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda.h>
#include <cuda_runtime.h>

__global__ void tanh_kernel(float* __restrict__ y, long total) {
    long idx = blockIdx.x * blockDim.x + threadIdx.x;
    long stride = (long)blockDim.x * gridDim.x;
    for (; idx < total; idx += stride) {
        y[idx] = tanhf(y[idx]);
    }
}

void concat_cols(torch::Tensor x, torch::Tensor h, torch::Tensor out) {
    int rows = (int)x.size(0);
    size_t c1 = (size_t)x.size(1);
    size_t c2 = (size_t)h.size(1);
    size_t cols = c1 + c2;
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    cudaMemcpy2DAsync(out.data_ptr<float>(), cols * sizeof(float),
                      x.data_ptr<float>(), c1 * sizeof(float),
                      c1 * sizeof(float), rows, cudaMemcpyDeviceToDevice, stream);
    cudaMemcpy2DAsync(out.data_ptr<float>() + c1, cols * sizeof(float),
                      h.data_ptr<float>(), c2 * sizeof(float),
                      c2 * sizeof(float), rows, cudaMemcpyDeviceToDevice, stream);
}

void tanh_inplace(torch::Tensor y) {
    long total = y.numel();
    int threads = 256;
    int blocks = (int)((total + threads - 1) / threads);
    if (blocks > 4096) blocks = 4096;
    tanh_kernel<<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(y.data_ptr<float>(), total);
}
"""

_ext = load_inline(
    name="rnn_memcpy2d_tanh_ext_final",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["concat_cols", "tanh_inplace"],
    extra_cuda_cflags=["-O3"],
    verbose=False,
)

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((256, hidden_size))
        self.i2h = nn.Linear(input_size + hidden_size, hidden_size)
        self.h2o = nn.Linear(hidden_size, output_size)
        self._combined = None

    def forward(self, x: torch.Tensor, initial_hidden=None) -> torch.Tensor:
        h = initial_hidden.to(x.device) if initial_hidden is not None else self.hidden.to(x.device)

        if self._combined is None or self._combined.device != x.device or self._combined.shape[0] != x.shape[0]:
            self._combined = torch.empty((x.size(0), self.input_size + self.hidden_size), device=x.device, dtype=x.dtype)

        _ext.concat_cols(x, h, self._combined)
        y = torch.addmm(self.i2h.bias, self._combined, self.i2h.weight.t())
        _ext.tanh_inplace(y)
        self.hidden = y
        return torch.addmm(self.h2o.bias, self.hidden, self.h2o.weight.t())