distribute_transpiler.py 46.4 KB
Newer Older
1
#   Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
D
dzhwinter 已提交
2
#
D
dzhwinter 已提交
3 4 5
# 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
D
dzhwinter 已提交
6
#
D
dzhwinter 已提交
7
#     http://www.apache.org/licenses/LICENSE-2.0
D
dzhwinter 已提交
8
#
D
dzhwinter 已提交
9 10 11 12 13 14
# 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.

T
typhoonzero 已提交
15
from __future__ import print_function
16

T
typhoonzero 已提交
17
import math
18 19 20

import distributed_splitter as splitter
import framework
T
typhoonzero 已提交
21
from framework import Program, default_main_program, Variable, Parameter
22
from . import core
23 24 25 26

LOOKUP_TABLE_TYPE = "lookup_table"
LOOKUP_TABLE_GRAD_TYPE = "lookup_table_grad"
RPC_CLIENT_VAR_NAME = "RPC_CLIENT_VAR"
T
done  
typhoonzero 已提交
27

28 29
GLOBAL_BLOCK_IDX = 0

T
done  
typhoonzero 已提交
30

T
typhoonzero 已提交
31 32 33 34 35 36
class VarBlock:
    def __init__(self, varname, offset, size):
        self.varname = varname
        # NOTE: real offset is offset * size
        self.offset = offset
        self.size = size
T
done  
typhoonzero 已提交
37

T
typhoonzero 已提交
38 39
    def __str__(self):
        return "%s:%d:%d" % (self.varname, self.offset, self.size)
T
done  
typhoonzero 已提交
40 41


42
class UnionFind(object):
43
    """ Union-find data structure.
44

45
    Union-find is a data structure that keeps track of a set of elements partitioned
46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92
    into a number of disjoint (non-overlapping) subsets.

    Reference:
    https://en.wikipedia.org/wiki/Disjoint-set_data_structure

    Args:
      elements(list): The initialize element list.
    """

    def __init__(self, elementes=None):
        self._parents = []  # index -> parent index
        self._index = {}  # element -> index
        self._curr_idx = 0
        if not elementes:
            elementes = []
        for ele in elementes:
            self._parents.append(self._curr_idx)
            self._index.update({ele: self._curr_idx})
            self._curr_idx += 1

    def find(self, x):
        # Find the root index of given element x,
        # execute the path compress while findind the root index
        if not x in self._index:
            return -1
        idx = self._index[x]
        while idx != self._parents[idx]:
            t = self._parents[idx]
            self._parents[idx] = self._parents[t]
            idx = t
        return idx

    def union(self, x, y):
        # Union two given element
        x_root = self.find(x)
        y_root = self.find(y)

        if x_root == y_root:
            return
        self._parents[x_root] = y_root

    def is_connected(self, x, y):
        # If two given elements have the same root index,
        # then they are connected.
        return self.find(x) == self.find(y)


93 94 95 96
def same_or_split_var(p_name, var_name):
    return p_name == var_name or p_name.startswith(var_name + ".block")


T
typhoonzero 已提交
97 98 99 100 101
def split_dense_variable(var_list,
                         pserver_count,
                         min_block_size=1024,
                         max_block_size=1048576):
    """
102
        We may need to split dense tensor to one or more blocks and put
T
typhoonzero 已提交
103 104
        them equally onto parameter server. One block is a sub-tensor
        aligned by dim[0] of the tensor.
105

T
typhoonzero 已提交
106 107
        We need to have a minimal block size so that the calculations in
        the parameter server side can gain better performance. By default
108 109
        minimum block size is 1024. The max block size is used to prevent
        very large blocks that may cause send error.
110 111
        :return: A list of VarBlocks. Each VarBlock specifies a shard of
           the var.
T
typhoonzero 已提交
112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129
    """
    blocks = []
    for var in var_list:
        split_count = pserver_count
        var_numel = reduce(lambda x, y: x * y, var.shape)
        max_pserver_count = int(math.floor(var_numel / float(min_block_size)))
        if max_pserver_count == 0:
            max_pserver_count = 1
        if max_pserver_count < pserver_count:
            split_count = max_pserver_count
        block_size = int(math.ceil(var_numel / float(split_count)))

        if len(var.shape) >= 2:
            # align by dim1(width)
            dim1 = reduce(lambda x, y: x * y, var.shape[1:])
            remains = block_size % dim1
            if remains != 0:
                block_size += dim1 - remains
130
        # update split_count after aligning
T
typhoonzero 已提交
131 132 133 134 135 136 137 138 139
        split_count = int(math.ceil(var_numel / float(block_size)))
        for block_id in xrange(split_count):
            curr_block_size = min(block_size, var_numel - (
                (block_id) * block_size))
            block = VarBlock(var.name, block_id, curr_block_size)
            blocks.append(str(block))
    return blocks


T
done  
typhoonzero 已提交
140 141 142 143
class DistributeTranspiler:
    def transpile(self,
                  optimize_ops,
                  params_grads,
T
typhoonzero 已提交
144
                  trainer_id,
T
done  
typhoonzero 已提交
145 146 147
                  program=None,
                  pservers="127.0.0.1:6174",
                  trainers=1,
148
                  split_method=splitter.round_robin):
T
done  
typhoonzero 已提交
149
        """
150 151
            Transpile the program to distributed data-parallelism programs.
            The main_program will be transformed to use a remote parameter server
T
done  
typhoonzero 已提交
152
            to do parameter optimization. And the optimization graph will be put
153
            into a parameter server program.
T
done  
typhoonzero 已提交
154

155
            Use different methods to split trainable variables to different
T
done  
typhoonzero 已提交
156 157
            parameter servers.

T
typhoonzero 已提交
158 159 160 161 162 163 164 165 166 167 168 169 170 171 172
            Steps to transpile trainer:
            1. split variable to multiple blocks, aligned by product(dim[1:]) (width).
            2. rename splited grad variables to add trainer_id suffix ".trainer_%d".
            3. modify trainer program add split_op to each grad variable.
            4. append send_op to send splited variables to server and fetch
               params(splited blocks or origin param) from server.
            5. append concat_op to merge splited blocks to update local weights.

            Steps to transpile pserver:
            1. create new program for parameter server.
            2. create params and grad variables that assigned to current server instance.
            3. create a sub-block in the server side program
            4. append ops that should run on current server instance.
            5. add listen_and_serv op

T
done  
typhoonzero 已提交
173
            :param optimize_ops: op list of optimization, should be the
174
                                    return value of Optimizer.minimize
T
done  
typhoonzero 已提交
175
            :type optimize_ops: list
T
typhoonzero 已提交
176 177 178 179
            :param params_grads: list of tuple(weight, gradient)
            :type params_grads: list
            :param trainer_id: one unique id for each trainer in a job.
            :type trainer_id: int
T
typhoonzero 已提交
180
            :param program: program to transpile, default is default_main_program
T
typhoonzero 已提交
181
            :type program: Program
T
done  
typhoonzero 已提交
182 183
            :param pservers: parameter server endpoints like "m1:6174,m2:6174"
            :type pservers: string
T
typhoonzero 已提交
184 185 186 187 188
            :param trainers: total number of workers/trainers in the job
            :type trainers: int
            :param split_method: A function to determin how to split variables
                to different servers equally.
            :type split_method: function
T
done  
typhoonzero 已提交
189
        """
