Sim Computing: construcción del compilador DSA AI basado en TVM

Este artículo se publicó por primera vez desde la cuenta pública HyperAI Super Neural WeChat ~

Hola a todos, soy Dan Xiaoqiang de Sim Computing. Hoy compartiré con ustedes cómo admitir NPU en TVM con tres colegas.

El problema esencial que resuelve el compilador DSA es que se deben implementar diferentes modelos en hardware, y se utilizan varios métodos de optimización de nivel de abstracción para hacer que el modelo llene el chip tanto como sea posible, es decir, para comprimir las burbujas. En cuanto a cómo programar, el triángulo de programación descrito por Halide es la esencia de este problema.

¿Cuál es el problema principal que el compilador DSA está tratando de resolver? Primero, abstraemos una arquitectura DSA.Como se muestra en la figura, habana, Ascend e IPU son instancias de esta arquitectura abstracta. Cada núcleo en un núcleo de computación general tiene unidades de computación vectoriales, escalares y tensoriales. Desde la perspectiva de las operaciones de instrucciones y la granularidad de los datos, muchos DSA pueden tender a utilizar instrucciones de granularidad relativamente gruesa, como instrucciones vectoriales y tensoriales bidimensionales y tridimensionales, y muchos hardware utilizan instrucciones de granularidad fina, como instrucciones unidimensionales. SIMD y VLIW. Algunas de las dependencias entre instrucciones están controladas por el software a través de la exposición explícita de la interfaz y otras están controladas por el propio hardware. La memoria es una memoria multinivel, en su mayoría memoria de borrador. El paralelismo tiene varias granularidades y dimensiones de paralelismo, como el paralelismo de flujo, el paralelismo de clúster, el paralelismo de múltiples núcleos y el paralelismo de canalización entre diferentes componentes informáticos. Para admitir este tipo de arquitectura, desde la perspectiva de los desarrolladores de compiladores, se presentan diferentes requisitos para los compiladores de IA de los aspectos de la arquitectura mencionada anteriormente, que ampliaremos más adelante.

Desde el punto de vista del usuario, en primer lugar, debe haber un compilador estable y generalizado, de modo que se puedan compilar con éxito tantos modelos u operadores como sea posible. Además, los usuarios esperan que el compilador pueda proporcionar una interfaz programable para implementar algoritmos. y operadores Personalización para garantizar que el trabajo innovador en algunos algoritmos clave se pueda llevar a cabo de forma independiente. Finalmente, los equipos o amigos como nosotros también prestarán atención a: cómo usar TVM para construir compiladores de IA, cómo administrar códigos TVM de desarrollo propio y de código abierto, cómo construir un CI eficiente, etc. Esto es lo que vamos a compartir hoy, a continuación mi colega hablará sobre la parte de compilación y optimización.

Schim Computing Wang Chengke: proceso de compilación y optimización de DSA

Esta parte es compartida por el ingeniero informático de CIM, Wang Chengke, en el acto.

Primero, permítanme presentarles la descripción general del proceso de la práctica de compilación de Sim.

En vista de las características arquitectónicas que acabamos de mencionar, construimos un pase de optimización de desarrollo propio basado en la estructura de datos de TVM y reutilizamos TVM para formar una implementación de un nuevo modo: tensorturbo.

Hemos visto una arquitectura DSA más clásica, que generalmente proporciona algunos núcleos de cómputo multinúcleo de capa vectorial y matriz personalizada de alta eficiencia, tiene un mecanismo de caché de capa múltiple que lo combina y también proporciona ejecución de módulo múltiple que puede ejecutarse en unidad paralela. En consecuencia, debemos abordar los siguientes problemas:

• Segmente los cálculos de datos, enlace núcleos de manera eficiente y vectorice instrucciones personalizadas de manera eficiente;

• Administre finamente el caché limitado en el chip y realice la búsqueda previa de datos correspondiente en diferentes niveles de caché;

• Optimice la canalización de varias etapas para la ejecución de varios módulos y esfuércese por obtener una mejor relación de aceleración.

La parte roja aquí (arriba) muestra la parte con alta reutilización de TVM en todo el proceso. Las optimizaciones generales relacionadas con la capa implementadas en el relé se pueden reutilizar directamente. Además, el mayor grado de reutilización se basa en TensorIR y la implementación del operador. parte de LLIR personalizado. Las optimizaciones personalizadas relacionadas con las funciones de hardware como las que acabamos de mencionar requieren más trabajo de autoinvestigación.

En primer lugar, veamos un trabajo de desarrollo propio sobre la capa. Preste atención al diagrama de flujo de cálculo típico en el extremo izquierdo, y puede ver que, de arriba a abajo, la ocupación general de caché y cálculo está disminuyendo, mostrando un estado de pirámide invertida. Para la primera mitad, cuando la escala del modelo es grande, debemos centrarnos en resolver el problema de la residencia de caché en el chip; para la segunda mitad, cuando la escala del modelo es relativamente pequeña, debemos abordar el problema de la baja utilización de unidades de cómputo. Si simplemente ajusta el tamaño del modelo, como ajustar el tamaño del lote, un tamaño de lote más pequeño puede tener una latencia más baja y el rendimiento correspondiente se reducirá; de manera similar, un tamaño de lote más grande conducirá a una latencia más alta, pero es posible para mejorar el rendimiento general.

那么我们就可以用图调度来解决这个问题。首先，允许一个比较大的 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 对其进行排序。

我们知道优化套路本质是如何把数据搬移隐藏在计算中，无非是对操作照此标准进行模拟，最后计算代价。

P3: Además de las reglas de fusión predeterminadas admitidas por TVM, Sim ha creado nuevas reglas de fusión, como la fusión única de diferentes personalizaciones de hardware en la capa de cálculo.

R3: Con respecto a la fusión, en realidad hay dos niveles, el primero, el búfer, y el segundo, la fusión en bucle. El método de fusión TVM en realidad está dirigido a este último. Sim básicamente siguió el patrón de fusión de TVM que mencionaste, pero hizo algunas restricciones.