forked from ccf-ai-infra/GPUCodeForces
263 lines
9.6 KiB
Python
263 lines
9.6 KiB
Python
|
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 (H * W) % 4 == 0, "Instance size (H * W) must be a multiple of 4"
|
|||
|
|
|
|||
|
|
|
|||
|
|
class ModelNew(nn.Module):
|
|||
|
|
"""
|
|||
|
|
EvoNorm-S0/B0 的 CUDA 优化实现
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
def __init__(self, evonorm_gamma, evonorm_beta, evonorm_v=None, use_b0=False):
|
|||
|
|
super().__init__()
|
|||
|
|
self.gamma = nn.Parameter(evonorm_gamma.clone().view(1, C, 1, 1))
|
|||
|
|
self.beta = nn.Parameter(evonorm_beta.clone().view(1, C, 1, 1))
|
|||
|
|
self.eps = EPS
|
|||
|
|
self.nonlinear = (evonorm_v is not None)
|
|||
|
|
self.use_b0 = use_b0
|
|||
|
|
|
|||
|
|
if self.nonlinear:
|
|||
|
|
self.v = nn.Parameter(evonorm_v.clone().view(1, C, 1, 1))
|
|||
|
|
else:
|
|||
|
|
self.register_parameter('v', None)
|
|||
|
|
|
|||
|
|
if self.use_b0:
|
|||
|
|
self.register_buffer('running_var', torch.ones(1, C, 1, 1))
|
|||
|
|
self.momentum = 0.1
|
|||
|
|
|
|||
|
|
self._compile_cuda_kernel()
|
|||
|
|
|
|||
|
|
def _compile_cuda_kernel(self):
|
|||
|
|
cpp_source = """
|
|||
|
|
#include <torch/extension.h>
|
|||
|
|
|
|||
|
|
torch::Tensor evonorm_forward_cuda(
|
|||
|
|
torch::Tensor input,
|
|||
|
|
torch::Tensor mean,
|
|||
|
|
torch::Tensor var,
|
|||
|
|
torch::Tensor gamma,
|
|||
|
|
torch::Tensor beta,
|
|||
|
|
torch::Tensor v,
|
|||
|
|
float eps,
|
|||
|
|
bool nonlinear,
|
|||
|
|
int N, int C, int H, int W);
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
cuda_source = """
|
|||
|
|
#include <cuda_runtime.h>
|
|||
|
|
#include <device_launch_parameters.h>
|
|||
|
|
#include <cmath>
|
|||
|
|
|
|||
|
|
// 优化: 使用快速数学函数
|
|||
|
|
#define FAST_DIV(a, b) __fdividef(a, b)
|
|||
|
|
#define FAST_EXP(x) __expf(x)
|
|||
|
|
|
|||
|
|
// 优化 1: Sigmoid 快速计算（使用查找表或优化公式）
|
|||
|
|
__device__ __forceinline__ float fast_sigmoid(float x) {
|
|||
|
|
// 使用快速除法和指数
|
|||
|
|
return FAST_DIV(1.0f, 1.0f + FAST_EXP(-x));
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// 优化 2: 向量化 sigmoid 计算
|
|||
|
|
__device__ __forceinline__ float4 sigmoid_vec(float4 x, float v_val) {
|
|||
|
|
float4 result;
|
|||
|
|
result.x = fast_sigmoid(x.x * v_val);
|
|||
|
|
result.y = fast_sigmoid(x.y * v_val);
|
|||
|
|
result.z = fast_sigmoid(x.z * v_val);
|
|||
|
|
result.w = fast_sigmoid(x.w * v_val);
|
|||
|
|
return result;
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
__global__ void evonorm_apply_kernel(
|
|||
|
|
const float* __restrict__ x,
|
|||
|
|
const float* __restrict__ mean,
|
|||
|
|
const float* __restrict__ var,
|
|||
|
|
const float* __restrict__ gamma,
|
|||
|
|
const float* __restrict__ beta,
|
|||
|
|
const float* __restrict__ v,
|
|||
|
|
float* __restrict__ y,
|
|||
|
|
float eps,
|
|||
|
|
bool nonlinear,
|
|||
|
|
int N, int C, int H, int W
|
|||
|
|
) {
|
|||
|
|
const int nc_idx = blockIdx.x;
|
|||
|
|
if (nc_idx >= N * C) return;
|
|||
|
|
|
|||
|
|
const int n_idx = nc_idx / C;
|
|||
|
|
const int c_idx = nc_idx % C;
|
|||
|
|
|
|||
|
|
// 优化 3: 使用 __ldg() 读取只读全局内存
|
|||
|
|
const float m = __ldg(&mean[nc_idx]);
|
|||
|
|
const float variance = __ldg(&var[nc_idx]);
|
|||
|
|
|
|||
|
|
// 优化 4: 预计算常量
|
|||
|
|
const float inv_std = rsqrtf(variance + eps); // rsqrtf 比 1.0f/sqrtf 快
|
|||
|
|
|
|||
|
|
const float g = __ldg(&gamma[c_idx]);
|
|||
|
|
const float b = __ldg(&beta[c_idx]);
|
|||
|
|
const float v_val = nonlinear ? __ldg(&v[c_idx]) : 0.0f;
|
|||
|
|
|
|||
|
|
const int instance_size = H * W;
|
|||
|
|
const int instance_offset = n_idx * C * instance_size + c_idx * instance_size;
|
|||
|
|
const float* x_ptr = x + instance_offset;
|
|||
|
|
float* y_ptr = y + instance_offset;
|
|||
|
|
|
|||
|
|
const int instance_size_div4 = instance_size / 4;
|
|||
|
|
const float4* x4_ptr = reinterpret_cast<const float4*>(x_ptr);
|
|||
|
|
float4* y4_ptr = reinterpret_cast<float4*>(y_ptr);
|
|||
|
|
|
|||
|
|
const int BLOCK_SIZE = 256;
|
|||
|
|
|
|||
|
|
// 优化 5: 循环展开（处理 2 个 float4 每次迭代）
|
|||
|
|
const int items_per_thread = (instance_size_div4 + BLOCK_SIZE - 1) / BLOCK_SIZE;
|
|||
|
|
const int base_idx = threadIdx.x;
|
|||
|
|
|
|||
|
|
#pragma unroll 2
|
|||
|
|
for (int i = 0; i < items_per_thread; ++i) {
|
|||
|
|
int idx = base_idx + i * BLOCK_SIZE;
|
|||
|
|
if (idx < instance_size_div4) {
|
|||
|
|
// 优化 6: 使用 __ldg() 读取输入（如果对齐）
|
|||
|
|
float4 x_val = x4_ptr[idx];
|
|||
|
|
float4 y_val;
|
|||
|
|
|
|||
|
|
// 归一化: (x - m) / std
|
|||
|
|
// 注意: 不使用 volatile，因为统计量已在 Python 端计算
|
|||
|
|
float x_norm_x = (x_val.x - m) * inv_std;
|
|||
|
|
float x_norm_y = (x_val.y - m) * inv_std;
|
|||
|
|
float x_norm_z = (x_val.z - m) * inv_std;
|
|||
|
|
float x_norm_w = (x_val.w - m) * inv_std;
|
|||
|
|
|
|||
|
|
// 仿射变换: x_norm * g + b (使用 FMA)
|
|||
|
|
float y_affine_x = fmaf(x_norm_x, g, b);
|
|||
|
|
float y_affine_y = fmaf(x_norm_y, g, b);
|
|||
|
|
float y_affine_z = fmaf(x_norm_z, g, b);
|
|||
|
|
float y_affine_w = fmaf(x_norm_w, g, b);
|
|||
|
|
|
|||
|
|
// 非线性门控
|
|||
|
|
if (nonlinear) {
|
|||
|
|
// 优化 7: 向量化 sigmoid 计算
|
|||
|
|
float sigmoid_x = fast_sigmoid(x_val.x * v_val);
|
|||
|
|
float sigmoid_y = fast_sigmoid(x_val.y * v_val);
|
|||
|
|
float sigmoid_z = fast_sigmoid(x_val.z * v_val);
|
|||
|
|
float sigmoid_w = fast_sigmoid(x_val.w * v_val);
|
|||
|
|
|
|||
|
|
y_val.x = y_affine_x * sigmoid_x;
|
|||
|
|
y_val.y = y_affine_y * sigmoid_y;
|
|||
|
|
y_val.z = y_affine_z * sigmoid_z;
|
|||
|
|
y_val.w = y_affine_w * sigmoid_w;
|
|||
|
|
} else {
|
|||
|
|
y_val.x = y_affine_x;
|
|||
|
|
y_val.y = y_affine_y;
|
|||
|
|
y_val.z = y_affine_z;
|
|||
|
|
y_val.w = y_affine_w;
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
y4_ptr[idx] = y_val;
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// ============================================================
|
|||
|
|
// C++ Wrapper
|
|||
|
|
// ============================================================
|
|||
|
|
torch::Tensor evonorm_forward_cuda(
|
|||
|
|
torch::Tensor input,
|
|||
|
|
torch::Tensor mean,
|
|||
|
|
torch::Tensor var,
|
|||
|
|
torch::Tensor gamma,
|
|||
|
|
torch::Tensor beta,
|
|||
|
|
torch::Tensor v,
|
|||
|
|
float eps,
|
|||
|
|
bool nonlinear,
|
|||
|
|
int N, int C, int H, int W
|
|||
|
|
) {
|
|||
|
|
input = input.contiguous();
|
|||
|
|
auto output = torch::empty_like(input);
|
|||
|
|
|
|||
|
|
const int BLOCK_SIZE = 256;
|
|||
|
|
dim3 blocks(N * C);
|
|||
|
|
dim3 threads(BLOCK_SIZE);
|
|||
|
|
|
|||
|
|
// 优化 8: 使用 CUDA stream（可选）
|
|||
|
|
evonorm_apply_kernel<<<blocks, threads>>>(
|
|||
|
|
input.data_ptr<float>(),
|
|||
|
|
mean.data_ptr<float>(),
|
|||
|
|
var.data_ptr<float>(),
|
|||
|
|
gamma.data_ptr<float>(),
|
|||
|
|
beta.data_ptr<float>(),
|
|||
|
|
nonlinear ? v.data_ptr<float>() : nullptr,
|
|||
|
|
output.data_ptr<float>(),
|
|||
|
|
eps,
|
|||
|
|
nonlinear,
|
|||
|
|
N, C, H, W
|
|||
|
|
);
|
|||
|
|
|
|||
|
|
return output;
|
|||
|
|
}
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
# 优化 9: 使用更激进的编译选项
|
|||
|
|
self.evonorm_op = load_inline(
|
|||
|
|
name="evonorm_cuda_optimized_v4",
|
|||
|
|
cpp_sources=cpp_source,
|
|||
|
|
cuda_sources=cuda_source,
|
|||
|
|
functions=["evonorm_forward_cuda"],
|
|||
|
|
extra_cuda_cflags=[
|
|||
|
|
"-O3",
|
|||
|
|
"--use_fast_math", # 启用快速数学（可能略微降低精度但提升性能）
|
|||
|
|
"-lineinfo" # 便于性能分析
|
|||
|
|
],
|
|||
|
|
verbose=False
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
|||
|
|
if x.dtype != torch.float32 or not x.is_cuda:
|
|||
|
|
x = x.to("cuda", dtype=torch.float32)
|
|||
|
|
|
|||
|
|
N, C, H, W = x.size()
|
|||
|
|
|
|||
|
|
if self.use_b0:
|
|||
|
|
# EvoNorm-B0
|
|||
|
|
if self.training:
|
|||
|
|
mean = x.mean(dim=[2, 3], keepdim=True)
|
|||
|
|
var = x.var(dim=[2, 3], keepdim=True, unbiased=False)
|
|||
|
|
|
|||
|
|
with torch.no_grad():
|
|||
|
|
batch_var = var.mean(dim=0, keepdim=True)
|
|||
|
|
self.running_var = (1 - self.momentum) * self.running_var + self.momentum * batch_var
|
|||
|
|
else:
|
|||
|
|
mean = x.mean(dim=[2, 3], keepdim=True)
|
|||
|
|
var = self.running_var.expand(N, C, 1, 1)
|
|||
|
|
else:
|
|||
|
|
# EvoNorm-S0
|
|||
|
|
# 优化 10: 融合计算 E[x^2] 和 E[x] 可以考虑自定义 CUDA kernel
|
|||
|
|
x_sq_mean = torch.mean(x * x, dim=[2, 3], keepdim=True)
|
|||
|
|
x_mean = torch.mean(x, dim=[2, 3], keepdim=True)
|
|||
|
|
var = x_sq_mean - x_mean * x_mean
|
|||
|
|
mean = torch.zeros_like(x_mean)
|
|||
|
|
|
|||
|
|
gamma_view = self.gamma.data.view(C).contiguous()
|
|||
|
|
beta_view = self.beta.data.view(C).contiguous()
|
|||
|
|
|
|||
|
|
if self.nonlinear:
|
|||
|
|
v_view = self.v.data.view(C).contiguous()
|
|||
|
|
else:
|
|||
|
|
v_view = torch.zeros(C, device=x.device, dtype=torch.float32)
|
|||
|
|
|
|||
|
|
return self.evonorm_op.evonorm_forward_cuda(
|
|||
|
|
x.contiguous(),
|
|||
|
|
mean.contiguous().view(N, C),
|
|||
|
|
var.contiguous().view(N, C),
|
|||
|
|
gamma_view,
|
|||
|
|
beta_view,
|
|||
|
|
v_view,
|
|||
|
|
self.eps,
|
|||
|
|
self.nonlinear,
|
|||
|
|
N, C, H, W
|
|||
|
|
)
|