import torch
from torch.utils.cpp_extension import load_inline

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

__device__ __forceinline__ float logcosh(float x){
    float ax = fabsf(x);
    if(ax < 1e-3f){
        float x2 = x * x;
        float x4 = x2 * x2;
        return 0.5f * x2 - (1.0f/12.0f) * x4;
    }
    float t = __expf(-2.0f * ax);
    if(t < 1e-5f){
        return ax + t - 0.6931471805599453f;
    }
    return ax + __logf(1.0f + t) - 0.6931471805599453f;
}

__device__ __forceinline__ float sigmoid_stable(float x){
    if(x >= 0.0f){
        float e = __expf(-x);
        return 1.0f / (1.0f + e);
    } else {
        float e = __expf(x);
        return e / (1.0f + e);
    }
}

__global__ __launch_bounds__(1024) void logcosh_affine_gate_kernel(const float* __restrict__ x, const float* __restrict__ scale, const float* __restrict__ bias, float* __restrict__ y, int B, int D, float alpha, float beta){
    long long tid = blockIdx.x * blockDim.x + threadIdx.x;
    long long stride = blockDim.x * gridDim.x;
    long long total = (long long)B * D;
    for(long long i = tid; i < total; i += stride){
        int d = (int)(i % D);
        float xv = __ldg(x + i);
        float z = fmaf(xv, __ldg(scale + d), __ldg(bias + d));
        float t = fmaf(logcosh(z), alpha, beta);
        float g = 1.0f / (1.0f + __expf(-t));
        y[i] = xv * g;
    }
}

torch::Tensor logcosh_affine_gate_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, torch::Tensor alpha, torch::Tensor beta){
    auto xc = x.contiguous();
    auto sc = scale.contiguous();
    auto bc = bias.contiguous();
    auto y = torch::empty_like(xc);
    int B = (int)xc.size(0);
    int D = (int)xc.size(1);
    float a = alpha.item<float>();
    float be = beta.item<float>();
    int block = 1024;
    long long total = (long long)B * D;
    int grid = (int)std::min<long long>(65535LL, (total + block - 1) / block);
    logcosh_affine_gate_kernel<<<grid, block>>>(xc.data_ptr<float>(), sc.data_ptr<float>(), bc.data_ptr<float>(), y.data_ptr<float>(), B, D, a, be);
    return y;
}

__global__ __launch_bounds__(512) void logcosh_affine_gate_kernel_v2(const float* __restrict__ x, const float* __restrict__ scale, const float* __restrict__ bias, float* __restrict__ y, int B, int D, float alpha, float beta){
    int b = blockIdx.x;
    int lane = blockIdx.y * blockDim.x + threadIdx.x;
    int stride = blockDim.x * gridDim.y;
    int row_start = b * D;
    const float* xr = x + row_start;
    float* yr = y + row_start;
    int aligned = ((((long long)xr & 15LL) == 0) && (((long long)yr & 15LL) == 0) && (((long long)scale & 15LL) == 0) && (((long long)bias & 15LL) == 0) && ((D & 3) == 0));
    if(aligned){
        int D4 = (D / 4) * 4;
        #pragma unroll 12
        for(int i = lane * 4; i < D4; i += stride * 4){
            float4 xv = reinterpret_cast<const float4*>(xr)[i / 4];
            float4 sv = reinterpret_cast<const float4*>(scale)[i / 4];
            float4 bv = reinterpret_cast<const float4*>(bias)[i / 4];
            float4 yv;
            float z0 = fmaf(xv.x, sv.x, bv.x);
            float z1 = fmaf(xv.y, sv.y, bv.y);
            float z2 = fmaf(xv.z, sv.z, bv.z);
            float z3 = fmaf(xv.w, sv.w, bv.w);
            float t0 = fmaf(logcosh(z0), alpha, beta);
            float t1 = fmaf(logcosh(z1), alpha, beta);
            float t2 = fmaf(logcosh(z2), alpha, beta);
            float t3 = fmaf(logcosh(z3), alpha, beta);
            float g0 = __fdividef(1.0f, 1.0f + __expf(-t0));
            float g1 = __fdividef(1.0f, 1.0f + __expf(-t1));
            float g2 = __fdividef(1.0f, 1.0f + __expf(-t2));
            float g3 = __fdividef(1.0f, 1.0f + __expf(-t3));
            yv.x = xv.x * g0;
            yv.y = xv.y * g1;
            yv.z = xv.z * g2;
            yv.w = xv.w * g3;
            reinterpret_cast<float4*>(yr)[i / 4] = yv;
        }
        #pragma unroll 12
        for(int i = D4 + lane; i < D; i += stride){
            float z = fmaf(__ldg(xr + i), __ldg(scale + i), __ldg(bias + i));
            float t = fmaf(logcosh(z), alpha, beta);
            float g = __fdividef(1.0f, 1.0f + __expf(-t));
            yr[i] = xr[i] * g;
        }
    } else {
        #pragma unroll 12
        for(int i = lane; i < D; i += stride){
            float z = fmaf(__ldg(xr + i), __ldg(scale + i), __ldg(bias + i));
            float t = fmaf(logcosh(z), alpha, beta);
            float g = __fdividef(1.0f, 1.0f + __expf(-t));
            yr[i] = xr[i] * g;
        }
    }
}

