| import os |
| import argparse |
| from itertools import count, combinations |
| from ..autograd.utils import CodeTemplate, write, uninplace_api_name |
| from ..autograd.gen_autograd import load_aten_declarations |
| from collections import OrderedDict |
| |
| template_path = os.path.join(os.path.dirname(__file__), 'templates') |
| |
| ATEN_DISPATCH_CPP = CodeTemplate.from_file(template_path + '/aten_dispatch.cpp') |
| ATEN_INTERNED_STRINGS_H = CodeTemplate.from_file(template_path + '/aten_interned_strings.h') |
| ATEN_SCHEMA_CPP = CodeTemplate.from_file(template_path + '/aten_schema.cpp') |
| |
| ATTR_METHOD_MAP = { |
| 'int64_t': 'i', |
| 'IntList': 'is', |
| 'Scalar': 't', |
| 'bool': 'i', |
| 'double': 'f', |
| 'std::array<bool,2>': 'is', |
| 'std::array<bool,3>': 'is', |
| 'std::array<bool,4>': 'is', |
| } |
| |
| TYPE_CASTS = { |
| 'std::array<bool,2>': 'as_bool_array<2>', |
| 'std::array<bool,3>': 'as_bool_array<3>', |
| 'std::array<bool,4>': 'as_bool_array<4>', |
| 'Scalar': 'Scalar', |
| 'IntList': 'std::vector<int64_t>', |
| } |
| |
| KW_ASSIGNMENT = CodeTemplate("""\ |
| auto ${name} = ${type_cast}(node->${method}(Symbol::attr("${name}")));\ |
| """) |
| |
| POS_ASSIGNMENT = CodeTemplate("""\ |
| auto ${name} = tensor_as<${type}>(std::move(peek(stack, ${i}, ${N})));\ |
| """) |
| |
| POS_INTLIST_ASSIGNMENT = CodeTemplate("""\ |
| auto ${name}_tensor = peek(stack, ${i}, ${N}); |
| if (${name}_tensor.dim() == 0) |
| ${name}_tensor = ${name}_tensor.expand(${size}); |
| auto ${name} = tensor_as<at::IntList>(std::move(${name}_tensor));\ |
| """) |
| |
| CALL_NAMESPACE = CodeTemplate("at::${name}(${args})") |
| CALL_METHOD = CodeTemplate("(${first}).${name}(${args})") |
| |
| CONSTRUCTOR = CodeTemplate("""\ |
| {"${descriptor}", [](Node *node) { |
| ${kw_assignments} |
| return TensorOp([=](Stack & stack) { |
| autograd::profiler::RecordFunction record("${name}"); |
| AutoGPU device_guard(deviceForInputs(stack, ${num_dynamic_inputs})); |
| ${pos_assignments} |
| auto result = ${call}; |
| drop(stack, ${num_dynamic_inputs}); |
| pack(stack, std::move(result)); |
| return 0; |
| }, "${name}", ${num_dynamic_inputs}, ${num_outputs}); |
| }}, |
| """) |
| |
| |
| def is_magic_method(api_name): |
| return api_name.startswith('__') and api_name.endswith('__') |
| |
| |
| def is_jit_op(decl): |
| # we currently only support vararg tensor lists when they are the _first_ argument |
| # and the only tensor argument |
| arguments = decl['arguments'] |
| has_tensorlist = any(arg['simple_type'] == 'TensorList' for arg in arguments) |
| num_tensor_args = sum(map(is_tensor_arg, arguments)) |
| if has_tensorlist and (num_tensor_args != 1 or arguments[0]['simple_type'] != 'TensorList'): |
| return False |
| |
| uses_tensors = any(arg['simple_type'] in {'Tensor', 'TensorList'} for arg in decl['arguments']) or \ |
| 'Tensor' in decl['method_of'] |
| return ((not decl['api_name'].endswith('_') or is_magic_method(decl['api_name'])) and |
| not decl['name'].endswith('_out') and |
| not any(arg['simple_type'] == 'Generator' for arg in decl['arguments']) and |
| not any(arg['simple_type'] == 'SparseTensor' for arg in decl['arguments']) and |
| not any(arg['simple_type'] == 'Storage' for arg in decl['arguments']) and |
| not any(arg['simple_type'] == 'ScalarType' for arg in decl['arguments']) and |
| not any(arg['simple_type'] == 'optional<ScalarType>' for arg in decl['arguments']) and |
| not any(arg['simple_type'] == 'Type' for arg in decl['arguments']) and |
| uses_tensors) |
| |
| |
| # Scalar overloads like add(Tensor self, Scalar other) are not supported atm. |
| # TODO: Why are they not supported? |
| skip_scalar_overload = { |
| 'lt-2': [1], 'gt-2': [1], 'le-2': [1], 'ge-2': [1], 'eq-2': [1], 'ne-2': [1], |
| 'pow-2': [0, 1], 'add-3': [1], 'sub-3': [1], 'mul-2': [1], 'div-2': [1], |
| 'fmod-2': [1], 'remainder-2': [1], '__and__-2': [1], '__or__-2': [1], |
| '__iand__-2': [1], '__ior__-2': [1], '__xor__-2': [1], '__ixor__-2': [1], |
| '__lshift__-2': [1], '__ilshift__-2': [1], '__rshift__-2': [1], '__irshift__-2': [1], |
| } |
| |
| |
| def is_tensor_arg(arg): |
| return arg['simple_type'] in {'Tensor', 'TensorList'} |
| |
| |
| def is_sized_intlist_arg(arg): |
| """Returns True for arguments declared as IntList[k], but False for IntList.""" |
| return (arg['simple_type'] == 'IntList') and ('size' in arg) |
| |
| |
| def gen_jit_dispatch(declarations, out): |
| ops = {} |
| |
| def get_invocation(decl, args): |
| if 'namespace' in decl['method_of']: |
| return CALL_NAMESPACE.substitute(name=decl['name'], args=args) |
| else: |
| return CALL_METHOD.substitute(name=decl['name'], first=args[0], args=args[1:]) |
| |
| def emit_decl_variant(decl, is_positional_arg, has_tensorlist): |
| # is_positional_arg is a boolean list the same length as decl['arguments'] |
| # that indicates if the argument should come from the postional list |
| # of inputs. If false, the argument comes from the constant attributes |
| kw_assignments = [] |
| attr_names = [] |
| pos_assignments = [] |
| arguments = [] |
| |
| if has_tensorlist: |
| kw_assignments.append('size_t varargs_length = node->inputs().size();') |
| # arguments look like: [tensor list], arg1, arg2, arg3 |
| # we use peek(<i>, static_inputs) to read the non-vararg inputs |
| # from the end of the stack |
| static_inputs = sum(is_positional_arg) - 1 |
| num_dynamic_inputs = 'varargs_length' |
| else: |
| static_inputs = sum(is_positional_arg) |
| num_dynamic_inputs = static_inputs |
| |
| real_inputs = count() |
| for i, arg in enumerate(decl['arguments']): |
| # XXX: we currently support only TensorList ops that have a TensorList as |
| # the first argument, that is then followed by a number of positional args. |
| if arg['simple_type'] == 'TensorList': |
| arguments.append('peekSlice(stack, 0, varargs_length - {}, varargs_length)'.format(static_inputs)) |
| elif is_tensor_arg(arg): |
| arguments.append('std::move(peek(stack, {}, {}))'.format(next(real_inputs), static_inputs)) |
| elif is_positional_arg[i]: |
| template_kwargs = dict(type=arg['simple_type'], |
| name=arg['name'], |
| i=next(real_inputs), |
| N=static_inputs) |
| |
| if is_sized_intlist_arg(arg): |
| assign = POS_INTLIST_ASSIGNMENT.substitute(size=arg['size'], |
| **template_kwargs) |
| else: |
| assign = POS_ASSIGNMENT.substitute(**template_kwargs) |
| |
| pos_assignments.append(assign) |
| arguments.append(arg['name']) |
| else: |
| simple_type = arg['simple_type'] |
| |
| attr_method = ATTR_METHOD_MAP[simple_type] |
| assign = KW_ASSIGNMENT.substitute(type_cast=TYPE_CASTS.get(simple_type, simple_type), |
| name=arg['name'], |
| method=attr_method) |
| kw_assignments.append(assign) |
| attr_names.append('{}_{}'.format(arg['name'], attr_method)) |
| arguments.append(arg['name']) |
| |
| call = get_invocation(decl, arguments) |
| |
| # Descriptor is a unique identifier for a particular overload of an op. |
| attr_names = sorted(attr_names) |
| num_inputs = '*' if has_tensorlist else static_inputs |
| descriptor = '-'.join([decl['name'], str(num_inputs)] + attr_names) |
| |
| # If there are two overloads with the same descriptor, that differ only by a type of a |
| # single argument, where one of them takes a tensor, while another one takes an |
| # at::Scalar as a positional scalar arg, then prefer the tensor overload. |
| # It should get broadcasted correctly. |
| if descriptor in skip_scalar_overload: |
| if any(decl['arguments'][idx]['simple_type'] == 'Scalar' |
| for idx in skip_scalar_overload[descriptor]): |
| return |
| |
| returns = decl['returns'] |
| all_scalars = all(r['dynamic_type'] != 'TensorList' for r in returns) |
| num_outputs = str(len(returns)) if all_scalars else 'UNKNOWN_OUTPUTS' |
| |
| constructor = CONSTRUCTOR.substitute(descriptor=descriptor, name=decl['name'], |
| call=call, |
| kw_assignments=kw_assignments, |
| pos_assignments=pos_assignments, |
| num_dynamic_inputs=num_dynamic_inputs, |
| num_outputs=num_outputs) |
| |
| assert descriptor not in ops, descriptor |
| ops[descriptor] = constructor |
| |
| def emit_decl(decl): |
| arguments = decl['arguments'] |
| has_tensorlist = any(arg['simple_type'] == 'TensorList' for arg in arguments) |
| num_tensor_args = sum(map(is_tensor_arg, arguments)) |
| |
| # Right now, we generate dispatch methods that either take all non-tensor arguments |
| # as attributes, or don't use any attributes at all. In the future we might want to |
| # have something in the middle too (might be useful for e.g. constant propagation |
| # into attributes, as that would allow us to avoid reparsing tensors into scalar |
| # args at every invocation). |
| |
| all_arguments_are_inputs = tuple(True for _ in arguments) |
| only_tensors_are_inputs = tuple(is_tensor_arg(arg) for arg in arguments) |
| |
| # NB: if there are no scalar args then both options on LHS are equivalent, so deduplicate them. |
| for variant in {all_arguments_are_inputs, only_tensors_are_inputs}: |
| emit_decl_variant(decl, variant, has_tensorlist) |
| |
| # We need to add methods implemented manually in TensorImpl |
| tensor_impl_methods = [{ |
| 'name': name, |
| 'api_name': name, |
| 'method_of': ['Tensor'], |
| 'arguments': [{'name': 'self', 'simple_type': 'Tensor'}], |
| 'returns': [{'name': 'result', 'type': 'int64_t', 'dynamic_type': 'int64_t'}], |
| } for name in ['sizes', 'strides', 'dim']] |
| aten_decls = load_aten_declarations(declarations) + tensor_impl_methods |
| jit_decls = [d for d in aten_decls if is_jit_op(d)] |
| |
| for decl in jit_decls: |
| emit_decl(decl) |
| |
| # Sort the generated snippets to ensure that the generation is deterministic |
| env = { |
| 'constructors': sorted(ops.values()), |
| } |
| write(out, 'aten_dispatch.cpp', ATEN_DISPATCH_CPP, env) |
| |
| emit_schema(jit_decls, out) |
| |
| # NB: Operate on aten_decls, not jit_decls, because VariableType is |
| # a client for these symbols as well |
| # NB: This means we DON'T generate interned strings for inplace ops. |
| # Change this when you do! |
| # NB: Keep this code synchronized with the code in |
| # tool/autograd/gen_variable_type.py |
| # NB: Some operations have inplace versions, but NOT non-inplace |
| # versions! Thus uninplace_api_name() is mandatory (if you remove |
| # it, you will get missing symbols.) |
| names = set(uninplace_api_name(decl['api_name']) for decl in aten_decls) |
| # NB: This grabs non keyword arguments too, but it's harmless |
| attrs = set(arg['name'] for decl in aten_decls for arg in decl['arguments']) |
| strings_env = { |
| 'aten_symbols': ["_(aten, {}) \\".format(n) for n in sorted(names)], |
| 'attr_symbols': ["_(attr, {}) \\".format(n) for n in sorted(attrs)] |
| } |
| write(out, 'aten_interned_strings.h', ATEN_INTERNED_STRINGS_H, strings_env) |
| |
| |
| def emit_schema(jit_decls, out): |
| |
| # see [aten_schema encoding] for how this gets translated to C++ object |
| |
| names = OrderedDict() |
| types = OrderedDict() |
| tensors = OrderedDict() |
| attributes = OrderedDict() |
| |
| env = { |
| 'arguments': [], |
| 'operators': [], |
| 'n_operators': len(jit_decls), |
| } |
| |
| # de-duplicate v strings and return the index in to d where v will occur |
| def interned(d, v): |
| v = v + ", " |
| if v not in d: |
| d[v] = len(d) |
| return d[v] |
| |
| def get_name(name): |
| return interned(names, '"{}"'.format(name)) |
| |
| def emit_arg(arg, is_return): |
| n = get_name(arg['name']) |
| if arg.get('type') == 'TensorList': |
| typ = 'ListType::ofTensors()' |
| else: |
| typ = 'DynamicType::get()' |
| tensor = 'at::nullopt' |
| attribute = 'at::nullopt' |
| if not is_return: |
| if is_tensor_arg(arg): |
| if 'default' in arg and arg['default'] == '{}': |
| tensor = 'at::Tensor()' |
| else: |
| data = 'at::nullopt' if not is_sized_intlist_arg(arg) else str(arg['size']) |
| attribute = 'AttributeInfo{{ AttributeKind::{}, {} }}'.format(ATTR_METHOD_MAP[arg['simple_type']], data) |
| if 'default' in arg: |
| value = arg['default'] |
| # conversion in yaml turns string 'true' into python bool |
| # we need it to turn into |
| value = str(value).lower() if type(value) == bool else value |
| tensor = 'as_tensor({}({}))'.format(arg['simple_type'], value) |
| d = interned(tensors, tensor) |
| a = interned(attributes, attribute) |
| t = interned(types, typ) |
| comment = '// Argument("{}", {}, {}, {})'.format(arg['name'], tensor, attribute, typ) |
| env['arguments'].append("{{ {}, {}, {}, {} }}, {} ".format(n, t, d, a, comment)) |
| |
| def emit(decl): |
| n = get_name(decl['name']) |
| for a in decl['arguments']: |
| emit_arg(a, False) |
| for a in decl['returns']: |
| emit_arg(a, True) |
| n_args = len(decl['arguments']) |
| n_returns = len(decl['returns']) |
| env['operators'].append('{{ {}, {}, {} }}, // FunctionSchema("{}", <{} arguments>, <{} returns>) '.format( |
| n, n_args, n_returns, decl['name'], n_args, n_returns)) |
| |
| for decl in jit_decls: |
| emit(decl) |
| |
| env['names'] = list(names.keys()) |
| env['tensors'] = list(tensors.keys()) |
| env['attributes'] = list(attributes.keys()) |
| env['types'] = list(types.keys()) |
| |
| write(out, 'aten_schema.cpp', ATEN_SCHEMA_CPP, env) |
| |
| |
| def main(): |
| parser = argparse.ArgumentParser( |
| description='Generate JIT op dispatch') |
| parser.add_argument('declarations', metavar='DECL', |
| help='path to Declarations.yaml') |
| parser.add_argument('out', metavar='OUT', |
| help='path to output directory') |
| args = parser.parse_args() |
| gen_jit_dispatch(args.declarations, args.out) |
| |
| |
| if __name__ == '__main__': |
| main() |
