GPUCodeForces/S1/20/poissonnllloss_cuda.py

import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline

BATCH_SIZE = 4096
FEATURE_DIM = 512

# --- 损失函数的参数 ---
LOG_INPUT = True
FULL = False
EPS = 1e-8


# -------------------------------------------------------------

class ModelNew(nn.Module):

    def __init__(self):
        super().__init__()
        # 将 Python 端的常量存储为实例属性
        self.log_input = LOG_INPUT
        self.full = FULL
        self.eps = EPS
        self._compile_cuda_kernel()

    def _compile_cuda_kernel(self):
        cpp_source = """
        #include <torch/extension.h>

        // C++ 接口
        torch::Tensor poisson_nll_forward_cuda(
            torch::Tensor input,
            torch::Tensor target,
            bool log_input,
            bool full,
            float eps_val
        );
        """
        cuda_source = """
        #include <torch/extension.h> 
        #include <cuda_runtime.h>
        #include <cmath> // for logf, expf
        #include <float.h>

        // 块大小
        #define BLOCK_SIZE 256
        // 定义 PI
        const float PI = 3.141592653589793f;

        /*
         * PoissonNLLLoss 融合核函数
         * 这是一个标准的并行归约核函数。
         * 每个线程处理多个元素 (Grid-Stride Loop)，计算它们的 loss 并累加。
         * 然后执行一个块内归约 (Block-level reduction)。
         * C++ host 端对所有块的和再次求和，然后除以 N 得到 'mean'。
         */
        __global__ void poisson_nll_fused_kernel(
            const float* __restrict__ input_data,
            const float* __restrict__ target_data,
            float* __restrict__ block_loss_sums_out, // (grid_size,)
            int n_elements,
            bool log_input,
            bool full,
            float eps_val
        ) {
            __shared__ float s_data[BLOCK_SIZE];

            float thread_loss_sum = 0.0f;
            int grid_stride = gridDim.x * blockDim.x;

            // Grid-Stride Loop 遍历所有元素
            for (int idx = blockIdx.x * blockDim.x + threadIdx.x;
                 idx < n_elements;
                 idx += grid_stride)
            {
                float x = input_data[idx]; // 'input'
                float y = target_data[idx]; // 'target'
                float loss_val = 0.0f;

                // --- 核心 Loss 计算 ---
                if (log_input) {
                    // loss = exp(input) - target * input
                    loss_val = expf(x) - y * x;
                } else {
                    // loss = input - target * log(input + eps)
                    loss_val = x - y * logf(x + eps_val);
                }

                // --- 'full' 模式的附加项 ---
                // (target * log(target) - target + 0.5 * log(2 * pi * target))
                if (full && y > 1.0f) {
                    float stirling_term = y * logf(y) - y + 0.5f * logf(2.0f * PI * y);
                    loss_val += stirling_term;
                }

                // 累加该线程处理的所有元素的 loss
                thread_loss_sum += loss_val;
            }

            // --- 块内归约 (Sum) ---
            s_data[threadIdx.x] = thread_loss_sum;
            __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) {
                block_loss_sums_out[blockIdx.x] = s_data[0];
            }
        }

        // C++ 封装函数
        torch::Tensor poisson_nll_forward_cuda(
            torch::Tensor input, 
            torch::Tensor target,
            bool log_input,
            bool full,
            float eps_val
        ) {
            // 检查
            TORCH_CHECK(input.is_cuda(), "input must be a CUDA tensor");
            TORCH_CHECK(target.is_cuda(), "target must be a CUDA tensor");

            input = input.contiguous();
            target = target.contiguous();

            const int n_elements = input.numel();
            TORCH_CHECK(target.numel() == n_elements, "input and target must have the same number of elements");

            if (n_elements == 0) {
                return torch::tensor(0.0f, input.options());
            }

            // 分配一个张量来保存每个块的部分和
            const int block_size = BLOCK_SIZE;
            const int grid_size = std::max(1, (n_elements + block_size - 1) / block_size);
            auto block_loss_sums = torch::empty({grid_size}, input.options());

            // 启动 CUDA 核函数
            poisson_nll_fused_kernel<<<grid_size, block_size>>>(
                input.data_ptr<float>(),
                target.data_ptr<float>(),
                block_loss_sums.data_ptr<float>(),
                n_elements,
                log_input,
                full,
                eps_val
            );

            // 核函数返回后，对所有块的和进行求和，然后除以总元素数
            // (reduction='mean')
            return block_loss_sums.sum() / n_elements;
        }
        """

        # JIT (Just-In-Time) 编译
        self.pnl_op = load_inline(
            name="poisson_nll_op_v1",
            cpp_sources=cpp_source,
            cuda_sources=cuda_source,
            functions=["poisson_nll_forward_cuda"],
            extra_cuda_cflags=["-O3"],
            verbose=False
        )

    def forward(self, input_tensor: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
        # 调用我们编译好的 CUDA C++ 函数
        return self.pnl_op.poisson_nll_forward_cuda(
            input_tensor,
            target,
            self.log_input,
            self.full,
            self.eps
        )