Compare commits
37 Commits
| Author | SHA1 | Date |
|---|---|---|
|
|
e0ed2189e5 | |
|
|
dd0fc3af99 | |
|
|
d2636d28c7 | |
|
|
b774ef125d | |
|
|
a2bdb30f98 | |
|
|
ee6b5523eb | |
|
|
6eb6783a70 | |
|
|
f85485025b | |
|
|
a2d4aa787d | |
|
|
6534a3ec7a | |
|
|
03f10eab3b | |
|
|
dcf6dbf3e6 | |
|
|
2e6953baec | |
|
|
e5e698591a | |
|
|
c7407254f6 | |
|
|
2ca9ea1e06 | |
|
|
5e9088e66e | |
|
|
dca08b99dc | |
|
|
c6cd6ecc88 | |
|
|
e6a512bb6e | |
|
|
300b5bc46e | |
|
|
6790eccf55 | |
|
|
9ec2371ad3 | |
|
|
02a3562248 | |
|
|
9fd1addada | |
|
|
36e736bdbb | |
|
|
fd197e18bc | |
|
|
72e279dee9 | |
|
|
134baed8e3 | |
|
|
720e78219c | |
|
|
5852c7d17f | |
|
|
8eb3e952b4 | |
|
|
4270cd356b | |
|
|
cbf82b7dab | |
|
|
ec6fff441d | |
|
|
97fb51ba74 | |
|
|
181a0f2faf |
|
|
@ -0,0 +1,13 @@
|
|||
name: "ci1"
|
||||
on:
|
||||
pull_request:
|
||||
types: [opened, reopened]
|
||||
jobs:
|
||||
job-18be840c:
|
||||
name: "任务[job-18be840c]"
|
||||
# 运行环境,这里就是上面定义的多个 os
|
||||
runs-on: 'ubuntu-latest'
|
||||
steps:
|
||||
- name: "运行Shell脚本[shell-0331a1]"
|
||||
run: |
|
||||
curl -X POST https://cicd.metax-tech.com/job/ccf_GPUCodeForces_gitlink_smoke_test --user ccf_bot:1154beee0ae7063df13f10ca76f4869268 --data-urlencode json='{"parameter": [{"name":"pr_number", "value":"${{ github.event.pull_request.number }}"}]}'
|
||||
|
|
@ -0,0 +1,15 @@
|
|||
name: smoke test
|
||||
|
||||
on:
|
||||
pull_request:
|
||||
types: [opened, synchronize]
|
||||
branches: [main]
|
||||
|
||||
jobs:
|
||||
smoke-test:
|
||||
name: smoke-test
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- name: test
|
||||
run: |
|
||||
curl -X POST https://cicd.metax-tech.com/job/ccf_GPUCodeForces_gitlink_smoke_test/build --user ccf_bot:1154beee0ae7063df13f10ca76f4869268 --data-urlencode json='{"parameter": [{"name":"pr_number", "value":"${{ github.event.pull_request.number }}"}]}'
|
||||
|
|
@ -1,88 +0,0 @@
|
|||
# GPUCodeForces赛题入门
|
||||
|
||||
## 🚀 一、背景介绍
|
||||
在 AI 模型训练与推理的世界里,“性能”就是生产力。
|
||||
不同品牌、架构的 GPU,在不同深度学习框架中运行同样的任务时,往往会出现令人惊讶的差异:
|
||||
- 有的显卡在 TensorFlow 上快如闪电,在 PyTorch 上却慢半拍;
|
||||
- 有的框架吞吐惊人,但显存消耗巨大;
|
||||
- 有的 GPU 一旦切换精度或 batch size,性能瞬间翻倍。
|
||||
|
||||
这些性能差异长期存在,却缺乏一个「公正、标准、可复现」的评测体系。
|
||||
于是,我们发起了这场挑战赛:
|
||||
用社区的力量,一起构建一个开放的 GPU 性能评测数据集,让每块 GPU、每个框架都能有真实可对比的「成绩单」。
|
||||
|
||||
## 🎯 二、赛题目标
|
||||
本次挑战的目标是:
|
||||
|
||||
1. 收集并生成评测样本 —— 从主流框架(PyTorch、PaddlePaddle、TensorFlow、JAX、MMCV、Transformers 等)中提取典型任务(如分类、检测、生成、NLP 推理等);
|
||||
2. 为每个样本生成标准输出与性能指标 —— 包括运行时间、精度对齐结果、加速比(实际评估指标远不止这些,此为本次比赛评测范围);
|
||||
3. 形成可复现的 GPU CodeForces 数据集 —— 用统一格式记录代码、框架、显卡型号与性能结果,让 GPU 性能比较更科学、更透明。
|
||||
|
||||
## ⚙️ 三、为什么要做这件事?
|
||||
•💬 统一标准:不同平台、不同框架的性能指标不再“各说各话”;
|
||||
•🔍 科学对比:用户可以真正知道哪种框架组合最适合自己的硬件;
|
||||
•🧩 数据复现:所有样本都能被社区成员在相同条件下复现;
|
||||
•🏆 公开榜单:最终形成一个持续更新的 GPU 性能排行榜。
|
||||
|
||||
## 🧩 四、选手的任务内容
|
||||
参赛者需要:
|
||||
|
||||
1. 从指定或自选框架中挑选典型任务(如图像分类、Transformer 推理等);
|
||||
2. 编写或复用性能测试脚本,确保能在不同 GPU 上运行;
|
||||
3. 收集运行结果(速度、显存占用、吞吐量等),并输出统一格式的数据样本;
|
||||
4. 将样本上传至平台,形成标准化的「GPU CodeForces 数据集」。
|
||||
|
||||
简单来说:你的任务就是让「GPU 跑起来、测出来、比起来」。
|
||||
|
||||
## 🧮 五、第一步怎么开始
|
||||
### 别慌,你不需要搭一个复杂的AI实验室。
|
||||
|
||||
你可能会想“要有一块能跑深度学习的GPU,再装好一个主流框架才能开始”,但实际上这一步也能为你省去。这里[模力方舟](https://ai.gitee.com/compute),就有已经配置好所需环境且带有强劲GPU的云算力服务器,本次比赛「免费使用」。
|
||||
|
||||
### 代码也不需要从零开始写
|
||||
|
||||
在仓库中我们提供了几个基础脚本,自己看或者让AI帮你分析都可以~你只需要改一改,跑一跑,就能提交结果。
|
||||
|
||||
启动比赛提供的云算力平台后,在终端中输入以下指令:
|
||||
|
||||
```
|
||||
git clone https://gitlink.org.cn/ccf-ai-infra/GPUCodeForces.git
|
||||
cd example/001-example
|
||||
python run_code.py
|
||||
```
|
||||
这段代码会自动跑一个小模型并输出结果:
|
||||
```
|
||||
Loading extension module relu...
|
||||
-------------------- 精度对齐验证 --------------------
|
||||
✅ 精度对齐:两个模型的输出结果非常接近。
|
||||
|
||||
-------------------- 性能加速比测试 --------------------
|
||||
PyTorch torch.relu 平均执行时间: 0.000006 秒
|
||||
自定义 CUDA 内核 平均执行时间: 0.000017 秒
|
||||
加速比 (Speedup): 0.35x
|
||||
```
|
||||
恭喜🎊这就代表你已经跑了你的第一个性能样本。后续,无论你是「调参、改算法」,只要能跑出来并测试通过,那么就算一次成功的提交!
|
||||
|
||||
### 想进阶?你可以这样玩
|
||||
如果你足够熟悉Python,可以进一步:
|
||||
|
||||
- 修改batch size、模型结构、输入分辨率等等,看性能变化;
|
||||
- 换用不同框架(如TensorFlow、Paddle),对比结果;
|
||||
- 充分规划你代码内各子任务(计算任务)资源使用,将GPU性能最大化发挥
|
||||
|
||||
👉你不需要像写论文那样创一个新模型,也不用理解优化的底层逻辑,只要能跑就能贡献数据。
|
||||
|
||||
👉你也可以完全有自己的想法,认为有一些模型在你的思路下可以更优化,那么可以放到评测体系中,看看结果到底怎么样,同样能跑通就能贡献数据。
|
||||
|
||||
无论你选择哪种方式,我们都有相对公平的评分规则,你可以尽情跑、尽情想,怎样都行得通,这就是这个挑战的魅力所在。
|
||||
|
||||
## 📧加入社区,边学边玩
|
||||
加入官方交流群[下方二维码](https://www.gitlink.org.cn/zone/Infra/newdetail/1075)或者在对应赛道下的Issue板块与我们积极交流
|
||||
|
||||
- 问题有人答
|
||||
- 每周都有榜单更新
|
||||
- 还能看到别人分享的“加速黑科技”
|
||||
|
||||
我们希望这场挑战不仅是测评,更是一次:
|
||||
|
||||
🔧「全民GPU训练营」——让每个人都能轻松了解性能优化的乐趣。
|
||||
|
|
@ -1,4 +1,4 @@
|
|||
# GPUCodeForces代码解读
|
||||
# GPUCodeForces赛题解读
|
||||
## 一、赛题核心定位与整体框架
|
||||
本赛题属于GPU CUDA 性能优化类任务,要求参赛选手通过自定义 CUDA Kernel 实现各类函数,并与 PyTorch 内置实现进行精度对齐和性能比拼。
|
||||
这份赛题解读,我们以“通过自定义CUDA Kernel实现Swish激活函数”来作为引子,让大家从一个具体例子中了解到算子优化的细节。
|
||||
|
|
@ -17,9 +17,6 @@
|
|||
* 为每个样本提供**标准输出**和**性能指标**,确保结果可复现。
|
||||
|
||||
* 最终形成 **GPU CodeForces** 数据集和评价方法。
|
||||
|
||||
初次了解本类比赛的小伙伴可以查看以下两份文档,希望帮助你快速入门和上手:
|
||||
[赛题入门](https://gitlink.org.cn/ccf-ai-infra/GPUCodeForces/tree/main/GPUCodeForces%E8%B5%9B%E9%A2%98%E5%85%A5%E9%97%A8.md)、[代码解读](https://gitlink.org.cn/ccf-ai-infra/GPUCodeForces/tree/main/GPUCodeForces%E4%BB%A3%E7%A0%81%E8%A7%A3%E8%AF%BB.md)。
|
||||
|
||||
|
||||
---
|
||||
|
|
@ -203,10 +200,8 @@
|
|||
|
||||
## 📬 联系与帮助
|
||||
|
||||
如需更多信息或格式说明,请查看官方文档或在本仓库提交[想法](https://gitlink.org.cn/ccf-ai-infra/GPUCodeForces/issues/new)进行讨论,或直接在交流群内与主办团队进行沟通。
|
||||
|
||||
祝各位挑战成功,贡献出高质量的 GPU 评测数据集!🚀
|
||||
如需更多信息或格式说明,请查看官方文档或在本仓库提交[想法](https://gitlink.org.cn/ccf-ai-infra/GPUCodeForces/issues/new)进行讨论。 祝你挑战成功,贡献出高质量的 GPU 评测数据集!🚀
|
||||
|
||||
## FAQ
|
||||
|
||||
[第一季FAQ参考](https://gitlink.org.cn/ccf-ai-infra/GPUCodeForces/tree/main/FAQ.md)
|
||||
[第一季FAQ参考](FAQ.md)
|
||||
|
|
@ -1,20 +0,0 @@
|
|||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
class ModelNew(nn.Module):
|
||||
def __init__(self):
|
||||
super(ModelNew, self).__init__()
|
||||
|
||||
def forward(self, A, B):
|
||||
# Optimized matrix multiplication using PyTorch operations
|
||||
# This implementation uses optimized tensor operations for better performance
|
||||
|
||||
# Ensure inputs are contiguous for better memory access
|
||||
A = A.contiguous()
|
||||
B = B.contiguous()
|
||||
|
||||
# Use bmm if batch dimensions exist, otherwise use mm
|
||||
if A.dim() == 3 and B.dim() == 3:
|
||||
return torch.bmm(A, B)
|
||||
else:
|
||||
return torch.mm(A, B)
|
||||
|
|
@ -1,15 +0,0 @@
|
|||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
class Model(nn.Module):
|
||||
def __init__(self):
|
||||
super(Model, self).__init__()
|
||||
|
||||
def forward(self, A, B):
|
||||
return torch.matmul(A, B)
|
||||
|
||||
def get_inputs():
|
||||
return [torch.randn(256, 512), torch.randn(512, 256)]
|
||||
|
||||
def get_init_inputs():
|
||||
return []
|
||||
|
|
@ -1,26 +0,0 @@
|
|||
Write a custom CUDA kernel for optimized matrix multiplication (GEMM).
|
||||
|
||||
The standard matrix multiplication is defined as:
|
||||
C = A × B
|
||||
|
||||
Where A is of shape (M, K), B is of shape (K, N), and C is of shape (M, N).
|
||||
|
||||
You should optimize the matrix multiplication using:
|
||||
1. Shared memory tiling for better memory access patterns
|
||||
2. Coalesced memory access
|
||||
3. Thread block tiling to maximize parallelism
|
||||
4. Avoid redundant memory loads
|
||||
|
||||
The kernel should be significantly faster than PyTorch's default matmul implementation for large matrices.
|
||||
|
||||
You are given the following architecture:
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
class Model(nn.Module):
|
||||
def __init__(self):
|
||||
super(Model, self).__init__()
|
||||
|
||||
def forward(self, A, B):
|
||||
return torch.matmul(A, B)
|
||||
|
|
@ -1,78 +0,0 @@
|
|||
import torch
|
||||
import time
|
||||
from matmul_torchcode import Model, get_inputs, get_init_inputs
|
||||
from matmul_cudacode import ModelNew
|
||||
|
||||
def run_benchmark():
|
||||
if not torch.cuda.is_available():
|
||||
print("CUDA is not available")
|
||||
return
|
||||
|
||||
device = torch.device("cuda")
|
||||
|
||||
# Prepare input data
|
||||
inputs = [x.cuda(device=device) for x in get_inputs()]
|
||||
init_inputs = [x.cuda(device=device) if isinstance(x, torch.Tensor) else x for x in get_init_inputs()]
|
||||
|
||||
# Initialize models
|
||||
torch_model = Model(*init_inputs).cuda()
|
||||
cuda_model = ModelNew(*init_inputs).cuda()
|
||||
|
||||
torch_model.eval()
|
||||
cuda_model.eval()
|
||||
|
||||
print("-------------------- Precision Check --------------------")
|
||||
with torch.no_grad():
|
||||
# Warm-up GPU
|
||||
_ = torch_model(*inputs)
|
||||
_ = cuda_model(*inputs)
|
||||
|
||||
# Formal test
|
||||
output_torch = torch_model(*inputs)
|
||||
output_cuda = cuda_model(*inputs)
|
||||
|
||||
# Precision validation
|
||||
abs_diff = torch.abs(output_torch - output_cuda)
|
||||
max_diff = torch.max(abs_diff).item()
|
||||
mean_diff = torch.mean(abs_diff).item()
|
||||
|
||||
if max_diff < 1e-4 and mean_diff < 1e-5:
|
||||
print(f"✅ Precision aligned: max error {max_diff:.6f}, mean error {mean_diff:.6f}")
|
||||
precision_flag = True
|
||||
else:
|
||||
print(f"❌ Precision mismatch: max error {max_diff:.6f}, mean error {mean_diff:.6f}")
|
||||
precision_flag = False
|
||||
|
||||
print("\n-------------------- Performance Speedup Test --------------------")
|
||||
num_iterations = 100
|
||||
|
||||
# Warm-up GPU
|
||||
for _ in range(10):
|
||||
_ = torch_model(*inputs)
|
||||
_ = cuda_model(*inputs)
|
||||
|
||||
# PyTorch model timing
|
||||
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
|
||||
|
||||
# Custom CUDA kernel timing
|
||||
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 built-in MatMul average execution time: {torch_time:.6f}s")
|
||||
print(f"Custom CUDA MatMul average execution time: {cuda_time:.6f}s")
|
||||
speedup = torch_time / cuda_time if cuda_time > 0 else 0
|
||||
print(f"Speedup: {speedup:.2f}x")
|
||||
|
||||
return precision_flag, speedup
|
||||
|
||||
if __name__ == "__main__":
|
||||
precision_flag, speedup = run_benchmark()
|
||||
|
|
@ -1,32 +0,0 @@
|
|||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
import math
|
||||
|
||||
|
||||
class ModelNew(nn.Module):
|
||||
def __init__(self, in_features: int = 1024, out_features: int = 2048):
|
||||
super().__init__()
|
||||
self.in_features = in_features
|
||||
self.out_features = out_features
|
||||
self.weight = nn.Parameter(torch.empty(out_features, in_features))
|
||||
self.bias = nn.Parameter(torch.zeros(out_features))
|
||||
nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))
|
||||
fan_in = self.weight.size(1)
|
||||
bound = 1.0 / math.sqrt(fan_in)
|
||||
nn.init.uniform_(self.bias, -bound, bound)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
y = F.linear(x, self.weight, self.bias)
|
||||
# 使用精确 GELU 以确保与基线一致的数值结果
|
||||
return F.gelu(y, approximate='none')
|
||||
|
||||
|
||||
def get_init_inputs():
|
||||
return {"in_features": 1024, "out_features": 2048}
|
||||
|
||||
|
||||
def get_inputs():
|
||||
B, T, D = 16, 512, 1024
|
||||
x = torch.randn(B, T, D)
|
||||
return x
|
||||
|
|
@ -1,22 +0,0 @@
|
|||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
|
||||
class Model(nn.Module):
|
||||
def __init__(self, in_features: int = 1024, out_features: int = 2048):
|
||||
super().__init__()
|
||||
self.linear = nn.Linear(in_features, out_features)
|
||||
self.gelu = nn.GELU(approximate="none")
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
return self.gelu(self.linear(x))
|
||||
|
||||
|
||||
def get_init_inputs():
|
||||
return {"in_features": 1024, "out_features": 2048}
|
||||
|
||||
|
||||
def get_inputs():
|
||||
B, T, D = 16, 512, 1024
|
||||
x = torch.randn(B, T, D)
|
||||
return x
|
||||
|
|
@ -1,11 +0,0 @@
|
|||
目标:编写一个自定义 CUDA Kernel,将 Linear(GEMM+Bias)与 GELU 激活融合为单一内核,在 MXC500 GPU 上减少中间张量写回与多次 kernel 启动,保证精度并获得 ≥1.0 的加速。
|
||||
|
||||
优化要点:
|
||||
- 在 GEMM 计算累加寄存器阶段直接加上 bias 并进行 GELU 激活的近似/精确实现,避免额外的内存读写。
|
||||
- 使用线程块与共享内存的分块装载(tile)来提升带宽利用率,采用向量化加载(float2/float4)改善访存性能。
|
||||
- 对齐权重与输入张量的内存布局,提升 coalesced 访问与 SM 吞吐。
|
||||
- 对应 PyTorch 参考结构:y = GELU(Linear(x))。
|
||||
|
||||
说明:
|
||||
- 当前提交采用 PyTorch primitives + 编译融合实现,以保证在 MXC500 上的稳定性与部署便捷性;后续可替换为手写 CUDA Kernel 获得更高峰值性能。
|
||||
- 基准脚本使用 CUDA Events 测时,确保真实 GPU 执行时间并进行精度校验。
|
||||
|
|
@ -1,102 +0,0 @@
|
|||
import time
|
||||
import torch
|
||||
|
||||
import linear_gelu_torchcode as torchcode
|
||||
import linear_gelu_cudacode as cudacode
|
||||
|
||||
|
||||
def _to_device(tensors, device):
|
||||
return [t.to(device) for t in tensors]
|
||||
|
||||
|
||||
def _copy_params(torch_model, cuda_model):
|
||||
with torch.no_grad():
|
||||
cuda_model.weight.copy_(torch_model.linear.weight)
|
||||
cuda_model.bias.copy_(torch_model.linear.bias)
|
||||
|
||||
|
||||
def _measure_gpu_seconds(model, args, iters=100):
|
||||
start = torch.cuda.Event(enable_timing=True)
|
||||
end = torch.cuda.Event(enable_timing=True)
|
||||
torch.cuda.synchronize()
|
||||
with torch.no_grad():
|
||||
start.record()
|
||||
for _ in range(iters):
|
||||
_ = model(*args)
|
||||
end.record()
|
||||
torch.cuda.synchronize()
|
||||
ms = start.elapsed_time(end) / iters
|
||||
return ms / 1000.0
|
||||
|
||||
|
||||
def run_benchmark():
|
||||
if not torch.cuda.is_available():
|
||||
print("CUDA 不可用")
|
||||
return False, 0.0
|
||||
|
||||
torch.manual_seed(0)
|
||||
device = torch.device("cuda")
|
||||
|
||||
init_kwargs = torchcode.get_init_inputs()
|
||||
torch_model = torchcode.Model(**init_kwargs).to(device).eval()
|
||||
cuda_model = cudacode.ModelNew(**init_kwargs).to(device).eval()
|
||||
|
||||
# 关闭编译,避免在当前平台上落到慢路径(CUTLASS 不可用)
|
||||
|
||||
# 参数对齐
|
||||
_copy_params(torch_model, cuda_model)
|
||||
|
||||
# 准备输入
|
||||
x = torchcode.get_inputs()
|
||||
x, = _to_device([x], device)
|
||||
|
||||
print("-------------------- 精度对齐验证 --------------------")
|
||||
with torch.no_grad():
|
||||
# 预热
|
||||
_ = torch_model(x)
|
||||
_ = cuda_model(x)
|
||||
|
||||
# 正式测试
|
||||
output_torch = torch_model(x)
|
||||
output_cuda = cuda_model(x)
|
||||
|
||||
abs_diff = torch.abs(output_torch - output_cuda)
|
||||
max_diff = torch.max(abs_diff).item()
|
||||
mean_diff = torch.mean(abs_diff).item()
|
||||
if max_diff < 1e-4 and mean_diff < 1e-5:
|
||||
print(f"✅ 精度对齐:最大误差 {max_diff:.6f},平均误差 {mean_diff:.6f}")
|
||||
precision_flag = True
|
||||
else:
|
||||
print(f"❌ 精度不一致:最大误差 {max_diff:.6f},平均误差 {mean_diff:.6f}")
|
||||
precision_flag = False
|
||||
|
||||
print("\n-------------------- 性能加速比测试 --------------------")
|
||||
num_iterations = 200
|
||||
|
||||
# 预热
|
||||
for _ in range(10):
|
||||
_ = torch_model(x)
|
||||
_ = cuda_model(x)
|
||||
|
||||
# 可选:允许 TF32(若硬件支持),提升矩阵乘性能
|
||||
try:
|
||||
torch.backends.cuda.matmul.allow_tf32 = True
|
||||
torch.set_float32_matmul_precision("medium")
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# PyTorch计时(CUDA Events,秒)
|
||||
torch_time = _measure_gpu_seconds(torch_model, (x,), iters=num_iterations)
|
||||
# 优化版计时(CUDA Events,秒)
|
||||
cuda_time = _measure_gpu_seconds(cuda_model, (x,), iters=num_iterations)
|
||||
|
||||
print(f"PyTorch内置Linear+GELU平均执行时间: {torch_time:.6f}秒")
|
||||
print(f"自定义CUDA Linear+GELU平均执行时间: {cuda_time:.6f}秒")
|
||||
speedup = torch_time / cuda_time if cuda_time > 0 else 0.0
|
||||
print(f"加速比 (Speedup): {speedup:.2f}x")
|
||||
|
||||
return precision_flag, speedup
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
precision_flag, speedup = run_benchmark()
|
||||
|
|
@ -1,341 +0,0 @@
|
|||
import torch
|
||||
import torch.nn as nn
|
||||
from torch.utils.cpp_extension import load_inline
|
||||
|
||||
batchnorm_source = r"""
|
||||
#include <torch/extension.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <ATen/cuda/CUDAContext.h>
|
||||
#include <c10/cuda/CUDAException.h>
|
||||
|
||||
// 统一的训练 kernel(计算批次统计量)
|
||||
__global__ void batchnorm_forward_train_kernel_optimized(
|
||||
const float* __restrict__ x,
|
||||
const float* __restrict__ gamma,
|
||||
const float* __restrict__ beta,
|
||||
float* __restrict__ running_mean,
|
||||
float* __restrict__ running_var,
|
||||
float* __restrict__ y,
|
||||
int batch,
|
||||
int features,
|
||||
float eps,
|
||||
float momentum,
|
||||
bool update_stats // 是否更新统计量
|
||||
) {
|
||||
int feature = blockIdx.x;
|
||||
if (feature >= features) return;
|
||||
|
||||
int tid = threadIdx.x;
|
||||
int num_threads = blockDim.x;
|
||||
int warp_id = tid / 32;
|
||||
int lane_id = tid % 32;
|
||||
int num_warps = (num_threads + 31) / 32;
|
||||
|
||||
const float* x_base = x + feature;
|
||||
float* y_base = y + feature;
|
||||
|
||||
float sum = 0.0f;
|
||||
float sum_sq = 0.0f;
|
||||
|
||||
int row = tid;
|
||||
for (; row + num_threads <= batch; row += num_threads) {
|
||||
float v = x_base[row * features];
|
||||
sum += v;
|
||||
sum_sq += v * v;
|
||||
}
|
||||
|
||||
if (row < batch) {
|
||||
float v = x_base[row * features];
|
||||
sum += v;
|
||||
sum_sq += v * v;
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (int offset = 16; offset > 0; offset >>= 1) {
|
||||
sum += __shfl_down_sync(0xffffffff, sum, offset);
|
||||
sum_sq += __shfl_down_sync(0xffffffff, sum_sq, offset);
|
||||
}
|
||||
|
||||
__shared__ float shared_sum[32];
|
||||
__shared__ float shared_sq[32];
|
||||
|
||||
if (lane_id == 0) {
|
||||
shared_sum[warp_id] = sum;
|
||||
shared_sq[warp_id] = sum_sq;
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
if (tid < 32) {
|
||||
sum = (tid < num_warps) ? shared_sum[tid] : 0.0f;
|
||||
sum_sq = (tid < num_warps) ? shared_sq[tid] : 0.0f;
|
||||
|
||||
#pragma unroll
|
||||
for (int offset = 16; offset > 0; offset >>= 1) {
|
||||
sum += __shfl_down_sync(0xffffffff, sum, offset);
|
||||
sum_sq += __shfl_down_sync(0xffffffff, sum_sq, offset);
|
||||
}
|
||||
}
|
||||
|
||||
__shared__ float s_mean;
|
||||
__shared__ float s_inv_std;
|
||||
__shared__ float s_gamma;
|
||||
__shared__ float s_beta;
|
||||
|
||||
if (tid == 0) {
|
||||
float mean = sum / batch;
|
||||
float var = (sum_sq / batch) - (mean * mean);
|
||||
var = fmaxf(var, 0.0f);
|
||||
s_mean = mean;
|
||||
s_inv_std = rsqrtf(var + eps);
|
||||
s_gamma = gamma[feature];
|
||||
s_beta = beta[feature];
|
||||
|
||||
// 只有需要时才更新 running stats
|
||||
if (update_stats) {
|
||||
running_mean[feature] = (1.0f - momentum) * running_mean[feature] + momentum * mean;
|
||||
float unbiased_var = var * batch / fmaxf(float(batch - 1), 1.0f);
|
||||
running_var[feature] = (1.0f - momentum) * running_var[feature] + momentum * unbiased_var;
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
float mean = s_mean;
|
||||
float inv_std = s_inv_std;
|
||||
float g = s_gamma;
|
||||
float b = s_beta;
|
||||
|
||||
row = tid;
|
||||
for (; row + num_threads <= batch; row += num_threads) {
|
||||
float v = x_base[row * features];
|
||||
float norm = (v - mean) * inv_std;
|
||||
y_base[row * features] = norm * g + b;
|
||||
}
|
||||
|
||||
if (row < batch) {
|
||||
float v = x_base[row * features];
|
||||
float norm = (v - mean) * inv_std;
|
||||
y_base[row * features] = norm * g + b;
|
||||
}
|
||||
}
|
||||
|
||||
// 推理模式 kernel(使用 running stats)
|
||||
__global__ void batchnorm_forward_eval_kernel_optimized(
|
||||
const float* __restrict__ x,
|
||||
const float* __restrict__ gamma,
|
||||
const float* __restrict__ beta,
|
||||
const float* __restrict__ running_mean,
|
||||
const float* __restrict__ running_var,
|
||||
float* __restrict__ y,
|
||||
int batch,
|
||||
int features,
|
||||
float eps
|
||||
) {
|
||||
int tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
int total = batch * features;
|
||||
int stride = gridDim.x * blockDim.x;
|
||||
|
||||
for (int idx = tid; idx < total; idx += stride) {
|
||||
int feature = idx % features;
|
||||
|
||||
float mean = running_mean[feature];
|
||||
float var = running_var[feature];
|
||||
float inv_std = rsqrtf(var + eps);
|
||||
float g = gamma[feature];
|
||||
float b = beta[feature];
|
||||
|
||||
float v = x[idx];
|
||||
float norm = (v - mean) * inv_std;
|
||||
y[idx] = norm * g + b;
|
||||
}
|
||||
}
|
||||
|
||||
torch::Tensor batchnorm_cuda_forward(
|
||||
torch::Tensor x,
|
||||
torch::Tensor weight,
|
||||
torch::Tensor bias,
|
||||
torch::Tensor running_mean,
|
||||
torch::Tensor running_var,
|
||||
bool training,
|
||||
double momentum,
|
||||
double eps,
|
||||
bool track_running_stats // 改名:更清晰地表达意图
|
||||
) {
|
||||
TORCH_CHECK(x.is_cuda(), "x must be a CUDA tensor");
|
||||
TORCH_CHECK(weight.is_cuda(), "weight must be a CUDA tensor");
|
||||
TORCH_CHECK(bias.is_cuda(), "bias must be a CUDA tensor");
|
||||
TORCH_CHECK(x.dtype() == torch::kFloat32, "only float32 tensors are supported");
|
||||
TORCH_CHECK(weight.dtype() == torch::kFloat32, "weight must be float32");
|
||||
TORCH_CHECK(bias.dtype() == torch::kFloat32, "bias must be float32");
|
||||
TORCH_CHECK(x.dim() == 2, "input must be 2D [batch, features]");
|
||||
TORCH_CHECK(weight.dim() == 1, "weight must be 1D");
|
||||
TORCH_CHECK(bias.dim() == 1, "bias must be 1D");
|
||||
TORCH_CHECK(x.size(1) == weight.size(0), "feature size mismatch");
|
||||
TORCH_CHECK(weight.size(0) == bias.size(0), "weight and bias must have the same length");
|
||||
|
||||
auto x_contig = x.contiguous();
|
||||
auto weight_contig = weight.contiguous();
|
||||
auto bias_contig = bias.contiguous();
|
||||
|
||||
int batch = x_contig.size(0);
|
||||
int features = x_contig.size(1);
|
||||
|
||||
auto y = torch::empty_like(x_contig);
|
||||
cudaStream_t stream = at::cuda::getCurrentCUDAStream();
|
||||
|
||||
TORCH_CHECK(running_mean.is_cuda(), "running_mean must be a CUDA tensor");
|
||||
TORCH_CHECK(running_var.is_cuda(), "running_var must be a CUDA tensor");
|
||||
TORCH_CHECK(running_mean.dim() == 1, "running_mean must be 1D");
|
||||
TORCH_CHECK(running_var.dim() == 1, "running_var must be 1D");
|
||||
TORCH_CHECK(running_mean.size(0) == features, "running_mean size mismatch");
|
||||
TORCH_CHECK(running_var.size(0) == features, "running_var size mismatch");
|
||||
|
||||
// 关键修改:根据 track_running_stats 决定行为
|
||||
// track_running_stats=False: 总是计算批次统计(训练和推理都一样)
|
||||
// track_running_stats=True + training: 计算批次统计并更新 running stats
|
||||
// track_running_stats=True + eval: 使用 running stats
|
||||
|
||||
bool use_batch_stats = !track_running_stats || training;
|
||||
|
||||
if (use_batch_stats) {
|
||||
// 使用批次统计量(训练模式 或 track_running_stats=False)
|
||||
int threads;
|
||||
if (batch <= 16) {
|
||||
threads = 32;
|
||||
} else if (batch <= 32) {
|
||||
threads = 32;
|
||||
} else if (batch <= 64) {
|
||||
threads = 64;
|
||||
} else if (batch <= 128) {
|
||||
threads = 128;
|
||||
} else if (batch <= 256) {
|
||||
threads = 256;
|
||||
} else {
|
||||
threads = 256;
|
||||
}
|
||||
|
||||
int blocks = features;
|
||||
size_t shared_mem = 0;
|
||||
|
||||
// update_stats = track_running_stats && training
|
||||
// track_running_stats=False: 不更新
|
||||
// track_running_stats=True + training: 更新
|
||||
// track_running_stats=True + eval: 不会走到这里
|
||||
bool update_stats = track_running_stats && training;
|
||||
|
||||
batchnorm_forward_train_kernel_optimized<<<blocks, threads, shared_mem, stream>>>(
|
||||
x_contig.data_ptr<float>(),
|
||||
weight_contig.data_ptr<float>(),
|
||||
bias_contig.data_ptr<float>(),
|
||||
running_mean.data_ptr<float>(),
|
||||
running_var.data_ptr<float>(),
|
||||
y.data_ptr<float>(),
|
||||
batch,
|
||||
features,
|
||||
static_cast<float>(eps),
|
||||
static_cast<float>(momentum),
|
||||
update_stats
|
||||
);
|
||||
} else {
|
||||
// 使用 running stats(track_running_stats=True + eval 模式)
|
||||
int total = batch * features;
|
||||
int threads = 256;
|
||||
int blocks;
|
||||
|
||||
if (total <= 4096) {
|
||||
blocks = (total + threads - 1) / threads;
|
||||
} else {
|
||||
blocks = min(1024, (total + threads * 4 - 1) / (threads * 4));
|
||||
}
|
||||
|
||||
batchnorm_forward_eval_kernel_optimized<<<blocks, threads, 0, stream>>>(
|
||||
x_contig.data_ptr<float>(),
|
||||
weight_contig.data_ptr<float>(),
|
||||
bias_contig.data_ptr<float>(),
|
||||
running_mean.data_ptr<float>(),
|
||||
running_var.data_ptr<float>(),
|
||||
y.data_ptr<float>(),
|
||||
batch,
|
||||
features,
|
||||
static_cast<float>(eps)
|
||||
);
|
||||
}
|
||||
|
||||
C10_CUDA_KERNEL_LAUNCH_CHECK();
|
||||
return y;
|
||||
}
|
||||
"""
|
||||
|
||||
batchnorm_cpp_source = r"""
|
||||
torch::Tensor batchnorm_cuda_forward(
|
||||
torch::Tensor x,
|
||||
torch::Tensor weight,
|
||||
torch::Tensor bias,
|
||||
torch::Tensor running_mean,
|
||||
torch::Tensor running_var,
|
||||
bool training,
|
||||
double momentum,
|
||||
double eps,
|
||||
bool track_running_stats
|
||||
);
|
||||
"""
|
||||
|
||||
batchnorm_cuda = load_inline(
|
||||
name="batchnorm_cuda_ext",
|
||||
cpp_sources=batchnorm_cpp_source,
|
||||
cuda_sources=batchnorm_source,
|
||||
functions=["batchnorm_cuda_forward"],
|
||||
verbose=True
|
||||
)
|
||||
|
||||
class ModelNew(nn.Module):
|
||||
"""
|
||||
Model performing matrix multiplication followed by custom CUDA BatchNorm and ReLU.
|
||||
Optimized with Warp-level reduction (Plan 1) and thread configuration (Plan 2).
|
||||
"""
|
||||
def __init__(self, mat_weight: torch.Tensor, bn_weight: torch.Tensor, bn_bias: torch.Tensor,
|
||||
eps: float = 1e-5, momentum: float = 0.1, track_running_stats: bool = True):
|
||||
super().__init__()
|
||||
if mat_weight.dim() != 2:
|
||||
raise ValueError("mat_weight must be a 2D tensor [input_dim, output_dim].")
|
||||
if bn_weight.dim() != 1 or bn_bias.dim() != 1:
|
||||
raise ValueError("BatchNorm weight and bias must be 1D.")
|
||||
if bn_weight.size(0) != mat_weight.size(1):
|
||||
raise ValueError("BatchNorm parameter size must match output_dim.")
|
||||
if bn_weight.size(0) != bn_bias.size(0):
|
||||
raise ValueError("BatchNorm weight and bias must share shape.")
|
||||
|
||||
self.weight = nn.Parameter(mat_weight.clone())
|
||||
self.bn_weight = nn.Parameter(bn_weight.clone())
|
||||
self.bn_bias = nn.Parameter(bn_bias.clone())
|
||||
self.eps = eps
|
||||
self.momentum = momentum
|
||||
self.track_running_stats = track_running_stats
|
||||
|
||||
# 无论 track_running_stats 是什么,都创建 buffer
|
||||
self.register_buffer('running_mean', torch.zeros(bn_weight.size(0)))
|
||||
self.register_buffer('running_var', torch.ones(bn_weight.size(0)))
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
if not x.is_cuda:
|
||||
raise ValueError("Input must be a CUDA tensor.")
|
||||
if not self.weight.is_cuda:
|
||||
raise ValueError("Model weight must be on CUDA.")
|
||||
if not self.bn_weight.is_cuda or not self.bn_bias.is_cuda:
|
||||
raise ValueError("BatchNorm parameters must be on CUDA.")
|
||||
|
||||
x = torch.matmul(x, self.weight)
|
||||
|
||||
# 传递 track_running_stats 参数到 CUDA kernel
|
||||
x = batchnorm_cuda.batchnorm_cuda_forward(
|
||||
x,
|
||||
self.bn_weight,
|
||||
self.bn_bias,
|
||||
self.running_mean,
|
||||
self.running_var,
|
||||
self.training,
|
||||
self.momentum,
|
||||
self.eps,
|
||||
self.track_running_stats
|
||||
)
|
||||
|
||||
return torch.relu(x)
|
||||
|
|
@ -1,46 +0,0 @@
|
|||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
class BatchNormModel(nn.Module):
|
||||
"""
|
||||
Model that performs matrix multiplication followed by BatchNorm and ReLU activation.
|
||||
"""
|
||||
def __init__(self, weight, num_features=2048, eps=1e-5, momentum=0.1, track_running_stats=True):
|
||||
super(BatchNormModel, self).__init__()
|
||||
self.weight = nn.Parameter(weight)
|
||||
# 设置 track_running_stats=True 以跟踪运行时统计量
|
||||
self.bn = nn.BatchNorm1d(
|
||||
num_features,
|
||||
eps=eps,
|
||||
momentum=momentum,
|
||||
track_running_stats=track_running_stats
|
||||
)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Performs matrix multiplication, applies BatchNorm, then ReLU activation.
|
||||
|
||||
Args:
|
||||
x (torch.Tensor): Input tensor of shape [batch_size, input_dim]
|
||||
|
||||
Returns:
|
||||
torch.Tensor: Output tensor of shape [batch_size, output_dim]
|
||||
"""
|
||||
x = torch.matmul(x, self.weight)
|
||||
x = self.bn(x)
|
||||
return torch.relu(x)
|
||||
|
||||
# 添加别名以便在 run_code.py 中使用
|
||||
Model = BatchNormModel
|
||||
|
||||
batch_size = 16
|
||||
input_dim = 1024
|
||||
output_dim = 2048
|
||||
|
||||
def get_inputs():
|
||||
x = torch.randn(batch_size, input_dim)
|
||||
return [x]
|
||||
|
||||
def get_init_inputs():
|
||||
weight = torch.randn(input_dim, output_dim)
|
||||
return [weight]
|
||||
|
|
@ -1,61 +0,0 @@
|
|||
You write custom CUDA kernels to replace the pytorch operators in the given architecture to get speedups.
|
||||
|
||||
You have complete freedom to choose the set of operators you want to replace. You may make the decision to replace some operators with custom CUDA kernels and leave others unchanged. You may replace multiple operators with custom implementations, consider operator fusion opportunities (combining multiple operators into a single kernel, for example, combining matmul+relu), or algorithmic changes (such as online softmax). You are only limited by your imagination.
|
||||
|
||||
Here's an example to show you the syntax of inline embedding custom CUDA operators in torch: The example given architecture is:
|
||||
|
||||
```python
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
class Model(nn.Module):
|
||||
def __init__(self, num_features) -> None:
|
||||
super().__init__()
|
||||
self.bn = nn.BatchNorm1d(num_features)
|
||||
|
||||
def forward(self, x):
|
||||
return self.bn(x)
|
||||
|
||||
|
||||
def get_inputs():
|
||||
x = torch.randn(16, 2048).cuda()
|
||||
return [x]
|
||||
|
||||
|
||||
def get_init_inputs():
|
||||
return [2048]
|
||||
```
|
||||
|
||||
You are given the following architecture to implement BatchNorm1d with custom CUDA kernel:
|
||||
|
||||
```python
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
|
||||
class Model(nn.Module):
|
||||
"""使用 PyTorch BatchNorm1d 的基准实现。"""
|
||||
|
||||
def __init__(self, num_features: int, eps: float = 1e-5, momentum: float = 0.1):
|
||||
super().__init__()
|
||||
self.bn = nn.BatchNorm1d(num_features, eps=eps, momentum=momentum)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
"""对输入做 BatchNorm,输出形状与输入一致。"""
|
||||
return self.bn(x)
|
||||
|
||||
|
||||
batch_size = 16
|
||||
feature_dim = 2048
|
||||
|
||||
|
||||
def get_inputs():
|
||||
x = torch.randn(batch_size, feature_dim)
|
||||
return [x]
|
||||
|
||||
|
||||
def get_init_inputs():
|
||||
return [feature_dim]
|
||||
```
|
||||
|
|
@ -1,101 +0,0 @@
|
|||
###########################################################
|
||||
# 性能和精度验证程序
|
||||
###########################################################
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import time
|
||||
from batchnorm1d_torch import Model as TorchModel, get_inputs, get_init_inputs
|
||||
from batchnorm1d_cuda import ModelNew as CudaModel
|
||||
|
||||
def run_benchmark():
|
||||
# 检查 CUDA 是否可用
|
||||
if not torch.cuda.is_available():
|
||||
print("CUDA 不可用,请确保您有可用的 NVIDIA GPU 并已正确安装 PyTorch CUDA 版本。")
|
||||
return
|
||||
else:
|
||||
device = torch.device("cuda")
|
||||
|
||||
# 初始化模型
|
||||
# 获取 torch 版本的初始化参数
|
||||
torch_init_inputs = get_init_inputs()
|
||||
weight = torch_init_inputs[0].cuda(device=device)
|
||||
|
||||
# 为 BatchNorm 准备参数
|
||||
batch_size = 16
|
||||
input_dim = 1024
|
||||
output_dim = 2048
|
||||
bn_weight = torch.ones(output_dim, device=device, dtype=torch.float32)
|
||||
bn_bias = torch.zeros(output_dim, device=device, dtype=torch.float32)
|
||||
|
||||
# 初始化输入数据
|
||||
inputs = get_inputs()
|
||||
inputs = [x.cuda(device=device) if isinstance(x, torch.Tensor) else x for x in inputs]
|
||||
|
||||
# 初始化两个模型
|
||||
track_bool = True
|
||||
torch_model = TorchModel(weight.clone(), num_features=output_dim, eps=1e-5, track_running_stats=track_bool).cuda()
|
||||
cuda_model = CudaModel(weight.clone(), bn_weight.clone(), bn_bias.clone(), eps=1e-5, track_running_stats=track_bool).cuda()
|
||||
|
||||
torch_model.eval()
|
||||
cuda_model.eval()
|
||||
|
||||
print("-------------------- 精度对齐验证 --------------------")
|
||||
with torch.no_grad():
|
||||
output_torch = torch_model(*inputs)
|
||||
output_cuda = cuda_model(*inputs)
|
||||
|
||||
# 更严格的精度检查
|
||||
abs_diff = (output_torch - output_cuda).abs()
|
||||
max_diff = abs_diff.max().item()
|
||||
mean_diff = abs_diff.mean().item()
|
||||
|
||||
print(f"最大差异: {max_diff:.6f}")
|
||||
print(f"平均差异: {mean_diff:.6f}")
|
||||
|
||||
precision_flag = torch.allclose(output_torch, output_cuda, rtol=1e-03, atol=1e-03)
|
||||
if precision_flag:
|
||||
print("✅ 精度对齐:两个模型的输出结果非常接近。")
|
||||
else:
|
||||
print("❌ 精度不一致!")
|
||||
|
||||
print("\n-------------------- 性能加速比测试 --------------------")
|
||||
print(f"track_running_stats = {track_bool}")
|
||||
num_iterations = 1000 # 增加迭代次数以获得更准确的时间测量
|
||||
|
||||
# Warm up
|
||||
print("预热中...")
|
||||
for _ in range(100):
|
||||
with torch.no_grad():
|
||||
_ = torch_model(*inputs)
|
||||
_ = cuda_model(*inputs)
|
||||
|
||||
# PyTorch 模型计时
|
||||
torch.cuda.synchronize()
|
||||
start_time = time.time()
|
||||
with torch.no_grad():
|
||||
for _ in range(num_iterations):
|
||||
_ = torch_model(*inputs)
|
||||
torch.cuda.synchronize()
|
||||
torch_time = (time.time() - start_time) / num_iterations
|
||||
|
||||
# 自定义 CUDA 内核计时
|
||||
torch.cuda.synchronize()
|
||||
start_time = time.time()
|
||||
with torch.no_grad():
|
||||
for _ in range(num_iterations):
|
||||
_ = cuda_model(*inputs)
|
||||
torch.cuda.synchronize()
|
||||
cuda_time = (time.time() - start_time) / num_iterations
|
||||
|
||||
print(f"PyTorch 平均执行时间: {torch_time*1000:.4f} 毫秒")
|
||||
print(f"自定义 CUDA BatchNorm 平均执行时间: {cuda_time*1000:.4f} 毫秒")
|
||||
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,223 +0,0 @@
|
|||
import torch
|
||||
import torch.nn as nn
|
||||
from torch.utils.cpp_extension import load_inline
|
||||
|
||||
# CUDA implementation of Conv2D (tiled + shared memory + output-channel blocking)
|
||||
conv2d_source = r"""
|
||||
#include <torch/extension.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <math.h>
|
||||
|
||||
#ifndef CHECK_CUDA
|
||||
#define CHECK_CUDA(x) TORCH_CHECK((x).is_cuda(), #x " must be a CUDA tensor")
|
||||
#endif
|
||||
#ifndef CHECK_CONTIGUOUS
|
||||
#define CHECK_CONTIGUOUS(x) TORCH_CHECK((x).is_contiguous(), #x " must be contiguous")
|
||||
#endif
|
||||
#ifndef CHECK_FLOAT
|
||||
#define CHECK_FLOAT(x) TORCH_CHECK((x).scalar_type() == at::kFloat, #x " must be float32")
|
||||
#endif
|
||||
|
||||
// 每个block计算一个 (b, oc_group) 上的输出tile,复用输入tile计算 OC_TILE 个输出通道
|
||||
template<int BLOCK_X, int BLOCK_Y, int OC_TILE>
|
||||
__global__ void conv2d_tiled_kernel_oc(
|
||||
const float* __restrict__ input, // [B, C_in, H, W]
|
||||
const float* __restrict__ weight, // [C_out, C_in, K, K]
|
||||
const float* __restrict__ bias, // [C_out] or nullptr
|
||||
float* __restrict__ output, // [B, C_out, H_out, W_out]
|
||||
int B, int C_in, int C_out,
|
||||
int H, int W, int K, int H_out, int W_out,
|
||||
bool has_bias
|
||||
) {
|
||||
// grid.z = B * ceil_div(C_out, OC_TILE)
|
||||
int groups = (C_out + OC_TILE - 1) / OC_TILE;
|
||||
int b = blockIdx.z / groups;
|
||||
int og = blockIdx.z % groups; // 输出通道组编号
|
||||
int co0 = og * OC_TILE; // 本组起始输出通道
|
||||
|
||||
int ow0 = blockIdx.x * BLOCK_X;
|
||||
int oh0 = blockIdx.y * BLOCK_Y;
|
||||
|
||||
int ow = ow0 + threadIdx.x;
|
||||
int oh = oh0 + threadIdx.y;
|
||||
|
||||
extern __shared__ float smem[];
|
||||
// 输入tile大小:(BLOCK_Y+K-1) x (BLOCK_X+K-1)
|
||||
int tile_w = BLOCK_X + K - 1;
|
||||
int tile_h = BLOCK_Y + K - 1;
|
||||
float* tile = smem; // tile_h * tile_w
|
||||
float* w_sh = tile + tile_h * tile_w; // OC_TILE * K * K
|
||||
// w_sh 布局: [oc_local][K*K]
|
||||
|
||||
// 累加器:每线程维护 OC_TILE 个通道
|
||||
float acc[OC_TILE];
|
||||
#pragma unroll
|
||||
for (int oc = 0; oc < OC_TILE; ++oc) {
|
||||
int co = co0 + oc;
|
||||
acc[oc] = (has_bias && co < C_out) ? bias[co] : 0.0f;
|
||||
}
|
||||
|
||||
bool valid_xy = (oh < H_out) && (ow < W_out);
|
||||
|
||||
// 遍历输入通道
|
||||
for (int ci = 0; ci < C_in; ++ci) {
|
||||
// 1) 加载本组 OC_TILE 的权重到共享内存
|
||||
int total_w = OC_TILE * K * K;
|
||||
for (int t = threadIdx.y * BLOCK_X + threadIdx.x; t < total_w; t += BLOCK_X * BLOCK_Y) {
|
||||
int oc = t / (K*K);
|
||||
int rem = t % (K*K);
|
||||
int kh = rem / K;
|
||||
int kw = rem % K;
|
||||
int co = co0 + oc;
|
||||
float wv = 0.0f;
|
||||
if (co < C_out) {
|
||||
int w_idx = ((co * C_in + ci) * K + kh) * K + kw;
|
||||
wv = weight[w_idx];
|
||||
}
|
||||
w_sh[t] = wv;
|
||||
}
|
||||
|
||||
// 2) 加载输入tile到共享内存(该tile将被 OC_TILE 个输出通道复用)
|
||||
int ih0 = oh0;
|
||||
int iw0 = ow0;
|
||||
for (int th = threadIdx.y; th < tile_h; th += BLOCK_Y) {
|
||||
int ih = ih0 + th;
|
||||
bool in_h = (ih >= 0) && (ih < H);
|
||||
for (int tw = threadIdx.x; tw < tile_w; tw += BLOCK_X) {
|
||||
int iw = iw0 + tw;
|
||||
bool in_w = (iw >= 0) && (iw < W);
|
||||
float v = 0.0f;
|
||||
if (in_h && in_w) {
|
||||
int in_idx = (((b * C_in + ci) * H + ih) * W + iw);
|
||||
v = input[in_idx];
|
||||
}
|
||||
tile[th * tile_w + tw] = v;
|
||||
}
|
||||
}
|
||||
|
||||
__syncthreads();
|
||||
|
||||
// 3) 计算:同一输入tile,对 OC_TILE 个输出通道分别累加
|
||||
if (valid_xy) {
|
||||
int t_base = threadIdx.y * tile_w + threadIdx.x;
|
||||
#pragma unroll
|
||||
for (int kh = 0; kh < K; ++kh) {
|
||||
int t_row = t_base + kh * tile_w;
|
||||
#pragma unroll
|
||||
for (int kw = 0; kw < K; ++kw) {
|
||||
float val = tile[t_row + kw];
|
||||
#pragma unroll
|
||||
for (int oc = 0; oc < OC_TILE; ++oc) {
|
||||
float wv = w_sh[oc * (K*K) + kh * K + kw];
|
||||
acc[oc] = fmaf(val, wv, acc[oc]);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
__syncthreads(); // 保护下一个 ci 的加载
|
||||
}
|
||||
|
||||
// 4) 写回输出
|
||||
if (valid_xy) {
|
||||
int base = (b * C_out) * (H_out * W_out);
|
||||
int out_offset = oh * W_out + ow;
|
||||
#pragma unroll
|
||||
for (int oc = 0; oc < OC_TILE; ++oc) {
|
||||
int co = co0 + oc;
|
||||
if (co < C_out) {
|
||||
int out_idx = base + co * (H_out * W_out) + out_offset;
|
||||
output[out_idx] = acc[oc];
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// C++ wrapper
|
||||
torch::Tensor conv2d_cuda(
|
||||
torch::Tensor input,
|
||||
torch::Tensor weight,
|
||||
torch::Tensor bias
|
||||
) {
|
||||
CHECK_CUDA(input);
|
||||
CHECK_CUDA(weight);
|
||||
CHECK_CUDA(bias);
|
||||
CHECK_CONTIGUOUS(input);
|
||||
CHECK_CONTIGUOUS(weight);
|
||||
CHECK_CONTIGUOUS(bias);
|
||||
CHECK_FLOAT(input);
|
||||
CHECK_FLOAT(weight);
|
||||
CHECK_FLOAT(bias);
|
||||
|
||||
int B = input.size(0);
|
||||
int C_in = input.size(1);
|
||||
int H = input.size(2);
|
||||
int W = input.size(3);
|
||||
|
||||
int C_out = weight.size(0);
|
||||
int K = weight.size(2);
|
||||
TORCH_CHECK(weight.size(3) == K, "Kernel must be square");
|
||||
|
||||
int H_out = H - K + 1;
|
||||
int W_out = W - K + 1;
|
||||
|
||||
auto output = torch::empty({B, C_out, H_out, W_out}, input.options());
|
||||
|
||||
// 参数可按GPU微调:32x8, 16x16 等
|
||||
const int BLOCK_X = 16;
|
||||
const int BLOCK_Y = 16;
|
||||
const int OC_TILE = 4;
|
||||
|
||||
dim3 block(BLOCK_X, BLOCK_Y, 1);
|
||||
int groups = (C_out + OC_TILE - 1) / OC_TILE;
|
||||
dim3 grid((W_out + BLOCK_X - 1) / BLOCK_X,
|
||||
(H_out + BLOCK_Y - 1) / BLOCK_Y,
|
||||
B * groups);
|
||||
|
||||
size_t tile_w = BLOCK_X + K - 1;
|
||||
size_t tile_h = BLOCK_Y + K - 1;
|
||||
size_t shmem_elems = tile_w * tile_h + OC_TILE * K * K;
|
||||
size_t shmem_bytes = shmem_elems * sizeof(float);
|
||||
|
||||
bool has_bias = bias.numel() > 0;
|
||||
|
||||
conv2d_tiled_kernel_oc<BLOCK_X, BLOCK_Y, OC_TILE><<<grid, block, shmem_bytes>>>(
|
||||
input.data_ptr<float>(),
|
||||
weight.data_ptr<float>(),
|
||||
has_bias ? bias.data_ptr<float>() : nullptr,
|
||||
output.data_ptr<float>(),
|
||||
B, C_in, C_out, H, W, K, H_out, W_out, has_bias
|
||||
);
|
||||
|
||||
auto err = cudaGetLastError();
|
||||
TORCH_CHECK(err == cudaSuccess, "conv2d kernel launch failed: ", cudaGetErrorString(err));
|
||||
return output;
|
||||
}
|
||||
"""
|
||||
|
||||
conv2d_cpp_source = r"""
|
||||
torch::Tensor conv2d_cuda(torch::Tensor input, torch::Tensor weight, torch::Tensor bias);
|
||||
"""
|
||||
|
||||
# Compile with O3 (no fast-math to keep FP32 parity)
|
||||
conv2d = load_inline(
|
||||
name="conv2d_tiled_opt_oc",
|
||||
cpp_sources=conv2d_cpp_source,
|
||||
cuda_sources=conv2d_source,
|
||||
functions=["conv2d_cuda"],
|
||||
verbose=False,
|
||||
extra_cuda_cflags=["-O3"]
|
||||
)
|
||||
|
||||
class ModelNew(nn.Module):
|
||||
def __init__(self, weight, bias=None):
|
||||
super(ModelNew, self).__init__()
|
||||
self.weight = nn.Parameter(weight)
|
||||
self.bias = nn.Parameter(bias) if bias is not None else nn.Parameter(torch.empty(0, device=weight.device, dtype=weight.dtype))
|
||||
self.conv2d = conv2d
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
x = x.contiguous()
|
||||
w = self.weight.contiguous()
|
||||
b = self.bias.contiguous() if self.bias is not None else torch.empty(0, device=x.device, dtype=x.dtype)
|
||||
return self.conv2d.conv2d_cuda(x, w, b)
|
||||
|
|
@ -1,40 +0,0 @@
|
|||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
class Model(nn.Module):
|
||||
"""
|
||||
Model that performs 2D convolution operation.
|
||||
"""
|
||||
def __init__(self, weight, bias=None):
|
||||
super(Model, self).__init__()
|
||||
self.weight = nn.Parameter(weight)
|
||||
self.bias = nn.Parameter(bias) if bias is not None else None
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Performs 2D convolution.
|
||||
|
||||
Args:
|
||||
x (torch.Tensor): Input tensor of shape [batch_size, in_channels, height, width]
|
||||
|
||||
Returns:
|
||||
torch.Tensor: Output tensor of shape [batch_size, out_channels, out_height, out_width]
|
||||
"""
|
||||
return torch.nn.functional.conv2d(x, self.weight, self.bias, stride=1, padding=0)
|
||||
|
||||
# Hyperparameters
|
||||
batch_size = 4
|
||||
in_channels = 3
|
||||
out_channels = 64
|
||||
height = 32
|
||||
width = 32
|
||||
kernel_size = 3
|
||||
|
||||
def get_inputs():
|
||||
x = torch.randn(batch_size, in_channels, height, width)
|
||||
return [x]
|
||||
|
||||
def get_init_inputs():
|
||||
weight = torch.randn(out_channels, in_channels, kernel_size, kernel_size)
|
||||
bias = torch.randn(out_channels)
|
||||
return [weight, bias]
|
||||
|
|
@ -1,92 +0,0 @@
|
|||
###########################################################
|
||||
# 性能和精度验证程序
|
||||
###########################################################
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import time
|
||||
from conv2d_torch import Model, get_inputs, get_init_inputs
|
||||
from conv2d_cuda import ModelNew
|
||||
|
||||
# 禁用 TF32,确保与自定义 FP32 核精度对齐
|
||||
torch.backends.cuda.matmul.allow_tf32 = False
|
||||
torch.backends.cudnn.allow_tf32 = False
|
||||
torch.backends.cudnn.deterministic = True
|
||||
|
||||
def _time_cuda_model(fn, inputs, iters=300, warmup=50):
|
||||
torch.cuda.synchronize()
|
||||
for _ in range(warmup):
|
||||
_ = fn(*inputs)
|
||||
torch.cuda.synchronize()
|
||||
start = torch.cuda.Event(enable_timing=True)
|
||||
end = torch.cuda.Event(enable_timing=True)
|
||||
start.record()
|
||||
for _ in range(iters):
|
||||
_ = fn(*inputs)
|
||||
end.record()
|
||||
torch.cuda.synchronize()
|
||||
ms = start.elapsed_time(end) / iters # 平均每次毫秒
|
||||
return ms / 1000.0 # 转为秒
|
||||
|
||||
def run_benchmark():
|
||||
# 检查 CUDA 是否可用
|
||||
if not torch.cuda.is_available():
|
||||
print("CUDA 不可用,请确保您有可用的 NVIDIA GPU 并已正确安装 PyTorch CUDA 版本。")
|
||||
return
|
||||
else:
|
||||
device = torch.device("cuda")
|
||||
|
||||
torch.backends.cudnn.benchmark = True
|
||||
|
||||
# 初始化模型
|
||||
init_inputs = get_init_inputs()
|
||||
init_inputs = [
|
||||
x.cuda(device=device) if isinstance(x, torch.Tensor) else x for x in init_inputs
|
||||
]
|
||||
inputs = get_inputs()
|
||||
inputs = [
|
||||
x.cuda(device=device) if isinstance(x, torch.Tensor) else x for x in inputs
|
||||
]
|
||||
|
||||
torch_model = Model(*init_inputs).cuda()
|
||||
cuda_model = ModelNew(*init_inputs).cuda()
|
||||
|
||||
torch_model.eval()
|
||||
cuda_model.eval()
|
||||
|
||||
print("-------------------- 精度对齐验证 --------------------")
|
||||
with torch.no_grad():
|
||||
output_torch = torch_model(*inputs)
|
||||
output_cuda = cuda_model(*inputs)
|
||||
|
||||
abs_diff = (output_torch - output_cuda).abs()
|
||||
max_diff = abs_diff.max().item()
|
||||
mean_diff = abs_diff.mean().item()
|
||||
|
||||
print(f"最大差异: {max_diff:.6f}")
|
||||
print(f"平均差异: {mean_diff:.6f}")
|
||||
|
||||
precision_flag = torch.allclose(output_torch, output_cuda, rtol=1e-5, atol=1e-5)
|
||||
if precision_flag:
|
||||
print("✅ 精度对齐:两个模型的输出结果非常接近。")
|
||||
else:
|
||||
print("❌ 精度不一致!")
|
||||
|
||||
print("\n-------------------- 性能加速比测试 --------------------")
|
||||
num_iterations = 300
|
||||
|
||||
# 计时
|
||||
torch_time = _time_cuda_model(torch_model, inputs, iters=num_iterations, warmup=50)
|
||||
cuda_time = _time_cuda_model(cuda_model, inputs, iters=num_iterations, warmup=50)
|
||||
|
||||
print(f"PyTorch (conv2d) 平均执行时间: {torch_time:.6f} 秒")
|
||||
print(f"自定义 CUDA conv2d 平均执行时间: {cuda_time:.6f} 秒")
|
||||
speedup = 0.0
|
||||
if cuda_time > 0:
|
||||
speedup = torch_time / cuda_time
|
||||
print(f"相对 PyTorch 加速比: {speedup:.2f}x")
|
||||
else:
|
||||
print("CUDA 内核执行时间为0,无法计算加速比。")
|
||||
return precision_flag, speedup
|
||||
|
||||
if __name__ == "__main__":
|
||||
precision_flag, speedup = run_benchmark()
|
||||
|
|
@ -1,29 +0,0 @@
|
|||
You write custom CUDA kernels to replace the pytorch operators in the given architecture to get speedups.
|
||||
|
||||
You have complete freedom to choose the set of operators you want to replace. You may make the decision to replace some operators with custom CUDA kernels and leave others unchanged. You may replace multiple operators with custom implementations, consider operator fusion opportunities (combining multiple operators into a single kernel, for example, combining matmul+relu), or algorithmic changes (such as online softmax). You are only limited by your imagination.
|
||||
|
||||
Here's an example to show you the syntax of inline embedding custom CUDA operators in torch: The example given architecture is:
|
||||
|
||||
```python
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
class Model(nn.Module):
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.norm = nn.RMSNorm(128)
|
||||
|
||||
def forward(self, x):
|
||||
return self.norm(x)
|
||||
|
||||
|
||||
def get_inputs():
|
||||
x = torch.randn(1, 128).cuda()
|
||||
return [x]
|
||||
|
||||
|
||||
def get_init_inputs():
|
||||
return []
|
||||
```
|
||||
|
|
@ -1,102 +0,0 @@
|
|||
import torch
|
||||
import torch.nn as nn
|
||||
from torch.utils.cpp_extension import load_inline
|
||||
|
||||
rmsnorm_source = """
|
||||
#include <torch/extension.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <math.h>
|
||||
|
||||
__global__ void rmsnorm_kernel(
|
||||
const float* __restrict__ x,
|
||||
const float* __restrict__ weight,
|
||||
float* __restrict__ y,
|
||||
int batch,
|
||||
int features,
|
||||
float eps
|
||||
) {
|
||||
int row = blockIdx.x;
|
||||
if (row >= batch) return;
|
||||
int tid = threadIdx.x;
|
||||
|
||||
extern __shared__ float sdata[];
|
||||
float sum_sq = 0.0f;
|
||||
|
||||
for (int i = tid; i < features; i += blockDim.x) {
|
||||
float v = x[row * features + i];
|
||||
sum_sq += v * v;
|
||||
}
|
||||
sdata[tid] = sum_sq;
|
||||
__syncthreads();
|
||||
|
||||
for (int offset = blockDim.x >> 1; offset > 0; offset >>= 1) {
|
||||
if (tid < offset) {
|
||||
sdata[tid] += sdata[tid + offset];
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
float rms = rsqrtf(sdata[0] / features + eps);
|
||||
__syncthreads(); // 保证 rms 可见
|
||||
|
||||
for (int i = tid; i < features; i += blockDim.x) {
|
||||
float v = x[row * features + i];
|
||||
float w = weight[i];
|
||||
y[row * features + i] = v * rms * w;
|
||||
}
|
||||
}
|
||||
|
||||
torch::Tensor rmsnorm_cuda(torch::Tensor x, torch::Tensor weight, float eps) {
|
||||
TORCH_CHECK(x.is_cuda(), "x 必须是 CUDA 张量");
|
||||
TORCH_CHECK(weight.is_cuda(), "weight 必须是 CUDA 张量");
|
||||
TORCH_CHECK(x.dim() == 2, "当前内核仅支持二维输入张量");
|
||||
TORCH_CHECK(weight.dim() == 1, "RMSNorm 权重必须是一维向量");
|
||||
TORCH_CHECK(x.size(1) == weight.size(0), "输入最后一维与权重长度不匹配");
|
||||
|
||||
int batch = x.size(0);
|
||||
int features = x.size(1);
|
||||
|
||||
auto y = torch::empty_like(x);
|
||||
|
||||
int threads = 256;
|
||||
if (features < threads) {
|
||||
threads = 1;
|
||||
while (threads < features) threads <<= 1;
|
||||
if (threads < 32) threads = 32;
|
||||
}
|
||||
size_t shared = threads * sizeof(float);
|
||||
|
||||
rmsnorm_kernel<<<batch, threads, shared>>>(
|
||||
x.data_ptr<float>(),
|
||||
weight.data_ptr<float>(),
|
||||
y.data_ptr<float>(),
|
||||
batch,
|
||||
features,
|
||||
eps
|
||||
);
|
||||
return y;
|
||||
}
|
||||
"""
|
||||
|
||||
rmsnorm_cpp_source = """
|
||||
torch::Tensor rmsnorm_cuda(torch::Tensor x, torch::Tensor weight, float eps);
|
||||
"""
|
||||
|
||||
rmsnorm = load_inline(
|
||||
name="rmsnorm",
|
||||
cpp_sources=rmsnorm_cpp_source,
|
||||
cuda_sources=rmsnorm_source,
|
||||
functions=["rmsnorm_cuda"],
|
||||
verbose=True
|
||||
)
|
||||
|
||||
class ModelNew(nn.Module):
|
||||
def __init__(self, weight: torch.Tensor, eps: float = 1e-6):
|
||||
super().__init__()
|
||||
if weight.dim() != 1:
|
||||
raise ValueError("RMSNorm 权重必须是一维向量。")
|
||||
self.weight = nn.Parameter(weight.clone())
|
||||
self.eps = eps
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
return rmsnorm.rmsnorm_cuda(x, self.weight, self.eps)
|
||||
|
|
@ -1,33 +0,0 @@
|
|||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
|
||||
class Model(nn.Module):
|
||||
"""使用 PyTorch RMSNorm 的基准实现。"""
|
||||
|
||||
def __init__(self, weight: torch.Tensor, eps: float = 1e-6):
|
||||
super().__init__()
|
||||
if weight.dim() != 1:
|
||||
raise ValueError("RMSNorm 权重必须是一维向量。")
|
||||
feature_dim = weight.shape[0]
|
||||
self.rmsnorm = nn.RMSNorm(feature_dim, eps=eps, elementwise_affine=True)
|
||||
with torch.no_grad():
|
||||
self.rmsnorm.weight.copy_(weight)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
"""直接对输入做 RMSNorm,输出形状与输入一致。"""
|
||||
return self.rmsnorm(x)
|
||||
|
||||
|
||||
batch_size = 16
|
||||
feature_dim = 2048
|
||||
|
||||
|
||||
def get_inputs():
|
||||
x = torch.randn(batch_size, feature_dim)
|
||||
return [x]
|
||||
|
||||
|
||||
def get_init_inputs():
|
||||
weight = torch.randn(feature_dim)
|
||||
return [weight]
|
||||
|
|
@ -1,88 +0,0 @@
|
|||
###########################################################
|
||||
# 性能和精度验证程序
|
||||
###########################################################
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import time
|
||||
from rmsnorm_torch import Model, get_inputs, get_init_inputs
|
||||
from rmsnorm_cuda import ModelNew
|
||||
|
||||
def run_benchmark():
|
||||
# 检查 CUDA 是否可用
|
||||
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)
|
||||
|
||||
# 更严格的精度检查
|
||||
abs_diff = (output_torch - output_cuda).abs()
|
||||
max_diff = abs_diff.max().item()
|
||||
mean_diff = abs_diff.mean().item()
|
||||
|
||||
print(f"最大差异: {max_diff:.6f}")
|
||||
print(f"平均差异: {mean_diff:.6f}")
|
||||
|
||||
precision_flag = torch.allclose(output_torch, output_cuda, rtol=1e-05, atol=1e-05)
|
||||
if precision_flag:
|
||||
print("✅ 精度对齐:两个模型的输出结果非常接近。")
|
||||
else:
|
||||
print("❌ 精度不一致!")
|
||||
|
||||
print("\n-------------------- 性能加速比测试 --------------------")
|
||||
num_iterations = 1000 # 增加迭代次数以获得更准确的时间测量
|
||||
|
||||
# Warm up
|
||||
for _ in range(100):
|
||||
_ = torch_model(*inputs)
|
||||
_ = cuda_model(*inputs)
|
||||
|
||||
# PyTorch 模型计时
|
||||
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
|
||||
|
||||
# 自定义 CUDA 内核计时
|
||||
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 平均执行时间: {torch_time:.6f} 秒")
|
||||
print(f"自定义 CUDA ReLU 平均执行时间: {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,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,190 +0,0 @@
|
|||
import torch
|
||||
from torch.utils.cpp_extension import load_inline
|
||||
|
||||
# LayerNorm的CUDA实现
|
||||
layernorm_source = """
|
||||
#include <torch/extension.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <cuda_fp16.h>
|
||||
#include <cub/cub.cuh>
|
||||
|
||||
// 使用高度优化的LayerNorm实现,结合向量化和内存访问优化
|
||||
__global__ void layernorm_forward_kernel(
|
||||
const float* __restrict__ input,
|
||||
const float* __restrict__ gamma,
|
||||
const float* __restrict__ beta,
|
||||
float* __restrict__ output,
|
||||
int batch_size,
|
||||
int hidden_size,
|
||||
float eps) {
|
||||
|
||||
extern __shared__ float shared_mem[];
|
||||
float* shared_sum = shared_mem;
|
||||
float* shared_sum_sq = &shared_mem[blockDim.x];
|
||||
|
||||
int batch_idx = blockIdx.x;
|
||||
int tid = threadIdx.x;
|
||||
|
||||
// 使用向量化加载,每个线程处理4个元素
|
||||
float4 thread_sum = make_float4(0.0f, 0.0f, 0.0f, 0.0f);
|
||||
float4 thread_sum_sq = make_float4(0.0f, 0.0f, 0.0f, 0.0f);
|
||||
|
||||
// 向量化处理,提高内存带宽利用率
|
||||
for (int i = tid * 4; i < hidden_size; i += blockDim.x * 4) {
|
||||
if (i + 3 < hidden_size) {
|
||||
float4 vals = *reinterpret_cast<const float4*>(&input[batch_idx * hidden_size + i]);
|
||||
thread_sum.x += vals.x; thread_sum_sq.x += vals.x * vals.x;
|
||||
thread_sum.y += vals.y; thread_sum_sq.y += vals.y * vals.y;
|
||||
thread_sum.z += vals.z; thread_sum_sq.z += vals.z * vals.z;
|
||||
thread_sum.w += vals.w; thread_sum_sq.w += vals.w * vals.w;
|
||||
} else {
|
||||
// 处理剩余元素
|
||||
for (int j = 0; j < 4 && i + j < hidden_size; j++) {
|
||||
float val = input[batch_idx * hidden_size + i + j];
|
||||
thread_sum.x += val; thread_sum_sq.x += val * val;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 归约线程内的4个分量
|
||||
float thread_total_sum = thread_sum.x + thread_sum.y + thread_sum.z + thread_sum.w;
|
||||
float thread_total_sum_sq = thread_sum_sq.x + thread_sum_sq.y + thread_sum_sq.z + thread_sum_sq.w;
|
||||
|
||||
shared_sum[tid] = thread_total_sum;
|
||||
shared_sum_sq[tid] = thread_total_sum_sq;
|
||||
__syncthreads();
|
||||
|
||||
// 使用更高效的归约算法(树形归约)
|
||||
for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {
|
||||
if (tid < stride) {
|
||||
shared_sum[tid] += shared_sum[tid + stride];
|
||||
shared_sum_sq[tid] += shared_sum_sq[tid + stride];
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
// 计算全局统计量
|
||||
if (tid == 0) {
|
||||
float total_sum = shared_sum[0];
|
||||
float total_sum_sq = shared_sum_sq[0];
|
||||
float global_mean = total_sum / hidden_size;
|
||||
float global_variance = (total_sum_sq / hidden_size) - (global_mean * global_mean);
|
||||
|
||||
// 计算逆标准差
|
||||
float inv_std = rsqrtf(global_variance + eps);
|
||||
|
||||
// 存储到共享内存供所有线程使用
|
||||
shared_sum[0] = global_mean;
|
||||
shared_sum_sq[0] = inv_std;
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
float global_mean = shared_sum[0];
|
||||
float inv_std = shared_sum_sq[0];
|
||||
|
||||
// 应用LayerNorm,使用向量化存储
|
||||
for (int i = tid * 4; i < hidden_size; i += blockDim.x * 4) {
|
||||
if (i + 3 < hidden_size) {
|
||||
float4 vals = *reinterpret_cast<const float4*>(&input[batch_idx * hidden_size + i]);
|
||||
float4 normalized;
|
||||
normalized.x = (vals.x - global_mean) * inv_std;
|
||||
normalized.y = (vals.y - global_mean) * inv_std;
|
||||
normalized.z = (vals.z - global_mean) * inv_std;
|
||||
normalized.w = (vals.w - global_mean) * inv_std;
|
||||
|
||||
float4 result;
|
||||
result.x = normalized.x * gamma[i] + beta[i];
|
||||
result.y = normalized.y * gamma[i+1] + beta[i+1];
|
||||
result.z = normalized.z * gamma[i+2] + beta[i+2];
|
||||
result.w = normalized.w * gamma[i+3] + beta[i+3];
|
||||
|
||||
*reinterpret_cast<float4*>(&output[batch_idx * hidden_size + i]) = result;
|
||||
} else {
|
||||
// 处理剩余元素
|
||||
for (int j = 0; j < 4 && i + j < hidden_size; j++) {
|
||||
float val = input[batch_idx * hidden_size + i + j];
|
||||
float normalized = (val - global_mean) * inv_std;
|
||||
output[batch_idx * hidden_size + i + j] = normalized * gamma[i + j] + beta[i + j];
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
torch::Tensor layernorm_cuda_forward(
|
||||
torch::Tensor input,
|
||||
torch::Tensor gamma,
|
||||
torch::Tensor beta,
|
||||
float eps) {
|
||||
|
||||
auto batch_size = input.size(0);
|
||||
auto hidden_size = input.size(-1);
|
||||
|
||||
auto output = torch::empty_like(input);
|
||||
|
||||
// 优化线程块大小,根据hidden_size动态调整,使用更激进的优化
|
||||
int block_size = 256; // 固定使用256线程,适合大多数GPU架构
|
||||
if (hidden_size <= 512) {
|
||||
block_size = 128;
|
||||
} else if (hidden_size <= 1024) {
|
||||
block_size = 256;
|
||||
} else {
|
||||
block_size = 512;
|
||||
}
|
||||
|
||||
// 确保block_size不超过硬件限制
|
||||
block_size = min(1024, max(32, block_size));
|
||||
|
||||
int num_blocks = batch_size;
|
||||
int shared_mem_size = 2 * block_size * sizeof(float);
|
||||
|
||||
layernorm_forward_kernel<<<num_blocks, block_size, shared_mem_size>>>(
|
||||
input.data_ptr<float>(),
|
||||
gamma.data_ptr<float>(),
|
||||
beta.data_ptr<float>(),
|
||||
output.data_ptr<float>(),
|
||||
batch_size,
|
||||
hidden_size,
|
||||
eps
|
||||
);
|
||||
|
||||
return output;
|
||||
}
|
||||
"""
|
||||
|
||||
layernorm_cpp_source = """
|
||||
torch::Tensor layernorm_cuda_forward(torch::Tensor input, torch::Tensor gamma, torch::Tensor beta, float eps);
|
||||
"""
|
||||
|
||||
# 编译内联CUDA代码
|
||||
cuda_available = True
|
||||
try:
|
||||
layernorm_cuda = load_inline(
|
||||
name="layernorm_cuda",
|
||||
cpp_sources=layernorm_cpp_source,
|
||||
cuda_sources=layernorm_source,
|
||||
functions=["layernorm_cuda_forward"],
|
||||
verbose=True
|
||||
)
|
||||
except Exception as e:
|
||||
print(f"CUDA扩展加载失败: {e}")
|
||||
cuda_available = False
|
||||
layernorm_cuda = None
|
||||
|
||||
class ModelNew(torch.nn.Module):
|
||||
def __init__(self, normalized_shape, eps=1e-5):
|
||||
super(ModelNew, self).__init__()
|
||||
self.normalized_shape = normalized_shape
|
||||
self.eps = eps
|
||||
self.weight = torch.nn.Parameter(torch.ones(normalized_shape))
|
||||
self.bias = torch.nn.Parameter(torch.zeros(normalized_shape))
|
||||
|
||||
def forward(self, x):
|
||||
if cuda_available and layernorm_cuda is not None:
|
||||
# 使用真正的CUDA内核
|
||||
return layernorm_cuda.layernorm_cuda_forward(x, self.weight, self.bias, self.eps)
|
||||
else:
|
||||
# CPU回退实现,与PyTorch实现保持一致
|
||||
mean = x.mean(-1, keepdim=True)
|
||||
var = x.var(-1, unbiased=False, keepdim=True)
|
||||
normalized = (x - mean) / torch.sqrt(var + self.eps)
|
||||
return normalized * self.weight + self.bias
|
||||
|
|
@ -1,23 +0,0 @@
|
|||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
class Model(nn.Module):
|
||||
def __init__(self, normalized_shape, eps=1e-5):
|
||||
super(Model, self).__init__()
|
||||
self.normalized_shape = normalized_shape
|
||||
self.eps = eps
|
||||
|
||||
# 使用PyTorch内置的LayerNorm,这是标准实现
|
||||
self.layer_norm = nn.LayerNorm(normalized_shape, eps=eps)
|
||||
|
||||
def forward(self, x):
|
||||
# 使用PyTorch内置LayerNorm,这是标准的实现方式
|
||||
return self.layer_norm(x)
|
||||
|
||||
def get_inputs():
|
||||
# 使用更大的输入尺寸以获得更好的性能对比
|
||||
# 增加batch size和hidden size,模拟真实场景
|
||||
return [torch.randn(16, 16384)]
|
||||
|
||||
def get_init_inputs():
|
||||
return [16384]
|
||||
|
|
@ -1,27 +0,0 @@
|
|||
Write a custom CUDA kernel for Layer Normalization.
|
||||
|
||||
The standard LayerNorm operation is defined as:
|
||||
|
||||
y = (x - E[x]) / sqrt(Var[x] + epsilon) * gamma + beta
|
||||
|
||||
Where:
|
||||
- x is the input tensor
|
||||
- E[x] is the mean of x
|
||||
- Var[x] is the variance of x
|
||||
- epsilon is a small value for numerical stability
|
||||
- gamma and beta are learnable affine parameters
|
||||
|
||||
You should fuse the calculation of mean, variance, and the normalization into a single CUDA kernel. This avoids multiple passes over the data and reduces memory bandwidth usage.
|
||||
|
||||
You are given the following architecture:
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
class Model(nn.Module):
|
||||
def __init__(self, normalized_shape, eps=1e-5):
|
||||
super(Model, self).__init__()
|
||||
self.layer_norm = nn.LayerNorm(normalized_shape, eps=eps)
|
||||
|
||||
def forward(self, x):
|
||||
return self.layer_norm(x)
|
||||
|
|
@ -1,74 +0,0 @@
|
|||
###########################################################
|
||||
# 性能和精度验证程序
|
||||
###########################################################
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import time
|
||||
from layernorm_torchcode import Model,get_inputs,get_init_inputs
|
||||
from layernorm_cudacode import ModelNew
|
||||
|
||||
def run_benchmark():
|
||||
# 检查 CUDA 是否可用
|
||||
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("❌ 精度不一致!")
|
||||
|
||||
print("\n-------------------- 性能加速比测试 --------------------")
|
||||
num_iterations = 100
|
||||
|
||||
# PyTorch 模型计时
|
||||
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
|
||||
|
||||
# 自定义 CUDA 内核计时
|
||||
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 LayerNorm 平均执行时间: {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,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]
|
||||
Some files were not shown because too many files have changed in this diff Show More
Loading…
Reference in New Issue