Python support register pass via PassDesc

[Avin0323]

PR types

New features

PR changes

Others

Describe

PR主要功能：针对fusion等子图替换场景，支持Python侧开发并注册Pass。

背景

• Pass是指输入一个深度学习计算图Graph，依照一定条件进行修改，输出修改后的Graph的过程；
• 当前PaddlePadle框架编写Pass代码存在以下问题：
  • 用户需要手写Graph的条件匹配、在Graph上的修改代码；
  • 对Graph操作需要深入底层框架代码，了解Graph的结构，并且知道相关Pass写法；
• 我们提出了针对fusion等子图替换类Pass的优化方案以支持用户在Python侧开发注册Pass，提升二次开发体验：
  • 用户只需要输入匹配和替换的子图描述，由深度学习框架编写的代码来生成匹配和替换的逻辑，不需要用户对Graph进行匹配和替换操作；
  • API级别的替换，用户可以通过Paddle的Python API构造子图，从而不需要知道Graph的结构，也能写Paddle的Graph Pass代码；

方案设计

子图替换类Pass规则描述

子图替换类型的Pass进行的图变换如下图所示（实线表示数据流，虚线表示变换前后的映射）：

针对图变换场景，描述一个Pass规则主要需要以下4部分内容：

1. 匹配子图：用于匹配原计算图的子图；
2. 替换子图：用于替换匹配子图的子图；
3. 输入/输出映射：指示图变换后数据映射关系；
4. 算子属性映射：指示图变换后算子属性的映射关系；

其中匹配/替换子图也可以看做是一个Graph，框架中Graph由ProgramDesc生成，这里我们也采用ProgramDesc表示子图。

针对如上的规则描述，我们使用protobuf创建PassDesc的数据结构来保存，详细内容见paddle/fluid/framework/pass_desc.proto定义文件。

用于Python侧Pass注册

Python侧支持使用两种方式用于开发Pass，类似定义program的方式来定义匹配/替换子图：

1. 直接使用Paddle Python API；
2. 使用辅助类型；

提供辅助类型开发的原因在于：

• 存在部分算子并未开放用户API，辅助类型用于可以“访问”这部分算子；
• 当替换子图中算子存在属性映射在匹配图中的情况时，辅助类型用于指示这些映射情况；

定义子图通过定义一个函数完成，函数参数为子图的输入数据，返回值为输出数据，一个示例如下：

// 使用Paddle Python API方式
def pattern(x, y1, y2):
        mul1 = paddle.matmul(x, y1)
        mul2 = paddle.matmul(x, y2)
        return mul1, mul2

// 使用辅助类型方式
def pattern(x, y1, y2):
        mul1 = ir.PassDesc.OP.matmul_v2(x, y1)
        mul2 = ir.PassDesc.OP.matmul_v2(x, y2)
        return mul1, mul2

Python侧开发Pass用户接口

Pass注册使用函数装饰器的方式完成，RegisterPass用于修改一个Pass注册函数，将函数名作为Pass类型，用于后续对应该Pass实例。

当定义子图需要指定InputSpec时，通过RegisterPass中添加input_specs参数完成。

Pass注册函数参数为空，返回值为一对或多对匹配、替换子图函数。

注册代码示例如下：

@ir.RegisterPass
def generate_add_n():
    def pattern(x, y, z):
        return paddle.add(paddle.add(x, y), z)

    def replace(x, y, z):
        return paddle.add_n([x, y, z])

    return pattern, replace

[paddle-bot-old]

Thanks for your contribution!
Please wait for the result of CI firstly. See Paddle CI Manual for details.

[zhhsplendid]

LGTM