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_src = r'''
#include <torch/extension.h>
#include <ATen/ops/addmm.h>

void sigmoid_scale_add_launcher(torch::Tensor y, double scale);

torch::Tensor fused_forward(torch::Tensor input, torch::Tensor weight, torch::Tensor bias, double scale) {
    // Equivalent to torch.nn.functional.linear(input, weight, bias), using cuBLAS via ATen.
    auto y = at::addmm(bias, input, weight.t(), 1.0, 1.0);
    sigmoid_scale_add_launcher(y, scale);
    return y;
}
'''

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

static __forceinline__ __device__ float post_op(float x, float scale) {
    return x + scale / (1.0f + __expf(-x));
}

__global__ void sigmoid_scale_add_kernel(float* __restrict__ y, long long n, float scale) {
    long long i = (long long)blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n) {
        float v = y[i];
        y[i] = post_op(v, scale);
    }
}

__global__ void sigmoid_scale_add_vec4_kernel(float4* __restrict__ y4, long long n4, float scale) {
    long long i = (long long)blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n4) {
        float4 v = y4[i];
        v.x = post_op(v.x, scale);
        v.y = post_op(v.y, scale);
        v.z = post_op(v.z, scale);
        v.w = post_op(v.w, scale);
        y4[i] = v;
    }
}

void sigmoid_scale_add_launcher(torch::Tensor y, double scale_d) {
    const long long n = (long long)y.numel();
    if (n == 0) return;
    const float scale = (float)scale_d;
    constexpr int threads = 256;
    auto stream = at::cuda::getCurrentCUDAStream();
    float* ptr = y.data_ptr<float>();
    uintptr_t addr = reinterpret_cast<uintptr_t>(ptr);
    if ((addr & 15) == 0 && (n & 3LL) == 0) {
        long long n4 = n >> 2;
        int blocks = (int)((n4 + threads - 1) / threads);
        sigmoid_scale_add_vec4_kernel<<<blocks, threads, 0, stream>>>(reinterpret_cast<float4*>(ptr), n4, scale);
    } else {
        int blocks = (int)((n + threads - 1) / threads);
        sigmoid_scale_add_kernel<<<blocks, threads, 0, stream>>>(ptr, n, scale);
    }
    C10_CUDA_KERNEL_LAUNCH_CHECK();
}
'''

_ext = load_inline(
    name='gemm_sigmoid_scale_residual_ext',
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=['fused_forward'],
    extra_cflags=['-O3'],
    extra_cuda_cflags=['-O3'],
    with_cuda=True,
    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 = float(scaling_factor)

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