# bcewithlogitsloss_cuda.py
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline

from bcewithlogitsloss_torch import BATCH_SIZE, FEATURE_DIM # 导入维度常量

N_ELEMENTS = BATCH_SIZE * FEATURE_DIM

class ModelNew(nn.Module):
    def __init__(self):
        super().__init__()
        self._compile_cuda_kernel()

    def _compile_cuda_kernel(self):
        cpp_source = """
        #include <torch/extension.h>
        torch::Tensor bce_forward_cuda(torch::Tensor logits, torch::Tensor target);
        """

        cuda_source = """
        #include <cuda_runtime.h>
        #include <cmath>
        #include <float.h> // For FLT_MAX

        #define BLOCK_SIZE 256
        #define WARP_SIZE 32

        __device__ __forceinline__ float warp_reduce_sum(float val) {
            for (int offset = WARP_SIZE / 2; offset > 0; offset /= 2)
                val += __shfl_down_sync(0xffffffff, val, offset);
            return val;
        }

        __global__ void bce_fused_kernel(
            const float* __restrict__ logits,
            const float* __restrict__ target,
            float* __restrict__ loss_sum_out, 
            int n_elements
        ) {
            __shared__ float s_data[BLOCK_SIZE]; 
            
    
            float thread_loss = 0.0f;
            int grid_stride = gridDim.x * blockDim.x;

            for (int idx = blockIdx.x * blockDim.x + threadIdx.x;
                 idx < n_elements;
                 idx += grid_stride)
            {
                float x = logits[idx];
                float y = target[idx];
                
                float max_val = fmaxf(0.0f, x);


                if (x >= 0) {
                    thread_loss += (x - x * y) + logf(1.0f + expf(-x));
                } else {
                    thread_loss += (-x * y) + logf(1.0f + expf(x));
                }
            }
            

            s_data[threadIdx.x] = thread_loss;
            __syncthreads();


            for (int offset = BLOCK_SIZE / 2; offset > 0; offset >>= 1) {
                if (threadIdx.x < offset) {
                    s_data[threadIdx.x] += s_data[threadIdx.x + offset];
                }
                __syncthreads();
            }


            if (threadIdx.x == 0) {
                loss_sum_out[blockIdx.x] = s_data[0]; 
            }
        }

        torch::Tensor bce_forward_cuda(torch::Tensor logits, torch::Tensor target) {
            TORCH_CHECK(logits.is_cuda(), "Input must be a CUDA tensor");
            logits = logits.contiguous();
            target = target.contiguous();

            int n_elements = logits.numel();

            const int block_size = BLOCK_SIZE;
            const int grid_size = (n_elements + block_size - 1) / block_size; 

            auto block_loss_sums = torch::empty({grid_size}, logits.options());

            bce_fused_kernel<<<grid_size, block_size>>>(
                logits.data_ptr<float>(),
                target.data_ptr<float>(),
                block_loss_sums.data_ptr<float>(),
                n_elements
            );
            
            return block_loss_sums.sum() / n_elements;
        }
        """
        
        self.bce_op = load_inline(
            name="bce_fused_op",
            cpp_sources=cpp_source,
            cuda_sources=cuda_source,
            functions=["bce_forward_cuda"],
            extra_cuda_cflags=["-O3", "--use_fast_math"],
            verbose=True
        )

    def forward(self, logits: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
        return self.bce_op.bce_forward_cuda(logits, target)