GPUCodeForces/S1/wut0n_#8/bce_cudacode.py

# 优化后的 bce_cudacode.py
import torch
from torch.utils.cpp_extension import load_inline

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

// 高效归约版本 - 使用两级归约
__global__ void bce_kernel_efficient(
    const float* __restrict__ input,
    const float* __restrict__ target,
    float* __restrict__ partial_sums,
    int size,
    float epsilon
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    
    // 局部累加器
    float local_bce = 0.0f;
    
    // 每个线程处理多个元素
    int stride = blockDim.x * gridDim.x;
    for (int i = idx; i < size; i += stride) {
        float xi = input[i];
        float yi = target[i];
        
        // 数值稳定性处理
        xi = max(min(xi, 1.0f - epsilon), epsilon);
        
        // BCE = -[y * log(x) + (1-y) * log(1-x)]
        float term1 = yi * logf(xi);
        float term2 = (1.0f - yi) * logf(1.0f - xi);
        local_bce -= (term1 + term2);
    }
    
    // 使用共享内存进行块内归约
    extern __shared__ float shared_mem[];
    shared_mem[threadIdx.x] = local_bce;
    
    __syncthreads();
    
    // 块内归约
    for (int stride = blockDim.x / 2; stride > 0; stride /= 2) {
        if (threadIdx.x < stride) {
            shared_mem[threadIdx.x] += shared_mem[threadIdx.x + stride];
        }
        __syncthreads();
    }
    
    // 每个block写入部分和到全局内存
    if (threadIdx.x == 0) {
        partial_sums[blockIdx.x] = shared_mem[0];
    }
}

// 快速版本 - 使用快速log近似
__global__ void bce_kernel_fast(
    const float* __restrict__ input,
    const float* __restrict__ target,
    float* __restrict__ bce_loss,
    int size,
    float epsilon
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    
    // 局部累加器
    float local_bce = 0.0f;
    
    // 每个线程处理多个元素
    int stride = blockDim.x * gridDim.x;
    for (int i = idx; i < size; i += stride) {
        float xi = input[i];
        float yi = target[i];
        
        // 数值稳定性处理
        xi = max(min(xi, 1.0f - epsilon), epsilon);
        
        // 使用快速log近似
        float log_x = __logf(xi);  // 使用CUDA内置快速log
        float log_1_x = __logf(1.0f - xi);
        
        local_bce -= (yi * log_x + (1.0f - yi) * log_1_x);
    }
    
    // 使用共享内存进行块内归约
    extern __shared__ float shared_mem[];
    shared_mem[threadIdx.x] = local_bce;
    
    __syncthreads();
    
    // 块内归约
    for (int stride = blockDim.x / 2; stride > 0; stride /= 2) {
        if (threadIdx.x < stride) {
            shared_mem[threadIdx.x] += shared_mem[threadIdx.x + stride];
        }
        __syncthreads();
    }
    
    // 只有第一个线程做原子操作
    if (threadIdx.x == 0) {
        atomicAdd(bce_loss, shared_mem[0]);
    }
}

torch::Tensor bce_cuda(torch::Tensor input, torch::Tensor target, std::string mode = "fast") {
    TORCH_CHECK(input.scalar_type() == torch::kFloat32, "Input must be float32");
    TORCH_CHECK(target.scalar_type() == torch::kFloat32, "Target must be float32");
    TORCH_CHECK(input.sizes() == target.sizes(), "Input and target must have same shape");
    
    auto input_contig = input.contiguous();
    auto target_contig = target.contiguous();
    int size = input_contig.numel();
    
    if (mode == "efficient") {
        // 高效归约版本 - 两级归约
        const int block_size = 256;
        int num_blocks = min(1024, (size + block_size - 1) / block_size);
        
        // 创建部分和数组
        auto partial_sums = torch::zeros({num_blocks}, input.options());
        
        size_t shared_mem = block_size * sizeof(float);
        bce_kernel_efficient<<<num_blocks, block_size, shared_mem>>>(
            input_contig.data_ptr<float>(),
            target_contig.data_ptr<float>(),
            partial_sums.data_ptr<float>(),
            size,
            1e-8f
        );
        
        // 在CPU上完成最终归约
        auto partial_sums_cpu = partial_sums.to(torch::kCPU);
        float total_loss = 0.0f;
        for (int i = 0; i < num_blocks; i++) {
            total_loss += partial_sums_cpu.data_ptr<float>()[i];
        }
        
        return torch::tensor({total_loss}, input.options());
        
    } else {
        // 快速版本 - 默认
        auto bce_loss = torch::zeros(1, input.options());
        
        const int block_size = 256;
        int num_blocks = min(1024, (size + block_size - 1) / block_size);
        size_t shared_mem = block_size * sizeof(float);
        
        bce_kernel_fast<<<num_blocks, block_size, shared_mem>>>(
            input_contig.data_ptr<float>(),
            target_contig.data_ptr<float>(),
            bce_loss.data_ptr<float>(),
            size,
            1e-8f
        );
        
        return bce_loss;
    }
}
"""

bce_cpp_source = """
torch::Tensor bce_cuda(torch::Tensor input, torch::Tensor target, std::string mode);
"""

# 编译CUDA代码
bce = load_inline(
    name="bce",
    cpp_sources=bce_cpp_source,
    cuda_sources=bce_source,
    functions=["bce_cuda"],
    extra_cuda_cflags=[
        "-O3", 
        "--use_fast_math",
        "-gencode=arch=compute_80,code=sm_80"
    ],
    verbose=True
)

class ModelNew(torch.nn.Module):
    def __init__(self, mode="fast"):
        super(ModelNew, self).__init__()
        self.mode = mode
        self.bce = bce

    def forward(self, input, target):
        return self.bce.bce_cuda(input, target, self.mode)