import torch
from torch.utils.cpp_extension import load_inline

source = """
#include <torch/extension.h>
#include <cuda_runtime.h>

__global__ void glu_kernel(const float* x, float* y, int D, long long rows) {
    long long idx = blockIdx.x * blockDim.x + threadIdx.x;
    long long stride = blockDim.x * gridDim.x;
    long long total = rows * (long long)D;
    for (long long t = idx; t < total; t += stride) {
        long long row = t / D;
        int col = (int)(t % D);
        long long base = row * (long long)(2 * D);
        float a = x[base + col];
        float b = x[base + D + col];
        float s = 1.0f / (1.0f + expf(-b));
        y[t] = a * s;
    }
}

torch::Tensor glu_cuda(torch::Tensor x) {
    auto x_contig = x.contiguous();
    long long rows = 1;
    for (int i = 0; i < x_contig.dim() - 1; ++i) rows *= x_contig.size(i);
    int D = (int)(x_contig.size(-1) / 2);
    auto y = torch::empty({rows, D}, x_contig.options());
    int block = 512;
    long long total = rows * (long long)D;
    long long grid = (total + block - 1) / block;
    grid = grid > 65535 ? 65535 : grid;
    glu_kernel<<<(int)grid, block>>>(x_contig.data_ptr<float>(), y.data_ptr<float>(), D, rows);
    return y;
}
"""

cpp_source = """
torch::Tensor glu_cuda(torch::Tensor x);
"""

ops = load_inline(
    name="glu",
    cpp_sources=cpp_source,
    cuda_sources=source,
    functions=["glu_cuda"],
    verbose=True
)

class ModelNew(torch.nn.Module):
    def __init__(self):
        super(ModelNew, self).__init__()
        self.ops = ops

    def forward(self, x):
        return self.ops.glu_cuda(x)