T
typhoonzero 已提交
190
        assert (callable(split_method))
T
done  
typhoonzero 已提交
191 192
        if program is None:
            program = default_main_program()
193 194
        self.origin_program = program
        self.trainer_num = trainers
T
typhoonzero 已提交
195
        self.optimize_ops = optimize_ops
T
typhoonzero 已提交
196 197 198 199
        # TODO(typhoonzero): currently trainer_id is fetched from cluster system
        # like Kubernetes, we should port this to use etcd later when developing
        # fluid distributed training with fault-tolerance.
        self.trainer_id = trainer_id
T
typhoonzero 已提交
200
        pserver_endpoints = pservers.split(",")
201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223
        self.pserver_endpoints = pserver_endpoints

        # process lookup_table_op
        # 1. check all lookup_table_op is distributed
        # 2. check all lookup_table_op share the same table.
        distributed_lookup_table_ops = []
        # support only one distributed_lookup_table now
        self.table_name = None
        for op in program.global_block().ops:
            if op.type == LOOKUP_TABLE_TYPE:
                if op.attrs['is_distributed'] is True:
                    if self.table_name is None:
                        self.table_name = op.input("W")[0]
                    if self.table_name != op.input("W")[0]:
                        raise RuntimeError("all distributed lookup_table_ops"
                                           " should have only one table")
                    distributed_lookup_table_ops.append(op)
                else:
                    if self.table_name is not None:
                        assert op.input("W")[0] != self.table_name

        self.has_distributed_lookup_table = len(
            distributed_lookup_table_ops) > 0
T
typhoonzero 已提交
224

225 226
        # step1: For large parameters and gradients, split them into smaller
        # blocks.
T
typhoonzero 已提交
227 228 229 230 231 232 233 234
        param_list = []
        grad_list = []
        for p, g in params_grads:
            # skip parameter marked not trainable
            if type(p) == Parameter and p.trainable == False:
                continue
            param_list.append(p)
            grad_list.append(g)
235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258

        if self.has_distributed_lookup_table:
            param_list = [
                param for param in param_list if param.name != self.table_name
            ]
            grad_list = [
                grad for grad in grad_list
                if grad.name != framework.grad_var_name(self.table_name)
            ]
            self.table_param_grad = [
                param_grad for param_grad in params_grads
                if param_grad[0].name == self.table_name
            ][0]
            table_grad_var = self.table_param_grad[1]
            self.table_grad_list = [
                program.global_block().create_var(
                    name="%s.trainer_%d.pserver_%d" %
                    (table_grad_var.name, trainer_id, index),
                    type=table_grad_var.type,
                    shape=table_grad_var.shape,
                    dtype=table_grad_var.dtype)
                for index in range(len(self.pserver_endpoints))
            ]

T
typhoonzero 已提交
259 260
        grad_blocks = split_dense_variable(grad_list, len(pserver_endpoints))
        param_blocks = split_dense_variable(param_list, len(pserver_endpoints))
261 262
        # step2: Create new vars for the parameters and gradients blocks and
        # add ops to do the split.
T
typhoonzero 已提交
263
        grad_var_mapping = self._append_split_op(program, grad_blocks)
264 265 266 267
        param_var_mapping = self._create_vars_from_blocklist(program,
                                                             param_blocks)
        # step3: Add gradients as send op inputs and parameters as send
        # op outputs.
T
typhoonzero 已提交
268
        send_inputs = []
T
typhoonzero 已提交
269
        send_outputs = []
T
typhoonzero 已提交
270 271 272 273 274 275
        for b in grad_blocks:  # append by order
            varname, block_id, _ = b.split(":")
            send_inputs.append(grad_var_mapping[varname][int(block_id)])
        for b in param_blocks:
            varname, block_id, _ = b.split(":")
            send_outputs.append(param_var_mapping[varname][int(block_id)])
276 277
        # let send_op know which endpoint to send which var to, eplist has the same
        # order as send_inputs.
T
typhoonzero 已提交
278
        eplist = split_method(send_inputs, pserver_endpoints)
279
        # create mapping of endpoint -> split var to create pserver side program
T
typhoonzero 已提交
280 281 282 283 284 285 286 287
        self.param_grad_ep_mapping = dict()
        for i, ep in enumerate(eplist):
            param = send_outputs[i]
            grad = send_inputs[i]
            if not self.param_grad_ep_mapping.has_key(ep):
                self.param_grad_ep_mapping[ep] = {"params": [], "grads": []}
            self.param_grad_ep_mapping[ep]["params"].append(param)
            self.param_grad_ep_mapping[ep]["grads"].append(grad)
T
typhoonzero 已提交
288

T
typhoonzero 已提交
289
        rpc_client_var = program.global_block().create_var(
290
            name=RPC_CLIENT_VAR_NAME,
T
typhoonzero 已提交
291
            persistable=True,
T
typhoonzero 已提交
292
            type=core.VarDesc.VarType.RAW)
T
typhoonzero 已提交
293

294
        # create send_op
T
typhoonzero 已提交
295
        program.global_block().append_op(
T
typhoonzero 已提交
296 297
            type="send",
            inputs={"X": send_inputs},
T
typhoonzero 已提交
298 299
            outputs={"Out": send_outputs,
                     "RPCClient": rpc_client_var},
T
typhoonzero 已提交
300
            attrs={"endpoints": pserver_endpoints,
T
typhoonzero 已提交
301
                   "epmap": eplist})
