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
|
||
) |