import torch
import torch.nn as nn

class Model(nn.Module):
    """
    A model that performs a matrix multiplication, scaling, and residual addition.

    Args:
        in_features (int): Number of input features.
        out_features (int): Number of output features.
        scaling_factor (float): Scaling factor to apply after matrix multiplication.
    """
    def __init__(self, in_features, out_features, scaling_factor):
        super(Model, self).__init__()
        self.matmul = nn.Linear(in_features, out_features)
        self.scaling_factor = scaling_factor

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

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

        Returns:
            torch.Tensor: Output tensor of shape (batch_size, out_features).
        """
        x = self.matmul(x)
        original_x = x.clone().detach()
        x = x * self.scaling_factor
        x = x + original_x
        return x

batch_size = 16384
in_features = 4096
out_features = 4096
scaling_factor = 0.5

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

def get_init_inputs():
    return [in_features, out_features, 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/cuda/CUDAContext.h>
#include <ATen/cuda/CUDABlas.h>
#include <cublas_v2.h>

void scale_bias_cuda(at::Tensor y, const at::Tensor& bias, float s);

at::Tensor linear_scaled_forward(const at::Tensor& x, const at::Tensor& weight, const at::Tensor& bias, double scaling_factor) {
    const int64_t M64 = x.size(0), K64 = x.size(1), N64 = weight.size(0);
    at::Tensor y = at::empty({M64, N64}, x.options());
    int M = (int)M64, N = (int)N64, K = (int)K64;
    float alpha = 1.0f, beta = 0.0f;
    cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle();
    cublasSetStream(handle, at::cuda::getCurrentCUDAStream());
    // Row-major y = x * weight^T. cuBLAS is column-major, so compute y^T = weight * x^T.
    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);
    scale_bias_cuda(y, bias, static_cast<float>(1.0 + scaling_factor));
    return y;
}
'''

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

__global__ void scale_bias_kernel4_2d(float4* __restrict__ y4, const float4* __restrict__ b4, int rows, int out4, float s) {
    int c4 = blockIdx.x * blockDim.x + threadIdx.x;
    int r = blockIdx.y;
    if (c4 < out4) {
        long long idx = (long long)r * out4 + c4;
        float4 v = y4[idx];
        float4 b = b4[c4];
        v.x = (v.x + b.x) * s;
        v.y = (v.y + b.y) * s;
        v.z = (v.z + b.z) * s;
        v.w = (v.w + b.w) * s;
        y4[idx] = v;
    }
}

void scale_bias_cuda(at::Tensor y, const at::Tensor& bias, float s) {
    int rows = (int)y.size(0);
    int out = (int)y.size(1);
    int out4 = out / 4;
    dim3 block(256);
    dim3 grid((out4 + block.x - 1) / block.x, rows);
    auto stream = at::cuda::getCurrentCUDAStream();
    scale_bias_kernel4_2d<<<grid, block, 0, stream>>>(reinterpret_cast<float4*>(y.data_ptr<float>()), reinterpret_cast<const float4*>(bias.data_ptr<float>()), rows, out4, s);
}
'''

_ext = load_inline(
    name="linear_scaled_ext_final",
    cpp_sources=[cpp_src],
    cuda_sources=[cuda_src],
    functions=["linear_scaled_forward"],
    with_cuda=True,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
    extra_ldflags=["-lcublas"],
    verbose=False,
)

class ModelNew(nn.Module):
    def __init__(self, in_features, out_features, scaling_factor):
        super(ModelNew, self).__init__()
        self.matmul = nn.Linear(in_features, out_features)
        self.scaling_factor = float(scaling_factor)
    def forward(self, x):
        return _ext.linear_scaled_forward(x, self.matmul.weight, self.matmul.bias, self.scaling_factor)