[Pytorch Delegated Backend] Save operator name and function name in (pytorch#57441)

Summary:
Pull Request resolved: pytorch#57441

debug info

Previous diffs did not save operator name in debug info. For delegated
backends that only idenfity op for profiling with debug handle, operator
name should be stores as well.
Furthermore to complete debug informaton also serialize function name.

Test Plan:
Existing lite interpreter and backend tests

Existing lite interpreter and backend tests

Imported from OSS

Differential Revision:
D28144581
D28144581

Reviewed By: raziel

Pulled By: kimishpatel

fbshipit-source-id: 415210f147530a53b444b07f1d6ee699a3570d99

Author: kimishpatel

test/cpp/jit/test_backend.cpp

@@ -190,7 +190,7 @@ TEST(BackendTestDebugInfo, TestCompiler) {
   lm._save_for_mobile(ss, ExtraFilesMap(), true);
   auto mlm = _load_for_mobile(ss);
   std::string error_pattern = R"(
-  Module hierarchy:top(backend_with_compiler_demoLoweredModule)
+  Module hierarchy:top(backend_with_compiler_demoLoweredModule).aten::add
 Traceback of TorchScript (most recent call last):
   File "<string>", line 5, in FunctionName_UNKNOWN
                 typed_inputs: List[Any] = [x, h, ]
@@ -244,7 +244,7 @@ TEST(BackendTestDebugInfo, TestExceptionStackForCompilerWithModuleHierarchy) {
   lm._save_for_mobile(ss, ExtraFilesMap(), true);
   auto mlm = _load_for_mobile(ss);
   std::string error_pattern = R"(
-  Module hierarchy:top(backend_with_compiler_demoLoweredModule).A0(A)
+  Module hierarchy:top(backend_with_compiler_demoLoweredModule).A0(A).aten::add
 Traceback of TorchScript (most recent call last):
   File "<string>", line 5, in FunctionName_UNKNOWN
                 typed_inputs: List[Any] = [x, y, ]
@@ -337,7 +337,7 @@ TEST(
    *
    */
   std::string error_pattern = R"(
-  Module hierarchy:top(backend_with_compiler_demoLoweredModule).B0(B).A0(A)
+  Module hierarchy:top(backend_with_compiler_demoLoweredModule).B0(B).A0(A).aten::add
 Traceback of TorchScript (most recent call last):
   File "<string>", line 5, in FunctionName_UNKNOWN
                 typed_inputs: List[Any] = [x, y, ]
@@ -424,7 +424,7 @@ TEST(BackendTestDebugInfo, TestExceptionStackForCompilerWithLoweredSubModule) {
   c._save_for_mobile(ss, ExtraFilesMap(), true);
   auto c_loaded = _load_for_mobile(ss);
   std::string error_pattern = R"(
-  Module hierarchy:top(C).A0(backend_with_compiler_demoLoweredModule)
+  Module hierarchy:top(C).A0(backend_with_compiler_demoLoweredModule).aten::add
 Traceback of TorchScript (most recent call last):
   File "<string>", line 3, in FunctionName_UNKNOWN
 
@@ -545,7 +545,7 @@ TEST(
    *
    *  */
   std::string error_pattern = R"(
-  Module hierarchy:top(C).A0(backend_with_compiler_demoLoweredModule).AA0(AA)
+  Module hierarchy:top(C).A0(backend_with_compiler_demoLoweredModule).AA0(AA).aten::add
 Traceback of TorchScript (most recent call last):
   File "<string>", line 3, in FunctionName_UNKNOWN
 

test/cpp/jit/test_cs_debug_info_serialization.cpp

@@ -25,21 +25,23 @@ namespace jit {
 
 namespace {
 bool validate_debug_info(
-    const DebugInfoPair& pre_serialize,
-    const DebugInfoPair& post_serialize) {
-  auto sr1 = pre_serialize.first;
-  auto sr2 = post_serialize.first;
+    const DebugInfoTuple& pre_serialize,
+    const DebugInfoTuple& post_serialize) {
+  auto sr1 = std::get<kDebugInfoTupleSourceRangeIndex>(pre_serialize);
+  auto sr2 = std::get<kDebugInfoTupleSourceRangeIndex>(post_serialize);
   if (sr1 != sr2) {
     return false;
   }
-  if (!pre_serialize.second.defined()) {
-    return !post_serialize.second.defined();
+  auto csptr1 = std::get<kDebugInfoTupleInlinedCSIndex>(pre_serialize);
+  auto csptr2 = std::get<kDebugInfoTupleInlinedCSIndex>(post_serialize);
+  if (!csptr1.defined()) {
+    return !csptr2.defined();
   }
-  if (!post_serialize.second.defined()) {
+  if (!csptr2.defined()) {
     return false;
   }
-  auto vec1 = pre_serialize.second->vec();
-  auto vec2 = post_serialize.second->vec();
+  auto vec1 = csptr1->vec();
+  auto vec2 = csptr2->vec();
   if (vec1.size() != vec2.size()) {
     return false;
   }

test/cpp/jit/test_lite_interpreter.cpp

@@ -496,8 +496,7 @@ TEST(LiteInterpreterTest, ModuleInfoBasic) {
     }
   }
 
-  std::unordered_set<std::string> expected_result({"top(M)"});
-  AT_ASSERT(module_debug_info_set == expected_result);
+  AT_ASSERT(module_debug_info_set.count("top(M).aten::mul"));
 }
 
 // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables)
@@ -559,8 +558,9 @@ TEST(LiteInterpreterTest, OneSubmoduleModuleInfo) {
     }
   }
 
-  std::set<std::string> expected_result({"top(B)", "top(B).A0(A)"});
-  AT_ASSERT(module_debug_info_set == expected_result);
+  AT_ASSERT(module_debug_info_set.count("top(B).aten::add"));
+  AT_ASSERT(module_debug_info_set.count("top(B).A0(A).aten::add"));
+  AT_ASSERT(module_debug_info_set.count("top(B).A0(A).aten::mul"));
 }
 
 // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables)
@@ -594,7 +594,6 @@ TEST(LiteInterpreterTest, TwoSubmodulesModuleInfo) {
       std::string module_info = bc.get_forward_method_debug_info(pc);
       if (!module_info.empty() &&
           (module_info.find("debug_handle") == std::string::npos)) {
-        std::cout << "Module info:" << module_info << std::endl;
         module_debug_info_set.insert(module_info);
       }
       ++pc;
@@ -603,9 +602,9 @@ TEST(LiteInterpreterTest, TwoSubmodulesModuleInfo) {
     }
   }
 
-  std::set<std::string> expected_result(
-      {"top(C)", "top(C).A0(A)", "top(C).B0(B)"});
-  AT_ASSERT(module_debug_info_set == expected_result);
+  AT_ASSERT(module_debug_info_set.count("top(C).aten::add"));
+  AT_ASSERT(module_debug_info_set.count("top(C).A0(A).aten::add"));
+  AT_ASSERT(module_debug_info_set.count("top(C).B0(B).aten::add"));
 }
 
 TEST(LiteInterpreterTest, GetRuntimeByteCodeVersion) {
@@ -790,9 +789,9 @@ TEST(LiteInterpreterTest, SequentialModuleInfo) {
   //   def forward(self, x):
   //     return self.A0.forward(self.B0.forward(x))
 
-  std::set<std::string> expected_result(
-      {"top(C)", "top(C).A0(A)", "top(C).B0(B)"});
-  AT_ASSERT(module_debug_info_set == expected_result);
+  AT_ASSERT(module_debug_info_set.count("top(C).prim::Return"));
+  AT_ASSERT(module_debug_info_set.count("top(C).A0(A).aten::add"));
+  AT_ASSERT(module_debug_info_set.count("top(C).B0(B).aten::add"));
 }
 
 // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables)
@@ -838,9 +837,9 @@ TEST(LiteInterpreterTest, HierarchyModuleInfo) {
   // "top(C).forward": for the add operator in top.
   // "top(C).B0(B).forward": for the add operator in B0.
   // "top(C).B0(B).forward.A0(A).forward": for the add operator in A0.
-  std::set<std::string> expected_result(
-      {"top(C)", "top(C).B0(B)", "top(C).B0(B).A0(A)"});
-  AT_ASSERT(module_debug_info_set == expected_result);
+  AT_ASSERT(module_debug_info_set.count("top(C).aten::add"));
+  AT_ASSERT(module_debug_info_set.count("top(C).B0(B).aten::add"));
+  AT_ASSERT(module_debug_info_set.count("top(C).B0(B).A0(A).aten::add"));
 }
 
 // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables)
@@ -898,9 +897,9 @@ TEST(LiteInterpreterTest, DuplicatedClassTypeModuleInfo) {
   // "top(B).A0(A).forward": for the add operator in A0.
   // "top(B).A1(A).forward": for the add operator in A1.
 
-  std::set<std::string> expected_result(
-      {"top(B)", "top(B).A0(A)", "top(B).A1(A)"});
-  AT_ASSERT(module_debug_info_set == expected_result);
+  AT_ASSERT(module_debug_info_set.count("top(B).aten::add"));
+  AT_ASSERT(module_debug_info_set.count("top(B).A0(A).aten::add"));
+  AT_ASSERT(module_debug_info_set.count("top(B).A1(A).aten::add"));
 }
 
 // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables)

test/cpp/lite_interpreter_runtime/delegated_submodule_with_debug_info.ptl

@@ -142,7 +142,7 @@ TEST(RunTimeTest, DelegateException) {
   inputs.emplace_back(torch::rand({13, 9}));
 
   std::string error_pattern = R"(
-  Module hierarchy:top(C).A0(backend_with_compiler_demoLoweredModule).AA0(AA)
+  Module hierarchy:top(C).A0(backend_with_compiler_demoLoweredModule).AA0(AA).aten::add
 Traceback of TorchScript (most recent call last):
   File "<string>", line 3, in FunctionName_UNKNOWN
 

test/cpp/lite_interpreter_runtime/test_lite_interpreter_runtime.cpp

@@ -20,7 +20,7 @@ int64_t BackendDebugInfoRecorder::getNextDebugHandle(const Node* node) {
   DebugHandleType debug_handle = unique_debug_handle_;
   const SourceRange& range = node->sourceRange();
   handles_to_inlined_callstack_ptrs_[debug_handle] =
-      std::make_pair(range, cs_ptr);
+      std::make_tuple(range, node->kind().toQualString(), cs_ptr);
   // This increment is with seq memory order.
   // Not trying to perf optimizing this for now.
   unique_debug_handle_++;

torch/csrc/jit/backends/backend_debug_handler.cpp

@@ -13,22 +13,22 @@ namespace jit {
  *  BackendDebugHandleManager is responsible for issuing debug handles to
  *  backends. Debug handles are associated with nodes of a graph.
  *  BackendDebugHandleManager also maintains a map
- *  [debug-handle, DebugInfoPair = {source range, inlined callstack ptr]} that
+ *  [debug-handle, DebugInfoTuple = {source range, inlined callstack ptr]} that
  *  will help generate a callstack for exception raised using debug handles.
  *  Effectively debug handles are something that is given to backend and later
  *  when an exception occurs in the backend, backend can tell, using debug
  *  handle, that an exception occurred here. Then the runtime can generate
  *  callstack correspoding to the exception.
  *  There are two parts to BackendDebugHandleManager:
  *  1. static std::atomic debug_handle
- *  2. Map of [debug-handle, DebugInfoPair]
+ *  2. Map of [debug-handle, DebugInfoTuple]
  *
  *  About 1:
  *  Why do they have to be unique. The reason is that by ensuring
  *  uniqueness of debug handles, we remove the burden of another layer of
  *  mapping where we need to say this set of debug handles were generated for
  *  this lowered module or this bytecode function. This simplifies the API for
- *  serialization since debug handles can uniquely identify DebugInfoPair.
+ *  serialization since debug handles can uniquely identify DebugInfoTuple.
  *  Thus simplifies the runtime API for throwing exception. Exception throwing
  *  only needs to know debug_handle and not which module or method threw it.
  *  There are 2 issues to keep in mind, though,for static std::atomic
@@ -40,8 +40,8 @@ namespace jit {
  *  done.
  *
  *  Now about 2:
- *  There are two usecases for [debug-handle, DebugInfoPair]
- *  A. During bytecode generation the DebugInfoPair corresponding to the nodes
+ *  There are two usecases for [debug-handle, DebugInfoTuple]
+ *  A. During bytecode generation the DebugInfoTuple corresponding to the nodes
  *  of the inlined graph being serialized, are stored in this object and a
  *  unique debug handle is returned. This unique debug handle is stored in
  *  mobile_debug info for pytorch lite models. It will be used for raising
@@ -52,29 +52,29 @@ namespace jit {
  *  the debug handles provide a way to map nodes of the graph to the model level
  *  debug info.
  *
- *  During byte-code model serialization, [debug-handle, DebugInfoPair] is
+ *  During byte-code model serialization, [debug-handle, DebugInfoTuple] is
  *  serialized. Now we know a. debug handles and b. how to map debug handles to
  *  model source code. Thus we can either do eager symbolication by converting
  *  debug handles to corresponding source code at runtime, or do lazy
  *  symbolicattion offline.
  *
- *  Note that it is not necessary to serialize [debug-handle, DebugInfoPair]
+ *  Note that it is not necessary to serialize [debug-handle, DebugInfoTuple]
  *  corresponding to lowered backend if the lowering process, that is
  *  preprocess/compile, and execution happens in the same session, then eager
  *  symbolication can be employed.
  *
  *  Now how does BackendDebugHandleManager capture all of the above?
  *  By providing two API.
  *  1. getNextDebugHandle which given a Node* returns a unique debug handle,
- *     that will uniquely identify DebugInfoPair.
+ *     that will uniquely identify DebugInfoTuple.
  *     and
  *  2. getCallStackPtrMap which returns the map
- *     [debug-handle, DebugInfoPair]
+ *     [debug-handle, DebugInfoTuple]
  *
  *  1 provides debug handles to backends and 2 provides runtime a way to map
  *  debug handles to source level debug info.
  *
- *  So why does debug handle map to DebugInfoPair = {source range and inlined
+ *  So why does debug handle map to DebugInfoTuple = {source range and inlined
  *  cs}? {debug_handle, source_range_tag, serialized_callstack} Take this
  *  example: class L(nn.Module): def __init__(self):
  *      ...
@@ -112,7 +112,7 @@ namespace jit {
 using DebugHandleType = int64_t;
 
 using BackendDebugInfoMapType =
-    std::unordered_map<DebugHandleType, DebugInfoPair>;
+    std::unordered_map<DebugHandleType, DebugInfoTuple>;
 
 /*
  * This class is used to generate debug info map.

torch/csrc/jit/backends/backend_debug_handler.h

@@ -175,6 +175,13 @@ struct TORCH_API InlinedCallStack : public c10::intrusive_ptr_target {
   }
 };
 
-using DebugInfoPair = std::pair<SourceRange, InlinedCallStackPtr>;
+// {source range, node name, InlinedCallStack}
+// We store node name because same debug infor will be used for
+// profiling as well, so we need to know op names as well.
+using DebugInfoTuple =
+    std::tuple<SourceRange, std::string, InlinedCallStackPtr>;
+constexpr size_t kDebugInfoTupleSourceRangeIndex{0};
+constexpr size_t kDebugInfoTupleNodeNameIndex{1};
+constexpr size_t kDebugInfoTupleInlinedCSIndex{2};
 } // namespace jit
 } // namespace torch

torch/csrc/jit/ir/scope.h

@@ -14,13 +14,15 @@ namespace jit {
 namespace {
 
 std::pair<std::vector<StackEntry>, std::string> getStackTraceWithModuleHierarchy(
-    const DebugInfoPair& source_callstack) {
+    const DebugInfoTuple& source_callstack) {
   constexpr size_t kSourceRange = 1;
   constexpr size_t kModuleInstanceInfo = 2;
   std::vector<StackEntry> entries;
 
-  const SourceRange& range = source_callstack.first;
-  InlinedCallStackPtr callstack_ptr = source_callstack.second;
+  const SourceRange& range =
+      std::get<kDebugInfoTupleSourceRangeIndex>(source_callstack);
+  InlinedCallStackPtr callstack_ptr =
+      std::get<kDebugInfoTupleInlinedCSIndex>(source_callstack);
   std::string module_info;
   if (!callstack_ptr) {
     // If not cs then top level node
@@ -70,7 +72,7 @@ std::pair<std::vector<StackEntry>, std::string> getStackTraceWithModuleHierarchy
 // will be TopM(A).MyModule(B).SomeModule(C).Conv2d(conv)
 // Source level stack information will be from model source code.
 std::pair<std::string, std::string> getStackTraceWithModuleHierarchy(
-    const std::vector<DebugInfoPair>& source_callstacks,
+    const std::vector<DebugInfoTuple>& source_callstacks,
     const std::string& root_scope_string,
     const std::string& top_module_type_name) {
   std::vector<StackEntry> stack_entries;
@@ -82,6 +84,12 @@ std::pair<std::string, std::string> getStackTraceWithModuleHierarchy(
     stack_entries.insert(stack_entries.end(), entries.begin(), entries.end());
     module_info += debug_info_pair.second;
   }
+  // Only last entry in the callstack will have a node name of interest.
+  // Rest are likely CallMethod/CallFunction nodes
+  auto last_entry = source_callstacks.back();
+  const std::string& node_name =
+      std::get<kDebugInfoTupleNodeNameIndex>(last_entry);
+  module_info += "." + node_name;
   std::ostringstream ss;
   ss << "Module hierarchy:" << module_info << "\n";
   format_stack_trace(ss, stack_entries);
@@ -177,7 +185,7 @@ std::pair<std::string, std::string> MobileDebugTable::
     getSourceDebugModuleHierarchyInfo(
         const std::vector<int64_t>& debug_handles,
         const std::string& top_module_type_name) const {
-  std::vector<DebugInfoPair> debug_infos;
+  std::vector<DebugInfoTuple> debug_infos;
   bool debug_handle_not_found{false};
   for (auto it = debug_handles.rbegin(); it != debug_handles.rend(); ++it) {
     auto debug_handle = *it;

torch/csrc/jit/mobile/debug_info.cpp

@@ -40,7 +40,7 @@ class MobileDebugTable {
   std::pair<std::string, std::string> getSourceDebugModuleHierarchyInfo(
       const std::vector<int64_t>& debug_handles,
       const std::string& top_module_type_name = "ModuleTypeUnknown") const;
-  ska::flat_hash_map<int64_t, DebugInfoPair> callstack_ptr_map_;
+  ska::flat_hash_map<int64_t, DebugInfoTuple> callstack_ptr_map_;
 };
 
 } // namespace jit