feat(profiler): record python api scope

GitOrigin-RevId: bd6929f33ffeca8d73f81156aee7bea2600e887c

imperative/python/megengine/traced_module/traced_module.py

@@ -290,6 +290,7 @@ def _wrap_mnode_call(orig_call):
                 active_module_tracer().current_scope()._add_input(self)
             rst = obj(*args, **kwargs)
         else:
+
             raise TypeError("'ModuleNode' object is not callable")
         return rst
 

imperative/python/megengine/utils/profiler.py

@@ -8,13 +8,16 @@ from typing import List
 from weakref import WeakSet
 
 from .. import _atexit
+from ..core._imperative_rt.core2 import Tensor as raw_tensor
 from ..core._imperative_rt.core2 import (
     cupti_available,
     disable_cupti,
     enable_cupti,
     full_sync,
     pop_scope,
+    pop_scope_with_type,
     push_scope,
+    push_scope_with_type,
     set_python_backtrace_enabled,
     start_profile,
     stop_profile,
@@ -33,6 +36,7 @@ class Profiler(ContextDecorator):
     Args:
         path: default path prefix for profiler to dump.
         with_backtrace: Whether to record backtrace information for ops.
+        with_scopes: Whether to keep more scopes to record record module/functional hierarchy. Enabling this option will slow down your program execution.
     
     Examples:
     
@@ -70,6 +74,7 @@ class Profiler(ContextDecorator):
         format: str = "chrome_timeline.json",
         formats: List[str] = None,
         with_backtrace: bool = False,
+        with_scopes: bool = False,
         **kwargs
     ) -> None:
         if not formats:
@@ -89,6 +94,8 @@ class Profiler(ContextDecorator):
             self._options[opt] = int(kwargs.pop(opt, optval))
         self._pid = "<PID>"
         self._dump_callback = None
+        self._api_patcher = None
+        self._with_scopes = with_scopes
         if self._options.get("enable_cupti", 0):
             if cupti_available():
                 enable_cupti()
@@ -110,6 +117,59 @@ class Profiler(ContextDecorator):
     def directory(self):
         return self._path
 
+    @property
+    def _patcher(self):
+        if self._api_patcher != None:
+            return self._api_patcher
+        from ..traced_module.module_tracer import Patcher, module_tracer
+        from ..module import Module
+
+        def wrap_tensormethod_and_functional(origin_fn):
+            def get_tensormeth_name(obj, func):
+                tp = obj if isinstance(obj, type) else type(obj)
+                if not issubclass(tp, raw_tensor):
+                    return None
+                for cls in tp.mro():
+                    for k, v in cls.__dict__.items():
+                        if v == func:
+                            return k
+                return None
+
+            @wraps(origin_fn)
+            def wrapped_fn(*args, **kwargs):
+                methname = (
+                    get_tensormeth_name(args[0], wrapped_fn) if len(args) > 0 else None
+                )
+                name, scope_type = (
+                    ("tensor." + methname, "tensor_method")
+                    if methname is not None
+                    else (origin_fn.__name__, "functional")
+                )
+                push_scope_with_type(name, scope_type)
+                rst = origin_fn(*args, **kwargs)
+                pop_scope_with_type(name, scope_type)
+                return rst
+
+            return wrapped_fn
+
+        def wrap_module_call(origin_fn):
+            @wraps(origin_fn)
+            def wrapped_fn(*args, **kwargs):
+                is_builtin_module = module_tracer.is_builtin(type(args[0]))
+                if not is_builtin_module:
+                    return origin_fn(*args, **kwargs)
+                name, scope_type = type(args[0]).__name__, "module"
+                push_scope_with_type(name, scope_type)
+                rst = origin_fn(*args, **kwargs)
+                pop_scope_with_type(name, scope_type)
+                return rst
+
+            return wrapped_fn
+
+        self._api_patcher = Patcher(wrap_tensormethod_and_functional)
+        self._api_patcher.patch_method(Module, "__call__", wrap_module_call)
+        return self._api_patcher
+
     @property
     def formats(self):
         return list(self._formats)
