MegEngine/python_module/megengine/_internal/comp_graph_tools.py

# -*- coding: utf-8 -*-
# MegEngine is Licensed under the Apache License, Version 2.0 (the "License")
#
# Copyright (c) 2014-2020 Megvii Inc. All rights reserved.
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

"""tools for graph manipulation"""

import collections

from . import mgb as _mgb


def get_dep_vars(var, var_type=None):
    """return :class:`.SymbolVar` of type ``var_type`` that input ``var``
    depands on. If ``var_type`` is None, return all types.

    :type var: an instance or iterable of :class:`.SymbolVar`
    :type var_type: ``str`` or an iterable of ``str``
    "rtype: list of :class:`.SymbolVar`
    """
    outputs = []
    memo = set()

    if isinstance(var, _mgb.SymbolVar):
        var = [var]

    if isinstance(var_type, str):
        var_type = [var_type]

    q = list(var)
    while q:
        v = q.pop()
        if v in memo:
            continue
        memo.add(v)
        q.extend(get_inputs(v))
        if var_type is not None:
            if get_type(v) in var_type:
                outputs.append(v)
        else:
            outputs.append(v)

    return outputs


def get_inputs(var):
    """get the inputs of owner opr of a variable

    :type var: :class:`.SymbolVar`
    :rtype: list of :class:`.SymbolVar`
    """
    assert isinstance(var, _mgb.SymbolVar)
    return _mgb._get_owner_opr_inputs(var)


def get_type(var):
    """get the type of owner opr of a variable

    :type var: :class:`.SymbolVar`
    :rtype: ``str``
    """
    assert isinstance(var, _mgb.SymbolVar)
    return _mgb._get_owner_opr_type(var)


def get_opr_type(opr):
    """get the type of a opr

    :type var: :class:`.Operator`
    :rtype: ``str``
    """
    assert isinstance(opr, _mgb.Operator)
    return _mgb._get_opr_type(opr)


def graph_traversal(outputs):
    """helper function to traverse the computing graph and reeturn enough useful information

    :param outputs: model outputs
    :type outputs: :class:`.Symbolvar`
    :return:  tuple (map_oprs, map_vars, var2oprs, opr2receivers, indegree2opr, opr2indegree)
        WHERE
        map_oprs is dict from opr_id to actual opr
        map_vars is dict from var_id to actual var
        var2oprs is dict from var to dest oprs along with index
        opr2receivers is dict from current opr to next opr
        indegree2opr is dict from in_degree to opr in computing graph
        opr2indegree is dict from opr in computing graph to in_degree

        (indegree2opr, opr2indegree) are only used in topological sort in get_oprs_seq function
    """
    # meta information for comp graph
    map_oprs = collections.defaultdict(set)
    map_vars = collections.defaultdict(set)

    var2oprs = collections.defaultdict(list)
    opr2receivers = collections.defaultdict(list)

    queue = list(map(lambda x: x.owner_opr, outputs))
    visited = set(map(lambda x: x.id, queue))

    # iterate through whole comp_graph, fill in meta information
    indegree2opr = collections.defaultdict(set)
    opr2indegree = {}

    idx = 0
    while idx < len(queue):
        cur_opr = queue[idx]
        map_oprs[cur_opr.id] = cur_opr

        idx += 1

        indegree = 0
        for var_idx, var in enumerate(cur_opr.inputs):
            map_vars[var.id] = var
            var2oprs[var.id].append((cur_opr.id, var_idx))

            pre_opr = var.owner_opr

            if pre_opr.id not in visited:
                visited.add(pre_opr.id)
                queue.append(pre_opr)

            indegree += 1
            opr2receivers[pre_opr.id].append(cur_opr.id)

        indegree2opr[indegree].add(cur_opr.id)
        opr2indegree[cur_opr.id] = indegree

    return map_oprs, map_vars, var2oprs, opr2receivers, indegree2opr, opr2indegree


