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>

torch::Tensor fused_linear_scale_residual_forward(torch::Tensor x,
                                                  torch::Tensor weight,
                                                  torch::Tensor bias,
                                                  double scaling_factor);
"""

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

#define CHECK_CUBLAS(call) do {                                      \
    cublasStatus_t _status = (call);                                 \
    TORCH_CHECK(_status == CUBLAS_STATUS_SUCCESS,                    \
                "cuBLAS error: ", static_cast<int>(_status));        \
} while (0)

__global__ void fill_scaled_bias_kernel(float* __restrict__ out,
                                         const float* __restrict__ bias,
                                         int64_t total,
                                         int64_t n,
                                         float scale) {
    int64_t idx = (int64_t)blockIdx.x * blockDim.x + threadIdx.x;
    int64_t stride = (int64_t)blockDim.x * gridDim.x;
    for (; idx < total; idx += stride) {
        out[idx] = scale * bias[idx % n];
    }
}

torch::Tensor fused_linear_scale_residual_forward(torch::Tensor x,
                                                  torch::Tensor weight,
                                                  torch::Tensor bias,
                                                  double scaling_factor) {
    TORCH_CHECK(x.is_cuda() && weight.is_cuda() && bias.is_cuda(), "all tensors must be CUDA tensors");
    TORCH_CHECK(x.scalar_type() == at::kFloat && weight.scalar_type() == at::kFloat && bias.scalar_type() == at::kFloat,
                "only float32 tensors are supported");
    TORCH_CHECK(x.is_contiguous() && weight.is_contiguous() && bias.is_contiguous(),
                "all tensors must be contiguous");
    TORCH_CHECK(x.dim() == 2 && weight.dim() == 2 && bias.dim() == 1, "bad tensor dimensions");

    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, "incompatible shapes");
    TORCH_CHECK(M64 <= INT_MAX && N64 <= INT_MAX && K64 <= INT_MAX, "matrix too large for cuBLAS int API");

    c10::cuda::CUDAGuard device_guard(x.device());
    auto out = torch::empty({M64, N64}, x.options());

    const float gamma = static_cast<float>(1.0 + scaling_factor);
    const int64_t total = M64 * N64;
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    const int threads = 256;
    int blocks = (int)((total + threads - 1) / threads);
    if (blocks > 65535) blocks = 65535;
    fill_scaled_bias_kernel<<<blocks, threads, 0, stream>>>(
        out.data_ptr<float>(), bias.data_ptr<float>(), total, N64, gamma);
    C10_CUDA_KERNEL_LAUNCH_CHECK();

    cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle();

    const int M = static_cast<int>(M64);
    const int N = static_cast<int>(N64);
    const int K = static_cast<int>(K64);
    const float alpha = gamma;
    const float beta = 1.0f;

    // Row-major C[M,N] is viewed as column-major C^T[N,M].
    // Compute C^T = gamma * W[N,K] * X^T[K,M] + scaled_bias, using W and X storage directly.
    cublasComputeType_t compute_type = at::globalContext().allowTF32CuBLAS()
        ? CUBLAS_COMPUTE_32F_FAST_TF32
        : CUBLAS_COMPUTE_32F_PEDANTIC;

    CHECK_CUBLAS(cublasGemmEx(handle,
                              CUBLAS_OP_T, CUBLAS_OP_N,
                              N, M, K,
                              &alpha,
                              weight.data_ptr<float>(), CUDA_R_32F, K,
                              x.data_ptr<float>(), CUDA_R_32F, K,
                              &beta,
                              out.data_ptr<float>(), CUDA_R_32F, N,
                              compute_type,
                              CUBLAS_GEMM_DEFAULT));
    return out;
}
"""

_ext = load_inline(
    name="fused_linear_scale_residual_ext_v1",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["fused_linear_scale_residual_forward"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
    extra_ldflags=["-lcublas"],
    with_cuda=True,
    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 = scaling_factor

    def forward(self, x):
        return _ext.fused_linear_scale_residual_forward(
            x.contiguous(),
            self.matmul.weight.contiguous(),
            self.matmul.bias.contiguous(),
            float(self.scaling_factor),
        )