302
        # step4: Concat the parameters splits together after recv.
T
typhoonzero 已提交
303
        for varname, splited_var in param_var_mapping.iteritems():
T
typhoonzero 已提交
304 305
            if len(splited_var) <= 1:
                continue
T
typhoonzero 已提交
306
            orig_param = program.global_block().vars[varname]
T
typhoonzero 已提交
307
            program.global_block().append_op(
T
typhoonzero 已提交
308
                type="concat",
T
typhoonzero 已提交
309
                inputs={"X": splited_var},
T
typhoonzero 已提交
310
                outputs={"Out": [orig_param]},
T
typhoonzero 已提交
311
                attrs={"axis": 0})
T
typhoonzero 已提交
312

313 314 315 316 317 318
        if self.has_distributed_lookup_table:
            self._replace_lookup_table_op_with_prefetch(program, rpc_client_var,
                                                        eplist)
            self._split_table_grad_and_add_send_vars(program, rpc_client_var,
                                                     pserver_endpoints)

T
typhoonzero 已提交
319 320
    def get_trainer_program(self):
        # remove optimize ops and add a send op to main_program
321 322
        self.origin_program.global_block().delete_ops(self.optimize_ops)
        self.origin_program.sync_with_cpp()
323
        # FIXME(typhoonzero): serialize once will fix error occurs when clone.
324 325
        self.origin_program.__str__()
        return self.origin_program
T
typhoonzero 已提交
326 327 328 329

    def get_pserver_program(self, endpoint):
        """
        Get pserver side program using the endpoint.
330
        TODO(panyx0718): Revisit this assumption. what if #blocks > #pservers.
T
typhoonzero 已提交
331 332 333 334 335 336
        NOTE: assume blocks of the same variable is not distributed
        on the same pserver, only change param/grad varnames for
        trainers to fetch.
        """
        # step1
        pserver_program = Program()
337
        # step2: Create vars to receive vars at parameter servers.
T
typhoonzero 已提交
338 339 340 341 342 343 344 345
        recv_inputs = []
        for v in self.param_grad_ep_mapping[endpoint]["params"]:
            self._clone_var(pserver_program.global_block(), v)
        for v in self.param_grad_ep_mapping[endpoint]["grads"]:
            # create vars for each trainer in global scope, so
            # we don't need to create them when grad arrives.
            # change client side var name to origin name by
            # removing ".trainer_%d" suffix
T
update  
typhoonzero 已提交
346 347 348 349 350 351

            suff_idx = v.name.find(".trainer_")
            if suff_idx >= 0:
                orig_var_name = v.name[:suff_idx]
            else:
                orig_var_name = v.name
T
typhoonzero 已提交
352 353 354 355 356 357 358 359 360
            # NOTE: single_trainer_var must be created for multi-trainer
            # case to merge grads from multiple trainers
            single_trainer_var = \
                pserver_program.global_block().create_var(
                    name=orig_var_name,
                    persistable=True,
                    type=v.type,
                    dtype=v.dtype,
                    shape=v.shape)
361 362
            if self.trainer_num > 1:
                for trainer_id in xrange(self.trainer_num):
T
typhoonzero 已提交
363 364 365 366 367 368 369 370 371
                    var = pserver_program.global_block().create_var(
                        name="%s.trainer_%d" % (orig_var_name, trainer_id),
                        persistable=False,
                        type=v.type,
                        dtype=v.dtype,
                        shape=v.shape)
                    recv_inputs.append(var)
            else:
                recv_inputs.append(single_trainer_var)
372

Q
qiaolongfei 已提交
373 374
        # step 3
        # each optimization op will has a optimize block
375 376
        optimize_block = None

T
typhoonzero 已提交
377
        # step 4
378
        # Create a union-find data structure from optimize ops,
T
typhoonzero 已提交
379 380 381
        # If two ops are connected, we could add these two ops
        # into one set.
        ufind = self._create_ufind(self.optimize_ops)
382
        # step 4.2
T
typhoonzero 已提交
383 384 385 386 387 388 389 390
        # Iterate through the ops and append optimize op which
        # located on current pserver
        opt_op_on_pserver = []
        for _, op in enumerate(self.optimize_ops):
            if self._is_opt_op(op) and self._is_opt_op_on_pserver(endpoint, op):
                opt_op_on_pserver.append(op)
        # step 4.3
        # Iterate through the ops, and if an op and the optimize ops
391
        # which located on current pserver are in one set, then
T
typhoonzero 已提交
392
        # append it into the sub program.
T
typhoonzero 已提交
393 394 395 396 397 398 399 400 401 402 403 404 405 406 407 408 409 410 411 412 413 414 415 416 417 418 419 420 421

        # We try to put optimization program run parallelly, assume
        # optimization program always looks like:
        #
        # prevop -> prevop -> opt op -> following op -> following op; ->
        # prevop -> prevop -> opt op -> following op -> following op; ->
        # global op -> global op
        #
        # we put operators that can run parallelly to many program blocks.
        # in above example, we seperate ops by the ";". Global ops must run
        # after all the optimize ops finished.

        global_ops = []
        # HACK: optimization global ops only used to scale beta1 and beta2
        # replace it with dependency engine.
        for op in self.optimize_ops:
            if op.type == "scale":
                for in_name in op.input_arg_names:
                    if in_name.startswith("beta1_pow_acc") or\
                        in_name.startswith("beta2_pow_acc"):
                        global_ops.append(op)

        def __append_optimize_op__(op, block):
            if self._is_opt_op(op):
                self._append_pserver_ops(block, op, endpoint,
                                         default_main_program())
            else:
                self._append_pserver_non_opt_ops(block, op)

422
        # append lr decay ops to the child block if exists
423
        lr_decay_block = None
424 425
        lr_ops = self._get_lr_ops()
        if len(lr_ops) > 0:
426
            lr_decay_block = pserver_program.create_block(GLOBAL_BLOCK_IDX)
427
            for _, op in enumerate(lr_ops):
428
                self._append_pserver_non_opt_ops(lr_decay_block, op)
429

T
typhoonzero 已提交
430
        # append op to the current block