torch::Tensor logcosh_affine_gate_cuda_v2(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, torch::Tensor alpha, torch::Tensor beta){
    auto xc = x.contiguous();
    auto sc = scale.contiguous();
    auto bc = bias.contiguous();
    auto y = torch::empty_like(xc);
    int B = (int)xc.size(0);
    int D = (int)xc.size(1);
    float a = alpha.item<float>();
    float be = beta.item<float>();
    int block = 512;
    int gy = std::min(32, std::max(1, (D + 1023) / 1024));
    dim3 grid(B, gy);
    logcosh_affine_gate_kernel_v2<<<grid, block>>>(xc.data_ptr<float>(), sc.data_ptr<float>(), bc.data_ptr<float>(), y.data_ptr<float>(), B, D, a, be);
    return y;
}

torch::Tensor logcosh_affine_gate_cuda_opt(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, torch::Tensor alpha, torch::Tensor beta){
    int D = (int)x.size(1);
    if((D & 3) == 0){
        return logcosh_affine_gate_cuda_v2(x, scale, bias, alpha, beta);
    }
    return logcosh_affine_gate_cuda(x, scale, bias, alpha, beta);
}
"""

cpp_source = """
torch::Tensor logcosh_affine_gate_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, torch::Tensor alpha, torch::Tensor beta);
torch::Tensor logcosh_affine_gate_cuda_v2(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, torch::Tensor alpha, torch::Tensor beta);
torch::Tensor logcosh_affine_gate_cuda_opt(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, torch::Tensor alpha, torch::Tensor beta);
"""

ops = load_inline(
    name="logcosh_affine_gate",
    cpp_sources=cpp_source,
    cuda_sources=source,
    functions=["logcosh_affine_gate_cuda","logcosh_affine_gate_cuda_v2","logcosh_affine_gate_cuda_opt"],
    extra_cuda_cflags=["-O3","--use_fast_math","-gencode=arch=compute_80,code=sm_80","-Xptxas","-O3,-dlcm=ca","-maxrregcount=64"],
    verbose=True
)

class ModelNew(torch.nn.Module):
    def __init__(self, scale: torch.Tensor, bias: torch.Tensor, alpha: float, beta: float):
        super(ModelNew, self).__init__()
        self.ops = ops
        self.register_buffer("scale", scale)
        self.register_buffer("bias", bias)
        self.register_buffer("alpha", torch.tensor(float(alpha), dtype=torch.float32))
        self.register_buffer("beta", torch.tensor(float(beta), dtype=torch.float32))

    def forward(self, x):
        return self.ops.logcosh_affine_gate_cuda_opt(x, self.scale, self.bias, self.alpha, self.beta)
