# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from collections import defaultdict

import numpy as np

# (TODO: GhostScreaming) It will be removed later.
from paddle.base import core
from paddle.distributed import fleet
from paddle.framework import Block, Program, in_dynamic_mode


class HybridParallelInferenceHelper:
    """
    A helper class to split program for inference with hybrid parallelism.

    Args:
        startup_program (Program): the startup program.
        main_program (Program): the main program.
        num_mp (int): number of model parallel degree. Default ``1``.
        num_pp (int): number of pipeline parallel degree. Default ``1``.
        micro_batch_size (int): number of micro batch size. Default ``1``.
        beam_size (int): number of beam search size. Default ``1``.
        init_comm (bool): whether if initialize communication group. Default ``True``.
        role_maker (RoleMakerBase or subclass): user custom define RoleMakerBase.
            If ``role_maker==None``, then use PaddleCloudRoleMaker. Default ``None``.

    Returns:
        None.

    Write Paradigm:

        .. code-block:: text
            :name: text-example1

            >>> # doctest: +REQUIRES(env:DISTRIBUTED, env:GPU)
            >>> import paddle
            >>> # while op pattern
            >>> with paddle.base.device_guard(f'{device}:all'):
            ...     # init global cond
            ...     max_len = paddle.full(shape=[1], dtype="int64", fill_value=10)
            ...     step_idx = paddle.full(shape=[1], dtype="int64", fill_value=0)
            ...     cond_int = paddle.full(shape=[1], dtype="int64", fill_value=0, name="cond_int")
            ...     cond = layers.cast(step_idx < max_len, dtype="bool")
            ...     while_op = layers.While(cond, is_test=True)

            ...     # init global lod_tensor_array for generation task
            ...     arr = paddle.tensor.array_write(data, step_idx)

            >>> with while_op.block():
            ...     with paddle.base.device_guard(f'{device}:all'):
            ...         # read data from global lod_tensor_array
            ...         element_in_arr = paddle.tensor.array_read(array=arr, i=step_idx)
            ...         # write placeholder data to global lod_tensor_array,
            ...         # it need for send_v2 of lod_tensor_array
            ...         paddle.increment(x=step_idx, value=1.0)
            ...         paddle.tensor.array_write(element_in_arr, i=step_idx, array=arr)
            ...     with paddle.base.device_guard(f'{device}:0'):
            ...         pass  # some code
            ...     with paddle.base.device_guard(f'{device}:1'):
            ...         pass  # some code
            ...     with paddle.base.device_guard(f'{device}:{num_pp - 1}'):
            ...         # generate some data in while block and write to global lod_tensor_array
            ...         # that they are read in next while step.
            ...         # we will using send_v2 to send global lod_tensor_array to other pipeline and sync
            ...         paddle.tensor.array_write(other_var, i=step_idx, array=arr)
            ...         # update cond and assign to cond_int, we will sync cond_int
            ...         layers.assign(layers.cast(cond, dtype="int32"), cond_int)
            ...     with paddle.base.device_guard(f'{model._device}:all'):
            ...         # the code below must at end of while block and exists in device:all
            ...         layers.assign(layers.cast(cond_int, dtype='bool'), cond)

            >>> with paddle.base.device_guard(f'{model._device}:all'):
            ...     # use a empty lod_tensor_array to clear lod_tensor_array
            ...     layers.assign(layers.create_array(data.dtype), arr)

    Examples:

        .. code-block:: pycon
            :name: code-example1

            >>> # doctest: +REQUIRES(env:DISTRIBUTED, env:GPU)
            >>> import os
            >>> import numpy as np
            >>> import paddle
            >>> import paddle.distributed.fleet as fleet
            >>> from paddle.distributed.fleet.utils import hybrid_parallel_inference
            >>> paddle.enable_static()
            >>> nranks = int(os.getenv("PADDLE_TRAINERS_NUM", 1))
            >>> rank = int(os.getenv("PADDLE_TRAINER_ID", 0))
            >>> dev_id = int(os.getenv("FLAGS_selected_gpus", 0))
            >>> main_program = paddle.static.Program()
            >>> startup_program = paddle.static.Program()
            >>> if nranks > 1:
            ...     dist_strategy = fleet.DistributedStrategy()
            ...     dist_strategy.without_graph_optimization = True
            ...     fleet.init(is_collective=True, strategy=dist_strategy)
            >>> device = "gpu"
            >>> with paddle.static.program_guard(main_program, startup_program):
            ...     with paddle.base.device_guard(f'{device}:0'):
            ...         X = paddle.static.data(name='X', shape=[None, 2], dtype='float32')
            ...     with paddle.base.device_guard(f'{device}:all'):
            ...         max_len = paddle.full(shape=[1], dtype="int64", fill_value=5, name="n")
            ...         step_idx = paddle.full(shape=[1], dtype="int64", fill_value=0, name="i")
            ...         data = paddle.tensor.array_write(X, step_idx)
            ...         cond_int = paddle.full(shape=[1], dtype="int64", fill_value=0, name="cond_int")
            ...         cond = paddle.less_than(x=step_idx, y=max_len)
            ...         while_op = paddle.static.nn.control_flow.While(cond, is_test=True)
            ...     with while_op.block():
            ...         with paddle.base.device_guard(f'{device}:all'):
            ...             input = paddle.tensor.array_read(array=data, i=step_idx)
            ...             paddle.increment(x=step_idx, value=1.0)
            ...             paddle.tensor.array_write(input, i=step_idx, array=data)
            ...         with paddle.base.device_guard(f'{device}:0'):
            ...             param_attr = paddle.ParamAttr(initializer=paddle.nn.initializer.Constant(1.0))
            ...             weight1 = paddle.static.create_parameter(shape=[2, 5], dtype='float32', attr=param_attr, is_bias=False)
            ...             hidden1 = paddle.matmul(input, weight1)
            ...         with paddle.base.device_guard(f'{device}:1'):
            ...             param_attr = paddle.ParamAttr(initializer=paddle.nn.initializer.Constant(2.0))
            ...             weight2 = paddle.static.create_parameter(shape=[5, 2], dtype='float32', attr=param_attr, is_bias=False)
            ...             hidden2 = paddle.matmul(hidden1, weight2)
            ...             paddle.tensor.array_write(hidden2, i=step_idx, array=data)
            ...             # update cond and assign to cond_int, we will sync cond_int
            ...             paddle.assign(paddle.less_than(x=step_idx, y=max_len), cond)
            ...             paddle.assign(paddle.cast(cond, dtype="int32"), cond_int)
            ...         with paddle.base.device_guard(f'{device}:all'):
            ...             # the code below must at end of while block and exists in device:all
            ...             paddle.assign(paddle.cast(cond_int, dtype='bool'), cond)
            ...     with paddle.base.device_guard(f'{device}:all'):
            ...         out = paddle.tensor.create_array(data.dtype)
            ...         paddle.assign(data, out)
            ...     with paddle.base.device_guard(f'{device}:all'):
            ...         # use a empty lod_tensor_array to clear lod_tensor_array
            ...         paddle.assign(paddle.tensor.create_array(data.dtype), data)
            >>> helper = hybrid_parallel_inference.HybridParallelInferenceHelper(
            ...     startup_program,
            ...     main_program,
            ...     micro_batch_size=2,
            ...     num_pp=2,
            ...     init_comm=nranks > 1,
            ... )
            >>> helper.gen_infer_program(['array_write_0.out'], ['cond_int.tmp_0'])
            >>> exe = paddle.static.Executor(paddle.CUDAPlace(dev_id))
            >>> exe.run(startup_program)
            >>> np.random.seed(2333)
            >>> for step in range(5):
            ...     init_data = np.random.uniform(low=0.0, high=1.0, size=[2, 2]).astype('float32')
            ...     [res] = exe.run(main_program, feed={"X": init_data}, fetch_list=[out])
            ...     print('-------- step', step, ' --------')
            ...     print(res)

    """

    def __init__(
        self,
        startup_program,
        main_program,
        num_mp=1,
        num_pp=1,
        micro_batch_size=1,
        beam_size=1,
        init_comm=True,
        role_maker=None,
    ):
        assert isinstance(startup_program, Program)
        assert isinstance(main_program, Program)

        self._device = None
        if core.is_compiled_with_cuda():
            self._device = "gpu"
        assert self._device, "Only gpu are supported."

        assert not in_dynamic_mode(), "Only static graph mode is supported."

        op_maker = core.op_proto_and_checker_maker
        self._op_role = op_maker.OpRole
        self._op_role_key = op_maker.kOpRoleAttrName()
        self._op_device_key = op_maker.kOpDeviceAttrName()

        self._param_device_map = {}

        self._pipeline_pair = []
        self._pipeline_pair_in_while = []
        self._pp_ring_map = {}
        self.ring_id = 20  # Just a magic number

        self.micro_batch_size = micro_batch_size
        self.beam_size = beam_size
        self.init_comm = init_comm

        self._output_var_to_op = None
        self._input_var_to_op = None
        self._main_program = main_program
        self._startup_program = startup_program

        if role_maker is None:
            self.role_maker = fleet.base.role_maker.PaddleCloudRoleMaker(
                is_collective=True
            )
        else:
            if isinstance(role_maker, fleet.base.role_maker.RoleMakerBase):
                assert role_maker._is_collective
                self.role_maker = role_maker

        # communication_group info
        self.mp_ring_id = 0
        self.global_ring_id = 1

        self.endpoints = self.role_maker._get_trainer_endpoints()
        self.current_endpoint = self.endpoints[self.role_maker._worker_index()]
        self.rank = self.role_maker._worker_index()
        self.nranks = self.role_maker._worker_num()
        assert num_mp * num_pp == self.nranks
        self.num_pp = num_pp
        self.num_mp = num_mp

        # global ring info
        self.global_endpoints = self.endpoints
        self.global_rank = self.rank
        self.global_nranks = self.nranks

        arr = np.arange(0, self.num_pp * self.num_mp).reshape(
            [self.num_pp, self.num_mp]
        )
        ipp, imp = np.where(arr == self.rank)
        ipp = ipp[0]
        imp = imp[0]
        self.mp_group = arr[ipp, :]
        self.pp_group = arr[:, imp]

        self._stage = ipp

    def _init_communication_group(self):
        dev_ids = []
        for pair in self._pipeline_pair:
            prev_id, cur_id = pair
            if prev_id not in dev_ids:
                dev_ids.append(prev_id)
            if cur_id not in dev_ids:
                dev_ids.append(cur_id)
        num_pp = len(dev_ids)
        num_pp = max(1, num_pp)
        assert num_pp == self.num_pp, (
            f'num_pp: {num_pp}, self.num_pp: {self.num_pp}'
        )

        collective_helper = fleet.meta_optimizers.common.CollectiveHelper(
            self.role_maker, wait_port=False
        )

        # Create global rings
        collective_helper._init_communicator(
            self._startup_program,
            self.current_endpoint,
            self.global_endpoints,
            self.global_rank,
            self.global_ring_id,
            True,
            self.global_ring_id,
            True,
        )

        # Create mp rings
        if self.num_mp > 1:
            mp_endpoints = [self.endpoints[mp_idx] for mp_idx in self.mp_group]
            mp_rank = next(
                idx
                for idx, mp_idx in enumerate(self.mp_group)
                if mp_idx == self.rank
            )
            collective_helper._init_communicator(
                self._startup_program,
                self.current_endpoint,
                mp_endpoints,
                mp_rank,
                self.mp_ring_id,
                True,
                self.global_ring_id,
                True,
            )

        # Create pipeline rings
        if self.num_pp > 1:
            for pair in self._pipeline_pair:
                pair_key = pair[0] * 1000 + pair[1]
                ring_id = self._pp_ring_map[pair_key]

                first_node = self.pp_group[pair[0]]
                second_node = self.pp_group[pair[1]]
                if self.rank != first_node and self.rank != second_node:
                    collective_helper._init_communicator(
                        self._startup_program,
                        None,
                        None,
                        None,
                        None,
                        False,
                        self.global_ring_id,
                        True,
                    )
                    continue

                pipeline_endpoints = [
                    self.endpoints[first_node],
                    self.endpoints[second_node],
                ]
                pipeline_rank = 0 if self.rank == first_node else 1
                collective_helper._init_communicator(
                    self._startup_program,
                    self.current_endpoint,
                    pipeline_endpoints,
                    pipeline_rank,
                    ring_id,
                    False,
                    self.global_ring_id,
                    True,
                )

    def _get_input_output_info(self, block):
        '''
        Get info of op input and output.
        '''
        # A map from output var to op which generate it.
        output_var_to_op = defaultdict(list)
        # A map from var to op which takes it as input.
        input_var_to_op = defaultdict(list)

        for index, op in enumerate(block.ops):
            for var_name in op.input_arg_names:
                input_var_to_op[var_name].append([op, index])
            for var_name in op.output_arg_names:
                output_var_to_op[var_name].append([op, index])

        return output_var_to_op, input_var_to_op

    def _update_param_device_map(self):
        """
        Get the device info for parameters.
        """
        params = [param.name for param in self._main_program.all_parameters()]
        for each_block in self._main_program.blocks:
            for op in each_block.ops:
                for var_name in op.input_arg_names:
                    if (
                        var_name not in params
                        or var_name in self._param_device_map
                    ):
                        continue
                    device = op.attr(self._op_device_key)

                    self._param_device_map[var_name] = device

    def _split_program(self, program, stage, block_idx):
        """
        Split a program and get the one with the given pipeline stage.

        Args:
            stage (int): pipeline stage
            block_idx (int): block index

        Returns:
            used_var_names (set): used var names in block_idx block
        """

        used_var_names = set()
        block = program.block(block_idx)
        op_idx = 0
        for op in list(block.ops):
            op_stage = op.attr(self._op_device_key).split(':')[1]
            # Copy ops whose op_device set to "gpu:all" to all sections.
            if op_stage == "all" or int(op_stage) == stage:
                op_idx += 1
                if op.type == "while":
                    sub_block_id = int(op.attr('sub_block').id)
                    sub_used_var_names = self._split_program(
                        program, stage, sub_block_id
                    )

                    used_var_names.update(sub_used_var_names)

                    input_idxs = []
                    input_arg_names = op.input("X")
                    for i, name in enumerate(input_arg_names):
                        if name not in sub_used_var_names:
                            input_idxs.append(i)
                    if len(input_idxs) > 0:
                        for i in reversed(input_idxs):
                            input_arg_names.pop(i)
                        op.desc.set_input("X", input_arg_names)

                    output_idxs = []
                    output_arg_names = op.output("Out")
                    for i, name in enumerate(output_arg_names):
                        if name not in sub_used_var_names:
                            output_idxs.append(i)
                    if len(output_idxs) > 0:
                        for i in reversed(output_idxs):
                            output_arg_names.pop(i)
                        op.desc.set_output("Out", output_arg_names)

                for var_name in op.input_arg_names + op.output_arg_names:
                    used_var_names.add(var_name)
            else:
                block._remove_op(op_idx)

        for var_name in list(block.vars.keys()):
            if var_name not in used_var_names:
                block._remove_var(var_name)

        return used_var_names

    #     def _find_post_op(self, index, var_name):
    #         """
    #         Find the post op that has variable named var_name as input.
    #         """
    #         # bugfix for uniform hybrid parallelism
    #         if '.cast_fp32' in var_name:
    #             var_name = var_name.replace('.cast_fp32', '')
    #         if '.cast_fp16' in var_name:
    #             var_name = var_name.replace('.cast_fp16', '')

    #         post_ops = self._input_var_to_op[var_name]
    #         if post_ops == None: return None
    #         result_op = None
    #         for post_op, post_idx in reversed(post_ops):
    #             if post_idx > index:
    #                 result_op = post_op
    #                 break
    #         return result_op

    def _find_prev_op(self, index, var_name):
        """
        Find the previous op of op with index that outputs
        variable named var_name.
        """
        prev_ops = self._output_var_to_op[var_name]
        if prev_ops is None:
            return None
        result_op = None
        for prev_op, prev_idx in reversed(prev_ops):
            if prev_idx < index:
                result_op = prev_op
                break
        return result_op

    def _add_op_device_attr(self, block):
        """
        Add op_device attribute for ops in block that have
        not that attribute set.

        Args:
            block (Block): the block to process.
        """
        assert isinstance(block, Block)

        # Ops should be copied to all pipeline stages.
        device_all_ops = [
            "create_py_reader",
            "read",
            "create_double_buffer_reader",
            "while",
        ]

        for op in block.ops:
            if op.type in device_all_ops:
                # We use "gpu:all" to represent an op should be put on all
                # pipeline stages, such as read ops. Note that: "gpu:all"
                # is only used by pipeline as an indicator.
                op._set_attr(self._op_device_key, self._device + ":all")
            if op.type == "while":
                sub_block_id = op.attr('sub_block').id
                sub_block = block.program.block(sub_block_id)
                self._add_op_device_attr(sub_block)

    def _check_validation(self, block):
        """
        Check whether ops in a block have both the op_device and the
        op_role attributes set.
        """
        assert isinstance(block, Block)

        pre_stage_id = None
        for op in block.ops:
            assert op.has_attr(self._op_role_key), (
                f"{op.type} has no {self._op_role_key} set ."
            )
            op_role = op.attr(self._op_role_key)
            assert op_role == int(self._op_role.Forward), (
                "Only forward is supported for inference."
            )
            if not op._has_kernel(op.type):
                assert op.type in [
                    "while",
                    "conditional_block",
                ], "The only supported op without kernel is while."
                sub_block_id = op.attr('sub_block').id
                sub_block = block.program.block(sub_block_id)
                self._check_validation(sub_block)
            assert op.has_attr(self._op_device_key), (
                f"{op.type} has no {self._op_device_key} set."
            )

            device = op.attr(self._op_device_key)
            assert device, f"{op.type} has no {self._op_device_key} set."
            if device.split(':')[1] == "all":
                continue

            dev_type = device.split(':')[0]
            assert dev_type == self._device
            stage_id = int(device.split(':')[1])
            pre_stage_id = stage_id

    def _insert_sendrecv_ops_for_boundaries(self, block, is_while_block):
        """
        Insert a pair of send and recv ops for every two
        consecutive ops on different devices.
        """
        # A map from var to device where op takes it as input,
        # avoiding multiple send and recv ops.
        input_var_to_device = {}

        extra_index_info = {
            'index': 0,
        }

        for index, op in enumerate(list(block.ops)):
            cur_device = op.attr(self._op_device_key)
            if cur_device.split(':')[-1] == "all":
                continue
            for var_name in op.input_arg_names:
                if not block.has_var(var_name) and block._find_var_recursive(
                    var_name
                ):
                    continue
                var = block.var(var_name)
                # skip data var
                if var.is_data:
                    continue
                prev_device = None
                generate_ops = self._output_var_to_op.get(var_name)
                if generate_ops is None:
                    if var_name not in self._param_device_map:
                        continue
                    prev_device = self._param_device_map[var_name]

                prev_op = self._find_prev_op(index, var_name)

                if not prev_device:
                    prev_device = (
                        prev_op.attr(self._op_device_key) if prev_op else None
                    )

                if prev_device is None or prev_device.split(":")[-1] == "all":
                    continue

                if prev_device == cur_device:
                    continue

                if var_name not in input_var_to_device:
                    input_var_to_device[var_name] = []
                if (cur_device, prev_device) in input_var_to_device[var_name]:
                    continue

                assert self._device == cur_device.split(':')[0], (
                    "More than one device type found."
                )
                device_type = cur_device.split(':')[0] + ':'

                def _insert_send_recv(cur_id, prev_id):
                    assert cur_id > prev_id
                    cur_dev = device_type + str(cur_id)
                    prev_dev = device_type + str(prev_id)
                    if (cur_dev, prev_dev) in input_var_to_device[var_name]:
                        return

                    if cur_id - prev_id > 1:
                        _insert_send_recv(cur_id - 1, prev_id)
                        _insert_send_recv(cur_id, cur_id - 1)
                        input_var_to_device[var_name].append(
                            (cur_dev, prev_dev)
                        )
                        return

                    assert cur_id - prev_id == 1
                    input_var_to_device[var_name].append((cur_dev, prev_dev))

                    op_role = op.attr(self._op_role_key)
                    var = block.vars[var_name]
                    pair = (prev_id, cur_id)
                    if (
                        is_while_block
                        and pair not in self._pipeline_pair_in_while
                    ):
                        self._pipeline_pair_in_while.append(pair)

                    # 1000 is just a magic number
                    pair_key = prev_id * 1000 + cur_id
                    if pair not in self._pipeline_pair:
                        self._pipeline_pair.append(pair)
                        self._pp_ring_map[pair_key] = self.ring_id
                        ring_id = self.ring_id
                        self.ring_id += 1
                    else:
                        ring_id = self._pp_ring_map[pair_key]

                    block._insert_op_without_sync(
                        index=index + extra_index_info['index'],
                        type='send_v2',
                        inputs={'X': var},
                        attrs={
                            self._op_device_key: prev_dev,
                            self._op_role_key: op_role,
                            'use_calc_stream': True,
                            'peer': 1,
                            'ring_id': ring_id,
                        },
                    )
                    extra_index_info['index'] += 1
                    var_shape = list(var.shape)
                    if var_shape[0] < 0:
                        if is_while_block:
                            var_shape[0] = (
                                self.micro_batch_size * self.beam_size
                            )
                        else:
                            var_shape[0] = self.micro_batch_size

                    block._insert_op_without_sync(
                        index=index + extra_index_info['index'],
                        type='recv_v2',
                        outputs={'Out': [var]},
                        attrs={
                            'out_shape': var_shape,
                            'dtype': var.dtype,
                            self._op_device_key: cur_dev,
                            self._op_role_key: op_role,
                            'use_calc_stream': True,
                            'peer': 0,
                            'ring_id': ring_id,
                        },
                    )
                    extra_index_info['index'] += 1

                _insert_send_recv(
                    int(cur_device.split(':')[1]),
                    int(prev_device.split(':')[1]),
                )
        block._sync_with_cpp()

    def _insert_sendrecv_ops_in_while_block(
        self,
        block,
        sync_in_while_lastpp2firstpp_var_names,
        sync_in_while_var_names,
        stage,
    ):
        dev_ids = []
        for pair in self._pipeline_pair_in_while:
            prev_id, cur_id = pair
            if prev_id not in dev_ids:
                dev_ids.append(prev_id)
            if cur_id not in dev_ids:
                dev_ids.append(cur_id)

        if len(dev_ids) == 0:
            return

        first_id = min(dev_ids)
        last_id = max(dev_ids)

        assert len(block.ops) > 2, (
            "It must have more than 2 ops in while sub block, "
            "layers.assign(layers.cast(cond_int, dtype='bool'), cond) must at end of while block, "
            "because nccl cannot send bool dtype var"
        )
        index = len(block.ops) - 2

        for prev_id in dev_ids:
            if prev_id == cur_id:
                continue
            assert cur_id > prev_id

            pair = (prev_id, cur_id)
            # 1000 is just a magic number
            pair_key = prev_id * 1000 + cur_id
            if pair not in self._pipeline_pair:
                self._pipeline_pair.append(pair)
                self._pp_ring_map[pair_key] = self.ring_id
                ring_id = self.ring_id
                self.ring_id += 1
            else:
                ring_id = self._pp_ring_map[pair_key]

            if cur_id == last_id and prev_id == first_id:
                var_names = (
                    sync_in_while_lastpp2firstpp_var_names
                    + sync_in_while_var_names
                )
            else:
                var_names = sync_in_while_var_names

            for var_name in var_names:
                var = block._var_recursive(var_name)
                if stage == cur_id:
                    block._insert_op_without_sync(
                        index=index,
                        type='send_v2',
                        inputs={'X': var},
                        attrs={
                            self._op_device_key: self._device
                            + ':'
                            + str(cur_id),
                            self._op_role_key: int(self._op_role.Forward),
                            'use_calc_stream': True,
                            'peer': 0,
                            'ring_id': ring_id,
                        },
                    )
                else:
                    var_shape = list(var.shape)
                    print(var_name)
                    if len(var.shape) > 0:
                        var_shape[0] = (
                            self.micro_batch_size
                            if var_shape[0] < 0
                            else var_shape[0]
                        )
                    block._insert_op_without_sync(
                        index=index,
                        type='recv_v2',
                        outputs={'Out': [var]},
                        attrs={
                            'out_shape': var_shape,
                            'dtype': var.dtype,
                            self._op_device_key: self._device
                            + ':'
                            + str(prev_id),
                            self._op_role_key: int(self._op_role.Forward),
                            'use_calc_stream': True,
                            'peer': 1,
                            'ring_id': ring_id,
                        },
                    )
                index += 1
        block._sync_with_cpp()

    def _get_while_block(self):
        """
        Get the while sub-block.
        """
        main_block = self._main_program.global_block()
        num_while = 0
        sub_block_id = None
        for op in main_block.ops:
            assert num_while < 2, "More than one while op found."
            if op.type == 'while':
                sub_block_id = op.attr('sub_block').id
                num_while += 1
        if sub_block_id:
            return op, self._main_program.block(sub_block_id)
        return None, None

    def gen_infer_program(
        self,
        sync_in_while_lastpp2firstpp_var_names=None,
        sync_in_while_var_names=None,
        debug=False,
    ):
        """
        Generate inference program.
        Params:
            sync_in_while_lastpp2firstpp_var_names (list(str)): the vars in the last pipeline
                that need to send var to first pipeline and exclude bool dtype var
            sync_in_while_var_names (list(str)): the vars sync among all pipeline in while block
                e.g cond. Note that cond cannot be bool dtype.
            debug (bool): the flag indicate debug
        """
        main_block = self._main_program.global_block()
        startup_block = self._startup_program.global_block()

        if debug:
            with open('main_program.txt', 'w') as f:
                f.write(str(self._main_program))
            with open('startup_program.txt', 'w') as f:
                f.write(str(self._startup_program))

        # step1: add op_device attribute for all ops
        self._add_op_device_attr(startup_block)
        self._check_validation(startup_block)
        self._add_op_device_attr(main_block)
        self._check_validation(main_block)

        # step2: add send/recv ops
        self._update_param_device_map()
        # step2.1: add send/recv for main_block
        out_var_to_op, in_var_to_op = self._get_input_output_info(main_block)
        self._output_var_to_op = out_var_to_op
        self._input_var_to_op = in_var_to_op
        self._insert_sendrecv_ops_for_boundaries(main_block, False)

        # step2.2: add send/recv for while_block
        while_op, while_block = self._get_while_block()
        if while_block:
            out_var_to_op, in_var_to_op = self._get_input_output_info(
                while_block
            )
            self._output_var_to_op = out_var_to_op
            self._input_var_to_op = in_var_to_op

            self._insert_sendrecv_ops_for_boundaries(while_block, True)

            self._insert_sendrecv_ops_in_while_block(
                while_block,
                sync_in_while_lastpp2firstpp_var_names,
                sync_in_while_var_names,
                self._stage,
            )

        # step3: split programs
        self._split_program(self._startup_program, self._stage, 0)
        self._split_program(self._main_program, self._stage, 0)

        if debug:
            with open(f'main_program.txt.{self.rank}', 'w') as f:
                f.write(str(self._main_program))
            with open(f'startup_program.txt.{self.rank}', 'w') as f:
                f.write(str(self._startup_program))

        if self.init_comm:
            self._init_communication_group()