431 432 433 434
        per_opt_block = None
        pre_block_idx = GLOBAL_BLOCK_IDX
        if lr_decay_block is not None:
            pre_block_idx = lr_decay_block.idx
T
typhoonzero 已提交
435
        for idx, opt_op in enumerate(opt_op_on_pserver):
436 437
            per_opt_block = pserver_program.create_block(pre_block_idx)
            if optimize_block is None:
Q
qiaolongfei 已提交
438
                # first optimize block
439
                optimize_block = per_opt_block
T
typhoonzero 已提交
440 441
            for _, op in enumerate(self.optimize_ops):
                # optimizer is connected to itself
442
                if ufind.is_connected(op, opt_op) and op not in global_ops:
T
typhoonzero 已提交
443 444 445
                    __append_optimize_op__(op, per_opt_block)

        # append global ops
446 447 448
        opt_state_block = None
        if global_ops:
            opt_state_block = pserver_program.create_block(per_opt_block.idx)
T
typhoonzero 已提交
449
        for glb_op in global_ops:
450
            __append_optimize_op__(glb_op, opt_state_block)
T
typhoonzero 已提交
451 452 453 454 455 456 457 458 459

        # NOT USED: single block version:
        #
        # for _, op in enumerate(self.optimize_ops):
        #     for _, opt_op in enumerate(opt_op_on_pserver):
        #         if ufind.is_connected(op, opt_op):
        #             __append_optimize_op__(glb_op, optimize_block)
        #             break

460 461 462 463
        # process distributed lookup_table
        prefetch_block = None
        if self.has_distributed_lookup_table:
            pserver_index = self.pserver_endpoints.index(endpoint)
464 465 466
            table_opt_block = self._create_table_optimize_block(
                pserver_index, pserver_program, opt_state_block or
                pserver_program.global_block())
467
            prefetch_block = self._create_prefetch_block(
468
                pserver_index, pserver_program, table_opt_block)
469 470 471 472 473 474 475 476 477

        # NOTE: if has_distributed_lookup_table is False, then prefetch_block will
        # not be executed, so it's safe to use optimize_block to hold the place
        if self.has_distributed_lookup_table:
            assert prefetch_block is not None
        else:
            assert prefetch_block is None
            prefetch_block = pserver_program.global_block()

T
typhoonzero 已提交
478 479 480 481 482 483 484 485
        # step5 append the listen_and_serv op
        pserver_program.global_block().append_op(
            type="listen_and_serv",
            inputs={'X': recv_inputs},
            outputs={},
            attrs={
                "OptimizeBlock": optimize_block,
                "endpoint": endpoint,
486 487
                "Fanin": self.trainer_num,
                "PrefetchBlock": prefetch_block
T
typhoonzero 已提交
488
            })
489

T
typhoonzero 已提交
490 491 492 493 494 495 496 497 498 499 500 501 502 503 504 505 506 507 508 509 510 511 512 513
        pserver_program.sync_with_cpp()
        return pserver_program

    def get_startup_program(self, endpoint, pserver_program):
        """
        Get startup program for current parameter server.
        Modify operator input variables if there are variables that
        were split to several blocks.
        """
        s_prog = Program()
        orig_s_prog = framework.default_startup_program()
        params = self.param_grad_ep_mapping[endpoint]["params"]

        def _get_splited_name_and_shape(varname):
            for idx, splited_param in enumerate(params):
                pname = splited_param.name
                if same_or_split_var(pname, varname) and varname != pname:
                    return pname, splited_param.shape
            return "", []

        # 1. create vars in pserver program to startup program
        pserver_vars = pserver_program.global_block().vars
        created_var_map = dict()
        for _, var in pserver_vars.iteritems():
T
update  
typhoonzero 已提交
514
            tmpvar = s_prog.global_block().clone_variable(var)
T
typhoonzero 已提交
515 516 517 518 519 520 521 522 523 524 525 526 527 528 529 530 531 532 533 534 535 536 537 538 539 540 541 542 543 544 545 546
            created_var_map[var.name] = tmpvar

        # 2. rename op outputs
        for op in orig_s_prog.global_block().ops:
            new_inputs = dict()
            new_outputs = dict()
            # do not append startup op if var is not on this pserver
            op_on_pserver = False
            for key in op.output_names:
                newname, _ = _get_splited_name_and_shape(op.output(key)[0])
                if newname:
                    op_on_pserver = True
                    new_outputs[key] = created_var_map[newname]
                elif op.output(key)[0] in pserver_vars:
                    op_on_pserver = True
                    new_outputs[key] = pserver_vars[op.output(key)[0]]

            # most startup program ops have no inputs
            new_inputs = self._get_input_map_from_op(pserver_vars, op)

            if op_on_pserver:
                if op.type in [
                        "gaussian_random", "fill_constant", "uniform_random"
                ]:
                    op.attrs["shape"] = new_outputs["Out"].shape
                s_prog.global_block().append_op(
                    type=op.type,
                    inputs=new_inputs,
                    outputs=new_outputs,
                    attrs=op.attrs)
        return s_prog

