forked from ccf-ai-infra/GPUCodeForces
223 lines
7.3 KiB
Python
223 lines
7.3 KiB
Python
import torch
|
||
import torch.nn as nn
|
||
from torch.utils.cpp_extension import load_inline
|
||
|
||
# CUDA implementation of Conv2D (tiled + shared memory + output-channel blocking)
|
||
conv2d_source = r"""
|
||
#include <torch/extension.h>
|
||
#include <cuda_runtime.h>
|
||
#include <math.h>
|
||
|
||
#ifndef CHECK_CUDA
|
||
#define CHECK_CUDA(x) TORCH_CHECK((x).is_cuda(), #x " must be a CUDA tensor")
|
||
#endif
|
||
#ifndef CHECK_CONTIGUOUS
|
||
#define CHECK_CONTIGUOUS(x) TORCH_CHECK((x).is_contiguous(), #x " must be contiguous")
|
||
#endif
|
||
#ifndef CHECK_FLOAT
|
||
#define CHECK_FLOAT(x) TORCH_CHECK((x).scalar_type() == at::kFloat, #x " must be float32")
|
||
#endif
|
||
|
||
// 每个block计算一个 (b, oc_group) 上的输出tile，复用输入tile计算 OC_TILE 个输出通道
|
||
template<int BLOCK_X, int BLOCK_Y, int OC_TILE>
|
||
__global__ void conv2d_tiled_kernel_oc(
|
||
const float* __restrict__ input, // [B, C_in, H, W]
|
||
const float* __restrict__ weight, // [C_out, C_in, K, K]
|
||
const float* __restrict__ bias, // [C_out] or nullptr
|
||
float* __restrict__ output, // [B, C_out, H_out, W_out]
|
||
int B, int C_in, int C_out,
|
||
int H, int W, int K, int H_out, int W_out,
|
||
bool has_bias
|
||
) {
|
||
// grid.z = B * ceil_div(C_out, OC_TILE)
|
||
int groups = (C_out + OC_TILE - 1) / OC_TILE;
|
||
int b = blockIdx.z / groups;
|
||
int og = blockIdx.z % groups; // 输出通道组编号
|
||
int co0 = og * OC_TILE; // 本组起始输出通道
|
||
|
||
int ow0 = blockIdx.x * BLOCK_X;
|
||
int oh0 = blockIdx.y * BLOCK_Y;
|
||
|
||
int ow = ow0 + threadIdx.x;
|
||
int oh = oh0 + threadIdx.y;
|
||
|
||
extern __shared__ float smem[];
|
||
// 输入tile大小：(BLOCK_Y+K-1) x (BLOCK_X+K-1)
|
||
int tile_w = BLOCK_X + K - 1;
|
||
int tile_h = BLOCK_Y + K - 1;
|
||
float* tile = smem; // tile_h * tile_w
|
||
float* w_sh = tile + tile_h * tile_w; // OC_TILE * K * K
|
||
// w_sh 布局: [oc_local][K*K]
|
||
|
||
// 累加器：每线程维护 OC_TILE 个通道
|
||
float acc[OC_TILE];
|
||
#pragma unroll
|
||
for (int oc = 0; oc < OC_TILE; ++oc) {
|
||
int co = co0 + oc;
|
||
acc[oc] = (has_bias && co < C_out) ? bias[co] : 0.0f;
|
||
}
|
||
|
||
bool valid_xy = (oh < H_out) && (ow < W_out);
|
||
|
||
// 遍历输入通道
|
||
for (int ci = 0; ci < C_in; ++ci) {
|
||
// 1) 加载本组 OC_TILE 的权重到共享内存
|
||
int total_w = OC_TILE * K * K;
|
||
for (int t = threadIdx.y * BLOCK_X + threadIdx.x; t < total_w; t += BLOCK_X * BLOCK_Y) {
|
||
int oc = t / (K*K);
|
||
int rem = t % (K*K);
|
||
int kh = rem / K;
|
||
int kw = rem % K;
|
||
int co = co0 + oc;
|
||
float wv = 0.0f;
|
||
if (co < C_out) {
|
||
int w_idx = ((co * C_in + ci) * K + kh) * K + kw;
|
||
wv = weight[w_idx];
|
||
}
|
||
w_sh[t] = wv;
|
||
}
|
||
|
||
// 2) 加载输入tile到共享内存（该tile将被 OC_TILE 个输出通道复用）
|
||
int ih0 = oh0;
|
||
int iw0 = ow0;
|
||
for (int th = threadIdx.y; th < tile_h; th += BLOCK_Y) {
|
||
int ih = ih0 + th;
|
||
bool in_h = (ih >= 0) && (ih < H);
|
||
for (int tw = threadIdx.x; tw < tile_w; tw += BLOCK_X) {
|
||
int iw = iw0 + tw;
|
||
bool in_w = (iw >= 0) && (iw < W);
|
||
float v = 0.0f;
|
||
if (in_h && in_w) {
|
||
int in_idx = (((b * C_in + ci) * H + ih) * W + iw);
|
||
v = input[in_idx];
|
||
}
|
||
tile[th * tile_w + tw] = v;
|
||
}
|
||
}
|
||
|
||
__syncthreads();
|
||
|
||
// 3) 计算：同一输入tile，对 OC_TILE 个输出通道分别累加
|
||
if (valid_xy) {
|
||
int t_base = threadIdx.y * tile_w + threadIdx.x;
|
||
#pragma unroll
|
||
for (int kh = 0; kh < K; ++kh) {
|
||
int t_row = t_base + kh * tile_w;
|
||
#pragma unroll
|
||
for (int kw = 0; kw < K; ++kw) {
|
||
float val = tile[t_row + kw];
|
||
#pragma unroll
|
||
for (int oc = 0; oc < OC_TILE; ++oc) {
|
||
float wv = w_sh[oc * (K*K) + kh * K + kw];
|
||
acc[oc] = fmaf(val, wv, acc[oc]);
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
__syncthreads(); // 保护下一个 ci 的加载
|
||
}
|
||
|
||
// 4) 写回输出
|
||
if (valid_xy) {
|
||
int base = (b * C_out) * (H_out * W_out);
|
||
int out_offset = oh * W_out + ow;
|
||
#pragma unroll
|
||
for (int oc = 0; oc < OC_TILE; ++oc) {
|
||
int co = co0 + oc;
|
||
if (co < C_out) {
|
||
int out_idx = base + co * (H_out * W_out) + out_offset;
|
||
output[out_idx] = acc[oc];
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
// C++ wrapper
|
||
torch::Tensor conv2d_cuda(
|
||
torch::Tensor input,
|
||
torch::Tensor weight,
|
||
torch::Tensor bias
|
||
) {
|
||
CHECK_CUDA(input);
|
||
CHECK_CUDA(weight);
|
||
CHECK_CUDA(bias);
|
||
CHECK_CONTIGUOUS(input);
|
||
CHECK_CONTIGUOUS(weight);
|
||
CHECK_CONTIGUOUS(bias);
|
||
CHECK_FLOAT(input);
|
||
CHECK_FLOAT(weight);
|
||
CHECK_FLOAT(bias);
|
||
|
||
int B = input.size(0);
|
||
int C_in = input.size(1);
|
||
int H = input.size(2);
|
||
int W = input.size(3);
|
||
|
||
int C_out = weight.size(0);
|
||
int K = weight.size(2);
|
||
TORCH_CHECK(weight.size(3) == K, "Kernel must be square");
|
||
|
||
int H_out = H - K + 1;
|
||
int W_out = W - K + 1;
|
||
|
||
auto output = torch::empty({B, C_out, H_out, W_out}, input.options());
|
||
|
||
// 参数可按GPU微调：32x8, 16x16 等
|
||
const int BLOCK_X = 16;
|
||
const int BLOCK_Y = 16;
|
||
const int OC_TILE = 4;
|
||
|
||
dim3 block(BLOCK_X, BLOCK_Y, 1);
|
||
int groups = (C_out + OC_TILE - 1) / OC_TILE;
|
||
dim3 grid((W_out + BLOCK_X - 1) / BLOCK_X,
|
||
(H_out + BLOCK_Y - 1) / BLOCK_Y,
|
||
B * groups);
|
||
|
||
size_t tile_w = BLOCK_X + K - 1;
|
||
size_t tile_h = BLOCK_Y + K - 1;
|
||
size_t shmem_elems = tile_w * tile_h + OC_TILE * K * K;
|
||
size_t shmem_bytes = shmem_elems * sizeof(float);
|
||
|
||
bool has_bias = bias.numel() > 0;
|
||
|
||
conv2d_tiled_kernel_oc<BLOCK_X, BLOCK_Y, OC_TILE><<<grid, block, shmem_bytes>>>(
|
||
input.data_ptr<float>(),
|
||
weight.data_ptr<float>(),
|
||
has_bias ? bias.data_ptr<float>() : nullptr,
|
||
output.data_ptr<float>(),
|
||
B, C_in, C_out, H, W, K, H_out, W_out, has_bias
|
||
);
|
||
|
||
auto err = cudaGetLastError();
|
||
TORCH_CHECK(err == cudaSuccess, "conv2d kernel launch failed: ", cudaGetErrorString(err));
|
||
return output;
|
||
}
|
||
"""
|
||
|
||
conv2d_cpp_source = r"""
|
||
torch::Tensor conv2d_cuda(torch::Tensor input, torch::Tensor weight, torch::Tensor bias);
|
||
"""
|
||
|
||
# Compile with O3 (no fast-math to keep FP32 parity)
|
||
conv2d = load_inline(
|
||
name="conv2d_tiled_opt_oc",
|
||
cpp_sources=conv2d_cpp_source,
|
||
cuda_sources=conv2d_source,
|
||
functions=["conv2d_cuda"],
|
||
verbose=False,
|
||
extra_cuda_cflags=["-O3"]
|
||
)
|
||
|
||
class ModelNew(nn.Module):
|
||
def __init__(self, weight, bias=None):
|
||
super(ModelNew, self).__init__()
|
||
self.weight = nn.Parameter(weight)
|
||
self.bias = nn.Parameter(bias) if bias is not None else nn.Parameter(torch.empty(0, device=weight.device, dtype=weight.dtype))
|
||
self.conv2d = conv2d
|
||
|
||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||
x = x.contiguous()
|
||
w = self.weight.contiguous()
|
||
b = self.bias.contiguous() if self.bias is not None else torch.empty(0, device=x.device, dtype=x.dtype)
|
||
return self.conv2d.conv2d_cuda(x, w, b) |