PR Category

CINN

PR Types

Improvements

Description

Pcard-76996

Done in this PR

实现新的多下游融合方案，利用 Loop/Value 轴之间的依赖路径来推导出上下游 Pattern 融合需要进行的循环变换，以下是该方案中的几个关键数据结构：

1. 实现 LoopAxisMapping 数据结构：保存 Pattern Loop 到每个输入/输出 Value 之间轴的变换及相应的逆变换，其中 Loop 对应了后端 Lowering 时的 Expr 表示，同时在构建每个初始 Op 的 LoopAxisMapping 时会考虑 0-d 的情况以及 1-size axis 的情况，如一个 ReduceOp 的 LoopAxisMapping 表示如下：

(%0) = "cinn_op.reduce_max" [id:38] (%1) {axis:[1],keepdim:true,stop_gradient:[true]} : (tensor<16x32x8xf32>) -> tensor<16x1x8xf32>(0x510f5c38)
LoopAxisMapping is: 
 input 0 :	[16, 32, 8], 0x510f5a50
  loop   :	[16, 1, 8, 32], reduce_axis_num: 1
output 0 :	[16, 1, 8], 0x510f5c20, use_count: 1
input2loop  0 : Transpose{perm=(0,2,1)} -> AppendAxis{axis=(1), shape=(1)}
loop2output 0 : DeleteAxis{axis=(3), shape=(32)}

2. 实现不同融合策略下的 LoopAxisMapping 融合：从上下游 Pattern 的 LoopAxisMapping 通过 Value 连接的路径推导出融合 Pattern 地 LoopAxisMapping
3. 实现 SimplifyTransformRoute：通过模拟实现变换路径的简化，能够避免后期融合 Pattern 推导出来的变换路径太长影响编译性能
4. 实现 GetValidLoopTransformRoute：尝试获取多下游融合中上下游之间有效的循环变换路径，从而判断是否能够进行融合，以及如何对进行变换融合

Todo

1. 启用多下游融合中 Reshape 变换，目前已经开发完成但是因为存在性能问题和精度问题 Ban 了，待修复后打开
2. 将 Reduce 单下游融合合入到当前多下游融合方案中
3. 移除旧方案相关的无用代码