forked from ccf-ai-infra/GPUCodeForces
Add conv2d operator
This commit is contained in:
parent
2f77913144
commit
182cf2797b
|
|
@ -0,0 +1,223 @@
|
|||
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)
|
||||
|
|
@ -0,0 +1,40 @@
|
|||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
class Model(nn.Module):
|
||||
"""
|
||||
Model that performs 2D convolution operation.
|
||||
"""
|
||||
def __init__(self, weight, bias=None):
|
||||
super(Model, self).__init__()
|
||||
self.weight = nn.Parameter(weight)
|
||||
self.bias = nn.Parameter(bias) if bias is not None else None
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Performs 2D convolution.
|
||||
|
||||
Args:
|
||||
x (torch.Tensor): Input tensor of shape [batch_size, in_channels, height, width]
|
||||
|
||||
Returns:
|
||||
torch.Tensor: Output tensor of shape [batch_size, out_channels, out_height, out_width]
|
||||
"""
|
||||
return torch.nn.functional.conv2d(x, self.weight, self.bias, stride=1, padding=0)
|
||||
|
||||
# Hyperparameters
|
||||
batch_size = 4
|
||||
in_channels = 3
|
||||
out_channels = 64
|
||||
height = 32
|
||||
width = 32
|
||||
kernel_size = 3
|
||||
|
||||
def get_inputs():
|
||||
x = torch.randn(batch_size, in_channels, height, width)
|
||||
return [x]
|
||||
|
||||
def get_init_inputs():
|
||||
weight = torch.randn(out_channels, in_channels, kernel_size, kernel_size)
|
||||
bias = torch.randn(out_channels)
|
||||
return [weight, bias]
|
||||
|
|
@ -0,0 +1,30 @@
|
|||
You write custom CUDA kernels to replace the pytorch operators in the given architecture to get speedups.
|
||||
|
||||
You have complete freedom to choose the set of operators you want to replace. You may make the decision to replace some operators with custom CUDA kernels and leave others unchanged. You may replace multiple operators with custom implementations, consider operator fusion opportunities (combining multiple operators into a single kernel, for example, combining matmul+relu), or algorithmic changes (such as online softmax). You are only limited by your imagination.
|
||||
|
||||
Here's an example to show you the syntax of inline embedding custom CUDA operators in torch: The example given architecture is:
|
||||
|
||||
```python
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
class Model(nn.Module):
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
|
||||
def forward(self, a, b):
|
||||
return a + b
|
||||
|
||||
|
||||
def get_inputs():
|
||||
# randomly generate input tensors based on the model architecture
|
||||
a = torch.randn(1, 128).cuda()
|
||||
b = torch.randn(1, 128).cuda()
|
||||
return [a, b]
|
||||
|
||||
|
||||
def get_init_inputs():
|
||||
# randomly generate tensors required for initialization based on the model architecture
|
||||
return []
|
||||
|
|
@ -0,0 +1,92 @@
|
|||
###########################################################
|
||||
# 性能和精度验证程序
|
||||
###########################################################
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import time
|
||||
from conv2d_torch import Model, get_inputs, get_init_inputs
|
||||
from conv2d_cuda import ModelNew
|
||||
|
||||
# 禁用 TF32，确保与自定义 FP32 核精度对齐
|
||||
torch.backends.cuda.matmul.allow_tf32 = False
|
||||
torch.backends.cudnn.allow_tf32 = False
|
||||
torch.backends.cudnn.deterministic = True
|
||||
|
||||
def _time_cuda_model(fn, inputs, iters=300, warmup=50):
|
||||
torch.cuda.synchronize()
|
||||
for _ in range(warmup):
|
||||
_ = fn(*inputs)
|
||||
torch.cuda.synchronize()
|
||||
start = torch.cuda.Event(enable_timing=True)
|
||||
end = torch.cuda.Event(enable_timing=True)
|
||||
start.record()
|
||||
for _ in range(iters):
|
||||
_ = fn(*inputs)
|
||||
end.record()
|
||||
torch.cuda.synchronize()
|
||||
ms = start.elapsed_time(end) / iters # 平均每次毫秒
|
||||
return ms / 1000.0 # 转为秒
|
||||
|
||||
def run_benchmark():
|
||||
# 检查 CUDA 是否可用
|
||||
if not torch.cuda.is_available():
|
||||
print("CUDA 不可用，请确保您有可用的 NVIDIA GPU 并已正确安装 PyTorch CUDA 版本。")
|
||||
return
|
||||
else:
|
||||
device = torch.device("cuda")
|
||||
|
||||
torch.backends.cudnn.benchmark = True
|
||||
|
||||
# 初始化模型
|
||||
init_inputs = get_init_inputs()
|
||||
init_inputs = [
|
||||
x.cuda(device=device) if isinstance(x, torch.Tensor) else x for x in init_inputs
|
||||
]
|
||||
inputs = get_inputs()
|
||||
inputs = [
|
||||
x.cuda(device=device) if isinstance(x, torch.Tensor) else x for x in inputs
|
||||
]
|
||||
|
||||
torch_model = Model(*init_inputs).cuda()
|
||||
cuda_model = ModelNew(*init_inputs).cuda()
|
||||
|
||||
torch_model.eval()
|
||||
cuda_model.eval()
|
||||
|
||||
print("-------------------- 精度对齐验证 --------------------")
|
||||
with torch.no_grad():
|
||||
output_torch = torch_model(*inputs)
|
||||
output_cuda = cuda_model(*inputs)
|
||||
|
||||
abs_diff = (output_torch - output_cuda).abs()
|
||||
max_diff = abs_diff.max().item()
|
||||
mean_diff = abs_diff.mean().item()
|
||||
|
||||
print(f"最大差异: {max_diff:.6f}")
|
||||
print(f"平均差异: {mean_diff:.6f}")
|
||||
|
||||
precision_flag = torch.allclose(output_torch, output_cuda, rtol=1e-5, atol=1e-5)
|
||||
if precision_flag:
|
||||
print("✅ 精度对齐：两个模型的输出结果非常接近。")
|
||||
else:
|
||||
print("❌ 精度不一致！")
|
||||
|
||||
print("\n-------------------- 性能加速比测试 --------------------")
|
||||
num_iterations = 300
|
||||
|
||||
# 计时
|
||||
torch_time = _time_cuda_model(torch_model, inputs, iters=num_iterations, warmup=50)
|
||||
cuda_time = _time_cuda_model(cuda_model, inputs, iters=num_iterations, warmup=50)
|
||||
|
||||
print(f"PyTorch (conv2d) 平均执行时间: {torch_time:.6f} 秒")
|
||||
print(f"自定义 CUDA conv2d 平均执行时间: {cuda_time:.6f} 秒")
|
||||
speedup = 0.0
|
||||
if cuda_time > 0:
|
||||
speedup = torch_time / cuda_time
|
||||
print(f"相对 PyTorch 加速比: {speedup:.2f}x")
|
||||
else:
|
||||
print("CUDA 内核执行时间为0，无法计算加速比。")
|
||||
return precision_flag, speedup
|
||||
|
||||
if __name__ == "__main__":
|
||||
precision_flag, speedup = run_benchmark()
|
||||
Loading…
Reference in New Issue