GPUCodeForces/S1/29/kldivloss_cuda.py

# kldivloss_cuda_ultra.py
import torch
from torch.utils.cpp_extension import load_inline
from kldivloss_torch import BATCH_SIZE, DIM 

class ModelNew(torch.nn.Module):
    
    def __init__(self):
        super().__init__()
        self._compile_cuda_kernel()

    def _compile_cuda_kernel(self):
        cpp_source = """
        #include <torch/extension.h>
        torch::Tensor kldiv_forward_cuda(torch::Tensor input_logits, torch::Tensor target_prob);
        """

        cuda_source = """
        #include <cuda_runtime.h>
        #include <cooperative_groups.h>
        namespace cg = cooperative_groups;

        #define BLOCK_SIZE 256
        #define VEC_SIZE 4
        #define WARP_SIZE 32

        // Ultra-fast warp reduction using cooperative groups
        __device__ __forceinline__ double warp_reduce_sum_cg(double val) {
            #pragma unroll
            for (int offset = 16; offset > 0; offset >>= 1) {
                val += __shfl_down_sync(0xffffffff, val, offset);
            }
            return val;
        }

        // Optimized block reduction
        __device__ __forceinline__ double block_reduce_sum(double val) {
            __shared__ double warp_sums[8];
            
            int lane = threadIdx.x & 31;
            int warp_id = threadIdx.x >> 5;
            
            val = warp_reduce_sum_cg(val);
            
            if (lane == 0) {
                warp_sums[warp_id] = val;
            }
            __syncthreads();
            
            if (warp_id == 0) {
                val = (lane < 8) ? warp_sums[lane] : 0.0;
                val = warp_reduce_sum_cg(val);
            }
            
            return val;
        }

        // Safe and fast KL term computation
        __device__ __forceinline__ double safe_kl_term(float p, float log_q) {
            // Avoid NaN: when p < eps, contribution is 0
            return (p > 1e-8f) ? ((double)p * (__logf(p) - log_q)) : 0.0;
        }

        // Version 1: Maximally unrolled with 8x vectorization
        __global__ void kldiv_kernel_v1(
            const float* __restrict__ input_logits,
            const float* __restrict__ target_prob,
            double* __restrict__ output_sum,
            int N_elements
        ) {
            double sum = 0.0;
            
            int N_vec = N_elements >> 2; // / 4
            int tid = (blockIdx.x * blockDim.x + threadIdx.x);
            int stride = gridDim.x * blockDim.x;
            
            const float4* __restrict__ in4 = (const float4*)input_logits;
            const float4* __restrict__ tgt4 = (const float4*)target_prob;

            // 8x unrolled loop
            for (int i = tid; i < N_vec; i += stride << 3) {
                #pragma unroll
                for (int u = 0; u < 8; u++) {
                    int idx = i + (u * stride);
                    if (idx < N_vec) {
                        float4 lq = in4[idx];
                        float4 p = tgt4[idx];
                        
                        sum += safe_kl_term(p.x, lq.x);
                        sum += safe_kl_term(p.y, lq.y);
                        sum += safe_kl_term(p.z, lq.z);
                        sum += safe_kl_term(p.w, lq.w);
                    }
                }
            }
            
            sum = block_reduce_sum(sum);
            
            if (threadIdx.x == 0) {
                output_sum[blockIdx.x] = sum;
            }
        }

        // Version 2: Register-tiled with manual prefetching
        __global__ void kldiv_kernel_v2(
            const float* __restrict__ input_logits,
            const float* __restrict__ target_prob,
            double* __restrict__ output_sum,
            int N_elements
        ) {
            double sum = 0.0;
            
            int N_vec = N_elements >> 2;
            int tid = blockIdx.x * blockDim.x + threadIdx.x;
            int stride = gridDim.x * blockDim.x;
            
            const float4* __restrict__ in4 = (const float4*)input_logits;
            const float4* __restrict__ tgt4 = (const float4*)target_prob;

            // Process 4 vectors per iteration (register tiling)
            for (int i = tid; i < N_vec; i += stride * 4) {
                float4 lq0, lq1, lq2, lq3;
                float4 p0, p1, p2, p3;
                
                // Load 4 vectors
                if (i < N_vec) { lq0 = in4[i]; p0 = tgt4[i]; }
                if (i + stride < N_vec) { lq1 = in4[i + stride]; p1 = tgt4[i + stride]; }
                if (i + stride*2 < N_vec) { lq2 = in4[i + stride*2]; p2 = tgt4[i + stride*2]; }
                if (i + stride*3 < N_vec) { lq3 = in4[i + stride*3]; p3 = tgt4[i + stride*3]; }
                
                // Compute
                if (i < N_vec) {
                    sum += safe_kl_term(p0.x, lq0.x) + safe_kl_term(p0.y, lq0.y);
                    sum += safe_kl_term(p0.z, lq0.z) + safe_kl_term(p0.w, lq0.w);
                }
                if (i + stride < N_vec) {
                    sum += safe_kl_term(p1.x, lq1.x) + safe_kl_term(p1.y, lq1.y);
                    sum += safe_kl_term(p1.z, lq1.z) + safe_kl_term(p1.w, lq1.w);
                }
                if (i + stride*2 < N_vec) {
                    sum += safe_kl_term(p2.x, lq2.x) + safe_kl_term(p2.y, lq2.y);
                    sum += safe_kl_term(p2.z, lq2.z) + safe_kl_term(p2.w, lq2.w);
                }
                if (i + stride*3 < N_vec) {
                    sum += safe_kl_term(p3.x, lq3.x) + safe_kl_term(p3.y, lq3.y);
                    sum += safe_kl_term(p3.z, lq3.z) + safe_kl_term(p3.w, lq3.w);
                }
            }
            
            sum = block_reduce_sum(sum);
            
            if (threadIdx.x == 0) {
                output_sum[blockIdx.x] = sum;
            }
        }

        torch::Tensor kldiv_forward_cuda(torch::Tensor input_logits, torch::Tensor target_prob) {
            TORCH_CHECK(input_logits.is_cuda() && target_prob.is_cuda(), "Inputs must be CUDA tensors");
            
            input_logits = input_logits.contiguous();
            target_prob = target_prob.contiguous();

            int N = input_logits.numel();
            
            if (N % VEC_SIZE != 0) {
                TORCH_CHECK(false, "Total elements must be divisible by 4");
            }

            const int block_size = BLOCK_SIZE;
            // Optimal grid size: balance between parallelism and reduction overhead
            const int grid_size = min(1024, (N / VEC_SIZE + block_size - 1) / block_size);
            
            auto partial_sum = torch::empty({grid_size}, input_logits.options().dtype(torch::kFloat64));
            
            // Choose best kernel based on problem size
            if (N >= 1048576) { // >= 1M elements, use v2 (register tiling)
                kldiv_kernel_v2<<<grid_size, block_size>>>(
                    input_logits.data_ptr<float>(),
                    target_prob.data_ptr<float>(),
                    partial_sum.data_ptr<double>(),
                    N
                );
            } else { // Use v1 (maximally unrolled)
                kldiv_kernel_v1<<<grid_size, block_size>>>(
                    input_logits.data_ptr<float>(),
                    target_prob.data_ptr<float>(),
                    partial_sum.data_ptr<double>(),
                    N
                );
            }
            
            // Final reduction on device
            double total = partial_sum.sum().item<double>();
            float result = (float)(total / N);

            return torch::tensor(result, input_logits.options());
        }
        """
        
        self.kldiv_op = load_inline(
            name="kldiv_ultra_op",
            cpp_sources=cpp_source,
            cuda_sources=cuda_source,
            functions=["kldiv_forward_cuda"],
            extra_cuda_cflags=["-O3", "--use_fast_math", "-lineinfo", "--maxrregcount=64"],
            verbose=True
        )

    def forward(self, input_logits: torch.Tensor, target_prob: torch.Tensor) -> torch.Tensor:
        return self.kldiv_op.kldiv_forward_cuda(input_logits, target_prob)