mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-06 12:20:52 +01:00
0c5cd8c2b9
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/45722 This diff does a bunch of things: 1. Introduces some abstractions as detailed in https://fb.quip.com/2oEzAR5MKqbD to help with selective build related codegen in multiple files. 2. Adds helper methods to combine operators, debug info, operator lists, etc... 3. Currently, the selective build machinery querying `op_registration_whitelist` directly at various places in the code. `op_registration_whitelist` is a list of allowed operator names (without overload name). We want to move to a world where the overload names are also included so that we can be more selective about which operators we include. To that effect, it makes sense to hide the checking logic in a separate abstraction and have the build use that abstraction instead of putting all this selective build specific logic in the code-generator itself. This change is attempting to do just that. 4. Updates generate_code, unboxing-wrapper codegen, and autograd codegen to accept the operator selector paradigm as opposed to a selected operator list. 5. Update `tools/code_analyzer/gen_op_registration_allowlist.py` to expose providing an actual structured operator dependency graph in addition to a serialized string. There are a bunch of structural changes as well: 1. `root_op_list.yaml` and `combined_op_list.yaml` are now actual YAML files (not a space separated list of operator names) 2. `generate_code.py` accepts only paths to operator list YAML files (both old style as well as new style) and not list of operator names on the command line as arguments 3. `gen.py` optionally also accepts a custom build related operators YAML path (this file has information about which operators to register in the generated library). ghstack-source-id: 114578753 (Note: this ignores all push blocking failures!) Test Plan: `buck test caffe2/test:selective_build` Generated YAML files after the change: {P143981979} {P143982025} {P143982056} Ensure that the generated files are same before and after the change: ``` [dhruvbird@devvm2490 /tmp/TypeDefault.cpp] find -name "*.cpp" | xargs md5sum d72c3d125baa7b77e4c5581bbc7110d2 ./after_change/gen_aten/TypeDefault.cpp 42353036c83ebc7620a7159235b9647f ./after_change/lite_predictor_lib_aten/TypeDefault.cpp d72c3d125baa7b77e4c5581bbc7110d2 ./before_change/gen_aten/TypeDefault.cpp 42353036c83ebc7620a7159235b9647f ./before_change/lite_predictor_lib_aten/TypeDefault.cpp ``` `VariableTypes_N.cpp` are generated the same both before and after the change: ``` [dhruvbird@devvm2490 /tmp/VariableType] find -name "*.cpp" | xargs -n 1 md5sum | sort 3be89f63fd098291f01935077a60b677 ./after/VariableType_2.cpp 3be89f63fd098291f01935077a60b677 ./before/VariableType_2.cpp 40a3e59d64e9dbe86024cf314f127fd6 ./after/VariableType_4.cpp 40a3e59d64e9dbe86024cf314f127fd6 ./before/VariableType_4.cpp a4911699ceda3c3a430f08c64e8243fd ./after/VariableType_1.cpp a4911699ceda3c3a430f08c64e8243fd ./before/VariableType_1.cpp ca9aa611fcb2a573a8cba4e269468c99 ./after/VariableType_0.cpp ca9aa611fcb2a573a8cba4e269468c99 ./before/VariableType_0.cpp e18f639ed23d802dc4a31cdba40df570 ./after/VariableType_3.cpp e18f639ed23d802dc4a31cdba40df570 ./before/VariableType_3.cpp ``` Reviewed By: ljk53 Differential Revision: D23837010 fbshipit-source-id: ad06b1756af5be25baa39fd801dfdf09bc565442
90 lines
2.8 KiB
Python
90 lines
2.8 KiB
Python
import re
|
|
import os
|
|
import yaml
|
|
from .nested_dict import nested_dict
|
|
|
|
|
|
__all__ = [
|
|
'CodeTemplate', 'IDENT_REGEX', 'YamlLoader', 'nested_dict',
|
|
'split_name_params', 'write',
|
|
]
|
|
|
|
from tools.codegen.code_template import CodeTemplate
|
|
|
|
# You should use these lines, rather than doing it manually.
|
|
# Especially if you see this error!
|
|
#
|
|
# File "/usr/local/lib/python2.7/dist-packages/yaml/__init__.py", line 69, in load
|
|
# loader = Loader(stream)
|
|
# TypeError: 'module' object is not callable
|
|
try:
|
|
# use faster C loader if available
|
|
from yaml import CLoader as YamlLoader
|
|
except ImportError:
|
|
from yaml import Loader as YamlLoader
|
|
|
|
GENERATED_COMMENT = CodeTemplate(
|
|
"@" + "generated from ${filename}")
|
|
|
|
# Matches "foo" in "foo, bar" but not "foobar". Used to search for the
|
|
# occurrence of a parameter in the derivative formula
|
|
IDENT_REGEX = r'(^|\W){}($|\W)'
|
|
|
|
|
|
# TODO: Use a real parser here; this will get bamboozled
|
|
# by signatures that contain things like std::array<bool, 2> (note the space)
|
|
def split_name_params(prototype):
|
|
name, overload_name, params = re.match(r'(\w+)(\.\w+)?\((.*)\)', prototype).groups()
|
|
return name, params.split(', ')
|
|
|
|
|
|
# When tracing, we record inplace operations as out-of-place operations,
|
|
# because we don't have a story for side effects in the IR yet.
|
|
#
|
|
# Doing this un-inplacing is a little delicate however; __and__ is NOT inplace!
|
|
# TODO: Do something more robust
|
|
def uninplace_api_name(api_name):
|
|
if api_name.endswith('_') and not api_name.endswith('__'):
|
|
api_name = api_name[:-1]
|
|
if api_name.endswith('_out'):
|
|
api_name = api_name[:-4]
|
|
return api_name
|
|
|
|
|
|
def write(dirname, name, template, env):
|
|
env['generated_comment'] = GENERATED_COMMENT.substitute(filename=template.filename)
|
|
path = os.path.join(dirname, name)
|
|
# See Note [Unchanging results for ninja]
|
|
try:
|
|
with open(path, 'r') as f:
|
|
old_val = f.read()
|
|
except IOError:
|
|
old_val = None
|
|
new_val = template.substitute(env)
|
|
if old_val != new_val:
|
|
with open(path, 'w') as f:
|
|
print("Writing {}".format(path))
|
|
f.write(new_val)
|
|
else:
|
|
print("Skipped writing {}".format(path))
|
|
|
|
def is_tensor_method(declaration):
|
|
return 'Tensor' in declaration['method_of']
|
|
|
|
def is_out_variant(decl):
|
|
return decl['name'].endswith('_out')
|
|
|
|
def op_name_with_overload(decl):
|
|
return decl['operator_name_with_overload']
|
|
|
|
def load_op_list_and_strip_overload(op_list, op_list_path):
|
|
if op_list is None and op_list_path is None:
|
|
return None
|
|
if op_list is None:
|
|
op_list = []
|
|
if op_list_path is not None:
|
|
with open(op_list_path, 'r') as f:
|
|
op_list += yaml.load(f, Loader=YamlLoader)
|
|
# strip out the overload part
|
|
return {opname.split('.', 1)[0] for opname in op_list}
|