547 548 549 550 551 552 553 554 555 556 557 558 559 560 561 562 563 564 565 566 567 568 569 570 571 572 573 574 575 576 577 578 579 580 581 582 583 584 585 586 587 588 589 590 591 592 593 594 595 596 597 598 599 600 601 602 603 604 605 606 607 608 609 610 611 612 613 614 615 616 617 618 619 620 621 622 623 624 625 626 627 628 629 630 631 632 633 634 635 636 637 638 639 640 641 642 643 644 645 646 647 648 649 650 651 652 653 654 655 656 657 658 659 660 661 662 663 664 665 666 667 668 669 670 671 672 673 674 675 676 677 678 679 680 681 682 683 684 685 686 687 688 689 690 691 692 693 694 695 696 697 698 699 700 701 702 703 704 705 706 707 708 709 710 711 712 713 714 715 716 717 718 719 720 721 722 723 724 725 726 727 728 729 730 731 732 733 734 735 736 737
    # transpiler function for dis lookup_table
    def _replace_lookup_table_op_with_prefetch(self, program, rpc_client_var,
                                               eplist):
        # 1. replace lookup_table_op with split_ids_op -> prefetch_op -> sum_op
        self.prefetch_input_vars = None
        self.prefetch_output_vars = None

        continue_search_lookup_table_op = True
        while continue_search_lookup_table_op:
            continue_search_lookup_table_op = False
            all_ops = program.global_block().ops
            for op in all_ops:
                if op.type == LOOKUP_TABLE_TYPE:
                    continue_search_lookup_table_op = True

                    op_index = list(all_ops).index(op)
                    ids_name = op.input("Ids")
                    out_name = op.output("Out")

                    if self.prefetch_input_vars is None:
                        ids_var = program.global_block().vars[ids_name[0]]
                        self.prefetch_input_vars = self.create_splited_vars(
                            source_var=ids_var,
                            block=program.global_block(),
                            tag="_prefetch_in_")
                    if self.prefetch_output_vars is None:
                        out_var = program.global_block().vars[out_name[0]]
                        self.prefetch_output_vars = self.create_splited_vars(
                            source_var=out_var,
                            block=program.global_block(),
                            tag="_prefetch_out_")

                    # insert split_ids_op
                    program.global_block().insert_op(
                        index=op_index,
                        type="split_ids",
                        inputs={
                            'Ids': [
                                program.global_block().vars[varname]
                                for varname in ids_name
                            ]
                        },
                        outputs={"Out": self.prefetch_input_vars})

                    # insert prefetch_op
                    program.global_block().insert_op(
                        index=op_index + 1,
                        type="prefetch",
                        inputs={'X': self.prefetch_input_vars},
                        outputs={
                            "Out": self.prefetch_output_vars,
                            "RPCClient": rpc_client_var
                        },
                        attrs={"epmap": eplist})

                    # insert concat_op
                    program.global_block().insert_op(
                        index=op_index + 2,
                        type="concat",
                        inputs={'X': self.prefetch_output_vars},
                        outputs={
                            "Out": [
                                program.global_block().vars[varname]
                                for varname in out_name
                            ]
                        },
                        attrs={"axis": 0})

                    # delete lookup_table_op
                    program.global_block().delete_ops([op])
                    program.sync_with_cpp()
                    # break for loop
                    break

    def _split_table_grad_and_add_send_vars(self, program, rpc_client_var,
                                            pserver_endpoints):
        # 2. add split_ids_op and send_vars_op to send gradient to pservers
        # there should only be one table_name
        all_ops = program.global_block().ops
        table_grad_name = framework.grad_var_name(self.table_name)
        for op in all_ops:
            if table_grad_name in op.output_arg_names:
                op_index = list(all_ops).index(op)
                # insert split_ids_op
                program.global_block().insert_op(
                    index=op_index + 1,
                    type="split_ids",
                    inputs={
                        'Ids': [program.global_block().vars[table_grad_name]]
                    },
                    outputs={"Out": self.table_grad_list})
                program.global_block().insert_op(
                    index=op_index + 2,
                    type="send_vars",
                    inputs={'X': self.table_grad_list},
                    outputs={"RPCClient": rpc_client_var},
                    attrs={"sync_send": True,
                           "epmap": pserver_endpoints})
                break

    def _create_prefetch_block(self, pserver_index, pserver_program,
                               optimize_block):
        # STEP: create prefetch block
        table_var = pserver_program.global_block().vars[self.table_name]
        prefetch_block = pserver_program.create_block(optimize_block.idx)
        trainer_ids = self.prefetch_input_vars[pserver_index]
        pserver_ids = pserver_program.global_block().create_var(
            name=trainer_ids.name,
            type=trainer_ids.type,
            shape=trainer_ids.shape,
            dtype=trainer_ids.dtype)
        trainer_out = self.prefetch_output_vars[pserver_index]
        pserver_out = pserver_program.global_block().create_var(
            name=trainer_out.name,
            type=trainer_out.type,
            shape=trainer_out.shape,
            dtype=trainer_out.dtype)
        prefetch_block.append_op(
            type=LOOKUP_TABLE_TYPE,
            inputs={'Ids': pserver_ids,
                    "W": table_var},
            outputs={"Out": pserver_out},
            attrs={
                "is_sparse": True,  # has no effect on lookup_table op
                "is_distributed": True,
                "padding_idx": -1
            })
        return prefetch_block

    def _create_table_optimize_block(self, pserver_index, pserver_program,
                                     append_block):
        def _clone_var(block, var, persistable=True):
            assert isinstance(var, Variable)
            return block.create_var(
                name=var.name,
                shape=var.shape,
                dtype=var.dtype,
                type=var.type,
                persistable=persistable)

        # STEP: create table optimize block
        # create table param and grad var in pserver program
        param_var = _clone_var(
            pserver_program.global_block(),
            self.origin_program.global_block().vars[self.table_name])
        grad_var = _clone_var(
            pserver_program.global_block(),
            self.origin_program.global_block().vars[framework.grad_var_name(
                self.table_name)],
            persistable=False)

        # create grad vars in pserver program
        table_grad_var = self.table_param_grad[1]
        table_grad_list = [
            pserver_program.global_block().create_var(
                name="%s.trainer_%d.pserver_%d" %
                (table_grad_var.name, index, pserver_index),
                type=table_grad_var.type,
                shape=table_grad_var.shape,
                dtype=table_grad_var.dtype) for index in range(self.trainer_num)
        ]

        # create table optimize block in pserver program
        table_opt_op = [
            op for op in self.optimize_ops
            if op.input("Param")[0] == self.table_name
        ][0]
        table_opt_block = pserver_program.create_block(append_block.idx)
        # only support sgd now
        assert table_opt_op.type == "sgd"

        # append sum op for table_grad_list
        table_opt_block.append_op(
            type="sum",
            inputs={"X": table_grad_list},
            outputs={"Out": [grad_var]})

        lr_var = pserver_program.global_block().vars[table_opt_op.input(
            "LearningRate")[0]]
        inputs = {
            "Param": [param_var],
            "Grad": [grad_var],
            "LearningRate": [lr_var]
        }
        outputs = {"ParamOut": [param_var]}
        table_opt_block.append_op(
            type=table_opt_op.type,
            inputs=inputs,
            outputs=outputs,
            attrs=table_opt_op.attrs)

738 739
        return table_opt_block

