[PIR]support while op translate

[winter-wang]

PR types

New features

PR changes

Others

Description

本 PR 用于支持将旧 Program 体系的 while算子 translate 为 Pir 的 WhileOp；
示例模型为一个 while 嵌套 while 的模型结构，模型代码如下：

import paddle
paddle.enable_static()

def cond(i, ten):
    return i < ten

def body(i, ten):
    i = i + 1
    i, ten = paddle.static.nn.while_loop(
        cond, lambda i, ten: [i + 3, ten], [i, ten]
    )
    return [i, ten]

main_program = paddle.static.default_main_program()
startup_program = paddle.static.default_startup_program()
with paddle.static.program_guard(main_program, startup_program):
    i = paddle.full(shape=[1], fill_value=0, dtype='int64')  # loop counter
    ten = paddle.full(shape=[1], fill_value=10, dtype='int64')  # loop length
    i, ten = paddle.static.nn.while_loop(cond, body, [i, ten])
with open('while_op_test.prog', 'wb') as f:
    f.write(main_program.desc.serialize_to_string())

旧 Program：

{ // block 0
    var fill_constant_1.tmp_0 : LOD_TENSOR.shape(1,).dtype(int64).stop_gradient(True)
    var fill_constant_3.tmp_0 : LOD_TENSOR.shape(1,).dtype(int64).stop_gradient(True)
    var tmp_0 : LOD_TENSOR.shape(1,).dtype(bool).stop_gradient(True)
    var _generated_var_1 : STEP_SCOPES)

    {Out=['fill_constant_1.tmp_0']} = fill_constant(inputs={ShapeTensor=[], ShapeTensorList=[], ValueTensor=[]}, dtype = 3, force_cpu = False, op_device = , op_namescope = /, op_role = 0, op_role_var = [], place_type = -1, shape = [1], str_value = 0, value = 0.0, with_quant_attr = False)
    {Out=['fill_constant_3.tmp_0']} = fill_constant(inputs={ShapeTensor=[], ShapeTensorList=[], ValueTensor=[]}, dtype = 3, force_cpu = False, op_device = , op_namescope = /, op_role = 0, op_role_var = [], place_type = -1, shape = [1], str_value = 10, value = 10.0, with_quant_attr = False)
    {Out=['tmp_0']} = less_than(inputs={X=['fill_constant_1.tmp_0'], Y=['fill_constant_3.tmp_0']}, axis = -1, force_cpu = False, op_device = , op_namescope = /, op_role = 0, op_role_var = [], with_quant_attr = False)
    {Out=['tmp_0', 'fill_constant_1.tmp_0'], StepScopes=['_generated_var_1']} = while(inputs={Condition=['tmp_0'], X=['fill_constant_3.tmp_0', 'fill_constant_1.tmp_0']}, is_test = False, op_device = , op_namescope = /, op_role = 0, op_role_var = [], sub_block = block[1], with_quant_attr = False)
}
{ // block 1
    var tmp_1 : LOD_TENSOR.shape(1,).dtype(int64).stop_gradient(True)
    var tmp_2 : LOD_TENSOR.shape(1,).dtype(bool).stop_gradient(True)
    var _generated_var_0 : STEP_SCOPES)
    var tmp_5 : LOD_TENSOR.shape(1,).dtype(bool).stop_gradient(True)

    {Out=['tmp_1']} = scale(inputs={ScaleTensor=[], X=['fill_constant_1.tmp_0']}, bias = 1.0, bias_after_scale = True, op_device = , op_namescope = /, op_role = 0, op_role_var = [], scale = 1.0, with_quant_attr = False)
    {Out=['tmp_2']} = less_than(inputs={X=['tmp_1'], Y=['fill_constant_3.tmp_0']}, axis = -1, force_cpu = False, op_device = , op_namescope = /, op_role = 0, op_role_var = [], with_quant_attr = False)
    {Out=['tmp_1', 'tmp_2'], StepScopes=['_generated_var_0']} = while(inputs={Condition=['tmp_2'], X=['fill_constant_3.tmp_0', 'tmp_1']}, is_test = False, op_device = , op_namescope = /, op_role = 0, op_role_var = [], sub_block = block[2], with_quant_attr = False)
    {Out=['tmp_5']} = less_than(inputs={X=['tmp_1'], Y=['fill_constant_3.tmp_0']}, axis = -1, force_cpu = False, op_device = , op_namescope = /, op_role = 0, op_role_var = [], with_quant_attr = False)
    {Out=['fill_constant_1.tmp_0']} = assign(inputs={X=['tmp_1']}, op_device = , op_namescope = /, op_role = 0, op_role_var = [], with_quant_attr = False)
    {Out=['tmp_0']} = assign(inputs={X=['tmp_5']}, op_device = , op_namescope = /, op_role = 0, op_role_var = [], with_quant_attr = False)
}
{ // block 2
    var tmp_3 : LOD_TENSOR.shape(1,).dtype(int64).stop_gradient(True)
    var tmp_4 : LOD_TENSOR.shape(1,).dtype(bool).stop_gradient(True)

    {Out=['tmp_3']} = scale(inputs={ScaleTensor=[], X=['tmp_1']}, bias = 3.0, bias_after_scale = True, op_device = , op_namescope = /, op_role = 0, op_role_var = [], scale = 1.0, with_quant_attr = False)
    {Out=['tmp_4']} = less_than(inputs={X=['tmp_3'], Y=['fill_constant_3.tmp_0']}, axis = -1, force_cpu = False, op_device = , op_namescope = /, op_role = 0, op_role_var = [], with_quant_attr = False)
    {Out=['tmp_1']} = assign(inputs={X=['tmp_3']}, op_device = , op_namescope = /, op_role = 0, op_role_var = [], with_quant_attr = False)
    {Out=['tmp_2']} = assign(inputs={X=['tmp_4']}, op_device = , op_namescope = /, op_role = 0, op_role_var = [], with_quant_attr = False)
}

