Compare commits
No commits in common. "main" and "main" have entirely different histories.
|
|
@ -1,120 +0,0 @@
|
|||
import torch
|
||||
from torch.utils.cpp_extension import load_inline
|
||||
|
||||
source = """
|
||||
#include <torch/extension.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <vector_types.h>
|
||||
|
||||
__global__ __launch_bounds__(256) void log1pabs_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){
|
||||
int row = blockIdx.x;
|
||||
int tid = threadIdx.x;
|
||||
int col_block = blockIdx.y;
|
||||
int stride4 = blockDim.x * 4;
|
||||
int col_idx = (col_block * stride4) + tid * 4;
|
||||
const float* xr = x + row * D;
|
||||
float* yr = y + row * D;
|
||||
for(int base = col_idx; base < D; base += stride4 * 2){
|
||||
if (base < D){
|
||||
float4 xv = reinterpret_cast<const float4*>(xr + base)[0];
|
||||
float4 sv = reinterpret_cast<const float4*>(scale + base)[0];
|
||||
float4 bv = reinterpret_cast<const float4*>(bias + base)[0];
|
||||
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 m0 = log1pf(fabsf(z0));
|
||||
float m1 = log1pf(fabsf(z1));
|
||||
float m2 = log1pf(fabsf(z2));
|
||||
float m3 = log1pf(fabsf(z3));
|
||||
float t0 = fmaf(alpha, m0, beta);
|
||||
float t1 = fmaf(alpha, m1, beta);
|
||||
float t2 = fmaf(alpha, m2, beta);
|
||||
float t3 = fmaf(alpha, m3, 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 + base)[0] = yv;
|
||||
}
|
||||
int base2 = base + stride4;
|
||||
if (base2 < D){
|
||||
float4 xv2 = reinterpret_cast<const float4*>(xr + base2)[0];
|
||||
float4 sv2 = reinterpret_cast<const float4*>(scale + base2)[0];
|
||||
float4 bv2 = reinterpret_cast<const float4*>(bias + base2)[0];
|
||||
float4 yv2;
|
||||
float z0b = fmaf(xv2.x, sv2.x, bv2.x);
|
||||
float z1b = fmaf(xv2.y, sv2.y, bv2.y);
|
||||
float z2b = fmaf(xv2.z, sv2.z, bv2.z);
|
||||
float z3b = fmaf(xv2.w, sv2.w, bv2.w);
|
||||
float mb0 = log1pf(fabsf(z0b));
|
||||
float mb1 = log1pf(fabsf(z1b));
|
||||
float mb2 = log1pf(fabsf(z2b));
|
||||
float mb3 = log1pf(fabsf(z3b));
|
||||
float tb0 = fmaf(alpha, mb0, beta);
|
||||
float tb1 = fmaf(alpha, mb1, beta);
|
||||
float tb2 = fmaf(alpha, mb2, beta);
|
||||
float tb3 = fmaf(alpha, mb3, beta);
|
||||
float gb0 = __fdividef(1.0f, 1.0f + __expf(-tb0));
|
||||
float gb1 = __fdividef(1.0f, 1.0f + __expf(-tb1));
|
||||
float gb2 = __fdividef(1.0f, 1.0f + __expf(-tb2));
|
||||
float gb3 = __fdividef(1.0f, 1.0f + __expf(-tb3));
|
||||
yv2.x = xv2.x * gb0;
|
||||
yv2.y = xv2.y * gb1;
|
||||
yv2.z = xv2.z * gb2;
|
||||
yv2.w = xv2.w * gb3;
|
||||
reinterpret_cast<float4*>(yr + base2)[0] = yv2;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
torch::Tensor log1pabs_affine_gate_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, double alpha, double 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 = (float)alpha;
|
||||
float be = (float)beta;
|
||||
int block = 256;
|
||||
int elements_per_thread = 4;
|
||||
int elements_per_block = block * elements_per_thread;
|
||||
int gy = (D + elements_per_block - 1) / elements_per_block;
|
||||
dim3 grid(B, gy);
|
||||
log1pabs_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;
|
||||
}
|
||||
"""
|
||||
|
||||
cpp_source = """
|
||||
#include <torch/extension.h>
|
||||
torch::Tensor log1pabs_affine_gate_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, double alpha, double beta);
|
||||
"""
|
||||
|
||||
ops = load_inline(
|
||||
name="log1pabs_affine_gate",
|
||||
cpp_sources=cpp_source,
|
||||
cuda_sources=source,
|
||||
functions=["log1pabs_affine_gate_cuda"],
|
||||
extra_cflags=["-O3","-std=c++17"],
|
||||
extra_cuda_cflags=["-O3","--use_fast_math","-std=c++17","-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.alpha = float(alpha)
|
||||
self.beta = float(beta)
|
||||
|
||||
def forward(self, x):
|
||||
return self.ops.log1pabs_affine_gate_cuda(x, self.scale, self.bias, self.alpha, self.beta)
|
||||
|
|
@ -1,10 +0,0 @@
|
|||
Operator: Log1pAbs-Affine-Gate (Fused CUDA Kernel)
|
||||
|
||||
Definition
|
||||
- z = x * scale + bias
|
||||
- m = log(1 + |z|)
|
||||
- g = sigmoid(alpha * m + beta)
|
||||
- y = x * g
|
||||
|
||||
Goal
|
||||
- Stable log1p fusion; target ≥1.30x speedup.
|
||||
|
|
@ -1,50 +0,0 @@
|
|||
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 Log1pAbs-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()
|
||||
|
|
@ -1,28 +0,0 @@
|
|||
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
|
||||
m = torch.log1p(torch.abs(z))
|
||||
g = torch.sigmoid(self.alpha * m + 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]
|
||||
|
|
@ -1,94 +0,0 @@
|
|||
import torch
|
||||
from torch.utils.cpp_extension import load_inline
|
||||
|
||||
source = """
|
||||
#include <torch/extension.h>
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
__global__ void square_sigmoid_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){
|
||||
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 4
|
||||
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 g0 = 1.0f / (1.0f + expf(-(alpha * (z0 * z0) + beta)));
|
||||
float g1 = 1.0f / (1.0f + expf(-(alpha * (z1 * z1) + beta)));
|
||||
float g2 = 1.0f / (1.0f + expf(-(alpha * (z2 * z2) + beta)));
|
||||
float g3 = 1.0f / (1.0f + expf(-(alpha * (z3 * z3) + beta)));
|
||||
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 4
|
||||
for(int i = D4 + lane; i < D; i += stride){
|
||||
float z = fmaf(xr[i], scale[i], bias[i]);
|
||||
float g = 1.0f / (1.0f + expf(-(alpha * (z * z) + beta)));
|
||||
yr[i] = xr[i] * g;
|
||||
}
|
||||
} else {
|
||||
#pragma unroll 4
|
||||
for(int i = lane; i < D; i += stride){
|
||||
float z = fmaf(xr[i], scale[i], bias[i]);
|
||||
float g = 1.0f / (1.0f + expf(-(alpha * (z * z) + beta)));
|
||||
yr[i] = xr[i] * g;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
torch::Tensor square_sigmoid_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 = 256;
|
||||
int gy = max(1, min((D + 4095) / 4096, 8));
|
||||
dim3 grid(B, gy);
|
||||
square_sigmoid_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;
|
||||
}
|
||||
"""
|
||||
|
||||
cpp_source = """
|
||||
torch::Tensor square_sigmoid_affine_gate_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, torch::Tensor alpha, torch::Tensor beta);
|
||||
"""
|
||||
|
||||
ops = load_inline(
|
||||
name="square_sigmoid_affine_gate",
|
||||
cpp_sources=cpp_source,
|
||||
cuda_sources=source,
|
||||
functions=["square_sigmoid_affine_gate_cuda"],
|
||||
extra_cuda_cflags=["-O3","--use_fast_math"],
|
||||
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.square_sigmoid_affine_gate_cuda(x, self.scale, self.bias, self.alpha, self.beta)
|
||||
|
|
@ -1,28 +0,0 @@
|
|||
融合算子：Square-Sigmoid-Affine-Gate（一次核内完成仿射、平方与 Sigmoid 门控，返回 y = x * σ(α * z^2 + β)，其中 z = x*scale + bias）。平方增强幅值差异并通过 Sigmoid 控制门控强度。
|
||||
|
||||
目标与定义
|
||||
- 输入张量：`x[B, D]`
|
||||
- 逐维参数：`scale[D]`、`bias[D]`
|
||||
- 标量超参：`alpha`、`beta`
|
||||
- 计算流程：`z = x*scale + bias`，`v = z*z`，`g = sigmoid(alpha*v + beta)`，`y = x * g`
|
||||
|
||||
参考实现（文件要求）
|
||||
- `torchcode.py`：PyTorch 参考 `Model`；统一的 `get_inputs()`/`get_init_inputs()`
|
||||
- `cudacode.py`：单核融合（仿射+平方+sigmoid+乘法）；`-O3 --use_fast_math`
|
||||
- `run_code.py`：迭代 100 次；`rtol=1e-03, atol=1e-06` 精度；打印加速比
|
||||
|
||||
CUDA 实现要点
|
||||
- 并行布局：`grid = B`；块内沿 D 合并访存
|
||||
- 对齐向量化：16 字节对齐且 `D%4==0` 时走 `float4`；否则标量回退
|
||||
- 指令优化：仿射用 `fmaf`；sigmoid 用 `expf`；循环 `#pragma unroll 4`
|
||||
- 溢出注意：`z^2` 对大幅值会放大，sigmoid 可缓和，但仍需避免中间溢出；使用 `float` 常规范围下问题不大
|
||||
- 线程配置：推荐 `block=1024`，按设备与规模微调
|
||||
|
||||
评估与目标
|
||||
- 精度：对齐 `rtol=1e-03, atol=1e-06`
|
||||
- 性能：≥1.0x 加速；对齐触发向量化时更佳
|
||||
|
||||
加分项（可选）
|
||||
- 尾元素处理与分支收敛优化
|
||||
- 每线程批量步长以提高吞吐与占用
|
||||
|
||||
|
|
@ -1,57 +0,0 @@
|
|||
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():
|
||||
output_torch = torch_model(*inputs)
|
||||
output_cuda = cuda_model(*inputs)
|
||||
precision_flag = torch.allclose(output_torch, output_cuda, rtol=1e-03, atol=1e-06)
|
||||
if precision_flag:
|
||||
print("✅ 精度对齐：两个模型的输出结果非常接近。")
|
||||
else:
|
||||
print("❌ 精度不一致！")
|
||||
diff = (output_torch - output_cuda).abs().max().item()
|
||||
print(f"最大绝对误差: {diff}")
|
||||
print(f"输出张量形状: torch={tuple(output_torch.shape)}, cuda={tuple(output_cuda.shape)}")
|
||||
print(f"数据类型: torch={output_torch.dtype}, cuda={output_cuda.dtype}")
|
||||
print(f"设备: torch={output_torch.device}, cuda={output_cuda.device}")
|
||||
|
||||
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 Square-Sigmoid-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()
|
||||
|
||||
|
|
@ -1,29 +0,0 @@
|
|||
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 = z * 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]
|
||||
|
||||
|
|
@ -1,94 +0,0 @@
|
|||
import torch
|
||||
from torch.utils.cpp_extension import load_inline
|
||||
|
||||
source = """
|
||||
#include <torch/extension.h>
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
__global__ void atan_sigmoid_mix_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){
|
||||
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 4
|
||||
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 g0 = 1.0f / (1.0f + expf(-(alpha * atanf(z0) + beta)));
|
||||
float g1 = 1.0f / (1.0f + expf(-(alpha * atanf(z1) + beta)));
|
||||
float g2 = 1.0f / (1.0f + expf(-(alpha * atanf(z2) + beta)));
|
||||
float g3 = 1.0f / (1.0f + expf(-(alpha * atanf(z3) + beta)));
|
||||
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 4
|
||||
for(int i = D4 + lane; i < D; i += stride){
|
||||
float z = fmaf(xr[i], scale[i], bias[i]);
|
||||
float g = 1.0f / (1.0f + expf(-(alpha * atanf(z) + beta)));
|
||||
yr[i] = xr[i] * g;
|
||||
}
|
||||
} else {
|
||||
#pragma unroll 4
|
||||
for(int i = lane; i < D; i += stride){
|
||||
float z = fmaf(xr[i], scale[i], bias[i]);
|
||||
float g = 1.0f / (1.0f + expf(-(alpha * atanf(z) + beta)));
|
||||
yr[i] = xr[i] * g;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
torch::Tensor atan_sigmoid_mix_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 = 256;
|
||||
int gy = max(1, min((D + 4095) / 4096, 8));
|
||||
dim3 grid(B, gy);
|
||||
atan_sigmoid_mix_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;
|
||||
}
|
||||
"""
|
||||
|
||||
cpp_source = """
|
||||
torch::Tensor atan_sigmoid_mix_gate_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, torch::Tensor alpha, torch::Tensor beta);
|
||||
"""
|
||||
|
||||
ops = load_inline(
|
||||
name="atan_sigmoid_mix_gate",
|
||||
cpp_sources=cpp_source,
|
||||
cuda_sources=source,
|
||||
functions=["atan_sigmoid_mix_gate_cuda"],
|
||||
extra_cuda_cflags=["-O3","--use_fast_math"],
|
||||
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.atan_sigmoid_mix_gate_cuda(x, self.scale, self.bias, self.alpha, self.beta)
|
||||
|
|
@ -1,27 +0,0 @@
|
|||
融合算子：Atan-Sigmoid-Mix-Gate（一次核内完成仿射与反正切混合门控，返回 y = x * σ(α * atan(z) + β)，其中 z = x*scale + bias）。atan 在大幅值区间趋于常数，有利于抑制过大输入的门控强度。
|
||||
|
||||
目标与定义
|
||||
- 输入张量：`x[B, D]`
|
||||
- 逐维参数：`scale[D]`、`bias[D]`
|
||||
- 标量超参：`alpha`、`beta`
|
||||
- 计算流程：`z = x*scale + bias`，`g = sigmoid(alpha*atan(z) + beta)`，`y = x * g`
|
||||
|
||||
参考实现（文件要求）
|
||||
- `torchcode.py`：PyTorch 参考 `Model`；统一接口
|
||||
- `cudacode.py`：单核融合（仿射+atan+sigmoid+乘法）；`-O3 --use_fast_math`
|
||||
- `run_code.py`：100 次迭代；精度 `rtol=1e-03, atol=1e-06`；打印加速比
|
||||
|
||||
CUDA 实现要点
|
||||
- 行并行：`grid = B`；块内沿 D 连续访存；一次遍历写回
|
||||
- 对齐向量化：16 字节对齐且 `D%4==0` 走 `float4`，否则标量回退
|
||||
- 指令优化：仿射用 `fmaf`；`atanf` 与 `expf` 走快速数学；循环 `#pragma unroll 4`
|
||||
- 线程配置：`block=1024` 起步，按设备试探最佳
|
||||
|
||||
评估与目标
|
||||
- 精度：满足 `rtol=1e-03, atol=1e-06`
|
||||
- 性能：≥1.0x 加速，向量化与融合带来优势
|
||||
|
||||
加分项（可选）
|
||||
- 尾元素处理与分支收敛优化
|
||||
- 根据分布特性调参 `alpha/beta` 增强门控稳定性（参考实现一致性优先）
|
||||
|
||||
|
|
@ -1,57 +0,0 @@
|
|||
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():
|
||||
output_torch = torch_model(*inputs)
|
||||
output_cuda = cuda_model(*inputs)
|
||||
precision_flag = torch.allclose(output_torch, output_cuda, rtol=1e-03, atol=1e-06)
|
||||
if precision_flag:
|
||||
print("✅ 精度对齐：两个模型的输出结果非常接近。")
|
||||
else:
|
||||
print("❌ 精度不一致！")
|
||||
diff = (output_torch - output_cuda).abs().max().item()
|
||||
print(f"最大绝对误差: {diff}")
|
||||
print(f"输出张量形状: torch={tuple(output_torch.shape)}, cuda={tuple(output_cuda.shape)}")
|
||||
print(f"数据类型: torch={output_torch.dtype}, cuda={output_cuda.dtype}")
|
||||
print(f"设备: torch={output_torch.device}, cuda={output_cuda.device}")
|
||||
|
||||
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 Atan-Sigmoid-Mix-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()
|
||||
|
||||
|
|
@ -1,28 +0,0 @@
|
|||
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
|
||||
g = torch.sigmoid(self.alpha * torch.atan(z) + 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]
|
||||
|
||||
|
|
@ -1,126 +0,0 @@
|
|||
import torch
|
||||
from torch.utils.cpp_extension import load_inline
|
||||
|
||||
source = """
|
||||
#include <torch/extension.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <vector_types.h>
|
||||
|
||||
__device__ __forceinline__ float huberf(float z){
|
||||
float az = fabsf(z);
|
||||
if (az <= 1.0f) return 0.5f * z * z;
|
||||
return az - 0.5f;
|
||||
}
|
||||
|
||||
__global__ __launch_bounds__(256) void huber_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){
|
||||
int row = blockIdx.x;
|
||||
int tid = threadIdx.x;
|
||||
int col_block = blockIdx.y;
|
||||
int stride4 = blockDim.x * 4;
|
||||
int col_idx = (col_block * stride4) + tid * 4;
|
||||
const float* xr = x + row * D;
|
||||
float* yr = y + row * D;
|
||||
for(int base = col_idx; base < D; base += stride4 * 2){
|
||||
if (base < D){
|
||||
float4 xv = reinterpret_cast<const float4*>(xr + base)[0];
|
||||
float4 sv = reinterpret_cast<const float4*>(scale + base)[0];
|
||||
float4 bv = reinterpret_cast<const float4*>(bias + base)[0];
|
||||
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 m0 = huberf(z0);
|
||||
float m1 = huberf(z1);
|
||||
float m2 = huberf(z2);
|
||||
float m3 = huberf(z3);
|
||||
float t0 = fmaf(alpha, m0, beta);
|
||||
float t1 = fmaf(alpha, m1, beta);
|
||||
float t2 = fmaf(alpha, m2, beta);
|
||||
float t3 = fmaf(alpha, m3, 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 + base)[0] = yv;
|
||||
}
|
||||
int base2 = base + stride4;
|
||||
if (base2 < D){
|
||||
float4 xv2 = reinterpret_cast<const float4*>(xr + base2)[0];
|
||||
float4 sv2 = reinterpret_cast<const float4*>(scale + base2)[0];
|
||||
float4 bv2 = reinterpret_cast<const float4*>(bias + base2)[0];
|
||||
float4 yv2;
|
||||
float z0b = fmaf(xv2.x, sv2.x, bv2.x);
|
||||
float z1b = fmaf(xv2.y, sv2.y, bv2.y);
|
||||
float z2b = fmaf(xv2.z, sv2.z, bv2.z);
|
||||
float z3b = fmaf(xv2.w, sv2.w, bv2.w);
|
||||
float mb0 = huberf(z0b);
|
||||
float mb1 = huberf(z1b);
|
||||
float mb2 = huberf(z2b);
|
||||
float mb3 = huberf(z3b);
|
||||
float tb0 = fmaf(alpha, mb0, beta);
|
||||
float tb1 = fmaf(alpha, mb1, beta);
|
||||
float tb2 = fmaf(alpha, mb2, beta);
|
||||
float tb3 = fmaf(alpha, mb3, beta);
|
||||
float gb0 = __fdividef(1.0f, 1.0f + __expf(-tb0));
|
||||
float gb1 = __fdividef(1.0f, 1.0f + __expf(-tb1));
|
||||
float gb2 = __fdividef(1.0f, 1.0f + __expf(-tb2));
|
||||
float gb3 = __fdividef(1.0f, 1.0f + __expf(-tb3));
|
||||
yv2.x = xv2.x * gb0;
|
||||
yv2.y = xv2.y * gb1;
|
||||
yv2.z = xv2.z * gb2;
|
||||
yv2.w = xv2.w * gb3;
|
||||
reinterpret_cast<float4*>(yr + base2)[0] = yv2;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
torch::Tensor huber_affine_gate_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, double alpha, double 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 = (float)alpha;
|
||||
float be = (float)beta;
|
||||
int block = 256;
|
||||
int elements_per_thread = 4;
|
||||
int elements_per_block = block * elements_per_thread;
|
||||
int gy = (D + elements_per_block - 1) / elements_per_block;
|
||||
dim3 grid(B, gy);
|
||||
huber_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;
|
||||
}
|
||||
"""
|
||||
|
||||
cpp_source = """
|
||||
#include <torch/extension.h>
|
||||
torch::Tensor huber_affine_gate_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, double alpha, double beta);
|
||||
"""
|
||||
|
||||
ops = load_inline(
|
||||
name="huber_affine_gate",
|
||||
cpp_sources=cpp_source,
|
||||
cuda_sources=source,
|
||||
functions=["huber_affine_gate_cuda"],
|
||||
extra_cflags=["-O3","-std=c++17"],
|
||||
extra_cuda_cflags=["-O3","--use_fast_math","-std=c++17","-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.alpha = float(alpha)
|
||||
self.beta = float(beta)
|
||||
|
||||
def forward(self, x):
|
||||
return self.ops.huber_affine_gate_cuda(x, self.scale, self.bias, self.alpha, self.beta)
|
||||
|
|
@ -1,10 +0,0 @@
|
|||
Operator: Huber-Affine-Gate (Fused CUDA Kernel)
|
||||
|
||||
Definition
|
||||
- z = x * scale + bias
|
||||
- m = 0.5*z^2 if |z|<=1 else |z| - 0.5
|
||||
- g = sigmoid(alpha * m + beta)
|
||||
- y = x * g
|
||||
|
||||
Goal
|
||||
- Piecewise smooth fusion; target ≥1.30x speedup.
|
||||
|
|
@ -1,50 +0,0 @@
|
|||
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 Huber-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()
|
||||
|
|
@ -1,29 +0,0 @@
|
|||
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
|
||||
az = torch.abs(z)
|
||||
m = torch.where(az <= 1.0, 0.5 * z * z, az - 0.5)
|
||||
g = torch.sigmoid(self.alpha * m + 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]
|
||||
|
|
@ -1,53 +0,0 @@
|
|||
import torch
|
||||
from torch.utils.cpp_extension import load_inline
|
||||
|
||||
source = """
|
||||
#include <torch/extension.h>
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
__global__ void bias_gelu_kernel(const float* x, const float* bias, float* y, int dim, long long total) {
|
||||
long long idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
long long stride = blockDim.x * gridDim.x;
|
||||
for (long long i = idx; i < total; i += stride) {
|
||||
int j = (int)(i % dim);
|
||||
float z = x[i] + bias[j];
|
||||
float c = 0.7978845608f;
|
||||
float p = z + 0.044715f * z * z * z;
|
||||
y[i] = 0.5f * z * (1.f + tanhf(c * p));
|
||||
}
|
||||
}
|
||||
|
||||
torch::Tensor bias_gelu_cuda(torch::Tensor x, torch::Tensor bias) {
|
||||
auto x_contig = x.contiguous();
|
||||
auto b_contig = bias.contiguous();
|
||||
auto y = torch::empty_like(x_contig);
|
||||
long long total = x_contig.numel();
|
||||
int dim = (int)x_contig.size(-1);
|
||||
int block = 512;
|
||||
long long grid = (total + block - 1) / block;
|
||||
grid = grid > 65535 ? 65535 : grid;
|
||||
bias_gelu_kernel<<<(int)grid, block>>>(x_contig.data_ptr<float>(), b_contig.data_ptr<float>(), y.data_ptr<float>(), dim, total);
|
||||
return y;
|
||||
}
|
||||
"""
|
||||
|
||||
cpp_source = """
|
||||
torch::Tensor bias_gelu_cuda(torch::Tensor x, torch::Tensor bias);
|
||||
"""
|
||||
|
||||
ops = load_inline(
|
||||
name="bias_gelu",
|
||||
cpp_sources=cpp_source,
|
||||
cuda_sources=source,
|
||||
functions=["bias_gelu_cuda"],
|
||||
verbose=True
|
||||
)
|
||||
|
||||
class ModelNew(torch.nn.Module):
|
||||
def __init__(self, bias: torch.Tensor):
|
||||
super(ModelNew, self).__init__()
|
||||
self.ops = ops
|
||||
self.register_buffer("bias", bias)
|
||||
|
||||
def forward(self, x):
|
||||
return self.ops.bias_gelu_cuda(x, self.bias)
|
||||
|
|
@ -1,7 +0,0 @@
|
|||
本目录展示一个避免常见 LayerNorm 的融合算子：Bias+GELU（tanh 近似）。
|
||||
|
||||
torchcode.py 提供 PyTorch 参考实现：`y = gelu(x + bias)`，使用 `approximate='tanh'` 以匹配 CUDA 近似。
|
||||
|
||||
cudacode.py 内含 `__global__ void bias_gelu_kernel(...)`，一次遍历完成加偏置与 GELU 计算，减少显存往返与内核启动次数。
|
||||
|
||||
run_code.py 负责精度和性能对比，迭代 100 次并输出平均耗时与加速比，精度以 `rtol=1e-03` 检验。
|
||||
|
|
@ -1,76 +0,0 @@
|
|||
###########################################################
|
||||
# 性能和精度验证程序
|
||||
###########################################################
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
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
|
||||
else:
|
||||
device = torch.device("cuda")
|
||||
|
||||
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)
|
||||
|
||||
precision_flag = torch.allclose(output_torch, output_cuda,rtol=1e-03)
|
||||
if precision_flag:
|
||||
print("✅ 精度对齐：两个模型的输出结果非常接近。")
|
||||
else:
|
||||
print("❌ 精度不一致！")
|
||||
diff = (output_torch - output_cuda).abs().max().item()
|
||||
print(f"最大绝对误差: {diff}")
|
||||
print(f"输出张量形状: torch={tuple(output_torch.shape)}, cuda={tuple(output_cuda.shape)}")
|
||||
print(f"数据类型: torch={output_torch.dtype}, cuda={output_cuda.dtype}")
|
||||
print(f"设备: torch={output_torch.device}, cuda={output_cuda.device}")
|
||||
|
||||
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 Bias+GELU 平均执行时间: {torch_time:.6f} 秒")
|
||||
print(f"自定义 CUDA 融合内核 平均执行时间: {cuda_time:.6f} 秒")
|
||||
speedup = 0
|
||||
if cuda_time > 0:
|
||||
speedup = torch_time / cuda_time
|
||||
print(f"加速比 (Speedup): {speedup:.2f}x")
|
||||
else:
|
||||
print("CUDA 内核执行时间为0，无法计算加速比。")
|
||||
return precision_flag,speedup
|
||||
|
||||
if __name__ == "__main__":
|
||||
precision_flag,speedup = run_benchmark()
|
||||
|
|
@ -1,22 +0,0 @@
|
|||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
class Model(nn.Module):
|
||||
def __init__(self, bias: torch.Tensor):
|
||||
super(Model, self).__init__()
|
||||
self.register_buffer("bias", bias)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
return F.gelu(x + self.bias, approximate='tanh')
|
||||
|
||||
batch_size = 16
|
||||
dim = 16384
|
||||
|
||||
def get_inputs():
|
||||
x = torch.randn(batch_size, dim)
|
||||
return [x]
|
||||
|
||||
def get_init_inputs():
|
||||
bias = torch.randn(dim)
|
||||
return [bias]
|
||||
|
|
@ -1,120 +0,0 @@
|
|||
import torch
|
||||
from torch.utils.cpp_extension import load_inline
|
||||
|
||||
source = """
|
||||
#include <torch/extension.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <vector_types.h>
|
||||
|
||||
__global__ __launch_bounds__(256) void gaussian_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){
|
||||
int row = blockIdx.x;
|
||||
int tid = threadIdx.x;
|
||||
int col_block = blockIdx.y;
|
||||
int stride4 = blockDim.x * 4;
|
||||
int col_idx = (col_block * stride4) + tid * 4;
|
||||
const float* xr = x + row * D;
|
||||
float* yr = y + row * D;
|
||||
for(int base = col_idx; base < D; base += stride4 * 2){
|
||||
if (base < D){
|
||||
float4 xv = reinterpret_cast<const float4*>(xr + base)[0];
|
||||
float4 sv = reinterpret_cast<const float4*>(scale + base)[0];
|
||||
float4 bv = reinterpret_cast<const float4*>(bias + base)[0];
|
||||
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 m0 = __expf(-(z0*z0));
|
||||
float m1 = __expf(-(z1*z1));
|
||||
float m2 = __expf(-(z2*z2));
|
||||
float m3 = __expf(-(z3*z3));
|
||||
float t0 = fmaf(alpha, m0, beta);
|
||||
float t1 = fmaf(alpha, m1, beta);
|
||||
float t2 = fmaf(alpha, m2, beta);
|
||||
float t3 = fmaf(alpha, m3, 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 + base)[0] = yv;
|
||||
}
|
||||
int base2 = base + stride4;
|
||||
if (base2 < D){
|
||||
float4 xv2 = reinterpret_cast<const float4*>(xr + base2)[0];
|
||||
float4 sv2 = reinterpret_cast<const float4*>(scale + base2)[0];
|
||||
float4 bv2 = reinterpret_cast<const float4*>(bias + base2)[0];
|
||||
float4 yv2;
|
||||
float z0b = fmaf(xv2.x, sv2.x, bv2.x);
|
||||
float z1b = fmaf(xv2.y, sv2.y, bv2.y);
|
||||
float z2b = fmaf(xv2.z, sv2.z, bv2.z);
|
||||
float z3b = fmaf(xv2.w, sv2.w, bv2.w);
|
||||
float mb0 = __expf(-(z0b*z0b));
|
||||
float mb1 = __expf(-(z1b*z1b));
|
||||
float mb2 = __expf(-(z2b*z2b));
|
||||
float mb3 = __expf(-(z3b*z3b));
|
||||
float tb0 = fmaf(alpha, mb0, beta);
|
||||
float tb1 = fmaf(alpha, mb1, beta);
|
||||
float tb2 = fmaf(alpha, mb2, beta);
|
||||
float tb3 = fmaf(alpha, mb3, beta);
|
||||
float gb0 = __fdividef(1.0f, 1.0f + __expf(-tb0));
|
||||
float gb1 = __fdividef(1.0f, 1.0f + __expf(-tb1));
|
||||
float gb2 = __fdividef(1.0f, 1.0f + __expf(-tb2));
|
||||
float gb3 = __fdividef(1.0f, 1.0f + __expf(-tb3));
|
||||
yv2.x = xv2.x * gb0;
|
||||
yv2.y = xv2.y * gb1;
|
||||
yv2.z = xv2.z * gb2;
|
||||
yv2.w = xv2.w * gb3;
|
||||
reinterpret_cast<float4*>(yr + base2)[0] = yv2;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
torch::Tensor gaussian_affine_gate_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, double alpha, double 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 = (float)alpha;
|
||||
float be = (float)beta;
|
||||
int block = 256;
|
||||
int elements_per_thread = 4;
|
||||
int elements_per_block = block * elements_per_thread;
|
||||
int gy = (D + elements_per_block - 1) / elements_per_block;
|
||||
dim3 grid(B, gy);
|
||||
gaussian_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;
|
||||
}
|
||||
"""
|
||||
|
||||
cpp_source = """
|
||||
#include <torch/extension.h>
|
||||
torch::Tensor gaussian_affine_gate_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, double alpha, double beta);
|
||||
"""
|
||||
|
||||
ops = load_inline(
|
||||
name="gaussian_affine_gate",
|
||||
cpp_sources=cpp_source,
|
||||
cuda_sources=source,
|
||||
functions=["gaussian_affine_gate_cuda"],
|
||||
extra_cflags=["-O3","-std=c++17"],
|
||||
extra_cuda_cflags=["-O3","--use_fast_math","-std=c++17","-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.alpha = float(alpha)
|
||||
self.beta = float(beta)
|
||||
|
||||
def forward(self, x):
|
||||
return self.ops.gaussian_affine_gate_cuda(x, self.scale, self.bias, self.alpha, self.beta)
|
||||
|
|
@ -1,10 +0,0 @@
|
|||
Operator: Gaussian-Affine-Gate (Fused CUDA Kernel)
|
||||
|
||||
Definition
|
||||
- z = x * scale + bias
|
||||
- m = exp(-z^2)
|
||||
- g = sigmoid(alpha * m + beta)
|
||||
- y = x * g
|
||||
|
||||
Goal
|
||||
- Elementwise fusion with fast exp; target ≥1.30x speedup.
|
||||
|
|
@ -1,50 +0,0 @@
|
|||
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 Gaussian-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()
|
||||
|
|
@ -1,28 +0,0 @@
|
|||
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
|
||||
m = torch.exp(-z*z)
|
||||
g = torch.sigmoid(self.alpha * m + 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]
|
||||
|
|
@ -1,124 +0,0 @@
|
|||
import torch
|
||||
from torch.utils.cpp_extension import load_inline
|
||||
|
||||
source = """
|
||||
#include <torch/extension.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <vector_types.h>
|
||||
|
||||
__device__ __forceinline__ float celuf(float x){
|
||||
return x >= 0.0f ? x : (__expf(x) - 1.0f);
|
||||
}
|
||||
|
||||
__global__ __launch_bounds__(256) void celu_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){
|
||||
int row = blockIdx.x;
|
||||
int tid = threadIdx.x;
|
||||
int col_block = blockIdx.y;
|
||||
int stride4 = blockDim.x * 4;
|
||||
int col_idx = (col_block * stride4) + tid * 4;
|
||||
const float* xr = x + row * D;
|
||||
float* yr = y + row * D;
|
||||
for(int base = col_idx; base < D; base += stride4 * 2){
|
||||
if (base < D){
|
||||
float4 xv = reinterpret_cast<const float4*>(xr + base)[0];
|
||||
float4 sv = reinterpret_cast<const float4*>(scale + base)[0];
|
||||
float4 bv = reinterpret_cast<const float4*>(bias + base)[0];
|
||||
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 m0 = celuf(z0);
|
||||
float m1 = celuf(z1);
|
||||
float m2 = celuf(z2);
|
||||
float m3 = celuf(z3);
|
||||
float t0 = fmaf(alpha, m0, beta);
|
||||
float t1 = fmaf(alpha, m1, beta);
|
||||
float t2 = fmaf(alpha, m2, beta);
|
||||
float t3 = fmaf(alpha, m3, 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 + base)[0] = yv;
|
||||
}
|
||||
int base2 = base + stride4;
|
||||
if (base2 < D){
|
||||
float4 xv2 = reinterpret_cast<const float4*>(xr + base2)[0];
|
||||
float4 sv2 = reinterpret_cast<const float4*>(scale + base2)[0];
|
||||
float4 bv2 = reinterpret_cast<const float4*>(bias + base2)[0];
|
||||
float4 yv2;
|
||||
float z0b = fmaf(xv2.x, sv2.x, bv2.x);
|
||||
float z1b = fmaf(xv2.y, sv2.y, bv2.y);
|
||||
float z2b = fmaf(xv2.z, sv2.z, bv2.z);
|
||||
float z3b = fmaf(xv2.w, sv2.w, bv2.w);
|
||||
float mb0 = celuf(z0b);
|
||||
float mb1 = celuf(z1b);
|
||||
float mb2 = celuf(z2b);
|
||||
float mb3 = celuf(z3b);
|
||||
float tb0 = fmaf(alpha, mb0, beta);
|
||||
float tb1 = fmaf(alpha, mb1, beta);
|
||||
float tb2 = fmaf(alpha, mb2, beta);
|
||||
float tb3 = fmaf(alpha, mb3, beta);
|
||||
float gb0 = __fdividef(1.0f, 1.0f + __expf(-tb0));
|
||||
float gb1 = __fdividef(1.0f, 1.0f + __expf(-tb1));
|
||||
float gb2 = __fdividef(1.0f, 1.0f + __expf(-tb2));
|
||||
float gb3 = __fdividef(1.0f, 1.0f + __expf(-tb3));
|
||||
yv2.x = xv2.x * gb0;
|
||||
yv2.y = xv2.y * gb1;
|
||||
yv2.z = xv2.z * gb2;
|
||||
yv2.w = xv2.w * gb3;
|
||||
reinterpret_cast<float4*>(yr + base2)[0] = yv2;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
torch::Tensor celu_affine_gate_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, double alpha, double 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 = (float)alpha;
|
||||
float be = (float)beta;
|
||||
int block = 256;
|
||||
int elements_per_thread = 4;
|
||||
int elements_per_block = block * elements_per_thread;
|
||||
int gy = (D + elements_per_block - 1) / elements_per_block;
|
||||
dim3 grid(B, gy);
|
||||
celu_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;
|
||||
}
|
||||
"""
|
||||
|
||||
cpp_source = """
|
||||
#include <torch/extension.h>
|
||||
torch::Tensor celu_affine_gate_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, double alpha, double beta);
|
||||
"""
|
||||
|
||||
ops = load_inline(
|
||||
name="celu_affine_gate",
|
||||
cpp_sources=cpp_source,
|
||||
cuda_sources=source,
|
||||
functions=["celu_affine_gate_cuda"],
|
||||
extra_cflags=["-O3","-std=c++17"],
|
||||
extra_cuda_cflags=["-O3","--use_fast_math","-std=c++17","-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.alpha = float(alpha)
|
||||
self.beta = float(beta)
|
||||
|
||||
def forward(self, x):
|
||||
return self.ops.celu_affine_gate_cuda(x, self.scale, self.bias, self.alpha, self.beta)
|
||||
|
|
@ -1,10 +0,0 @@
|
|||
Operator: CELU-Affine-Gate (Fused CUDA Kernel)
|
||||
|
||||
Definition
|
||||
- z = x * scale + bias
|
||||
- m = CELU(z, alpha=1)
|
||||
- g = sigmoid(alpha_g * m + beta)
|
||||
- y = x * g
|
||||
|
||||
Goal
|
||||
- Elementwise fusion with exp; target ≥1.30x speedup.
|
||||
|
|
@ -1,50 +0,0 @@
|
|||
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 CELU-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()
|
||||
|
|
@ -1,29 +0,0 @@
|
|||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
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
|
||||
m = F.celu(z, alpha=1.0)
|
||||
g = torch.sigmoid(self.alpha * m + 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]
|
||||
|
|
@ -1,133 +0,0 @@
|
|||
import torch
|
||||
from torch.utils.cpp_extension import load_inline
|
||||
|
||||
source = """
|
||||
#include <torch/extension.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <vector_types.h>
|
||||
|
||||
__device__ __forceinline__ float softplusf(float x){
|
||||
float ax = fabsf(x);
|
||||
return fmaxf(x, 0.0f) + __logf(1.0f + __expf(-ax));
|
||||
}
|
||||
|
||||
__global__ __launch_bounds__(256) void softplus3_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){
|
||||
int row = blockIdx.x;
|
||||
int tid = threadIdx.x;
|
||||
int col_block = blockIdx.y;
|
||||
int stride4 = blockDim.x * 4;
|
||||
int col_idx = (col_block * stride4) + tid * 4;
|
||||
const float* xr = x + row * D;
|
||||
float* yr = y + row * D;
|
||||
for(int base = col_idx; base < D; base += stride4 * 2){
|
||||
if (base < D){
|
||||
float4 xv = reinterpret_cast<const float4*>(xr + base)[0];
|
||||
float4 sv = reinterpret_cast<const float4*>(scale + base)[0];
|
||||
float4 bv = reinterpret_cast<const float4*>(bias + base)[0];
|
||||
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 s0 = softplusf(z0);
|
||||
float s1 = softplusf(z1);
|
||||
float s2 = softplusf(z2);
|
||||
float s3 = softplusf(z3);
|
||||
float m0 = s0 * s0 * s0;
|
||||
float m1 = s1 * s1 * s1;
|
||||
float m2 = s2 * s2 * s2;
|
||||
float m3 = s3 * s3 * s3;
|
||||
float t0 = fmaf(alpha, m0, beta);
|
||||
float t1 = fmaf(alpha, m1, beta);
|
||||
float t2 = fmaf(alpha, m2, beta);
|
||||
float t3 = fmaf(alpha, m3, 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 + base)[0] = yv;
|
||||
}
|
||||
int base2 = base + stride4;
|
||||
if (base2 < D){
|
||||
float4 xv2 = reinterpret_cast<const float4*>(xr + base2)[0];
|
||||
float4 sv2 = reinterpret_cast<const float4*>(scale + base2)[0];
|
||||
float4 bv2 = reinterpret_cast<const float4*>(bias + base2)[0];
|
||||
float4 yv2;
|
||||
float z0b = fmaf(xv2.x, sv2.x, bv2.x);
|
||||
float z1b = fmaf(xv2.y, sv2.y, bv2.y);
|
||||
float z2b = fmaf(xv2.z, sv2.z, bv2.z);
|
||||
float z3b = fmaf(xv2.w, sv2.w, bv2.w);
|
||||
float sb0 = softplusf(z0b);
|
||||
float sb1 = softplusf(z1b);
|
||||
float sb2 = softplusf(z2b);
|
||||
float sb3 = softplusf(z3b);
|
||||
float mb0 = sb0 * sb0 * sb0;
|
||||
float mb1 = sb1 * sb1 * sb1;
|
||||
float mb2 = sb2 * sb2 * sb2;
|
||||
float mb3 = sb3 * sb3 * sb3;
|
||||
float tb0 = fmaf(alpha, mb0, beta);
|
||||
float tb1 = fmaf(alpha, mb1, beta);
|
||||
float tb2 = fmaf(alpha, mb2, beta);
|
||||
float tb3 = fmaf(alpha, mb3, beta);
|
||||
float gb0 = __fdividef(1.0f, 1.0f + __expf(-tb0));
|
||||
float gb1 = __fdividef(1.0f, 1.0f + __expf(-tb1));
|
||||
float gb2 = __fdividef(1.0f, 1.0f + __expf(-tb2));
|
||||
float gb3 = __fdividef(1.0f, 1.0f + __expf(-tb3));
|
||||
yv2.x = xv2.x * gb0;
|
||||
yv2.y = xv2.y * gb1;
|
||||
yv2.z = xv2.z * gb2;
|
||||
yv2.w = xv2.w * gb3;
|
||||
reinterpret_cast<float4*>(yr + base2)[0] = yv2;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
torch::Tensor softplus3_affine_gate_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, double alpha, double 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 = (float)alpha;
|
||||
float be = (float)beta;
|
||||
int block = 256;
|
||||
int elements_per_thread = 4;
|
||||
int elements_per_block = block * elements_per_thread;
|
||||
int gy = (D + elements_per_block - 1) / elements_per_block;
|
||||
dim3 grid(B, gy);
|
||||
softplus3_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;
|
||||
}
|
||||
"""
|
||||
|
||||
cpp_source = """
|
||||
#include <torch/extension.h>
|
||||
torch::Tensor softplus3_affine_gate_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, double alpha, double beta);
|
||||
"""
|
||||
|
||||
ops = load_inline(
|
||||
name="softplus3_affine_gate",
|
||||
cpp_sources=cpp_source,
|
||||
cuda_sources=source,
|
||||
functions=["softplus3_affine_gate_cuda"],
|
||||
extra_cflags=["-O3","-std=c++17"],
|
||||
extra_cuda_cflags=["-O3","--use_fast_math","-std=c++17","-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.alpha = float(alpha)
|
||||
self.beta = float(beta)
|
||||
|
||||
def forward(self, x):
|
||||
return self.ops.softplus3_affine_gate_cuda(x, self.scale, self.bias, self.alpha, self.beta)
|
||||
|
|
@ -1,10 +0,0 @@
|
|||
Operator: Softplus^3-Affine-Gate (Fused CUDA Kernel)
|
||||
|
||||
Definition
|
||||
- z = x * scale + bias
|
||||
- m = softplus(z)^3
|
||||
- g = sigmoid(alpha * m + beta)
|
||||
- y = x * g
|
||||
|
||||
Goal
|
||||
- Elementwise fusion with fast softplus; target ≥1.30x speedup.
|
||||
|
|
@ -1,50 +0,0 @@
|
|||
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 Softplus^3-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()
|
||||
|
|
@ -1,30 +0,0 @@
|
|||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
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
|
||||
m = F.softplus(z)
|
||||
m = m * m * m
|
||||
g = torch.sigmoid(self.alpha * m + 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]
|
||||
|
|
@ -1,127 +0,0 @@
|
|||
import torch
|
||||
from torch.utils.cpp_extension import load_inline
|
||||
|
||||
source = """
|
||||
#include <torch/extension.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <vector_types.h>
|
||||
|
||||
__device__ __forceinline__ float softshrinkf(float x, float lambda_){
|
||||
if (x > lambda_) return x - lambda_;
|
||||
if (x < -lambda_) return x + lambda_;
|
||||
return 0.0f;
|
||||
}
|
||||
|
||||
__global__ __launch_bounds__(256) void softshrink_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){
|
||||
int row = blockIdx.x;
|
||||
int tid = threadIdx.x;
|
||||
int col_block = blockIdx.y;
|
||||
int stride4 = blockDim.x * 4;
|
||||
int col_idx = (col_block * stride4) + tid * 4;
|
||||
const float* xr = x + row * D;
|
||||
float* yr = y + row * D;
|
||||
const float lambda_ = 0.5f;
|
||||
for(int base = col_idx; base < D; base += stride4 * 2){
|
||||
if (base < D){
|
||||
float4 xv = reinterpret_cast<const float4*>(xr + base)[0];
|
||||
float4 sv = reinterpret_cast<const float4*>(scale + base)[0];
|
||||
float4 bv = reinterpret_cast<const float4*>(bias + base)[0];
|
||||
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 m0 = softshrinkf(z0, lambda_);
|
||||
float m1 = softshrinkf(z1, lambda_);
|
||||
float m2 = softshrinkf(z2, lambda_);
|
||||
float m3 = softshrinkf(z3, lambda_);
|
||||
float t0 = fmaf(alpha, m0, beta);
|
||||
float t1 = fmaf(alpha, m1, beta);
|
||||
float t2 = fmaf(alpha, m2, beta);
|
||||
float t3 = fmaf(alpha, m3, 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 + base)[0] = yv;
|
||||
}
|
||||
int base2 = base + stride4;
|
||||
if (base2 < D){
|
||||
float4 xv2 = reinterpret_cast<const float4*>(xr + base2)[0];
|
||||
float4 sv2 = reinterpret_cast<const float4*>(scale + base2)[0];
|
||||
float4 bv2 = reinterpret_cast<const float4*>(bias + base2)[0];
|
||||
float4 yv2;
|
||||
float z0b = fmaf(xv2.x, sv2.x, bv2.x);
|
||||
float z1b = fmaf(xv2.y, sv2.y, bv2.y);
|
||||
float z2b = fmaf(xv2.z, sv2.z, bv2.z);
|
||||
float z3b = fmaf(xv2.w, sv2.w, bv2.w);
|
||||
float mb0 = softshrinkf(z0b, lambda_);
|
||||
float mb1 = softshrinkf(z1b, lambda_);
|
||||
float mb2 = softshrinkf(z2b, lambda_);
|
||||
float mb3 = softshrinkf(z3b, lambda_);
|
||||
float tb0 = fmaf(alpha, mb0, beta);
|
||||
float tb1 = fmaf(alpha, mb1, beta);
|
||||
float tb2 = fmaf(alpha, mb2, beta);
|
||||
float tb3 = fmaf(alpha, mb3, beta);
|
||||
float gb0 = __fdividef(1.0f, 1.0f + __expf(-tb0));
|
||||
float gb1 = __fdividef(1.0f, 1.0f + __expf(-tb1));
|
||||
float gb2 = __fdividef(1.0f, 1.0f + __expf(-tb2));
|
||||
float gb3 = __fdividef(1.0f, 1.0f + __expf(-tb3));
|
||||
yv2.x = xv2.x * gb0;
|
||||
yv2.y = xv2.y * gb1;
|
||||
yv2.z = xv2.z * gb2;
|
||||
yv2.w = xv2.w * gb3;
|
||||
reinterpret_cast<float4*>(yr + base2)[0] = yv2;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
torch::Tensor softshrink_affine_gate_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, double alpha, double 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 = (float)alpha;
|
||||
float be = (float)beta;
|
||||
int block = 256;
|
||||
int elements_per_thread = 4;
|
||||
int elements_per_block = block * elements_per_thread;
|
||||
int gy = (D + elements_per_block - 1) / elements_per_block;
|
||||
dim3 grid(B, gy);
|
||||
softshrink_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;
|
||||
}
|
||||
"""
|
||||
|
||||
cpp_source = """
|
||||
#include <torch/extension.h>
|
||||
torch::Tensor softshrink_affine_gate_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, double alpha, double beta);
|
||||
"""
|
||||
|
||||
ops = load_inline(
|
||||
name="softshrink_affine_gate",
|
||||
cpp_sources=cpp_source,
|
||||
cuda_sources=source,
|
||||
functions=["softshrink_affine_gate_cuda"],
|
||||
extra_cflags=["-O3","-std=c++17"],
|
||||
extra_cuda_cflags=["-O3","--use_fast_math","-std=c++17","-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.alpha = float(alpha)
|
||||
self.beta = float(beta)
|
||||
|
||||
def forward(self, x):
|
||||
return self.ops.softshrink_affine_gate_cuda(x, self.scale, self.bias, self.alpha, self.beta)
|
||||
|
|
@ -1,10 +0,0 @@
|
|||
Operator: SoftShrink-Affine-Gate (Fused CUDA Kernel)
|
||||
|
||||
Definition
|
||||
- z = x * scale + bias
|
||||
- m = SoftShrink(z, lambda=0.5)
|
||||
- g = sigmoid(alpha * m + beta)
|
||||
- y = x * g
|
||||
|
||||
Goal
|
||||
- Fuse ops to reduce bandwidth and kernel launches; target ≥1.30x speedup.
|
||||
|
|
@ -1,50 +0,0 @@
|
|||
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 SoftShrink-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()
|
||||
|
|
@ -1,29 +0,0 @@
|
|||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
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
|
||||
m = F.softshrink(z, lambd=0.5)
|
||||
g = torch.sigmoid(self.alpha * m + 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]
|
||||
|
|
@ -1,125 +0,0 @@
|
|||
import torch
|
||||
from torch.utils.cpp_extension import load_inline
|
||||
|
||||
source = """
|
||||
#include <torch/extension.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <vector_types.h>
|
||||
|
||||
__device__ __forceinline__ float hardshrinkf(float x, float lambda_){
|
||||
return (x > lambda_ || x < -lambda_) ? x : 0.0f;
|
||||
}
|
||||
|
||||
__global__ __launch_bounds__(256) void hardshrink_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){
|
||||
int row = blockIdx.x;
|
||||
int tid = threadIdx.x;
|
||||
int col_block = blockIdx.y;
|
||||
int stride4 = blockDim.x * 4;
|
||||
int col_idx = (col_block * stride4) + tid * 4;
|
||||
const float* xr = x + row * D;
|
||||
float* yr = y + row * D;
|
||||
const float lambda_ = 0.5f;
|
||||
for(int base = col_idx; base < D; base += stride4 * 2){
|
||||
if (base < D){
|
||||
float4 xv = reinterpret_cast<const float4*>(xr + base)[0];
|
||||
float4 sv = reinterpret_cast<const float4*>(scale + base)[0];
|
||||
float4 bv = reinterpret_cast<const float4*>(bias + base)[0];
|
||||
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 m0 = hardshrinkf(z0, lambda_);
|
||||
float m1 = hardshrinkf(z1, lambda_);
|
||||
float m2 = hardshrinkf(z2, lambda_);
|
||||
float m3 = hardshrinkf(z3, lambda_);
|
||||
float t0 = fmaf(alpha, m0, beta);
|
||||
float t1 = fmaf(alpha, m1, beta);
|
||||
float t2 = fmaf(alpha, m2, beta);
|
||||
float t3 = fmaf(alpha, m3, 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 + base)[0] = yv;
|
||||
}
|
||||
int base2 = base + stride4;
|
||||
if (base2 < D){
|
||||
float4 xv2 = reinterpret_cast<const float4*>(xr + base2)[0];
|
||||
float4 sv2 = reinterpret_cast<const float4*>(scale + base2)[0];
|
||||
float4 bv2 = reinterpret_cast<const float4*>(bias + base2)[0];
|
||||
float4 yv2;
|
||||
float z0b = fmaf(xv2.x, sv2.x, bv2.x);
|
||||
float z1b = fmaf(xv2.y, sv2.y, bv2.y);
|
||||
float z2b = fmaf(xv2.z, sv2.z, bv2.z);
|
||||
float z3b = fmaf(xv2.w, sv2.w, bv2.w);
|
||||
float mb0 = hardshrinkf(z0b, lambda_);
|
||||
float mb1 = hardshrinkf(z1b, lambda_);
|
||||
float mb2 = hardshrinkf(z2b, lambda_);
|
||||
float mb3 = hardshrinkf(z3b, lambda_);
|
||||
float tb0 = fmaf(alpha, mb0, beta);
|
||||
float tb1 = fmaf(alpha, mb1, beta);
|
||||
float tb2 = fmaf(alpha, mb2, beta);
|
||||
float tb3 = fmaf(alpha, mb3, beta);
|
||||
float gb0 = __fdividef(1.0f, 1.0f + __expf(-tb0));
|
||||
float gb1 = __fdividef(1.0f, 1.0f + __expf(-tb1));
|
||||
float gb2 = __fdividef(1.0f, 1.0f + __expf(-tb2));
|
||||
float gb3 = __fdividef(1.0f, 1.0f + __expf(-tb3));
|
||||
yv2.x = xv2.x * gb0;
|
||||
yv2.y = xv2.y * gb1;
|
||||
yv2.z = xv2.z * gb2;
|
||||
yv2.w = xv2.w * gb3;
|
||||
reinterpret_cast<float4*>(yr + base2)[0] = yv2;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
torch::Tensor hardshrink_affine_gate_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, double alpha, double 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 = (float)alpha;
|
||||
float be = (float)beta;
|
||||
int block = 256;
|
||||
int elements_per_thread = 4;
|
||||
int elements_per_block = block * elements_per_thread;
|
||||
int gy = (D + elements_per_block - 1) / elements_per_block;
|
||||
dim3 grid(B, gy);
|
||||
hardshrink_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;
|
||||
}
|
||||
"""
|
||||
|
||||
cpp_source = """
|
||||
#include <torch/extension.h>
|
||||
torch::Tensor hardshrink_affine_gate_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, double alpha, double beta);
|
||||
"""
|
||||
|
||||
ops = load_inline(
|
||||
name="hardshrink_affine_gate",
|
||||
cpp_sources=cpp_source,
|
||||
cuda_sources=source,
|
||||
functions=["hardshrink_affine_gate_cuda"],
|
||||
extra_cflags=["-O3","-std=c++17"],
|
||||
extra_cuda_cflags=["-O3","--use_fast_math","-std=c++17","-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.alpha = float(alpha)
|
||||
self.beta = float(beta)
|
||||
|
||||
def forward(self, x):
|
||||
return self.ops.hardshrink_affine_gate_cuda(x, self.scale, self.bias, self.alpha, self.beta)
|
||||
|
|
@ -1,10 +0,0 @@
|
|||
Operator: HardShrink-Affine-Gate (Fused CUDA Kernel)
|
||||
|
||||
Definition
|
||||
- z = x * scale + bias
|
||||
- m = HardShrink(z, lambda=0.5)
|
||||
- g = sigmoid(alpha * m + beta)
|
||||
- y = x * g
|
||||
|
||||
Goal
|
||||
- Fuse ops to reduce memory passes and launches; target ≥1.30x speedup.
|
||||
|
|
@ -1,50 +0,0 @@
|
|||
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 HardShrink-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()
|
||||
|
|
@ -1,29 +0,0 @@
|
|||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
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
|
||||
m = F.hardshrink(z, lambd=0.5)
|
||||
g = torch.sigmoid(self.alpha * m + 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]
|
||||
|
|
@ -1,120 +0,0 @@
|
|||
import torch
|
||||
from torch.utils.cpp_extension import load_inline
|
||||
|
||||
source = """
|
||||
#include <torch/extension.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <vector_types.h>
|
||||
|
||||
__global__ __launch_bounds__(256) void laplacian_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){
|
||||
int row = blockIdx.x;
|
||||
int tid = threadIdx.x;
|
||||
int col_block = blockIdx.y;
|
||||
int stride4 = blockDim.x * 4;
|
||||
int col_idx = (col_block * stride4) + tid * 4;
|
||||
const float* xr = x + row * D;
|
||||
float* yr = y + row * D;
|
||||
for(int base = col_idx; base < D; base += stride4 * 2){
|
||||
if (base < D){
|
||||
float4 xv = reinterpret_cast<const float4*>(xr + base)[0];
|
||||
float4 sv = reinterpret_cast<const float4*>(scale + base)[0];
|
||||
float4 bv = reinterpret_cast<const float4*>(bias + base)[0];
|
||||
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 m0 = __expf(-fabsf(z0));
|
||||
float m1 = __expf(-fabsf(z1));
|
||||
float m2 = __expf(-fabsf(z2));
|
||||
float m3 = __expf(-fabsf(z3));
|
||||
float t0 = fmaf(alpha, m0, beta);
|
||||
float t1 = fmaf(alpha, m1, beta);
|
||||
float t2 = fmaf(alpha, m2, beta);
|
||||
float t3 = fmaf(alpha, m3, 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 + base)[0] = yv;
|
||||
}
|
||||
int base2 = base + stride4;
|
||||
if (base2 < D){
|
||||
float4 xv2 = reinterpret_cast<const float4*>(xr + base2)[0];
|
||||
float4 sv2 = reinterpret_cast<const float4*>(scale + base2)[0];
|
||||
float4 bv2 = reinterpret_cast<const float4*>(bias + base2)[0];
|
||||
float4 yv2;
|
||||
float z0b = fmaf(xv2.x, sv2.x, bv2.x);
|
||||
float z1b = fmaf(xv2.y, sv2.y, bv2.y);
|
||||
float z2b = fmaf(xv2.z, sv2.z, bv2.z);
|
||||
float z3b = fmaf(xv2.w, sv2.w, bv2.w);
|
||||
float mb0 = __expf(-fabsf(z0b));
|
||||
float mb1 = __expf(-fabsf(z1b));
|
||||
float mb2 = __expf(-fabsf(z2b));
|
||||
float mb3 = __expf(-fabsf(z3b));
|
||||
float tb0 = fmaf(alpha, mb0, beta);
|
||||
float tb1 = fmaf(alpha, mb1, beta);
|
||||
float tb2 = fmaf(alpha, mb2, beta);
|
||||
float tb3 = fmaf(alpha, mb3, beta);
|
||||
float gb0 = __fdividef(1.0f, 1.0f + __expf(-tb0));
|
||||
float gb1 = __fdividef(1.0f, 1.0f + __expf(-tb1));
|
||||
float gb2 = __fdividef(1.0f, 1.0f + __expf(-tb2));
|
||||
float gb3 = __fdividef(1.0f, 1.0f + __expf(-tb3));
|
||||
yv2.x = xv2.x * gb0;
|
||||
yv2.y = xv2.y * gb1;
|
||||
yv2.z = xv2.z * gb2;
|
||||
yv2.w = xv2.w * gb3;
|
||||
reinterpret_cast<float4*>(yr + base2)[0] = yv2;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
torch::Tensor laplacian_affine_gate_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, double alpha, double 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 = (float)alpha;
|
||||
float be = (float)beta;
|
||||
int block = 256;
|
||||
int elements_per_thread = 4;
|
||||
int elements_per_block = block * elements_per_thread;
|
||||
int gy = (D + elements_per_block - 1) / elements_per_block;
|
||||
dim3 grid(B, gy);
|
||||
laplacian_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;
|
||||
}
|
||||
"""
|
||||
|
||||
cpp_source = """
|
||||
#include <torch/extension.h>
|
||||
torch::Tensor laplacian_affine_gate_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, double alpha, double beta);
|
||||
"""
|
||||
|
||||
ops = load_inline(
|
||||
name="laplacian_affine_gate",
|
||||
cpp_sources=cpp_source,
|
||||
cuda_sources=source,
|
||||
functions=["laplacian_affine_gate_cuda"],
|
||||
extra_cflags=["-O3","-std=c++17"],
|
||||
extra_cuda_cflags=["-O3","--use_fast_math","-std=c++17","-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.alpha = float(alpha)
|
||||
self.beta = float(beta)
|
||||
|
||||
def forward(self, x):
|
||||
return self.ops.laplacian_affine_gate_cuda(x, self.scale, self.bias, self.alpha, self.beta)
|
||||
|
|
@ -1,10 +0,0 @@
|
|||
Operator: Laplacian-Affine-Gate (Fused CUDA Kernel)
|
||||
|
||||
Definition
|
||||
- z = x * scale + bias
|
||||
- m = exp(-|z|)
|
||||
- g = sigmoid(alpha * m + beta)
|
||||
- y = x * g
|
||||
|
||||
Goal
|
||||
- Elementwise fusion; target ≥1.30x speedup.
|
||||
|
|
@ -1,50 +0,0 @@
|
|||
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 Laplacian-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()
|
||||
|
|
@ -1,28 +0,0 @@
|
|||
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
|
||||
m = torch.exp(-torch.abs(z))
|
||||
g = torch.sigmoid(self.alpha * m + 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]
|
||||
|
|
@ -1,129 +0,0 @@
|
|||
import torch
|
||||
from torch.utils.cpp_extension import load_inline
|
||||
|
||||
source = """
|
||||
#include <torch/extension.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <vector_types.h>
|
||||
|
||||
__global__ __launch_bounds__(256) 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){
|
||||
int row = blockIdx.x;
|
||||
int tid = threadIdx.x;
|
||||
int col_block = blockIdx.y;
|
||||
int stride4 = blockDim.x * 4;
|
||||
int col_idx = (col_block * stride4) + tid * 4;
|
||||
const float* xr = x + row * D;
|
||||
float* yr = y + row * D;
|
||||
const float LOG2 = 0.6931471805599453094f;
|
||||
for(int base = col_idx; base < D; base += stride4 * 2){
|
||||
if (base < D){
|
||||
float4 xv = reinterpret_cast<const float4*>(xr + base)[0];
|
||||
float4 sv = reinterpret_cast<const float4*>(scale + base)[0];
|
||||
float4 bv = reinterpret_cast<const float4*>(bias + base)[0];
|
||||
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 a0 = fabsf(z0);
|
||||
float a1 = fabsf(z1);
|
||||
float a2 = fabsf(z2);
|
||||
float a3 = fabsf(z3);
|
||||
float m0 = a0 + __logf(1.0f + __expf(-2.0f * a0)) - LOG2;
|
||||
float m1 = a1 + __logf(1.0f + __expf(-2.0f * a1)) - LOG2;
|
||||
float m2 = a2 + __logf(1.0f + __expf(-2.0f * a2)) - LOG2;
|
||||
float m3 = a3 + __logf(1.0f + __expf(-2.0f * a3)) - LOG2;
|
||||
float t0 = fmaf(alpha, m0, beta);
|
||||
float t1 = fmaf(alpha, m1, beta);
|
||||
float t2 = fmaf(alpha, m2, beta);
|
||||
float t3 = fmaf(alpha, m3, 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 + base)[0] = yv;
|
||||
}
|
||||
int base2 = base + stride4;
|
||||
if (base2 < D){
|
||||
float4 xv2 = reinterpret_cast<const float4*>(xr + base2)[0];
|
||||
float4 sv2 = reinterpret_cast<const float4*>(scale + base2)[0];
|
||||
float4 bv2 = reinterpret_cast<const float4*>(bias + base2)[0];
|
||||
float4 yv2;
|
||||
float z0b = fmaf(xv2.x, sv2.x, bv2.x);
|
||||
float z1b = fmaf(xv2.y, sv2.y, bv2.y);
|
||||
float z2b = fmaf(xv2.z, sv2.z, bv2.z);
|
||||
float z3b = fmaf(xv2.w, sv2.w, bv2.w);
|
||||
float ab0 = fabsf(z0b);
|
||||
float ab1 = fabsf(z1b);
|
||||
float ab2 = fabsf(z2b);
|
||||
float ab3 = fabsf(z3b);
|
||||
float mb0 = ab0 + __logf(1.0f + __expf(-2.0f * ab0)) - LOG2;
|
||||
float mb1 = ab1 + __logf(1.0f + __expf(-2.0f * ab1)) - LOG2;
|
||||
float mb2 = ab2 + __logf(1.0f + __expf(-2.0f * ab2)) - LOG2;
|
||||
float mb3 = ab3 + __logf(1.0f + __expf(-2.0f * ab3)) - LOG2;
|
||||
float tb0 = fmaf(alpha, mb0, beta);
|
||||
float tb1 = fmaf(alpha, mb1, beta);
|
||||
float tb2 = fmaf(alpha, mb2, beta);
|
||||
float tb3 = fmaf(alpha, mb3, beta);
|
||||
float gb0 = __fdividef(1.0f, 1.0f + __expf(-tb0));
|
||||
float gb1 = __fdividef(1.0f, 1.0f + __expf(-tb1));
|
||||
float gb2 = __fdividef(1.0f, 1.0f + __expf(-tb2));
|
||||
float gb3 = __fdividef(1.0f, 1.0f + __expf(-tb3));
|
||||
yv2.x = xv2.x * gb0;
|
||||
yv2.y = xv2.y * gb1;
|
||||
yv2.z = xv2.z * gb2;
|
||||
yv2.w = xv2.w * gb3;
|
||||
reinterpret_cast<float4*>(yr + base2)[0] = yv2;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
torch::Tensor logcosh_affine_gate_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, double alpha, double 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 = (float)alpha;
|
||||
float be = (float)beta;
|
||||
int block = 256;
|
||||
int elements_per_thread = 4;
|
||||
int elements_per_block = block * elements_per_thread;
|
||||
int gy = (D + elements_per_block - 1) / elements_per_block;
|
||||
dim3 grid(B, gy);
|
||||
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;
|
||||
}
|
||||
"""
|
||||
|
||||
cpp_source = """
|
||||
#include <torch/extension.h>
|
||||
torch::Tensor logcosh_affine_gate_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, double alpha, double beta);
|
||||
"""
|
||||
|
||||
ops = load_inline(
|
||||
name="logcosh_affine_gate",
|
||||
cpp_sources=cpp_source,
|
||||
cuda_sources=source,
|
||||
functions=["logcosh_affine_gate_cuda"],
|
||||
extra_cflags=["-O3","-std=c++17"],
|
||||
extra_cuda_cflags=["-O3","--use_fast_math","-std=c++17","-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.alpha = float(alpha)
|
||||
self.beta = float(beta)
|
||||
|
||||
def forward(self, x):
|
||||
return self.ops.logcosh_affine_gate_cuda(x, self.scale, self.bias, self.alpha, self.beta)
|
||||
|
|
@ -1,10 +0,0 @@
|
|||
Operator: LogCosh-Affine-Gate (Fused CUDA Kernel)
|
||||
|
||||
Definition
|
||||
- z = x * scale + bias
|
||||
- m = log(cosh(z))
|
||||
- g = sigmoid(alpha * m + beta)
|
||||
- y = x * g
|
||||
|
||||
Goal
|
||||
- Numerically stable logcosh fusion; target ≥1.30x speedup.
|
||||
|
|
@ -1,50 +0,0 @@
|
|||
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()
|
||||
|
|
@ -1,29 +0,0 @@
|
|||
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
|
||||
az = torch.abs(z)
|
||||
m = az + torch.log1p(torch.exp(-2.0 * az)) - torch.log(torch.tensor(2.0))
|
||||
g = torch.sigmoid(self.alpha * m + 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]
|
||||
|
|
@ -1,133 +0,0 @@
|
|||
import torch
|
||||
from torch.utils.cpp_extension import load_inline
|
||||
|
||||
source = """
|
||||
#include <torch/extension.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <vector_types.h>
|
||||
|
||||
__device__ __forceinline__ float softplusf(float x){
|
||||
float ax = fabsf(x);
|
||||
return fmaxf(x, 0.0f) + __logf(1.0f + __expf(-ax));
|
||||
}
|
||||
|
||||
__global__ __launch_bounds__(256) void softplus_sqrt_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){
|
||||
int row = blockIdx.x;
|
||||
int tid = threadIdx.x;
|
||||
int col_block = blockIdx.y;
|
||||
int stride4 = blockDim.x * 4;
|
||||
int col_idx = (col_block * stride4) + tid * 4;
|
||||
const float* xr = x + row * D;
|
||||
float* yr = y + row * D;
|
||||
for(int base = col_idx; base < D; base += stride4 * 2){
|
||||
if (base < D){
|
||||
float4 xv = reinterpret_cast<const float4*>(xr + base)[0];
|
||||
float4 sv = reinterpret_cast<const float4*>(scale + base)[0];
|
||||
float4 bv = reinterpret_cast<const float4*>(bias + base)[0];
|
||||
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 s0 = softplusf(z0);
|
||||
float s1 = softplusf(z1);
|
||||
float s2 = softplusf(z2);
|
||||
float s3 = softplusf(z3);
|
||||
float m0 = sqrtf(s0);
|
||||
float m1 = sqrtf(s1);
|
||||
float m2 = sqrtf(s2);
|
||||
float m3 = sqrtf(s3);
|
||||
float t0 = fmaf(alpha, m0, beta);
|
||||
float t1 = fmaf(alpha, m1, beta);
|
||||
float t2 = fmaf(alpha, m2, beta);
|
||||
float t3 = fmaf(alpha, m3, 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 + base)[0] = yv;
|
||||
}
|
||||
int base2 = base + stride4;
|
||||
if (base2 < D){
|
||||
float4 xv2 = reinterpret_cast<const float4*>(xr + base2)[0];
|
||||
float4 sv2 = reinterpret_cast<const float4*>(scale + base2)[0];
|
||||
float4 bv2 = reinterpret_cast<const float4*>(bias + base2)[0];
|
||||
float4 yv2;
|
||||
float z0b = fmaf(xv2.x, sv2.x, bv2.x);
|
||||
float z1b = fmaf(xv2.y, sv2.y, bv2.y);
|
||||
float z2b = fmaf(xv2.z, sv2.z, bv2.z);
|
||||
float z3b = fmaf(xv2.w, sv2.w, bv2.w);
|
||||
float sb0 = softplusf(z0b);
|
||||
float sb1 = softplusf(z1b);
|
||||
float sb2 = softplusf(z2b);
|
||||
float sb3 = softplusf(z3b);
|
||||
float mb0 = sqrtf(sb0);
|
||||
float mb1 = sqrtf(sb1);
|
||||
float mb2 = sqrtf(sb2);
|
||||
float mb3 = sqrtf(sb3);
|
||||
float tb0 = fmaf(alpha, mb0, beta);
|
||||
float tb1 = fmaf(alpha, mb1, beta);
|
||||
float tb2 = fmaf(alpha, mb2, beta);
|
||||
float tb3 = fmaf(alpha, mb3, beta);
|
||||
float gb0 = __fdividef(1.0f, 1.0f + __expf(-tb0));
|
||||
float gb1 = __fdividef(1.0f, 1.0f + __expf(-tb1));
|
||||
float gb2 = __fdividef(1.0f, 1.0f + __expf(-tb2));
|
||||
float gb3 = __fdividef(1.0f, 1.0f + __expf(-tb3));
|
||||
yv2.x = xv2.x * gb0;
|
||||
yv2.y = xv2.y * gb1;
|
||||
yv2.z = xv2.z * gb2;
|
||||
yv2.w = xv2.w * gb3;
|
||||
reinterpret_cast<float4*>(yr + base2)[0] = yv2;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
torch::Tensor softplus_sqrt_affine_gate_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, double alpha, double 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 = (float)alpha;
|
||||
float be = (float)beta;
|
||||
int block = 256;
|
||||
int elements_per_thread = 4;
|
||||
int elements_per_block = block * elements_per_thread;
|
||||
int gy = (D + elements_per_block - 1) / elements_per_block;
|
||||
dim3 grid(B, gy);
|
||||
softplus_sqrt_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;
|
||||
}
|
||||
"""
|
||||
|
||||
cpp_source = """
|
||||
#include <torch/extension.h>
|
||||
torch::Tensor softplus_sqrt_affine_gate_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, double alpha, double beta);
|
||||
"""
|
||||
|
||||
ops = load_inline(
|
||||
name="softplus_sqrt_affine_gate",
|
||||
cpp_sources=cpp_source,
|
||||
cuda_sources=source,
|
||||
functions=["softplus_sqrt_affine_gate_cuda"],
|
||||
extra_cflags=["-O3","-std=c++17"],
|
||||
extra_cuda_cflags=["-O3","--use_fast_math","-std=c++17","-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.alpha = float(alpha)
|
||||
self.beta = float(beta)
|
||||
|
||||
def forward(self, x):
|
||||
return self.ops.softplus_sqrt_affine_gate_cuda(x, self.scale, self.bias, self.alpha, self.beta)
|
||||
|
|
@ -1,10 +0,0 @@
|
|||
Operator: SoftplusSqrt-Affine-Gate (Fused CUDA Kernel)
|
||||
|
||||
Definition
|
||||
- z = x * scale + bias
|
||||
- m = sqrt(softplus(z))
|
||||
- g = sigmoid(alpha * m + beta)
|
||||
- y = x * g
|
||||
|
||||
Goal
|
||||
- Fuse softplus and sqrt; target ≥1.30x speedup.
|
||||
|
|
@ -1,50 +0,0 @@
|
|||
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 SoftplusSqrt-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()
|
||||
|
|
@ -1,29 +0,0 @@
|
|||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
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
|
||||
m = torch.sqrt(F.softplus(z))
|
||||
g = torch.sigmoid(self.alpha * m + 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]
|
||||
|
|
@ -1,120 +0,0 @@
|
|||
import torch
|
||||
from torch.utils.cpp_extension import load_inline
|
||||
|
||||
source = """
|
||||
#include <torch/extension.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <vector_types.h>
|
||||
|
||||
__global__ __launch_bounds__(256) void arcsin_tanh_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){
|
||||
int row = blockIdx.x;
|
||||
int tid = threadIdx.x;
|
||||
int col_block = blockIdx.y;
|
||||
int stride4 = blockDim.x * 4;
|
||||
int col_idx = (col_block * stride4) + tid * 4;
|
||||
const float* xr = x + row * D;
|
||||
float* yr = y + row * D;
|
||||
for(int base = col_idx; base < D; base += stride4 * 2){
|
||||
if (base < D){
|
||||
float4 xv = reinterpret_cast<const float4*>(xr + base)[0];
|
||||
float4 sv = reinterpret_cast<const float4*>(scale + base)[0];
|
||||
float4 bv = reinterpret_cast<const float4*>(bias + base)[0];
|
||||
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 m0 = asinf(tanhf(z0));
|
||||
float m1 = asinf(tanhf(z1));
|
||||
float m2 = asinf(tanhf(z2));
|
||||
float m3 = asinf(tanhf(z3));
|
||||
float t0 = fmaf(alpha, m0, beta);
|
||||
float t1 = fmaf(alpha, m1, beta);
|
||||
float t2 = fmaf(alpha, m2, beta);
|
||||
float t3 = fmaf(alpha, m3, 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 + base)[0] = yv;
|
||||
}
|
||||
int base2 = base + stride4;
|
||||
if (base2 < D){
|
||||
float4 xv2 = reinterpret_cast<const float4*>(xr + base2)[0];
|
||||
float4 sv2 = reinterpret_cast<const float4*>(scale + base2)[0];
|
||||
float4 bv2 = reinterpret_cast<const float4*>(bias + base2)[0];
|
||||
float4 yv2;
|
||||
float z0b = fmaf(xv2.x, sv2.x, bv2.x);
|
||||
float z1b = fmaf(xv2.y, sv2.y, bv2.y);
|
||||
float z2b = fmaf(xv2.z, sv2.z, bv2.z);
|
||||
float z3b = fmaf(xv2.w, sv2.w, bv2.w);
|
||||
float mb0 = asinf(tanhf(z0b));
|
||||
float mb1 = asinf(tanhf(z1b));
|
||||
float mb2 = asinf(tanhf(z2b));
|
||||
float mb3 = asinf(tanhf(z3b));
|
||||
float tb0 = fmaf(alpha, mb0, beta);
|
||||
float tb1 = fmaf(alpha, mb1, beta);
|
||||
float tb2 = fmaf(alpha, mb2, beta);
|
||||
float tb3 = fmaf(alpha, mb3, beta);
|
||||
float gb0 = __fdividef(1.0f, 1.0f + __expf(-tb0));
|
||||
float gb1 = __fdividef(1.0f, 1.0f + __expf(-tb1));
|
||||
float gb2 = __fdividef(1.0f, 1.0f + __expf(-tb2));
|
||||
float gb3 = __fdividef(1.0f, 1.0f + __expf(-tb3));
|
||||
yv2.x = xv2.x * gb0;
|
||||
yv2.y = xv2.y * gb1;
|
||||
yv2.z = xv2.z * gb2;
|
||||
yv2.w = xv2.w * gb3;
|
||||
reinterpret_cast<float4*>(yr + base2)[0] = yv2;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
torch::Tensor arcsin_tanh_affine_gate_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, double alpha, double 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 = (float)alpha;
|
||||
float be = (float)beta;
|
||||
int block = 256;
|
||||
int elements_per_thread = 4;
|
||||
int elements_per_block = block * elements_per_thread;
|
||||
int gy = (D + elements_per_block - 1) / elements_per_block;
|
||||
dim3 grid(B, gy);
|
||||
arcsin_tanh_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;
|
||||
}
|
||||
"""
|
||||
|
||||
cpp_source = """
|
||||
#include <torch/extension.h>
|
||||
torch::Tensor arcsin_tanh_affine_gate_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, double alpha, double beta);
|
||||
"""
|
||||
|
||||
ops = load_inline(
|
||||
name="arcsin_tanh_affine_gate",
|
||||
cpp_sources=cpp_source,
|
||||
cuda_sources=source,
|
||||
functions=["arcsin_tanh_affine_gate_cuda"],
|
||||
extra_cflags=["-O3","-std=c++17"],
|
||||
extra_cuda_cflags=["-O3","--use_fast_math","-std=c++17","-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.alpha = float(alpha)
|
||||
self.beta = float(beta)
|
||||
|
||||
def forward(self, x):
|
||||
return self.ops.arcsin_tanh_affine_gate_cuda(x, self.scale, self.bias, self.alpha, self.beta)
|
||||
|
|
@ -1,10 +0,0 @@
|
|||
Operator: ArcSinTanh-Affine-Gate (Fused CUDA Kernel)
|
||||
|
||||
Definition
|
||||
- z = x * scale + bias
|
||||
- m = asin(tanh(z))
|
||||
- g = sigmoid(alpha * m + beta)
|
||||
- y = x * g
|
||||
|
||||
Goal
|
||||
- Domain-safe fusion; target ≥1.30x speedup.
|
||||
|
|
@ -1,50 +0,0 @@
|
|||
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 ArcSinTanh-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()
|
||||
|
|
@ -1,29 +0,0 @@
|
|||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
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
|
||||
m = torch.asin(torch.tanh(z))
|
||||
g = torch.sigmoid(self.alpha * m + 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]
|
||||
|
|
@ -1,51 +0,0 @@
|
|||
import torch
|
||||
from torch.utils.cpp_extension import load_inline
|
||||
|
||||
source = """
|
||||
#include <torch/extension.h>
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
__global__ void bias_silu_kernel(const float* x, const float* bias, float* y, int dim, long long total) {
|
||||
long long idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
long long stride = blockDim.x * gridDim.x;
|
||||
for (long long i = idx; i < total; i += stride) {
|
||||
int j = (int)(i % dim);
|
||||
float z = x[i] + bias[j];
|
||||
y[i] = z / (1.0f + expf(-z));
|
||||
}
|
||||
}
|
||||
|
||||
torch::Tensor bias_silu_cuda(torch::Tensor x, torch::Tensor bias) {
|
||||
auto x_contig = x.contiguous();
|
||||
auto b_contig = bias.contiguous();
|
||||
auto y = torch::empty_like(x_contig);
|
||||
long long total = x_contig.numel();
|
||||
int dim = (int)x_contig.size(-1);
|
||||
int block = 512;
|
||||
long long grid = (total + block - 1) / block;
|
||||
grid = grid > 65535 ? 65535 : grid;
|
||||
bias_silu_kernel<<<(int)grid, block>>>(x_contig.data_ptr<float>(), b_contig.data_ptr<float>(), y.data_ptr<float>(), dim, total);
|
||||
return y;
|
||||
}
|
||||
"""
|
||||
|
||||
cpp_source = """
|
||||
torch::Tensor bias_silu_cuda(torch::Tensor x, torch::Tensor bias);
|
||||
"""
|
||||
|
||||
ops = load_inline(
|
||||
name="bias_silu",
|
||||
cpp_sources=cpp_source,
|
||||
cuda_sources=source,
|
||||
functions=["bias_silu_cuda"],
|
||||
verbose=True
|
||||
)
|
||||
|
||||
class ModelNew(torch.nn.Module):
|
||||
def __init__(self, bias: torch.Tensor):
|
||||
super(ModelNew, self).__init__()
|
||||
self.ops = ops
|
||||
self.register_buffer("bias", bias)
|
||||
|
||||
def forward(self, x):
|
||||
return self.ops.bias_silu_cuda(x, self.bias)
|
||||
|
|
@ -1,5 +0,0 @@
|
|||
Bias+SiLU（Swish）融合：一次内核完成加偏置与 SiLU 激活，减少内核次数与显存往返。
|
||||
|
||||
torchcode.py：参考实现 `y = silu(x + bias)`。
|
||||
cudacode.py：`__global__ void bias_silu_kernel(...)` 完成融合计算。
|
||||
run_code.py：比较精度与性能（100 次迭代，`rtol=1e-03`）。
|
||||
|
|
@ -1,76 +0,0 @@
|
|||
###########################################################
|
||||
# 性能和精度验证程序
|
||||
###########################################################
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
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
|
||||
else:
|
||||
device = torch.device("cuda")
|
||||
|
||||
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)
|
||||
|
||||
precision_flag = torch.allclose(output_torch, output_cuda,rtol=1e-03)
|
||||
if precision_flag:
|
||||
print("✅ 精度对齐：两个模型的输出结果非常接近。")
|
||||
else:
|
||||
print("❌ 精度不一致！")
|
||||
diff = (output_torch - output_cuda).abs().max().item()
|
||||
print(f"最大绝对误差: {diff}")
|
||||
print(f"输出张量形状: torch={tuple(output_torch.shape)}, cuda={tuple(output_cuda.shape)}")
|
||||
print(f"数据类型: torch={output_torch.dtype}, cuda={output_cuda.dtype}")
|
||||
print(f"设备: torch={output_torch.device}, cuda={output_cuda.device}")
|
||||
|
||||
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 Bias+SiLU 平均执行时间: {torch_time:.6f} 秒")
|
||||
print(f"自定义 CUDA 融合内核 平均执行时间: {cuda_time:.6f} 秒")
|
||||
speedup = 0
|
||||
if cuda_time > 0:
|
||||
speedup = torch_time / cuda_time
|
||||
print(f"加速比 (Speedup): {speedup:.2f}x")
|
||||
else:
|
||||
print("CUDA 内核执行时间为0，无法计算加速比。")
|
||||
return precision_flag,speedup
|
||||
|
||||
if __name__ == "__main__":
|
||||
precision_flag,speedup = run_benchmark()
|
||||
|
|
@ -1,22 +0,0 @@
|
|||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
class Model(nn.Module):
|
||||
def __init__(self, bias: torch.Tensor):
|
||||
super(Model, self).__init__()
|
||||
self.register_buffer("bias", bias)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
return F.silu(x + self.bias)
|
||||
|
||||
batch_size = 16
|
||||
dim = 16384
|
||||
|
||||
def get_inputs():
|
||||
x = torch.randn(batch_size, dim)
|
||||
return [x]
|
||||
|
||||
def get_init_inputs():
|
||||
bias = torch.randn(dim)
|
||||
return [bias]
|
||||
|
|
@ -1,132 +0,0 @@
|
|||
import torch
|
||||
from torch.utils.cpp_extension import load_inline
|
||||
|
||||
source = """
|
||||
#include <torch/extension.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <vector_types.h>
|
||||
|
||||
__device__ __forceinline__ float sigmoidf(float x){
|
||||
return __fdividef(1.0f, 1.0f + __expf(-x));
|
||||
}
|
||||
|
||||
__global__ __launch_bounds__(256) void sigmoid_slope_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){
|
||||
int row = blockIdx.x;
|
||||
int tid = threadIdx.x;
|
||||
int col_block = blockIdx.y;
|
||||
int stride4 = blockDim.x * 4;
|
||||
int col_idx = (col_block * stride4) + tid * 4;
|
||||
const float* xr = x + row * D;
|
||||
float* yr = y + row * D;
|
||||
for(int base = col_idx; base < D; base += stride4 * 2){
|
||||
if (base < D){
|
||||
float4 xv = reinterpret_cast<const float4*>(xr + base)[0];
|
||||
float4 sv = reinterpret_cast<const float4*>(scale + base)[0];
|
||||
float4 bv = reinterpret_cast<const float4*>(bias + base)[0];
|
||||
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 s0 = sigmoidf(z0);
|
||||
float s1 = sigmoidf(z1);
|
||||
float s2 = sigmoidf(z2);
|
||||
float s3 = sigmoidf(z3);
|
||||
float m0 = s0 * (1.0f - s0);
|
||||
float m1 = s1 * (1.0f - s1);
|
||||
float m2 = s2 * (1.0f - s2);
|
||||
float m3 = s3 * (1.0f - s3);
|
||||
float t0 = fmaf(alpha, m0, beta);
|
||||
float t1 = fmaf(alpha, m1, beta);
|
||||
float t2 = fmaf(alpha, m2, beta);
|
||||
float t3 = fmaf(alpha, m3, beta);
|
||||
float g0 = sigmoidf(t0);
|
||||
float g1 = sigmoidf(t1);
|
||||
float g2 = sigmoidf(t2);
|
||||
float g3 = sigmoidf(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 + base)[0] = yv;
|
||||
}
|
||||
int base2 = base + stride4;
|
||||
if (base2 < D){
|
||||
float4 xv2 = reinterpret_cast<const float4*>(xr + base2)[0];
|
||||
float4 sv2 = reinterpret_cast<const float4*>(scale + base2)[0];
|
||||
float4 bv2 = reinterpret_cast<const float4*>(bias + base2)[0];
|
||||
float4 yv2;
|
||||
float z0b = fmaf(xv2.x, sv2.x, bv2.x);
|
||||
float z1b = fmaf(xv2.y, sv2.y, bv2.y);
|
||||
float z2b = fmaf(xv2.z, sv2.z, bv2.z);
|
||||
float z3b = fmaf(xv2.w, sv2.w, bv2.w);
|
||||
float sb0 = sigmoidf(z0b);
|
||||
float sb1 = sigmoidf(z1b);
|
||||
float sb2 = sigmoidf(z2b);
|
||||
float sb3 = sigmoidf(z3b);
|
||||
float mb0 = sb0 * (1.0f - sb0);
|
||||
float mb1 = sb1 * (1.0f - sb1);
|
||||
float mb2 = sb2 * (1.0f - sb2);
|
||||
float mb3 = sb3 * (1.0f - sb3);
|
||||
float tb0 = fmaf(alpha, mb0, beta);
|
||||
float tb1 = fmaf(alpha, mb1, beta);
|
||||
float tb2 = fmaf(alpha, mb2, beta);
|
||||
float tb3 = fmaf(alpha, mb3, beta);
|
||||
float gb0 = sigmoidf(tb0);
|
||||
float gb1 = sigmoidf(tb1);
|
||||
float gb2 = sigmoidf(tb2);
|
||||
float gb3 = sigmoidf(tb3);
|
||||
yv2.x = xv2.x * gb0;
|
||||
yv2.y = xv2.y * gb1;
|
||||
yv2.z = xv2.z * gb2;
|
||||
yv2.w = xv2.w * gb3;
|
||||
reinterpret_cast<float4*>(yr + base2)[0] = yv2;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
torch::Tensor sigmoid_slope_affine_gate_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, double alpha, double 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 = (float)alpha;
|
||||
float be = (float)beta;
|
||||
int block = 256;
|
||||
int elements_per_thread = 4;
|
||||
int elements_per_block = block * elements_per_thread;
|
||||
int gy = (D + elements_per_block - 1) / elements_per_block;
|
||||
dim3 grid(B, gy);
|
||||
sigmoid_slope_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;
|
||||
}
|
||||
"""
|
||||
|
||||
cpp_source = """
|
||||
#include <torch/extension.h>
|
||||
torch::Tensor sigmoid_slope_affine_gate_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, double alpha, double beta);
|
||||
"""
|
||||
|
||||
ops = load_inline(
|
||||
name="sigmoid_slope_affine_gate",
|
||||
cpp_sources=cpp_source,
|
||||
cuda_sources=source,
|
||||
functions=["sigmoid_slope_affine_gate_cuda"],
|
||||
extra_cflags=["-O3","-std=c++17"],
|
||||
extra_cuda_cflags=["-O3","--use_fast_math","-std=c++17","-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.alpha = float(alpha)
|
||||
self.beta = float(beta)
|
||||
|
||||
def forward(self, x):
|
||||
return self.ops.sigmoid_slope_affine_gate_cuda(x, self.scale, self.bias, self.alpha, self.beta)
|
||||
|
|
@ -1,10 +0,0 @@
|
|||
Operator: SigmoidSlope-Affine-Gate (Fused CUDA Kernel)
|
||||
|
||||
Definition
|
||||
- z = x * scale + bias
|
||||
- m = sigmoid(z) * (1 - sigmoid(z))
|
||||
- g = sigmoid(alpha * m + beta)
|
||||
- y = x * g
|
||||
|
||||
Goal
|
||||
- Fuse derivative-like feature; target ≥1.30x speedup.
|
||||
|
|
@ -1,50 +0,0 @@
|
|||
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 SigmoidSlope-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()
|
||||
|
|
@ -1,29 +0,0 @@
|
|||
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
|
||||
s = torch.sigmoid(z)
|
||||
m = s * (1.0 - s)
|
||||
g = torch.sigmoid(self.alpha * m + 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]
|
||||
|
|
@ -1,128 +0,0 @@
|
|||
import torch
|
||||
from torch.utils.cpp_extension import load_inline
|
||||
|
||||
source = """
|
||||
#include <torch/extension.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <vector_types.h>
|
||||
|
||||
__global__ __launch_bounds__(256) void tanh_slope_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){
|
||||
int row = blockIdx.x;
|
||||
int tid = threadIdx.x;
|
||||
int col_block = blockIdx.y;
|
||||
int stride4 = blockDim.x * 4;
|
||||
int col_idx = (col_block * stride4) + tid * 4;
|
||||
const float* xr = x + row * D;
|
||||
float* yr = y + row * D;
|
||||
for(int base = col_idx; base < D; base += stride4 * 2){
|
||||
if (base < D){
|
||||
float4 xv = reinterpret_cast<const float4*>(xr + base)[0];
|
||||
float4 sv = reinterpret_cast<const float4*>(scale + base)[0];
|
||||
float4 bv = reinterpret_cast<const float4*>(bias + base)[0];
|
||||
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 = tanhf(z0);
|
||||
float t1 = tanhf(z1);
|
||||
float t2 = tanhf(z2);
|
||||
float t3 = tanhf(z3);
|
||||
float m0 = 1.0f - t0 * t0;
|
||||
float m1 = 1.0f - t1 * t1;
|
||||
float m2 = 1.0f - t2 * t2;
|
||||
float m3 = 1.0f - t3 * t3;
|
||||
float tt0 = fmaf(alpha, m0, beta);
|
||||
float tt1 = fmaf(alpha, m1, beta);
|
||||
float tt2 = fmaf(alpha, m2, beta);
|
||||
float tt3 = fmaf(alpha, m3, beta);
|
||||
float g0 = __fdividef(1.0f, 1.0f + __expf(-tt0));
|
||||
float g1 = __fdividef(1.0f, 1.0f + __expf(-tt1));
|
||||
float g2 = __fdividef(1.0f, 1.0f + __expf(-tt2));
|
||||
float g3 = __fdividef(1.0f, 1.0f + __expf(-tt3));
|
||||
yv.x = xv.x * g0;
|
||||
yv.y = xv.y * g1;
|
||||
yv.z = xv.z * g2;
|
||||
yv.w = xv.w * g3;
|
||||
reinterpret_cast<float4*>(yr + base)[0] = yv;
|
||||
}
|
||||
int base2 = base + stride4;
|
||||
if (base2 < D){
|
||||
float4 xv2 = reinterpret_cast<const float4*>(xr + base2)[0];
|
||||
float4 sv2 = reinterpret_cast<const float4*>(scale + base2)[0];
|
||||
float4 bv2 = reinterpret_cast<const float4*>(bias + base2)[0];
|
||||
float4 yv2;
|
||||
float z0b = fmaf(xv2.x, sv2.x, bv2.x);
|
||||
float z1b = fmaf(xv2.y, sv2.y, bv2.y);
|
||||
float z2b = fmaf(xv2.z, sv2.z, bv2.z);
|
||||
float z3b = fmaf(xv2.w, sv2.w, bv2.w);
|
||||
float tb0 = tanhf(z0b);
|
||||
float tb1 = tanhf(z1b);
|
||||
float tb2 = tanhf(z2b);
|
||||
float tb3 = tanhf(z3b);
|
||||
float mb0 = 1.0f - tb0 * tb0;
|
||||
float mb1 = 1.0f - tb1 * tb1;
|
||||
float mb2 = 1.0f - tb2 * tb2;
|
||||
float mb3 = 1.0f - tb3 * tb3;
|
||||
float ttb0 = fmaf(alpha, mb0, beta);
|
||||
float ttb1 = fmaf(alpha, mb1, beta);
|
||||
float ttb2 = fmaf(alpha, mb2, beta);
|
||||
float ttb3 = fmaf(alpha, mb3, beta);
|
||||
float gb0 = __fdividef(1.0f, 1.0f + __expf(-ttb0));
|
||||
float gb1 = __fdividef(1.0f, 1.0f + __expf(-ttb1));
|
||||
float gb2 = __fdividef(1.0f, 1.0f + __expf(-ttb2));
|
||||
float gb3 = __fdividef(1.0f, 1.0f + __expf(-ttb3));
|
||||
yv2.x = xv2.x * gb0;
|
||||
yv2.y = xv2.y * gb1;
|
||||
yv2.z = xv2.z * gb2;
|
||||
yv2.w = xv2.w * gb3;
|
||||
reinterpret_cast<float4*>(yr + base2)[0] = yv2;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
torch::Tensor tanh_slope_affine_gate_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, double alpha, double 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 = (float)alpha;
|
||||
float be = (float)beta;
|
||||
int block = 256;
|
||||
int elements_per_thread = 4;
|
||||
int elements_per_block = block * elements_per_thread;
|
||||
int gy = (D + elements_per_block - 1) / elements_per_block;
|
||||
dim3 grid(B, gy);
|
||||
tanh_slope_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;
|
||||
}
|
||||
"""
|
||||
|
||||
cpp_source = """
|
||||
#include <torch/extension.h>
|
||||
torch::Tensor tanh_slope_affine_gate_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, double alpha, double beta);
|
||||
"""
|
||||
|
||||
ops = load_inline(
|
||||
name="tanh_slope_affine_gate",
|
||||
cpp_sources=cpp_source,
|
||||
cuda_sources=source,
|
||||
functions=["tanh_slope_affine_gate_cuda"],
|
||||
extra_cflags=["-O3","-std=c++17"],
|
||||
extra_cuda_cflags=["-O3","--use_fast_math","-std=c++17","-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.alpha = float(alpha)
|
||||
self.beta = float(beta)
|
||||
|
||||
def forward(self, x):
|
||||
return self.ops.tanh_slope_affine_gate_cuda(x, self.scale, self.bias, self.alpha, self.beta)
|
||||
|
|
@ -1,10 +0,0 @@
|
|||
Operator: TanhSlope-Affine-Gate (Fused CUDA Kernel)
|
||||
|
||||
Definition
|
||||
- z = x * scale + bias
|
||||
- m = 1 - tanh(z)^2
|
||||
- g = sigmoid(alpha * m + beta)
|
||||
- y = x * g
|
||||
|
||||
Goal
|
||||
- Fuse derivative-like feature; target ≥1.30x speedup.
|
||||
|
|
@ -1,50 +0,0 @@
|
|||
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 TanhSlope-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()
|
||||
|
|
@ -1,30 +0,0 @@
|
|||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
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
|
||||
t = torch.tanh(z)
|
||||
m = 1.0 - t * t
|
||||
g = torch.sigmoid(self.alpha * m + 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]
|
||||
|
|
@ -1,120 +0,0 @@
|
|||
import torch
|
||||
from torch.utils.cpp_extension import load_inline
|
||||
|
||||
source = """
|
||||
#include <torch/extension.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <vector_types.h>
|
||||
|
||||
__global__ __launch_bounds__(256) void rationalclip_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){
|
||||
int row = blockIdx.x;
|
||||
int tid = threadIdx.x;
|
||||
int col_block = blockIdx.y;
|
||||
int stride4 = blockDim.x * 4;
|
||||
int col_idx = (col_block * stride4) + tid * 4;
|
||||
const float* xr = x + row * D;
|
||||
float* yr = y + row * D;
|
||||
for(int base = col_idx; base < D; base += stride4 * 2){
|
||||
if (base < D){
|
||||
float4 xv = reinterpret_cast<const float4*>(xr + base)[0];
|
||||
float4 sv = reinterpret_cast<const float4*>(scale + base)[0];
|
||||
float4 bv = reinterpret_cast<const float4*>(bias + base)[0];
|
||||
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 m0 = __fdividef(z0, 1.0f + z0*z0);
|
||||
float m1 = __fdividef(z1, 1.0f + z1*z1);
|
||||
float m2 = __fdividef(z2, 1.0f + z2*z2);
|
||||
float m3 = __fdividef(z3, 1.0f + z3*z3);
|
||||
float t0 = fmaf(alpha, m0, beta);
|
||||
float t1 = fmaf(alpha, m1, beta);
|
||||
float t2 = fmaf(alpha, m2, beta);
|
||||
float t3 = fmaf(alpha, m3, 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 + base)[0] = yv;
|
||||
}
|
||||
int base2 = base + stride4;
|
||||
if (base2 < D){
|
||||
float4 xv2 = reinterpret_cast<const float4*>(xr + base2)[0];
|
||||
float4 sv2 = reinterpret_cast<const float4*>(scale + base2)[0];
|
||||
float4 bv2 = reinterpret_cast<const float4*>(bias + base2)[0];
|
||||
float4 yv2;
|
||||
float z0b = fmaf(xv2.x, sv2.x, bv2.x);
|
||||
float z1b = fmaf(xv2.y, sv2.y, bv2.y);
|
||||
float z2b = fmaf(xv2.z, sv2.z, bv2.z);
|
||||
float z3b = fmaf(xv2.w, sv2.w, bv2.w);
|
||||
float mb0 = __fdividef(z0b, 1.0f + z0b*z0b);
|
||||
float mb1 = __fdividef(z1b, 1.0f + z1b*z1b);
|
||||
float mb2 = __fdividef(z2b, 1.0f + z2b*z2b);
|
||||
float mb3 = __fdividef(z3b, 1.0f + z3b*z3b);
|
||||
float tb0 = fmaf(alpha, mb0, beta);
|
||||
float tb1 = fmaf(alpha, mb1, beta);
|
||||
float tb2 = fmaf(alpha, mb2, beta);
|
||||
float tb3 = fmaf(alpha, mb3, beta);
|
||||
float gb0 = __fdividef(1.0f, 1.0f + __expf(-tb0));
|
||||
float gb1 = __fdividef(1.0f, 1.0f + __expf(-tb1));
|
||||
float gb2 = __fdividef(1.0f, 1.0f + __expf(-tb2));
|
||||
float gb3 = __fdividef(1.0f, 1.0f + __expf(-tb3));
|
||||
yv2.x = xv2.x * gb0;
|
||||
yv2.y = xv2.y * gb1;
|
||||
yv2.z = xv2.z * gb2;
|
||||
yv2.w = xv2.w * gb3;
|
||||
reinterpret_cast<float4*>(yr + base2)[0] = yv2;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
torch::Tensor rationalclip_affine_gate_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, double alpha, double 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 = (float)alpha;
|
||||
float be = (float)beta;
|
||||
int block = 256;
|
||||
int elements_per_thread = 4;
|
||||
int elements_per_block = block * elements_per_thread;
|
||||
int gy = (D + elements_per_block - 1) / elements_per_block;
|
||||
dim3 grid(B, gy);
|
||||
rationalclip_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;
|
||||
}
|
||||
"""
|
||||
|
||||
cpp_source = """
|
||||
#include <torch/extension.h>
|
||||
torch::Tensor rationalclip_affine_gate_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, double alpha, double beta);
|
||||
"""
|
||||
|
||||
ops = load_inline(
|
||||
name="rationalclip_affine_gate",
|
||||
cpp_sources=cpp_source,
|
||||
cuda_sources=source,
|
||||
functions=["rationalclip_affine_gate_cuda"],
|
||||
extra_cflags=["-O3","-std=c++17"],
|
||||
extra_cuda_cflags=["-O3","--use_fast_math","-std=c++17","-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.alpha = float(alpha)
|
||||
self.beta = float(beta)
|
||||
|
||||
def forward(self, x):
|
||||
return self.ops.rationalclip_affine_gate_cuda(x, self.scale, self.bias, self.alpha, self.beta)
|
||||
|
|
@ -1,10 +0,0 @@
|
|||
Operator: RationalClip-Affine-Gate (Fused CUDA Kernel)
|
||||
|
||||
Definition
|
||||
- z = x * scale + bias
|
||||
- m = z / (1 + z^2)
|
||||
- g = sigmoid(alpha * m + beta)
|
||||
- y = x * g
|
||||
|
||||
Goal
|
||||
- Rational clipping fusion; target ≥1.30x speedup.
|
||||
|
|
@ -1,50 +0,0 @@
|
|||
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 RationalClip-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()
|
||||
|
|
@ -1,28 +0,0 @@
|
|||
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
|
||||
m = z / (1.0 + z * z)
|
||||
g = torch.sigmoid(self.alpha * m + 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]
|
||||
|
|
@ -1,120 +0,0 @@
|
|||
import torch
|
||||
from torch.utils.cpp_extension import load_inline
|
||||
|
||||
source = """
|
||||
#include <torch/extension.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <vector_types.h>
|
||||
|
||||
__global__ __launch_bounds__(256) void cosine_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){
|
||||
int row = blockIdx.x;
|
||||
int tid = threadIdx.x;
|
||||
int col_block = blockIdx.y;
|
||||
int stride4 = blockDim.x * 4;
|
||||
int col_idx = (col_block * stride4) + tid * 4;
|
||||
const float* xr = x + row * D;
|
||||
float* yr = y + row * D;
|
||||
for(int base = col_idx; base < D; base += stride4 * 2){
|
||||
if (base < D){
|
||||
float4 xv = reinterpret_cast<const float4*>(xr + base)[0];
|
||||
float4 sv = reinterpret_cast<const float4*>(scale + base)[0];
|
||||
float4 bv = reinterpret_cast<const float4*>(bias + base)[0];
|
||||
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 m0 = cosf(z0);
|
||||
float m1 = cosf(z1);
|
||||
float m2 = cosf(z2);
|
||||
float m3 = cosf(z3);
|
||||
float t0 = fmaf(alpha, m0, beta);
|
||||
float t1 = fmaf(alpha, m1, beta);
|
||||
float t2 = fmaf(alpha, m2, beta);
|
||||
float t3 = fmaf(alpha, m3, 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 + base)[0] = yv;
|
||||
}
|
||||
int base2 = base + stride4;
|
||||
if (base2 < D){
|
||||
float4 xv2 = reinterpret_cast<const float4*>(xr + base2)[0];
|
||||
float4 sv2 = reinterpret_cast<const float4*>(scale + base2)[0];
|
||||
float4 bv2 = reinterpret_cast<const float4*>(bias + base2)[0];
|
||||
float4 yv2;
|
||||
float z0b = fmaf(xv2.x, sv2.x, bv2.x);
|
||||
float z1b = fmaf(xv2.y, sv2.y, bv2.y);
|
||||
float z2b = fmaf(xv2.z, sv2.z, bv2.z);
|
||||
float z3b = fmaf(xv2.w, sv2.w, bv2.w);
|
||||
float mb0 = cosf(z0b);
|
||||
float mb1 = cosf(z1b);
|
||||
float mb2 = cosf(z2b);
|
||||
float mb3 = cosf(z3b);
|
||||
float tb0 = fmaf(alpha, mb0, beta);
|
||||
float tb1 = fmaf(alpha, mb1, beta);
|
||||
float tb2 = fmaf(alpha, mb2, beta);
|
||||
float tb3 = fmaf(alpha, mb3, beta);
|
||||
float gb0 = __fdividef(1.0f, 1.0f + __expf(-tb0));
|
||||
float gb1 = __fdividef(1.0f, 1.0f + __expf(-tb1));
|
||||
float gb2 = __fdividef(1.0f, 1.0f + __expf(-tb2));
|
||||
float gb3 = __fdividef(1.0f, 1.0f + __expf(-tb3));
|
||||
yv2.x = xv2.x * gb0;
|
||||
yv2.y = xv2.y * gb1;
|
||||
yv2.z = xv2.z * gb2;
|
||||
yv2.w = xv2.w * gb3;
|
||||
reinterpret_cast<float4*>(yr + base2)[0] = yv2;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
torch::Tensor cosine_affine_gate_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, double alpha, double 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 = (float)alpha;
|
||||
float be = (float)beta;
|
||||
int block = 256;
|
||||
int elements_per_thread = 4;
|
||||
int elements_per_block = block * elements_per_thread;
|
||||
int gy = (D + elements_per_block - 1) / elements_per_block;
|
||||
dim3 grid(B, gy);
|
||||
cosine_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;
|
||||
}
|
||||
"""
|
||||
|
||||
cpp_source = """
|
||||
#include <torch/extension.h>
|
||||
torch::Tensor cosine_affine_gate_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, double alpha, double beta);
|
||||
"""
|
||||
|
||||
ops = load_inline(
|
||||
name="cosine_affine_gate",
|
||||
cpp_sources=cpp_source,
|
||||
cuda_sources=source,
|
||||
functions=["cosine_affine_gate_cuda"],
|
||||
extra_cflags=["-O3","-std=c++17"],
|
||||
extra_cuda_cflags=["-O3","--use_fast_math","-std=c++17","-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.alpha = float(alpha)
|
||||
self.beta = float(beta)
|
||||
|
||||
def forward(self, x):
|
||||
return self.ops.cosine_affine_gate_cuda(x, self.scale, self.bias, self.alpha, self.beta)
|
||||
|
|
@ -1,10 +0,0 @@
|
|||
Operator: Cosine-Affine-Gate (Fused CUDA Kernel)
|
||||
|
||||
Definition
|
||||
- z = x * scale + bias
|
||||
- m = cos(z)
|
||||
- g = sigmoid(alpha * m + beta)
|
||||
- y = x * g
|
||||
|
||||
Goal
|
||||
- Elementwise fusion with periodic feature; target ≥1.30x speedup.
|
||||
|
|
@ -1,50 +0,0 @@
|
|||
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 Cosine-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()
|
||||
|
|
@ -1,28 +0,0 @@
|
|||
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
|
||||
m = torch.cos(z)
|
||||
g = torch.sigmoid(self.alpha * m + 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]
|
||||
|
|
@ -1,120 +0,0 @@
|
|||
import torch
|
||||
from torch.utils.cpp_extension import load_inline
|
||||
|
||||
source = """
|
||||
#include <torch/extension.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <vector_types.h>
|
||||
|
||||
__global__ __launch_bounds__(256) void sine_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){
|
||||
int row = blockIdx.x;
|
||||
int tid = threadIdx.x;
|
||||
int col_block = blockIdx.y;
|
||||
int stride4 = blockDim.x * 4;
|
||||
int col_idx = (col_block * stride4) + tid * 4;
|
||||
const float* xr = x + row * D;
|
||||
float* yr = y + row * D;
|
||||
for(int base = col_idx; base < D; base += stride4 * 2){
|
||||
if (base < D){
|
||||
float4 xv = reinterpret_cast<const float4*>(xr + base)[0];
|
||||
float4 sv = reinterpret_cast<const float4*>(scale + base)[0];
|
||||
float4 bv = reinterpret_cast<const float4*>(bias + base)[0];
|
||||
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 m0 = sinf(z0);
|
||||
float m1 = sinf(z1);
|
||||
float m2 = sinf(z2);
|
||||
float m3 = sinf(z3);
|
||||
float t0 = fmaf(alpha, m0, beta);
|
||||
float t1 = fmaf(alpha, m1, beta);
|
||||
float t2 = fmaf(alpha, m2, beta);
|
||||
float t3 = fmaf(alpha, m3, 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 + base)[0] = yv;
|
||||
}
|
||||
int base2 = base + stride4;
|
||||
if (base2 < D){
|
||||
float4 xv2 = reinterpret_cast<const float4*>(xr + base2)[0];
|
||||
float4 sv2 = reinterpret_cast<const float4*>(scale + base2)[0];
|
||||
float4 bv2 = reinterpret_cast<const float4*>(bias + base2)[0];
|
||||
float4 yv2;
|
||||
float z0b = fmaf(xv2.x, sv2.x, bv2.x);
|
||||
float z1b = fmaf(xv2.y, sv2.y, bv2.y);
|
||||
float z2b = fmaf(xv2.z, sv2.z, bv2.z);
|
||||
float z3b = fmaf(xv2.w, sv2.w, bv2.w);
|
||||
float mb0 = sinf(z0b);
|
||||
float mb1 = sinf(z1b);
|
||||
float mb2 = sinf(z2b);
|
||||
float mb3 = sinf(z3b);
|
||||
float tb0 = fmaf(alpha, mb0, beta);
|
||||
float tb1 = fmaf(alpha, mb1, beta);
|
||||
float tb2 = fmaf(alpha, mb2, beta);
|
||||
float tb3 = fmaf(alpha, mb3, beta);
|
||||
float gb0 = __fdividef(1.0f, 1.0f + __expf(-tb0));
|
||||
float gb1 = __fdividef(1.0f, 1.0f + __expf(-tb1));
|
||||
float gb2 = __fdividef(1.0f, 1.0f + __expf(-tb2));
|
||||
float gb3 = __fdividef(1.0f, 1.0f + __expf(-tb3));
|
||||
yv2.x = xv2.x * gb0;
|
||||
yv2.y = xv2.y * gb1;
|
||||
yv2.z = xv2.z * gb2;
|
||||
yv2.w = xv2.w * gb3;
|
||||
reinterpret_cast<float4*>(yr + base2)[0] = yv2;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
torch::Tensor sine_affine_gate_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, double alpha, double 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 = (float)alpha;
|
||||
float be = (float)beta;
|
||||
int block = 256;
|
||||
int elements_per_thread = 4;
|
||||
int elements_per_block = block * elements_per_thread;
|
||||
int gy = (D + elements_per_block - 1) / elements_per_block;
|
||||
dim3 grid(B, gy);
|
||||
sine_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;
|
||||
}
|
||||
"""
|
||||
|
||||
cpp_source = """
|
||||
#include <torch/extension.h>
|
||||
torch::Tensor sine_affine_gate_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, double alpha, double beta);
|
||||
"""
|
||||
|
||||
ops = load_inline(
|
||||
name="sine_affine_gate",
|
||||
cpp_sources=cpp_source,
|
||||
cuda_sources=source,
|
||||
functions=["sine_affine_gate_cuda"],
|
||||
extra_cflags=["-O3","-std=c++17"],
|
||||
extra_cuda_cflags=["-O3","--use_fast_math","-std=c++17","-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.alpha = float(alpha)
|
||||
self.beta = float(beta)
|
||||
|
||||
def forward(self, x):
|
||||
return self.ops.sine_affine_gate_cuda(x, self.scale, self.bias, self.alpha, self.beta)
|
||||
|
|
@ -1,10 +0,0 @@
|
|||
Operator: Sine-Affine-Gate (Fused CUDA Kernel)
|
||||
|
||||
Definition
|
||||
- z = x * scale + bias
|
||||
- m = sin(z)
|
||||
- g = sigmoid(alpha * m + beta)
|
||||
- y = x * g
|
||||
|
||||
Goal
|
||||
- Elementwise fusion with periodic feature; target ≥1.30x speedup.
|
||||
|
|
@ -1,50 +0,0 @@
|
|||
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 Sine-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()
|
||||
|
|
@ -1,28 +0,0 @@
|
|||
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
|
||||
m = torch.sin(z)
|
||||
g = torch.sigmoid(self.alpha * m + 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]
|
||||
|
|
@ -1,124 +0,0 @@
|
|||
import torch
|
||||
from torch.utils.cpp_extension import load_inline
|
||||
|
||||
source = """
|
||||
#include <torch/extension.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <vector_types.h>
|
||||
|
||||
__device__ __forceinline__ float asinh_stable(float z){
|
||||
return __logf(z + sqrtf(1.0f + z*z));
|
||||
}
|
||||
|
||||
__global__ __launch_bounds__(256) void arcsinh_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){
|
||||
int row = blockIdx.x;
|
||||
int tid = threadIdx.x;
|
||||
int col_block = blockIdx.y;
|
||||
int stride4 = blockDim.x * 4;
|
||||
int col_idx = (col_block * stride4) + tid * 4;
|
||||
const float* xr = x + row * D;
|
||||
float* yr = y + row * D;
|
||||
for(int base = col_idx; base < D; base += stride4 * 2){
|
||||
if (base < D){
|
||||
float4 xv = reinterpret_cast<const float4*>(xr + base)[0];
|
||||
float4 sv = reinterpret_cast<const float4*>(scale + base)[0];
|
||||
float4 bv = reinterpret_cast<const float4*>(bias + base)[0];
|
||||
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 m0 = asinh_stable(z0);
|
||||
float m1 = asinh_stable(z1);
|
||||
float m2 = asinh_stable(z2);
|
||||
float m3 = asinh_stable(z3);
|
||||
float t0 = fmaf(alpha, m0, beta);
|
||||
float t1 = fmaf(alpha, m1, beta);
|
||||
float t2 = fmaf(alpha, m2, beta);
|
||||
float t3 = fmaf(alpha, m3, 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 + base)[0] = yv;
|
||||
}
|
||||
int base2 = base + stride4;
|
||||
if (base2 < D){
|
||||
float4 xv2 = reinterpret_cast<const float4*>(xr + base2)[0];
|
||||
float4 sv2 = reinterpret_cast<const float4*>(scale + base2)[0];
|
||||
float4 bv2 = reinterpret_cast<const float4*>(bias + base2)[0];
|
||||
float4 yv2;
|
||||
float z0b = fmaf(xv2.x, sv2.x, bv2.x);
|
||||
float z1b = fmaf(xv2.y, sv2.y, bv2.y);
|
||||
float z2b = fmaf(xv2.z, sv2.z, bv2.z);
|
||||
float z3b = fmaf(xv2.w, sv2.w, bv2.w);
|
||||
float mb0 = asinh_stable(z0b);
|
||||
float mb1 = asinh_stable(z1b);
|
||||
float mb2 = asinh_stable(z2b);
|
||||
float mb3 = asinh_stable(z3b);
|
||||
float tb0 = fmaf(alpha, mb0, beta);
|
||||
float tb1 = fmaf(alpha, mb1, beta);
|
||||
float tb2 = fmaf(alpha, mb2, beta);
|
||||
float tb3 = fmaf(alpha, mb3, beta);
|
||||
float gb0 = __fdividef(1.0f, 1.0f + __expf(-tb0));
|
||||
float gb1 = __fdividef(1.0f, 1.0f + __expf(-tb1));
|
||||
float gb2 = __fdividef(1.0f, 1.0f + __expf(-tb2));
|
||||
float gb3 = __fdividef(1.0f, 1.0f + __expf(-tb3));
|
||||
yv2.x = xv2.x * gb0;
|
||||
yv2.y = xv2.y * gb1;
|
||||
yv2.z = xv2.z * gb2;
|
||||
yv2.w = xv2.w * gb3;
|
||||
reinterpret_cast<float4*>(yr + base2)[0] = yv2;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
torch::Tensor arcsinh_affine_gate_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, double alpha, double 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 = (float)alpha;
|
||||
float be = (float)beta;
|
||||
int block = 256;
|
||||
int elements_per_thread = 4;
|
||||
int elements_per_block = block * elements_per_thread;
|
||||
int gy = (D + elements_per_block - 1) / elements_per_block;
|
||||
dim3 grid(B, gy);
|
||||
arcsinh_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;
|
||||
}
|
||||
"""
|
||||
|
||||
cpp_source = """
|
||||
#include <torch/extension.h>
|
||||
torch::Tensor arcsinh_affine_gate_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, double alpha, double beta);
|
||||
"""
|
||||
|
||||
ops = load_inline(
|
||||
name="arcsinh_affine_gate",
|
||||
cpp_sources=cpp_source,
|
||||
cuda_sources=source,
|
||||
functions=["arcsinh_affine_gate_cuda"],
|
||||
extra_cflags=["-O3","-std=c++17"],
|
||||
extra_cuda_cflags=["-O3","--use_fast_math","-std=c++17","-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.alpha = float(alpha)
|
||||
self.beta = float(beta)
|
||||
|
||||
def forward(self, x):
|
||||
return self.ops.arcsinh_affine_gate_cuda(x, self.scale, self.bias, self.alpha, self.beta)
|
||||
|
|
@ -1,10 +0,0 @@
|
|||
Operator: ArcSinh-Affine-Gate (Fused CUDA Kernel)
|
||||
|
||||
Definition
|
||||
- z = x * scale + bias
|
||||
- m = asinh(z)
|
||||
- g = sigmoid(alpha * m + beta)
|
||||
- y = x * g
|
||||
|
||||
Goal
|
||||
- Stable asinh fusion; target ≥1.30x speedup.
|
||||
|
|
@ -1,50 +0,0 @@
|
|||
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 ArcSinh-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()
|
||||
|
|
@ -1,28 +0,0 @@
|
|||
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
|
||||
m = torch.asinh(z)
|
||||
g = torch.sigmoid(self.alpha * m + 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]
|
||||
|
|
@ -1,132 +0,0 @@
|
|||
import torch
|
||||
from torch.utils.cpp_extension import load_inline
|
||||
|
||||
source = """
|
||||
#include <torch/extension.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <vector_types.h>
|
||||
|
||||
__device__ __forceinline__ float sigmoidf(float x){
|
||||
return __fdividef(1.0f, 1.0f + __expf(-x));
|
||||
}
|
||||
|
||||
__global__ __launch_bounds__(256) void sigmoid_squared_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){
|
||||
int row = blockIdx.x;
|
||||
int tid = threadIdx.x;
|
||||
int col_block = blockIdx.y;
|
||||
int stride4 = blockDim.x * 4;
|
||||
int col_idx = (col_block * stride4) + tid * 4;
|
||||
const float* xr = x + row * D;
|
||||
float* yr = y + row * D;
|
||||
for(int base = col_idx; base < D; base += stride4 * 2){
|
||||
if (base < D){
|
||||
float4 xv = reinterpret_cast<const float4*>(xr + base)[0];
|
||||
float4 sv = reinterpret_cast<const float4*>(scale + base)[0];
|
||||
float4 bv = reinterpret_cast<const float4*>(bias + base)[0];
|
||||
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 s0 = sigmoidf(z0);
|
||||
float s1 = sigmoidf(z1);
|
||||
float s2 = sigmoidf(z2);
|
||||
float s3 = sigmoidf(z3);
|
||||
float m0 = s0 * s0;
|
||||
float m1 = s1 * s1;
|
||||
float m2 = s2 * s2;
|
||||
float m3 = s3 * s3;
|
||||
float t0 = fmaf(alpha, m0, beta);
|
||||
float t1 = fmaf(alpha, m1, beta);
|
||||
float t2 = fmaf(alpha, m2, beta);
|
||||
float t3 = fmaf(alpha, m3, beta);
|
||||
float g0 = sigmoidf(t0);
|
||||
float g1 = sigmoidf(t1);
|
||||
float g2 = sigmoidf(t2);
|
||||
float g3 = sigmoidf(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 + base)[0] = yv;
|
||||
}
|
||||
int base2 = base + stride4;
|
||||
if (base2 < D){
|
||||
float4 xv2 = reinterpret_cast<const float4*>(xr + base2)[0];
|
||||
float4 sv2 = reinterpret_cast<const float4*>(scale + base2)[0];
|
||||
float4 bv2 = reinterpret_cast<const float4*>(bias + base2)[0];
|
||||
float4 yv2;
|
||||
float z0b = fmaf(xv2.x, sv2.x, bv2.x);
|
||||
float z1b = fmaf(xv2.y, sv2.y, bv2.y);
|
||||
float z2b = fmaf(xv2.z, sv2.z, bv2.z);
|
||||
float z3b = fmaf(xv2.w, sv2.w, bv2.w);
|
||||
float sb0 = sigmoidf(z0b);
|
||||
float sb1 = sigmoidf(z1b);
|
||||
float sb2 = sigmoidf(z2b);
|
||||
float sb3 = sigmoidf(z3b);
|
||||
float mb0 = sb0 * sb0;
|
||||
float mb1 = sb1 * sb1;
|
||||
float mb2 = sb2 * sb2;
|
||||
float mb3 = sb3 * sb3;
|
||||
float tb0 = fmaf(alpha, mb0, beta);
|
||||
float tb1 = fmaf(alpha, mb1, beta);
|
||||
float tb2 = fmaf(alpha, mb2, beta);
|
||||
float tb3 = fmaf(alpha, mb3, beta);
|
||||
float gb0 = sigmoidf(tb0);
|
||||
float gb1 = sigmoidf(tb1);
|
||||
float gb2 = sigmoidf(tb2);
|
||||
float gb3 = sigmoidf(tb3);
|
||||
yv2.x = xv2.x * gb0;
|
||||
yv2.y = xv2.y * gb1;
|
||||
yv2.z = xv2.z * gb2;
|
||||
yv2.w = xv2.w * gb3;
|
||||
reinterpret_cast<float4*>(yr + base2)[0] = yv2;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
torch::Tensor sigmoid_squared_affine_gate_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, double alpha, double 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 = (float)alpha;
|
||||
float be = (float)beta;
|
||||
int block = 256;
|
||||
int elements_per_thread = 4;
|
||||
int elements_per_block = block * elements_per_thread;
|
||||
int gy = (D + elements_per_block - 1) / elements_per_block;
|
||||
dim3 grid(B, gy);
|
||||
sigmoid_squared_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;
|
||||
}
|
||||
"""
|
||||
|
||||
cpp_source = """
|
||||
#include <torch/extension.h>
|
||||
torch::Tensor sigmoid_squared_affine_gate_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, double alpha, double beta);
|
||||
"""
|
||||
|
||||
ops = load_inline(
|
||||
name="sigmoid_squared_affine_gate",
|
||||
cpp_sources=cpp_source,
|
||||
cuda_sources=source,
|
||||
functions=["sigmoid_squared_affine_gate_cuda"],
|
||||
extra_cflags=["-O3","-std=c++17"],
|
||||
extra_cuda_cflags=["-O3","--use_fast_math","-std=c++17","-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.alpha = float(alpha)
|
||||
self.beta = float(beta)
|
||||
|
||||
def forward(self, x):
|
||||
return self.ops.sigmoid_squared_affine_gate_cuda(x, self.scale, self.bias, self.alpha, self.beta)
|
||||
|
|
@ -1,10 +0,0 @@
|
|||
Operator: SigmoidSquared-Affine-Gate (Fused CUDA Kernel)
|
||||
|
||||
Definition
|
||||
- z = x * scale + bias
|
||||
- m = sigmoid(z)^2
|
||||
- g = sigmoid(alpha * m + beta)
|
||||
- y = x * g
|
||||
|
||||
Goal
|
||||
- Fuse power of logistic feature; target ≥1.30x speedup.
|
||||
|
|
@ -1,50 +0,0 @@
|
|||
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 SigmoidSquared-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()
|
||||
|
|
@ -1,29 +0,0 @@
|
|||
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
|
||||
m = torch.sigmoid(z)
|
||||
m = m * m
|
||||
g = torch.sigmoid(self.alpha * m + 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]
|
||||
|
|
@ -1,128 +0,0 @@
|
|||
import torch
|
||||
from torch.utils.cpp_extension import load_inline
|
||||
|
||||
source = """
|
||||
#include <torch/extension.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <vector_types.h>
|
||||
|
||||
__global__ __launch_bounds__(256) void softsign_squared_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){
|
||||
int row = blockIdx.x;
|
||||
int tid = threadIdx.x;
|
||||
int col_block = blockIdx.y;
|
||||
int stride4 = blockDim.x * 4;
|
||||
int col_idx = (col_block * stride4) + tid * 4;
|
||||
const float* xr = x + row * D;
|
||||
float* yr = y + row * D;
|
||||
for(int base = col_idx; base < D; base += stride4 * 2){
|
||||
if (base < D){
|
||||
float4 xv = reinterpret_cast<const float4*>(xr + base)[0];
|
||||
float4 sv = reinterpret_cast<const float4*>(scale + base)[0];
|
||||
float4 bv = reinterpret_cast<const float4*>(bias + base)[0];
|
||||
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 s0 = __fdividef(z0, 1.0f + fabsf(z0));
|
||||
float s1 = __fdividef(z1, 1.0f + fabsf(z1));
|
||||
float s2 = __fdividef(z2, 1.0f + fabsf(z2));
|
||||
float s3 = __fdividef(z3, 1.0f + fabsf(z3));
|
||||
float m0 = s0 * s0;
|
||||
float m1 = s1 * s1;
|
||||
float m2 = s2 * s2;
|
||||
float m3 = s3 * s3;
|
||||
float t0 = fmaf(alpha, m0, beta);
|
||||
float t1 = fmaf(alpha, m1, beta);
|
||||
float t2 = fmaf(alpha, m2, beta);
|
||||
float t3 = fmaf(alpha, m3, 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 + base)[0] = yv;
|
||||
}
|
||||
int base2 = base + stride4;
|
||||
if (base2 < D){
|
||||
float4 xv2 = reinterpret_cast<const float4*>(xr + base2)[0];
|
||||
float4 sv2 = reinterpret_cast<const float4*>(scale + base2)[0];
|
||||
float4 bv2 = reinterpret_cast<const float4*>(bias + base2)[0];
|
||||
float4 yv2;
|
||||
float z0b = fmaf(xv2.x, sv2.x, bv2.x);
|
||||
float z1b = fmaf(xv2.y, sv2.y, bv2.y);
|
||||
float z2b = fmaf(xv2.z, sv2.z, bv2.z);
|
||||
float z3b = fmaf(xv2.w, sv2.w, bv2.w);
|
||||
float sb0 = __fdividef(z0b, 1.0f + fabsf(z0b));
|
||||
float sb1 = __fdividef(z1b, 1.0f + fabsf(z1b));
|
||||
float sb2 = __fdividef(z2b, 1.0f + fabsf(z2b));
|
||||
float sb3 = __fdividef(z3b, 1.0f + fabsf(z3b));
|
||||
float mb0 = sb0 * sb0;
|
||||
float mb1 = sb1 * sb1;
|
||||
float mb2 = sb2 * sb2;
|
||||
float mb3 = sb3 * sb3;
|
||||
float tb0 = fmaf(alpha, mb0, beta);
|
||||
float tb1 = fmaf(alpha, mb1, beta);
|
||||
float tb2 = fmaf(alpha, mb2, beta);
|
||||
float tb3 = fmaf(alpha, mb3, beta);
|
||||
float gb0 = __fdividef(1.0f, 1.0f + __expf(-tb0));
|
||||
float gb1 = __fdividef(1.0f, 1.0f + __expf(-tb1));
|
||||
float gb2 = __fdividef(1.0f, 1.0f + __expf(-tb2));
|
||||
float gb3 = __fdividef(1.0f, 1.0f + __expf(-tb3));
|
||||
yv2.x = xv2.x * gb0;
|
||||
yv2.y = xv2.y * gb1;
|
||||
yv2.z = xv2.z * gb2;
|
||||
yv2.w = xv2.w * gb3;
|
||||
reinterpret_cast<float4*>(yr + base2)[0] = yv2;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
torch::Tensor softsign_squared_affine_gate_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, double alpha, double 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 = (float)alpha;
|
||||
float be = (float)beta;
|
||||
int block = 256;
|
||||
int elements_per_thread = 4;
|
||||
int elements_per_block = block * elements_per_thread;
|
||||
int gy = (D + elements_per_block - 1) / elements_per_block;
|
||||
dim3 grid(B, gy);
|
||||
softsign_squared_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;
|
||||
}
|
||||
"""
|
||||
|
||||
cpp_source = """
|
||||
#include <torch/extension.h>
|
||||
torch::Tensor softsign_squared_affine_gate_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, double alpha, double beta);
|
||||
"""
|
||||
|
||||
ops = load_inline(
|
||||
name="softsign_squared_affine_gate",
|
||||
cpp_sources=cpp_source,
|
||||
cuda_sources=source,
|
||||
functions=["softsign_squared_affine_gate_cuda"],
|
||||
extra_cflags=["-O3","-std=c++17"],
|
||||
extra_cuda_cflags=["-O3","--use_fast_math","-std=c++17","-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.alpha = float(alpha)
|
||||
self.beta = float(beta)
|
||||
|
||||
def forward(self, x):
|
||||
return self.ops.softsign_squared_affine_gate_cuda(x, self.scale, self.bias, self.alpha, self.beta)
|
||||
|
|
@ -1,10 +0,0 @@
|
|||
Operator: SoftSignSquared-Affine-Gate (Fused CUDA Kernel)
|
||||
|
||||
Definition
|
||||
- z = x * scale + bias
|
||||
- m = (z / (1 + |z|))^2
|
||||
- g = sigmoid(alpha * m + beta)
|
||||
- y = x * g
|
||||
|
||||
Goal
|
||||
- Squared softsign fusion; target ≥1.30x speedup.
|
||||
|
|
@ -1,50 +0,0 @@
|
|||
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 SoftSignSquared-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()
|
||||
|
|
@ -1,29 +0,0 @@
|
|||
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
|
||||
s = z / (1.0 + torch.abs(z))
|
||||
m = s * s
|
||||
g = torch.sigmoid(self.alpha * m + 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]
|
||||
|
|
@ -1,56 +0,0 @@
|
|||
import torch
|
||||
from torch.utils.cpp_extension import load_inline
|
||||
|
||||
source = """
|
||||
#include <torch/extension.h>
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
__global__ void glu_kernel(const float* x, float* y, int D, long long rows) {
|
||||
long long idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
long long stride = blockDim.x * gridDim.x;
|
||||
long long total = rows * (long long)D;
|
||||
for (long long t = idx; t < total; t += stride) {
|
||||
long long row = t / D;
|
||||
int col = (int)(t % D);
|
||||
long long base = row * (long long)(2 * D);
|
||||
float a = x[base + col];
|
||||
float b = x[base + D + col];
|
||||
float s = 1.0f / (1.0f + expf(-b));
|
||||
y[t] = a * s;
|
||||
}
|
||||
}
|
||||
|
||||
torch::Tensor glu_cuda(torch::Tensor x) {
|
||||
auto x_contig = x.contiguous();
|
||||
long long rows = 1;
|
||||
for (int i = 0; i < x_contig.dim() - 1; ++i) rows *= x_contig.size(i);
|
||||
int D = (int)(x_contig.size(-1) / 2);
|
||||
auto y = torch::empty({rows, D}, x_contig.options());
|
||||
int block = 512;
|
||||
long long total = rows * (long long)D;
|
||||
long long grid = (total + block - 1) / block;
|
||||
grid = grid > 65535 ? 65535 : grid;
|
||||
glu_kernel<<<(int)grid, block>>>(x_contig.data_ptr<float>(), y.data_ptr<float>(), D, rows);
|
||||
return y;
|
||||
}
|
||||
"""
|
||||
|
||||
cpp_source = """
|
||||
torch::Tensor glu_cuda(torch::Tensor x);
|
||||
"""
|
||||
|
||||
ops = load_inline(
|
||||
name="glu",
|
||||
cpp_sources=cpp_source,
|
||||
cuda_sources=source,
|
||||
functions=["glu_cuda"],
|
||||
verbose=True
|
||||
)
|
||||
|
||||
class ModelNew(torch.nn.Module):
|
||||
def __init__(self):
|
||||
super(ModelNew, self).__init__()
|
||||
self.ops = ops
|
||||
|
||||
def forward(self, x):
|
||||
return self.ops.glu_cuda(x)
|
||||
|
|
@ -1,5 +0,0 @@
|
|||
GLU（Gated Linear Unit）融合：将 `a * sigmoid(b)` 一次内核完成，其中输入最后一维大小为 2D，输出为 D。
|
||||
|
||||
torchcode.py：参考实现，按最后维切分为 `a` 与 `b`。
|
||||
cudacode.py：`__global__ void glu_kernel(...)`，按行访问并融合计算，减少内核次数与显存往返。
|
||||
run_code.py：比较精度与性能（100 次迭代，`rtol=1e-03`）。
|
||||
|
|
@ -1,76 +0,0 @@
|
|||
###########################################################
|
||||
# 性能和精度验证程序
|
||||
###########################################################
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
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
|
||||
else:
|
||||
device = torch.device("cuda")
|
||||
|
||||
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)
|
||||
|
||||
precision_flag = torch.allclose(output_torch, output_cuda,rtol=1e-03)
|
||||
if precision_flag:
|
||||
print("✅ 精度对齐：两个模型的输出结果非常接近。")
|
||||
else:
|
||||
print("❌ 精度不一致！")
|
||||
diff = (output_torch - output_cuda).abs().max().item()
|
||||
print(f"最大绝对误差: {diff}")
|
||||
print(f"输出张量形状: torch={tuple(output_torch.shape)}, cuda={tuple(output_cuda.shape)}")
|
||||
print(f"数据类型: torch={output_torch.dtype}, cuda={output_cuda.dtype}")
|
||||
print(f"设备: torch={output_torch.device}, cuda={output_cuda.device}")
|
||||
|
||||
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 GLU 平均执行时间: {torch_time:.6f} 秒")
|
||||
print(f"自定义 CUDA 融合内核 平均执行时间: {cuda_time:.6f} 秒")
|
||||
speedup = 0
|
||||
if cuda_time > 0:
|
||||
speedup = torch_time / cuda_time
|
||||
print(f"加速比 (Speedup): {speedup:.2f}x")
|
||||
else:
|
||||
print("CUDA 内核执行时间为0，无法计算加速比。")
|
||||
return precision_flag,speedup
|
||||
|
||||
if __name__ == "__main__":
|
||||
precision_flag,speedup = run_benchmark()
|
||||
|
|
@ -1,23 +0,0 @@
|
|||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
class Model(nn.Module):
|
||||
def __init__(self):
|
||||
super(Model, self).__init__()
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
D = x.size(-1) // 2
|
||||
a = x[..., :D]
|
||||
b = x[..., D:]
|
||||
return a * torch.sigmoid(b)
|
||||
|
||||
batch_size = 16
|
||||
dim_half = 16384
|
||||
dim = dim_half * 2
|
||||
|
||||
def get_inputs():
|
||||
x = torch.randn(batch_size, dim)
|
||||
return [x]
|
||||
|
||||
def get_init_inputs():
|
||||
return []
|
||||
|
|
@ -1,120 +0,0 @@
|
|||
import torch
|
||||
from torch.utils.cpp_extension import load_inline
|
||||
|
||||
source = """
|
||||
#include <torch/extension.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <vector_types.h>
|
||||
|
||||
__global__ __launch_bounds__(256) void exponential_linear_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){
|
||||
int row = blockIdx.x;
|
||||
int tid = threadIdx.x;
|
||||
int col_block = blockIdx.y;
|
||||
int stride4 = blockDim.x * 4;
|
||||
int col_idx = (col_block * stride4) + tid * 4;
|
||||
const float* xr = x + row * D;
|
||||
float* yr = y + row * D;
|
||||
for(int base = col_idx; base < D; base += stride4 * 2){
|
||||
if (base < D){
|
||||
float4 xv = reinterpret_cast<const float4*>(xr + base)[0];
|
||||
float4 sv = reinterpret_cast<const float4*>(scale + base)[0];
|
||||
float4 bv = reinterpret_cast<const float4*>(bias + base)[0];
|
||||
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 m0 = z0 * __expf(-fabsf(z0));
|
||||
float m1 = z1 * __expf(-fabsf(z1));
|
||||
float m2 = z2 * __expf(-fabsf(z2));
|
||||
float m3 = z3 * __expf(-fabsf(z3));
|
||||
float t0 = fmaf(alpha, m0, beta);
|
||||
float t1 = fmaf(alpha, m1, beta);
|
||||
float t2 = fmaf(alpha, m2, beta);
|
||||
float t3 = fmaf(alpha, m3, 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 + base)[0] = yv;
|
||||
}
|
||||
int base2 = base + stride4;
|
||||
if (base2 < D){
|
||||
float4 xv2 = reinterpret_cast<const float4*>(xr + base2)[0];
|
||||
float4 sv2 = reinterpret_cast<const float4*>(scale + base2)[0];
|
||||
float4 bv2 = reinterpret_cast<const float4*>(bias + base2)[0];
|
||||
float4 yv2;
|
||||
float z0b = fmaf(xv2.x, sv2.x, bv2.x);
|
||||
float z1b = fmaf(xv2.y, sv2.y, bv2.y);
|
||||
float z2b = fmaf(xv2.z, sv2.z, bv2.z);
|
||||
float z3b = fmaf(xv2.w, sv2.w, bv2.w);
|
||||
float mb0 = z0b * __expf(-fabsf(z0b));
|
||||
float mb1 = z1b * __expf(-fabsf(z1b));
|
||||
float mb2 = z2b * __expf(-fabsf(z2b));
|
||||
float mb3 = z3b * __expf(-fabsf(z3b));
|
||||
float tb0 = fmaf(alpha, mb0, beta);
|
||||
float tb1 = fmaf(alpha, mb1, beta);
|
||||
float tb2 = fmaf(alpha, mb2, beta);
|
||||
float tb3 = fmaf(alpha, mb3, beta);
|
||||
float gb0 = __fdividef(1.0f, 1.0f + __expf(-tb0));
|
||||
float gb1 = __fdividef(1.0f, 1.0f + __expf(-tb1));
|
||||
float gb2 = __fdividef(1.0f, 1.0f + __expf(-tb2));
|
||||
float gb3 = __fdividef(1.0f, 1.0f + __expf(-tb3));
|
||||
yv2.x = xv2.x * gb0;
|
||||
yv2.y = xv2.y * gb1;
|
||||
yv2.z = xv2.z * gb2;
|
||||
yv2.w = xv2.w * gb3;
|
||||
reinterpret_cast<float4*>(yr + base2)[0] = yv2;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
torch::Tensor exponential_linear_affine_gate_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, double alpha, double 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 = (float)alpha;
|
||||
float be = (float)beta;
|
||||
int block = 256;
|
||||
int elements_per_thread = 4;
|
||||
int elements_per_block = block * elements_per_thread;
|
||||
int gy = (D + elements_per_block - 1) / elements_per_block;
|
||||
dim3 grid(B, gy);
|
||||
exponential_linear_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;
|
||||
}
|
||||
"""
|
||||
|
||||
cpp_source = """
|
||||
#include <torch/extension.h>
|
||||
torch::Tensor exponential_linear_affine_gate_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, double alpha, double beta);
|
||||
"""
|
||||
|
||||
ops = load_inline(
|
||||
name="exponential_linear_affine_gate",
|
||||
cpp_sources=cpp_source,
|
||||
cuda_sources=source,
|
||||
functions=["exponential_linear_affine_gate_cuda"],
|
||||
extra_cflags=["-O3","-std=c++17"],
|
||||
extra_cuda_cflags=["-O3","--use_fast_math","-std=c++17","-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.alpha = float(alpha)
|
||||
self.beta = float(beta)
|
||||
|
||||
def forward(self, x):
|
||||
return self.ops.exponential_linear_affine_gate_cuda(x, self.scale, self.bias, self.alpha, self.beta)
|
||||
|
|
@ -1,10 +0,0 @@
|
|||
Operator: ExponentialLinear-Affine-Gate (Fused CUDA Kernel)
|
||||
|
||||
Definition
|
||||
- z = x * scale + bias
|
||||
- m = z * exp(-|z|)
|
||||
- g = sigmoid(alpha * m + beta)
|
||||
- y = x * g
|
||||
|
||||
Goal
|
||||
- Decay-weighted linear fusion; target ≥1.30x speedup.
|
||||
|
|
@ -1,50 +0,0 @@
|
|||
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 ExponentialLinear-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()
|
||||
|
|
@ -1,28 +0,0 @@
|
|||
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
|
||||
m = z * torch.exp(-torch.abs(z))
|
||||
g = torch.sigmoid(self.alpha * m + 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]
|
||||
Some files were not shown because too many files have changed in this diff Show More
Loading…
Reference in New Issue