@@ -171,9 +231,14 @@ class Profiler(ContextDecorator):
 
     def __enter__(self):
         self.start()
+        if self._with_scopes:
+            self._patcher.__enter__()
 
     def __exit__(self, val, tp, trace):
         self.stop()
+        if self._with_scopes and self._api_patcher is not None:
+            self._api_patcher.__exit__(val, tp, trace)
+        self._api_patcher = None
 
     def __call__(self, func):
         func = super().__call__(func)

imperative/python/src/tensor.cpp

@@ -948,6 +948,29 @@ void init_tensor(py::module m) {
         channel->pop_scope(name);
         Transformation::pop_scope(name);
     });
+    std::unordered_map<std::string, ScopeType> str2scopetype = {
+            {"default", ScopeType::DEFAULT},
+            {"module", ScopeType::MODULE},
+            {"tensor_method", ScopeType::TENSOR_METHOD},
+            {"functional", ScopeType::FUNCTIONAL},
+            {"backward", ScopeType::BACKWARD}};
+
+    m.def("push_scope_with_type",
+          [channel, str2scopetype](std::string name, std::string type) {
+              if (str2scopetype.find(type) == str2scopetype.end()) {
+                  throw py::value_error("unsupport scope type");
+              } else {
+                  channel->push_scope(name, str2scopetype.find(type)->second);
+              }
+          });
+    m.def("pop_scope_with_type",
+          [channel, str2scopetype](std::string name, std::string type) {
+              if (str2scopetype.find(type) == str2scopetype.end()) {
+                  throw py::value_error("unsupport scope type");
+              } else {
+                  channel->pop_scope(name, str2scopetype.find(type)->second);
+              }
+          });
     m.def("start_profile", [channel](imperative::Profiler::options_t options) {
         channel->sync();
         imperative::Profiler::load_options(std::move(options));

imperative/src/impl/interpreter/commands.h

@@ -4,14 +4,14 @@
 #include <unordered_set>
 #include <variant>
 
+#include "./stack_manager.h"
+#include "./tensor_info.h"
 #include "megbrain/imperative/backtrace.h"
 #include "megbrain/imperative/op_def.h"
+#include "megbrain/imperative/profiler.h"
 #include "megbrain/imperative/utils/to_string.h"
 #include "megbrain/tensor.h"
 
-#include "./stack_manager.h"
-#include "./tensor_info.h"
-
 namespace mgb::imperative {
 
 namespace interpreter::intl {
@@ -125,6 +125,7 @@ struct StopStep {
 
 struct PushScope {
     std::string scope_name;
+    ScopeType type;
 
     template <typename TFunctor>
     void get_props(TFunctor&& functor) const {
@@ -136,7 +137,7 @@ struct PushScope {
 
 struct PopScope {
     std::string scope_name;
-
+    ScopeType type;
     template <typename TFunctor>
     void get_props(TFunctor&& functor) const {
         functor("scope_name", scope_name);

imperative/src/impl/interpreter/interpreter_impl.cpp

@@ -1356,7 +1356,7 @@ void ChannelImpl::process_one_task(Command& icmd) {
         } else if constexpr (std::is_same_v<T, StopStep>) {
             MGB_RECORD_EVENT(StopStepEvent);
         } else if constexpr (std::is_same_v<T, PushScope>) {
-            MGB_RECORD_EVENT(ScopeEvent, cmd.scope_name);
+            MGB_RECORD_EVENT(ScopeEvent, cmd.scope_name, cmd.type);
         } else if constexpr (std::is_same_v<T, PopScope>) {
             MGB_RECORD_EVENT(ScopeFinishEvent, cmd.scope_name);
         } else if constexpr (std::is_same_v<T, StartRegen>) {
@@ -1461,15 +1461,15 @@ void ChannelImpl::stop_step() {
              get_channel_state().stack_manager.dump()});
 }
 
-void ChannelImpl::push_scope(std::string name) {
+void ChannelImpl::push_scope(std::string name, ScopeType type) {
     MGB_LOCK_GUARD(m_spin);
     assert_available();
     auto& state = get_channel_state();
     state.stack_manager.enter(name);
-    MGB_RECORD_EVENT(ScopeEvent, name);
+    MGB_RECORD_EVENT(ScopeEvent, name, type);
     if (Profiler::is_profiling()) {
         m_worker.add_task(
-                {Profiler::next_id(), PushScope{name},
+                {Profiler::next_id(), PushScope{name, type},
                  get_channel_state().stack_manager.dump()});
     } else {
         m_worker.add_task({
@@ -1479,15 +1479,15 @@ void ChannelImpl::push_scope(std::string name) {
     }
 }
 
-void ChannelImpl::pop_scope(std::string name) {
+void ChannelImpl::pop_scope(std::string name, ScopeType type) {
     MGB_LOCK_GUARD(m_spin);
     assert_available();
     auto& state = get_channel_state();
     state.stack_manager.exit(name);
-    MGB_RECORD_EVENT(ScopeFinishEvent, name);
+    MGB_RECORD_EVENT(ScopeFinishEvent, name, type);
     if (Profiler::is_profiling()) {
         m_worker.add_task(
-                {Profiler::next_id(), PopScope{name},
+                {Profiler::next_id(), PopScope{name, type},
                  get_channel_state().stack_manager.dump()});
     } else {
         m_worker.add_task({

imperative/src/impl/interpreter/interpreter_impl.h

@@ -65,8 +65,8 @@ struct ChannelImpl : Interpreter::Channel, NonCopyableObj, NonMoveableObj {
     void stop_profile() override;
     void stop_step() override;
 
-    void push_scope(std::string) override;
-    void pop_scope(std::string) override;
+    void push_scope(std::string, ScopeType type = ScopeType::DEFAULT) override;
+    void pop_scope(std::string, ScopeType type = ScopeType::DEFAULT) override;
 
     BackTraceInfoPtr& get_backtrace() override;
     void set_backtrace(BackTraceInfoPtr bt) override;

imperative/src/impl/profiler/events.h

@@ -139,7 +139,10 @@ DEF_EVENT(WorkerException, {});
 
 DEF_EVENT(ShapeInfer, { bool success; });
 
-DEF_DUR_EVENT(Scope, { std::string name; });
+DEF_DUR_EVENT(Scope, {
+    std::string name;
+    ScopeType type = ScopeType::DEFAULT;
+});
 
 DEF_DUR_EVENT(Sync, { Trace trace; });
 

imperative/src/impl/profiler/memory_chunk.cpp

@@ -5,11 +5,10 @@
 #include "megbrain/imperative/utils/to_string.h"
 #include "megbrain/utils/debug.h"
 
+#include "./events.h"
 #include "./formats.h"
 #include "./states.h"
 
-#include "./events.h"
-
 namespace mgb::imperative::profiler {
 
 class XMLWriter {

imperative/src/impl/transformations/eval.cpp

@@ -135,6 +135,12 @@ ValueRefList InterpreterTransformation::apply_transformation(
         DeviceTensorND dev_tensor;
         dev_tensor.copy_from(m_channel->get_dev_tensor(input.handle()->handle()));
         return m_value_type.make(share_handle(m_channel->put(dev_tensor, {})));
+    } else if (auto push_scope = op.as<PushScope>()) {
+        m_channel->push_scope(push_scope->name, push_scope->type);
+        return {};
+    } else if (auto pop_scope = op.as<PopScope>()) {
+        m_channel->pop_scope(pop_scope->name, pop_scope->type);
+        return {};
     } else {
         return op.fallback(inputs);
     }

imperative/src/impl/transformations/grad.cpp

@@ -2,11 +2,11 @@
 
 #include <variant>
 
+#include <range/v3/all.hpp>
 #include "megbrain/imperative/graph_cache.h"
+#include "megbrain/imperative/profiler.h"
 #include "megbrain/imperative/resource_manager.h"
 
-#include <range/v3/all.hpp>
-
 namespace mgb {
 namespace imperative {
 
@@ -232,7 +232,14 @@ void GradKey::backward() {
                 for (auto&& slot : grad_fn->m_slots) {
                     *iter++ = slot.m_grad;
                 }
+                std::string name = op ? op->name() + "Backward" : "CustomBackward";
+                if (Profiler::is_profiling()) {
+                imperative::apply(PushScope(name, ScopeType::BACKWARD), Span<ValueRef>(nullptr, nullptr));
+                }
                 backward(grads, grad_receiver);
+                if (Profiler::is_profiling()) {
+                imperative::apply(PopScope(name, ScopeType::BACKWARD), Span<ValueRef>(nullptr, nullptr));
+                }
             }
         }, grad_fn->m_backward);
         // clang-format on

imperative/src/impl/transformations/lazy.cpp

@@ -172,6 +172,8 @@ ValueRefList LazyEvalTransformation::apply_transformation(
         } else {
             return imperative::apply(op, inputs);
         }
+    } else if (op.is<PushScope>() || op.is<PopScope>()) {
+        return {};
     } else {
         return op.fallback(inputs);
     }

imperative/src/include/megbrain/imperative/basic_operators.h

@@ -5,6 +5,7 @@
 
 #include "megbrain/imperative/op_def.h"
 #include "megbrain/imperative/operator.h"
+#include "megbrain/imperative/profiler.h"
 #include "megbrain/imperative/utils/data_format.h"
 #include "megbrain/imperative/utils/helper.h"
 #include "megbrain/imperative/utils/value_shape.h"
@@ -206,6 +207,24 @@ public:
     std::string to_string() const override;
 };
 
+class PushScope final : public OperatorImpl<PushScope> {
+public:
+    std::string name;
+    ScopeType type;
+    PushScope(std::string name, ScopeType type) : name{std::move(name)}, type{type} {};
+    std::string raw_type() const { return "PushScope"; }
+    std::string to_string() const override { return "PushScope"; }
+};
+
+class PopScope final : public OperatorImpl<PopScope> {
+public:
+    std::string name;
+    ScopeType type;
+    PopScope(std::string name, ScopeType type) : name{std::move(name)}, type{type} {};
+    std::string raw_type() const { return "PopScope"; }
+    std::string to_string() const override { return "PopScope"; }
+};
+
 class DupTensor final : public OperatorImpl<DupTensor, Operator::IdentityLike> {
 public:
     std::string to_string() const override { return "DupTensor"; }

imperative/src/include/megbrain/imperative/interpreter.h

@@ -5,6 +5,7 @@
 
 #include "./backtrace.h"
 #include "megbrain/imperative/op_def.h"
+#include "megbrain/imperative/profiler.h"
 
 namespace mgb::imperative::interpreter {
 
@@ -63,12 +64,13 @@ struct Interpreter {
         virtual void stop_profile() = 0;
         virtual void stop_step() = 0;
 
-        virtual void push_scope(std::string name) = 0;
-        virtual void pop_scope(std::string name) = 0;
-
         virtual BackTraceInfoPtr& get_backtrace() = 0;
         virtual void set_backtrace(BackTraceInfoPtr bt) = 0;
         virtual void clear_backtrace() = 0;
+        virtual void push_scope(
+                std::string name, ScopeType type = ScopeType::DEFAULT) = 0;
+        virtual void pop_scope(
+                std::string name, ScopeType type = ScopeType::DEFAULT) = 0;
     };
 
     virtual std::unique_ptr<Channel> create_channel() = 0;

imperative/src/include/megbrain/imperative/profiler.h

@@ -41,6 +41,14 @@ public:
     static std::shared_ptr<CompNode::Event> record_device(CompNode device);
 };
 
+enum ScopeType {
+    DEFAULT = 0,
+    MODULE = 1,
+    TENSOR_METHOD = 2,
+    FUNCTIONAL = 3,
+    BACKWARD = 4,
+};
+
 class Profiler {
 public:
     struct Record {