IREE项目中编码属性广播映射信息不足的问题分析

背景介绍

在IREE编译器项目中，EncodingAttr（编码属性）是处理张量操作编码的重要数据结构。当前实现中使用的bcastMap（广播映射）机制在编码传播过程中存在信息不足的问题，这限制了编码系统的表达能力，特别是在处理复杂张量操作时。

问题描述

当前的EncodingAttr使用bcastMap来表示广播操作，这种设计在原型阶段足够使用，但随着编码传播机制的复杂化，单一的广播映射已经无法满足需求。特别是在进行多次映射组合时，bcastMap无法完整保留原始操作的所有必要信息。

技术分析

当前实现的问题

现有实现中，编码属性包含以下关键部分：

• operand_index：操作数索引
• op_type：操作类型（如矩阵乘法）
• element_types：元素类型列表
• user_indexing_maps：用户索引映射列表
• bcast_map：广播映射

这种设计的主要局限性在于：

1. 广播映射与用户索引映射分离，导致信息不连贯
2. 在传播过程中进行映射组合时会丢失原始维度信息
3. 难以准确推断B/M/N/K等关键维度
4. 无法有效确定各维度的内部平铺大小

解决方案设计

改进方案是将广播映射整合到用户索引映射中，形成嵌套的映射结构。具体来说：

1. 将bcastMap转换为用户索引映射列表的一部分
2. 使用AffineMap::compose方法递归组合映射
3. 保留原始映射信息以支持维度推断

例如，原始编码：

#encoding1 = #iree_encoding.encoding<
  user_indexing_maps = [#map, #map1, #map2],
  bcast_map = #map3>

改进后的编码：

#encoding1 = #iree_encoding.encoding<
  user_indexing_maps = [#map, [#map1, #map3], #map2]>

实现考量

信息保留的重要性

在映射组合过程中，必须保留原始映射信息，原因包括：

1. 准确识别批处理维度（出现在所有索引映射中的维度）
2. 支持维度推断算法
3. 保持操作语义的完整性

技术实现细节

1. 使用嵌套的映射列表而非单一映射
2. 延迟映射组合时机，仅在必要时进行
3. 保持原始映射结构以支持分析
4. 提供组合映射的访问接口

影响范围

这一改进将影响：

1. 编码属性定义和解析
2. 编码传播算法
3. 维度推断逻辑
4. 平铺策略决策

结论

将EncodingAttr从bcastMap迁移到映射列表的设计，解决了编码传播中的信息不足问题，为IREE编译器处理更复杂的张量操作提供了坚实的基础。这一改进不仅增强了系统的表达能力，还保持了必要的原始信息以支持后续优化决策。