forked from ccf-ai-infra/GPUCodeForces
308 lines
9.9 KiB
Python
308 lines
9.9 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-8
|
||
assert (C * H * W) % 4 == 0
|
||
|
||
|
||
class ModelNew(nn.Module):
|
||
def __init__(self):
|
||
super().__init__()
|
||
self.block_size = 512
|
||
self.eps = EPS
|
||
self.register_buffer('temp_buffer', torch.zeros((N, 5), dtype=torch.float32))
|
||
self._compile_cuda_kernel()
|
||
|
||
def _compile_cuda_kernel(self):
|
||
cpp_source = """
|
||
#include <torch/extension.h>
|
||
#include <cuda_runtime.h>
|
||
|
||
void pearson_sum_cuda(
|
||
torch::Tensor x,
|
||
torch::Tensor y,
|
||
torch::Tensor temp_buffer,
|
||
int N,
|
||
int D);
|
||
"""
|
||
|
||
cuda_source = f"""
|
||
#include <cuda_runtime.h>
|
||
#include <cmath>
|
||
#include <torch/types.h>
|
||
|
||
#define BLOCK_SIZE {self.block_size}
|
||
#define WARP_SIZE 32
|
||
#define ILP 4
|
||
|
||
// 线程内归约
|
||
__inline__ __device__ float warp_reduce_sum(float val) {{
|
||
#pragma unroll
|
||
for (int offset = WARP_SIZE / 2; offset > 0; offset /= 2) {{
|
||
val += __shfl_down_sync(0xffffffff, val, offset);
|
||
}}
|
||
return val;
|
||
}}
|
||
|
||
// 分块计算5个统计量的Kernel函数
|
||
__global__ __launch_bounds__(BLOCK_SIZE)
|
||
void pearson_split_kernel(
|
||
const float* __restrict__ x,
|
||
const float* __restrict__ y,
|
||
float* __restrict__ temp_buffer,
|
||
int D_vec_total
|
||
) {{
|
||
const int n_idx = blockIdx.y;
|
||
const int split_idx = blockIdx.x;
|
||
const int num_splits = gridDim.x;
|
||
|
||
const int chunk_size = (D_vec_total + num_splits - 1) / num_splits;
|
||
const int start_idx = split_idx * chunk_size;
|
||
const int end_idx = min(start_idx + chunk_size, D_vec_total);
|
||
|
||
if (start_idx >= D_vec_total) return;
|
||
|
||
// D_vec_total 是 D/4，所以批量偏移量是 n_idx * D_vec_total * 4
|
||
const int64_t batch_offset = (int64_t)n_idx * D_vec_total * 4;
|
||
|
||
// float4 指针：按4个float（16字节）的块访问数据
|
||
const float4* curr_x = reinterpret_cast<const float4*>(x + batch_offset) + start_idx + threadIdx.x;
|
||
const float4* curr_y = reinterpret_cast<const float4*>(y + batch_offset) + start_idx + threadIdx.x;
|
||
const float4* end_ptr = reinterpret_cast<const float4*>(x + batch_offset) + end_idx;
|
||
|
||
// 线程局部累加器（ILP=4）
|
||
float sum_x[ILP];
|
||
float sum_y[ILP];
|
||
float sum_xx[ILP];
|
||
float sum_yy[ILP];
|
||
float sum_xy[ILP];
|
||
|
||
#pragma unroll
|
||
for (int k=0; k<ILP; ++k) {{
|
||
sum_x[k] = 0.0f;
|
||
sum_y[k] = 0.0f;
|
||
sum_xx[k] = 0.0f;
|
||
sum_yy[k] = 0.0f;
|
||
sum_xy[k] = 0.0f;
|
||
}}
|
||
|
||
const int stride = BLOCK_SIZE * ILP;
|
||
|
||
// 循环展开和向量化加载主循环
|
||
while (curr_x + (ILP - 1) * BLOCK_SIZE < end_ptr) {{
|
||
float4 r_x[ILP];
|
||
float4 r_y[ILP];
|
||
|
||
#pragma unroll
|
||
for (int k = 0; k < ILP; ++k) {{
|
||
r_x[k] = __ldg(curr_x + k * BLOCK_SIZE);
|
||
r_y[k] = __ldg(curr_y + k * BLOCK_SIZE);
|
||
}}
|
||
|
||
#pragma unroll
|
||
for (int k = 0; k < ILP; ++k) {{
|
||
// x
|
||
float vx = r_x[k].x; float vy = r_y[k].x;
|
||
sum_x[k] += vx; sum_y[k] += vy;
|
||
sum_xx[k] += vx * vx; sum_yy[k] += vy * vy; sum_xy[k] += vx * vy;
|
||
|
||
// y
|
||
vx = r_x[k].y; vy = r_y[k].y;
|
||
sum_x[k] += vx; sum_y[k] += vy;
|
||
sum_xx[k] += vx * vx; sum_yy[k] += vy * vy; sum_xy[k] += vx * vy;
|
||
|
||
// z
|
||
vx = r_x[k].z; vy = r_y[k].z;
|
||
sum_x[k] += vx; sum_y[k] += vy;
|
||
sum_xx[k] += vx * vx; sum_yy[k] += vy * vy; sum_xy[k] += vx * vy;
|
||
|
||
// w
|
||
vx = r_x[k].w; vy = r_y[k].w;
|
||
sum_x[k] += vx; sum_y[k] += vy;
|
||
sum_xx[k] += vx * vx; sum_yy[k] += vy * vy; sum_xy[k] += vx * vy;
|
||
}}
|
||
|
||
curr_x += stride;
|
||
curr_y += stride;
|
||
}}
|
||
|
||
// 处理剩余部分
|
||
while (curr_x < end_ptr) {{
|
||
float4 vx = __ldg(curr_x);
|
||
float4 vy = __ldg(curr_y);
|
||
|
||
float v1 = vx.x; float v2 = vy.x;
|
||
sum_x[0] += v1; sum_y[0] += v2;
|
||
sum_xx[0] += v1*v1; sum_yy[0] += v2*v2; sum_xy[0] += v1*v2;
|
||
|
||
v1 = vx.y; v2 = vy.y;
|
||
sum_x[0] += v1; sum_y[0] += v2;
|
||
sum_xx[0] += v1*v1; sum_yy[0] += v2*v2; sum_xy[0] += v1*v2;
|
||
|
||
v1 = vx.z; v2 = vy.z;
|
||
sum_x[0] += v1; sum_y[0] += v2;
|
||
sum_xx[0] += v1*v1; sum_yy[0] += v2*v2; sum_xy[0] += v1*v2;
|
||
|
||
v1 = vx.w; v2 = vy.w;
|
||
sum_x[0] += v1; sum_y[0] += v2;
|
||
sum_xx[0] += v1*v1; sum_yy[0] += v2*v2; sum_xy[0] += v1*v2;
|
||
|
||
curr_x += BLOCK_SIZE;
|
||
curr_y += BLOCK_SIZE;
|
||
}}
|
||
|
||
// 局部求和
|
||
float l_x = 0.0f, l_y = 0.0f, l_xx = 0.0f, l_yy = 0.0f, l_xy = 0.0f;
|
||
#pragma unroll
|
||
for (int k=0; k<ILP; ++k) {{
|
||
l_x += sum_x[k];
|
||
l_y += sum_y[k];
|
||
l_xx += sum_xx[k];
|
||
l_yy += sum_yy[k];
|
||
l_xy += sum_xy[k];
|
||
}}
|
||
|
||
// Warp 归约
|
||
l_x = warp_reduce_sum(l_x);
|
||
l_y = warp_reduce_sum(l_y);
|
||
l_xx = warp_reduce_sum(l_xx);
|
||
l_yy = warp_reduce_sum(l_yy);
|
||
l_xy = warp_reduce_sum(l_xy);
|
||
|
||
// Block 归约
|
||
__shared__ float s_x[BLOCK_SIZE / WARP_SIZE];
|
||
__shared__ float s_y[BLOCK_SIZE / WARP_SIZE];
|
||
__shared__ float s_xx[BLOCK_SIZE / WARP_SIZE];
|
||
__shared__ float s_yy[BLOCK_SIZE / WARP_SIZE];
|
||
__shared__ float s_xy[BLOCK_SIZE / WARP_SIZE];
|
||
|
||
const int lane_id = threadIdx.x % WARP_SIZE;
|
||
const int warp_id = threadIdx.x / WARP_SIZE;
|
||
|
||
if (lane_id == 0) {{
|
||
s_x[warp_id] = l_x;
|
||
s_y[warp_id] = l_y;
|
||
s_xx[warp_id] = l_xx;
|
||
s_yy[warp_id] = l_yy;
|
||
s_xy[warp_id] = l_xy;
|
||
}}
|
||
__syncthreads();
|
||
|
||
if (warp_id == 0) {{
|
||
float b_x = 0.0f, b_y = 0.0f, b_xx = 0.0f, b_yy = 0.0f, b_xy = 0.0f;
|
||
if (lane_id < (BLOCK_SIZE / WARP_SIZE)) {{
|
||
b_x = s_x[lane_id];
|
||
b_y = s_y[lane_id];
|
||
b_xx = s_xx[lane_id];
|
||
b_yy = s_yy[lane_id];
|
||
b_xy = s_xy[lane_id];
|
||
}}
|
||
|
||
b_x = warp_reduce_sum(b_x);
|
||
b_y = warp_reduce_sum(b_y);
|
||
b_xx = warp_reduce_sum(b_xx);
|
||
b_yy = warp_reduce_sum(b_yy);
|
||
b_xy = warp_reduce_sum(b_xy);
|
||
|
||
// 将结果原子加到 temp_buffer 中
|
||
if (lane_id == 0) {{
|
||
float* dst = temp_buffer + n_idx * 5;
|
||
atomicAdd(&dst[0], b_x);
|
||
atomicAdd(&dst[1], b_y);
|
||
atomicAdd(&dst[2], b_xx);
|
||
atomicAdd(&dst[3], b_yy);
|
||
atomicAdd(&dst[4], b_xy);
|
||
}}
|
||
}}
|
||
}}
|
||
|
||
// Host Wrapper for Kernel Launch (保持不变)
|
||
void pearson_sum_cuda(
|
||
torch::Tensor x,
|
||
torch::Tensor y,
|
||
torch::Tensor temp_buffer,
|
||
int N,
|
||
int D)
|
||
{{
|
||
int D_vec = D / 4;
|
||
|
||
int device_id;
|
||
cudaGetDevice(&device_id);
|
||
int sm_count;
|
||
cudaDeviceGetAttribute(&sm_count, cudaDevAttrMultiProcessorCount, device_id);
|
||
|
||
int target_blocks = sm_count * 4;
|
||
int splits = (target_blocks + N - 1) / N;
|
||
|
||
int max_splits = (D_vec + 1024 - 1) / 1024;
|
||
if (splits > max_splits) splits = max_splits;
|
||
if (splits < 1) splits = 1;
|
||
if (splits > 512) splits = 512;
|
||
|
||
dim3 blocks(splits, N);
|
||
dim3 threads(BLOCK_SIZE);
|
||
|
||
pearson_split_kernel<<<blocks, threads>>>(
|
||
x.data_ptr<float>(),
|
||
y.data_ptr<float>(),
|
||
temp_buffer.data_ptr<float>(),
|
||
D_vec
|
||
);
|
||
}}
|
||
"""
|
||
|
||
self.op = load_inline(
|
||
name='icc_cuda_opt_v1',
|
||
cpp_sources=cpp_source,
|
||
cuda_sources=cuda_source,
|
||
functions=['pearson_sum_cuda'],
|
||
extra_cuda_cflags=[
|
||
'-O3',
|
||
'--use_fast_math',
|
||
'-Xptxas=-v'
|
||
],
|
||
verbose=False
|
||
)
|
||
|
||
def forward(self, x, y):
|
||
|
||
if not x.is_contiguous(): x = x.contiguous()
|
||
if not y.is_contiguous(): y = y.contiguous()
|
||
|
||
if not x.is_cuda: x = x.cuda()
|
||
if not y.is_cuda: y = y.cuda()
|
||
|
||
N, C, H, W = x.size()
|
||
D = C * H * W
|
||
|
||
self.temp_buffer.zero_()
|
||
|
||
self.op.pearson_sum_cuda(
|
||
x,
|
||
y,
|
||
self.temp_buffer,
|
||
N,
|
||
D
|
||
)
|
||
|
||
sum_x = self.temp_buffer[:, 0]
|
||
sum_y = self.temp_buffer[:, 1]
|
||
sum_xx = self.temp_buffer[:, 2]
|
||
sum_yy = self.temp_buffer[:, 3]
|
||
sum_xy = self.temp_buffer[:, 4]
|
||
|
||
mean_x = sum_x / D
|
||
mean_y = sum_y / D
|
||
|
||
cov_sum = sum_xy - D * mean_x * mean_y
|
||
|
||
x_var_sum = sum_xx - D * mean_x * mean_x
|
||
|
||
y_var_sum = sum_yy - D * mean_y * mean_y
|
||
|
||
numerator = 2 * cov_sum
|
||
denominator = x_var_sum + y_var_sum
|
||
|
||
return numerator / (denominator + self.eps) |