def serialize_allreduce_ops()

in blocksparse/nccl.py [0:0]

• 53 lines of code
• 16 McCabe index (conditional complexity)

def serialize_allreduce_ops(graph_targets, serialize_inputs=True, print_dag=False):

    # Traverse all graph_targets through their inputs and:
    # Build a mutable dag of dict()'s' with ops as keys and their input ops as values (as set() elements)
    # For ops with no inputs, add to the ready to scheudle list.
    dag     = dict()
    ready   = list()
    queue   = deque([t.op for t in graph_targets])
    visited = set()
    while queue:
        op = queue.popleft()
        if op not in visited:
            visited.add(op)
            inputs = _get_parents_set(op)
            if len(inputs):
                dag[op] = inputs
                # add parents to queue in deterministc order (not python set ordering)
                queue.extend(_get_parents_list(op))
            else:
                ready.append(op)

    # Implement topological sorting found here:
    # https://en.wikipedia.org/wiki/Topological_sorting
    # Pick out AllreduceNccl ops and append them to a list in the order we'd like them scheduled.
    waves    = list()
    nccl_ops = list()
    while len(ready):
        ready_new = list()
        for ready_op in ready:
            for child_op in _get_children_list(ready_op):
                child_inputs = dag.get(child_op)
                if child_inputs is not None:
                    if ready_op in child_inputs:
                        child_inputs.remove(ready_op)
                        if len(child_inputs) == 0:
                            ready_new.append(child_op)
                            dag.pop(child_op)
                            if child_op.type == "AllreduceNccl":
                                nccl_ops.append(child_op)
        waves.append(ready)
        ready = ready_new

    if len(dag):
        raise ValueError("Error: graph_targets have at least one cycle")

    # We could serialize all ops within each wave.
    # Instead, just serialize the ops that are the inputs to the nccl ops.
    # Don't serialize the nccl ops themselves since they are async.
    # We just need them to be scheduled in a consistent order.
    prev_op = None
    for nccl_op in nccl_ops:
        if serialize_inputs:
            input_op = nccl_op.inputs[0].op
            if prev_op is not None:
                input_op._add_control_input(prev_op)
            prev_op = input_op
        else:
            if prev_op is not None:
                nccl_op._add_control_input(prev_op)
            prev_op = nccl_op

    if print_dag:
        f = open(print_dag, 'w') if type(print_dag) is str else sys.stdout
        for wave in waves:
            for op in sorted(wave, key=lambda op: (op.type, op.name)):
                print(op.type, op.name, op.outputs[0].dtype, op.outputs[0].shape, file=f)
            print("", file=f)
        if f is not sys.stdout:
            f.close()