Skip to content

Latest commit

 

History

History
 
 

jit

Folders and files

NameName
Last commit message
Last commit date

parent directory

..
batched
batched
 
 
fuser
fuser
 
 
passes
passes
 
 
script
script
 
 
README.md
README.md
 
 
alias_info.h
alias_info.h
 
 
argument_spec.h
argument_spec.h
 
 
assertions.h
assertions.h
 
 
attributes.h
attributes.h
 
 
autodiff.cpp
autodiff.cpp
 
 
autodiff.h
autodiff.h
 
 
catch_utils.hpp
catch_utils.hpp
 
 
code_template.h
code_template.h
 
 
constants.cpp
constants.cpp
 
 
constants.h
constants.h
 
 
custom_operator.h
custom_operator.h
 
 
dynamic_dag.h
dynamic_dag.h
 
 
export.cpp
export.cpp
 
 
export.h
export.h
 
 
function_schema.h
function_schema.h
 
 
generic_if.h
generic_if.h
 
 
graph_executor.cpp
graph_executor.cpp
 
 
graph_executor.h
graph_executor.h
 
 
graph_node_list.h
graph_node_list.h
 
 
import.cpp
import.cpp
 
 
import.h
import.h
 
 
init.cpp
init.cpp
 
 
init.h
init.h
 
 
interned_strings.h
interned_strings.h
 
 
interned_strings_class.h
interned_strings_class.h
 
 
interpreter.cpp
interpreter.cpp
 
 
interpreter.h
interpreter.h
 
 
ir.cpp
ir.cpp
 
 
ir.h
ir.h
 
 
ivalue.h
ivalue.h
 
 
named_value.h
named_value.h
 
 
node_hashing.cpp
node_hashing.cpp
 
 
node_hashing.h
node_hashing.h
 
 
operator.cpp
operator.cpp
 
 
operator.h
operator.h
 
 
pybind.h
pybind.h
 
 
pybind_utils.h
pybind_utils.h
 
 
python_arg_flatten.cpp
python_arg_flatten.cpp
 
 
python_arg_flatten.h
python_arg_flatten.h
 
 
python_interpreter.cpp
python_interpreter.cpp
 
 
python_ir.cpp
python_ir.cpp
 
 
python_ir.h
python_ir.h
 
 
python_tracer.cpp
python_tracer.cpp
 
 
python_tracer.h
python_tracer.h
 
 
register_prim_ops.cpp
register_prim_ops.cpp
 
 
register_special_ops.cpp
register_special_ops.cpp
 
 
resource_guard.h
resource_guard.h
 
 
scope.cpp
scope.cpp
 
 
scope.h
scope.h
 
 
source_location.h
source_location.h
 
 
source_range.h
source_range.h
 
 
stack.h
stack.h
 
 
symbolic_variable.h
symbolic_variable.h
 
 
tracer.cpp
tracer.cpp
 
 
tracer.h
tracer.h
 
 
tracing_state.h
tracing_state.h
 
 
type.cpp
type.cpp
 
 
type.h
type.h
 
 
variable_tensor_list.h
variable_tensor_list.h
 
 

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:
      Subgraph
      The graph of fused operators. Its inputs and outputs should match the number of inputs and outputs to the FusionGroup operator.
    • input: 1 - ∞ (same as inputs of Subgraph)
    • output: 1 - ∞ (same as outputs of Subgraph)