MLIR 字节码

本文翻译自:MLIR Bytecode Format
https://mlir.llvm.org/docs/BytecodeFormat/



MLIR : Multi-Level IR Compiler Framework


一、关于 MLIR 字节码格式

本文档描述了 MLIR 字节码格式及其编码。此格式已版本化且稳定:我们不会破坏兼容性,也就是说,dialect 应该能够反序列化 较旧的字节码。同样,我们支持反向部署,以便可以针对较旧版本的格式。

也就是说,重要的是 要认识到 字节码格式的承诺是 假设不可变的方言(immutable dialects):格式允许向后和向前兼容,但前提是方言中没有任何变化(操作、类型、属性定义)。

方言 可以选择通过 BytecodeDialectInterface 来处理自己的版本控制。方言中暴露了一些钩子,以允许管理编码到字节码文件中的版本。

版本是延迟加载的,允许在解码输入 IR 时检索版本信息,并为存在版本的每种方言提供机会,通过该upgradeFromVersion方法在解析后执行 IR 升级。

对于方言可以编码哪种信息来模拟其版本控制,没有任何限制。


Magic Number 魔数

MLIR 使用以下四字节幻数来指示字节码文件:

‘[‘M’ 8 , ‘L’ 8 , ‘ï’ 8 , ‘R’ 8 ]’

十六进制:

‘[‘4D’ 8 , ‘4C’ 8 , ‘EF’ 8 , ‘52’ 8 ]’


二、格式概述

MLIR 字节码文件由字节流和其顶层的一些简单的结构概念组成。


原语


固定宽度整数
  byte ::= `0x00`...`0xFF`

固定宽度整数是 已知字节大小的 无符号整数。值以小端字节顺序存储。

TODO:根据需要添加更大的固定宽度整数。


可变宽度整数

可变宽度整数(或称VarInts)为整数提供了一种紧凑的表示形式。

每个编码的 VarInt 由一到九个字节组成,它们共同表示一个 64 位值。

MLIR 字节码对 VarInt 使用 PrefixVarInt 编码。此编码是 LEB128(“Little-Endian Base 128”) 编码的一种变体,其中编码的每个字节为值提供最多 7 位,剩余的位用于存储指示用于编码的字节数的标签。

这意味着小的无符号整数(小于 2^7)可以存储在一个字节中,最大 2^14 的无符号整数可以存储在两个字节中,等等。

编码的第一个字节在低位中包含一个长度前缀。此前缀是0的位序列,后跟终结符 1,即字节的结尾。

“0”位数表示用于编码值的附加字节数(不包括前缀字节)。第一个字节中的所有剩余位以及附加字节中的所有位提供整数的值。以下是前缀字节的各种可能编码:

xxxxxxx1:  7 value bits, the encoding uses 1 byte
xxxxxx10: 14 value bits, the encoding uses 2 bytes
xxxxx100: 21 value bits, the encoding uses 3 bytes
xxxx1000: 28 value bits, the encoding uses 4 bytes
xxx10000: 35 value bits, the encoding uses 5 bytes
xx100000: 42 value bits, the encoding uses 6 bytes
x1000000: 49 value bits, the encoding uses 7 bytes
10000000: 56 value bits, the encoding uses 8 bytes
00000000: 64 value bits, the encoding uses 9 bytes

有符号可变宽度整数

有符号可变宽度整数值的编码方式与 varints类似,但采用 zigzag 编码

此编码使用值的低位 来表示符号,从而可以更有效地编码负数。

如果使用普通 varint编码负值,它将被视为非常大的无符号值。

使用 zigzag 编码允许值中的活动位数更少,从而导致较小的编码。

以下是生成 zigzag 编码的基本计算:

(value << 1) ^ (value >> 63)

字符串

字符串是具有相关长度的字符块。


部分

section {
  idAndIsAligned: byte // id | (hasAlign << 7)
  length: varint,

  alignment: varint?,
  padding: byte[], // Padding bytes are always `0xCB`.

  data: byte[]
}

