forked from ccf-ai-infra/GPUCodeForces
finish logcosh-affinegate #58
This commit is contained in:
parent
10eed82956
commit
cb3102003e
|
|
@ -0,0 +1,165 @@
|
|||
import torch
|
||||
from torch.utils.cpp_extension import load_inline
|
||||
|
||||
source = """
|
||||
#include <torch/extension.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <algorithm>
|
||||
|
||||
__device__ __forceinline__ float logcosh(float x){
|
||||
float ax = fabsf(x);
|
||||
if(ax < 1e-3f){
|
||||
float x2 = x * x;
|
||||
float x4 = x2 * x2;
|
||||
return 0.5f * x2 - (1.0f/12.0f) * x4;
|
||||
}
|
||||
float t = __expf(-2.0f * ax);
|
||||
if(t < 1e-5f){
|
||||
return ax + t - 0.6931471805599453f;
|
||||
}
|
||||
return ax + __logf(1.0f + t) - 0.6931471805599453f;
|
||||
}
|
||||
|
||||
__device__ __forceinline__ float sigmoid_stable(float x){
|
||||
if(x >= 0.0f){
|
||||
float e = __expf(-x);
|
||||
return 1.0f / (1.0f + e);
|
||||
} else {
|
||||
float e = __expf(x);
|
||||
return e / (1.0f + e);
|
||||
}
|
||||
}
|
||||
|
||||
__global__ __launch_bounds__(1024) void logcosh_affine_gate_kernel(const float* __restrict__ x, const float* __restrict__ scale, const float* __restrict__ bias, float* __restrict__ y, int B, int D, float alpha, float beta){
|
||||
long long tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
long long stride = blockDim.x * gridDim.x;
|
||||
long long total = (long long)B * D;
|
||||
for(long long i = tid; i < total; i += stride){
|
||||
int d = (int)(i % D);
|
||||
float xv = __ldg(x + i);
|
||||
float z = fmaf(xv, __ldg(scale + d), __ldg(bias + d));
|
||||
float t = fmaf(logcosh(z), alpha, beta);
|
||||
float g = 1.0f / (1.0f + __expf(-t));
|
||||
y[i] = xv * g;
|
||||
}
|
||||
}
|
||||
|
||||
torch::Tensor logcosh_affine_gate_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, torch::Tensor alpha, torch::Tensor beta){
|
||||
auto xc = x.contiguous();
|
||||
auto sc = scale.contiguous();
|
||||
auto bc = bias.contiguous();
|
||||
auto y = torch::empty_like(xc);
|
||||
int B = (int)xc.size(0);
|
||||
int D = (int)xc.size(1);
|
||||
float a = alpha.item<float>();
|
||||
float be = beta.item<float>();
|
||||
int block = 1024;
|
||||
long long total = (long long)B * D;
|
||||
int grid = (int)std::min<long long>(65535LL, (total + block - 1) / block);
|
||||
logcosh_affine_gate_kernel<<<grid, block>>>(xc.data_ptr<float>(), sc.data_ptr<float>(), bc.data_ptr<float>(), y.data_ptr<float>(), B, D, a, be);
|
||||
return y;
|
||||
}
|
||||
|
||||
__global__ __launch_bounds__(512) void logcosh_affine_gate_kernel_v2(const float* __restrict__ x, const float* __restrict__ scale, const float* __restrict__ bias, float* __restrict__ y, int B, int D, float alpha, float beta){
|
||||
int b = blockIdx.x;
|
||||
int lane = blockIdx.y * blockDim.x + threadIdx.x;
|
||||
int stride = blockDim.x * gridDim.y;
|
||||
int row_start = b * D;
|
||||
const float* xr = x + row_start;
|
||||
float* yr = y + row_start;
|
||||
int aligned = ((((long long)xr & 15LL) == 0) && (((long long)yr & 15LL) == 0) && (((long long)scale & 15LL) == 0) && (((long long)bias & 15LL) == 0) && ((D & 3) == 0));
|
||||
if(aligned){
|
||||
int D4 = (D / 4) * 4;
|
||||
#pragma unroll 12
|
||||
for(int i = lane * 4; i < D4; i += stride * 4){
|
||||
float4 xv = reinterpret_cast<const float4*>(xr)[i / 4];
|
||||
float4 sv = reinterpret_cast<const float4*>(scale)[i / 4];
|
||||
float4 bv = reinterpret_cast<const float4*>(bias)[i / 4];
|
||||
float4 yv;
|
||||
float z0 = fmaf(xv.x, sv.x, bv.x);
|
||||
float z1 = fmaf(xv.y, sv.y, bv.y);
|
||||
float z2 = fmaf(xv.z, sv.z, bv.z);
|
||||
float z3 = fmaf(xv.w, sv.w, bv.w);
|
||||
float t0 = fmaf(logcosh(z0), alpha, beta);
|
||||
float t1 = fmaf(logcosh(z1), alpha, beta);
|
||||
float t2 = fmaf(logcosh(z2), alpha, beta);
|
||||
float t3 = fmaf(logcosh(z3), alpha, beta);
|
||||
float g0 = __fdividef(1.0f, 1.0f + __expf(-t0));
|
||||
float g1 = __fdividef(1.0f, 1.0f + __expf(-t1));
|
||||
float g2 = __fdividef(1.0f, 1.0f + __expf(-t2));
|
||||
float g3 = __fdividef(1.0f, 1.0f + __expf(-t3));
|
||||
yv.x = xv.x * g0;
|
||||
yv.y = xv.y * g1;
|
||||
yv.z = xv.z * g2;
|
||||
yv.w = xv.w * g3;
|
||||
reinterpret_cast<float4*>(yr)[i / 4] = yv;
|
||||
}
|
||||
#pragma unroll 12
|
||||
for(int i = D4 + lane; i < D; i += stride){
|
||||
float z = fmaf(__ldg(xr + i), __ldg(scale + i), __ldg(bias + i));
|
||||
float t = fmaf(logcosh(z), alpha, beta);
|
||||
float g = __fdividef(1.0f, 1.0f + __expf(-t));
|
||||
yr[i] = xr[i] * g;
|
||||
}
|
||||
} else {
|
||||
#pragma unroll 12
|
||||
for(int i = lane; i < D; i += stride){
|
||||
float z = fmaf(__ldg(xr + i), __ldg(scale + i), __ldg(bias + i));
|
||||
float t = fmaf(logcosh(z), alpha, beta);
|
||||
float g = __fdividef(1.0f, 1.0f + __expf(-t));
|
||||
yr[i] = xr[i] * g;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
torch::Tensor logcosh_affine_gate_cuda_v2(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, torch::Tensor alpha, torch::Tensor beta){
|
||||
auto xc = x.contiguous();
|
||||
auto sc = scale.contiguous();
|
||||
auto bc = bias.contiguous();
|
||||
auto y = torch::empty_like(xc);
|
||||
int B = (int)xc.size(0);
|
||||
int D = (int)xc.size(1);
|
||||
float a = alpha.item<float>();
|
||||
float be = beta.item<float>();
|
||||
int block = 512;
|
||||
int gy = std::min(32, std::max(1, (D + 1023) / 1024));
|
||||
dim3 grid(B, gy);
|
||||
logcosh_affine_gate_kernel_v2<<<grid, block>>>(xc.data_ptr<float>(), sc.data_ptr<float>(), bc.data_ptr<float>(), y.data_ptr<float>(), B, D, a, be);
|
||||
return y;
|
||||
}
|
||||
|
||||
torch::Tensor logcosh_affine_gate_cuda_opt(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, torch::Tensor alpha, torch::Tensor beta){
|
||||
int D = (int)x.size(1);
|
||||
if((D & 3) == 0){
|
||||
return logcosh_affine_gate_cuda_v2(x, scale, bias, alpha, beta);
|
||||
}
|
||||
return logcosh_affine_gate_cuda(x, scale, bias, alpha, beta);
|
||||
}
|
||||
"""
|
||||
|
||||
cpp_source = """
|
||||
torch::Tensor logcosh_affine_gate_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, torch::Tensor alpha, torch::Tensor beta);
|
||||
torch::Tensor logcosh_affine_gate_cuda_v2(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, torch::Tensor alpha, torch::Tensor beta);
|
||||
torch::Tensor logcosh_affine_gate_cuda_opt(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, torch::Tensor alpha, torch::Tensor beta);
|
||||
"""
|
||||
|
||||
ops = load_inline(
|
||||
name="logcosh_affine_gate",
|
||||
cpp_sources=cpp_source,
|
||||
cuda_sources=source,
|
||||
functions=["logcosh_affine_gate_cuda","logcosh_affine_gate_cuda_v2","logcosh_affine_gate_cuda_opt"],
|
||||
extra_cuda_cflags=["-O3","--use_fast_math","-gencode=arch=compute_80,code=sm_80","-Xptxas","-O3,-dlcm=ca","-maxrregcount=64"],
|
||||
verbose=True
|
||||
)
|
||||
|
||||
class ModelNew(torch.nn.Module):
|
||||
def __init__(self, scale: torch.Tensor, bias: torch.Tensor, alpha: float, beta: float):
|
||||
super(ModelNew, self).__init__()
|
||||
self.ops = ops
|
||||
self.register_buffer("scale", scale)
|
||||
self.register_buffer("bias", bias)
|
||||
self.register_buffer("alpha", torch.tensor(float(alpha), dtype=torch.float32))
|
||||
self.register_buffer("beta", torch.tensor(float(beta), dtype=torch.float32))
|
||||
|
||||
def forward(self, x):
|
||||
return self.ops.logcosh_affine_gate_cuda_opt(x, self.scale, self.bias, self.alpha, self.beta)
|
||||
|
|
@ -0,0 +1,19 @@
|
|||
Objective: Replace PyTorch operations with a single fused CUDA kernel for LogCosh-Affine-Gate to achieve ≥1.3x speedup while matching outputs within rtol=1e-3.
|
||||
|
||||
Constraints:
|
||||
- Keep the 4-file structure identical to `example` folder: `cudacode.py`, `torchcode.py`, `run_code.py`, `prompt.txt`.
|
||||
- The fused kernel must compute: `z = x*scale + bias; v = log(cosh(z)); g = sigmoid(alpha*v + beta); y = x*g`.
|
||||
- Use numerically stable math and avoid NaNs for typical random inputs.
|
||||
- Exploit GPU multi-threading aggressively and minimize memory traffic via fusion.
|
||||
|
||||
Design Guidelines:
|
||||
- Launch one block per row with 256–512 threads, and set `grid.y` to split long rows for higher SM occupancy.
|
||||
- Prefer vectorized loads/stores (`float4`) on aligned paths, and fallback to scalar path otherwise.
|
||||
- Use FMA for affine and gate preparation to reduce instruction count and rounding.
|
||||
- Use fast intrinsic `__expf` for the sigmoid path, verify accuracy under rtol=1e-3.
|
||||
- Keep reductions out of the hot path; this kernel is purely elementwise and memory-bound.
|
||||
|
||||
Benchmark Setup:
|
||||
- Batch size: 16, Dimension: 16384.
|
||||
- Measure average latency over 100 iterations for both PyTorch and the fused CUDA kernel.
|
||||
- Report precision alignment and speedup. Target speedup ≥1.3x.
|
||||
|
|
@ -0,0 +1,50 @@
|
|||
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().eval()
|
||||
cuda_model = ModelNew(*init_inputs).cuda().eval()
|
||||
|
||||
print("-------------------- 精度对齐验证 --------------------")
|
||||
with torch.no_grad():
|
||||
output_torch = torch_model(*inputs)
|
||||
output_cuda = cuda_model(*inputs)
|
||||
precision_flag = torch.allclose(output_torch, output_cuda, rtol=1e-03)
|
||||
if precision_flag:
|
||||
print("✅ 精度对齐：两个模型的输出结果非常接近。")
|
||||
else:
|
||||
print("❌ 精度不一致！")
|
||||
|
||||
print("\n-------------------- 性能加速比测试 --------------------")
|
||||
num_iterations = 100
|
||||
torch.cuda.synchronize(); start_time = time.time()
|
||||
for _ in range(num_iterations):
|
||||
_ = torch_model(*inputs)
|
||||
torch.cuda.synchronize(); torch_time = (time.time() - start_time) / num_iterations
|
||||
|
||||
torch.cuda.synchronize(); start_time = time.time()
|
||||
for _ in range(num_iterations):
|
||||
_ = cuda_model(*inputs)
|
||||
torch.cuda.synchronize(); cuda_time = (time.time() - start_time) / num_iterations
|
||||
|
||||
print(f"PyTorch LogCosh-Affine-Gate 平均执行时间: {torch_time:.6f} 秒")
|
||||
print(f"自定义 CUDA 融合内核 平均执行时间: {cuda_time:.6f} 秒")
|
||||
speedup = torch_time / cuda_time if cuda_time > 0 else 0
|
||||
if cuda_time > 0:
|
||||
print(f"加速比 (Speedup): {speedup:.2f}x")
|
||||
else:
|
||||
print("CUDA 内核执行时间为0，无法计算加速比。")
|
||||
return precision_flag, speedup
|
||||
|
||||
if __name__ == "__main__":
|
||||
run_benchmark()
|
||||
|
|
@ -0,0 +1,28 @@
|
|||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
class Model(nn.Module):
|
||||
def __init__(self, scale: torch.Tensor, bias: torch.Tensor, alpha: float, beta: float):
|
||||
super(Model, self).__init__()
|
||||
self.register_buffer("scale", scale)
|
||||
self.register_buffer("bias", bias)
|
||||
self.register_buffer("alpha", torch.tensor(float(alpha), dtype=torch.float32))
|
||||
self.register_buffer("beta", torch.tensor(float(beta), dtype=torch.float32))
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
z = x * self.scale + self.bias
|
||||
v = torch.log(torch.cosh(z))
|
||||
g = torch.sigmoid(self.alpha * v + self.beta)
|
||||
return x * g
|
||||
|
||||
batch_size = 16
|
||||
dim = 16384
|
||||
|
||||
def get_inputs():
|
||||
x = torch.randn(batch_size, dim)
|
||||
return [x]
|
||||
|
||||
def get_init_inputs():
|
||||
scale = torch.randn(dim)
|
||||
bias = torch.randn(dim)
|
||||
return [scale, bias, 1.0, 0.0]
|
||||
Loading…
Reference in New Issue