GPUCodeForces/S1/uucoco_#14/softmin_cuda.py

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

# --- Hyperparameters ---
N, D = 16, 512
DIM = 1
BLOCK_SIZE = 512  # 通常设置为 D 或 D 的因子，用于 Reduction


class ModelNew(nn.Module):

    def __init__(self, dim=DIM):
        super().__init__()
        self.dim = dim
        self.block_size = BLOCK_SIZE
        self._compile_cuda_kernel()

    def _compile_cuda_kernel(self):
        cpp_source = """
        #include <torch/extension.h>

        torch::Tensor softmin_forward_cuda(
            torch::Tensor input,
            int dim
        );
        """

        cuda_source = f"""
        #include <cuda_runtime.h>
        #include <math.h>

        #define BLOCK_SIZE {self.block_size}
        #define MAX_FLOAT -3.402823466e+38F // 负的最大浮点数

        /*
         * Pass 1 Kernel: 查找稳定的最大值 m = max_j(-x_j)
         * 每个 Block 处理 N 行中的一行。
         */
        __global__ void softmin_max_kernel(
            const float* __restrict__ input,
            float* __restrict__ max_neg_output, // 存储 max_j(-x_j)
            int N, 
            int D
        ) {{
            // 共享内存用于存储负输入的局部最大值
            __shared__ float s_max[BLOCK_SIZE];

            const int n_idx = blockIdx.x; // 当前处理的向量/行
            const int tid = threadIdx.x;

            const float* p_in = input + (int64_t)n_idx * D;
            float block_max = MAX_FLOAT;

            // 1. Grid-Stride Loop: 找出线程负责区域的局部最大值
            for (int d = tid; d < D; d += BLOCK_SIZE) {{
                // 注意：这里是对 -x 进行最大值搜索
                block_max = fmaxf(block_max, -p_in[d]);
            }}

            s_max[tid] = block_max;
            __syncthreads();


            // 2. 共享内存并行规约 (找出 Block 内全局最大值)
            for (int offset = BLOCK_SIZE / 2; offset > 0; offset >>= 1) {{
                if (tid < offset) {{
                    s_max[tid] = fmaxf(s_max[tid], s_max[tid + offset]);
                }}
                __syncthreads();
            }}

            // 3. 存储最终最大值 m = max_j(-x_j)
            if (tid == 0) {{
                max_neg_output[n_idx] = s_max[0];
            }}
        }}


        /*
         * Pass 2 Kernel: 计算 Softmin 最终结果
         */
        __global__ void softmin_sum_output_kernel(
            const float* __restrict__ input,
            const float* __restrict__ max_neg_input, // Pass 1 的结果: m
            float* __restrict__ output,
            int N, 
            int D
        ) {{
            __shared__ float s_sum[BLOCK_SIZE];

            const int n_idx = blockIdx.x;
            const int tid = threadIdx.x;

            const float* p_in = input + (int64_t)n_idx * D;
            float* p_out = output + (int64_t)n_idx * D;

            const float max_neg_x = max_neg_input[n_idx]; // 当前向量的最大负值 m

            float thread_sum = 0.0f;

            // 1. Grid-Stride Loop: 计算指数和 (Sum Exp)
            for (int d = tid; d < D; d += BLOCK_SIZE) {{
                // 稳定计算: exp(-x_d - m)
                // 这一步也计算了最终的分子 exp(-x_i - m)
                float exp_val = expf(-p_in[d] - max_neg_x);
                thread_sum += exp_val;
                p_out[d] = exp_val; // 临时存储分子
            }}

            s_sum[tid] = thread_sum;
            __syncthreads();


            // 2. 共享内存并行规约 (求 Block 内全局总和)
            for (int offset = BLOCK_SIZE / 2; offset > 0; offset >>= 1) {{
                if (tid < offset) {{
                    s_sum[tid] += s_sum[tid + offset];
                }}
                __syncthreads();
            }}

            const float total_sum = s_sum[0]; // 分母: Z = sum_j exp(-x_j - m)
            __syncthreads(); // 确保所有线程都拿到了最终总和


            // 3. 计算最终 Softmin 结果 (Softmin = 分子 / 分母)
            for (int d = tid; d < D; d += BLOCK_SIZE) {{
                // p_out[d] 现在存储着 exp(-x_d - m)
                p_out[d] = p_out[d] / total_sum;
            }}
        }}


        torch::Tensor softmin_forward_cuda(
            torch::Tensor input,
            int dim
        ) {{

            TORCH_CHECK(input.is_cuda(), "Input must be CUDA tensor");
            TORCH_CHECK(input.dim() == 2, "CUDA Kernel only supports 2D (N, D) input");
            TORCH_CHECK(dim == 1, "CUDA Kernel only supports dim=1");
            TORCH_CHECK(input.scalar_type() == torch::kFloat32, "Input must be float32");

            input = input.contiguous();

            const int N = input.size(0);
            const int D = input.size(1);

            // 临时张量用于存储 Pass 1 的结果 (Max)
            auto max_neg_output = torch::empty({{N}}, input.options());

            // 输出张量 C
            auto output = torch::empty_like(input);

            // Grid 尺寸: N 个 Block, 每个 Block 处理 N 行中的一行
            dim3 grid_dim(N);
            dim3 block_dim(BLOCK_SIZE);

            // --- Pass 1: 计算 Max_j(-x_j) ---
            softmin_max_kernel<<<grid_dim, block_dim>>>(
                input.data_ptr<float>(),
                max_neg_output.data_ptr<float>(),
                N, D
            );

            // --- Pass 2: 计算 Sum 和最终 Output ---
            softmin_sum_output_kernel<<<grid_dim, block_dim>>>(
                input.data_ptr<float>(),
                max_neg_output.data_ptr<float>(),
                output.data_ptr<float>(),
                N, D
            );

            return output;
        }}
        """

        nvcc_flags = ['-O3']

        self.softmin_op = load_inline(
            name="softmin_opt_two_pass",
            cpp_sources=cpp_source,
            cuda_sources=cuda_source,
            functions=["softmin_forward_cuda"],
            extra_cuda_cflags=nvcc_flags,
            verbose=False
        )

    def forward(self, input: torch.Tensor) -> torch.Tensor:
        if not input.is_cuda: input = input.cuda()
        input_cont = input.contiguous()

        return self.softmin_op.softmin_forward_cuda(
            input_cont,
            self.dim
        )