节是一种在字节码内对数据进行分组的机制。它们支持延迟处理,这对于无序数据处理、延迟加载等非常有用。每个节都包含一个节 ID,其高位表示该节是否有对齐要求、长度(允许跳过该节)和可选对齐。当存在对齐时,节数据前可能会出现可变数量的填充字节 (0xCB)。节的对齐必须是 2 的幂。


三、MLIR 编码

鉴于 MLIR 的通用结构,字节码编码实际上相当简单。它有效地映射到 MLIR 的核心组件。


1、顶层结构

字节码的顶层结构包含 4 字节“魔法数字”、版本号、以空字符结尾的生产者字符串和部分列表。

目前,每个部分在字节码文件中仅应出现一次。

bytecode {
  magic: "MLïR",
  version: varint,
  producer: string,
  sections: section[]
}

2、字符串部分

strings {
  numStrings: varint,
  reverseStringLengths: varint[],
  stringData: byte[]
}

字符串部分包含 字节码中 引用的字符串表,可更轻松地实现 字符串共享。

此部分首先使用字符串总数进行编码,然后以 相反的顺序 使用每个单独字符串的大小进行编码。

其余编码包含一个 blob,其中包含连接在一起的所有字符串。


3、方言部分

字节码的方言部分包含编码 IR 中引用的所有方言,以及有关所引用的方言组成部分的一些信息。

dialect_section {
  numDialects: varint,
  dialectNames: varint[],
  numTotalOpNames: varint,
  opNames: op_name_group[]
}

op_name_group {
  dialect: varint // (dialectID << 1) | (hasVersion),
  version : dialect_version_section
  numOpNames: varint,
  opNames: varint[]
}

dialect_version_section {
  size: varint,
  version: byte[]
}

方言被编码为varint包含字符串部分中名称字符串的索引,以及指示方言是否已版本化的标志。

操作名称按方言分组编码,每组包含方言、操作名称的数量以及字符串部分中每个名称的索引数组。版本被编码为嵌套部分。


4、属性/类型部分

属性和类型使用两个部分进行编码 ,一个部分(attr_type_section)包含实际的编码表示,另一个部分(attr_type_offset_section)包含 每个编码属性/类型到前一个部分的偏移量。

此结构允许属性和类型始终按需延迟加载。

attr_type_section {
  attrs: attribute[],
  types: type[]
}
attr_type_offset_section {
  numAttrs: varint,
  numTypes: varint,
  offsets: attr_type_offset_group[]
}

attr_type_offset_group {
  dialect: varint,
  numElements: varint,
  offsets: varint[] // (offset << 1) | (hasCustomEncoding)
}

attribute {
  encoding: ...
}
type {
  encoding: ...
}

上面的每一个attr_type_offset_section 中的 offset 都是属性或类型的编码大小,以及一个标志 用来指示编码 是 使用文本汇编格式 还是自定义字节码编码。

我们避免使用直接偏移量 attr_type_section,因为较小的相对偏移量可以提供更有效的压缩。

属性和类型按方言分组,attr_type_offset_group偏移量部分中的每个都包含相应的父方言、元素数量以及组内每个元素的偏移量。


4.1 属性/类型编码

从抽象的角度来说,属性/类型可以通过两种可能的方式之一进行编码:通过其汇编格式,或者通过自定义方言定义的编码。


4.1.1 汇编格式回退

如果方言未定义对属性或类型进行编码的方法,则将使用该属性或类型的文本汇编格式作为后备。例如,类型!bytecode.type将被编码为以空字符结尾的字符串“!bytecode.type”。这确保可以对每个属性和类型进行编码,即使所属方言尚未选择更高效的序列化。

TODO:我们不应该在这里重复编码方言名称,而应该使用对父方言的引用。


4.1.2 方言定义编码

