GPUCodeForces/S1/uucoco_#57/AngularLoss_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=1e-6):
        super().__init__()
        self.alpha = alpha
        self.smooth = smooth
        self._compile_cuda_kernel()

    def _compile_cuda_kernel(self):
        cpp_source = """
        torch::Tensor angular_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 angular_loss_kernel(
            const double* __restrict__ logits,
            const double* __restrict__ targets,
            double* __restrict__ intersection_out,
            double* __restrict__ union_out,
            double* __restrict__ bce_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_union[256];
            __shared__ double s_bce[256];

            double local_inter = 0.0;
            double local_union = 0.0;
            double local_bce = 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_union += p + y;

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

                local_bce += bce;
            }

            s_inter[tid] = local_inter;
            s_union[tid] = local_union;
            s_bce[tid] = local_bce;
            __syncthreads();

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

            if (tid == 0) {
                intersection_out[batch_idx] = s_inter[0];
                union_out[batch_idx] = s_union[0];
                bce_out[batch_idx] = s_bce[0];
            }
        }

        torch::Tensor angular_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 union_out = torch::zeros({batch_size}, Z_c.options());
            auto bce_out = torch::zeros({batch_size}, Z_c.options());

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

            angular_loss_kernel<<<blocks, threads>>>(
                Z_c.data_ptr<double>(),
                Y_c.data_ptr<double>(),
                intersection_out.data_ptr<double>(),
                union_out.data_ptr<double>(),
                bce_out.data_ptr<double>(),
                batch_size,
                feature_dim
            );

            return torch::cat({intersection_out, union_out, bce_out}, 0);
        }
        """

        self.op = load_inline(
            name="angular_loss_v1",
            cpp_sources=cpp_source,
            cuda_sources=cuda_source,
            functions=["angular_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.angular_loss_cuda(logits, targets_f, self.smooth, batch_size)

        intersection = comp_flat[:batch_size]
        union = comp_flat[batch_size:2 * batch_size]
        bce_sum = comp_flat[2 * batch_size:]

        angle = 1.0 - (2.0 * intersection + self.smooth) / (union + self.smooth)
        angular_loss = (angle * (1.0 - angle)).mean()
        bce_loss = bce_sum.sum() / (batch_size * feature_dim)
        return self.alpha * angular_loss + (1.0 - self.alpha) * bce_loss