IREE 3.5.0 版本发布：编译器与运行时系统的重大升级

IREE（Intermediate Representation Execution Environment）是一个专注于机器学习模型部署的开源项目，它通过多层中间表示（IR）将高级机器学习模型转换为高效的可执行代码，支持多种硬件后端。IREE 的设计理念强调模块化和可扩展性，使其能够适应从嵌入式设备到数据中心的各种计算场景。

近日，IREE 发布了 3.5.0 版本，带来了多项编译器优化和运行时系统增强。本文将深入解析这一版本的主要技术亮点。

编译器架构的重大改进

AMD GPU 支持扩展

3.5.0 版本显著增强了对 AMD GPU 的支持，新增了 Radeon 9060XT 和 Radeon PRO AI R9070 GPU 的兼容性。特别值得注意的是新增的 gfx950 目标定义，它引入了多种新型矩阵融合乘加运算（MFMA），这些运算在深度学习工作负载中表现尤为出色。

在底层实现上，开发团队优化了数据传输策略对中等尺寸扩展形状 FP8 的支持，并移除了动态 M 维度检查，这些改进直接提升了矩阵运算在 AMD GPU 上的执行效率。

布局转换优化

新引入的 CombineLayoutTransformation 机制能够将转置、重塑和切片操作合并为单一的 iree_linalg_ext.map_scatter 操作。这一优化不仅简化了中间表示，还能显著减少内存访问开销。后续的增强使其进一步支持将 tensor.pad 操作折叠到 map_scatter 操作中，完善了布局转换的整体流程。

编码专业化处理

SpecializeEncodings 通道的处理能力得到显著提升，现在能够更好地处理基于填充的编码方案。这一改进使得非可序列化编码能够转换为可序列化形式，为模型部署提供了更大的灵活性。同时，开发团队还优化了调度创建机制，限制了填充编码的应用范围，并引入了新的逻辑来防止与填充编码相关的 set_encoding 和 unset_encoding 操作被提升。

运行时系统的创新特性

条件化执行支持

3.5.0 版本引入了 hal.executable.export 条件区域，这是一个重要的架构创新。它允许在每个调度站点基于设备能力和工作负载参数做出执行决策，为异构计算环境提供了更精细的控制能力。

内存管理增强

在内存管理方面，新增了用户定义的 IREE_ALLOCATOR_SYSTEM 支持，允许外部覆盖分配器控制函数。同时加入了 mimalloc v3 作为可选系统分配器的支持，开发者可以通过设置编译标志静态链接 mimalloc 到 iree::base 中，这对内存敏感型应用尤其有价值。

设备异构性支持

虽然仍处于开发初期，但 3.5.0 版本已开始支持异构设备架构。这项功能将允许编译后的程序在兼容设备间分配缓冲区，并通过信号量同步操作。当前阶段主要支持 CPU-only 配置，但为未来的多设备协同计算奠定了基础。

类型系统与算子优化

新型数据类型支持

该版本扩展了对多种新型数据类型的支持，包括 F8E8M0FNU 类型的确认为有效 HAL 元素类型，以及多种缩放 MFMA 类型（f8E8M0FNU、f6E3M2FN 和 f4E2M1FN）的加入。这些数据类型特别适合新兴的机器学习工作负载，为模型压缩和量化提供了更多选择。

线性代数扩展改进

Linalg Extensions Dialect 在这一版本中获得了多项改进，包括对 map_scatter 操作的平铺接口实现、移除 GatherOp 的区域限制，以及新增的排序操作规范化模式。这些改进共同提升了高级线性代数操作的表达能力和执行效率。

性能优化与调试工具

开发团队在 3.5.0 版本中投入了大量精力优化性能热点。例如，通过限制 GPUAllocPrivateMemoryForDPSOps 通道为纯张量语义，解决了近期实现中的性能问题。同时，优化了注意力操作，移除了掩码操作数中的单位维度，减少了不必要的计算开销。

在调试方面，增强了 DumpDispatchGraph 通道的功能，使其能够更好地处理模型级别的程序，并添加了亲和性信息输出，为性能分析和调优提供了更多线索。

总结

IREE 3.5.0 版本在编译器优化、运行时系统、数据类型支持和性能调优等多个方面都取得了显著进展。特别是对 AMD GPU 的增强支持和条件化执行机制的引入，为机器学习模型的高效部署提供了更强大的工具链。随着异构设备支持的逐步完善，IREE 正在成为一个更加全面和灵活的模型部署解决方案，能够满足从边缘设备到云端服务器的各种计算需求。

这些改进不仅提升了框架的性能和灵活性，也为开发者提供了更多优化模型执行的选择。随着项目的持续发展，IREE 有望在机器学习部署领域扮演越来越重要的角色。