除了汇编格式回退之外,方言还可以为其属性和类型提供自定义编码。自定义编码非常有益,因为它们明显更小,读写速度更快。

方言可以通过实现 来选择提供自定义编码 BytecodeDialectInterface。此接口提供钩子,即 readAttribute/readTypewriteAttribute/ writeType,供字节码读取器和写入器使用。这些钩子提供了一个读取器和写入器实现,可用于以底层字节码格式对各种构造进行编码。此接口的一个独特功能是方言可以选择仅以自定义字节码格式对其属性和类型的子集进行编码,这可以简化添加尚未完全成型的新组件或实验性组件的过程。

在实现字节码接口时,方言负责编码的所有方面。这包括指示正在编码哪种属性或类型;字节码读取器只会知道它遇到了给定方言的属性或类型,而不会编码任何其他信息。因此,一种常见的编码习惯用法是使用前导代码 varint来指示属性或类型的编码方式。


5、资源部分

资源使用两个部分进行编码 ,一个部分(resource_section)包含实际的编码表示,另一个部分(resource_offset_section)包含每个编码资源到前一个部分的偏移量。

resource_section {
  resources: resource[]
}
resource {
  value: resource_bool | resource_string | resource_blob
}
resource_bool {
  value: byte
}
resource_string {
  value: varint
}
resource_blob {
  alignment: varint,
  size: varint,
  padding: byte[],
  blob: byte[]
}

resource_offset_section {
  numExternalResourceGroups: varint,
  resourceGroups: resource_group[]
}
resource_group {
  key: varint,
  numResources: varint,
  resources: resource_info[]
}
resource_info {
  key: varint,
  size: varint
  kind: byte,
}

资源按提供者(外部实体或方言)分组,每个资源的resource_group偏移部分包含相应的提供者、元素数量以及组内每个元素的信息。

对于每个元素,我们记录键、值类型和编码大小。

我们避免使用直接偏移量resource_section,因为较小的相对偏移量可以提供更有效的压缩。


6、IR 部分

IR 部分包含字节码内操作的编码形式。

ir_section {
  block: block; // Single block without arguments.
}

6.1 操作编码
op {
  name: varint,
  encodingMask: byte,
  location: varint,

  attrDict: varint?,

  numResults: varint?,
  resultTypes: varint[],

  numOperands: varint?,
  operands: varint[],

  numSuccessors: varint?,
  successors: varint[],

  numUseListOrders: varint?,
  useListOrders: uselist[],

  regionEncoding: varint?, // (numRegions << 1) | (isIsolatedFromAbove)

  // regions are stored in a section if isIsolatedFromAbove
  regions: (region | region_section)[]
}

uselist {
  indexInRange: varint?,
  useListEncoding: varint, // (numIndices << 1) | (isIndexPairEncoding)
  indices: varint[]
}

操作的编码很重要,因为这通常是字节码中最常见的结构。每种类型的操作都使用一种编码。鉴于这种普遍性,操作的许多字段都是可选的。该encodingMask字段是一个位掩码,指示操作的哪些组件存在。


地点

该位置被编码为属性表内位置的索引。


属性

如果操作具有属性,则对属性表内的操作属性字典的索引进行编码。


结果

如果操作有结果,则对结果的数量和类型表中结果类型的索引进行编码。


操作数

如果操作有操作数,则对操作数的数量和每个操作数的值索引进行编码。该值索引是该值的定义从第一个祖先隔离区域的开始的相对顺序。


后继者

如果操作有后继,则对后继的数量和父区域内的后继块的索引进行编码。


使用列表顺序

参考使用列表顺序被假定为通过 IR 的预排序遍历获得的所有操作数的全局枚举的逆序。

此顺序是通过逐个操作构建操作块自然获得的。

但是,某些转换可能会根据此参考顺序打乱使用列表。

如果任何操作结果的使用列表顺序未根据参考使用列表顺序排序,则会发出一种编码,以便在解析字节码后可以重建这种顺序。

