import torch
from torch.utils.cpp_extension import load_inline

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

__device__ inline float warp_reduce_sum(float v){
    for(int offset=16; offset>0; offset/=2){
        v += __shfl_down_sync(0xffffffff, v, offset);
    }
    return v;
}

__global__ void bce_loss_row_kernel(const float* __restrict__ logits, const float* __restrict__ targets, float* __restrict__ y, int B, int D){
    int b = blockIdx.x;
    int tid = threadIdx.x;
    int lane = tid & 31;
    int warp_id = tid >> 5;
    int warps = blockDim.x >> 5;
    const float* lr = logits + b * D;
    const float* tr = targets + b * D;
    __shared__ float ssum[32];
    float local_sum = 0.0f;
    for(int i = tid; i < D; i += blockDim.x){
        float x = lr[i];
        float t = tr[i];
        float ax = fabsf(x);
        float loss = fmaxf(x, 0.0f) - x * t + log1pf(expf(-ax));
        local_sum += loss;
    }
    float wsum = warp_reduce_sum(local_sum);
    if(lane == 0) ssum[warp_id] = wsum;
    __syncthreads();
    float block_sum = (tid < warps) ? ssum[tid] : 0.0f;
    if(warp_id == 0){
        float s = warp_reduce_sum(block_sum);
        if(lane == 0) ssum[0] = s;
    }
    __syncthreads();
    if(tid == 0){
        y[b] = ssum[0] / (float)D;
    }
}

torch::Tensor bce_loss_row_cuda(torch::Tensor logits, torch::Tensor targets){
    auto lc = logits.contiguous();
    auto tc = targets.contiguous();
    int B = (int)lc.size(0);
    int D = (int)lc.size(1);
    auto y = torch::empty({B}, lc.options());
    int block = 256;
    dim3 grid(B);
    bce_loss_row_kernel<<<grid, block>>>(lc.data_ptr<float>(), tc.data_ptr<float>(), y.data_ptr<float>(), B, D);
    return y;
}
"""

cpp_source = """
torch::Tensor bce_loss_row_cuda(torch::Tensor logits, torch::Tensor targets);
"""

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

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

    def forward(self, logits, targets):
        return self.ops.bce_loss_row_cuda(logits, targets)

