forked from ccf-ai-infra/GPUCodeForces
optimized softmax-affine (use group) #37
This commit is contained in:
parent
10eed82956
commit
aff8b2bf19
|
|
@ -0,0 +1,110 @@
|
|||
import torch
|
||||
from torch.utils.cpp_extension import load_inline
|
||||
|
||||
source = """
|
||||
#include <torch/extension.h>
|
||||
|
||||
__global__ void group_softmax_affine_kernel(const float* x, const float* scale, const float* bias, float* y, int N, int C, int H, int W, int group) {
|
||||
int tid = threadIdx.x;
|
||||
int stride = blockDim.x;
|
||||
int total = N * H * W;
|
||||
extern __shared__ float sh[];
|
||||
for (int pos = blockIdx.x; pos < total; pos += gridDim.x) {
|
||||
int n = pos / (H * W);
|
||||
int hw = pos % (H * W);
|
||||
int h = hw / W;
|
||||
int w = hw % W;
|
||||
int G = (C + group - 1) / group;
|
||||
float* sX = sh;
|
||||
float* sZ = sh + C;
|
||||
float* sExp = sh + 2 * C;
|
||||
float* sM = sh + 3 * C;
|
||||
float* sDen = sM + G;
|
||||
for (int c = tid; c < C; c += stride) {
|
||||
long long idx = ((long long)n * C + c) * (long long)H * (long long)W + (long long)h * W + w;
|
||||
float xv = x[idx];
|
||||
sX[c] = xv;
|
||||
sZ[c] = xv * scale[c] + bias[c];
|
||||
}
|
||||
__syncthreads();
|
||||
if (tid == 0) {
|
||||
for (int gi = 0; gi < G; ++gi) {
|
||||
int s = gi * group;
|
||||
int e = s + group; if (e > C) e = C;
|
||||
float m = -INFINITY;
|
||||
for (int c = s; c < e; ++c) {
|
||||
float z = sZ[c];
|
||||
if (z > m) m = z;
|
||||
}
|
||||
sM[gi] = m;
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
for (int c = tid; c < C; c += stride) {
|
||||
int gi = c / group;
|
||||
float e = expf(sZ[c] - sM[gi]);
|
||||
sExp[c] = e;
|
||||
}
|
||||
__syncthreads();
|
||||
if (tid == 0) {
|
||||
for (int gi = 0; gi < G; ++gi) {
|
||||
int s = gi * group;
|
||||
int e = s + group; if (e > C) e = C;
|
||||
float den = 0.0f;
|
||||
for (int c = s; c < e; ++c) den += sExp[c];
|
||||
sDen[gi] = den;
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
for (int c = tid; c < C; c += stride) {
|
||||
int gi = c / group;
|
||||
long long idx = ((long long)n * C + c) * (long long)H * (long long)W + (long long)h * W + w;
|
||||
float g = sExp[c] / sDen[gi];
|
||||
y[idx] = g * sX[c];
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
}
|
||||
|
||||
torch::Tensor group_softmax_affine_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, torch::Tensor group) {
|
||||
auto xc = x.contiguous();
|
||||
auto sc = scale.contiguous();
|
||||
auto bc = bias.contiguous();
|
||||
auto y = torch::empty_like(xc);
|
||||
int N = (int)xc.size(0);
|
||||
int C = (int)xc.size(1);
|
||||
int H = (int)xc.size(2);
|
||||
int W = (int)xc.size(3);
|
||||
int g = group.item<int>();
|
||||
int block = C <= 1024 ? C : 1024;
|
||||
int grid = (N * H * W);
|
||||
if (grid > 65535) grid = 65535;
|
||||
int G = (C + g - 1) / g;
|
||||
size_t shmem = (size_t)(3 * C + G) * sizeof(float) + (size_t)block * sizeof(float);
|
||||
group_softmax_affine_kernel<<<grid, block, shmem>>>(xc.data_ptr<float>(), sc.data_ptr<float>(), bc.data_ptr<float>(), y.data_ptr<float>(), N, C, H, W, g);
|
||||
return y;
|
||||
}
|
||||
"""
|
||||
|
||||
cpp_source = """
|
||||
torch::Tensor group_softmax_affine_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, torch::Tensor group);
|
||||
"""
|
||||
|
||||
ops = load_inline(
|
||||
name="group_softmax_affine",
|
||||
cpp_sources=cpp_source,
|
||||
cuda_sources=source,
|
||||
functions=["group_softmax_affine_cuda"],
|
||||
verbose=True
|
||||
)
|
||||
|
||||
class ModelNew(torch.nn.Module):
|
||||
def __init__(self, scale: torch.Tensor, bias: torch.Tensor, group: int):
|
||||
super(ModelNew, self).__init__()
|
||||
self.ops = ops
|
||||
self.register_buffer("scale", scale)
|
||||
self.register_buffer("bias", bias)
|
||||
self.register_buffer("group", torch.tensor(int(group), dtype=torch.int32))
|
||||
|
||||
def forward(self, x: torch.Tensor):
|
||||
return self.ops.group_softmax_affine_cuda(x, self.scale, self.bias, self.group)
|
||||
|
|
@ -0,0 +1,2 @@
|
|||
You write custom CUDA kernels to replace PyTorch operators for speedups.
|
||||
Implement Group-Softmax-Affine gating on NCHW: Given x[N,C,H,W], per-channel scale[C], bias[C], and an integer group size g, compute z = x*scale[c] + bias[c], then for each (n,h,w) apply softmax within channel groups of size g: g_k = softmax(z over channels in group k), and output y = g_k * x for those channels. The CUDA kernel should process one spatial location per block, cache x and z in shared memory, compute per-group max/sub-exp/sum, and write gated outputs in one pass. Provide a PyTorch reference with nn.Parameter scale, bias, and group as Parameter. Accuracy must match within rtol=1e-3.
|
||||
|
|
@ -0,0 +1,52 @@
|
|||
import torch
|
||||
import time
|
||||
from torchcode import Model, get_inputs, get_init_inputs
|
||||
from cudacode import ModelNew
|
||||
|
||||
def run_benchmark():
|
||||
if not torch.cuda.is_available():
|
||||
print("CUDA 不可用，请确保您有可用的 NVIDIA GPU 并已正确安装 PyTorch CUDA 版本。")
|
||||
return
|
||||
device = torch.device("cuda")
|
||||
|
||||
init_inputs = [x.cuda(device=device) if isinstance(x, torch.Tensor) else x for x in get_init_inputs()]
|
||||
inputs = [x.cuda(device=device) if isinstance(x, torch.Tensor) else x for x in get_inputs()]
|
||||
|
||||
torch_model = Model(*init_inputs).cuda()
|
||||
cuda_model = ModelNew(*init_inputs).cuda()
|
||||
torch_model.eval(); cuda_model.eval()
|
||||
|
||||
print("-------------------- 精度对齐验证 --------------------")
|
||||
with torch.no_grad():
|
||||
out_torch = torch_model(*inputs)
|
||||
out_cuda = cuda_model(*inputs)
|
||||
flag = torch.allclose(out_torch, out_cuda, rtol=1e-03)
|
||||
if flag:
|
||||
print("✅ 精度对齐：两个模型的输出结果非常接近。")
|
||||
else:
|
||||
print("❌ 精度不一致！")
|
||||
print(f"最大绝对误差: {(out_torch - out_cuda).abs().max().item()}" )
|
||||
|
||||
print("\n-------------------- 性能加速比测试 --------------------")
|
||||
iters = 30
|
||||
torch.cuda.synchronize(); t0 = time.time()
|
||||
for _ in range(iters):
|
||||
_ = torch_model(*inputs)
|
||||
torch.cuda.synchronize(); t_torch = (time.time() - t0) / iters
|
||||
|
||||
torch.cuda.synchronize(); t0 = time.time()
|
||||
for _ in range(iters):
|
||||
_ = cuda_model(*inputs)
|
||||
torch.cuda.synchronize(); t_cuda = (time.time() - t0) / iters
|
||||
|
||||
print(f"PyTorch Group-Softmax-Affine 平均执行时间: {t_torch:.6f} 秒")
|
||||
print(f"自定义 CUDA 融合内核 平均执行时间: {t_cuda:.6f} 秒")
|
||||
sp = t_torch / t_cuda if t_cuda > 0 else 0
|
||||
if t_cuda > 0:
|
||||
print(f"加速比 (Speedup): {sp:.2f}x")
|
||||
else:
|
||||
print("CUDA 内核执行时间为0，无法计算加速比。")
|
||||
return flag, sp
|
||||
|
||||
if __name__ == "__main__":
|
||||
run_benchmark()
|
||||
|
|
@ -0,0 +1,36 @@
|
|||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
class Model(nn.Module):
|
||||
def __init__(self, scale: torch.Tensor, bias: torch.Tensor, group: int):
|
||||
super(Model, self).__init__()
|
||||
self.scale = nn.Parameter(scale)
|
||||
self.bias = nn.Parameter(bias)
|
||||
self.register_buffer("group", torch.tensor(int(group), dtype=torch.int32))
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
g = int(self.group.item())
|
||||
z = x * self.scale.view(1,-1,1,1) + self.bias.view(1,-1,1,1)
|
||||
C = z.size(1)
|
||||
G = (C + g - 1) // g
|
||||
y = torch.empty_like(x)
|
||||
for gi in range(G):
|
||||
s = gi * g
|
||||
e = min(s + g, C)
|
||||
m = torch.amax(z[:, s:e], dim=1, keepdim=True)
|
||||
ex = torch.exp(z[:, s:e] - m)
|
||||
den = torch.sum(ex, dim=1, keepdim=True)
|
||||
y[:, s:e] = (ex / den) * x[:, s:e]
|
||||
return y
|
||||
|
||||
N, C, H, W = 4, 96, 32, 32
|
||||
|
||||
def get_inputs():
|
||||
x = torch.randn(N, C, H, W)
|
||||
return [x]
|
||||
|
||||
def get_init_inputs():
|
||||
scale = torch.randn(C)
|
||||
bias = torch.randn(C)
|
||||
group = 16
|
||||
return [scale, bias, group]
|
||||
Loading…
Reference in New Issue