T
typhoonzero 已提交
740 741 742 743 744 745
    # ====================== private transpiler functions =====================
    def _create_vars_from_blocklist(self,
                                    program,
                                    block_list,
                                    add_trainer_suffix=False):
        """
746
        Create vars for each split.
T
typhoonzero 已提交
747 748
        NOTE: only grads need to be named for different trainers, use
              add_trainer_suffix to rename the grad vars.
749
        :return: A dict mapping from original var name to each var split.
T
typhoonzero 已提交
750
        """
T
typhoonzero 已提交
751
        block_map = dict()
T
typhoonzero 已提交
752
        var_mapping = dict()
T
typhoonzero 已提交
753 754 755 756 757 758
        for block_str in block_list:
            varname, offset, size = block_str.split(":")
            if not block_map.has_key(varname):
                block_map[varname] = []
            block_map[varname].append((long(offset), long(size)))
        for varname, splited in block_map.iteritems():
T
typhoonzero 已提交
759
            orig_var = program.global_block().var(varname)
T
typhoonzero 已提交
760
            if len(splited) == 1:
T
typhoonzero 已提交
761 762 763 764 765 766 767 768 769
                if add_trainer_suffix:
                    new_var_name = "%s.trainer_%d" % \
                        (orig_var.name, self.trainer_id)
                    program.global_block().rename_var(varname, new_var_name)
                    var_mapping[varname] = \
                        [program.global_block().var(new_var_name)]
                else:
                    var_mapping[varname] = \
                        [program.global_block().var(orig_var.name)]
T
typhoonzero 已提交
770
                continue
T
typhoonzero 已提交
771 772

            var_mapping[varname] = []
T
typhoonzero 已提交
773 774 775 776
            orig_shape = orig_var.shape
            orig_dim1_flatten = 1
            if len(orig_shape) >= 2:
                orig_dim1_flatten = reduce(lambda x, y: x * y, orig_shape[1:])
T
typhoonzero 已提交
777

T
typhoonzero 已提交
778
            for i, block in enumerate(splited):
T
typhoonzero 已提交
779
                size = block[1]
T
typhoonzero 已提交
780 781 782 783
                rows = size / orig_dim1_flatten
                splited_shape = [rows]
                if len(orig_shape) >= 2:
                    splited_shape.extend(orig_shape[1:])
T
typhoonzero 已提交
784 785 786 787 788 789 790
                new_var_name = ""
                if add_trainer_suffix:
                    new_var_name = "%s.block%d.trainer_%d" % \
                        (varname, i, self.trainer_id)
                else:
                    new_var_name = "%s.block%d" % \
                        (varname, i)
T
typhoonzero 已提交
791
                var = program.global_block().create_var(
T
typhoonzero 已提交
792 793
                    name=new_var_name,
                    persistable=False,
T
typhoonzero 已提交
794
                    dtype=orig_var.dtype,
795
                    type=orig_var.type,
T
typhoonzero 已提交
796
                    shape=splited_shape)  # flattend splited var
T
typhoonzero 已提交
797
                var_mapping[varname].append(var)
T
typhoonzero 已提交
798
            program.global_block().sync_with_cpp()
T
typhoonzero 已提交
799
        return var_mapping
T
done  
typhoonzero 已提交
800

801 802 803 804 805 806 807 808 809 810 811
    def create_splited_vars(self, source_var, block, tag):
        return [
            block.create_var(
                name=str(source_var.name + tag + str(index)),
                type=source_var.type,
                shape=source_var.shape,
                dtype=source_var.dtype)
            for index in range(len(self.pserver_endpoints))
        ]

    def _clone_var(self, block, var, persistable=True):
T
done  
typhoonzero 已提交
812 813 814 815 816 817 818
        assert isinstance(var, Variable)
        return block.create_var(
            name=var.name,
            shape=var.shape,
            dtype=var.dtype,
            type=var.type,
            lod_level=var.lod_level,
819
            persistable=persistable)
T
done  
typhoonzero 已提交
820

T
typhoonzero 已提交
821
    def _append_split_op(self, program, gradblocks):
822
        # Split variables that need to be split and append respective ops
T
typhoonzero 已提交
823
        add_suffix = False
824
        if self.trainer_num > 1:
T
typhoonzero 已提交
825
            add_suffix = True
T
typhoonzero 已提交
826
        var_mapping = self._create_vars_from_blocklist(
T
typhoonzero 已提交
827
            program, gradblocks, add_trainer_suffix=add_suffix)
T
typhoonzero 已提交
828
        for varname, splited_vars in var_mapping.iteritems():
T
typhoonzero 已提交
829 830
            # variable that don't need to split have empty splited_vars
            if len(splited_vars) <= 1:
T
typhoonzero 已提交
831
                continue
T
typhoonzero 已提交
832
            orig_var = program.global_block().vars[varname]
T
typhoonzero 已提交
833
            if orig_var.type == core.VarDesc.VarType.SELECTED_ROWS:
834 835 836 837 838 839 840 841
                height_sections = []
                for v in splited_vars:
                    height_sections.append(v.shape[0])
                program.global_block().append_op(
                    type="split_selected_rows",
                    inputs={"X": orig_var},
                    outputs={"Out": splited_vars},
                    attrs={"height_sections": height_sections})
T
typhoonzero 已提交
842
            elif orig_var.type == core.VarDesc.VarType.LOD_TENSOR:
843 844 845 846
                sections = []
                for v in splited_vars:
                    sections.append(v.shape[0])
                program.global_block().append_op(
T
typhoonzero 已提交
847
                    type="split_byref",
848 849 850 851 852 853 854
                    inputs={"X": orig_var},
                    outputs={"Out": splited_vars},
                    attrs={"sections": sections}  # assume split evenly
                )
            else:
                AssertionError("Variable type should be in set "
                               "[LOD_TENSOR, SELECTED_ROWS]")
T
typhoonzero 已提交
855
        return var_mapping
T
done  
typhoonzero 已提交
856

