[Pytorch Mobile] Combing instructions and debug hanles in single struct (#62418)

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

Debug handles have one to one correspondence with instruction, so just
combine them in one.

Test Plan:
CI

Imported from OSS

Reviewed By: raziel

Differential Revision: D29993661

fbshipit-source-id: 125c7163174cf66624dd95f110fdc8208fea8a07

diff --git a/torch/csrc/jit/mobile/function.cpp b/torch/csrc/jit/mobile/function.cpp
index 03d507f..0775a55 100644 (file)
--- a/torch/csrc/jit/mobile/function.cpp
+++ b/torch/csrc/jit/mobile/function.cpp
@@ -27,8 +27,8 @@ void Function::append_instruction(OpCode op, int X, int N, int64_t dbg_handle) {
       isOpSupportedInMobile(op),
       toString(op),
       " is not supported in mobile module.");
-  code_->instructions_.emplace_back(op, X, N);
-  code_->debug_handles_.emplace_back(dbg_handle);
+  code_->instructions_with_handles_.emplace_back(
+      Instruction(op, X, N), dbg_handle);
 }
 
 bool Function::append_operator(
@@ -143,9 +143,9 @@ void Function::set_register_size(size_t size) {
 int64_t Function::get_debug_handle(size_t pc) const {
   TORCH_CHECK(code_, "Valid code must exist.");
   TORCH_CHECK(
-      pc < code_->debug_handles_.size(),
+      pc < code_->instructions_with_handles_.size(),
       "Module debug info index out of boundary.");
-  return code_->debug_handles_[pc];
+  return code_->instructions_with_handles_[pc].debug_handle;
 }
 
 void Function::setSchema(c10::FunctionSchema schema) {
@@ -177,7 +177,11 @@ const std::shared_ptr<Code> Function::get_code() const {
 
 int64_t Function::getExceptionDebugHandle() const {
   size_t pc = getInterpretersExceptionPC();
-  return (pc < code_->debug_handles_.size()) ? code_->debug_handles_[pc] : -1;
+  // we dont do bounds check given that pc is obtained
+  // via internal method of getInterpretersExceptionPC
+  // which returns the PC of where the interpreter is.
+  // Although .at will do bounds check anyway.
+  return code_->instructions_with_handles_.at(pc).debug_handle;
 }
 
 } // namespace mobile

diff --git a/torch/csrc/jit/mobile/interpreter.cpp b/torch/csrc/jit/mobile/interpreter.cpp
index 98cc594..d82d84e 100644 (file)
--- a/torch/csrc/jit/mobile/interpreter.cpp
+++ b/torch/csrc/jit/mobile/interpreter.cpp
@@ -54,7 +54,7 @@ bool InterpreterState::run(Stack& stack) {
   size_t pc = 0;
   while (true) {
     try {
-      Instruction inst = code_->instructions_[pc];
+      Instruction inst = code_->instructions_with_handles_[pc].instruction;
 
       //    std::cout << "RUNNING " << pc << " " << code_->instructions_[pc];
       //    if (inst.op == OP) {

diff --git a/torch/csrc/jit/mobile/interpreter.h b/torch/csrc/jit/mobile/interpreter.h
index b89b73b..11f7f30 100644 (file)
--- a/torch/csrc/jit/mobile/interpreter.h
+++ b/torch/csrc/jit/mobile/interpreter.h
@@ -9,12 +9,18 @@ namespace jit {
 namespace mobile {
 using Stack = std::vector<c10::IValue>;
 using DebugHandle = int64_t;
+struct InstructionWithDebugHandle {
+  InstructionWithDebugHandle(Instruction inst, DebugHandle handle)
+      : instruction(inst), debug_handle(handle) {}
+  Instruction instruction;
+  DebugHandle debug_handle;
+};
+
 // NOLINTNEXTLINE(cppcoreguidelines-pro-type-member-init)
 struct Code {
   // TODO: Combine instructions and debug handles vector
   // into std::vector<<std::pair<Instruction, DebugHandle>>
-  std::vector<Instruction> instructions_;
-  std::vector<DebugHandle> debug_handles_;
+  std::vector<InstructionWithDebugHandle> instructions_with_handles_;
   std::vector<c10::OperatorName> op_names_;
   std::vector<std::function<void(Stack&)>> operators_;
   std::vector<c10::IValue> constants_;