走向通用张量处理器（二）：layout 溯源
我认为能否高效处理任意 layout 是从专用模型加速器到通用张量处理器的重要标志之一。社区中很多从软件代数角度介绍 tile / layout ，但鲜有内容从硬件设计捋一捋 layout 。先前的博文主要侧重 layout 抽象符号化建模[1]，本文将简单梳理下 layout 产生的第一性原理以及相应软硬件设计方法。
layout从何而来：不对等的位置
Layout，或又称 memory layout / buffer layout。从名称上可以见得和存储脱离不了关系，粗糙来说，layout 是指一个多维张量在存储上的位置分布，位置则包括地址/ bank / bit 顺序等等。
为什么一个多维张量的位置如此重要呢？本质在于对于某个处理而言位置分布并不是平权的。（1）在传统单线程架构里，虽然从程序上来看不同地址是平权的，但 2D DDR 的行刷新机制如果处理按顺序访存效率更优；而在（2）并行张量处理器，一些并行计算单元将控制流映射到空间上，从而使得对访存数据位置有了特殊要求。如下图，output stationary 的 systolic array （OS-SA）对矩阵 a 每次按一行读取，而对矩阵 b 每次按一列读取。
计算和存储在张量处理下步调失调
本文接着按“并行计算单元有特殊的访存顺序”这一基础假设继续分析。虽然张量可以指任意多维张量，但一般实践中只需要考虑建模成二维矩阵即可，更多维度一般以控制流实现，我们以 二维矩阵为基础数据对象 继续建模：
- 计算侧：计算单元和layout耦合： 访存顺序即读和写的顺序，比如上述的 4 x 4 OS-SA，完成计算 \(\text{Einops}(A[\text{Token}][\text{ChIn}],B[\text{ChOut}][\text{ChIn}]\rightarrow C[\text{Token}][\text{ChOut}])\) 。A 矩阵每次按 4 个 Token 维度读取，B 矩阵每次按 4 个 ChOut 维度读取，C 矩阵写出的时候，由于 buffer 带宽一般和其余变量统一也是每次读 4 个，16 个数据既可以按每次 4 个 Token 或者 4 个 ChOut 维度写入，这就是以上例子 systolic array 对读和写的顺序要求；
- 存储侧：变量存储和layout耦合： 高效的访问顺序由变量在存储中的分布决定，比如对于一个张量 \([Token][Channel]\)，一个简单的单 bank sram 每个 line 映射着 Channel 维度，访问不同的 Token 维度必须访问不同的地址，并行访问效率低下。
- 计算和存储：不存在单射关系： 计算和变量组成计算图。同一个变量，可能是拥有多个计算算子的 consumer，每个算子高效访问顺序不一致；producer 的写顺序并和 consumer 的高效读顺序不一致。
硬件的实现机制：解耦的两种思路
Layout 问题根源在于计算和存储耦合，设计者要想的则是怎么解耦合。解耦合是“通用”“转换”的同义词，要在哪里插入一个转换的“桥”。根据“桥”插入位置区分，有两种处理思路：
- 分离操作，在 Memory 内转换： layout 转换单元：添加一个 layout 转换单元，先在存储内转换变量 layout ，再执行计算操作；
- 融合操作，在 Dataflow 上转换： 在计算和存储之间插入一个 stream 的转换单元，交互是动态做布局转换，一般是由 multi-bank + crossbar 实现，相应的读和写的位置都可以加入。
注意，这里只是分析实现布局解耦硬件需要什么实现，还未曾涉及具体硬件控制流怎么生成，也没有讨论任意不同布局之间并非都可以高效转换，难免遇到 bank conflict 。这就涉及软件栈设计以及一些数学分析，数学分析前文博客/网络资料已经很多读者可自行查阅。
软件栈的改动：提升 layout 地位
计算和layout耦合，存储也和layout耦合，进而推导出计算和存储耦合，这又违背了计算存储之间的映射关系。在张量处理语境下，存储和计算需要当作一个问题设计。 但layout问题显然不是架构中的新问题，为什么之前没有太多呼声提高 layout 在编程体系中的地位？
存算分离本意将变量（存储）和算子解耦相互正交以提高通用性和灵活性，编程抽象以“功能”为第一公民，而layout是“效率”的公民，因此体系架构允许在隐藏layout抽象下维持功能正确。并行计算大概可以分为两种算子，（1）一种每条 lane 完全并行互不干扰的，以 element-wise 为例；（2）另一种 lane 之间互相有交互，比如 reduce operation / gemm 等。个人观点以往 SIMT 架构假设并行程序大部分为（1），大部分时候无需关心layout问题即可实现效率较优，而少部分时候需要处理（2）情况再按需投入一些编程人力。但随着深度学习算子比例逐渐往 （2）迁移，越来越多layout复杂度被泄露到编程层，调整编程体系又提上议程。
耦合！耦合！耦合！
到这里为止，本文也只是尽可能少引入具体设计和公式定性分析，且敲定 layout 设计需要的努力远远比本文所提得多，不同的精度位宽怎么处理、怎么侧重 workload 以及设计软件栈、电路特性又会带来什么影响…… 快速的发展、极致的性能需求、层层耦合的系统，把控住亿万黄金的算力猛兽需要多少思索和心血，考题范围似乎早已远远超出架构设计。最后以我在《辩经系列之一：NPU or GPGPU》[2] 文章下的评论结尾。
浙公网安备 33010602011771号