建模网站素材,南宁哪家公司建设网站比较好,怎样开网店卖别人的东西,wordpress 评论页面Triton 是一种用于并行编程的语言和编译器。它旨在提供一个基于 Python 的编程环境#xff0c;以高效编写自定义 DNN 计算内核#xff0c;并能够在现代 GPU 硬件上以最大吞吐量运行。
更多 Triton 中文文档可访问 →https://triton.hyper.ai/
在本教程中#xff0c;你将使…Triton 是一种用于并行编程的语言和编译器。它旨在提供一个基于 Python 的编程环境以高效编写自定义 DNN 计算内核并能够在现代 GPU 硬件上以最大吞吐量运行。
更多 Triton 中文文档可访问 →https://triton.hyper.ai/
在本教程中你将使用 Triton 编写一个简单的向量相加 (vector addition) 程序。
你将了解
Triton 的基本编程模型用于定义 Triton 内核的 triton.jit 装饰器 (decorator)验证和基准测试自定义算子与原生参考实现的最佳实践
计算内核
import torch
import triton
import triton.language as tltriton.jit
def add_kernel(x_ptr, # *Pointer* to first input vector. 指向第一个输入向量的指针。y_ptr, # *Pointer* to second input vector. 指向第二个输入向量的指针。output_ptr, # *Pointer* to output vector. 指向输出向量的指针。n_elements, # Size of the vector. 向量的大小。BLOCK_SIZE: tl.constexpr, # Number of elements each program should process. 每个程序应处理的元素数量。# NOTE: constexpr so it can be used as a shape value. 注意constexpr 因此它可以用作形状值。):# There are multiple programs processing different data. We identify which program# 有多个“程序”处理不同的数据。需要确定是哪一个程序pid tl.program_id(axis0) # We use a 1D launch grid so axis is 0. 使用 1D 启动网格因此轴为 0。# This program will process inputs that are offset from the initial data.# 该程序将处理相对初始数据偏移的输入。# For instance, if you had a vector of length 256 and block_size of 64, the programs would each access the elements [0:64, 64:128, 128:192, 192:256].# 例如如果有一个长度为 256, 块大小为 64 的向量程序将各自访问 [0:64, 64:128, 128:192, 192:256] 的元素。# Note that offsets is a list of pointers:# 注意 offsets 是指针列表block_start pid * BLOCK_SIZEoffsets block_start tl.arange(0, BLOCK_SIZE)# Create a mask to guard memory operations against out-of-bounds accesses.# 创建掩码以防止内存操作超出边界访问。mask offsets n_elements# Load x and y from DRAM, masking out any extra elements in case the input is not a multiple of the block size.# 从 DRAM 加载 x 和 y如果输入不是块大小的整数倍则屏蔽掉任何多余的元素。x tl.load(x_ptr offsets, maskmask)y tl.load(y_ptr offsets, maskmask)output x y# Write x y back to DRAM.# 将 x y 写回 DRAM。tl.store(output_ptr offsets, output, maskmask)创建一个辅助函数从而 (1) 生成 z 张量(2) 用适当的 grid/block sizes 将上述内核加入队列
def add(x: torch.Tensor, y: torch.Tensor):# We need to preallocate the output.# 需要预分配输出。output torch.empty_like(x)assert x.is_cuda and y.is_cuda and output.is_cudan_elements output.numel()# The SPMD launch grid denotes the number of kernel instances that run in parallel.# SPMD 启动网格表示并行运行的内核实例的数量。# It is analogous to CUDA launch grids. It can be either Tuple[int], or Callable(metaparameters) - Tuple[int].# 它类似于 CUDA 启动网格。它可以是 Tuple[int]也可以是 Callable(metaparameters) - Tuple[int]。# In this case, we use a 1D grid where the size is the number of blocks:# 在这种情况下使用 1D 网格其中大小是块的数量grid lambda meta: (triton.cdiv(n_elements, meta[BLOCK_SIZE]), )# NOTE:# 注意# - Each torch.tensor object is implicitly converted into a pointer to its first element.# - 每个 torch.tensor 对象都会隐式转换为其第一个元素的指针。# - triton.jited functions can be indexed with a launch grid to obtain a callable GPU kernel.# - triton.jit 函数可以通过启动网格索引来获得可调用的 GPU 内核。# - Dont forget to pass meta-parameters as keywords arguments.# - 不要忘记以关键字参数传递元参数。add_kernel[grid](x, y, output, n_elements, BLOCK_SIZE1024)# We return a handle to z but, since torch.cuda.synchronize() hasnt been called, the kernel is still running asynchronously at this point.# 返回 z 的句柄但由于 torch.cuda.synchronize() 尚未被调用此时内核仍在异步运行。return output使用上述函数计算两个 torch.tensor 对象的 element-wise sum并测试其正确性
torch.manual_seed(0)
size 98432
x torch.rand(size, devicecuda)
y torch.rand(size, devicecuda)
output_torch x y
output_triton add(x, y)
print(output_torch)
print(output_triton)
print(fThe maximum difference between torch and triton is f{torch.max(torch.abs(output_torch - output_triton))})Out:
tensor([1.3713, 1.3076, 0.4940, ..., 0.6724, 1.2141, 0.9733], devicecuda:0)tensor([1.3713, 1.3076, 0.4940, ..., 0.6724, 1.2141, 0.9733], devicecuda:0)The maximum difference between torch and triton is 0.0现在准备就绪。
基准测试
在 size 持续增长的向量上对自定义算子进行基准测试从而比较其与 PyTorch 的性能差异。为了方便操作Triton 提供了一系列内置工具允许开发者简洁地绘制自定义算子在不同问题规模 (problem sizes) 下的的性能图。
triton.testing.perf_report(triton.testing.Benchmark(x_names[size], # Argument names to use as an x-axis for the plot. 用作绘图 x 轴的参数名称。x_vals[2**i for i in range(12, 28, 1)], # Different possible values for x_name. x_name 的不同可能值。x_logTrue, # x axis is logarithmic. x 轴为对数。line_argprovider, # Argument name whose value corresponds to a different line in the plot. 参数名称其值对应于绘图中的不同线条。line_vals[triton, torch], # Possible values for line_arg. line_arg 的可能值。line_names[Triton, Torch], # Label name for the lines. 线条的标签名称。styles[(blue, -), (green, -)], # Line styles. 线条样式。ylabelGB/s, # Label name for the y-axis. y 轴标签名称。plot_namevector-add-performance, # Name for the plot. Used also as a file name for saving the plot. 绘图名称。也用作保存绘图的文件名。args{}, # Values for function arguments not in x_names and y_name. 不在 x_names 和 y_name 中的函数参数值。))
def benchmark(size, provider):x torch.rand(size, devicecuda, dtypetorch.float32)y torch.rand(size, devicecuda, dtypetorch.float32)quantiles [0.5, 0.2, 0.8]if provider torch:ms, min_ms, max_ms triton.testing.do_bench(lambda: x y, quantilesquantiles)if provider triton:ms, min_ms, max_ms triton.testing.do_bench(lambda: add(x, y), quantilesquantiles)gbps lambda ms: 3 * x.numel() * x.element_size() / ms * 1e-6return gbps(ms), gbps(max_ms), gbps(min_ms)运行上述装饰函数 (decorated function)。输入查看性能数据输入 show_plotsTrue 绘制结果 以及/或者输入 save_path/path/to/results/ 将其与原始 CSV 数据一起保存到磁盘
benchmark.run(print_dataTrue, show_plotsTrue)out: Download Jupyter notebook: 01-vector-add.ipynb Download Python source code: 01-vector-add.py Download zipped: 01-vector-add.zip
