Google Git
Sign in
android / platform / external / pytorch / c8b246abf3 / . / torch / csrc / jit
tree: fecfd8c16e8c4b2e5521e135933dc2eb239e6023 [path history] [tgz]
  1. batched/
  2. passes/
  3. script/
  4. argument_spec.h
  5. assertions.h
  6. attributes.h
  7. autodiff.cpp
  8. autodiff.h
  9. code_template.h
  10. constants.cpp
  11. constants.h
  12. custom_operator.h
  13. export.cpp
  14. export.h
  15. function_schema.h
  16. fusion_compiler.cpp
  17. fusion_compiler.h
  18. generic_if.h
  19. graph_executor.cpp
  20. graph_executor.h
  21. graph_node_list.h
  22. import.cpp
  23. import.h
  24. init.cpp
  25. init.h
  26. interned_strings.cpp
  27. interned_strings.h
  28. interned_strings_class.h
  29. interpreter.cpp
  30. interpreter.h
  31. ir.cpp
  32. ir.h
  33. ivalue.cpp
  34. ivalue.h
  35. named_value.h
  36. operator.cpp
  37. operator.h
  38. pybind.h
  39. pybind_utils.h
  40. python_arg_flatten.cpp
  41. python_arg_flatten.h
  42. python_interpreter.cpp
  43. python_ir.cpp
  44. python_ir.h
  45. python_tracer.cpp
  46. python_tracer.h
  47. README.md
  48. register_prim_ops.cpp
  49. register_symbols.cpp
  50. resource_guard.h
  51. serialization.h
  52. source_location.h
  53. source_range.h
  54. stack.h
  55. symbolic_variable.h
  56. test_jit.cpp
  57. tracer.cpp
  58. tracer.h
  59. type.cpp
  60. type.h
  61. variable_tensor_list.h
torch/csrc/jit/README.md

jit

The jit directory contains infrastructure for a just-in-time compiler for PyTorch and associated ‘script’ subset of python it can execute directly.

The JIT compiler has several phases.

  1. Parsing - An AST (defined in tree_views.h) is generated either by parsing a string of python-like code (jit/script/parser.h) or by translation from the Python AST (jit/frontend.py). This phase only checks for syntactic correctness and for use of the syntactic subset of python that the script supports.

  2. Semantic Checking/Specialization - We lower the AST into an IR Graph object. In this phase we check that variables are in scope and resolve any free variables to python objects. When we find free variables that are python objects, or references to non-first-class values such as modules, we temporarily represent them as SugaredValue objects. This phase then de-sugars these values by e.g. inserting a PythonOp into the graph to call a python function.

  3. Optimizations - A GraphExecutor works on an initial Graph object, performing optimizations, possibly differentiating it, and possibly specializing it to a particular size.

  4. Translation to Instructions - to execute a graph, it is lowered by the interpreter into a linear list of Instruction objects.

  5. Execution - the interpreter reads the instruction stream, executing ATen operations and any generated code fragments.

Well-known functions

Ordinarily, when defining a compiler you want the set of functions to be user extensible; e.g., a user can add to the set of defined functions by defining an appropriate autograd Function. However, there are some functions where we want to make assumptions about their semantics, because we are going to write optimizations over them or insert them into the program. Such functions are “well-known” functions, because the JIT compiler knows about them, and a user implementation must abide by the contract (sometimes implicitly) specified by the compiler.

A well-known function is usually implemented in several parts:

  • First, we pre-intern the string (interned_strings.h) that identifies the node. This allows us to more conveniently refer to these operators without having to first do a lookup through the intern table.

  • If we generate this operator during optimizations, we will often have a helper function in Graph (ir.h) for creating the operator. This is the easiest way to find out, in code, what attributes we assume for an operator.

  • There is a runtime interpretation of the operator in torch/csrc/autograd/functions/interpreter.cpp, which specifies how we actually interpret programs that contain such an operator.

So, whence the specifications! For the most part, we are following the ONNX operator specification to determine the semantics of our operators. However, there are a few other well-known functions which are specific to PyTorch.

  • FusionGroup

    A fusion group takes some number of input tensors, applies a graph Subgraph to them, producing the returned tensors of the subgraph. Operationally, operators inside a FusionGroup are fused into a single kernel, so that their intermediate results are never materialized. Not all operators support fusion:

    • attribute:
    • input: 1 - ∞ (same as inputs of Subgraph)
    • output: 1 - ∞ (same as outputs of Subgraph)