def get_oprs_seq(outputs, prune_reshape=False):
    """get oprs in some topological order for a dumped model

    :param outputs: model outputs
    :param prune_reshape: whether to prune the operators useless during inference
    :return: opr list with some correct execution order
    """

    def topological_sort(map_oprs, opr2receivers, indegree2opr, opr2indegree):
        # generate an execution order with topological sort algorithm
        oprs_seq = []
        nr_remain = len(map_oprs)
        while indegree2opr[0]:
            opr_id = indegree2opr[0].pop()
            opr = map_oprs[opr_id]
            nr_remain -= 1

            # skip const value generation operator
            if get_opr_type(opr) != "ImmutableTensor":
                oprs_seq.append(opr)

            for post_id in opr2receivers[opr_id]:
                indegree = opr2indegree[post_id]
                indegree2opr[indegree].remove(post_id)

                indegree -= 1
                indegree2opr[indegree].add(post_id)
                opr2indegree[post_id] = indegree

        assert nr_remain == 0, "there are {} remaining nodes; cyclic graph?".format(
            nr_remain
        )
        return oprs_seq

    # reshape op definition: reshape(input_tensor, dest_shape) -> output_tensor
    # when inferencing, shape of output_tensor is already known, so one can prune some operators related to dest_shape in the loaded graph
    def prune_reshape_oprs(outputs, oprs_seq, var2oprs):
        def iterative_pruning(cur_opr, post_opr, marked_opr_ids):
            useless = True
            for oup in cur_opr.outputs:
                if "workspace" not in oup.name:
                    var_idx = post_opr.inputs.index(oup)
                    var2oprs[oup.id].remove((post_opr.id, var_idx))
                    useless = useless and (len(var2oprs[oup.id]) == 0)

            if useless:
                marked_opr_ids.append(cur_opr.id)

                for inp in cur_opr.inputs:
                    iterative_pruning(inp.owner_opr, cur_opr, marked_opr_ids)

        reshape_vars = get_dep_vars(outputs, "Reshape")
        reshape_oprs = [var.owner_opr for var in reshape_vars]

        marked_opr_ids = []
        for reshape_opr in reshape_oprs:
            iterative_pruning(
                reshape_opr.inputs[1].owner_opr, reshape_opr, marked_opr_ids
            )

        # filter out all marked oprs
        return list(filter(lambda x: x.id not in marked_opr_ids, oprs_seq))

    map_oprs, _, var2oprs, opr2receivers, indegree2opr, opr2indegree = graph_traversal(
        outputs
    )
    oprs_seq = topological_sort(map_oprs, opr2receivers, indegree2opr, opr2indegree)
    if prune_reshape is True:
        oprs_seq = prune_reshape_oprs(outputs, oprs_seq, var2oprs.copy())
    return oprs_seq


def replace_vars(dst, varmap):
    """replace vars in the graph

    :param dst: target vars representing the graph
    :type dst: list of :class:`.SymbolVar`
    :param varmap: the map that specifies how to replace the vars
    :type varmap: dict that maps from src var to dst var

    :return: new vars that correspond to ``dst`` with all the dependencies
        replaced
    :rtype: list of :class:`.SymbolVar`
    """
    dst_vec = _mgb._VectorSymbolVar()
    repl_src_vec = _mgb._VectorSymbolVar()
    repl_dst_vec = _mgb._VectorSymbolVar()
    for i in dst:
        assert isinstance(i, _mgb.SymbolVar)
        dst_vec.push_back(i)

    for i, j in getattr(varmap, "items", lambda: varmap)():
        assert isinstance(i, _mgb.SymbolVar)
        assert isinstance(j, _mgb.SymbolVar)
        repl_src_vec.push_back(i)
        repl_dst_vec.push_back(j)

    return _mgb._replace_vars(repl_src_vec, repl_dst_vec, dst_vec)


def replace_oprs(dst, oprmap):
    """Replace operators in the graph. Roughly equivalent to

    :param dst: target vars representing the graph
    :type dst: list of :class:`.SymbolVar`
    :param oprmap: the map that specifies how to replace the operators
    :type oprmap: dict that maps from src operator to dst operator

    :return: new vars that correspond to ``dst`` with all the dependencies
        replaced
    :rtype: list of :class:`.SymbolVar`
    """
    dst_vec = _mgb._VectorSymbolVar()
    repl_src_vec = _mgb._VectorOperator()
    repl_dst_vec = _mgb._VectorOperator()
    for i in dst:
        assert isinstance(i, _mgb.SymbolVar)
        dst_vec.push_back(i)

    for i, j in getattr(oprmap, "items", lambda: oprmap)():
        assert isinstance(i, _mgb.Operator)
        assert isinstance(j, _mgb.Operator)
        repl_src_vec.push_back(i)
        repl_dst_vec.push_back(j)

    return _mgb._replace_oprs(repl_src_vec, repl_dst_vec, dst_vec)


def set_priority_to_id(dest_vars):
    """For all oprs in the subgraph constructed by dest_vars
       set its priority to id if its original priority is zero
    :param dest_vars: target vars representing the graph
    """
    dest_vec = _mgb._VectorSymbolVar()
    for i in dest_vars:
        assert isinstance(i, _mgb.SymbolVar)
        dest_vec.push_back(i)
    _mgb._set_priority_to_id(dest_vec)