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_cuda(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, double scaling);
"""

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

#define CHECK_CUDA(x) TORCH_CHECK(x.is_cuda(), #x " must be a CUDA tensor")
#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
#define CHECK_FLOAT(x) TORCH_CHECK(x.scalar_type() == at::kFloat, #x " must be float32")

__global__ void bias_swish_scale_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 v = y[idx] + bias[col];
        float s = 1.0f / (1.0f + expf(-v));
        y[idx] = (v * s) * scale;
    }
}

torch::Tensor linear_swish_cuda(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, double scaling) {
    CHECK_CUDA(x);
    CHECK_CUDA(weight);
    CHECK_CUDA(bias);
    CHECK_CONTIGUOUS(x);
    CHECK_CONTIGUOUS(weight);
    CHECK_CONTIGUOUS(bias);
    CHECK_FLOAT(x);
    CHECK_FLOAT(weight);
    CHECK_FLOAT(bias);

    TORCH_CHECK(x.dim() == 2, "x must be 2D");
    TORCH_CHECK(weight.dim() == 2, "weight must be 2D");

    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, "incompatible weight shape");
    TORCH_CHECK(bias.numel() == N64, "incompatible bias shape");
    TORCH_CHECK(M64 <= INT_MAX && N64 <= INT_MAX && K64 <= INT_MAX, "dimensions too large");

    const int M = (int)M64;
    const int N = (int)N64;
    const int K = (int)K64;

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

    cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle();
    cublasSetStream(handle, at::cuda::getCurrentCUDAStream());
    // Keep full fp32 SGEMM semantics for fp32-precision correctness.
    cublasSetMathMode(handle, CUBLAS_DEFAULT_MATH);

    const float alpha = 1.0f;
    const float beta = 0.0f;

    // Row-major Y[M,N] is interpreted as column-major C[N,M].
    // Computes C = W * X^T, i.e. Y = X * W^T.
    cublasStatus_t st = 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(st == CUBLAS_STATUS_SUCCESS, "cublasSgemm failed");

    dim3 block(256);
    dim3 grid((N + block.x - 1) / block.x, M);
    bias_swish_scale_kernel<<<grid, block, 0, at::cuda::getCurrentCUDAStream()>>>(
        y.data_ptr<float>(), bias.data_ptr<float>(), M, N, (float)scaling);
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return y;
}
"""

_ext = load_inline(
    name="linear_swish_cublas_fused_ext",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["linear_swish_cuda"],
    with_cuda=True,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3", "--use_fast_math"],
    extra_ldflags=["-lcublas"],
    verbose=False,
)


class ModelNew(nn.Module):
    """
    Linear layer followed by Swish activation and scalar scaling, with the
    GEMM launched through cuBLAS and the bias/activation/scale fused in one
    CUDA kernel.
    """
    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_cuda(x, self.matmul.weight, self.matmul.bias, float(self.scaling_factor))