import torch
from torch.utils.cpp_extension import load_inline

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

__global__ void logsigmoid_affine_gate_kernel(const float* x, const float* scale, const float* bias, float* y, int B, int D){
    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 b = (int)(i / D);
        int d = (int)(i % D);
        float z = x[i] * scale[d] + bias[d];
        float g = -log1pf(expf(-z));
        y[i] = x[i] * g;
    }
}

torch::Tensor logsigmoid_affine_gate_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias){
    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);
    int block = 1024;
    int grid = (int)((long long)B * D / block);
    if(grid < 1) grid = 1; if(grid > 65535) grid = 65535;
    logsigmoid_affine_gate_kernel<<<grid, block>>>(xc.data_ptr<float>(), sc.data_ptr<float>(), bc.data_ptr<float>(), y.data_ptr<float>(), B, D);
    return y;
}
"""

cpp_source = """
torch::Tensor logsigmoid_affine_gate_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias);
"""

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

class ModelNew(torch.nn.Module):
    def __init__(self, scale: torch.Tensor, bias: torch.Tensor):
        super(ModelNew, self).__init__()
        self.ops = ops
        self.register_buffer("scale", scale)
        self.register_buffer("bias", bias)

    def forward(self, x: torch.Tensor):
        return self.ops.logsigmoid_affine_gate_cuda(x, self.scale, self.bias)
