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 os
_tmpdir = os.path.abspath('tmp_compile')
os.makedirs(_tmpdir, exist_ok=True)
os.environ['TMPDIR'] = _tmpdir
os.environ['TEMP'] = _tmpdir
os.environ['TMP'] = _tmpdir

import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline

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

void sigmoid_sum_launcher(const float* y, float* out, int M, int N, cudaStream_t stream);

torch::Tensor linear_sigmoid_sum(torch::Tensor x, torch::Tensor weight, torch::Tensor bias) {
    const int M = (int)x.size(0);
    const int K = (int)x.size(1);
    const int N = (int)weight.size(0);

    auto y = 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);

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

    // Row-major y[M,N] = x[M,K] * weight[N,K]^T, expressed as column-major GEMM:
    // C_col(N,M) = weight_col(K,N)^T * x_col(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,
                 y.data_ptr<float>(), CUDA_R_32F, N,
                 CUBLAS_COMPUTE_32F_FAST_TF32,
                 CUBLAS_GEMM_DEFAULT_TENSOR_OP);

    sigmoid_sum_launcher(y.data_ptr<float>(), out.data_ptr<float>(), M, N, stream);
    return out;
}
'''

cuda_sources = r'''
#include <cuda_runtime.h>
#include <math.h>

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

#pragma unroll 4
    for (int i = tid; i < N; i += blockDim.x) {
        float v = p[i];
        sum += 1.0f / (1.0f + __expf(-v));
    }

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

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

    if (tid == 0) out[row] = smem[0];
}

void sigmoid_sum_launcher(const float* y, float* out, int M, int N, cudaStream_t stream) {
    sigmoid_sum_kernel<<<M, 256, 0, stream>>>(y, out, N);
}
'''

_ext = load_inline(
    name='linear_sigmoid_sum_final_ext',
    cpp_sources=cpp_sources,
    cuda_sources=cuda_sources,
    functions=['linear_sigmoid_sum'],
    verbose=False,
    extra_cflags=['-O3'],
    extra_cuda_cflags=['-O3', '--use_fast_math'],
    extra_ldflags=['-lcublas'],
)

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.linear_sigmoid_sum(x, self.linear.weight, self.linear.bias)