【护航实习计划】SOT 支持动态 Shape 功能拓展

[diadestiny]

问题描述 Please describe your issue

背景

SOT将来是Paddle动转静的默认状态，但是目前在出现动态Shape的网络中会有重复转写的性能问题。因此 SOT 支持动态Shape是一个重要的功能。

目前的SOT面对动态Shape的网络处理方式如下：

1. SOT 会对每个函数的PyCode记录一个转写次数的结构。
2. 当一个函数需要转写，并且已经转写的次数 > 20，那么说明这个函数大概率是一个动态Shape，因此会直接 fallback 到动态图下执行。

通过上述的fallback机制，我们会识别出那些动态Shape的函数，并放弃他们的静态化。牺牲静态化换取转写时间。
显然上面的方法不是一个非常高效的形式，我们希望做到的最终状态应该是：

def func(x):
    return x * 2 

x1 = paddle.randn((2,2))
x2 = paddle.randn((2,4))
x3 = paddle.randn((2,5))

func(x1) # 针对 x1.shape=(2,2) 进行转写。
func(x2) # 针对 x2.shape和x1的shape，我们直接转写一个 (2, -1) 的 Program
func(x3) # 命中第二次转写的 Program，直接复用Program即可

这也是动态Shape的基本思想。

方案文档【待楷浩完善】

为了解决上述的问题，我设计了一个方案来达到背景中描述的理想状态。目前这个任务将由护航实习生 @diadestiny
主要方案思想分为下面几个大模块：

1. 动态Shape的检测模块：我们需要在 guard 机制中进行历史shape的记录和动态Shape的推测。
2. SymbolicVariable 继承体系：目前Variable中只有TensorVariable是会被作为Symoblic结构记录在SIR中，当我们引入了动态Shape之后，我们需要支持常量的SymbolicTrace。因此我们需要抽象出这个具有Symbolic行为的Variable类别，使用类似的方法处理SymbolicConstantVariable和TensorVariable。
3. 修改OpCodeExecutor类的模拟执行行为，针对SymbolicVariable体系下的对象进行SIR组网，对于Guard而言，我们只是直接使用它的ConstantValue即可。

其实上述方案的核心是：我们会有一个新的变量类别，在组网的时候，使用SIR中存储的Symbolic变量，但是在模拟的时候，还是使用Constant来计算Guard等信息。

项目规划和MileStone

这里将进行项目拆解：

MileStone1：熟悉SOT流程，具备基本的SOT debug能力。
成果交付：

1. 独立修复SOT框架单测的bug3个。
2. 在动转静内部做一个SOT流程分享，产出SOT流程文档一份，需要重点关注的模块如下：
   1. SOT的Guard机制
   2. SOT中的OpCodeExecutor的实现方法
   3. Variable体系
   4. SOT前端到pir::Program的构件流程
   5. 静态图下的动态Shape机制（新IR下）

MileStone2：掌握SOT的动态Shape方案，并完善细节。
成果交付：

1. 产出一个完整的SOT动态Shape文档

MileStone3：按照方案文档开发，在特定demo组网上跑通动态Shape

MileStone4：完善细节，在Bert上跑通动态Shape方案

年前任务

MileStone1

[2742195759]

#61486

[2742195759]

下载最新的Paddle版本，然后运行下面的代码就是SOT模式的动转静了，可以按照LOG来追溯一下整个SOT的流程。

import os
os.environ['MIN_GRAPH_SIZE']='0'
os.environ['FLAGS_enable_pir_api']='True'
os.environ['SOT_LOG_LEVEL']='3' # open this to access the sot logs

import paddle

@paddle.jit.to_static
def fn(x):
    tmp = 2 * x 
    tmp = paddle.nn.functional.relu(tmp)
    return 3 *  tmp

input_ = paddle.ones([100], dtype="float32")
input_.stop_gradient = False
output = fn(input_)
#print (fn.get_concrete_program(input_)[0].main_program)