[PyTorch] Add Expanded call stack to nodes (#108426)

RihamSelim · pytorchmergebot · commit c657d9ecc555 · 2023-09-13T17:48:47.000Z
Summary: To get a Node's call stack we currently loop on the InlinedCallStack graph and follow the "callee" chain. Since the node's inlined stack does not change we can optimize this but expanding the node's inlined stack once and reusing it. This is particularly useful when reading the node's stack from another process (e.g. BPF) as it simplified the memory traversal process. The new data structure (NodeSourceInfo) only holds pointers to the function name and file name variables, and assumes these objects will be alive throughout the lifetime of the process. Each Node has an extended attribute that has an index to a vector of stack frames `expanded_node_stacks_` `node_stack_attr_symbol_` is only needed to make accessing the stack vector index attribute easier from BPF. Test Plan: - Performance Impact: The cost of expanding the call stack is between 500 - 1000 ns and happens only per instruction node at initialization time. - Verified using BPF Program in subsequent diffs Reviewed By: zdevito Differential Revision: D46578700 Pull Request resolved: #108426 Approved by: https://github.com/zdevito
diff --git a/aten/src/ATen/core/interned_strings.h b/aten/src/ATen/core/interned_strings.h
@@ -340,7 +340,8 @@ namespace c10 {
   _(attr, output_layouts)            \
   _(attr, allowzero)                 \
   _(attr, seen_none)                 \
-  _(attr, overload_name)
+  _(attr, overload_name)             \
+  _(attr, node_stack_idx)
 
 enum class _keys : unique_t {
     #define DEFINE_KEY(ns, s) ns##_##s,
diff --git a/torch/csrc/jit/ir/scope.cpp b/torch/csrc/jit/ir/scope.cpp
@@ -106,35 +106,49 @@ InlinedCallStackPtr InlinedCallStack::intrusive_from_this() {
 }
 
 InlinedCallStack::InlinedCallStack(Function* fn, SourceRange source_range)
-    : fn_(fn), source_range_(std::move(source_range)) {
-  if (fn_) {
-    set_function_name(fn_->name());
-  }
-}
+    : fn_(fn),
+      fn_name_(fn_ ? fn_->name() : ""),
+      source_range_(std::move(source_range)) {}
 
 InlinedCallStack::InlinedCallStack(
     Function* fn,
     SourceRange source_range,
     c10::optional<ModuleInstanceInfo> module_instance_info)
     : fn_(fn),
+      fn_name_(fn_ ? fn_->name() : ""),
       source_range_(std::move(source_range)),
-      module_instance_info_(std::move(module_instance_info)) {
-  if (fn_) {
-    set_function_name(fn_->name());
-  }
-}
+      module_instance_info_(std::move(module_instance_info)) {}
+
+InlinedCallStack::InlinedCallStack(
+    Function* fn,
+    SourceRange source_range,
+    c10::optional<ModuleInstanceInfo> module_instance_info,
+    std::string& function_name)
+    : fn_(fn),
+      fn_name_(std::move(function_name)),
+      source_range_(std::move(source_range)),
+      module_instance_info_(std::move(module_instance_info)) {}
 
 InlinedCallStack::InlinedCallStack(
     InlinedCallStackPtr callee,
     Function* fn,
     SourceRange source_range)
     : callee_(std::move(callee)),
       fn_(fn),
-      source_range_(std::move(source_range)) {
-  if (fn_) {
-    set_function_name(fn_->name());
-  }
-}
+      fn_name_(fn_ ? fn_->name() : ""),
+      source_range_(std::move(source_range)) {}
+
+InlinedCallStack::InlinedCallStack(
+    InlinedCallStackPtr callee,
+    Function* fn,
+    SourceRange source_range,
+    c10::optional<ModuleInstanceInfo> module_instance_info,
+    std::string& function_name)
+    : callee_(std::move(callee)),
+      fn_(fn),
+      fn_name_(std::move(function_name)),
+      source_range_(std::move(source_range)),
+      module_instance_info_(std::move(module_instance_info)) {}
 
 InlinedCallStack::InlinedCallStack(
     InlinedCallStackPtr callee,
@@ -143,12 +157,9 @@ InlinedCallStack::InlinedCallStack(
     c10::optional<ModuleInstanceInfo> module_instance_info)
     : callee_(std::move(callee)),
       fn_(fn),
+      fn_name_(fn_ ? fn_->name() : ""),
       source_range_(std::move(source_range)),
-      module_instance_info_(std::move(module_instance_info)) {
-  if (fn_) {
-    set_function_name(fn_->name());
-  }
-}
+      module_instance_info_(std::move(module_instance_info)) {}
 
 c10::optional<InlinedCallStackPtr> InlinedCallStack::callee() const {
   return callee_;
@@ -170,11 +181,7 @@ Function* InlinedCallStack::function() const {
   return fn_;
 }
 
-void InlinedCallStack::set_function_name(std::string fn_name) {
-  fn_name_ = std::move(fn_name);
-}
-
-std::string InlinedCallStack::function_name() const {
+const std::string& InlinedCallStack::function_name() const {
   return fn_name_;
 }
 
diff --git a/torch/csrc/jit/ir/scope.h b/torch/csrc/jit/ir/scope.h
@@ -134,7 +134,7 @@ struct TORCH_API InlinedCallStack : public c10::intrusive_ptr_target {
   // fn_name does not give you access to the same information that Function*
   // does, however in mobile/delegated backend runtime we use InlindedCallStack
   // for exception stack and for that purpose fn_name_ suffices.
-  std::string fn_name_;
+  const std::string fn_name_;
   SourceRange source_range_;
   InlinedCallStackPtr intrusive_from_this();
   c10::optional<ModuleInstanceInfo> module_instance_info_;
@@ -149,6 +149,13 @@ struct TORCH_API InlinedCallStack : public c10::intrusive_ptr_target {
       SourceRange source_range,
       c10::optional<ModuleInstanceInfo> module_instance_info);
 
+  // Constructor for a leaf callstack node.
+  InlinedCallStack(
+      Function* fn,
+      SourceRange source_range,
+      c10::optional<ModuleInstanceInfo> module_instance_info,
+      std::string& function_name);
+
   // Constructor for an inner callstack node.
   InlinedCallStack(
       InlinedCallStackPtr callee,
@@ -161,6 +168,13 @@ struct TORCH_API InlinedCallStack : public c10::intrusive_ptr_target {
       SourceRange source_range,
       c10::optional<ModuleInstanceInfo> module_instance_info);
 
+  InlinedCallStack(
+      InlinedCallStackPtr callee,
+      Function* fn,
+      SourceRange source_range,
+      c10::optional<ModuleInstanceInfo> module_instance_info,
+      std::string& function_name);
+
   // Return next element in the callstack list.
   c10::optional<InlinedCallStackPtr> callee() const;
 
@@ -172,9 +186,7 @@ struct TORCH_API InlinedCallStack : public c10::intrusive_ptr_target {
 
   Function* function() const;
 
-  void set_function_name(std::string fn_name);
-
-  std::string function_name() const;
+  const std::string& function_name() const;
 
   // Return callstack as a vector of [Function, SourceRange] pairs.
   std::vector<InlinedCallStackEntry> vec();
diff --git a/torch/csrc/jit/runtime/interpreter/code_impl.h b/torch/csrc/jit/runtime/interpreter/code_impl.h
@@ -59,15 +59,26 @@ struct WithCurrentNode {
   Node* old_value_;
 };
 
+struct NodeSourceInfo {
+  const char* func_name_;
+  const char* file_name_;
+  size_t line_;
+  NodeSourceInfo() : func_name_(nullptr), file_name_(nullptr), line_(0) {}
+};
+
 struct CodeImpl {
   friend struct InterpreterState;
   std::vector<Instruction> instructions_;
 
+  const c10::unique_t node_stack_attr_symbol_ =
+      static_cast<c10::unique_t>(attr::node_stack_idx);
+  // Expanded inlined stacks as pointers to values in inlined call stack.
+  std::vector<std::vector<NodeSourceInfo>> expanded_node_stacks_;
+
   // same length as instructions.
   // what node in the graph cause this
   // instruction to be emitted?
   std::vector<Node*> instructions_source_;
-
   std::vector<IValue> constant_table_;
   std::vector<Operation> operator_table_;
 #ifndef NDEBUG
@@ -213,13 +224,62 @@ struct CodeImpl {
     return instructions_source_;
   }
 
+  NodeSourceInfo getSourceInfoFromSourceRange(const SourceRange& range) {
+    NodeSourceInfo nodeSource;
+    SourceRange r = range;
+
+    if (range.source()) {
+      if (auto orig = range.source()->findSourceRangeThatGenerated(r)) {
+        r = *orig;
+      }
+    }
+    if (r.source()) {
+      auto lineno = r.source()->lineno_for_offset(r.start());
+      nodeSource.line_ = r.source()->lineno_to_source_lineno(lineno);
+      if (r.source()->filename()) {
+        nodeSource.file_name_ = r.source()->filename().value().c_str();
+      }
+    }
+    return nodeSource;
+  }
+
+  void expandInlinedNodeStack(
+      const InlinedCallStackPtr& cs,
+      std::vector<NodeSourceInfo>* expandedstack) {
+    auto nodeSourceInfo = getSourceInfoFromSourceRange(cs->source_range());
+    nodeSourceInfo.func_name_ = cs->function_name().c_str();
+    expandedstack->emplace_back(nodeSourceInfo);
+
+    if (cs->callee()) {
+      expandInlinedNodeStack(cs->callee().value(), expandedstack);
+    }
+  }
+
+  void getNodeStack(
+      const Node* node,
+      std::vector<NodeSourceInfo>* expandedstack) {
+    if (current_node_->callstack()) {
+      expandInlinedNodeStack(current_node_->callstack().value(), expandedstack);
+    }
+    auto nodeSourceInfo = getSourceInfoFromSourceRange(node->sourceRange());
+    expandedstack->emplace_back(nodeSourceInfo);
+  }
+
   void insertInstruction(OpCode op, int64_t X = 0, uint64_t N = 0) {
     instructions_.emplace_back(
         op,
         safe_narrow_cast<int32_t, int64_t>(X),
         safe_narrow_cast<uint16_t, uint64_t>(N));
     instructions_source_.emplace_back(current_node_);
 
+    if (!current_node_->hasAttribute(attr::node_stack_idx)) {
+      std::vector<NodeSourceInfo> expandedStack;
+      getNodeStack(current_node_, &expandedStack);
+      int insertIdx = expanded_node_stacks_.size();
+      expanded_node_stacks_.emplace_back(expandedStack);
+      current_node_->i_(attr::node_stack_idx, insertIdx);
+    }
+
     // check that we didn't accidentally emit nodes out of topological order
     if (op == OP) {
       if (last_inserted_op_ != nullptr && current_node_ != last_inserted_op_ &&
diff --git a/torch/csrc/jit/serialization/callstack_debug_info_serialization.cpp b/torch/csrc/jit/serialization/callstack_debug_info_serialization.cpp
@@ -155,12 +155,11 @@ InlinedCallStackPtr InlinedCallStackDeserializer::deserialize(
   InlinedCallStackPtr cs_ptr;
   if (callee) {
     cs_ptr = c10::make_intrusive<InlinedCallStack>(
-        callee, nullptr, source_range, module_instance_info);
+        callee, nullptr, source_range, module_instance_info, function_name);
   } else {
     cs_ptr = c10::make_intrusive<InlinedCallStack>(
-        nullptr, source_range, module_instance_info);
+        nullptr, source_range, module_instance_info, function_name);
   }
-  cs_ptr->set_function_name(function_name);
   cached_inlined_callstacks_[tup] = cs_ptr;
   // Invoking move constructor
   // It is not clear if copy-ellision can happen since
