Sim Computing: construção do compilador DSA AI baseado em TVM

Este artigo foi publicado pela primeira vez na conta pública HyperAI Super Neural WeChat ~

Olá a todos, sou Dan Xiaoqiang da Sim Computing. Hoje vou compartilhar com vocês como dar suporte ao NPU no TVM com três colegas.

O problema essencial resolvido pelo compilador DSA é que diferentes modelos precisam ser implantados em hardware, e vários métodos de otimização em nível de abstração são usados ​​para fazer com que o modelo preencha o chip o máximo possível, ou seja, para comprimir as bolhas. Com relação a como agendar, o triângulo de agendamento descrito por Halide é a essência desse problema.

Qual é o principal problema que o compilador DSA está tentando resolver? Primeiro, abstraímos uma arquitetura DSA. Conforme mostrado na figura, habana, Ascend e IPU são todas instanciações dessa arquitetura abstrata. Cada núcleo em um núcleo de computação geral possui unidades de computação vetorial, escalar e tensorial. Do ponto de vista das operações de instrução e da granularidade dos dados, muitos DSAs podem tender a usar instruções relativamente granulares, como instruções bidimensionais e tridimensionais de vetores e tensores, e muitos hardwares usam instruções granulares finas, como unidimensionais SIMD e VLIW. Algumas das dependências entre as instruções são controladas pelo software por meio da exposição explícita da interface e outras são controladas pelo próprio hardware. A memória é uma memória de vários níveis, principalmente memória de rascunho. O paralelismo tem várias granularidades e dimensões de paralelismo, como paralelismo de fluxo, paralelismo de cluster, paralelismo de vários núcleos e paralelismo de pipeline entre diferentes componentes de computação. Para dar suporte a esse tipo de arquitetura, do ponto de vista dos desenvolvedores de compiladores, diferentes requisitos são apresentados para compiladores de IA a partir dos aspectos da arquitetura mencionada acima, que expandiremos posteriormente.

Do ponto de vista do usuário, em primeiro lugar, deve haver um compilador estável e generalizado, para que o maior número possível de modelos ou operadores possam ser compilados com sucesso. Além disso, os usuários esperam que o compilador possa fornecer uma interface programável para implementar algoritmos e operadores Customização para garantir que o trabalho inovador em alguns algoritmos-chave possa ser realizado de forma independente. Por fim, equipes ou amigos como nós também prestarão atenção em: como usar o TVM para construir compiladores de IA, como gerenciar códigos TVM autodesenvolvidos e de código aberto, como construir um CI eficiente, etc. É isso que vamos compartilhar hoje. A seguir, meu colega falará sobre a parte de compilação e otimização.

Schim Computing Wang Chengke: Processo de compilação e otimização do DSA

Esta parte é compartilhada pelo engenheiro de computação CIM Wang Chengke no local.

Primeiro, deixe-me apresentar a visão geral do processo da prática de compilação do Sim.

Tendo em vista os recursos arquitetônicos mencionados, construímos um passo de otimização autodesenvolvido com base na estrutura de dados TVM e reutilizamos o TVM para formar uma nova implementação de modo: tensorturbo.

Vimos uma arquitetura DSA mais clássica, que geralmente fornece alguns núcleos de computação multinúcleos de matriz e camada vetorial personalizados de alta eficiência, possui um mecanismo de cache multicamada que combina com ele e também fornece execução multimódulo que pode ser executada em unidade paralela. Assim, precisamos abordar as seguintes questões:

• Cálculos de dados de segmentos, vincular núcleos com eficiência e vetorizar instruções personalizadas com eficiência;

• Gerencie com precisão o cache limitado no chip e execute a pré-busca de dados correspondente em diferentes níveis de cache;

• Otimize o pipeline de vários estágios para execução de vários módulos e esforce-se para obter uma melhor taxa de aceleração.

A parte vermelha aqui (acima) mostra a parte com alta reutilização de TVM em todo o processo. As otimizações gerais relacionadas à camada implementadas no relé podem ser reutilizadas diretamente. Além disso, o maior grau de reutilização é baseado no TensorIR e na implementação do operador parte do LLIR personalizado. Otimizações personalizadas relacionadas a recursos de hardware como os que acabamos de mencionar exigem mais trabalho de autopesquisa.

Em primeiro lugar, vejamos um trabalho autodesenvolvido sobre a camada. Preste atenção ao diagrama de fluxo de cálculo típico na extrema esquerda e você poderá ver que, de cima para baixo, a ocupação geral do cache e do cálculo está diminuindo, mostrando um estado de pirâmide invertida. Para a primeira metade, quando a escala do modelo é grande, precisamos nos concentrar em resolver o problema de residência do cache no chip; para a segunda metade, quando a escala do modelo é relativamente pequena, precisamos lidar com o problema de baixa utilização de unidades de computação. Se você simplesmente ajustar o tamanho do modelo, como ajustar o tamanho do lote, um tamanho de lote menor pode obter uma latência menor e a taxa de transferência correspondente será reduzida; da mesma forma, um tamanho de lote maior levará a uma latência maior, mas é possível para melhorar o rendimento geral.

那么我们就可以用图调度来解决这个问题。首先，允许一个比较大的 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: Além das regras de fusão padrão suportadas pelo TVM, o Sim criou novas regras de fusão, como a fusão única de diferentes personalizações de hardware na camada de cálculo.

A3: Em relação à fusão, na verdade existem dois níveis, o primeiro, buffer e o segundo, fusão de loop. O método de fusão TVM é realmente voltado para o último. Sim basicamente seguiu o padrão de fusão TVM que você mencionou, mas fez algumas restrições.