[Pytorch Profiler] Add support for adding module hierarchy to (#61792)
authorKimish Patel <kimishpatel@fb.com>
Sat, 14 Aug 2021 04:37:57 +0000 (21:37 -0700)
committerFacebook GitHub Bot <facebook-github-bot@users.noreply.github.com>
Sat, 14 Aug 2021 04:39:10 +0000 (21:39 -0700)
commit54f2eb6e7e313d84abfe7e1f3781998979732be2
tree333c8b438d483e7b13d03ef3e719cae97ea9e0a5
parent385b082854b47cfc5c6cdadf9d07aa5f338c3c1b
[Pytorch Profiler] Add support for adding module hierarchy to (#61792)

Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/61792

KinetoEvent

This PR adds module hierarchy information to events.
What is module hierarchy information attached to events?
During profiling a TorchScript module, when events are added, we ask JIT
what is the module hierarchy associated with the node being
executed. At the time of execution of that node, there might be multiple
frames in the stack of interpreter. For each frame, we find
corresponding node and the corresponding module hierarchy is queried.
Module hierarchy corresponding to the node is associated with node's
InlinedCallStack. InlinedCallStack of node tracks the path via which the
node is inlined. Thus during the inlining process we annotate
module information corresponding to the CallMethod nodes being inlined.

With this PR, chrome trace will contain additional metadata:
"Module Hierarchy". This can look like this:
TOP(ResNet)::forward.SELF(ResNet)::_forward_impl.layer1(Sequential)::forward.0(BasicBlock)::forward.conv1(Conv2d)::forward.SELF(Conv2d)::_conv_forward
It contains module instance, type name and the method name in the
callstack.

Test Plan:
test_profiler

Imported from OSS

Reviewed By: raziel, ilia-cher

Differential Revision: D29745442

fbshipit-source-id: dc8dfaf7c5b8ab256ff0b2ef1e5ec265ca366528
13 files changed:
test/cpp/c10d/ProcessGroupNCCLTest.cpp
test/test_profiler.py
torch/_C/_autograd.pyi
torch/autograd/profiler.py
torch/autograd/profiler_legacy.py
torch/csrc/autograd/init.cpp
torch/csrc/autograd/profiler_kineto.cpp
torch/csrc/autograd/profiler_kineto.h
torch/csrc/autograd/profiler_legacy.h
torch/csrc/jit/runtime/interpreter.cpp
torch/csrc/jit/runtime/interpreter.h
torch/distributed/rpc/server_process_global_profiler.py
torch/profiler/profiler.py