GPUCodeForces/S1/uucoco_#59/CrossEntropyDiceLoss_cuda.py

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


class ModelNew(nn.Module):
    def __init__(self, alpha=0.5, smooth=1.0):
        super().__init__()
        self.alpha = alpha
        self.smooth = smooth
        self._compile_cuda_kernel()

    def _compile_cuda_kernel(self):
        cpp_source = """
        torch::Tensor ce_dice_loss_cuda(torch::Tensor logits, torch::Tensor targets, double smooth, int batch_size);
        """

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

        __device__ __forceinline__ double sigmoid_d(double x) {
            if (x >= 0.0) {
                double z = exp(-x);
                return 1.0 / (1.0 + z);
            } else {
                double z = exp(x);
                return z / (1.0 + z);
            }
        }

        __device__ __forceinline__ double log_sigmoid_d(double x) {
            if (x >= 0.0) {
                return -log(1.0 + exp(-x));
            } else {
                return x - log(1.0 + exp(x));
            }
        }

        __global__ void ce_dice_loss_kernel(
            const double* __restrict__ logits,
            const double* __restrict__ targets,
            double* __restrict__ intersection_out,
            double* __restrict__ sum_probs_out,
            double* __restrict__ sum_targets_out,
            double* __restrict__ ce_out,
            const int batch_size,
            const int feature_dim) 
        {
            const int batch_idx = blockIdx.x;
            const int tid = threadIdx.x;
            const int stride = blockDim.x;

            if (batch_idx >= batch_size) return;

            __shared__ double s_inter[256];
            __shared__ double s_probs[256];
            __shared__ double s_targets[256];
            __shared__ double s_ce[256];

            double local_inter = 0.0;
            double local_probs = 0.0;
            double local_targets = 0.0;
            double local_ce = 0.0;

            const int offset = batch_idx * feature_dim;

            for (int i = tid; i < feature_dim; i += stride) {
                double z = logits[offset + i];
                double y = targets[offset + i];

                double p = sigmoid_d(z);

                local_inter += p * y;
                local_probs += p;
                local_targets += y;

                double log_p = log_sigmoid_d(z);
                double log_1mp = log_sigmoid_d(-z);
                double ce = -(y * log_p + (1.0 - y) * log_1mp);

                local_ce += ce;
            }

            s_inter[tid] = local_inter;
            s_probs[tid] = local_probs;
            s_targets[tid] = local_targets;
            s_ce[tid] = local_ce;
            __syncthreads();

            for (int s = blockDim.x / 2; s > 0; s >>= 1) {
                if (tid < s) {
                    s_inter[tid] += s_inter[tid + s];
                    s_probs[tid] += s_probs[tid + s];
                    s_targets[tid] += s_targets[tid + s];
                    s_ce[tid] += s_ce[tid + s];
                }
                __syncthreads();
            }

            if (tid == 0) {
                intersection_out[batch_idx] = s_inter[0];
                sum_probs_out[batch_idx] = s_probs[0];
                sum_targets_out[batch_idx] = s_targets[0];
                ce_out[batch_idx] = s_ce[0];
            }
        }

        torch::Tensor ce_dice_loss_cuda(torch::Tensor logits, torch::Tensor targets, double smooth, int batch_size) {
            auto Z_c = logits.contiguous();
            auto Y_c = targets.contiguous();

            const int feature_dim = Z_c.size(1);

            auto intersection_out = torch::zeros({batch_size}, Z_c.options());
            auto sum_probs_out = torch::zeros({batch_size}, Z_c.options());
            auto sum_targets_out = torch::zeros({batch_size}, Z_c.options());
            auto ce_out = torch::zeros({batch_size}, Z_c.options());

            const int threads = 256;
            const int blocks = batch_size;

            ce_dice_loss_kernel<<<blocks, threads>>>(
                Z_c.data_ptr<double>(),
                Y_c.data_ptr<double>(),
                intersection_out.data_ptr<double>(),
                sum_probs_out.data_ptr<double>(),
                sum_targets_out.data_ptr<double>(),
                ce_out.data_ptr<double>(),
                batch_size,
                feature_dim
            );

            return torch::cat({intersection_out, sum_probs_out, sum_targets_out, ce_out}, 0);
        }
        """

        self.op = load_inline(
            name="ce_dice_loss_v1",
            cpp_sources=cpp_source,
            cuda_sources=cuda_source,
            functions=["ce_dice_loss_cuda"],
            extra_cuda_cflags=["-O3"],
            verbose=False
        )

    def forward(self, logits, targets):
        targets_f = targets.to(logits.dtype)

        batch_size = logits.size(0)
        feature_dim = logits.size(1)

        comp_flat = self.op.ce_dice_loss_cuda(logits, targets_f, self.smooth, batch_size)

        intersection = comp_flat[:batch_size]
        sum_probs = comp_flat[batch_size:2 * batch_size]
        sum_targets = comp_flat[2 * batch_size:3 * batch_size]
        ce_sum = comp_flat[3 * batch_size:]

        dice = (2.0 * intersection + self.smooth) / (sum_probs + sum_targets + self.smooth)
        dice_loss = 1.0 - dice.mean()

        ce_loss = ce_sum.sum() / (batch_size * feature_dim)

        return self.alpha * dice_loss + (1.0 - self.alpha) * ce_loss