IREE体系结构介绍

IREE体系结构介绍

基于MLIR的端到端编译器与运行时，可将ML模型下译到统一的IR，该IR可扩展到数据中心，也可扩展到移动与边缘部署。

中间表示执行环境

IREE关键特性

1）标准与基于社区

A）采用Vulkan、SPIR-V、WebGPU等，与OSS社区合作

2）基于编译

A）使用编译器弥合级别语义差距并生成最佳任务/作业

时间表（即ML的自动化任务系统中间件）

3）全面的

A）一个统一的IR，用于表示可调度的可执行文件与调度逻辑

实现整个程序的优化

4）可扩展

A）与其他加速器用户合作，了解资源限制，对多样化的使用与部署场景友好

IREE核心编译流程

 

IREE运行时

IREE没有捆绑所有内容的传统运行时。IREE提供了一个几乎零成本的虚拟机，用于解释从ML模型编译的主机调度操作。它只是为工作负载大小计算执行轻量级数学运算，并执行任务调度。

 

HAL：受Vulkan启发的硬件抽象层

1）CPU、GPU等的通用抽象

A）全部具有多级内存/计算层次结构

B）所有这些都是为了以平铺的方式进行计算

在主机上

1）构建管道以利用调度层次结构

A）提交（工作负载+粗粒度同步）

B）命令缓存（工作负载+细粒度同步）

2）Dispatches(→ GPU；→ CPU）

在可执行文件中（在硬件上；目标CodeGen）

1）工作组(→ GPU核心；→ CPU线程）

A）子组(→ GPU-SIMT；→ CPU：SIMD）

B）Instr(→ GPU线程；→ CPU:通道）

 

CPU与GPU计算与内存（抽象插图）

HAL IR示例

HAL具有映射到新一代显式GPU API（如Vulkan）的调度操作。这些操作有效地将Vulkan C API为用于自动转换的编译器IR开放，以便通过编译器实现最佳实践的编码。这就是在可执行文件与调度之间实现具体（基于绑定的）ABI。

//第10章/build_hal.executable.c

module {

  hal.executable @executable_module {

    hal.interface @abi {

 hal.interface.binding @ret, set=0, binding=0, type="StorageBuffer", access="Read"

...

 }

    hal.executable.binary {data = dense<[...]>: vector<1620×i8>, format = "SPIR-V"}...

 func @main(%arg0: !iree.ref<!hal.buffer>, %arg1: !iree.ref<!hal.buffer>) -> !iree.ref<!hal.buffer> {

 %dev = hal.ex.shared_device: !iree.ref<!hal.device>

 %allocator = hal.device.allocator %dev: !iree.ref<!hal.allocator>

 %buffer = hal.allocator.allocate %allocator,..., shape=[...], element_size=4: !iree.ref<!hal.buffer>

 %cmd = hal.command_buffer.create %dev, "OneShot", "Transfer|Dispatch": !iree.ref<!hal.command_buffer>

 hal.command_buffer.begin %cmd

 hal.ex.push_descriptor_set %cmd,...

 hal.device.switch(%dev: !hal.device)

 #hal.device.match.id<"vulkan*">: !iree.ref<!hal.buffer> {

...

 %exe = hal.executable.look_up,...

 hal.command_buffer.dispatch %cmd, %exe, entry_point=0, workgroup_xyz=[%c1, %c5, %c1]

 }

 %memory_barrier = hal.make_memory_barrier "DispatchWrite", "DispatchRead": tuple<i32, i32>

 hal.command_buffer.execution_barrier %cmd, "CommandRetire", "CommandIssue",

 memory_barriers=[%memory_barrier]

...

 hal.command_buffer.end %cmd

 hal.ex.submit_and_wait %dev, %cmd

 return %buffer: !iree.ref<!hal.buffer>

Vulkan现状与路线图

一般方法

1）投资基础设施，为通用性与性能奠定坚实基础

A）优先选择有利于泛型模型与/或体系结构的任务；

B）旨在为所有情况提供合理良好的默认解决方案；

C）为超级用户超级特定案例敞开大门。

2）MLIR中内置SPIR-V CodeGen，IREE中内置Vulkan运行时

A）可以在各种硬件上编译与执行许多视觉与语言模型；

B）瞄准Vulkan计算着色器与核心Vulkan计算机API子集。

3）到目前为止，重点关注广泛适用的编译优化

A）无手动/自动调谐；使用一组启发式方法与默认参数。

各种GPU的Transformer模型

	FP32 Model	GPU/FLOps	IREE延迟	比较延迟
移动	MobileBERT	ARM Mali G78 (Pixel 6)/2T	120ms	TFLite OpenCL 123ms
笔记本电脑	miniLM	Apple M1 Max/10.4T	11.6ms	TF-Metal 16.99ms
桌面	miniLM	AMD RX 5700XT/9.7T	8ms
服务器	miniLM	NVIDIA Tesla V100/15.7T	6.3ms

移动GPU上的模型

	FP32 Model	IREE延迟	TFLite延迟(缓存)	TFLite延迟 (结构)
谷歌Pixel6	MobileBERT	120ms	172ms	123ms
谷歌Pixel6	MobileNetV2	9ms	8ms	6ms
谷歌Pixel6	DeepLabV3	12ms	12.1ms	9.2ms
谷歌Pixel6	PoseNet	15ms	14.4ms	8ms

路线图与任务

1）一般功能特征，例如，

A）较小的位宽（fp16、int8等）；

B）减少初始化开销。

2）一般的优化，例如，

A）更好的融合，更好的缓存布局，支持纹理；

B）搜索，自动调谐；

3）更多的平台。

A）SPIR-V CodeGen+WebGPU HAL→ Web平台；

B）SPIR-V CodeGen+金属HAL→ 苹果平台。

 

参考文献链接

https://www.khronos.org/assets/uploads/developers/presentations/IREE_targeting_Vulkan_Zhang_May22.pdf