GPUCodeForces/S1/uucoco_#6/HellingerDistance_cuda.py

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

N, C, H, W = 32, 64, 56, 56
EPS = 1e-6

assert (C * H * W) % 4 == 0, "Instance size (C*H*W) must be a multiple of 4"


class ModelNew(nn.Module):

    def __init__(self):
        super().__init__()
        self.block_size = 512
        self.eps = EPS
        self._compile_cuda_kernel()

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

        void hellinger_sum_cuda(
            torch::Tensor x,
            torch::Tensor y,
            torch::Tensor out,
            float eps,
            int N,
            int D);
        """

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

        #define BLOCK_SIZE {self.block_size}
        #define WARP_SIZE 32
        #define ILP 4  // 每个线程处理 4 个 float4 (16个 float)

        // Warp 归约
        __inline__ __device__ float warp_reduce_sum(float val) {{
            #pragma unroll
            for (int offset = 16; offset > 0; offset /= 2) {{
                val += __shfl_down_sync(0xffffffff, val, offset);
            }}
            return val;
        }}


        __global__ __launch_bounds__(BLOCK_SIZE)
        void hellinger_split_kernel(
            const float* __restrict__ x,
            const float* __restrict__ y,
            float* __restrict__ out,
            float eps,
            int D_vec_total  // D / 4
        ) {{
            // 1. 任务映射: Grid(Split, Batch)
            const int n_idx = blockIdx.y;     // Batch Index
            const int split_idx = blockIdx.x; // Split Index
            const int num_splits = gridDim.x;

            // 2. 计算分块范围
            const int chunk_size = (D_vec_total + num_splits - 1) / num_splits;
            const int start_idx = split_idx * chunk_size;
            const int end_idx = min(start_idx + chunk_size, D_vec_total);

            if (start_idx >= D_vec_total) return;

            // 3. 指针设置
            const int64_t batch_offset = (int64_t)n_idx * D_vec_total * 4;

            const float4* curr_x = reinterpret_cast<const float4*>(x + batch_offset) + start_idx + threadIdx.x;
            const float4* curr_y = reinterpret_cast<const float4*>(y + batch_offset) + start_idx + threadIdx.x;
            const float4* end_ptr = reinterpret_cast<const float4*>(x + batch_offset) + end_idx;

            // 4. 累加器
            float sum[ILP];
            #pragma unroll
            for (int k=0; k<ILP; ++k) sum[k] = 0.0f;

            const int stride = BLOCK_SIZE * ILP;

            // 5. 主循环 (Pointer Chasing)
            while (curr_x + (ILP - 1) * BLOCK_SIZE < end_ptr) {{
                float4 r_x[ILP];
                float4 r_y[ILP];

                // Load
                #pragma unroll
                for (int k = 0; k < ILP; ++k) {{
                    r_x[k] = __ldg(curr_x + k * BLOCK_SIZE);
                    r_y[k] = __ldg(curr_y + k * BLOCK_SIZE);
                }}

                // Compute: (sqrt(x) - sqrt(y))^2
                #pragma unroll
                for (int k = 0; k < ILP; ++k) {{
                    float sx_x = sqrtf(r_x[k].x + eps); float sy_x = sqrtf(r_y[k].x + eps);
                    float diff_x = sx_x - sy_x; sum[k] += diff_x * diff_x;

                    float sx_y = sqrtf(r_x[k].y + eps); float sy_y = sqrtf(r_y[k].y + eps);
                    float diff_y = sx_y - sy_y; sum[k] += diff_y * diff_y;

                    float sx_z = sqrtf(r_x[k].z + eps); float sy_z = sqrtf(r_y[k].z + eps);
                    float diff_z = sx_z - sy_z; sum[k] += diff_z * diff_z;

                    float sx_w = sqrtf(r_x[k].w + eps); float sy_w = sqrtf(r_y[k].w + eps);
                    float diff_w = sx_w - sy_w; sum[k] += diff_w * diff_w;
                }}

                curr_x += stride;
                curr_y += stride;
            }}

            // 6. 尾部循环
            while (curr_x < end_ptr) {{
                float4 vx = __ldg(curr_x);
                float4 vy = __ldg(curr_y);

                float d1 = sqrtf(vx.x + eps) - sqrtf(vy.x + eps);
                float d2 = sqrtf(vx.y + eps) - sqrtf(vy.y + eps);
                float d3 = sqrtf(vx.z + eps) - sqrtf(vy.z + eps);
                float d4 = sqrtf(vx.w + eps) - sqrtf(vy.w + eps);

                sum[0] += d1*d1 + d2*d2 + d3*d3 + d4*d4;

                curr_x += BLOCK_SIZE;
                curr_y += BLOCK_SIZE;
            }}

            // 7. 汇总
            float local_sum = 0.0f;
            #pragma unroll
            for (int k=0; k<ILP; ++k) local_sum += sum[k];

            // 8. Warp 归约
            float warp_sum = warp_reduce_sum(local_sum);

            // 9. Semi-Sync Reduction (减少原子操作)
            __shared__ float s_warp_sums[BLOCK_SIZE / WARP_SIZE];
            const int lane_id = threadIdx.x % WARP_SIZE;
            const int warp_id = threadIdx.x / WARP_SIZE;

            if (lane_id == 0) {{
                s_warp_sums[warp_id] = warp_sum;
            }}
            __syncthreads();

            if (warp_id == 0) {{
                float block_val = 0.0f;
                if (lane_id < (BLOCK_SIZE / WARP_SIZE)) {{
                    block_val = s_warp_sums[lane_id];
                }}
                block_val = warp_reduce_sum(block_val);

                if (lane_id == 0) {{
                    atomicAdd(&out[n_idx], block_val);
                }}
            }}
        }}

        void hellinger_sum_cuda(
            torch::Tensor x,
            torch::Tensor y,
            torch::Tensor out,
            float eps,
            int N,
            int D) 
        {{
            int D_vec = D / 4;

            int device_id;
            cudaGetDevice(&device_id);
            int sm_count;
            cudaDeviceGetAttribute(&sm_count, cudaDevAttrMultiProcessorCount, device_id);


            int target_blocks = sm_count * 4; 
            int splits = (target_blocks + N - 1) / N;


            int max_splits = (D_vec + 1024 - 1) / 1024;
            if (splits > max_splits) splits = max_splits;
            if (splits < 1) splits = 1;

            if (splits > 512) splits = 512;

            dim3 blocks(splits, N);
            dim3 threads(BLOCK_SIZE);

            hellinger_split_kernel<<<blocks, threads>>>(
                x.data_ptr<float>(),
                y.data_ptr<float>(),
                out.data_ptr<float>(),
                eps,
                D_vec
            );
        }}
        """

        self.op = load_inline(
            name='hellinger_cuda_opt_v1',
            cpp_sources=cpp_source,
            cuda_sources=cuda_source,
            functions=['hellinger_sum_cuda'],
            extra_cuda_cflags=[
                '-O3',
                '--use_fast_math',
                '-Xptxas=-v'
            ],
            verbose=False
        )

    def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
        if not x.is_contiguous(): x = x.contiguous()
        if not y.is_contiguous(): y = y.contiguous()

        N, C, H, W = x.size()
        D = C * H * W

        out = torch.zeros(N, device=x.device, dtype=torch.float32)

        self.op.hellinger_sum_cuda(
            x,
            y,
            out,
            self.eps,
            N,
            D
        )

        return torch.sqrt(out + self.eps) * 0.70710678