本文首发自 HyperAI超神经微信公众号~
咱们好我是来自希姆核算的淡孝强，今日我将和三位搭档一起来给咱们共享怎么在 TVM 上支撑 NPU。
DSA 编译器处理的实质问题便是不同的模型需求布置到硬件上，利用各种笼统层级的优化手段，使得模型尽量打满芯片，也便是要紧缩气泡。关于怎么去调度，Halide 描绘的调度三角形是这个问题的实质。
DSA 编译器要处理的首要问题是什么？首要咱们笼统一个 DSA 的架构，如图所示，habana、Ascend 以及 IPU 上都是这个笼统架构的实例化。一般核算核里每个核有向量、标量以及张量的核算单元。从指令的操作和数据粒度来看，不少 DSA 或许倾向于运用相对粗粒度的指令，例如二维三维的向量和张量的指令，也有不少硬件运用细粒度的指令，例如一维的 SIMD 和 VLIW。指令间的依赖，有一些是经过显式的接口露出让软件来操控，有的是硬件自己操控。内存是多级内存，大多是 scratchpad memory。并行有各种粒度和维度的并行，比方 stream 并行、cluster 并行、多核并行以及不同核算部件之间的流水并行。
从用户的视点看，首要要有一个稳定和泛化的编译器，尽量多的模型或许算子都能够编译成功，别的，用户希望编译器能够供给一个可编程的界面来进行算法和算子的自定义，以确保能够独立展开一些要害算法的立异作业。最终，关于相似咱们这样的团队或许友商也会重视：怎么用 TVM 构建 AI 编译器，比方怎样办理自研和开源的 TVM 代码，怎么建立一套高效 CI 等等。这便是咱们今日要共享的内容，下面由我的搭档来讲编译优化的部分。
希姆核算王成科：DSA 的编译优化流程
本部分为希姆核算工程师王成科现场共享。
首要介绍一下希姆编译实践的全体流程概况。
针对方才说到的架构特性，咱们依据 TVM 数据结构构建了自研的优化 pass 再加上对 TVM 的复用，组成了一个新模式完成：tensorturbo。
咱们看到一个比较经典的 DSA 架构，一般会供给一些高效、定制矩阵以及向量层的多核核算中心，拥有一个与之相配合的多层缓存机制，一起也会供给可并行运转的多模组履行单元。相应的咱们需求处理以下问题：
-
对数据核算进行切分、高效绑核以及定制指令的高效向量化；
-
精密办理有限的片上缓存，在不同的缓存等级上做相应的数据预取；
-
优化多模组履行的多级流水线，力求得到一个比较好的加快比。
这儿赤色部分（上图）显现的是整个流程里对 TVM 复用比较高的部分，在 relay 上完成的图层相关比较通用的优化能够直接复用，别的复用程度比较高的是依据 TensorIR 和 custom LLIR 的算子完成部分。像咱们方才说到的跟硬件特性相关的定制优化，则需求更多自研作业。
首要咱们来看在图层上自研的一项作业。
那么咱们就能够用图调度来处理这个问题。首要，允许一个比较大的 batch size 输入，确保全程对核算的利用率比较高，然后对整图做一个存储剖析，加上切分和调度战略，使得模型的前半部分成果能够更好地缓存在片上，一起完成核算中心利用率较高的成果。实践来看全体能够完成 latency 和 throughput 都表现较好成果（具体能够重视 OSDI 23 希姆文章：Effectively Scheduling Computational Graphs of Deep Neural Networks toward Their Domain-Specific Accelerators，6 月份可获取链接检查）。
下面介绍别的一个软流水的加快作业。
由此，也会引入一个新的问题，比方当多个履行单元对缓存的一起读写并发数要高于其时缓存可支撑的并发数时，就会产生竞争，这个问题会导致访存功率成倍下降，也便是 Bank Conflict 问题。对此，咱们会在编译时静态地对流水线进行模拟，提取抵触目标，结合 cost model 对分配地址进行交换和平移，能够极大地下降该问题的影响。
咱们也留意到，上述作业中对数据结构的改造以及相关规划思想与现在 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 上正确履行。
从这个比如能够看到，起先 IR 是不能发射任何指令的，经过优化后，最终能够发射两条 Tensor 指令且一切的循环轴都能够发射指令。
希姆核算袁晟：DSA 的自定义算子
本部分为希姆核算工程师袁晟现场共享。
首要，咱们知道算子开发现在碰到了四个大问题：
-
需求支撑的神经网络算子许多，进行归类后根底算子有 100 多个；
-
由于硬件架构不停迭代，相应指令以及算子参与的逻辑都需求进行变更；
-
功能考虑。算子交融(local memory, share memory)以及我前边说到的图算信息传递（切分等）；
-
算子需求开放给用户，用户只能进入软件进行自定义算子。
我首要分成了以下三个方面介绍。首要是图算子，图算子是依据 relay api，把它裁剪成根底的语言算子。
以下图为例：
接下来是一个比较杂乱的算子，依据 NPU 的状况下，咱们会发现 topk、nms 等带操控流的算法，带许多标量核算，现在用 compute/schedule 很难描绘，为处理这个问题，希姆供给一个相似 library 库。相当于在 library 库先编译杂乱的逻辑，然后经过跟 IR Builder 结合的方法，把整个算子的逻辑输出。
再接下来对算子开发比较有挑战性的便是交融算子。现在面对两个爆破性问题，榜首不知道怎么将自己的算子和其他算子组合，第二个能够看到 NPU 里有许多 memory 层级，出现爆破式 memory 层级的交融。希姆 LLB 会有 shared memory 和 local memory 等交融的组合，依据这种状况下，咱们也供给一个主动生成结构，先依据图层给的调度信息，刺进数据搬移操作，再依据 schedule 里 master op 和 salve op 提炼 schedule info，最终依据其时指令的约束等问题做一个后处理。
总结
本部分为希姆核算工程师淡孝强现场共享。
这是希姆依据 TVM 搭的 CI，这上面跑了 200 多个模型以及十分多的单元测验。MR 不抢 CI 资源的状况下，提交一个代码需求 40 多分钟。核算量很大，大约挂了 20 多张自研核算卡以及一些 CPU 机器。
-
大部分运用自研的 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 的套路，可是做了一些约束。
