GPUCodeForces/S1/6/layernorm_cuda.py

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

# LayerNorm CUDA 实现 - 增强优化版本
layernorm_source = """
#include <torch/extension.h>
#include <cuda_runtime.h>
#include <math.h>

#define WARP_SIZE 32

// Warp级归约函数
__device__ __forceinline__ float warp_reduce_sum(float val) {
    for (int offset = WARP_SIZE / 2; offset > 0; offset >>= 1) {
        val += __shfl_down_sync(0xffffffff, val, offset);
    }
    return val;
}

__global__ void layernorm_kernel_optimized(
    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__ float s_warp_sums[32];  // 支持最多1024个线程
    __shared__ float s_warp_sum_sqs[32];

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

    // 第一步：并行计算均值和方差
    float thread_sum = 0.0f;
    float thread_sum_sq = 0.0f;

    // 使用向量化加载（如果特征数是4的倍数）
    if (features % 4 == 0) {
        for (int i = tid * 4; i < features; i += blockDim.x * 4) {
            float4 vec = *reinterpret_cast<const float4*>(x_row + i);
            thread_sum += vec.x + vec.y + vec.z + vec.w;
            thread_sum_sq += vec.x * vec.x + vec.y * vec.y + vec.z * vec.z + vec.w * vec.w;
        }
    } else {
        // 标量版本
        for (int i = tid; i < features; i += blockDim.x) {
            float v = x_row[i];
            thread_sum += v;
            thread_sum_sq += v * v;
        }
    }

    // Warp级归约
    float warp_sum = warp_reduce_sum(thread_sum);
    float warp_sum_sq = warp_reduce_sum(thread_sum_sq);

    // 将warp结果写入共享内存
    if (lane_id == 0) {
        s_warp_sums[warp_id] = warp_sum;
        s_warp_sum_sqs[warp_id] = warp_sum_sq;
    }
    __syncthreads();

    // Block级归约（在第一个warp中完成）
    if (warp_id == 0) {
        float block_sum = (lane_id < num_warps) ? s_warp_sums[lane_id] : 0.0f;
        float block_sum_sq = (lane_id < num_warps) ? s_warp_sum_sqs[lane_id] : 0.0f;

        block_sum = warp_reduce_sum(block_sum);
        block_sum_sq = warp_reduce_sum(block_sum_sq);

        if (lane_id == 0) {
            float mean = block_sum / features;
            float var = (block_sum_sq / features) - (mean * mean);
            s_mean = mean;
            s_inv_std = rsqrtf(fmaxf(var, 0.0f) + 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;
    }

    // 确保线程数是warp大小的倍数
    threads = (threads + WARP_SIZE - 1) / WARP_SIZE * WARP_SIZE;
    threads = min(threads, features);

    // 计算共享内存大小
    size_t shared_mem = 2 * ((threads + WARP_SIZE - 1) / WARP_SIZE) * sizeof(float) + 2 * sizeof(float);

    layernorm_kernel_optimized<<<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);
"""

# 编译 CUDA 代码
layernorm = load_inline(
    name="layernorm",
    cpp_sources=layernorm_cpp_source,
    cuda_sources=layernorm_source,
    functions=["layernorm_cuda"],
    extra_cuda_cflags=["-O3", "--use_fast_math"],
    verbose=True
)


class ModelNew(nn.Module):
    def __init__(self, eps: float = 1e-5):
        super(ModelNew, self).__init__()
        self.eps = eps
        # 在forward中动态确定特征维度
        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 = nn.Parameter(torch.ones(feature_dim, device=x.device))
            self.bias = nn.Parameter(torch.zeros(feature_dim, device=x.device))
            self._initialized = True



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