共计 6191 个字符，预计需要花费 16 分钟才能阅读完成。
本文首发自 HyperAI 超神经微信公众号~
大家好我是来自希姆计算的淡孝强，明天我将和三位共事一起来给大家分享如何在 TVM 上反对 NPU。
DSA 编译器解决的实质问题就是不同的模型须要部署到硬件上，利用各种形象层级的优化伎俩，使得模型尽量打满芯片，也就是要压缩气泡。对于怎么去调度，Halide 形容的调度三角形是这个问题的实质。
DSA 编译器要解决的次要问题是什么？首先咱们形象一个 DSA 的架构，如图所示，habana、Ascend 以及 IPU 上都是这个形象架构的实例化。个别计算核里每个核有向量、标量以及张量的计算单元。从指令的操作和数据粒度来看，不少 DSA 可能偏向于应用绝对粗粒度的指令，例如二维三维的向量和张量的指令，也有不少硬件应用细粒度的指令，例如一维的 SIMD 和 VLIW。指令间的依赖，有一些是通过显式的接口裸露让软件来管制，有的是硬件本人管制。内存是多级内存，大多是 scratchpad memory。并行有各种粒度和维度的并行，比方 stream 并行、cluster 并行、多核并行以及不同计算部件之间的流水并行。
要反对这样一类架构，从编译器开发者的角度看，是从上述体系结构几个方面对 AI 编译器提出不同的需要，这部分前面咱们会开展。
从用户的角度看，首先要有一个稳固和泛化的编译器，尽量多的模型或者算子都能够编译胜利，另外，用户心愿编译器能够提供一个可编程的界面来进行算法和算子的自定义，以确保能够独立发展一些要害算法的翻新工作。最初，对于相似咱们这样的团队或者友商也会关注：怎么用 TVM 构建 AI 编译器，比方怎么治理自研和开源的 TVM 代码，怎么搭建一套高效 CI 等等。这就是咱们明天要分享的内容，上面由我的共事来讲编译优化的局部。
希姆计算王成科：DSA 的编译优化流程
本局部为希姆计算工程师王成科现场分享。
首先介绍一下希姆编译实际的整体流程详情。
针对方才提到的架构个性，咱们基于 TVM 数据结构构建了自研的优化 pass 再加上对 TVM 的复用，组成了一个新模式实现：tensorturbo。
咱们看到一个比拟经典的 DSA 架构，个别会提供一些高效、定制矩阵以及向量层的多核计算外围，领有一个与之相配合的多层缓存机制，同时也会提供可并行运行的多模组执行单元。相应的咱们须要解决以下问题：
- 对数据计算进行切分、高效绑核以及定制指令的高效向量化；
- 精密治理无限的片上缓存，在不同的缓存等级上做相应的数据预取；
- 优化多模组执行的多级流水线，力求失去一个比拟好的减速比。
这里红色局部（上图）显示的是整个流程里对 TVM 复用比拟高的局部，在 relay 上实现的图层相干比拟通用的优化能够间接复用，另外复用水平比拟高的是基于 TensorIR 和 custom LLIR 的算子实现局部。像咱们方才提到的跟硬件个性相干的定制优化，则须要更多自研工作。
首先咱们来看在图层上自研的一项工作。
关注最右边这张比拟典型的计算流图，能够看到从上到下，整体对缓存的占用及对计算的占用都在一直缩小，出现倒金字塔的状态。对于前半部分，模型规模较大时，咱们须要着重解决片上缓存驻留的问题；而后半局部，在模型规模比拟小的时候，须要解决计算单元利用率较低的问题。如果简略调整模型规模，比方调整 batch size，较小的 batch size 能够失去较低的 latency，而相应的 throughput 会有所升高；同样较大的 batch size 会导致 latency 较高，然而有可能进步整体 throughput。
那么咱们就能够用图调度来解决这个问题。首先，容许一个比拟大的 batch size 输出，保障全程对计算的利用率比拟高，而后对整图做一个存储剖析，加上切分和调度策略，使得模型的前半部分后果能够更好地缓存在片上，同时实现计算外围利用率较高的后果。实际来看整体能够实现 latency 和 throughput 都体现较好后果（具体能够关注 OSDI 23 希姆文章：Effectively Scheduling Computational Graphs of Deep Neural Networks toward Their Domain-Specific Accelerators，6 月份可获取链接查看）。
上面介绍另外一个软流水的减速工作。
关注右上图，实现了一个比拟 native 的四级流水线，但显著不是一个高效的流水线。个别高效的流水线，应该是通过几次迭代后，四个执行单元都能够同步并行起来，那么这须要做一些工作，包含 L1 及 L0 上的切分、L1 上跨层的数据预取以及 L0 层级上的 double buffer 操作。通过这些工作咱们能够实现像右下图所展现的，减速比拟高的流水线。
由此，也会引入一个新的问题，比方当多个执行单元对缓存的同时读写并发数要高于以后缓存可反对的并发数时，就会产生竞争，这个问题会导致访存效率成倍降落，也就是 Bank Conflict 问题。对此，咱们会在编译时动态地对流水线进行模仿，提取抵触对象，联合 cost model 对调配地址进行替换和平移，能够极大地升高该问题的影响。
有了各种 pass 之后，能够以一个简略的 Top-Down 形式把它们组合起来，沿着左图中彩色流程，就失去了一个性能上可行的编译 pipeline。然而实际中发现很多问题，包含思远提到的 pass 与 pass 之间的相互影响、短少交互逻辑，图层与算子之间短少沟通逻辑等。能够看到左图中红色局部批示的流程，实际中发现每个门路或者它们的组合都会导致编译失败。如何让其鲁棒性更强？希姆在每个可能失败的 pass 中提供一个反馈门路，在图层和算子之间引入了交互逻辑，进行预剖析、prelower 操作，同时在重点局部引入一些迭代调优机制，最终失去一个泛化性较高且调优能力比拟强的整体 pipeline 实现。
咱们也留意到，上述工作中对数据结构的革新以及相干设计思维与目前 TVM Unity 设计有较多相似之处，咱们也期待 Relax 可能带来更多可能性。
这里展现的是希姆在编译流程中更加细节的 pass，从左到右就是逐层递加的过程，其中红色局部是对 TVM 复用比拟高的，越凑近硬件个性局部会有更多的定制 pass。
上面持续对其中的局部模块进行具体介绍。
希姆计算刘飞：DSA 的向量化和张量化
本局部为希姆计算工程师刘飞现场分享。
这个章节将开展介绍希姆向量化和张量化工作。从指令粒度思考，指令粒度越粗，越靠近 Tensor IR 的多层 loop 表白，所以向量化张量化难度越小，相同，指令粒度越细，难度也就越大，咱们的 NPU 指令，反对一维 / 二维 / 三维的 tensor 数据计算。希姆也思考过原生 TVM tensorize 过程，但思考到 Compute Tensorize 对简单表达能力无限，例如对 if condition 这种简单表达式做 Tensorized 就比拟艰难，而且做 Tensorized 向量化后，无奈 schedule。
另外过后 TensorIR Tensorize 在开发当中，不能满足开发需要，所以希姆提供了本人的一套指令向量化流程，咱们称之为指令发射。这套流程下咱们反对了大略 120 条 Tensor 指令，包含各种维度的指令等。
咱们的指令流程大略分为三个模块：
- 发射前的优化解决。对循环轴的变换，为指令发射提供更多的发射条件和可能；
- 指令发射模块。剖析循环的后果和信息，抉择一个最优的指令生成形式；
- 指令发射后的模块。对指定发射解决失败之后解决，保障在 CPU 上正确执行。
上面是指令发射前的优化和解决模块，都是由一组优化 pass 组成，其中 IfPromotion 是把妨碍循环轴发射的 if 语句尽量外提，PreProcess 是把没有对应指令的 operator 做拆分解决，LoopShift 是对循环轴边界为归一化，LoopCallapse 是对间断的循环轴作尽可能的合并，LoopPartition 是做 if 相干的循环轴拆分，还有 LoopFission 是对循环内多个 store 语句的决裂。
从这个例子能够看到，起初 IR 是不能发射任何指令的，通过优化后，最初能够发射两条 Tensor 指令且所有的循环轴都可能发射指令。
再就是指令发射模块。首先，指令发射模块会循环剖析循环中的构造，从中获取 Optype、dtype、bufferAcess 等信息，有这些信息之后，指令辨认会辨认进去循环轴可能会发射哪几种指令。因为一种 IR 构造可能对应多种 NPU 指令，所以咱们会把所有可能发射的指令都辨认进去，由 VectorEngine 搜索引擎去依据指令的 alignment、reshape 等一系列信息去搜寻每种指令发射的可能性，最初再由 CostModel 做计算，找到最优发射模式进行发射。
最初就是指令发射后处理模块。次要是对指令发射失败的 tir 做解决，保障其能在 CPU 上正确运行。还有一些非凡指令，希姆须要在算法前端打一些标记，指令发射模块通过这些标记加上本人的 IR 剖析，正确地发射相应的指令。
以上是希姆整个 DSA 张量化和向量化的流程，咱们也在一些方向上做了摸索，比方微内核的计划，也是最近探讨比拟热烈的方向。它的根本思维是把一个计算过程分成两层，一层用组合微内核的模式去拼接，另一种用搜寻的形式去寻找，最终把两层的后果做拼接，抉择一个最优后果。这样其劣势是充分利用硬件资源的同时，升高搜寻的复杂度，进步搜寻效率。
希姆也在微内核上做了相干摸索，但思考到微内核计划与当初的解决方案相比，并没有在性能等方面有较大晋升，所以目前希姆在微内核方向还属于摸索阶段。
希姆计算袁晟：DSA 的自定义算子
本局部为希姆计算工程师袁晟现场分享。
首先，咱们晓得算子开发目前碰到了四个大问题：
- 须要反对的神经网络算子很多，进行归类后根底算子有 100 多个；
- 因为硬件架构不停迭代，相应指令以及算子参加的逻辑都须要进行变更；
- 性能思考。算子交融 (local memory, share memory) 以及我前边提到的图算信息传递（切分等）；
- 算子须要凋谢给用户，用户只能进入软件进行自定义算子。
我次要分成了以下三个方面介绍。首先是图算子，图算子是基于 relay api，把它裁剪成根底的语言算子。
以下图为例：
第二是元算子，所谓的元算子是基于 TVM Topi 用 compute/schedule 形容算子算法逻辑和循环变换相干逻辑。咱们在开发算子时，会发现很多算子的 schedule 是能够复用的，基于这种状况下，希姆提供了一套相似 schedule 的模板。当初，咱们把算子分成很多类，基于这些类，新的算子就会大量复用 schedule 模板。
接下来是一个比较复杂的算子，基于 NPU 的状况下，大家会发现 topk、nms 等带控制流的算法，带很多标量计算，目前用 compute/schedule 很难形容，为解决这个问题，希姆提供一个相似 library 库。相当于在 library 库先编译简单的逻辑，而后通过跟 IR Builder 联合的形式，把整个算子的逻辑输入。
接下来是算子的切分。对于 NPU，绝对 GPU 和 CPU 状况下，TVM 每条指令都会操作间断内存块，同时会有 memory size 限度。同时，在这种状况下，搜寻空间不大。基于这些问题，希姆提供了解决办法，首先，会有一个候选集，把可行的解题放到候选集里，其次，对可行性进行解释，次要思考性能要求以及 NPU 指令限度，最初，会引入 cost function，其中会思考算子特色以及可能用到的计算单元特色。
再接下来对算子开发比拟有挑战性的就是交融算子。目前面临两个爆炸性问题，第一不晓得如何将本人的算子和其余算子组合，第二个能够看到 NPU 里有很多 memory 层级，呈现爆炸式 memory 层级的交融。希姆 LLB 会有 shared memory 和 local memory 等交融的组合，基于这种状况下，咱们也提供一个主动生成框架，先依据图层给的调度信息，插入数据搬移操作，再依据 schedule 里 master op 和 salve op 提炼 schedule info，最初依据以后指令的限度等问题做一个后处理。
最初次要展现希姆反对的算子。ONNX 算子大略是 124 个，目前反对的大略是 112 个，占比 90.3%，同时希姆有一套随机测试，能够测试大质数、交融组合以及一些 pattern 交融组合。
总结
本局部为希姆计算工程师淡孝强现场分享。
这是希姆基于 TVM 搭的 CI，这下面跑了 200 多个模型以及十分多的单元测试。MR 不抢 CI 资源的状况下，提交一个代码须要 40 多分钟。计算量很大，大略挂了 20 多张自研计算卡以及一些 CPU 机器。
总结，这是希姆的架构图，如下所示：
成果来看，性能失去很大晋升，另外主动生成与另一个手写模型的团队对标的话，基本上能够达到他们的 90% 以上。
这是希姆代码的状况，右边是 TVM 和自研代码如何治理，TVM 是作为 third_party 里的数据结构来应用，希姆有本人的 source 和 python 的货色，如果咱们须要对 TVM 进行更改，就在 patch 文件夹中对 TVM 进行改变。这里有三个准则：
- 大部分应用自研的 pass，也自研了 Custom module；
- patch 会限度少批改 TVM 源代码，能 upstream 就及时 upstream；
- 定期跟 TVM 社区做同步，更新最新代码到仓库中。
整个代码量也如上图所示。
总结：
- 咱们基于 TVM 端到端反对希姆一代二代芯片；
- 基于 relay 和 tir 实现所有的编译优化需要；
- 基于 tir 实现了 100+ 条向量张量指令的主动生成；
- 基于 TVM 实现了自定义算子计划；
- 模型一代反对 160+，二代曾经使能 20+；
- 模型性能靠近手写极致。
Q&A
Q1：我对交融算子比拟感兴趣，它如何跟 TVM 的 tir 联合？
A1：对于右图，同一个算级，第一，如果算子有两个 input 一个 output，那算子状态就有 27 种。第二，各种各样的算子连接时，scope 有可能是三个之一，所以咱们不会假如有固定 pattern。那么如何在 TVM 上实现？首先依据图层调度，决定前后 add 和两头 scope 在哪里，图层是一个非常复杂的过程，输入的后果是决定算子存在于哪个缓存以及可用缓存有多少。有了这个调度的后果，咱们在算子层进行主动交融算子生成，比方咱们依据 scope 信息进行主动插入数据搬移的操作，实现数据流的构建。
schedule info 里边和 TVM 原生的机制很相似，交融过程中须要思考每个 member scope 所用的大小，所以这里就是 TVM 原生的货色，咱们只是用了一个特地的框架，将其集成到这里，让它自动化。
do schedule 在此基础上，把开发者所须要的 schedule 做进去，可能也会有一些后处理。
Q2：不便走漏 CostModel 更多细节吗？cost function 是依据算子层面的 feature 还是依据硬件层面的个性联合设计的？
A1：大思路曾经在这了，首先生成一个候选集，生成过程跟 NPL 构造相干，而后会有剪枝的过程，思考指令限度以及后边的优化，多核、double buffer 等，最初有一个 cost function 对其进行排序。
咱们晓得优化套路实质是如何把数据搬移暗藏在计算中，无非是对操作照此规范进行模仿，最初计算代价。
Q3：除了 TVM 反对的默认交融规定，希姆有没有产生新的交融规定，比方在计算图层联合不同硬件定制的特有交融。
A3：对于交融，实际上有两个档次，第一，buffer，第二，loop 交融。TVM 交融形式实际上是针对后一种。希姆根本沿用你说的 TVM 交融 pattern 的套路，然而做了一些限度。
