feat(mge/dump): add more options for megbrain_graph.dump_graph

GitOrigin-RevId: 848e3c87c183e716c750a2744f0dc27a50e5186a

imperative/python/megengine/core/tensor/megbrain_graph.py

@@ -11,6 +11,7 @@ import json
 import threading
 import weakref
 from concurrent.futures import Future, ThreadPoolExecutor
+from typing import Dict, List, Union
 
 import numpy as np
 
@@ -85,6 +86,97 @@ class Graph(_imperative_rt.ComputingGraph):
         return self._wrap(_imperative_rt.make_h2d(self, device, dtype, shape, name))
 
 
+class VarNode(TensorBase):
+    def __init__(self, node: _imperative_rt.VarNode):
+        self._node = node
+        if hasattr(self.graph, "_var_cache"):
+            self.graph._var_cache[node] = self
+
+    @property
+    def graph(self) -> Graph:
+        return self._node.graph
+
+    @property
+    def op(self):
+        if hasattr(self.graph, "_wrap"):
+            return self.graph._wrap(self._node.owner)
+        else:
+            return self._node.owner
+
+    @property
+    def name(self):
+        return self._node.name
+
+    @property
+    def id(self):
+        return self._node.id
+
+    @name.setter
+    def name(self, name):
+        self._node.name = name
+
+    @property
+    def dtype(self):
+        return self._node.dtype
+
+    @property
+    def device(self):
+        return as_device(self._node.comp_node)
+
+    @property
+    def shape(self):
+        return self._node.shape
+
+    @property
+    def value(self):
+        return self._node.value
+
+
+class OpNode:
+    def __init__(self, node: _imperative_rt.OperatorNode):
+        self._node = node
+        if hasattr(self.graph, "_op_cache"):
+            self.graph._op_cache[node] = self
+
+    @property
+    def graph(self) -> Graph:
+        return self._node.graph
+
+    @property
+    def name(self):
+        return self._node.name
+
+    @property
+    def id(self):
+        return self._node.id
+
+    @name.setter
+    def name(self, name):
+        self._node.name = name
+
+    @property
+    def inputs(self):
+        if hasattr(self.graph, "_wrap"):
+            return tuple(map(self.graph._wrap, self._node.inputs))
+        else:
+            return self._node.inputs
+
+    @property
+    def outputs(self):
+        if hasattr(self.graph, "_wrap"):
+            return tuple(map(self.graph._wrap, self._node.outputs))
+        else:
+            return self._node.outputs
+
+    @property
+    def params(self):
+        return json.loads(self._node.params)
+
+    @property
+    def type(self):
+        return self._node.type
+
+
 def optimize_for_inference(dest_vars, **kwargs):
     r"""Applies optimize_for_inference pass for computing graph.
 
@@ -162,8 +254,100 @@ def optimize_for_inference(dest_vars, **kwargs):
     return [VarNode(i) for i in res_vars]
 
 
-def dump_graph(*args):
-    return _imperative_rt.dump_graph([i._node for i in args])
+CompGraphDumpResult = collections.namedtuple(
+    "CompGraphDumpResult",
+    [
+        "nr_opr",
+        "tot_bytes",
+        "tensor_value_bytes",
+        "content_hash",
+        "inputs",
+        "outputs",
+        "params",
+    ],
+)
+
+
+def dump_graph(
+    output_vars: Union[Dict[str, VarNode], List[VarNode]],
+    *,
+    keep_var_name: int = 1,
+    keep_param_name: bool = False,
+    keep_opr_priority: bool = False,
+    strip_info_file=None,
+):
+    """serialize the computing graph of `output_vars` and get byte result.
+
+    :param output_vars: output variables which are the graph's end point.
+
+        .. note::
+
+            The underlying C++ API only accepts a var list. If a dict is given,
+            the vars would be renamed to the given names.
+
+    :param keep_var_name: level for keeping variable names:
+
+        * 0: none of the names are kept
+        * 1: (default)keep names of output vars
+        * 2: keep names of all (output and internal) vars
+    :param keep_param_name: whether to keep param names, so param values can be
+        easily manipulated after loading model
+    :param keep_opr_priority: whether to keep priority setting for operators
+    :param strip_info_file: a string for path or a file handler. if is not None,
+        then the dump information for code strip would be written to ``strip_info_file``
+    :return: dump result as byte string, and an instance of namedtuple
+        :class:`CompGraphDumpResult`, whose fields are:
+
+            * ``nr_opr`` number of operators dumped
+            * ``tot_bytes`` total bytes for the whole graph
+            * ``tensor_value_bytes`` bytes consumed for dumping tensor values
+            * ``inputs`` names of input tensors
+            * ``params`` list of names of dumped params
+            * ``outputs`` names of output vars
+    """
+    ov = []
+    if isinstance(output_vars, dict):
+        used_vars = set()
+        for name, var in output_vars.items():
+            assert isinstance(var, VarNode), "bad output var: {!r}".format(var)
+            assert var.id not in used_vars, (
+                "var name is associated with a var object, so we can not have "
+                "two names given to the same var: {}".format(var)
+            )
+            used_vars.add(var.id)
+            var.name = name
+            ov.append(var._node)
+    else:
+        for var in output_vars:
+            assert isinstance(var, VarNode), "bad output var: {!r}".format(var)
+            ov.append(var._node)
+
+    stat = []
+    inputs = []
+    outputs = []
+    params = []
+
+    dump_content = _imperative_rt.dump_graph(
+        ov,
+        keep_var_name,
+        keep_param_name,
+        keep_opr_priority,
+        stat,
+        inputs,
+        outputs,
+        params,
+    )
+
+    dump_info = CompGraphDumpResult(*stat, inputs, outputs, params)
+
+    if strip_info_file is not None:
+        if isinstance(strip_info_file, str):
+            strip_info_file = open(strip_info_file, "w")
+        strip_info = json.loads(_imperative_rt.get_info_for_strip(ov))
+        strip_info["hash"] = dump_info.content_hash
+        json.dump(strip_info, strip_info_file)
+
+    return dump_content, dump_info
 
 
 CompGraphLoadResult = collections.namedtuple(
@@ -193,97 +377,6 @@ def load_graph(fpath):
     return CompGraphLoadResult(cg, dict(output_vars_map), output_vars_list)
 
 
-class VarNode(TensorBase):
-    def __init__(self, node: _imperative_rt.VarNode):
-        self._node = node
-        if hasattr(self.graph, "_var_cache"):
-            self.graph._var_cache[node] = self
-
-    @property
-    def graph(self) -> Graph:
-        return self._node.graph
-
-    @property
-    def op(self):
-        if hasattr(self.graph, "_wrap"):
-            return self.graph._wrap(self._node.owner)
-        else:
-            return self._node.owner
-
-    @property
-    def name(self):
-        return self._node.name
-
-    @property
-    def id(self):
-        return self._node.id
-
-    @name.setter
-    def name(self, name):
-        self._node.name = name
-
-    @property
-    def dtype(self):
-        return self._node.dtype
-
-    @property
-    def device(self):
-        return as_device(self._node.comp_node)
-
-    @property
-    def shape(self):
-        return self._node.shape
-
-    @property
-    def value(self):
-        return self._node.value
-
-
-class OpNode:
-    def __init__(self, node: _imperative_rt.OperatorNode):
-        self._node = node
-        if hasattr(self.graph, "_op_cache"):
-            self.graph._op_cache[node] = self
-
-    @property
-    def graph(self) -> Graph:
-        return self._node.graph
-
-    @property
-    def name(self):
-        return self._node.name
-
-    @property
-    def id(self):
-        return self._node.id
-
-    @name.setter
-    def name(self, name):
-        self._node.name = name
-
-    @property
-    def inputs(self):
-        if hasattr(self.graph, "_wrap"):
-            return tuple(map(self.graph._wrap, self._node.inputs))
-        else:
-            return self._node.inputs
-
-    @property
-    def outputs(self):
-        if hasattr(self.graph, "_wrap"):
-            return tuple(map(self.graph._wrap, self._node.outputs))
-        else:
-            return self._node.outputs
-
-    @property
-    def params(self):
-        return json.loads(self._node.params)
-
-    @property
-    def type(self):
-        return self._node.type
-
-
 def _wrap(x):
     if isinstance(x, collections.abc.Sequence):
         return type(x)(map(_wrap, x))

imperative/python/megengine/jit/tracing.py

@@ -589,7 +589,9 @@ class trace:
         if isinstance(file, str):
             permission = "wb" if append == False else "ab"
             file = open(file, permission)
-        file.write(G.dump_graph(*dest_vars))
+        dump_content, dump_info = G.dump_graph(dest_vars)
+        file.write(dump_content)
+        return dump_info
 
     def _process_inputs(self, *args, **kwargs):
         if self._untraced:

imperative/python/src/graph_rt.cpp

@@ -27,6 +27,7 @@ namespace py = pybind11;
 
 using namespace mgb;
 using namespace imperative;
+namespace ser = mgb::serialization;
 
 using _OptimizeForInferenceOptions = mgb::gopt::OptimizeForInferenceOptions;
 using _LayoutTransform = _OptimizeForInferenceOptions::LayoutTransform;
@@ -183,7 +184,6 @@ void init_graph_rt(py::module m) {
             return "Opr:" + opr->name();
         });
 
-
     py::class_<cg::AsyncExecutable>(m, "AsyncExecutable")
         .def("execute", &cg::AsyncExecutable::execute, py::call_guard<py::gil_scoped_release>())
         .def("wait", &cg::AsyncExecutable::wait, py::call_guard<py::gil_scoped_release>());
@@ -206,15 +206,6 @@ void init_graph_rt(py::module m) {
                 }))
         .def("get", [](_CompGraphProfilerImpl& profiler) { return profiler._get_result(); });
 
-    m.def("dump_graph", [](const std::vector<VarNode*>& dest_vars) {
-        using namespace mgb::serialization;
-        std::vector<uint8_t> buf;
-        auto dumper = GraphDumper::make(OutputFile::make_vector_proxy(&buf));
-        SymbolVarArray symvars(dest_vars.begin(), dest_vars.end());
-        dumper->dump(symvars);
-        return py::bytes(reinterpret_cast<const char*>(&buf[0]), buf.size());
-    });
-
     auto GraphOptimizeOptions = py::class_<_OptimizeForInferenceOptions>(m, "GraphOptimizeOptions")
         .def(py::init())
         .def_readwrite("f16_io_f32_comp", &_OptimizeForInferenceOptions::f16_io_f32_comp)
@@ -245,23 +236,91 @@ void init_graph_rt(py::module m) {
         return vars;
     });
 
+    m.def("get_info_for_strip", [](const std::vector<VarNode*>& dest_vars) {
+        std::unordered_set<const char*> opr_types, dtype_names, elemwise_modes;
+        auto on_opr = [&](cg::OperatorNodeBase *opr) {
+            if (ser::GraphDumper::should_remove_in_dump(opr))
+                return;
+            opr_types.insert(opr->dyn_typeinfo()->name);
+            for (auto i : opr->output())
+                dtype_names.insert(i->dtype().name());
+            if (opr->same_type<opr::Elemwise>()) {
+                auto mode = opr->cast_final<opr::Elemwise>().param().mode;
+                elemwise_modes.insert(
+                        megdnn::Elemwise::ModeTrait::from_mode(mode).name);
+            }
+        };
+        cg::DepOprIter opr_iter{on_opr};
+        for (auto i : dest_vars)
+            opr_iter.add(i->owner_opr());
+
+        auto to_json = [](const std::unordered_set<const char*> &v) {
+            std::vector<std::string> vs(v.begin(), v.end());
+            std::sort(vs.begin(), vs.end());
+            auto ret = json::Array::make();
+            for (auto &&i : vs)
+                ret->add(json::String::make(i));
+            return ret;
+        };
+
+        return json::Object::make({
+            {"opr_types", to_json(opr_types)},
+            {"dtypes", to_json(dtype_names)},
+            {"elemwise_modes", to_json(elemwise_modes)},
+        });
+    });
+
+    m.def("dump_graph", [](
+        const std::vector<VarNode*>& dest_vars,
+        int keep_var_name,
+        bool keep_param_name,
+        bool keep_opr_priority,
+        py::list& stat,
+        py::list& inputs,
+        py::list& outputs,
+        py::list& params
+    ) {
+        std::vector<uint8_t> buf;
+        auto dumper = ser::GraphDumper::make(ser::OutputFile::make_vector_proxy(&buf));
+        SymbolVarArray symvars(dest_vars.begin(), dest_vars.end());
+
+        ser::GraphDumper::DumpConfig config{keep_var_name, keep_param_name,
+                                       keep_opr_priority};
+
+        auto rst = dumper->dump(symvars, config);
+        for (auto i : rst.inputs) {
+            inputs.append(py::cast(i));
+        }
+        for (auto i : rst.outputs) {
+            outputs.append(py::cast(i));
+        }
+        for (auto i : rst.params) {
+            params.append(py::cast(i));
+        }
+        auto rst_stat = std::vector{
+            rst.nr_opr, rst.tot_bytes, rst.tensor_value_bytes, rst.content_hash
+        };
+        for (auto i : rst_stat) {
+            stat.append(py::cast(i));
+        }
+        return py::bytes(reinterpret_cast<const char*>(&buf[0]), buf.size());
+    });
 
-    m.def("load_graph", [](std::string& buf, py::list& _output_var_map, py::list& _output_var_list) {
-        using namespace mgb::serialization;
-        auto file = InputFile::make_mem_proxy(buf.c_str(), buf.length());
-        auto format = GraphLoader::identify_graph_dump_format(*file);
-        auto loader = GraphLoader::make(std::move(file), format.val());
-        GraphLoader::LoadConfig config;
+    m.def("load_graph", [](
+        std::string& buf,
+        py::list& output_var_map,
+        py::list& output_var_list
+    ) {
+        auto file = ser::InputFile::make_mem_proxy(buf.c_str(), buf.length());
+        auto format = ser::GraphLoader::identify_graph_dump_format(*file);
+        auto loader = ser::GraphLoader::make(std::move(file), format.val());
+        ser::GraphLoader::LoadConfig config;
         auto rst = loader->load(config);
-        std::vector<std::pair<std::string, SymbolVar>> output_var_map;
-        SymbolVarArray output_var_list;
-        output_var_map = {rst.output_var_map.begin(), rst.output_var_map.end()};
-        output_var_list = std::move(rst.output_var_list);
-        for (auto i : output_var_list){
-            _output_var_list.append(i.node());
+        for (auto i : rst.output_var_map) {
+            output_var_map.append(py::make_tuple(i.first, i.second.node()));
         }
-        for (auto i : output_var_map){
-            _output_var_map.append(py::make_tuple(i.first,i.second.node()));
+        for (auto i : rst.output_var_list) {
+            output_var_list.append(i.node());
         }
         std::unordered_map<HostTensorND*, const std::string*> tensor2name;
         for (const auto& pair : rst.tensor_map) {
@@ -277,8 +336,8 @@ void init_graph_rt(py::module m) {
             h2d.output(0)->name(*it->second);
         };
         cg::DepOprIter iter{cb};
-        for (const auto& var : output_var_list) {
-            iter.add(var.node()->owner_opr());
+        for (const auto& var : rst.output_var_list) {
+            iter.add(var);
         }
         return rst.graph;
 

imperative/python/test/unit/test_tracing.py

@@ -12,12 +12,10 @@ from tempfile import mkstemp
 import numpy as np
 import pytest
 
-import megengine
-import megengine.module as M
+import megengine.core.tensor.megbrain_graph as G
 from megengine import cgtools, tensor
 from megengine.core._trace_option import set_tensor_shape
 from megengine.core.ops import builtin as ops
-from megengine.core.tensor import megbrain_graph as G
 from megengine.core.tensor.core import apply
 from megengine.core.tensor.raw_tensor import as_raw_tensor
 from megengine.functional import exp, log
@@ -121,7 +119,10 @@ def test_dump():
         np.testing.assert_equal(f(as_raw_tensor(a), as_raw_tensor(b)).numpy(), y)
 
     file = io.BytesIO()
-    f.dump(file)
+    dump_info = f.dump(file)
+    assert dump_info.nr_opr == 3
+    np.testing.assert_equal(dump_info.inputs, ["h2d[0]", "h2d[2]"])
+    np.testing.assert_equal(dump_info.outputs, ["ADD(h2d[0],h2d[2])[4]"])
     file.seek(0)
     result = load_and_inference(file, [a, b])
     np.testing.assert_equal(result[0], y)