import torch
from torch.utils.cpp_extension import load_inline

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

__global__ void hardswish_affine_gate_kernel(const float* x, const float* scale, const float* bias, float* y, int B, int D, float alpha, float beta){
    int b = blockIdx.x;
    int tid = threadIdx.x;
    int stride = blockDim.x;
    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 4
        for(int i = tid * 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 = fminf(fmaxf(alpha * z0 + beta + 3.0f, 0.0f), 6.0f) * (1.0f / 6.0f);
            float t1 = fminf(fmaxf(alpha * z1 + beta + 3.0f, 0.0f), 6.0f) * (1.0f / 6.0f);
            float t2 = fminf(fmaxf(alpha * z2 + beta + 3.0f, 0.0f), 6.0f) * (1.0f / 6.0f);
            float t3 = fminf(fmaxf(alpha * z3 + beta + 3.0f, 0.0f), 6.0f) * (1.0f / 6.0f);
            yv.x = xv.x * t0;
            yv.y = xv.y * t1;
            yv.z = xv.z * t2;
            yv.w = xv.w * t3;
            reinterpret_cast<float4*>(yr)[i / 4] = yv;
        }
        #pragma unroll 4
        for(int i = D4 + tid; i < D; i += stride){
            float z = fmaf(xr[i], scale[i], bias[i]);
            float t = fminf(fmaxf(alpha * z + beta + 3.0f, 0.0f), 6.0f) * (1.0f / 6.0f);
            yr[i] = xr[i] * t;
        }
    } else {
        #pragma unroll 4
        for(int i = tid; i < D; i += stride){
            float z = fmaf(xr[i], scale[i], bias[i]);
            float t = fminf(fmaxf(alpha * z + beta + 3.0f, 0.0f), 6.0f) * (1.0f / 6.0f);
            yr[i] = xr[i] * t;
        }
    }
}

torch::Tensor hardswish_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;
    int grid = B;
    hardswish_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 = """
torch::Tensor hardswish_affine_gate_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, torch::Tensor alpha, torch::Tensor beta);
"""

ops = load_inline(
    name="hardswish_affine_gate",
    cpp_sources=cpp_source,
    cuda_sources=source,
    functions=["hardswish_affine_gate_cuda"],
    extra_cuda_cflags=["-O3","--use_fast_math"],
    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.hardswish_affine_gate_cuda(x, self.scale, self.bias, self.alpha, self.beta)

