Fix InlinedCallStack annotation to account for module calling its own (pytorch#61791)

Summary:
Pull Request resolved: pytorch#61791

methods from forward

During inlining we attached InlinedCallstack to nodes being inlined. In
the process we attach moodule information as well, such that if
CallMethod is being inlined we know which class instance and class type
the method belongs to. However, CallMethod can be calling a method of
the same object to which the graph belongs. e.g.:

```
def forward(self, input):
  x = input + 10
  return forward_impl_(x, input)
```
Here forward_impl is method defined on the same class in which forward
is defined. Existing module hierarchy annotation will mislabel this as
unknown instance since the method is not associated with output of
GetAttr node (it would be we had called self.conv.forward_impl_ for
example).
Change in this PR reconciles this by creating a placeholder name "SELF"
for module instance indicating that you can traverse InlinedCallStack
backwards to find first node with name != SELF, which would be the name
of the object.
e.g.:
TOP(ResNet)::forward.SELF(ResNet)::_forward_impl.layer1(Sequential)::forward.0(BasicBlock)::forward.conv1(Conv2d)::forward.SELF(Conv2d)::_conv_forward

Test Plan:
Add test

Imported from OSS

Reviewed By: larryliu0820

Differential Revision: D29745443

fbshipit-source-id: 1525e41df53913341c4c36a56772454782a0ba93

Author: kimishpatel

test/cpp/jit/test_misc.cpp

@@ -2102,6 +2102,49 @@ TEST(InlinedCallStackTest, BlockAnnotation) {
   ASSERT_NE(mul_ss.str().find("return x * y"), std::string::npos);
 }
 
+// NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables)
+TEST(InlinedCallStackTest, SelfCallMethods) {
+  Module a("A");
+  a.define(R"(
+    def my_new_method(self, x):
+      return x * 3
+    def forward_impl_(self, x, y):
+      return self.my_new_method(x) + y
+    def forward(self, x, y):
+      y = y + 2
+      return self.forward_impl_(x, y)
+  )");
+  Module b("B");
+  b.define(R"(
+    def forward(self, x):
+      return x + 2
+  )");
+  Module c("C");
+  c.register_module("A0", a);
+  c.register_module("B0", b);
+  c.define(R"(
+    def call_b(self, x):
+      return self.B0.forward(x)
+    def forward(self, x, y):
+      return self.A0.forward(x, y) + self.call_b(x)
+  )");
+
+  auto graph = c.get_method("forward").function().optimized_graph();
+  std::unordered_map<std::string, size_t> module_hierarchies;
+  for (Node* n : graph->nodes()) {
+    auto hierarchy = torch::jit::utils::getNodesModuleHierarchy(*n);
+    if (module_hierarchies.count(hierarchy) == 0) {
+      module_hierarchies[hierarchy] = 0;
+    }
+    module_hierarchies[hierarchy] += 1;
+  }
+  ASSERT_EQ(module_hierarchies["A0(A)"], 2);
+  ASSERT_EQ(module_hierarchies["A0(A).SELF(A).SELF(A)"], 2);
+  ASSERT_EQ(module_hierarchies["A0(A).SELF(A)"], 1);
+  ASSERT_EQ(module_hierarchies["SELF(C)"], 1);
+  ASSERT_EQ(module_hierarchies["SELF(C).B0(B)"], 1);
+}
+
 TEST(AutogradSymbolsTest, Basic) {
   Symbol sym = Symbol::fromQualString("aten::test_symbol");
   Graph graph;

torch/csrc/jit/ir/ir.cpp

@@ -24,6 +24,29 @@
 namespace torch {
 namespace jit {
 
+namespace utils {
+std::string getNodesModuleHierarchy(const Node& n) {
+  if (!n.callstack().has_value()) {
+    return std::string();
+  }
+  InlinedCallStackPtr callstack_ptr = n.callstack().value();
+  std::string module_hierarchy;
+  for (auto& entry : callstack_ptr->vec()) {
+    const auto& opt_module_info = std::get<kModuleInstanceInfo>(entry);
+    if (opt_module_info.has_value()) {
+      const auto& module_instance_info = opt_module_info.value();
+      if (!module_hierarchy.empty()) {
+        module_hierarchy.append(".");
+      }
+      module_hierarchy.append(utils::get_module_info(module_instance_info));
+    } else {
+      module_hierarchy += ".UNKNOWN_INSTANCE(UNKNOWN_TYPE)";
+    }
+  }
+  return module_hierarchy;
+}
+} // namespace utils
+
 namespace {
 
 // Constants relating to maintaining the topological index of nodes.
@@ -2059,10 +2082,18 @@ std::vector<Value*> inlineCallTo(
     if (to_replace->input(0)->node()->kind() == prim::GetAttr) {
       module_instance_info = c10::make_optional(ModuleInstanceInfo(
           class_type_ptr, to_replace->input(0)->node()->s(attr::name)));
+    } else if (
+        to_replace->owningGraph()->inputs().size() > 0 &&
+        to_replace->input(0) == to_replace->owningGraph()->inputs()[0]) {
+      // This CallMethod must correspond to method of the same object
+      // to which this graph belongs.
+      module_instance_info =
+          c10::make_optional(ModuleInstanceInfo(class_type_ptr, "SELF"));
     } else {
-      std::string instance_name_unknown("INSTANCE_NAME_UNKNOWN");
+      // Not sure if it is possible to come here ever.
+      // TODO: Remove this else. Or add assert
       module_instance_info = c10::make_optional(
-          ModuleInstanceInfo(class_type_ptr, instance_name_unknown));
+          ModuleInstanceInfo(class_type_ptr, "INSTANCE_NAME_UNKNOWN"));
     }
   }
 

torch/csrc/jit/ir/ir.h

@@ -32,6 +32,9 @@ using pyobj_list = std::vector<THPObjectPtr>;
 
 namespace torch {
 namespace jit {
+namespace utils {
+TORCH_API std::string getNodesModuleHierarchy(const Node& n);
+} // namespace utils
 class AliasDb;
 
 using ::c10::Argument;

torch/csrc/jit/ir/scope.h

@@ -10,9 +10,10 @@
 namespace torch {
 namespace jit {
 struct ModuleInstanceInfo;
+constexpr size_t kModuleInstanceInfo = 2;
+
 namespace utils {
-TORCH_API std::string get_module_info(
-    const ModuleInstanceInfo& module_instance_info);
+std::string get_module_info(const ModuleInstanceInfo& module_instance_info);
 } // namespace utils
 
 // Scope is a node of a trie that represents the tree of nested scopes.

torch/csrc/jit/mobile/debug_info.cpp

@@ -16,10 +16,6 @@ namespace {
 std::pair<std::vector<StackEntry>, std::string> getStackTraceWithModuleHierarchy(
     const DebugInfoTuple& source_callstack,
     const std::string& caller_name) {
-  // NOLINTNEXTLINE(clang-diagnostic-unused-variable)
-  constexpr size_t kSourceRange = 1;
-  // NOLINTNEXTLINE(clang-diagnostic-unused-variable)
-  constexpr size_t kModuleInstanceInfo = 2;
   std::vector<StackEntry> entries;
 
   const SourceRange& range =

torch/csrc/jit/python/python_ir.cpp

@@ -569,28 +569,7 @@ void initPythonIRBindings(PyObject* module_) {
       .def("output", [](Node& n) { return n.output(); })
       .def(
           "getModuleHierarchy",
-          [](Node& n) {
-            if (!n.callstack().has_value()) {
-              return std::string();
-            }
-            InlinedCallStackPtr callstack_ptr = n.callstack().value();
-            std::string module_info;
-            for (auto& entry : callstack_ptr->vec()) {
-              const auto& opt_module_info =
-                  std::get<kModuleInstanceInfo>(entry);
-              if (opt_module_info.has_value()) {
-                const auto& module_instance_info = opt_module_info.value();
-                if (!module_info.empty()) {
-                  module_info.append(".");
-                }
-                module_info.append(
-                    utils::get_module_info(module_instance_info));
-              } else {
-                module_info += ".UNKNOWN_INSTANCE(UNKNOWN_TYPE)";
-              }
-            }
-            return module_info;
-          })
+          [](Node& n) { return torch::jit::utils::getNodesModuleHierarchy(n); })
       .NS(addInput)
       .NS(replaceInput)
       .NS(replaceInputWith)

torch/csrc/jit/python/python_ir.h

@@ -6,8 +6,6 @@
 namespace torch {
 namespace jit {
 
-constexpr size_t kModuleInstanceInfo = 2;
-
 void initPythonIRBindings(PyObject* module);
 
 // execute a Python function, used for Ops we can't optimize but that we want to