diff --git a/python_module/CMakeLists.txt b/python_module/CMakeLists.txt index ba2539f7b0de846e23e875efc0c759ffc5be104d..25d08a7b68202de5e80b9191aefe8249911fc2b0 100644 --- a/python_module/CMakeLists.txt +++ b/python_module/CMakeLists.txt @@ -10,6 +10,9 @@ find_package(Numpy REQUIRED) find_package(SWIG REQUIRED) set(SWIG_SRC src/swig/mgb.i) set(CMAKE_SWIG_FLAGS -Wall -threads -py3 -modern -DSWIGWORDSIZE64) +if(MGB_ENABLE_JSON) + set(CMAKE_SWIG_FLAGS ${CMAKE_SWIG_FLAGS} -DMGB_ENABLE_JSON) +endif() set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-unused-parameter") if(MGE_WITH_DISTRIBUTED) diff --git a/python_module/src/swig/comp_graph.i b/python_module/src/swig/comp_graph.i index 25c58949be0234f7be7034a335b2497b1315e3ec..99154dadaeadeac95988acb02a1baa082b03d5e0 100644 --- a/python_module/src/swig/comp_graph.i +++ b/python_module/src/swig/comp_graph.i @@ -69,6 +69,15 @@ class CompGraph { return reinterpret_cast(&$self->get()); } + std::string get_dynamic_info() const { +#ifdef MGB_ENABLE_JSON + auto jsonstr = self->get().get_dynamic_info(); + return jsonstr->to_string(); +#else + return std::string(""); +#endif + } + std::string __repr__() const { auto &&graph = $self->get(); return mgb::ssprintf("", graph.id(), &graph); diff --git a/python_module/test/unit/core/test_dynamic_profiling.py b/python_module/test/unit/core/test_dynamic_profiling.py new file mode 100644 index 0000000000000000000000000000000000000000..eb7db2a12f2188e3c06b1977f741b3bf98d32494 --- /dev/null +++ b/python_module/test/unit/core/test_dynamic_profiling.py @@ -0,0 +1,53 @@ +# -*- 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. +import json + +import numpy as np + +import megengine.functional as F +from megengine import graph, tensor + + +def test_dynmaic_profiling(): + sz = 16 + + cg = graph.get_default_graph() + + x = tensor(np.arange(0, sz, dtype=np.float32)) + y = F.relu(x) + + str1 = cg.get_dynamic_info() + if str1 == "": + return + json_str1 = json.loads(str1) + + z = F.add_update(x, y) + + json_str2 = json.loads(cg.get_dynamic_info()) + + diff = lambda l1, l2: [x for x in l1 if x not in l2] + + jdiff = diff(json_str2, json_str1) + assert len(jdiff) == 1, "add_update operator should produce only one opr internally" + + dest_key = list(jdiff[0].keys())[0] + assert ( + jdiff[0][dest_key]["output"][0]["memory"] == sz * 4 + ), "output of add_update operator has wrong allocated size" + + # check add_update is inplace or not + dest_ptr = jdiff[0][dest_key]["output"][0]["dev_ptr"] + + found = False + for li in json_str1: + if "0" in li.keys(): + src_ptr = li["0"]["output"][0]["dev_ptr"] + found = dest_ptr == src_ptr + + assert found == True, "add_update is not inplace" diff --git a/src/core/impl/graph/cg_impl.cpp b/src/core/impl/graph/cg_impl.cpp index 9c34748bde10d2da48cd4b27705c02c7a046de83..590e558a2393dc0a4817154c043f15e46b77fe81 100644 --- a/src/core/impl/graph/cg_impl.cpp +++ b/src/core/impl/graph/cg_impl.cpp @@ -726,4 +726,48 @@ std::string ComputingGraphImpl::VarReceiverInfo::to_string() const { allow_empty_value); } +#if MGB_ENABLE_JSON +std::shared_ptr ComputingGraphImpl::get_dynamic_info() const { + auto make_var_json = [](VarNode* single_var) { + auto &&cur_mem_plan = single_var->mem_plan(); + if (cur_mem_plan.valid()) + return json::Object::make({ + {"name", json::String::make(single_var->name())}, + {"memory", json::Number::make(cur_mem_plan.chunk().size())}, + {"dev_ptr", json::NumberInt::make( + reinterpret_cast(single_var->dev_tensor().raw_ptr()))} + }); + else + return json::Object::make({ + {"name", json::String::make(single_var->name())}, + {"memory", json::Null::make()}, + {"dev_ptr", json::Null::make()} + }); + }; + + auto objlist = json::Array::make(); + + for(auto &opri: m_opr_refkeeper){ + auto cur_opr = opri.get(); + + auto objptr = json::Object::make(); + auto &&objbody = *objptr; + + objbody["name"] = json::String::make(cur_opr->name()); + + auto jvars = json::Array::make(); + for(auto &outputi: cur_opr->output()){ + jvars->add(make_var_json(outputi)); + } + objbody["output"] = jvars; + + auto obj = json::Object::make({{std::to_string(cur_opr->id()), objptr}}); + + objlist->add(obj); + } + + return objlist; +} +#endif // MGB_ENABLE_JSON + // vim: syntax=cpp.doxygen foldmethod=marker foldmarker=f{{{,f}}} diff --git a/src/core/impl/graph/cg_impl.h b/src/core/impl/graph/cg_impl.h index 82f0fca94e42203552fb53d8045178ac34f2e572..e99ddec393161016734c5ab0d65afa30700013ee 100644 --- a/src/core/impl/graph/cg_impl.h +++ b/src/core/impl/graph/cg_impl.h @@ -146,6 +146,10 @@ public: return m_var_receiver.at(var); } +#if MGB_ENABLE_JSON + std::shared_ptr get_dynamic_info() const override; +#endif + VarNode* find_var_by_id(size_t id) const override; TopoSorter& topo_sorter() { return components().topo_sorter; } diff --git a/src/core/include/megbrain/graph/cg.h b/src/core/include/megbrain/graph/cg.h index 963cf416c597a762bc8f3bd305cbaa9df1f6d2e2..4f83e297705a827302f1ee190529294785ce93e9 100644 --- a/src/core/include/megbrain/graph/cg.h +++ b/src/core/include/megbrain/graph/cg.h @@ -17,6 +17,10 @@ #include "megbrain/graph/seq_comp_node_opt.h" #include "megbrain/utils/event.h" +#if MGB_ENABLE_JSON +#include "megbrain/utils/json.h" +#endif + namespace mgb { namespace cg { @@ -171,6 +175,10 @@ class ComputingGraph : public std::enable_shared_from_this, virtual const VarReceiverInfo& var_receiver_in_current_comp_seq( const VarNode *var) const = 0; +#if MGB_ENABLE_JSON + virtual std::shared_ptr get_dynamic_info() const = 0; +#endif + /*! * \brief find var node by its ID *