pytorch/torch/csrc/jit
Kimish Patel f4a921600a [PyTorch, Mobile] Serialization format change for source range (#54284)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/54284

In order to bring mobile deployment, via lite interpreter, on feature
parity with JIT, with respect model level debug information we must make
model level debug information available to mobile runtime.
At the moment, model level debug information is stored in SourceRange
which associates node's of graph to where the come from in original
python source code.
This information is serialized as part of debug_pkl and deserialized
when JIT loads the model and reads the model code.
On lite interpreter, we do not have access to all the functionality of
JIT and hence we cannot load model in the same way as JIT, by reading
code, constructing module hierarchy and graph corresponding module
methods etc. Instead in, lite interpreter, only bytecode corresonding to
the compiled graph, Code, is saved.
Thus in order to annotate OPs in the bytecode with equivalent
SourceRange information we do the following:
1. During model serialization, we create a unique tag for each source
range of the model.
2. Create a map of <SourceRange, tag>
3. During debug_pkl serialization we save tag along with SourceRange, on
top of byte offset.
4. During bytecode generation, the methods of the top module are
lowered. During this process methods are inlined. In the inlined graph,
when the node of a graph is lowered to bytecode, we query node's source
range and look it up against the map.
5. Resulting source range tag is serialized in module_debug_info.
6. During model deserialization, we read all the debug_pkl records in
the archieve and create a map of <tag, SourceRange>
7. This map can be used to find source code information.

During mobile runtime:
1. We read all the debug_pkl records and create <tag=debug_handle,
SourceRange> map.
   1.1 This map, MobileDebugInfo, is a member of mobile Module.
2. Interpreter catches appropriate exceptions and sets the thread local
debug handle and rethrows the exception.
3. In Function's run method we catch exception and query current debug
handle where the exception happened.
4. Query MobileDebugInfo with debug handle to retrieve source range and
augment error with source range info.

This information is still incomplete as it does not contain entire
callstack.

In the following diffs we will serialize InlinedCallStack directly.

Note that compilation is gated by SYMBOLICATE_MOBILE_DEBUG_HANDLE macro,
so that mobile builds can avoid building MobileDebugInfo, source range
and source range pickler/unpickler. Later we will add path where, if
building without debug support stack trace will contain only debug
handles. They can be symbolicated later.

Test Plan:
Ported bunch of source range tests from test_jit.py. Added on more test
in test_lite_interpreter.py

Imported from OSS

Reviewed By: raziel

Differential Revision: D27174722

fbshipit-source-id: a7b7c6088ce16dec37e823c7fefa4f0b61047e12
2021-05-04 09:19:27 -07:00
..
api Make PyTorch code-base clang-tidy compliant (#56892) 2021-04-28 14:10:25 -07:00
backends Make PyTorch code-base clang-tidy compliant (#56892) 2021-04-28 14:10:25 -07:00
codegen s/InplaceOrView/ADInplaceOrView/g (#57372) 2021-05-01 22:56:18 -07:00
cuda
docs
frontend [PyTorch, Mobile] Serialization format change for source range (#54284) 2021-05-04 09:19:27 -07:00
ir Make PyTorch code-base clang-tidy compliant (#56892) 2021-04-28 14:10:25 -07:00
mobile [PyTorch, Mobile] Serialization format change for source range (#54284) 2021-05-04 09:19:27 -07:00
passes [JIT] Enable conv-add-relu fusion as a part of frozen graph optimization (#56580) 2021-04-30 15:29:38 -07:00
python Add pybind type caster for c10::Device (#57292) 2021-05-01 16:11:10 -07:00
runtime [Static Runtime] Revamp op schema check (#57521) 2021-05-04 02:48:04 -07:00
serialization [PyTorch, Mobile] Serialization format change for source range (#54284) 2021-05-04 09:19:27 -07:00
tensorexpr [TensorExpr] Add methods for inspecting generated code in TensorExprKernel. (#57074) 2021-05-03 20:02:28 -07:00
testing Make PyTorch code-base clang-tidy compliant (#56892) 2021-04-28 14:10:25 -07:00
jit_log.cpp
jit_log.h
jit_opt_limit.cpp
jit_opt_limit.h
OVERVIEW.md
README.md
resource_guard.h

PyTorch JIT

This folder contains (most of) the C++ code for the PyTorch JIT, a language and compiler stack for executing PyTorch models portably and efficiently. To learn more about the JIT from a user perspective, please consult our reference documentation and tutorials.

A brief summary of the source tree:

  • OVERVIEW.md: High-level technical overview of the JIT.
  • frontend/: Taking PyTorch modules in Python and translating them into the JIT IR.
  • ir/: Core IR abstractions.
  • runtime/: Interpreter, graph execution, and JIT operators.
  • codegen/: Generating efficient, hardware-specific code for JIT subgraphs.
  • serialization/: Saving and loading modules.
  • api/: Any user-facing C++ or Python interfaces.
  • python/: Binding stuff into Python or accessing information from the Python environment.
  • testing/: Utilities and helpers for testing.
  • mobile/: Mobile-specific implementations of runtime components.
  • passes/: IR-to-IR passes, generally for optimization and lowering.
  • generated/: This folder is generated by the PyTorch build, and contains bindings for native PyTorch operators into the JIT.

Refer to each folder for more in-depth documentation.

Other relevant parts of the codebase not contained here:

  • aten/src/ATen/core: contains JIT code re-used by other elements of the runtime system (eager, mobile, etc.)