├── doc ├── pic │ ├── end-to-end-time.png │ └── end-to-end-time-b1.png └── 1-unified-ir-defs.md ├── .gitignore └── README.md /doc/pic/end-to-end-time.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/deathwings602/Unified-IR/HEAD/doc/pic/end-to-end-time.png -------------------------------------------------------------------------------- /doc/pic/end-to-end-time-b1.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/deathwings602/Unified-IR/HEAD/doc/pic/end-to-end-time-b1.png -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Prerequisites 2 | *.d 3 | 4 | # Compiled Object files 5 | *.slo 6 | *.lo 7 | *.o 8 | *.obj 9 | 10 | # Precompiled Headers 11 | *.gch 12 | *.pch 13 | 14 | # Compiled Dynamic libraries 15 | *.so 16 | *.dylib 17 | *.dll 18 | 19 | # Fortran module files 20 | *.mod 21 | *.smod 22 | 23 | # Compiled Static libraries 24 | *.lai 25 | *.la 26 | *.a 27 | *.lib 28 | 29 | # Executables 30 | *.exe 31 | *.out 32 | *.app 33 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Unified IR 2 | 3 | ## About 4 | 5 | **Uinfied IR** 是一套面向人工智能应用的统一中间表示（Unified Intermediate Representation），基于此可以方便的在不同的硬件平台上对人工智能应用做统一的编译优化和代码生成。随着深度学习在我们的日常生活中得到越来越广泛的应用，除英伟达外的各大硬件厂商也推出了自己的硬件加速平台用于大模型的训练与推理，比如：寒武纪、华为和沐曦。由于不同硬件平台的架构和编程接口不同，我们需要分别为它们编写庞大的深度学习算子库，这种依赖专家知识的方法通常不能得到所有算子的最佳实现。而且，各类对人工智能应用的编译工作往往和平台紧密耦合，难以在不同平台间进行迁移。因此，我们迫切需要一套能够对人工智能应用进行跨平台编译优化的统一方法。 6 | 7 | **Unified IR** 正是为了解决这一困境而提出的，它是一个指令级的描述张量程序的中间表示，其设计的基本单元是对硬件指令的细粒度抽象，尽可能利用了各类硬件的架构和编程接口的共性，从而可以用一个统一的表示来描述各类硬件上的张量程序实现。通过在这样一个统一表示上用诸如 **IntelliGen** 这样的深度学习编译器对算子实现进行编译优化，然后将优化后的中间表示经过不同的后端生成对应硬件平台上的代码，我们可以很方便的为不同的硬件平台生成性能优异的算子实现。同时，这也意味着在 **Unified IR** 上所做的任何优化都可以无缝的迁移到所有支持的硬件平台上。 8 | 9 | **Unified IR** 的优势有： 10 | + 用一种统一的方式来描述算子在不同硬件平台上的实现，从而可以实现对张量程序的跨平台编译优化； 11 | + 面向人工智能应用设计，相比于 OpenCL 这类具有强大表达能力的语言，更加充分的利用了深度学习算子的规整性，因此能挖掘并利用更多的优化机会，在表达能力和易优化性上取得了很好的平衡； 12 | + 在语义上贴近硬件的底层指令，可以十分方便的将其编译为不同硬件平台上的代码，具有很强的拓展性。 13 | 14 | ## Supported Platforms 15 | 16 | 目前，**Unified IR** 已经支持了NVIDIA、寒武纪和华为昇腾平台，我们通过深度学习编译器 **IntelliGen** 对张量程序的中间表示进行统一的优化，然后为不同的硬件平台实现了从 **Unified IR** 到对应平台代码生成后端，在包括Llama、GPT2的主流大模型中都取得了明显的加速。这样的结果表明，**Unified IR** 不仅可以为上层的人工智能应用提供恰当的抽象，从而能挖掘出更多的优化机会，还可以充分描述底层不同硬件平台的特性，能获得跨平台的优异性能。 17 | 18 | ### Results On Ascend 910B 19 | 20 | ![Ascend 910B上的加速效果，batch size = 1](https://github.com/deathwings602/Unified-IR/blob/main/doc/pic/end-to-end-time-b1.png) 21 | 22 | (batch size = 1) 23 | 24 | ![Ascend 910B上的加速效果，batch size = 4](https://github.com/deathwings602/Unified-IR/blob/main/doc/pic/end-to-end-time.png) 25 | 26 | (batch size = 4) 27 | ## Getting Started 28 | 29 | 在这里查看有关 **Unified IR** 的定义：[Introduction To Unified IR](https://github.com/deathwings602/Unified-IR/blob/main/doc/1-unified-ir-defs.md) 30 | 31 | 相关论文：[PowerFusion: A Tensor Compiler with Explicit Data Movement Description and Instruction-level Graph IR 32 | ](https://arxiv.org/abs/2307.04995) 33 | -------------------------------------------------------------------------------- /doc/1-unified-ir-defs.md: -------------------------------------------------------------------------------- 1 | # Introduction To Unified IR 2 | 3 | ## Overview 4 | 5 | **Unified IR** 是一个描述张量程序的统一中间表示，旨在为人工智能应用提供跨平台的统一编译优化。每个张量程序都对应于 **Unified IR** 中的一个 **Kernel Graph**，描述了不同的设备算子在张量程序中的数据依赖关系；**Kernel Graph** 中部分节点为自定义算子，表示为一个 **iGraph**，这是一个对描述算子实现的指令级中间表示。 6 | 7 | ## Kernel Graph 8 | 9 | 在 **Unified IR** 中，**Kernel Graph** 用于描述一个张量程序，其中的每个节点表示一个在目标硬件平台上执行的算子，每条边表示在不同算子之间共享的张量。所有的张量都存储在设备的DRAM中，这是因为不同的算子无法通过寄存器文件或者SRAM共享数据。每个节点所表示的算子可以是对应硬件预定义的算子，如矩阵乘法和卷积，也可以是由 **iGraph** 描述的自定义算子，**iGraph** 允许我们在 **Kernel Graph** 层级上做算子融合等算子间（inter-operator）优化。 10 | 11 | ## iGraph 12 | 13 | **iGraph** 从硬件指令的角度描述了算子的实现，它的形式是一个数据流图，描述了数据在不同内存层级间的移动和硬件指令对数据的处理。图中的每个节点为 **iOp**，表示在硬件的最小并行单元上的一条硬件指令或是一个指令序列，每条边是一个 **iSlice**，表示数据切片，是上述的 **iOp** 所作用的基本单元。**iGraph** 还包含硬件并行架构信息和对算子实现并行方式的描述。 14 | 15 | ### Hardware Model 16 | 17 | 考虑到AI加速硬件内在的相似性，**Unified IR** 提出了一套高层次的对硬件并行架构的建模，从而可以用一套统一的方式描述不同的硬件的特性。 18 | 19 | 我们定义 **最小并行单元** 为目标硬件中可以独立执行的最小单元，这在不同的硬件中有不同的含义，在NVIDIA系列的GPU上对应于SMSP（可以独立的调度warp），而在Ascend 910系列芯片中则对应于一个AI Core。在部分硬件中，多个最小并行单元可以组成一组协同计算，之间可以通过SRAM来进行通信，我们称这样的一组最小并行单元为 **并行组**。例如在CUDA编程中，一个线程块包含多个Warp，它们都在同一个SM上执行，彼此之间可以进行同步和通过共享内存同步数据。而一个AI加速硬件包含多个上述的 **并行组**。 20 | 21 | 对内存层级也可以做类似的划分，按照访问速度从高到低可以分为 **REG**、**SRAM** 和 **DRAM**。**REG** 层级的内存由每个最小并行单元独占，访问速度最快；**SRAM** 层级的内存则可以在不同的并行组之间共享，访问速度次之；**DRAM** 在不同的算子之间共享，访问速度最慢。 22 | 23 | 因此，根据上述的 “硬件 -> 并行组 -> 最小并行单元” 的抽象，我们可以用二元组 (numGroup, numUnit) 来描述硬件的并行架构，分别表示一个硬件中有多少个并行组，一个并行组中有多少个最小并行单元。由于CUDA的编程接口对底层的硬件架构做了进一步的抽象，一个核函数的grid size和block size均可变，我们这套模型在描述CUDA算子时需要做一些微妙的调整：我们令numGroup对应于算子实现中的grid size，而numUnit对应于线程块中的Warp数量。此外，Ascend 910 系列芯片的AI Core之间不能进行通信，因此就不存在上述的 **并行组** 层级，这对应于numUnit为1的退化情况。 24 | 25 | 类似的，我们也可以将具体硬件的内存层级按照其和并行架构的关系映射到上述的 “REG -> SRAM -> DRAM” 层级上。对于NVIDIA的GPU来说，寄存器文件对应 **REG**，共享内存对应 **SRAM**，全局内存对应 **DRAM**。 26 | 27 | ### Parallel Description 28 | 29 | **iGraph** 中的并行方式描述包含两个维度，分别是并行维度和循环维度，用元组 (numParallel, numLoop) 表示。前者表示算子实现中将计算任务在 numParallel 个最小并行单元中并行执行，并为每一个最小并行单元分配一个取值于 [0, numParallel) 的和硬件并行架构相容的并行编号。例如，对于硬件架构由 (108, 4) 来描述的硬件，其 numParallel = $108 \times 4$ = $432$，第1个并行组中的4个最小并行单元的编号为0、1、2、3，第2个并行组为4、5、6、7，以此类推。编号用于为不同的单元分配计算任务，其相容性让我们可以在并行组的层级上对数据的复用进行优化。后者表示了每个并行处理单元内会做长度为numLoop的串行循环，执行由 **iGraph** 内的 **iOp** 所定义的指令序列。其语义如下： 30 | 31 | ```python 32 | for parallel_id in 0...numParallel: # Parallel 33 | for loop_id in 0...numLoop: # Series 34 | loopBody(parallel_id, loop_id) 35 | ``` 36 | 37 | ### iSlice 38 | 39 | **iSlice** 用于表示每个并行处理单元在每个循环中所需要处理的数据切片，描述了数据如何在不同的并行单元和不同的循环上进行划分。所有 **iSlice** 都基于一个 **iPointer**，**iPointer** 表示在某内存层级上的一块连续内存，包含所在层级、名称、数据类型和长度等属性。其中，`dram` 上的iPointer在全局唯一，`sram` 上的iPointer在每个并行组内唯一，而 `reg` 上的iPointer则为每个最小并行单元独占。**iSlice** 则表示在它所基于的 **iPointer** 所指向的连续内存上的一个二维切片，由一组连续的内存片段组成，具有属性 `Shape`、`Stride` 和 `Offset`。`Shape` 是一个二元组，表示该内存切片的形状，`Stride` 也是一个二元组，表示在不同维度上的步长，`Offset` 则是一个关于 `parallel_id` 和 `loop_id` 的仿射函数，用于表示每一个最小并行单元在每个循环中所处理的内存切片的的首地址相对于其所基于的iPointer的偏移，用于表示输入输出数据和中间变量在不同并行单元和循环上的划分。 40 | 41 | ``` 42 | iPointer ::= iPointer(ID, Hierarchy, DType, Size) 43 | ID ::= Int 44 | Hierarchy ::= reg | sram | dram 45 | DType ::= fp64 | fp32 | fp16 ... 46 | Size ::= Int 47 | 48 | iSlice ::= iSlice(iPointer, Offset, Shape, Stride) 49 | Offset ::= Affine Function of (parallel_id, loop_id) 50 | Shape ::= (Int, Int) 51 | Stride ::= (Int, Int) 52 | ``` 53 | 54 | 举个例子，假设一个 **iGraph** 被用于表示一个张量加法算子，输入张量的形状为`(8， 1024)`，并且 `(numParallel, numLoop) = (64, 2)`，则这个输入张量可以用一个 **iPointer** 描述： 55 | ```c 56 | a = iPointer(0, dram, fp32, 8192) 57 | ``` 58 | 我们在各个最小并行单元间的各个循环间均匀分配所需要处理的输入数据，则数据的划分可以用 **iSlice** 描述： 59 | ```c 60 | aSlice = iSlice(a, [&](pid, lid) { return 4096 * lid + 128 * pid; }, (2, 32), (32, 1)) 61 | ``` 62 | 假设我们在算子实现中，需要先将数据读入寄存器中再进行计算，则可以按照如下描述这些寄存器，可以注意到寄存器是每个最小并行单元所独占的，所以 `b` 位于 `reg` 上，并且大小为64。 63 | ```c 64 | b = iPointer(1, reg, fp32, 64) 65 | bSlice = iPointer(b, [&](pid, lid) { return 0; }, (2, 32), (32, 1)) 66 | ``` 67 | 68 | ### iOp 69 | 70 | **iOp** 是对硬件指令的抽象，其输入输出都是上面定义的 **iSlice**，表示在内存切片上的硬件指令操作。作为对硬件指令的抽象，一个 **iOp** 通常可以在一个最小并行单元上由一条硬件指令或者一个硬件指令序列完成。**iOp** 分为访存iOp和计算iOp，分别用来表示访存操作和计算。 71 | 72 | #### 访存iOp 73 | 访存iOp分为两种，一种是移动（Move），另一种是同步（Sync）。 74 | 75 | 移动iOp表示在不同的内存层级间移动内存切片，或者在同一内存层级内对内存切片的布局进行改变，对应硬件指令上的内存相关指令，其语义是将位于 `from` 上的内存切片 `s1` 搬运到位于 `to` 上的内存切片 `t1` 上，并且搬运的数据类型为 `dtype`。实际上，由于 `s1` 和 `t1` 作为内存切片不仅包含其基于的指针、地址偏移、形状和步长等信息，也包含了该内存切片所在层级和内部的数据类型，因此 `move` 指令中的属性均可以从输入输出中推导出来，这里将其显式定义出来是为了易读性。如无特殊说明，下面介绍的其他 **iOp** 定义也遵循这一约定。 76 | 77 | ```c 78 | move.from.to.dtype t1, s1 79 | 80 | .from = {.dram, .sram, .reg} 81 | .to = {.dram, .sram, .reg} 82 | .dtype = {.fp64, .fp32, .fp16, .bf16, .fp8} 83 | 84 | // examples 85 | move.dram.sram.fp32 t1, s1 86 | move.sram.reg.fp16 t1, s1 87 | ``` 88 | 89 | 同步iOp则表示需要在某个内存层级上对相应的并行结构做同步，从而保证数据依赖的正确性。可以注意到 `sync` 指令要求 `iPointer(t1) = iPointer(s1)`，这是因为只有当算子实现中写入和读出的内存是重叠的时候才需要同步。例如，`s1` 和 `t1` 都对应于同一个 `dram` 上的iPointer，且有相同的形状和步长，但是二者的地址偏移不同，这就意味着一个最小并行单元所读取的切片可能是由别的单元写入的，因此需要进行同步才能保证数据依赖的正确性。 90 | 91 | ```c 92 | sync.scope t1, s1 93 | 94 | .scope = {.dram, .sram, .reg} 95 | s.t. iPointer(t1) = iPointer(s1) 96 | 97 | // examples 98 | sync.sram t1, s1 99 | ``` 100 | 101 | #### 计算iOp 102 | 103 | 计算iOp有几种类型，有 `unary`、`binary`、`broadcast`、`reduce` 和 `identity`。`unary` 为单目运算，语法如下，其中 `name` 表示运算类型，`dtype` 表示数据类型，`params` 为可选参数，用于 `muls`、`leaky_relu` 等指令。 104 | 105 | ```c 106 | unary.name.dtype t1, s1, [params] 107 | 108 | .name = {.muls, .adds, .divs, .subs, .exp, .log, .sin, .cos, .relu, ...} 109 | .dtype = {.fp64, .fp32, .fp16, .bf16, .fp8} 110 | 111 | // examples 112 | unary.muls.fp32 t1, s1 [2.0f, ] 113 | unary.exp.fp32 t1, s1 114 | ``` 115 | 116 | `binary` 为双目运算，其定义和 `unary` 类似。 117 | 118 | ```c 119 | binary.name.dtype t1, s1, s2 120 | 121 | .name = {.mul, .add, .div, .sub, ...} 122 | .dtype = {.fp64, .fp32, .fp16, .bf16, .fp8} 123 | 124 | // examples 125 | binary.mul.fp32 t1, s1, s2 126 | binary.sub.fp32 t1, s1, s2 127 | ``` 128 | 129 | `broadcast` 和 `reduce` 表示广播和规约iOp，其中 `dtype` 表示数据类型，`axis` 表示在进行规约或者广播的轴。而 `scope` 表示在什么范围内进行规约或者广播，其中 `unit` 表示在每个最小并行单元内，而 `group` 表示在同一并行组内的多个最小并行单元之间，`op` 表示用于规约的计算类型。对于 `group` 层级上的规约或者广播涉及到不同并行单元，因此需要在 `sram` 上通信，所以需要在 `params` 中指定一个位于 `sram` 上的iSlice作为buffer来辅助，另外参数groupSize指定的进行广播或者规约的组的大小，需要满足 `numUnit % groupSize == 0`。 130 | 131 | ```c 132 | broadcast.axis.scope.dtype t1, s1, [params] 133 | 134 | .axis = {.col, .row} 135 | .dtype = {.fp64, .fp32, .fp16, .bf16, .fp8} 136 | .scope = {.unit, .group} 137 | 138 | // example 139 | broadcast.row.unit.fp32 t1, s1 // Shape(s1) = (4, 1), Shape(t1) = (4, 128) 140 | broadcast.col.unit.fp32 t1, s1 // Shape(s1) = (1, 128), Shape(t1) = (4, 128) 141 | 142 | broadcast.row.group.fp32 t1, s1, [buffer=b1, groupSize=4] // Shape(s1) = (4, 1), Shape(t1) = (4, 128), numUnit = 8 143 | broadcast.col.group.fp32 t1, s1, [buffer=b1, groupSize=4] // Shape(s1) = (1, 128), Shape(t1) = (4, 128), numUnit = 8 144 | ``` 145 | 146 | 对于 `group` 层级上的规约，其最终规约的结果会保存在 `parallel_id % groupSize == 0` 的最小并行单元上；而对于 `group` 层级上的广播，其广播的值来自 `parallel_id % groupSize == 0` 的最小并行单元。 147 | 148 | ```c 149 | reduce.op.axis.scope.dtype t1, s1, [params] 150 | 151 | .op = {.max, .min, .add, .mul} 152 | .axis = {.col, .row} 153 | .dtype = {.fp64, .fp32, .fp16, .bf16, .fp8} 154 | .scope = {.unit, .group} 155 | 156 | // examples 157 | reduce.add.row.unit.fp32 t1, s1 // Shape(s1) = (4, 128), Shape(t1) = (4, 1) 158 | reduce.add.col.unit.fp32 t1, s1 // Shape(s1) = (4, 128), Shape(t1) = (1, 128) 159 | 160 | reduce.add.row.group.fp32 t1, s1, [buffer=b1, groupSize=4] // Shape(s1) = (4, 128), Shape(t1) = (1, 128), numUnit = 8 161 | reduce.add.col.group.fp32 t1, s1, [buffer=b1, groupSize=4] // Shape(s1) = (4, 128), Shape(t1) = (1, 128), numUnit = 8 162 | ``` 163 | 164 | 除了以上这些iOp之外，还有一类不表示具体计算但是对Unified IR的表达性十分重要的iOp，称为 `identity`。顾名思义，它表示输入输出的两组iSlice在物理内存上是同一块，可以用于在iGraph中表达Reshape、Split、Concat等逻辑。`identity` 指令要求输入输出的 `iSlice` 表示同一块的物理内存。 165 | 166 | ```c 167 | identity [ts], [ss] 168 | 169 | // example 170 | identity t1, s1 // Shape(t1) = (4, 64), Shape(s1) = (2, 128) 171 | identity t1, [s1, s2] // Shape(t1) = (4, 64), Shape(s1) = Shape(s2) = (4, 32) 172 | identity [t1, t2], s1 // Shape(t1) = Shape(t2) = (4, 32), Shape(s1) = (4, 64) 173 | ``` --------------------------------------------------------------------------------