forked from ccf-ai-infra/GPUCodeForces
286 lines
9.1 KiB
Python
286 lines
9.1 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>
|
|
|
|
#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;
|
|
}}
|
|
|
|
__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;
|
|
|
|
const int64_t batch_offset = (int64_t)n_idx * D_vec_total * 4;
|
|
|
|
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;
|
|
|
|
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) {{
|
|
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;
|
|
|
|
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;
|
|
|
|
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;
|
|
|
|
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];
|
|
}}
|
|
|
|
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);
|
|
|
|
__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);
|
|
|
|
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);
|
|
}}
|
|
}}
|
|
}}
|
|
|
|
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='pearson_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
|
|
|
|
numerator = sum_xy - D * mean_x * mean_y
|
|
var_x = sum_xx - D * mean_x * mean_x
|
|
var_y = sum_yy - D * mean_y * mean_y
|
|
|
|
return numerator / (torch.sqrt(var_x * var_y) + self.eps) |