翻译之后的 Program为：

{
  (%0) = "pd_op.full" () {dtype:(pd_op.DataType)int64,place:(pd_op.Place)Place(undefined:0),shape:(pd_op.IntArray)[1],value:(Float)0} : () -> pd_op.tensor<1xi64>
  (%1) = "pd_op.full" () {dtype:(pd_op.DataType)int64,is_persisable:[false],place:(pd_op.Place)Place(undefined:0),shape:(pd_op.IntArray)[1],stop_gradient:[true],value:(Float)10} : () -> pd_op.tensor<1xi64>
  (%2) = "pd_op.less_than" (%0, %1) {is_persisable:[false],stop_gradient:[true]} : (pd_op.tensor<1xi64>, pd_op.tensor<1xi64>) -> pd_op.tensor<1xb>
  (%3) = "pd_op.while"(%2) [%0] { 
     ^body(%arg0):
        (%4) = "pd_op.full" () {dtype:(pd_op.DataType)float32,place:(pd_op.Place)Place(cpu),shape:(pd_op.IntArray)[1],stop_gradient:[true],value:(Float)1} : () -> pd_op.tensor<1xf32>
        (%5) = "pd_op.scale" (%arg0, %4) {bias:(Float)1,bias_after_scale:true} : (pd_op.tensor<1xi64>, pd_op.tensor<1xf32>) -> pd_op.tensor<1xi64>
        (%6) = "pd_op.less_than" (%5, %1) {is_persisable:[false],stop_gradient:[true]} : (pd_op.tensor<1xi64>, pd_op.tensor<1xi64>) -> pd_op.tensor<1xb>
        (%7) = "pd_op.while"(%6) [%5] { 
          ^body(%arg1):
              (%8) = "pd_op.full" () {dtype:(pd_op.DataType)float32,place:(pd_op.Place)Place(cpu),shape:(pd_op.IntArray)[1],stop_gradient:[true],value:(Float)1} : () -> pd_op.tensor<1xf32>
              (%9) = "pd_op.scale" (%arg1, %8) {bias:(Float)3,bias_after_scale:true,is_persisable:[false],stop_gradient:[true]} : (pd_op.tensor<1xi64>, pd_op.tensor<1xf32>) -> pd_op.tensor<1xi64>
              (%10)= "pd_op.less_than" (%9, %1) {is_persisable:[false],stop_gradient:[true]} : (pd_op.tensor<1xi64>, pd_op.tensor<1xi64>) -> pd_op.tensor<1xb>
              () = "cf.yield" (%10, %9) {} : (pd_op.tensor<1xb>, pd_op.tensor<1xi64>) -> 
          }
        (%11) = "pd_op.less_than" (%7, %1) {is_persisable:[false],stop_gradient:[true]} : (pd_op.tensor<1xi64>, pd_op.tensor<1xi64>) -> pd_op.tensor<1xb>
        () = "cf.yield" (%11, %7) {} : (pd_op.tensor<1xb>, pd_op.tensor<1xi64>) -> 
  }
}

Other

Pcard-67164

[Aurelius84]

LGTM

[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.