GPUCodeForces/S1/44/multilabelmarginloss_cuda.py

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

cpp_source = """
#include <torch/extension.h>
#include <string>

// 函数前向声明
torch::Tensor multi_label_margin_loss_cuda_forward(
    const torch::Tensor& input, 
    const torch::Tensor& target,
    const std::string& reduction
);
"""

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

#define BLOCK_SIZE 256
// 设置一个合理的单个样本最大正类标签数，用于共享内存数组
#define MAX_POSITIVE_LABELS 32

template <typename T>
__global__ void multi_label_margin_loss_kernel(
    T* output,          // (N)
    const T* input,     // (N, C)
    const long* target, // (N, C)
    const int N,
    const int C)
{
    int sample_idx = blockIdx.x;
    if (sample_idx >= N) return;

    __shared__ long positive_indices[MAX_POSITIVE_LABELS];
    __shared__ int num_positives;

    // 线程 0 初始化共享内存计数器
    if (threadIdx.x == 0) {
        num_positives = 0;
    }
    __syncthreads();

    for (int i = threadIdx.x; i < C; i += blockDim.x) {
        long label = target[sample_idx * C + i];
        if (label != -1) {
            int index = atomicAdd(&num_positives, 1);
            if (index < MAX_POSITIVE_LABELS) {
                positive_indices[index] = label;
            }
        }
    }
    __syncthreads();

    __shared__ T sdata[BLOCK_SIZE];
    int tid = threadIdx.x;
    T my_sum = 0.0f;
    const T* input_row = input + sample_idx * C;

    for (int neg_class_idx = tid; neg_class_idx < C; neg_class_idx += blockDim.x) {
        // 检查当前类别是否为正类
        bool is_positive = false;
        for (int j = 0; j < num_positives; ++j) {
            if (positive_indices[j] == neg_class_idx) {
                is_positive = true;
                break;
            }
        }

        // 如果是负类，则计算与所有正类的损失
        if (!is_positive) {
            T x_neg = input_row[neg_class_idx];
            for (int j = 0; j < num_positives; ++j) {
                long pos_class_idx = positive_indices[j];
                T x_pos = input_row[pos_class_idx];
                T loss_term = 1.0f - (x_pos - x_neg);
                if (loss_term > 0) {
                    my_sum += loss_term;
                }
            }
        }
    }
    sdata[tid] = my_sum;
    __syncthreads();

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

    if (tid == 0) {
        output[sample_idx] = sdata[0] / C;
    }
}

torch::Tensor multi_label_margin_loss_cuda_forward(
    const torch::Tensor& input, 
    const torch::Tensor& target,
    const std::string& reduction)
{
    TORCH_CHECK(input.is_cuda() && target.is_cuda(), "Tensors must be on CUDA");
    TORCH_CHECK(input.dim() == 2, "Input must be 2D");
    TORCH_CHECK(target.dim() == 2, "Target must be 2D");
    TORCH_CHECK(input.size(0) == target.size(0), "Batch sizes must match");
    TORCH_CHECK(input.is_contiguous() && target.is_contiguous(), "Tensors must be contiguous");

    const int N = input.size(0);
    const int C = input.size(1);
    
    auto options = torch::TensorOptions().device(input.device()).dtype(input.dtype());
    auto sample_losses = torch::empty({N}, options);
    
    dim3 grid(N);
    dim3 block(BLOCK_SIZE);

    AT_DISPATCH_FLOATING_TYPES(input.scalar_type(), "multi_label_margin_loss_kernel", ([&] {{
        multi_label_margin_loss_kernel<scalar_t><<<grid, block>>>(
            sample_losses.data_ptr<scalar_t>(),
            input.data_ptr<scalar_t>(),
            target.data_ptr<long>(),
            N, C
        );
    }}));
    
    if (reduction == "none") {
        return sample_losses;
    } else if (reduction == "sum") {
        return sample_losses.sum();
    } else {{ // "mean"
        return sample_losses.mean();
    }}
}
"""

class ModelNew(nn.Module):
    def __init__(self, reduction='mean'):
        super(ModelNew, self).__init__()
        self.reduction = reduction
        
        self.op = load_inline(
            name='multi_label_margin_loss_op',
            cpp_sources=cpp_source,
            cuda_sources=cuda_source,
            functions=['multi_label_margin_loss_cuda_forward'],
            verbose=False
        )

    def forward(self, input_tensor: torch.Tensor, target_tensor: torch.Tensor) -> torch.Tensor:
        return self.op.multi_label_margin_loss_cuda_forward(
            input_tensor, 
            target_tensor,
            self.reduction
        )