IREE项目中消除冗余编码调度的优化方案

在IREE编译器项目中，开发团队发现当后端不支持数据平铺(data-tiling)时，延迟物化路径(late materialization path)的性能会受到显著影响。问题的根源在于某些设置/取消编码(set/unset_encoding)操作会形成自己的调度(dispatch)，导致额外的内核启动和不必要的数据拷贝。

问题背景

在编译器优化过程中，IREE会处理张量的编码操作。当后端不支持数据平铺时，这些编码操作会变成冗余操作，但仍然会产生额外的调度和内存拷贝开销。特别是在全局优化(GlobalOpt)阶段设置的编码，这种冗余操作会一直存在，直到完全支持数据平铺融合。

解决方案设计

开发团队设计了一个系统性的解决方案，主要包含以下几个关键步骤：

1. 引入新的编码操作：创建flow.tensor.encode和stream.tensor.encode操作，作为编码操作的统一表示。

2. 添加直接降级路径：实现从flow.tensor.encode到stream.tensor.encode的直接转换。

3. 转换现有操作：将现有的set_encoding调度识别并转换为新的flow.tensor.encode操作。

4. 实现编码物化过程：在Stream方言中添加MaterializeEncoding过程，位于EncodeHostTensors之后。这个过程会：

   • 进行编码特化
   • 解析布局
   • 当张量编码产生相同布局时，折叠掉冗余操作

5. 清理依赖关系：最终目标是切断代码生成对Flow方言的依赖，将这些操作迁移到专门的方言中。

技术实现细节

在实现过程中，开发团队遇到并解决了一些技术挑战：

1. 亲和性分析集成：新操作需要正确集成到IREE的亲和性分析系统中。这涉及到将操作添加到AffinityOpAttrExternalModel列表中，确保它们能参与正确的调度分配。

2. 布局一致性检查：实现了isSameLayout方法作为编码属性接口的一部分，用于判断编码操作是否实际改变了张量布局。

3. 操作折叠优化：当编码操作不改变实际布局时，通过优化过程将其完全消除，避免产生任何运行时开销。

4. 冗余取消编码处理：扩展方案以处理unset_encoding调度，确保全局优化阶段设置的编码也能被正确优化。

性能影响

这一优化显著改善了不支持数据平铺后端的性能表现，特别是：

• 消除了不必要的内核启动
• 减少了冗余内存拷贝
• 保持了原有功能完整性
• 为未来支持完整数据平铺融合奠定了基础

总结

IREE团队通过引入新的编码操作表示和相应的优化过程，系统性地解决了冗余编码调度导致的性能问题。这一改进不仅提升了当前不支持数据平铺后端的性能，还为未来的功能扩展提供了更清晰的操作抽象和优化框架。