T
typhoonzero 已提交
857 858 859 860
    def _get_optimizer_input_shape(self, op_type, varkey, orig_shape,
                                   param_shape):
        """
        Returns the shape for optimizer inputs that need to be reshaped when
861
        Param and Grad is split to multiple servers.
T
typhoonzero 已提交
862 863 864 865 866 867 868 869 870 871 872 873 874 875 876 877 878 879 880 881 882 883
        """
        # HACK(typhoonzero): Should use functions of corresponding optimizer in
        # optimizer.py to get the shape, do not  bind this in the transpiler.
        if op_type == "adam":
            if varkey in ["Moment1", "Moment2"]:
                return param_shape
        elif op_type == "adagrad":
            if varkey == "Moment":
                return param_shape
        elif op_type == "adamax":
            if varkey in ["Moment", "InfNorm"]:
                return param_shape
        elif op_type == "momentum":
            if varkey == "Velocity":
                return param_shape
        elif op_type == "":
            if varkey == "Moment":
                return param_shape
        elif op_type == "sgd":
            pass
        return orig_shape

T
typhoonzero 已提交
884 885 886 887 888
    def _orig_varname(self, varname):
        suff_idx = varname.find(".trainer_")
        orig_var_name = ""
        if suff_idx >= 0:
            orig_var_name = varname[:suff_idx]
T
typhoonzero 已提交
889 890
        else:
            orig_var_name = varname
T
typhoonzero 已提交
891 892
        return orig_var_name

893 894
    def _append_pserver_ops(self, optimize_block, opt_op, endpoint,
                            origin_program):
895
        program = optimize_block.program
T
typhoonzero 已提交
896
        pserver_block = program.global_block()
T
typhoonzero 已提交
897
        new_inputs = dict()
T
typhoonzero 已提交
898 899
        # update param/grad shape first, then other inputs like
        # moment can use the updated shape
T
typhoonzero 已提交
900
        for key in opt_op.input_names:
T
typhoonzero 已提交
901 902 903
            if key == "Grad":
                grad_block = None
                for g in self.param_grad_ep_mapping[endpoint]["grads"]:
T
typhoonzero 已提交
904
                    if same_or_split_var(
T
typhoonzero 已提交
905 906
                            self._orig_varname(g.name),
                            self._orig_varname(opt_op.input(key)[0])):
T
typhoonzero 已提交
907 908 909 910 911 912
                        grad_block = g
                        break
                if not grad_block:
                    # do not append this op if current endpoint
                    # is not dealing with this grad block
                    return
T
typhoonzero 已提交
913 914
                merged_var = \
                    pserver_block.vars[self._orig_varname(grad_block.name)]
915
                if self.trainer_num > 1:
T
typhoonzero 已提交
916
                    vars2merge = []
917
                    for i in xrange(self.trainer_num):
T
typhoonzero 已提交
918 919 920 921
                        per_trainer_name = "%s.trainer_%d" % \
                        (self._orig_varname(grad_block.name), i)
                        vars2merge.append(pserver_block.vars[per_trainer_name])

922
                    optimize_block.append_op(
T
done  
typhoonzero 已提交
923 924 925
                        type="sum",
                        inputs={"X": vars2merge},
                        outputs={"Out": merged_var})
926
                    # TODO(panyx0718): What if it's SELECTED_ROWS.
927 928 929 930 931
                    if not merged_var.type == core.VarDesc.VarType.SELECTED_ROWS:
                        optimize_block.append_op(
                            type="scale",
                            inputs={"X": merged_var},
                            outputs={"Out": merged_var},
932
                            attrs={"scale": 1.0 / float(self.trainer_num)})
T
typhoonzero 已提交
933 934 935 936 937
                new_inputs[key] = merged_var
            elif key == "Param":
                # param is already created on global program
                param_block = None
                for p in self.param_grad_ep_mapping[endpoint]["params"]:
T
typhoonzero 已提交
938
                    if same_or_split_var(p.name, opt_op.input(key)[0]):
T
typhoonzero 已提交
939 940 941 942
                        param_block = p
                        break
                if not param_block:
                    return
T
typhoonzero 已提交
943
                tmpvar = pserver_block.create_var(
T
typhoonzero 已提交
944
                    name=param_block.name,
T
typhoonzero 已提交
945
                    persistable=True,
T
typhoonzero 已提交
946 947 948
                    dtype=param_block.dtype,
                    shape=param_block.shape)
                new_inputs[key] = tmpvar
949
            elif key == "LearningRate":
950
                # learning rate variable has already be created by non-optimize op,
951
                # don't create it once again.
952 953 954 955 956 957 958 959 960 961 962
                lr_varname = opt_op.input(key)[0]
                if pserver_block.vars.has_key(lr_varname):
                    new_inputs[key] = pserver_block.vars[opt_op.input(key)[0]]
                else:
                    origin_var = origin_program.global_block().vars[lr_varname]
                    tmpvar = pserver_block.create_var(
                        name=origin_var.name,
                        persistable=origin_var.persistable,
                        dtype=origin_var.dtype,
                        shape=origin_var.shape)
                    new_inputs[key] = tmpvar
T
typhoonzero 已提交
963

T
typhoonzero 已提交
964
        for key in opt_op.input_names:
965 966
            new_shape = None
            if key in ["Param", "Grad", "LearningRate"]:
T
typhoonzero 已提交
967
                continue
968
            var = self.origin_program.global_block().vars[opt_op.input(key)[0]]
T
typhoonzero 已提交
969 970 971 972
            # update accumulator variable shape
            param_shape = new_inputs["Param"].shape
            new_shape = self._get_optimizer_input_shape(opt_op.type, key,
                                                        var.shape, param_shape)
T
typhoonzero 已提交
973
            tmpvar = pserver_block.create_var(
T
typhoonzero 已提交
974 975 976 977 978
                name=var.name,
                persistable=var.persistable,
                dtype=var.dtype,
                shape=new_shape)
            new_inputs[key] = tmpvar
T
typhoonzero 已提交
979

980
        # change output's ParamOut variable
981 982
        outputs = self._get_output_map_from_op(
            self.origin_program.global_block().vars, opt_op)
983
        outputs["ParamOut"] = new_inputs["Param"]
T
typhoonzero 已提交
984

985
        optimize_block.append_op(
T
typhoonzero 已提交
986 987
            type=opt_op.type,
            inputs=new_inputs,
T
typhoonzero 已提交
988
            outputs=outputs,
T
typhoonzero 已提交
989 990
            attrs=opt_op.attrs)

991 992
    def _append_pserver_non_opt_ops(self, optimize_block, opt_op):
        program = optimize_block.program
993
        # Append the ops for parameters that do not need to be optimized/updated
994 995
        inputs = self._get_input_map_from_op(
            self.origin_program.global_block().vars, opt_op)
