GPUCodeForces/S1/36/multimarginloss_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_margin_loss_cuda_forward(
    const torch::Tensor& input, 
    const torch::Tensor& target,
    int p,
    float margin,
    const std::string& reduction
);
"""

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

#define BLOCK_SIZE 256

template <typename T>
__global__ void multi_margin_loss_kernel(
    T* output,          // (N)
    const T* input,     // (N, C)
    const long* target, // (N)
    const int N,
    const int C,
    const int p,
    const T margin)
{
    // Each block processes one sample
    int sample_idx = blockIdx.x;
    if (sample_idx >= N) return;

    // --- Shared memory for broadcasting target index and correct class score ---
    __shared__ long y_shared;
    __shared__ T x_correct_shared;

    // Thread 0 of each block loads the critical data
    if (threadIdx.x == 0) {
        long y_idx = target[sample_idx];
        y_shared = y_idx;
        x_correct_shared = input[sample_idx * C + y_idx];
    }
    __syncthreads();
    
    // All threads in the block now have access to y_shared and x_correct_shared
    long y = y_shared;
    T x_correct = x_correct_shared;
    
    __shared__ T sdata[BLOCK_SIZE];
    int tid = threadIdx.x;
    T my_sum = 0.0f;

    // --- Grid-stride loop for this block to iterate over all classes ---
    for (int class_idx = tid; class_idx < C; class_idx += blockDim.x) {
        if (class_idx == y) {
            continue; // Skip the target class
        }
        
        T x_other = input[sample_idx * C + class_idx];
        T loss_term = margin - x_correct + x_other;

        if (loss_term > 0) {
            if (p == 2) {
                loss_term *= loss_term;
            }
            my_sum += loss_term;
        }
    }

    sdata[tid] = my_sum;
    __syncthreads();

    // --- Intra-block reduction to sum up all thread-local sums ---
    for (int s = blockDim.x / 2; s > 0; s >>= 1) {
        if (tid < s) {
            sdata[tid] += sdata[tid + s];
        }
        __syncthreads();
    }

    // --- Thread 0 writes the final result for this sample ---
    if (tid == 0) {
        output[sample_idx] = sdata[0] / C;
    }
}


torch::Tensor multi_margin_loss_cuda_forward(
    const torch::Tensor& input, 
    const torch::Tensor& target,
    int p,
    float margin,
    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() == 1, "Target must be 1D");
    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);
    
    // Launch one block per sample
    dim3 grid(N);
    dim3 block(BLOCK_SIZE);

    AT_DISPATCH_FLOATING_TYPES(input.scalar_type(), "multi_margin_loss_kernel", ([&] {
        multi_margin_loss_kernel<scalar_t><<<grid, block>>>(
            sample_losses.data_ptr<scalar_t>(),
            input.data_ptr<scalar_t>(),
            target.data_ptr<long>(),
            N, C, p, static_cast<scalar_t>(margin)
        );
    }));
    
    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, p=1, margin=1.0, reduction='mean'):
        super(ModelNew, self).__init__()
        self.p = p
        self.margin = margin
        self.reduction = reduction
        
        self.op = load_inline(
            name='multi_margin_loss_op',
            cpp_sources=cpp_source,
            cuda_sources=cuda_source,
            functions=['multi_margin_loss_cuda_forward'],
            verbose=False
        )

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