import torch
from torch.utils.cpp_extension import load_inline

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

__device__ __forceinline__ float seluf(float x){
    const float lambda_ = 1.0507009873554804934193349852946f;
    const float alpha_ = 1.6732632423543772848170429916717f;
    float v = x >= 0.0f ? x : alpha_ * (__expf(x) - 1.0f);
    return lambda_ * v;
}

__global__ __launch_bounds__(256) void selu_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){
    int row = blockIdx.x;
    int tid = threadIdx.x;
    int col_block = blockIdx.y;
    int stride4 = blockDim.x * 4;
    int col_idx = (col_block * stride4) + tid * 4;
    const float* xr = x + row * D;
    float* yr = y + row * D;
    for(int base = col_idx; base < D; base += stride4 * 2){
        if (base < D){
            float4 xv = reinterpret_cast<const float4*>(xr + base)[0];
            float4 sv = reinterpret_cast<const float4*>(scale + base)[0];
            float4 bv = reinterpret_cast<const float4*>(bias + base)[0];
            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 m0 = seluf(z0);
            float m1 = seluf(z1);
            float m2 = seluf(z2);
            float m3 = seluf(z3);
            float t0 = fmaf(alpha, m0, beta);
            float t1 = fmaf(alpha, m1, beta);
            float t2 = fmaf(alpha, m2, beta);
            float t3 = fmaf(alpha, m3, 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 + base)[0] = yv;
        }
        int base2 = base + stride4;
        if (base2 < D){
            float4 xv2 = reinterpret_cast<const float4*>(xr + base2)[0];
            float4 sv2 = reinterpret_cast<const float4*>(scale + base2)[0];
            float4 bv2 = reinterpret_cast<const float4*>(bias + base2)[0];
            float4 yv2;
            float z0b = fmaf(xv2.x, sv2.x, bv2.x);
            float z1b = fmaf(xv2.y, sv2.y, bv2.y);
            float z2b = fmaf(xv2.z, sv2.z, bv2.z);
            float z3b = fmaf(xv2.w, sv2.w, bv2.w);
            float mb0 = seluf(z0b);
            float mb1 = seluf(z1b);
            float mb2 = seluf(z2b);
            float mb3 = seluf(z3b);
            float tb0 = fmaf(alpha, mb0, beta);
            float tb1 = fmaf(alpha, mb1, beta);
            float tb2 = fmaf(alpha, mb2, beta);
            float tb3 = fmaf(alpha, mb3, beta);
            float gb0 = __fdividef(1.0f, 1.0f + __expf(-tb0));
            float gb1 = __fdividef(1.0f, 1.0f + __expf(-tb1));
            float gb2 = __fdividef(1.0f, 1.0f + __expf(-tb2));
            float gb3 = __fdividef(1.0f, 1.0f + __expf(-tb3));
            yv2.x = xv2.x * gb0;
            yv2.y = xv2.y * gb1;
            yv2.z = xv2.z * gb2;
            yv2.w = xv2.w * gb3;
            reinterpret_cast<float4*>(yr + base2)[0] = yv2;
        }
    }
}

torch::Tensor selu_affine_gate_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, double alpha, double 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 = (float)alpha;
    float be = (float)beta;
    int block = 256;
    int elements_per_thread = 4;
    int elements_per_block = block * elements_per_thread;
    int gy = (D + elements_per_block - 1) / elements_per_block;
    dim3 grid(B, gy);
    selu_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;
}
"""

cpp_source = """
#include <torch/extension.h>
torch::Tensor selu_affine_gate_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, double alpha, double beta);
"""

ops = load_inline(
    name="selu_affine_gate",
    cpp_sources=cpp_source,
    cuda_sources=source,
    functions=["selu_affine_gate_cuda"],
    extra_cflags=["-O3","-std=c++17"],
    extra_cuda_cflags=["-O3","--use_fast_math","-std=c++17","-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.alpha = float(alpha)
        self.beta = float(beta)

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