import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Simple model that performs a matrix multiplication, applies Swish activation, and scales the result.
    """
    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):
        x = self.matmul(x)
        x = x * torch.sigmoid(x)  # Swish activation
        x = x * self.scaling_factor
        return x

batch_size = 128
in_features = 32768
out_features = 32768
scaling_factor = 2.0

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 linear_swish_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, double scale_d);
'''

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

static cublasHandle_t get_handle(){
    static cublasHandle_t h=nullptr;
    static std::once_flag f;
    std::call_once(f, [](){
        cublasCreate(&h);
        cublasSetMathMode(h, CUBLAS_DEFAULT_MATH);
    });
    return h;
}

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

__global__ void bias_swish_scale_vec4(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 idx4 = row * N4 + col4;
        float4 yy = reinterpret_cast<float4*>(y)[idx4];
        float4 bb = reinterpret_cast<const float4*>(bias)[col4];
        yy.x = swish_scale(yy.x + bb.x, scale);
        yy.y = swish_scale(yy.y + bb.y, scale);
        yy.z = swish_scale(yy.z + bb.z, scale);
        yy.w = swish_scale(yy.w + bb.w, scale);
        reinterpret_cast<float4*>(y)[idx4] = yy;
    }
}

__global__ void bias_swish_scale_scalar(float* __restrict__ y,
                                        const float* __restrict__ bias,
                                        int total, int N, float scale){
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if(i < total){
        float v = y[i] + bias[i % N];
        y[i] = swish_scale(v, scale);
    }
}

torch::Tensor linear_swish_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, double scale_d){
    int64_t M64 = x.size(0), K64 = x.size(1), N64 = weight.size(0);
    auto y = torch::empty({M64, N64}, x.options());
    int M = (int)M64, N = (int)N64, K = (int)K64;

    float alpha = 1.0f, beta = 0.0f;
    cublasHandle_t h = get_handle();
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    cublasSetStream(h, stream);

    cublasStatus_t st = cublasSgemm(h, 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(st == CUBLAS_STATUS_SUCCESS, "cublasSgemm failed");

    float scale = (float)scale_d;
    if((N & 3) == 0){
        int N4 = N >> 2;
        bias_swish_scale_vec4<<<dim3((N4 + 255) / 256, M), dim3(256), 0, stream>>>(
            y.data_ptr<float>(), bias.data_ptr<float>(), M, N4, scale);
    } else {
        int total = M * N;
        bias_swish_scale_scalar<<<(total + 255) / 256, 256, 0, stream>>>(
            y.data_ptr<float>(), bias.data_ptr<float>(), total, N, scale);
    }
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return y;
}
'''

_ext = load_inline(name='linear_swish_ext_final',
                   cpp_sources=_cpp_src,
                   cuda_sources=_cuda_src,
                   functions=['linear_swish_forward'],
                   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 = scaling_factor

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