import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Model implementing the pattern "Gemm_Sigmoid_Scaling_ResidualAdd".
    """
    def __init__(self, input_size, hidden_size, scaling_factor):
        super(Model, self).__init__()
        self.gemm = nn.Linear(input_size, hidden_size)
        self.scaling_factor = scaling_factor

    def forward(self, x):
        """
        Forward pass of the model.

        Args:
            x (torch.Tensor): Input tensor of shape (batch_size, input_size).

        Returns:
            torch.Tensor: Output tensor of shape (batch_size, hidden_size).
        """
        x = self.gemm(x)
        original_x = x
        x = torch.sigmoid(x)
        x = x * self.scaling_factor
        x = x + original_x
        return x

batch_size = 1024
input_size = 8192
hidden_size = 8192
scaling_factor = 2.0

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

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

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

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

torch::Tensor fused_linear_sigmoid_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, double scale);
'''

_cuda_source = r'''
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <cuda_runtime.h>
#include <cublas_v2.h>
#include <climits>

__device__ __forceinline__ float act(float z, float scale) {
    return z + scale * (1.0f / (1.0f + __expf(-z)));
}

__global__ void epilogue_vec4_kernel(float* __restrict__ y, const float* __restrict__ bias, int m, int n4, float scale) {
    int col4 = blockIdx.x * blockDim.x + threadIdx.x;
    int row = blockIdx.y;
    if (row < m && col4 < n4) {
        int base4 = row * n4 + col4;
        float4 v = reinterpret_cast<float4*>(y)[base4];
        float4 b = reinterpret_cast<const float4*>(bias)[col4];
        v.x = act(v.x + b.x, scale);
        v.y = act(v.y + b.y, scale);
        v.z = act(v.z + b.z, scale);
        v.w = act(v.w + b.w, scale);
        reinterpret_cast<float4*>(y)[base4] = v;
    }
}

__global__ void epilogue_scalar_kernel(float* __restrict__ y, const float* __restrict__ bias, int m, int n, float scale) {
    int col = blockIdx.x * blockDim.x + threadIdx.x;
    int row = blockIdx.y;
    if (row < m && col < n) {
        int idx = row * n + col;
        float z = y[idx] + bias[col];
        y[idx] = act(z, scale);
    }
}

torch::Tensor fused_linear_sigmoid_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, double scale_d) {
    TORCH_CHECK(x.is_cuda() && weight.is_cuda() && bias.is_cuda(), "all tensors must be CUDA");
    TORCH_CHECK(x.scalar_type() == torch::kFloat32 && weight.scalar_type() == torch::kFloat32 && bias.scalar_type() == torch::kFloat32, "float32 only");
    TORCH_CHECK(x.dim() == 2 && weight.dim() == 2 && bias.dim() == 1, "bad dims");
    const int64_t m64 = x.size(0);
    const int64_t k64 = x.size(1);
    const int64_t n64 = weight.size(0);
    TORCH_CHECK(weight.size(1) == k64 && bias.size(0) == n64, "bad sizes");
    TORCH_CHECK(x.is_contiguous() && weight.is_contiguous() && bias.is_contiguous(), "tensors must be contiguous");
    TORCH_CHECK(m64 <= INT_MAX && n64 <= INT_MAX && k64 <= INT_MAX, "sizes too large");
    int m = (int)m64, n = (int)n64, k = (int)k64;
    c10::cuda::CUDAGuard device_guard(x.device());
    auto y = torch::empty({m64, n64}, x.options());

    cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle();
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    cublasSetStream(handle, stream);
    const float alpha = 1.0f;
    const float beta = 0.0f;
    cublasStatus_t stat = cublasSgemm(handle, CUBLAS_OP_T, CUBLAS_OP_N,
                                      n, m, k,
                                      &alpha,
                                      weight.data_ptr<float>(), k,
                                      x.data_ptr<float>(), k,
                                      &beta,
                                      y.data_ptr<float>(), n);
    TORCH_CHECK(stat == CUBLAS_STATUS_SUCCESS, "cublasSgemm failed");

    float scale = (float)scale_d;
    if ((n & 3) == 0) {
        int n4 = n >> 2;
        dim3 block(256);
        dim3 grid((n4 + block.x - 1) / block.x, m);
        epilogue_vec4_kernel<<<grid, block, 0, stream>>>(y.data_ptr<float>(), bias.data_ptr<float>(), m, n4, scale);
    } else {
        dim3 block(256);
        dim3 grid((n + block.x - 1) / block.x, m);
        epilogue_scalar_kernel<<<grid, block, 0, stream>>>(y.data_ptr<float>(), bias.data_ptr<float>(), m, n, scale);
    }
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return y;
}
'''

_ext = load_inline(
    name='fused_linear_sigmoid_residual_ext_final',
    cpp_sources=_cpp_source,
    cuda_sources=_cuda_source,
    functions=['fused_linear_sigmoid_forward'],
    extra_cuda_cflags=['-O3', '--use_fast_math'],
    extra_cflags=['-O3'],
    extra_ldflags=['-lcublas'],
    verbose=False,
)

class ModelNew(nn.Module):
    def __init__(self, input_size, hidden_size, scaling_factor):
        super(ModelNew, self).__init__()
        self.gemm = nn.Linear(input_size, hidden_size)
        self.scaling_factor = scaling_factor

    def forward(self, x):
        return _ext.fused_linear_sigmoid_forward(x, self.gemm.weight, self.gemm.bias, float(self.scaling_factor))