IREE项目中EncodingAttr广播映射信息不足问题分析

背景介绍

在IREE编译器项目中，EncodingAttr是用于描述张量编码属性的重要数据结构。它包含了操作数索引、操作类型、元素类型以及用户索引映射等信息。其中，bcastMap(广播映射)属性最初设计用于处理广播操作，但随着项目发展，这种设计被发现无法满足更复杂的编码传播需求。

问题本质

当前EncodingAttr中的bcastMap属性存在信息不足的问题。具体表现为：

1. 广播映射与用户索引映射分离存储，导致在编码传播过程中无法完整表达映射关系
2. 当进行多次映射组合时，原始设计无法保留足够的维度信息
3. 在推断B/M/N/K维度和确定各维度的内部平铺大小时，信息可能丢失

技术分析

现有设计局限性

现有设计将广播映射(bcastMap)与用户索引映射(user_indexing_maps)分开存储。例如：

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

这种分离式存储导致在编码传播过程中，当需要组合多个映射时，无法完整表达映射关系链。

改进方案

解决方案是将广播映射整合到用户索引映射中，形成映射链。例如：

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

这种设计允许通过AffineMap::compose方法递归解析映射关系，同时保留了原始映射信息。

关键考量

改进方案需要特别注意：

1. 维度信息保留：必须保留原始收缩索引映射以推断B/M/N/K维度
2. 批量维度识别：需要确保能够识别出现在所有索引映射中的维度(批量维度)
3. 映射组合顺序：需要明确映射组合的顺序和方式

实现影响

这一改进将影响：

1. 编码属性定义：需要修改EncodingAttr的定义方式
2. 编码传播逻辑：需要调整编码传播过程中的映射处理方式
3. 维度推断机制：需要确保维度推断仍然能够正确工作
4. 平铺大小确定：需要保证内部平铺大小的计算不受影响

技术价值

这一改进将为IREE项目带来以下好处：

1. 更强大的编码表达能力：能够处理更复杂的映射组合场景
2. 更准确的维度推断：保留原始映射信息确保维度推断正确性
3. 更灵活的编码传播：支持多级映射组合的传播过程
4. 更好的可扩展性：为未来可能的编码扩展奠定基础

总结

IREE项目中EncodingAttr的广播映射改进是编码系统演进的重要一步。通过将广播映射整合到用户索引映射链中，不仅解决了当前的信息不足问题，还为未来的编码功能扩展提供了更强大的基础架构。这一改进需要谨慎处理映射组合和维度信息保留等关键问题，以确保不影响现有的编码传播和优化逻辑。