GPUCodeForces/S1/35/softmarginloss_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 soft_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>
#include <cmath>

#define BLOCK_SIZE 256

// ----------------------------------------------------------------------------
// __device__ function for stable SoftMarginLoss calculation
// ----------------------------------------------------------------------------
template <typename T>
__device__ __forceinline__ T stable_soft_margin_loss(T input, T target) {
    T val = -target * input;
    if (val > 0) {
        return val + logf(1.0f + expf(-val));
    } else {
        return logf(1.0f + expf(val));
    }
}

// ----------------------------------------------------------------------------
// Element-wise Kernel for reduction='none'
// ----------------------------------------------------------------------------
template <typename T>
__global__ void soft_margin_loss_elementwise_kernel(
    T* output,
    const T* input,
    const T* target,
    int64_t n_elements) 
{
    int64_t idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= n_elements) return;
    output[idx] = stable_soft_margin_loss(input[idx], target[idx]);
}

// ----------------------------------------------------------------------------
// Reduction Kernel: Stage 1 (calculate loss and reduce within blocks)
// ----------------------------------------------------------------------------
template <typename T>
__global__ void soft_margin_loss_reduce_kernel_stage1(
    T* block_results,
    const T* input,
    const T* target,
    int64_t n_elements)
{
    __shared__ T sdata[BLOCK_SIZE];
    
    int64_t tid = threadIdx.x;
    int64_t i = blockIdx.x * blockDim.x + tid;
    
    T my_sum = 0.0f;
    // Grid-stride loop to process all elements
    while (i < n_elements) {
        my_sum += stable_soft_margin_loss(input[i], target[i]);
        i += gridDim.x * blockDim.x;
    }
    sdata[tid] = my_sum;
    __syncthreads(); 

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

    if (tid == 0) {
        block_results[blockIdx.x] = sdata[0];
    }
}

torch::Tensor soft_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.sizes() == target.sizes(), "Input and target shapes must match");
    TORCH_CHECK(input.is_contiguous() && target.is_contiguous(), "Tensors must be contiguous");

    const int64_t n_elements = input.numel();
    const auto scalar_type = input.scalar_type();

    if (reduction == "none") {
        auto output = torch::empty_like(input);
        const int num_blocks = (n_elements + BLOCK_SIZE - 1) / BLOCK_SIZE;
        AT_DISPATCH_FLOATING_TYPES(scalar_type, "soft_margin_loss_elementwise", ([&] {
            soft_margin_loss_elementwise_kernel<scalar_t><<<num_blocks, BLOCK_SIZE>>>(
                output.data_ptr<scalar_t>(),
                input.data_ptr<scalar_t>(),
                target.data_ptr<scalar_t>(),
                n_elements);
        }));
        return output;
    } 
    else // 'sum' or 'mean'
    {
        // Limit grid size to avoid creating a massive intermediate tensor
        int max_grid_size = 4096; 
        int num_blocks = std::min(max_grid_size, (int)((n_elements + BLOCK_SIZE - 1) / BLOCK_SIZE));
        
        auto options = torch::TensorOptions().device(input.device()).dtype(input.dtype());
        auto block_results = torch::empty({num_blocks}, options);

        AT_DISPATCH_FLOATING_TYPES(scalar_type, "soft_margin_loss_reduce_stage1", ([&] {
            soft_margin_loss_reduce_kernel_stage1<scalar_t><<<num_blocks, BLOCK_SIZE>>>(
                block_results.data_ptr<scalar_t>(),
                input.data_ptr<scalar_t>(),
                target.data_ptr<scalar_t>(),
                n_elements);
        }));
        
        torch::Tensor total_sum = block_results.sum();
        
        if (reduction == "mean") {
            return total_sum / n_elements;
        }
        return total_sum;
    }
}
"""

class ModelNew(nn.Module):
    """
    使用自定义 CUDA 内核进行优化的 SoftMarginLoss 模型。
    """
    def __init__(self, reduction='mean'):
        super(ModelNew, self).__init__()
        self.reduction = reduction
        
        self.soft_margin_loss_op = load_inline(
            name='soft_margin_loss_op',
            cpp_sources=cpp_source,
            cuda_sources=cuda_source,
            functions=['soft_margin_loss_cuda_forward'],
            verbose=False
        )

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