IREE AI编译器关键模块分析

IREE AI编译器设计大纲

概述

输入方言 - 量化
- 利用量化转换实现训练和推理时的量化，支持原数据类型运行，未来将探索直接与前端接口实现量化计算。
flow方言特性与优化
- 减少flow.stream读回操作，利用协程隐藏延迟，实现其在 CFG 中的线程化，对flow.dispatch进行谓词化处理，去重flow.executable，重新考虑 CSE 优化，根据设备特性进行操作放置。
hal层功能拓展计划
- 允许目标指定hal.interface，支持目标特定调度专业化，跟踪缓冲区使用，实现批处理可执行文件缓存和预编译、目标感知的压缩，优化命令缓冲区状态，管理资源时间线，使用瞬态张量环形缓冲区，在模块 ABI 上定义时间线信号量，采用类 GPU 的 CPU 调度策略。
vm虚拟机功能增强方向
- 引入协程支持批处理和协作调度，实现与 LLVM IR 的转换，考虑增加更多类型支持，探索在加速器上执行 VM 的方式。

详情

本文涵盖了 IREE 在设计过程中和未来版本计划中的各种特性，包括输入方言、flow、hal和vm等方面的设计规划。

输入方言 - 量化
- 当前计划使用量化转换来实现类型的训练和推理时量化，以保留最大精度，同时支持原始未量化浮点数运行，便于向量化过渡。
- 未来希望超越转换导向的量化方法，直接与具有足够定义类型系统的前端接口，以直接表示精确量化（及其他压缩）计算，减少对编译器端类型推断转换的依赖。
flow：数据和执行流建模
- 避免flow.stream的读回操作：多数现有flow.tensor.load.*操作（读回操作）将在实现 HLO 张量到基本类型转换后被移除。对于仍需读回的情况，IREE 会警告性能问题，鼓励调整输入模型。IREE VM 可通过协程有效隐藏读回延迟，例如对于动态副本（如 top-k + gather 操作），可通过合适的原语扩展，实现在同一流内计算索引和更新张量，避免主机往返。
- flow.stream在控制流图（CFG）中的线程化：当前flow.ex.stream.fragment是临时实现，为使流在更大并发范围内有效建模，需能跨 CFG 分支移动。转换为flow方言时，会遍历 CFG 并尝试在无外部依赖时将flow.stream值线程化，从而将整个流降低到一个命令缓冲区，无需主机往返。
- flow.dispatch的谓词化：对于执行依赖于先前调度结果的情况，flow.cond_dispatch允许提供条件来确定是否实际执行调度。对于支持命令缓冲区谓词化的目标（如 D3D12），可避免主机往返；对于不支持的目标（如 Vulkan，虽缺乏原生支持，但 Nvidia 通过扩展支持），可通过间接调度模拟谓词化，以减少开销。在flow级别建模谓词化，可降低到 HAL 时具有目标感知的谓词化语义，并融合间接调度工作组计数计算，减少开销。
- flow.executable的去重：在flow方言中，可利用 IR 树差异和 MLIR 规范化传递对目标无关的可执行文件进行去重，减少调度执行中的重复。
- 重新生成公共子表达式消除（CSE）后的表达式：CSE 虽常见，但在某些情况下（如广播操作被 CSE 且结果被独立调度使用），可能引入假依赖和额外分配。此时应在调度区域内重新生成广播，减少计算资源成本和中间张量需求，在多设备执行时需更谨慎平衡此优化。
- 设备放置：在flow方言中，可拆分流并安全调整操作，目标执行后端可根据设备限制（如最大在飞内存、最大调度深度和能力）进行操作。对于异构配置，可通过属性指定操作、调度和流应降低到的设备类别，约束求解可使用通用启发式方法、基于基准的配置文件引导数据库或机器学习获得的特征等。
hal：硬件抽象层和多架构可执行文件
- 允许目标指定hal.interface：hal.interface操作指定调度程序和设备之间的 ABI，包含缓冲区绑定和其他非缓冲区数据。目标后端可根据配置提供自己的接口，同一hal.executable可有多个接口，调度程序可根据接口差异生成适当的 HAL 操作。
- 目标特定的调度专业化：flow方言虽尝试融合操作，但并非所有后端都能将区域调度为单个调度。通过扩展目标后端的调度接口，后端可根据需要发出多个hal.executable和流命令，减少运行时分配和虚假依赖。调度专业化可根据调度参数（如归约形状）而变化，折叠和规范化可消除部分开销。
- 缓冲区使用跟踪：使用flow方言中 MLIR tensor的 SSA 形式值语义，可跟踪缓冲区使用情况，分析传递可标记张量，使hal方言分配缓冲区时选择合适内存类型和使用位，减少不必要的移动，传统系统使用启发式方法可能导致额外开销，而 IREE 可精确控制。
- 批处理可执行文件缓存和预编译：对于需要运行时预处理可执行文件的目标（如 SPIR-V 或 MSL），IREE HAL 基于 Vulkan 的管道缓存提供缓存和批编译机制。可对模块入口点进行可达性分析，预编译所需可执行文件，支持多线程编译，提高效率，模块可使用零个或多个作用域缓存，缓存可由宿主应用程序检索和保存。
- 目标感知的可执行文件压缩：将可执行文件表示为 IR 后，可应用后编译压缩技术，如针对 SPIR-V 可使用 SMOL-V 等压缩技术，结合批处理可执行文件缓存和预编译，可有效减少二进制大小。
- 目标感知的常量压缩：IREE 设计旨在实现高效的目标和上下文感知的大常量压缩，可重用 GPU 硬件压缩格式、ML 加速器特定格式或低比特深度量化技术，灵感来自 Crunch 和 Basis Universal 等格式，可能利用 GPU 硬件采样器进行解压。
- 命令缓冲区状态去重：IREE HAL 类似 Vulkan，大多使用不可变状态对象，但仍有少量状态入口点。对描述符集绑定和推送描述符等命令进行规范化和代码移动，可减少 IR、API 和执行开销。
- 资源时间线：IREE 调度程序的资源时间线概念允许重叠在飞调用，通过为可写资源分配时间线信号量，利用缓冲区使用跟踪和同步域信息，可有效同步资源，通过 IR 转换扩大时间窗口，提高重叠性，但对于资源间接和动态资源形状等情况可能需要其他技术辅助。
- 瞬态张量环形缓冲区：执行期间多数缓冲区不超出使用范围，可使用环形缓冲区（或双缓冲变体）存储瞬态张量数据和其他数据，通过 IR 计算动态形状张量大小，无需复杂运行时打包，可控制最大并发或内存使用，通过代码运动进行规划，减少寄存器压力，提高操作数量。
- 模块 ABI 上的时间线信号量：跨模块函数调用应能定义时间线信号量，自动为导出函数添加信号量，调用时填充，使调用自然链接内部异步工作，结合 VM 协程支持，可在等待和信号信号量之间交错主机执行，也可提供同步包装器，核心系统围绕单一系统支持的原语设计，避免额外复杂性。
- 类 GPU 的 CPU 调度：传统多线程方法在处理 IREE 的某些工作负载时可能成为瓶颈，IREE 将 CPU 核心视为 GPU 计算单元，通过flow和hal明确调度重叠和工作组大小，可避免管道气泡和不可预测调度。使用类似 marl 的调度器，即使仅针对 CPU，这种调度方式也有益，且对异构目标调度代码可能可共享。
vm：轻量级虚拟机
- 用于批处理和协作调度的协程：VM 当前缺少协程功能，协程可在模块内实现多在飞调用，无需复杂多线程逻辑。多数情况下，有时间线信号量暴露给调用者时无需在 VM 中 yield。对于无法移除的主机读回情况，编译器可发出显式 yield 点，VM 运行时遇到 yield 点会暂停协程，直到满足条件。唤醒协程可由应用程序提供回调或使用辅助线程，利用协程可提高吞吐量，但不降低每次调用延迟。此外，基于协程的蜂窝批处理可进一步减少延迟，通过识别可分区和贪婪调度的小均匀工作，实现批处理或降低相关调用成本，具体逻辑可内置于模块中，设计工作仍需确定如何在 IR 中表示，长期来看是主要研究领域之一。
- 降低到 LLVM IR：对于无需动态模块加载的场景，可将 VM IR 降低到 LLVM IR，将vm.call操作转换为llvm::CallInst，实现运行时解析函数指针，启用异构 / 运行时确定设备的灵活性、可插拔诊断和后端组合，还可扩展到 “无运行时模式”，减少代码大小。
- 改进类型支持：VM 目前仅支持i32和vm.ref<T>两种类型，未来可能引入f32、list/dict和vector<4xf32>等类型，以支持更复杂计算和提高与其他语言（如 Python）的兼容性。
- 间接命令缓冲区 / 在加速器上执行：尽管 IREE 使用多种技巧减少主机往返，但命令记录和提交仍在主机 CPU 上。对于低功耗始终在线计算或分支行为明显的应用，决策逻辑应尽可能靠近执行管道实时运行。IREE VM 设计为可在设备上安全协作运行，可通过将 VM IR 降低到 LLVM IR、转换为目标特定形式或直接在设备上执行 VM 字节码等方式，原型化设备上的完整使用，减少主机和设备调度的紧密耦合。