|
2 | 2 | CINN 提供了图层优化、低层 kernel 调优等编译优化技术加速计算图的运行效率,我们通过对 Paddle 训练框架执行过程进行一定升级改造,使得它能够联合 CINN 提升整体训练性能。本文介绍 Paddle 侧接入 CINN 的流程方案,包括框架编译期和运行时两方面的逻辑。 |
3 | 3 |
|
4 | 4 | ## 编译期修改计算图 |
5 | | - 这部分的工作是圈定计算图中可被 CINN 编译优化的子图, 并将待编译的子图替换为大 Op, 使得框架运行时通过执行器调度驱动后续的子图编译和执行,它是在编译期模型网络 ProgramDesc 转换为 Graph 之后,在框架 Graph Pass 阶段插入一个[build_cinn_pass](https://github.com/PaddlePaddle/Paddle/blob/develop/paddle/fluid/framework/paddle2cinn/build_cinn_pass.cc)来实现的,具体可分为如下两个步骤: |
| 5 | + 这部分的工作是圈定计算图中可被 CINN 编译优化的子图, 并将待编译的子图替换为大 Op, 使得框架运行时通过执行器调度驱动后续的子图编译和执行,它是在编译期模型网络 ProgramDesc 转换为 Graph 之后,在框架 Graph Pass 阶段插入一个[build_cinn_pass](https://github.com/PaddlePaddle/Paddle/blob/develop/paddle/fluid/pir/transforms/build_cinn_pass.cc)来实现的,具体可分为如下: |
6 | 6 |
|
7 | 7 | - **可编译算子搜索与子图标记**: 首先搜索识别原始计算图中哪些算子可以支持 CINN, 然后合并相邻选中算子聚合成若干个待编译子图,过程如下示例图: |
8 | 8 |
|
|
11 | 11 | Pass 运行结束后计算图形态表现为 cinn_launch_op 之间通过其它常规算子连通,但彼此互不邻接。 |
12 | 12 |
|
13 | 13 | ## 运行时编译(JIT)与执行编译结果 |
14 | | - 子图的编译和计算都在运行期完成,逻辑设计实现在[cinn_launch_op](https://github.com/PaddlePaddle/Paddle/blob/develop/paddle/fluid/operators/cinn/cinn_launch_op.h)算子中, 每个 Op 负责一个子图并分别独立运行,从框架执行器视角来看,它与常规算子 kernel 一致,同样都是在运行期调度执行该 kernel, 因此保证了与现有运行逻辑兼容。具体到 kernel 的逻辑上, 它先后执行子图编译和子图计算两阶段: |
15 | | - - **子图编译**: 首先将前述步骤保存在 CinnCompiler 中的子图对象通过 cache key 取出,通过 CinnCompiler::Compile 接口编译子图, 编译成功后获得编译结果结构体[CinnCompiledObject](https://github.com/PaddlePaddle/Paddle/blob/develop/paddle/fluid/framework/paddle2cinn/cinn_compiler.h#L56), 它包括编译过程收集到的信息、必要的上下文环境、以及进行具体计算的可执行对象。关键代码片段如下: |
| 14 | + 子图的编译和计算都在运行期完成,逻辑设计实现在[cinn_launch_op](https://github.com/PaddlePaddle/Paddle/blob/develop/paddle/fluid/pir/dialect/operator/ir/manual_op.h)算子中, 每个 Op 负责一个子图并分别独立运行,从框架执行器视角来看,它与常规算子 kernel 一致,同样都是在运行期调度执行该 kernel, 因此保证了与现有运行逻辑兼容。具体到 kernel 的逻辑上, 它先后执行子图编译和子图计算两阶段: |
| 15 | + - **子图编译**: 首先将前述步骤保存在 CinnCompiler 中的子图对象通过 cache key 取出,通过 CinnCompiler::Compile 接口编译子图, 编译成功后获得编译结果结构体[CinnCompiledObject](https://github.com/PaddlePaddle/Paddle/blob/develop/paddle/cinn/hlir/framework/pir_compiler.h), 它包括编译过程收集到的信息、必要的上下文环境、以及进行具体计算的可执行对象。关键代码片段如下: |
16 | 16 | ``` |
17 | 17 | std::map<std::string, const phi::DenseTensor*> inputs_name2tensor; // 输入变量信息 |
18 | 18 |
|
@@ -55,15 +55,15 @@ Pass 运行结束后计算图形态表现为 cinn_launch_op 之间通过其它 |
55 | 55 | } |
56 | 56 | ``` |
57 | 57 |
|
58 | | -上述数据映射,计算图转换等编译结果上下文的逻辑主要实现在[cinn_launch_context](https://github.com/PaddlePaddle/Paddle/blob/develop/paddle/fluid/operators/cinn/cinn_launch_context.h)类中,它由 CinnCompiler 产出编译结果时进行构造,并作为子图计算时的辅助结构体进行使用,主要功能函数接口见该类接口和实现的说明注释。 |
| 58 | +上述数据映射,计算图转换等编译结果上下文的逻辑主要实现在[cinn_launch_context](https://github.com/PaddlePaddle/Paddle/blob/develop/paddle/cinn/hlir/framework/pir_compiler.h)类中,它由 CinnCompiler 产出编译结果时进行构造,并作为子图计算时的辅助结构体进行使用,主要功能函数接口见该类接口和实现的说明注释。 |
59 | 59 |
|
60 | 60 | ## 显存管理 |
61 | 61 |
|
62 | 62 | 对应于上述两种不同的执行调度方式,我们也分别实现了不同的策略来进行显存管理。 |
63 | 63 |
|
64 | 64 | **Paddle 执行器调度的显存管理策略**: 此种执行方式下首先将可执行序列回转成了 Paddle Graph,为此可以直接应用框架侧已有的显存回收与复用策略,但是需要额外处理的问题是:每个子图包裹在一个 cinn_launch 算子中,算子外还有一个主图,计算图呈现嵌套结构,子图的输入/输出变量可能会被主图中其它算子使用,因此需要将子图外部变量(输入/输出)的使用信息透传至子图内,以避免外部变量在子图计算时被提前释放。 |
65 | 65 |
|
66 | | - 对于这个问题,在使用 PE 调度方式时我们新增了一个[share_varinfo_into_cinn_pass](https://github.com/PaddlePaddle/Paddle/blob/develop/paddle/fluid/framework/ir/memory_optimize_pass/share_varinfo_into_cinn_pass.cc)来实现将外部变量的`MemOptVarInfo`信息同步到子图内中,而新执行器调度方式下则是通过设置`ExecutionConfig::skip_gc_vars`来将子图外部变量的回收延迟至主图计算调度时进行。 |
| 66 | + 对于这个问题,在使用 PE 调度方式时我们新增了一个[share_varinfo_into_cinn_pass](https://github.com/PaddlePaddle/Paddle/blob/develop/paddle/cinn/hlir/framework/pir_compiler.h)来实现将外部变量的`MemOptVarInfo`信息同步到子图内中,而新执行器调度方式下则是通过设置`ExecutionConfig::skip_gc_vars`来将子图外部变量的回收延迟至主图计算调度时进行。 |
67 | 67 |
|
68 | 68 | **CINN runtime 调度的显存管理策略**: 它是依次串行调度可执行序列(CINN instruction),instruction 计算时要求输入/输出数据的内存空间均已分配完成,最粗暴的方式是整个子图计算前将所有数据集中分配,计算完成后再集中释放临时中间数据, 但是这种方式显然会导致显存使用量急剧攀升。为此,我们参考了 Paddle 框架显存回收的思想,通过分析变量的使用关系,在 CINN 生成的子图可执行序列中插入一类特殊的 BufferMalloc/BufferFree 指令来完成即时分配/释放显存的工作。 |
69 | 69 |
|
|
0 commit comments