mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 12:21:27 +01:00
efefd1d7cf
Summary: This is a series of two commits that should probably be read separately. They are stacked on top of #9018 since the second commit requires it for correctness. Commit 1 ======= This commit is the first in a series that will clean up how we handle declaring operators and intrinsics in the JIT to make it more modular and readable. This introduces readable declarations that can be used to register operators and switches gen_jit_dispatch to generate this schema. A follow up PR will remove the dispatch keys like "add-3" and resolve ops directly based on the registered schema, further simplifying the generation process. * Switches schema over to parsed declarations, in the future this will allow something like: ``` registry.register_intrinsic("foo(Tensor a, Tensor b) -> Tensor", [](Stack& stack) { ... }) ``` This will allow the scalable registration of intrinsics for lists, tuples, and other ops, as long as meta-data for these ops (e.g. derivatives and size propagation routines). The declarations resemble those used by PythonArgParser but have been singificantly cleaned up to minimize the number of types that can appear in the declaration. We should strive to get the other parts of PyTorch switched over to this restricted declaration set when possible, but it is too much to do in a single PR. My hope is that eventually we will use a very similar language to describe declarations in C10, and this can serve as a guide for that. Parsing is done using the script lexer, so it is very robust to whitespace and extensible for future types. This removes the other way we encoded schema, and makes it easier to see what schema are registered. Current generated declarations: https://gist.github.com/zdevito/a96a17766fb3a098d69a91ee00abaaf6 * Switches how we handle attempting to use an integer in the place of a fixed-sized int list, such as in conv (e.g. 'int[3] stride=1'). Now that we can statically distinguish between int and Tensor, we handle the expansion as an implicit conversion in the compiler. This allows us to simplify the interpreter since it no longer needs to handle the conversion itself. * Schema declarations have been changed so that they match the type system in the IR exactly. In particular, attribute_info which was used by liftConstantAttributes has been dropped and constant attributes are lifted purely based on the type of the input. Type conversions in compiler have been simplified due to this change. * Error highlighting in ErrorReport now only reports at most 20 lines of code, to make reading where an error occurred easier. Commit 2 ======= This commit unifies aten_dispatch and aten_schema into a single Operator object that both contains schema and implementation information. In the future we can use this object to also contain functionality like shape prop and autodiff needed by all operators. Operators are registered globally, and dispatch logic uses the schema information to figure out which variant to use. Descriptor keys, a frequent source of inscrutable debug errors, have been removed. * Introduce Operator, to replace TensorOp. Unlike TensorOp, we use Operator for all op implementations, including primitives that may occur in the graphs. The only exceptions are ops that are only known to the interpreter like jumps, and GraphExecutors where we need to record additional debug info. * Adds a global registry for Operator implementations. aten_dispatch.cpp turns into register_aten_ops.cpp, which registers all the Operators for aten with the operator registry. register_prim_ops.cpp now contains the implementations for primitive operators that used to be in the interpreter. This means that it is now safe to use `getOperation(node)` to lookup the true interpreter function for the node, which will simplify const-propagation passes. * Remove addInterpreterOpHandler in favor of global operator registry. * Instead of descriptors, we match Node arguments directly against FunctionSchema describing expected inputs in `matchSchema`. `matchSchema` knows how parse both attributes and positional inputs from a node and match it to the appropriate registered operator. Debug error messages when we try to run an invalid operator are significantly improved: they now automatically display the schema for the op with the same name that are registered. * Merge aten_schema into regsiter_aten_ops. Each Operator takes a string schema which is parsed to determine when to dispatch to that op. * Cleans up gen_jit_dispatch.py now that we do not need to write out descriptors. In particular, skip_scalar_overloads can be removed since Richard's code sorts declarations to put Tensor, Tensor declarations first. * remove matchSchemaAndLiftConstantAttributes and use emitBuiltinCall instead to remove code duplication * refactor stack manipulation functions into a separate header file. Pull Request resolved: https://github.com/pytorch/pytorch/pull/8885 Reviewed By: jamesr66a Differential Revision: D8751048 Pulled By: zdevito fbshipit-source-id: 312aabfbf88307c5f6ab947b6caf691468b94557
113 lines
4.1 KiB
Python
113 lines
4.1 KiB
Python
import argparse
|
|
import os
|
|
import sys
|
|
|
|
source_files = {'.py', '.cpp', '.h'}
|
|
|
|
DECLARATIONS_PATH = 'torch/lib/tmp_install/share/ATen/Declarations.yaml'
|
|
|
|
|
|
# TODO: This is a little inaccurate, because it will also pick
|
|
# up setup_helper scripts which don't affect code generation
|
|
def all_generator_source():
|
|
r = []
|
|
for directory, _, filenames in os.walk('tools'):
|
|
for f in filenames:
|
|
if os.path.splitext(f)[1] in source_files:
|
|
full = os.path.join(directory, f)
|
|
r.append(full)
|
|
return sorted(r)
|
|
|
|
|
|
inputs = [
|
|
'torch/lib/THNN.h',
|
|
'torch/lib/THCUNN.h',
|
|
'torch/lib/tmp_install/share/ATen/Declarations.yaml',
|
|
'tools/autograd/derivatives.yaml',
|
|
'tools/autograd/deprecated.yaml',
|
|
]
|
|
|
|
outputs = [
|
|
'torch/csrc/autograd/generated/Functions.cpp',
|
|
'torch/csrc/autograd/generated/Functions.h',
|
|
'torch/csrc/autograd/generated/python_functions.cpp',
|
|
'torch/csrc/autograd/generated/python_functions.h',
|
|
'torch/csrc/autograd/generated/python_nn_functions.cpp',
|
|
'torch/csrc/autograd/generated/python_nn_functions.h',
|
|
'torch/csrc/autograd/generated/python_nn_functions_dispatch.h',
|
|
'torch/csrc/autograd/generated/python_variable_methods.cpp',
|
|
'torch/csrc/autograd/generated/python_variable_methods_dispatch.h',
|
|
'torch/csrc/autograd/generated/variable_factories.h',
|
|
'torch/csrc/autograd/generated/VariableType.cpp',
|
|
'torch/csrc/autograd/generated/VariableType.h',
|
|
'torch/csrc/jit/generated/register_aten_ops.cpp',
|
|
'torch/csrc/jit/generated/operator.cpp',
|
|
]
|
|
|
|
|
|
def generate_code_ninja(w):
|
|
all_inputs = all_generator_source() + inputs
|
|
cmd = "{} {}".format(sys.executable, 'tools/setup_helpers/generate_code.py')
|
|
w.writer.build(
|
|
outputs, 'do_cmd', all_inputs,
|
|
variables={
|
|
'cmd': cmd,
|
|
# Note [Unchanging results for ninja]
|
|
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
|
# generate_code.py will avoid bumping the timestamp on its
|
|
# output files if the contents of the generated file did not
|
|
# change. To let Ninja take advantage of this, it must stat
|
|
# the output files after the build. See
|
|
# https://groups.google.com/forum/#!topic/ninja-build/rExDmgDL2oc
|
|
# for a more detailed discussion.
|
|
'restat': True,
|
|
})
|
|
|
|
|
|
def generate_code(ninja_global=None,
|
|
declarations_path=None,
|
|
nn_path=None,
|
|
install_dir=None):
|
|
# if ninja is enabled, we just register this file as something
|
|
# ninja will need to call if needed
|
|
if ninja_global is not None:
|
|
return generate_code_ninja(ninja_global)
|
|
|
|
# cwrap depends on pyyaml, so we can't import it earlier
|
|
root = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
|
sys.path.insert(0, root)
|
|
from tools.autograd.gen_autograd import gen_autograd
|
|
from tools.jit.gen_jit_dispatch import gen_jit_dispatch
|
|
|
|
from tools.nnwrap import generate_wrappers as generate_nn_wrappers
|
|
|
|
# Build THNN/THCUNN.cwrap and then THNN/THCUNN.cpp. These are primarily
|
|
# used by the legacy NN bindings.
|
|
generate_nn_wrappers(nn_path, install_dir, 'tools/cwrap/plugins/templates')
|
|
|
|
# Build ATen based Variable classes
|
|
autograd_gen_dir = install_dir or 'torch/csrc/autograd/generated'
|
|
jit_gen_dir = install_dir or 'torch/csrc/jit/generated'
|
|
for d in (autograd_gen_dir, jit_gen_dir):
|
|
if not os.path.exists(d):
|
|
os.makedirs(d)
|
|
gen_autograd(declarations_path or DECLARATIONS_PATH, autograd_gen_dir, 'tools/autograd')
|
|
gen_jit_dispatch(declarations_path or DECLARATIONS_PATH, jit_gen_dir, 'tools/jit/templates')
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description='Autogenerate code')
|
|
parser.add_argument('--declarations-path')
|
|
parser.add_argument('--nn-path')
|
|
parser.add_argument('--ninja-global')
|
|
parser.add_argument('--install_dir')
|
|
options = parser.parse_args()
|
|
generate_code(options.ninja_global,
|
|
options.declarations_path,
|
|
options.nn_path,
|
|
options.install_dir)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|