import torch
from torch.utils.cpp_extension import load_inline

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

__global__ void gtu_kernel(const float* x, float* y, int D, long long rows) {
    int b = blockIdx.x;
    int tid = threadIdx.x;
    int stride = blockDim.x;
    long long row_start = b * (long long)(2 * D);
    const float* a_row = x + row_start;
    const float* b_row = x + row_start + D;
    float* y_row = y + b * D;
    int aligned = ((((long long)a_row & 15LL) == 0) && (((long long)b_row & 15LL) == 0) && (((long long)y_row & 15LL) == 0) && ((D & 3) == 0));
    if (aligned) {
        int D4 = (D / 4) * 4;
        #pragma unroll 4
        for (int i = tid * 4; i < D4; i += stride * 4) {
            float4 av = reinterpret_cast<const float4*>(a_row)[i / 4];
            float4 bv = reinterpret_cast<const float4*>(b_row)[i / 4];
            float4 yv;
            yv.x = tanhf(av.x) * (1.0f / (1.0f + expf(-bv.x)));
            yv.y = tanhf(av.y) * (1.0f / (1.0f + expf(-bv.y)));
            yv.z = tanhf(av.z) * (1.0f / (1.0f + expf(-bv.z)));
            yv.w = tanhf(av.w) * (1.0f / (1.0f + expf(-bv.w)));
            reinterpret_cast<float4*>(y_row)[i / 4] = yv;
        }
        #pragma unroll 4
        for (int i = D4 + tid; i < D; i += stride) {
            float a = a_row[i];
            float b = b_row[i];
            float ta = tanhf(a);
            float sb = 1.0f / (1.0f + expf(-b));
            y_row[i] = ta * sb;
        }
    } else {
        #pragma unroll 4
        for (int i = tid; i < D; i += stride) {
            float a = a_row[i];
            float b = b_row[i];
            float ta = tanhf(a);
            float sb = 1.0f / (1.0f + expf(-b));
            y_row[i] = ta * sb;
        }
    }
}

torch::Tensor gtu_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 = 1024;
    int grid = (int)rows;
    gtu_kernel<<<grid, block>>>(x_contig.data_ptr<float>(), y.data_ptr<float>(), D, rows);
    return y;
}
"""

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

ops = load_inline(
    name="gtu",
    cpp_sources=cpp_source,
    cuda_sources=source,
    functions=["gtu_cuda"],
    extra_cuda_cflags=["-O3","--use_fast_math"],
    verbose=True
)

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

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