用Roofline模型去分析pytorch和Triton算子
本文演示了如何用Roofline模型去分析pytorch和Triton算子
遗留问题:NVIDIA Nsight Compute中的Peak Work是怎么算出来的,又不是峰值算力 -> Nsight Compute 是怎么计算Roofline的呢
1.参考链接
2.测试环境
+---------------------------------------------------------------------------------------+
| NVIDIA-SMI 535.161.07 Driver Version: 535.161.07 CUDA Version: 12.2 |
|-----------------------------------------+----------------------+----------------------+
| GPU Name Persistence-M | Bus-Id Disp.A | Volatile Uncorr. ECC |
| Fan Temp Perf Pwr:Usage/Cap | Memory-Usage | GPU-Util Compute M. |
| | | MIG M. |
|=========================================+======================+======================|
| 0 NVIDIA GeForce RTX 3060 On | 00000000:03:00.0 Off | N/A |
| 0% 48C P5 29W / 170W | 18MiB / 12288MiB | 0% Default |
| | | N/A |
+-----------------------------------------+----------------------+----------------------+
torch==2.3.1+cu121
3.安装相关依赖
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/cuda/lib64
export PATH=$PATH:/usr/local/cuda/bin
export CUDA_HOME=$CUDA_HOME:/usr/local/cuda
pip install pycuda
4.锁频
MAX_Graphics_CLOCK=`nvidia-smi -q -d SUPPORTED_CLOCKS | grep 'Graphics' | sed 's/[^0-9]//g' | sort -n | uniq | tail -n 1`
MAX_Memory_CLOCK=`nvidia-smi -q -d SUPPORTED_CLOCKS | grep 'Memory' | sed 's/[^0-9]//g' | sort -n | uniq | tail -n 1`
nvidia-smi -pm 1
nvidia-smi -lgc $MAX_Graphics_CLOCK,$MAX_Graphics_CLOCK
nvidia-smi -i 0 -ac $MAX_Memory_CLOCK,$MAX_Graphics_CLOCK
nvidia-smi -q -d CLOCK
5.获取理论算力
tee Theoretical_FLOPS.py <<-'EOF'
import pycuda.driver as cuda
import pycuda.autoinit
def get_gpu_compute_capability_and_clock_rate():
device = cuda.Device(0)
compute_capability = device.compute_capability()
clock_rate = device.get_attribute(cuda.device_attribute.CLOCK_RATE) # in kHz
sm_count = device.get_attribute(cuda.device_attribute.MULTIPROCESSOR_COUNT)
cores_per_sm = get_cuda_cores_per_sm(compute_capability)
return compute_capability, clock_rate, sm_count, cores_per_sm
def get_cuda_cores_per_sm(compute_capability):
major, minor = compute_capability
if major == 2:
return 32
elif major == 3:
return 192
elif major == 5:
return 128
elif major == 6 and minor in [0, 1]:
return 64
elif major == 6 and minor == 2:
return 128
elif major == 7 and minor in [0, 5]:
return 64
elif major == 7 and minor == 2:
return 64
elif major == 8 and minor in [0, 6]:
return 128
else:
raise ValueError("Unknown compute capability")
def calculate_theoretical_flops(clock_rate, sm_count, cores_per_sm):
clock_rate_hz = clock_rate * 1e3 # Convert kHz to Hz
flops = clock_rate_hz * sm_count * cores_per_sm * 2 # 2 FLOPs per clock per core (FMA)
return flops
compute_capability, clock_rate, sm_count, cores_per_sm = get_gpu_compute_capability_and_clock_rate()
theoretical_flops = calculate_theoretical_flops(clock_rate, sm_count, cores_per_sm)
print(f"GPU compute capability: {compute_capability}")
print(f"Clock rate (kHz): {clock_rate}")
print(f"Number of SMs: {sm_count}")
print(f"Cores per SM: {cores_per_sm}")
print(f"Theoretical FLOPS for float32: {theoretical_flops / 1e12} TFLOPS")
EOF
python Theoretical_FLOPS.py
输出
GPU compute capability: (8, 6)
Clock rate (kHz): 1852000
Number of SMs: 28
Cores per SM: 128
Theoretical FLOPS for float32: 13.275136 TFLOPS
6.创建测试脚本
tee roofline_model.py <<-'EOF'
import sys
import torch
import torch.nn as nn
import triton
import triton.language as tl
import math
import torch
import torch.nn as nn
from fvcore.nn import FlopCountAnalysis, ActivationCountAnalysis
import matplotlib.pyplot as plt
import numpy as np
from matplotlib.font_manager import FontProperties
import os
import argparse
# 定义一个测试模型
class SimpleModel(nn.Module):
def __init__(self,input_features,output_features):
super(SimpleModel, self).__init__()
self.fc1 = nn.Linear(input_features,output_features,bias=False)
def forward(self, x):
x = self.fc1(x)
return x
@triton.jit
def sgemm_kernel(
A, B, C,
M, N, K,
stride_am, stride_ak,
stride_bk, stride_bn,
stride_cm, stride_cn,
BLOCK_SIZE: tl.constexpr
):
""" Kernel for computing C = A @ B """
# Define the program ids
pid_m = tl.program_id(0)
pid_n = tl.program_id(1)
# Create base pointers for A and B and C
offs_am = pid_m * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
offs_bn = pid_n * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
offs_ak = tl.arange(0, BLOCK_SIZE)
a_ptrs = A + (stride_am * offs_am[:, None] + stride_ak * offs_ak[None, :])
b_ptrs = B + (stride_bk * offs_ak[:, None] + stride_bn * offs_bn[None, :])
# Initialize accumulator
acc = tl.zeros((BLOCK_SIZE, BLOCK_SIZE), dtype=tl.float32)
# Loop over K dimension
for k in range(0, K, BLOCK_SIZE):
a = tl.load(a_ptrs, mask=offs_am[:, None] < M)
b = tl.load(b_ptrs, mask=offs_bn[None, :] < N)
acc += tl.dot(a, b)
a_ptrs += BLOCK_SIZE * stride_ak
b_ptrs += BLOCK_SIZE * stride_bk
# Write back results
c_ptrs = C + stride_cm * offs_am[:, None] + stride_cn * offs_bn[None, :]
tl.store(c_ptrs, acc, mask=(offs_am[:, None] < M) & (offs_bn[None, :] < N))
class TritonLinear(nn.Module):
def __init__(self, in_features, out_features):
super(TritonLinear, self).__init__()
self.in_features = in_features
self.out_features = out_features
self.weight = nn.Parameter(torch.randn(out_features, in_features).float()).cuda()
def forward(self, x):
assert x.shape[1] == self.in_features
out = torch.empty((x.shape[0], self.out_features), device=x.device, dtype=x.dtype).cuda()
grid = lambda META: (math.ceil(x.shape[0] / META['BLOCK_SIZE']), math.ceil(self.out_features / META['BLOCK_SIZE']))
sgemm_kernel[grid](
x, self.weight, out,
x.shape[0], self.out_features, self.in_features,
x.stride(0), x.stride(1),
self.weight.stride(0), self.weight.stride(1),
out.stride(0), out.stride(1),
BLOCK_SIZE=64
)
return out
def main(args):
# 模型和输入数据
input_features = 8192
output_features = 8192
batch_size = 8192
model = SimpleModel(input_features,output_features)
input_data = torch.randn(batch_size, input_features)
test_count=10
# 计算 FLOPs 和内存访问量
flops = FlopCountAnalysis(model, input_data).total()*test_count
activations = ActivationCountAnalysis(model, input_data).total() + input_data.numel()
print("activations:",activations)
# 计算参数个数
params = sum(p.numel() for p in model.parameters())
# 内存访问量假定为 activations 和params 乘以 4 字节（假设 activations 和 params 是 float32 类型）
activation_memory_access = activations * 4
params_memory_access = params * 4
memory_access = activation_memory_access + params_memory_access
memory_access=memory_access*test_count
if args.triton_kernel:
model = TritonLinear(in_features=input_features, out_features=output_features)
else:
model=model.cuda()
input_data=input_data.float().cuda()
for i in range(5):
output = model(input_data)
torch.cuda.synchronize()
if args.warmup_only:
return
if False:
# 设置 CUDA 事件用于计算执行时间
start_event = torch.cuda.Event(enable_timing=True)
end_event = torch.cuda.Event(enable_timing=True)
start_event.record()
for _ in range(test_count):
output = model(input_data)
end_event.record()
torch.cuda.synchronize()
total_cuda_time = start_event.elapsed_time(end_event) / 1000 # 转换为秒
else:
# 使用 PyTorch Profiler 计算 FLOPs、内存访问和执行时间
with torch.profiler.profile(
activities=[torch.profiler.ProfilerActivity.CUDA]) as prof:
for _ in range(test_count):
output = model(input_data)
key_averages = prof.key_averages()
for ev in key_averages:
print(ev)
total_cuda_time = sum([event.self_cuda_time_total for event in key_averages if event.key.find("sgemm")>=0]) / 1e6 # 转换至秒
# FLOPs 转换至 GFLOPs
flops_measured_glops = flops / 1e9
# 内存带宽测量
memory_access_gb=memory_access/ 1e9
bandwidth_measured = memory_access_gb / total_cuda_time # 单位：GB/s
print("bandwidth_measured:",bandwidth_measured)
# GPU 的峰值性能和带宽
peak_performance = 13.275136 * 1e3 # 单位：GFLOPs
memory_bandwidth = 360.0 # 单位：GB/s
# 计算 Roofline 模型中的数据点
Io = np.logspace(-2,4,100) #GFLOPs/GB
performance = np.minimum(peak_performance, Io * memory_bandwidth) #不同计算密度下的最大FLOPs/S,上限为峰值算力peak_performance
# 绘制 Roofline 模型
plt.figure(figsize=(10, 6))
thresold=0.75
# 设置字体以支持中文
font_path = 'simsun.ttc' # 在这里替换为你的字体路径
font_prop = FontProperties(fname=font_path)
# Bandwidth Bound
x=Io[Io<(peak_performance / memory_bandwidth)]
plt.fill_between(x, np.minimum(peak_performance, x * memory_bandwidth)*thresold,
np.minimum(peak_performance, x * memory_bandwidth),
color='lightblue', alpha=0.6, label='Bandwidth Bound')
# Compute Bound
x2=Io[Io>=(peak_performance / memory_bandwidth)]
plt.fill_between(x2, np.minimum(peak_performance, x2 * memory_bandwidth)*thresold,
np.minimum(peak_performance, x2 * memory_bandwidth),
color='green', alpha=0.6, label='Compute Bound')
# 绘制低性能区域
plt.fill_between(Io, 0, np.minimum(peak_performance, Io * memory_bandwidth)*thresold,
color='gray', alpha=0.6, label='poor performance')
plt.axhline(y=peak_performance, color='b', linestyle='--',
label=f'峰值计算能力:{peak_performance/1e3:.2f}TFLOPs')
plt.axvline(x=peak_performance / memory_bandwidth, color='g', linestyle='--',
label=f'{peak_performance / memory_bandwidth:.2f}GFLOPs/GB')
plt.loglog(Io, performance, label='Roofline')
arithmetic_intensity_measured=flops_measured_glops/memory_access_gb #GFLOPs/GB(算法的静态属性)
point_y = arithmetic_intensity_measured*bandwidth_measured
plt.scatter(arithmetic_intensity_measured, point_y, c='r',
label=f'Measured Points {point_y/1e3:.2f} TFLOPs/sec {point_y*100/peak_performance:.2f}%')
plt.xlabel('操作强度 [GFLOPs/GB]', fontproperties=font_prop)
plt.ylabel('性能 [GFLOPs/sec]', fontproperties=font_prop)
plt.title('Roofline 模型', fontproperties=font_prop)
plt.legend(prop=font_prop)
# 保存图片而不显示
plt.savefig('roofline_model.png')
plt.close()
print(f"FLOPs: {flops} FLOPs")
print(f"内存访问量: {memory_access} 字节")
print(f"执行时间: {total_cuda_time:.4f} 秒")
print(f"理论值的:{point_y*100/peak_performance:.2f}%")
parser = argparse.ArgumentParser(description='Process some integers.')
parser.add_argument("--warmup_only", action="store_true", help="warmup_only")
parser.add_argument("--triton_kernel", action="store_true", help="triton_kernel")
args = parser.parse_args()
main(args)
EOF
7.运行测试程序生成Roofline图
python roofline_model.py
python roofline_model.py --triton_kernel
输出
FLOPs: 5497558138880 FLOPs
内存访问量: 8053063680 字节
执行时间: 1.3862 秒
理论值的:29.87%
FLOPs: 5497558138880 FLOPs
内存访问量: 8053063680 字节
执行时间: 1.0957 秒
理论值的:37.80%
8.NVIDIA Nsight Compute生成Roofline
/usr/local/cuda/bin/ncu -f --section SpeedOfLight_HierarchicalSingleRooflineChart \
--section ComputeWorkloadAnalysis --section MemoryWorkloadAnalysis \
--target-processes all --export roofline_report python roofline_model.py --warmup_only
/usr/local/cuda/bin/ncu -f --section SpeedOfLight_HierarchicalSingleRooflineChart \
--section ComputeWorkloadAnalysis --section MemoryWorkloadAnalysis \
--target-processes all --export roofline_triton_kernel_report python roofline_model.py --warmup_only --triton_kernel
9.效果图
A.nn.Linear
B.Triton实现
