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>
#include <vector>

void tanh_inplace_cuda(torch::Tensor t);

std::vector<torch::Tensor> rnn_forward_cuda(torch::Tensor x,
                                            torch::Tensor hidden,
                                            torch::Tensor w1,
                                            torch::Tensor b1,
                                            torch::Tensor w2,
                                            torch::Tensor b2,
                                            int64_t input_size) {
    TORCH_CHECK(x.is_cuda(), "x must be CUDA");
    TORCH_CHECK(hidden.is_cuda(), "hidden must be CUDA");
    TORCH_CHECK(w1.is_cuda() && b1.is_cuda() && w2.is_cuda() && b2.is_cuda(), "weights must be CUDA");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "only fp32 supported");

    // Avoid materializing torch.cat([x, hidden], 1).  The reference linear is
    // equivalent to two GEMMs accumulated into the same output.  This removes
    // a 256 x 32768 temporary and one full memory pass over it.
    auto wx = w1.narrow(1, 0, input_size);
    auto wh = w1.narrow(1, input_size, w1.size(1) - input_size);

    auto hnew = at::addmm(b1, x, wx.t());
    hnew.addmm_(hidden, wh.t(), 1.0, 1.0);

    tanh_inplace_cuda(hnew);

    auto out = at::addmm(b2, hnew, w2.t());
    return {out, hnew};
}
'''

_cuda_src = r'''
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAException.h>
#include <stdint.h>

__global__ void tanh_inplace_kernel(float* __restrict__ p, int64_t n) {
    int64_t i = (int64_t)blockIdx.x * blockDim.x + threadIdx.x;
    int64_t stride = (int64_t)blockDim.x * gridDim.x;
    for (; i < n; i += stride) {
        p[i] = tanhf(p[i]);
    }
}

void tanh_inplace_cuda(torch::Tensor t) {
    const int64_t n = t.numel();
    if (n == 0) return;
    const int threads = 256;
    int blocks = (int)((n + threads - 1) / threads);
    if (blocks > 65535) blocks = 65535;
    tanh_inplace_kernel<<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(t.data_ptr<float>(), n);
    C10_CUDA_KERNEL_LAUNCH_CHECK();
}
'''

_ext = load_inline(
    name="rnn_split_tanh_ext",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["rnn_forward_cuda"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3", "--use_fast_math"],
    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.tanh = nn.Tanh()

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

        # Ensure the common fast path: benchmark inputs and module parameters are
        # contiguous fp32 CUDA tensors.  The extension performs all heavy work and
        # returns both the RNN output and the updated hidden state.
        out, hnew = _ext.rnn_forward_cuda(
            x,
            self.hidden,
            self.i2h.weight,
            self.i2h.bias,
            self.h2o.weight,
            self.h2o.bias,
            self.input_size,
        )
        self.hidden = hnew
        return out