编码表示从参考操作数顺序到当前使用列表顺序的索引映射。

位标志用于检测此编码是否为索引对类型。当位标志设置为零时,元素表示引用列表的i使用i在当前使用列表中的位置。

当位标志设置为时1,编码表示索引对(i, j),这表示i引用列表位置处的使用已映射到当前使用列表中的位置j

当当前使用列表中只有不到一半的元素相对于参考使用列表被打乱时,使用索引对编码来减少字节码内存要求。


区域

如果操作有区域,则区域数量以及区域是否与上方隔离将一起编码在单个 varint 中。之后,每个区域都以内联方式进行编码。


6.2 区域编码
region {
  numBlocks: varint,

  numValues: varint?,
  blocks: block[]
}

区域首先用其中的块数进行编码。如果区域非空,则先编码区域内直接定义的值数,然后再编码区域的块数。


6.2 块编码
block {
  encoding: varint, // (numOps << 1) | (hasBlockArgs)
  arguments: block_arguments?, // Optional based on encoding
  ops : op[]
}

block_arguments {
  numArgs: varint?,
  args: block_argument[]
  numUseListOrders: varint?,
  useListOrders: uselist[],
}

block_argument {
  typeAndLocation: varint, // (type << 1) | (hasLocation)
  location: varint? // Optional, else unknown location
}

块由一组操作和块参数编码而成。第一个字段是将块中的操作数与指示块是否具有参数的标志相结合的编码。

使用列表顺序附加到块参数的方式与附加到操作结果的方式类似。


2024-05-30(四)

### TPU 和 MLIR 的集成及其硬件加速应用 Tensor Processing Unit (TPU) 是一种专为机器学习工作负载设计的定制化 ASIC 芯片[^1]。这些芯片旨在高效执行张量运算,这是许多现代神经网络的核心操作。 MLIR(多级中间表示)是一个编译器框架,用于构建可互操作和可重用的编译工具链组件[^2]。通过提供统一且灵活的中间表示形式,MLIR 改善了不同前端语言到后端目标架构之间的转换过程。 #### 集成方式 为了使模型能够在 TPU 上运行得更快更有效率,通常会经历以下几个阶段: - **模型定义**:开发者使用 TensorFlow 或其他支持 MLIR 导出功能的框架来创建并训练自己的模型。 - **图优化**:利用 MLIR 提供的各种变换机制对计算图进行优化处理,减少冗余计算节点数量以及调整内存访问模式等措施提升性能表现[^3]。 - **代码生成**:基于经过优化后的 IR 表达式序列,借助特定于 TPU 架构的知识库自动生成高效的汇编指令或者二进制文件格式以便加载至实际设备当中去执行推理任务。 ```cpp // 假设有一个简单的卷积层实现作为例子展示如何编写针对 TPU 的内核函数 void ConvLayer(const float* input, const float* filter, int stride_h, int stride_w, int pad_top, int pad_left, float* output) { // 这里省略具体实现细节... } ``` 这种紧密耦合不仅使得开发人员可以更容易地将自己的算法移植到新的平台上测试效果;同时也促进了整个生态系统内部各部分之间更好的协作关系形成良性循环促进技术进步与发展。 #### 性能优势 当 TPUs 与 MLIR 结合起来时,在多个方面都展现出了显著的优势: - 更高的吞吐量意味着单位时间内能够完成更多次预测请求; - 较低延迟有助于实现实时应用场景下的快速响应需求; - 减少功耗对于移动终端来说尤为重要因为它直接影响电池续航时间长短。 综上所述,TPU 和 MLIR 的结合为机器学习领域带来了强大的硬件加速能力,并简化了从高级描述向底层实现转化的过程。
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包

打赏作者

编程乐园

请我喝杯伯爵奶茶~!

¥1 ¥2 ¥4 ¥6 ¥10 ¥20
扫码支付:¥1
获取中
扫码支付

您的余额不足,请更换扫码支付或充值

打赏作者

实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值