Triton 是一种用于并行编程的语言和编译器。它旨在提供一个基于 Python 的编程环境，以高效编写自定义 DNN 计算内核，并能够在现代 GPU 硬件上以最大吞吐量运行。

更多 Triton 中文文档可访问 →https://triton.hyper.ai/

在本教程中，你将使用 Triton 编写一个简单的向量相加 (vector addition) 程序。

你将了解：

• Triton 的基本编程模型

• 用于定义 Triton 内核的 triton.jit 装饰器 (decorator)

• 验证和基准测试自定义算子与原生参考实现的最佳实践

计算内核

创建一个辅助函数从而： (1) 生成 z 张量，(2) 用适当的 grid/block sizes 将上述内核加入队列：

使用上述函数计算两个 torch.tensor 对象的 element-wise sum，并测试其正确性：

Out:

现在准备就绪。

基准测试

在 size 持续增长的向量上对自定义算子进行基准测试，从而比较其与 PyTorch 的性能差异。为了方便操作，Triton 提供了一系列内置工具，允许开发者简洁地绘制自定义算子在不同问题规模 (problem sizes) 下的的性能图。

运行上述装饰函数 (decorated function)。输入查看性能数据，输入 show_plots=True 绘制结果， 以及/或者输入 save_path='/path/to/results/' 将其与原始 CSV 数据一起保存到磁盘：

out: