跳转至

教程

六篇手把手,每篇结束时你都能得到一个跑得起来的东西。

前置语言 —— 至少读过 类型函数作用域与放置

本章与其余各章的区别

语言算子 是按能力组织的 —— 一页讲一个特性,方便你查。本章按任务组织:每页从零做出一件东西,而且沿途每一步都是一个你能跑的程序。

这意味着各页会有意重复彼此的特性。想要 pl.at 的完整面,去读 作用域与放置;想写一个 matmul,读 分块 matmul —— 它会用 pl.at,但不会把它讲全。

两条线

写算子(00–03)—— 怎么表达计算。

页面 你最终得到 预计耗时
第一个算子 一个可跑的逐元素 kernel,与 torch 对拍 约 20 分钟
规约与 softmax 一个数值稳定的 softmax 约 30 分钟
分块 matmul 一个 K 轴分块的 matmul 约 40 分钟
混合 kernel cube 与 vector 在同一个作用域里并发工作 约 40 分钟

塑形调度(04–05)—— 怎么控制运行时拿它做什么。

页面 你最终得到 预计耗时
塑形任务图 一个依赖图完全可控的多任务程序 约 30 分钟
调度调优 一套能套用到自己 kernel 上的度量流程 约 40 分钟

阅读顺序

00-elementwise ──► 01-reduction-softmax ──► 02-matmul ──► 03-mixed-kernel
      │                                                          │
      └──────────────────────────► 04-task-graph ──► 05-scheduling-tuning

算子线是累积的 —— 每页都假定你已掌握前一页的 tile 词汇。调度线只需要 00:它讲的是 kernel 之间那张图的形状,与任何单个 kernel 算什么无关。

你的算子跑在哪个单元上

一个 core group 配一个 cube 单元(AIC)与若干 vector 单元(AIV)。由哪个执行不是你逐次调用去选的,而是由算子本身决定:

算子族 单元 讲解位置
matmulmatmul_accgemv Cube(AIC) 分块 matmul
逐元素、规约、广播、cast Vector(AIV) 0001
tpush_to_aivtpop_from_aicaiv_shardaic_gather 两者,按构造 混合 kernel

只由单一算子族构成的 kernel 会占住一个单元、让另一个闲着。这个观察正是 混合 kernel 存在的理由;在它之前的各页写的都是单单元 kernel。完整清单见 算子

跑这些示例

每页都指名 examples/ 下的一个文件。RunConfig.platform 默认就是 "a2a3sim",所以它们都不需要真机:

python examples/beginner/02_elementwise.py
python examples/advanced/03_mixed_kernel.py --mode staged

多数配套文件直接沿用这个默认值。03_mixed_kernel.py 是例外:它用 --mode 在几种 split 形式间切换,用 --platform 改目标平台。

参见