GPUCodeForces/S1/31/GaussianNLLLoss_cuda.py

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


N_BATCH = 128
N_FEATURES = 512


FULL = False
EPS = 1e-6
REDUCTION = 'mean'


BLOCK_SIZE = 256




class ModelNew(nn.Module):

    def __init__(self, full=False, eps=1e-6, reduction='mean'):
        super().__init__()
        self.full = full
        self.eps = eps
        self.reduction = reduction
        self.reduction_str = reduction

        self.block_size = BLOCK_SIZE
        self._compile_cuda_kernel()

    def _compile_cuda_kernel(self):

        cpp_header = f"""
        #include <torch/extension.h>

        // C++ 接口
        torch::Tensor gaussian_nll_loss_forward_cuda(
            torch::Tensor input,
            torch::Tensor target,
            torch::Tensor var,
            bool full_flag,
            float eps_val,
            std::string reduction
        );
        """

        cuda_source = f"""
        #include <torch/extension.h> 
        #include <cuda_runtime.h>
        #include <cmath> // for logf, fmaxf

        #define BLOCK_SIZE {self.block_size}

        #define CONST_TERM (0.5f * 1.8378770664f)

        /*
         * GaussianNLLLoss 融合核函数 (Element-wise)
         */
        __global__ void gaussian_nll_loss_fused_kernel(
            const float* __restrict__ input_data,
            const float* __restrict__ target_data,
            const float* __restrict__ var_data,
            float* __restrict__ output_data,
            int N_total,
            bool full_flag,
            float eps_val
        ) {{
            for (int i = blockIdx.x * blockDim.x + threadIdx.x; 
                 i < N_total; 
                 i += gridDim.x * blockDim.x) 
            {{
                const float in = input_data[i];
                const float t = target_data[i];
                const float v = var_data[i];

                const float v_clamped = fmaxf(v, eps_val); 
                const float diff = in - t;
                float loss = 0.5f * (logf(v_clamped) + (diff * diff) / v_clamped);

                if (full_flag) {{
                    loss += CONST_TERM;
                }}

                output_data[i] = loss;
            }}
        }}

        // C++ 封装函数
        torch::Tensor gaussian_nll_loss_forward_cuda(
            torch::Tensor input,
            torch::Tensor target,
            torch::Tensor var,
            bool full_flag,
            float eps_val,
            std::string reduction
        ) {{
            TORCH_CHECK(input.is_cuda(), "input must be a CUDA tensor");
            TORCH_CHECK(input.is_contiguous(), "input must be contiguous");
            TORCH_CHECK(target.is_contiguous(), "target must be contiguous");
            TORCH_CHECK(var.is_contiguous(), "var must be contiguous");

            TORCH_CHECK(input.sizes() == target.sizes(), "input and target shape mismatch");
            TORCH_CHECK(input.sizes() == var.sizes(), "input and var shape mismatch");

            const int64_t N_total = input.numel();

            auto output = torch::empty_like(input);

            dim3 block_dim(BLOCK_SIZE);
            dim3 grid_dim((N_total + BLOCK_SIZE - 1) / BLOCK_SIZE);

            gaussian_nll_loss_fused_kernel<<<grid_dim, block_dim>>>(
                input.data_ptr<float>(),
                target.data_ptr<float>(),
                var.data_ptr<float>(),
                output.data_ptr<float>(),
                N_total,
                full_flag,
                eps_val
            );

            if (reduction == "mean") {{
                return output.mean();
            }} else if (reduction == "sum") {{
                return output.sum();
            }} else {{
                return output; // "none"
            }}
        }}
        """

        nvcc_flags = [
            '-O3',
            '--use_fast_math',
            '--expt-relaxed-constexpr'
        ]

        self.loss_op = load_inline(
            name="gaussian_nll_loss_op_v3_fixed_api",
            cpp_sources=cpp_header,
            cuda_sources=cuda_source,
            functions=["gaussian_nll_loss_forward_cuda"],
            extra_cuda_cflags=nvcc_flags,
            verbose=False
        )

    def forward(self, input: torch.Tensor, target: torch.Tensor, var: torch.Tensor) -> torch.Tensor:


        if target.size() != input.size():
            target = target.expand_as(input)
        if var.size() != input.size():
            var = var.expand_as(input)


        input_cont = input.contiguous()
        target_cont = target.contiguous()
        var_cont = var.contiguous()

        return self.loss_op.gaussian_nll_loss_forward_cuda(
            input_cont,
            target_cont,
            var_cont,
            self.full,
            self.eps,
            self.reduction_str
        )