[Dy2St][PIR] Replace output with inplace source

[SigureMo]

PR Category

Execute Infrastructure

PR Types

Bug fixes

Description

问题描述

SOT/AST+PIR 模型暴露问题，伪码描述如下：

def fn(x):
   mean1 = builtin.parameter("parameter_1")
   y, mean2 = pd_op.batch_norm_(x, mean1, ...)
   return y

对于这个 Program 来说，动转静前向输入包含原输入 x 和 parameter mean1，前向输出包含原输出 y 和反向所需的中间变量 mean2，问题就出在这个 mean1 和 mean2 是 inplace 的（实际上是 view 的，但 PIR 下是一样的）

由于 mean1 和 mean2 是 inplace 的，所以他们对应同一个 Variable，mean2 作为输出我们会添加 shadow_output 来确保能够取出来，但由于 mean2 实际上就是 mean1，它已经有作为输入的 parameter_1 这个名字了，实际上是不需要通过 shadow_output 来重命名的。而且，如果进行重命名，在第一个 step 我们在 share 到 scope 时需要使用 parameter_1 这个名字来 share，而第二个 step 因为已经 rename 成另一个名字了，这就会导致 share 失败。

解决方案

归根结底，问题出在一个输入 Value 经过若干 inplace 操作得到一个输出 Value，此时这个输出 Value 不应该添加 shadow_output，而是应该直接替换为输入 Value。为了达成这一点，我们需要在前反向拆分时，对 Program 中所有 inplace 链进行分析，比如下图：

这里仅标识出了 inplace Value，忽略非 inplace Value，并且按照时序拉直，比如这里的链路 A 中的 A1 作为输入，A4 作为输出，那么 A4 不应该添加 shadow_output。

具体替换结果如下图所示：

但是这里为什么不直接将 A4 从中间变量中删掉呢？因为考虑到动转静前向的输出侧并不仅仅包含中间变量，还包含原前向的输出，针对中间变量尚且可以删掉，但如果针对输出的话就不太好删掉了，因此我们对中间变量和输出变量采用相同的操作，就是替换为 inplace 的输入侧源 Value。

如果考虑前反向全部情况，根据下图动转静前反向的输入输出关系，实际上我们需要处理如下几种情况：

前向：

x | param -> outputs
x | param -> middles

反向：

x | param | middles | outputs | out_grads -> x_grads
x | param | middles | outputs | out_grads -> param_grads

但是由于前向输入 x 是用 data op 创建的，data op 在经过 lower 到 kernel pass 后会跟随一个 shadow_feed op，而 shadow_feed op 不是 inplace 的，会创建一个 inner var，也就是说前向输入 x 永远不会出现上述的同时作为输入和输出的问题，因为在一开始就有一个非 inplace 的操作，因此我们不需要处理前向输入，只需要处理 parameter。

类似地，反向全部输入都是通过 block kwarg 创建的，而 kwarg 也会在 lower 后跟随一个 shadow_feed，因此整个反向都不会出现上述问题，也就是说反向不需要处理。

因此需要处理的只有：

前向：

param -> outputs
param -> middles

但由于最开始是考虑这种情况的，即便将来 data/kwarg 后不插 shadow_feed 也能很容易实现。（因为目前看反向会有过多 shadow_feed，非常影响性能，也许之后这个 shadow_feed 会被删掉也说不定）

Note

虽然问题清楚了，但仍然没能找到一个简单的复现样例，因此没有添加单测，关于 inplace 的单测在 #62300 已经添加过了，但不能复现，本 PR 在模型上验证通过

[paddle-bot]

你的PR提交成功，感谢你对开源项目的贡献!
请关注后续CI自动化测试结果，详情请参考Paddle-CI手册。
Your PR has been submitted. Thanks for your contribution!
Please wait for the result of CI firstly. See Paddle CI Manual for details.

[2742195759]

inplace 是一个等价关系，可以直接上并查集，只要永远将右边作为root即可。但是inplace大概率是一个链，所以直接搜索似乎关系也不大。

[SigureMo]

在 Python 端已经有并查集，这里是为了简单且易于传递到 Python 端，这样也不用重新定义一个数据结构了，直接一个 list[list[Value]] 就可以传过去了，Python 端也可以复用，当然这个 PR 暂时没做

[2742195759]

LGTM