import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Simple model that performs a matrix multiplication, applies sigmoid, and sums the result.
    """
    def __init__(self, input_size, hidden_size):
        super(Model, self).__init__()
        self.linear = nn.Linear(input_size, hidden_size)

    def forward(self, x):
        """
        Args:
            x: Input tensor of shape (batch_size, input_size).

        Returns:
            Output tensor of shape (batch_size, 1).
        """
        x = self.linear(x)
        x = torch.sigmoid(x)
        x = torch.sum(x, dim=1, keepdim=True)
        return x

batch_size = 128
input_size = 32768
hidden_size = 32768

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

def get_init_inputs():
    return [input_size, hidden_size]

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

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

__global__ void sigmoid_sum_kernel(const float* __restrict__ tmp,
                                   const float* __restrict__ bias,
                                   float* __restrict__ out,
                                   int M, int N) {
    int row = blockIdx.x;
    int tid = threadIdx.x;
    float sum = 0.0f;
    const float* rowptr = tmp + (long long)row * N;

    for (int i = tid; i < N; i += blockDim.x) {
        float z = rowptr[i] + bias[i];
        sum += 1.0f / (1.0f + __expf(-z));
    }

    __shared__ float smem[256];
    smem[tid] = sum;
    __syncthreads();

    for (int s = 128; s >= 32; s >>= 1) {
        if (tid < s) smem[tid] += smem[tid + s];
        __syncthreads();
    }

    if (tid < 32) {
        float v = smem[tid];
        #pragma unroll
        for (int offset = 16; offset > 0; offset >>= 1) {
            v += __shfl_down_sync(0xffffffff, v, offset);
        }
        if (tid == 0) out[row] = v;
    }
}

torch::Tensor fused_linear_sigmoid_sum(torch::Tensor x, torch::Tensor weight, torch::Tensor bias) {
    c10::cuda::CUDAGuard device_guard(x.device());

    const int M = (int)x.size(0);
    const int K = (int)x.size(1);
    const int N = (int)weight.size(0);

    auto tmp = torch::empty({M, N}, x.options());
    auto out = torch::empty({M, 1}, x.options());

    cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle();
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    cublasSetStream(handle, stream);
    cublasSetPointerMode(handle, CUBLAS_POINTER_MODE_HOST);

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

    // Row-major tmp[M,N] = x[M,K] * weight[N,K]^T.
    // cuBLAS is column-major, so this is C_col[N,M] = weight^T[N,K] * x[K,M].
    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,
                 tmp.data_ptr<float>(), CUDA_R_32F, N,
                 CUBLAS_COMPUTE_32F_FAST_TF32,
                 CUBLAS_GEMM_DEFAULT_TENSOR_OP);

    sigmoid_sum_kernel<<<M, 256, 0, stream>>>(
        tmp.data_ptr<float>(), bias.data_ptr<float>(), out.data_ptr<float>(), M, N);

    return out;
}
"""

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

class ModelNew(nn.Module):
    def __init__(self, input_size, hidden_size):
        super(ModelNew, self).__init__()
        self.linear = nn.Linear(input_size, hidden_size)

    def forward(self, x):
        return _ext.fused_linear_sigmoid_sum(x, self.linear.weight, self.linear.bias)