GPUCodeForces/S1/uucoco_#97/meanstdnormalizeclip_cuda.py

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

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

#define WARP_SIZE 32

__inline__ __device__ float warp_reduce_sum(float val) {
    for (int offset = WARP_SIZE/2; offset > 0; offset /= 2) {
        val += __shfl_down_sync(0xffffffff, val, offset);
    }
    return val;
}

__global__ void normalize_clamp_kernel(
    const float* __restrict__ input,
    float* __restrict__ output,
    int batch_size,
    int dim,
    float min_val,
    float max_val,
    float eps) {

    int b = blockIdx.x;
    if (b >= batch_size) return;

    int tid = threadIdx.x;
    int stride = blockDim.x;

    const float* x_ptr = input + b * dim;
    float* out_ptr = output + b * dim;

    float sum = 0.0f;
    for (int i = tid; i < dim; i += stride) {
        sum += x_ptr[i];
    }

    sum = warp_reduce_sum(sum);

    __shared__ float shared_sum[32];
    int lane = tid % WARP_SIZE;
    int wid = tid / WARP_SIZE;

    if (lane == 0) {
        shared_sum[wid] = sum;
    }
    __syncthreads();

    if (tid < blockDim.x / WARP_SIZE) {
        sum = shared_sum[tid];
    } else {
        sum = 0.0f;
    }

    if (wid == 0) {
        sum = warp_reduce_sum(sum);
    }

    __shared__ float mean_shared;
    if (tid == 0) {
        mean_shared = sum / dim;
    }
    __syncthreads();

    float mean = mean_shared;

    float var_sum = 0.0f;
    for (int i = tid; i < dim; i += stride) {
        float diff = x_ptr[i] - mean;
        var_sum += diff * diff;
    }

    var_sum = warp_reduce_sum(var_sum);

    if (lane == 0) {
        shared_sum[wid] = var_sum;
    }
    __syncthreads();

    if (tid < blockDim.x / WARP_SIZE) {
        var_sum = shared_sum[tid];
    } else {
        var_sum = 0.0f;
    }

    if (wid == 0) {
        var_sum = warp_reduce_sum(var_sum);
    }

    __shared__ float std_shared;
    if (tid == 0) {
        float variance = var_sum / (dim - 1);
        std_shared = sqrtf(variance);
    }
    __syncthreads();

    float std = std_shared;

    for (int i = tid; i < dim; i += stride) {
        float norm = (x_ptr[i] - mean) / (std + eps);

        float clamped;
        if (norm < min_val) {
            clamped = min_val;
        } else if (norm > max_val) {
            clamped = max_val;
        } else {
            clamped = norm;
        }

        out_ptr[i] = clamped;
    }
}

torch::Tensor normalize_clamp_cuda(
    torch::Tensor x,
    float min_val,
    float max_val,
    float eps) {

    int batch_size = x.size(0);
    int dim = x.size(1);

    auto output = torch::empty_like(x);

    int threads = 256;

    normalize_clamp_kernel<<<batch_size, threads>>>(
        x.data_ptr<float>(),
        output.data_ptr<float>(),
        batch_size,
        dim,
        min_val,
        max_val,
        eps
    );

    return output;
}
"""

cpp_source = """
torch::Tensor normalize_clamp_cuda(
    torch::Tensor x,
    float min_val,
    float max_val,
    float eps);
"""

cuda_module = load_inline(
    name="normalize_clamp_module",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["normalize_clamp_cuda"],
    verbose=True
)


class ModelNew(nn.Module):
    def __init__(self, min_val, max_val, eps=1e-5):
        super(ModelNew, self).__init__()
        self.min_val = min_val
        self.max_val = max_val
        self.eps = eps

    def forward(self, x):
        return cuda_module.normalize_clamp_cuda(x, self.min_val, self.max_val, self.eps)