GPUCodeForces/S1/wut0n_#1/Layernorm_cudacode.py

# LayerNorm 向量化Welford

import torch
from torch.utils.cpp_extension import load_inline

# --- 核心 CUDA C++ 代码 ---
layernorm_source = """
#include <torch/extension.h>
#include <cuda_runtime.h>
#include <math.h>

#define WARP_SIZE 32

struct WelfordData {
    float mean;
    float m2;   // Sum of squares of differences from the current mean
    int count;
};

// Warp 级合并 WelfordData
__device__ __forceinline__ WelfordData warp_reduce_welford(WelfordData data) {
    for (int offset = WARP_SIZE / 2; offset > 0; offset >>= 1) {
        WelfordData other;
        other.mean = __shfl_down_sync(0xffffffff, data.mean, offset);
        other.m2 = __shfl_down_sync(0xffffffff, data.m2, offset);
        other.count = __shfl_down_sync(0xffffffff, data.count, offset);

        if (other.count > 0) {
            float delta = other.mean - data.mean;
            float new_count = data.count + other.count;
            data.m2 += other.m2 + delta * delta * data.count * other.count / new_count;
            data.mean += delta * other.count / new_count;
            data.count = new_count;
        }
    }
    return data;
}

__global__ void layernorm_kernel_welford_vectorized(
    const float* __restrict__ x,
    const float* __restrict__ weight,
    const float* __restrict__ bias,
    float* __restrict__ y,
    int batch,
    int features,
    float eps
) {
    int row = blockIdx.x;
    if (row >= batch) return;

    int tid = threadIdx.x;
    int warp_id = tid / WARP_SIZE;
    int lane_id = tid % WARP_SIZE;
    int num_warps = (blockDim.x + WARP_SIZE - 1) / WARP_SIZE;

    __shared__ float s_mean;
    __shared__ float s_inv_std;
    __shared__ WelfordData s_welford_results[32]; 

    const float* x_row = x + row * features;
    float* y_row = y + row * features;

    // 向量化单次遍历计算全局均值和方差 
    WelfordData local_data;
    local_data.mean = 0.0f;
    local_data.m2 = 0.0f;
    local_data.count = 0;

    // 使用向量化处理以提升性能
    if (features % 4 == 0) {
        for (int i = tid * 4; i < features; i += blockDim.x * 4) {
            float4 vec = *reinterpret_cast<const float4*>(x_row + i);
            local_data.count++;
            float delta = vec.x - local_data.mean;
            local_data.mean += delta / local_data.count;
            float delta2 = vec.x - local_data.mean;
            local_data.m2 += delta * delta2;

            local_data.count++;
            delta = vec.y - local_data.mean;
            local_data.mean += delta / local_data.count;
            delta2 = vec.y - local_data.mean;
            local_data.m2 += delta * delta2;

            local_data.count++;
            delta = vec.z - local_data.mean;
            local_data.mean += delta / local_data.count;
            delta2 = vec.z - local_data.mean;
            local_data.m2 += delta * delta2;

            local_data.count++;
            delta = vec.w - local_data.mean;
            local_data.mean += delta / local_data.count;
            delta2 = vec.w - local_data.mean;
            local_data.m2 += delta * delta2;
        }
    } else {
        for (int i = tid; i < features; i += blockDim.x) {
            float val = x_row[i];
            local_data.count++;
            float delta = val - local_data.mean;
            local_data.mean += delta / local_data.count;
            float delta2 = val - local_data.mean;
            local_data.m2 += delta * delta2;
        }
    }

    // Warp 级归约
    WelfordData warp_data = warp_reduce_welford(local_data);

    // 将 Warp 结果写入共享内存
    if (lane_id == 0) {
        s_welford_results[warp_id] = warp_data;
    }
    __syncthreads();

    // Block 级归约（在第一个 warp 中完成）
    if (warp_id == 0) {
        WelfordData block_data;
        block_data.mean = 0.0f;
        block_data.m2 = 0.0f;
        block_data.count = 0;

        if (lane_id < num_warps) {
            block_data = s_welford_results[lane_id];
        }

        block_data = warp_reduce_welford(block_data);

        if (lane_id == 0) {
            float mean = block_data.mean;
            float var = block_data.m2 / features;
            s_mean = mean;
            s_inv_std = rsqrtf(var + eps);
        }
    }
    __syncthreads();

    float mean = s_mean;
    float inv_std = s_inv_std;

    //应用归一化（向量化存储）
    if (features % 4 == 0) {
        for (int i = tid * 4; i < features; i += blockDim.x * 4) {
            float4 vec = *reinterpret_cast<const float4*>(x_row + i);
            float4 w_vec = *reinterpret_cast<const float4*>(weight + i);
            float4 b_vec = *reinterpret_cast<const float4*>(bias + i);

            vec.x = (vec.x - mean) * inv_std * w_vec.x + b_vec.x;
            vec.y = (vec.y - mean) * inv_std * w_vec.y + b_vec.y;
            vec.z = (vec.z - mean) * inv_std * w_vec.z + b_vec.z;
            vec.w = (vec.w - mean) * inv_std * w_vec.w + b_vec.w;

            *reinterpret_cast<float4*>(y_row + i) = vec;
        }
    } else {
        for (int i = tid; i < features; i += blockDim.x) {
            float v = x_row[i];
            float w = weight[i];
            float b = bias[i];
            y_row[i] = (v - mean) * inv_std * w + b;
        }
    }
}

torch::Tensor layernorm_cuda(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, float eps) {
    TORCH_CHECK(x.is_cuda(), "x 必须是 CUDA 张量");
    TORCH_CHECK(weight.is_cuda(), "weight 必须是 CUDA 张量");
    TORCH_CHECK(bias.is_cuda(), "bias 必须是 CUDA 张量");
    TORCH_CHECK(x.dim() == 2, "当前内核仅支持二维输入张量");
    TORCH_CHECK(weight.dim() == 1, "LayerNorm 权重必须是一维向量");
    TORCH_CHECK(bias.dim() == 1, "LayerNorm 偏置必须是一维向量");
    TORCH_CHECK(x.size(1) == weight.size(0), "输入最后一维与权重长度不匹配");
    TORCH_CHECK(weight.size(0) == bias.size(0), "权重和偏置长度必须相同");

    int batch = x.size(0);
    int features = x.size(1);

    auto y = torch::empty_like(x);

    //动态线程配置
    int threads;
    if (features <= 64) {
        threads = 64;
    } else if (features <= 256) {
        threads = 128;
    } else if (features <= 1024) {
        threads = 256;
    } else {
        threads = 512;
    }

    threads = (threads + WARP_SIZE - 1) / WARP_SIZE * WARP_SIZE;
    threads = min(threads, features);

    size_t shared_mem = sizeof(float) * 2 + sizeof(WelfordData) * ((threads + WARP_SIZE - 1) / WARP_SIZE);

    layernorm_kernel_welford_vectorized<<<batch, threads, shared_mem>>>(
        x.data_ptr<float>(),
        weight.data_ptr<float>(),
        bias.data_ptr<float>(),
        y.data_ptr<float>(),
        batch,
        features,
        eps
    );

    return y;
}
"""

layernorm_cpp_source = """
torch::Tensor layernorm_cuda(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, float eps);
"""

layernorm = load_inline(
    name="layernorm_welford_vectorized_final",
    cpp_sources=layernorm_cpp_source,
    cuda_sources=layernorm_source,
    functions=["layernorm_cuda"],
    extra_cuda_cflags=["-O3", "--use_fast_math"],
    verbose=True
)


class ModelNew(torch.nn.Module):
    def __init__(self, eps: float = 1e-5):
        super(ModelNew, self).__init__()
        self.eps = eps
        self.weight = None
        self.bias = None
        self.layernorm = layernorm
        self._initialized = False

    def forward(self, x):
        if self.weight is None:
            feature_dim = x.size(1)
            self.weight = torch.nn.Parameter(torch.ones(feature_dim, device=x.device))
            self.bias = torch.nn.Parameter(torch.zeros(feature_dim, device=x.device))
            self._initialized = True

        return self.layernorm.layernorm_cuda(x, self.weight, self.bias, self.eps)