写一个矩阵乘，计算过程往往没几行，前面的数据排布却能绕上好一阵。九齿过去的接口就有这个问题：分完块以后，还要广播、调整层次、去掉多余的维度，才能让计算拿到正确的数据。
九齿一直想让大家按串行的思路写计算，把并行化交给编译器。最近，我们把编程接口做了一轮简化。以矩阵乘为例，现在两个输入和一个输出各写一行，就能把分块和维度关系说清楚，块之间的对齐和广播交给编译器。
先看矩阵乘的变化
左矩阵是 $M \times K$，右矩阵是 $K \times N$，输出是 $M \times N$。沿三个方向分块，块大小分别记作 BM、BN、BK。
旧接口要这样安排数据：
def arrangement(lhs, rhs, output):
output = output.tile((BM, BN))
lhs = lhs.tile((BM, BK)).tile((1, -1))
lhs = lhs.expand((-1, output.shape[1]))
lhs.dtype = lhs.dtype.squeeze(0)
rhs = rhs.tile((BK, BN)).tile((-1, 1))
rhs = rhs.expand((output.shape[0], -1))
rhs.dtype = rhs.dtype.squeeze(1)
return lhs, rhs, output
这一串操作，是为了给每个输出块配好沿 K 方向要读取的输入：左矩阵的一行块、右矩阵的一列块。写的时候要想清楚，过段时间再回来读，也得重新捋一遍每次变换后的形状。[1]
现在，同样的排布可以这样写：
BM, BN, BK = 16, 64, 32
def arrangement(lhs, rhs, output):
return (
lhs.axes("m", "k").tile((BM, BK)),
rhs.axes("k", "n").tile((BK, BN)),
output.axes("m", "n").tile((BM, BN)),
)
三个张量，各自一行。原来额外写的两次 tile、两次 expand、两次 squeeze，都省掉了。[2]
axes 给维度标上名字。左矩阵是 m、k，右矩阵是 k、n，输出是 m、n，看代码就能读出它们的关系。
左矩阵的一行块要供不同输出列使用，右矩阵的一列块要供不同输出行使用，这些对应关系也就有了依据。编译器结合后面的计算，在编译阶段补上原来手写的对齐、广播和层次整理。 [3]
计算照着原来的思路写
数据排布和计算仍然分开写。比如想试另一组块大小，可以调整前面的 arrangement，后面的乘加逻辑继续复用：
import ninetoothed.language as ntl
def application(lhs, rhs, output):
accumulator = ntl.zeros(output.shape, dtype=ntl.float32)
for k in range(lhs.shape[0]):
accumulator = accumulator + ntl.dot(lhs[k], rhs[k])
output = accumulator.to(output.dtype)
每次取出一对输入块，做一次小矩阵乘，加到 accumulator 里，最后写出当前输出块。这里的 lhs[k] 和 rhs[k] 都是数据块，k 数的是第几个块。[2]
拿两个 $1024 \times 1024$ 的矩阵举例，按上面的块大小，输出会分成 1024 个 $16 \times 64$ 的小块。每个输出块沿 K 方向做 32 次块乘加，就能算出结果。我们写的是其中一块怎样顺序计算，九齿负责把它按输出块网格并行展开。 [1][3]
加偏置和激活也更直接
实际写算子时，矩阵乘后面经常还跟着偏置、残差和激活。加上这些计算，排布也不用重新绕一遍。
比如，偏置只有 N 个元素，残差和输出一样是 $M \times N$，对应的排布声明就是：
bias.axes("n").tile((BN,))
residual.axes("m", "n").tile((BM, BN))
偏置没有 m 这一维，每一行都要用到它，编译器会补上这层广播；残差则按 m、n 和输出块对齐。把这两个输入加进来后，最后的输出赋值可以直接换成：
output = ntl.maximum(
accumulator + bias + residual, 0
).to(output.dtype)
偏置、残差和 ReLU 直接接在矩阵乘后面，整个过程仍写在同一个 application 里。各个输入与输出的对应关系，也写在前面的排布声明中。[4]
接下来还想怎么简化
接下来，我们希望写算子能更接近写普通函数。
- 能从计算中推导出的维度关系，就交给编译器，不用再逐个声明；
- 已经写好的矩阵乘，也能直接调用，组合进新的算子，不必把排布和块循环再写一遍。
编译器负责处理组合时的数据排布，让开发者更多地按算法的思路写代码。
🔗 相关链接
开发分支：
https://github.com/InfiniTensor/ninetoothed/tree/add-named-axis-layouts
[4] 偏置与残差示例： https://github.com/InfiniTensor/ninetoothed/blob/f1148467f6b94e27347d442a94244686ab30873e/tests/test_named_layout.py
关注与交流
