# -*- 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)