996 997 998 999
        for varlist in inputs.itervalues():
            if not isinstance(varlist, list):
                varlist = [varlist]

T
typhoonzero 已提交
1000
            for var in varlist:
1001 1002
                if not program.global_block().vars.has_key(var.name):
                    program.global_block().create_var(
T
typhoonzero 已提交
1003 1004 1005 1006 1007
                        name=var.name,
                        persistable=var.persistable,
                        dtype=var.dtype,
                        shape=var.shape)

1008 1009
        outputs = self._get_output_map_from_op(
            self.origin_program.global_block().vars, opt_op)
T
typhoonzero 已提交
1010

1011 1012 1013 1014 1015
        for varlist in outputs.itervalues():
            if not isinstance(varlist, list):
                varlist = [varlist]

            for var in varlist:
T
update  
typhoonzero 已提交
1016
                program.global_block().clone_variable(var)
1017

1018
        optimize_block.append_op(
T
typhoonzero 已提交
1019
            type=opt_op.type,
T
typhoonzero 已提交
1020 1021
            inputs=inputs,
            outputs=outputs,
T
typhoonzero 已提交
1022 1023
            attrs=opt_op.attrs)

1024 1025 1026 1027
    def _is_op_connected(self, op1, op2):
        # If one op's input is another op's output or
        # one op's output is another op's input, we say
        # the two operator is connected.
T
typhoonzero 已提交
1028 1029 1030 1031 1032 1033 1034 1035 1036 1037 1038 1039 1040
        def _append_inname_remove_beta(varname_list):
            op_input_names = []
            for in_name in varname_list:
                # HACK: remove beta1 and beta2 to avoid let all
                # ops connected.
                if in_name.startswith("beta2_pow_acc") or \
                    in_name.startswith("beta1_pow_acc"):
                    continue
                else:
                    op_input_names.append(in_name)
            return op_input_names

        op1_input_names = _append_inname_remove_beta(op1.desc.input_arg_names())
T
typhoonzero 已提交
1041 1042
        op1_output_names = op1.desc.output_arg_names()

T
typhoonzero 已提交
1043
        op2_input_names = _append_inname_remove_beta(op2.desc.input_arg_names())
T
typhoonzero 已提交
1044
        op2_output_names = op2.desc.output_arg_names()
1045 1046 1047 1048 1049 1050 1051 1052 1053 1054 1055 1056 1057 1058 1059 1060 1061 1062 1063

        if set(op1_output_names) & set(op2_input_names) or \
           set(op1_input_names) & set(op2_output_names):
            return True
        return False

    def _create_ufind(self, optimize_ops):
        # Create a unit find data struct by optimize ops
        ufind = UnionFind(optimize_ops)
        for i in xrange(len(optimize_ops)):
            for j in xrange(i, len(optimize_ops)):
                op1 = optimize_ops[i]
                op2 = optimize_ops[j]
                if self._is_op_connected(op1, op2):
                    ufind.union(op1, op2)
        return ufind

    def _is_opt_op(self, op):
        # NOTE: It's a HACK implement.
1064
        # optimize op: SGDOptimize, MomentumOptimizer, AdamOptimizer and etc...
T
typhoonzero 已提交
1065 1066
        if "Param" in op.input_names and \
            "LearningRate" in op.input_names:
1067 1068 1069 1070 1071 1072 1073
            return True
        return False

    def _is_opt_op_on_pserver(self, endpoint, op):
        param_names = [
            p.name for p in self.param_grad_ep_mapping[endpoint]["params"]
        ]
T
typhoonzero 已提交
1074
        if op.input("Param")[0] in param_names:
1075 1076 1077
            return True
        else:
            for n in param_names:
T
typhoonzero 已提交
1078
                param = op.input("Param")[0]
T
typhoonzero 已提交
1079
                if same_or_split_var(n, param) and n != param:
1080 1081 1082
                    return True
            return False

T
typhoonzero 已提交
1083
    def _get_input_map_from_op(self, varmap, op):
1084
        """Returns a dict from op input name to the vars in varmap."""
T
typhoonzero 已提交
1085 1086 1087 1088 1089 1090 1091 1092 1093 1094 1095 1096
        iomap = dict()
        for key in op.input_names:
            vars = []
            for varname in op.input(key):
                vars.append(varmap[varname])
            if len(vars) == 1:
                iomap[key] = vars[0]
            else:
                iomap[key] = vars
        return iomap

    def _get_output_map_from_op(self, varmap, op):
1097
        """Returns a dict from op output name to the vars in varmap."""
T
typhoonzero 已提交
1098 1099 1100 1101 1102 1103 1104 1105 1106 1107
        iomap = dict()
        for key in op.output_names:
            vars = []
            for varname in op.output(key):
                vars.append(varmap[varname])
            if len(vars) == 1:
                iomap[key] = vars[0]
            else:
                iomap[key] = vars
        return iomap
1108 1109 1110 1111 1112 1113 1114 1115 1116 1117 1118

    def _get_lr_ops(self):
        lr_ops = []
        # find learning rate variables by optimize op
        lr_vars = set()
        for op in self.optimize_ops:
            if self._is_opt_op(op):
                lr_vars.add(op.input("LearningRate")[0])

        find_ops = []
        # find ops which output is lr var
1119
        block = self.origin_program.global_block()
1120 1121 1122 1123 1124
        for op in block.ops:
            if set(op.output_arg_names) & lr_vars:
                find_ops.append(op)
        # make a union find struct by the ops in default_main_program
        ufind = UnionFind(block.ops)
1125

1126 1127 1128 1129 1130 1131 1132 1133 1134 1135 1136 1137
        for op1 in block.ops:
            for op2 in block.ops:
                # NOTE: we need to skip all optimize ops, since it is connected
                # with forward/backward ops and lr ops, we only need the lr ops.
                if op1 != op2 and self._is_op_connected(op1, op2) and \
                    not self._is_opt_op(op1) and not self._is_opt_op(op2):
                    ufind.union(op1, op2)
        # find all ops which is related with lr var
        for op1 in block.ops:
            for op2 in find_ops:
                if ufind.is_connected(op1, op2):
                    lr_ops.append(op1)
1138 1139
                    # we only need to append op for once
                    break
1140
        return lr_ops