| import re |
| from copy import deepcopy |
| from function_wrapper import TYPE_FORMAL_GENERIC |
| import common_with_cwrap |
| |
| type_map = { |
| 'floating_point': [ |
| 'Float', |
| 'Double', |
| 'Half', |
| ], |
| 'integral': [ |
| 'Byte', |
| 'Char', |
| 'Short', |
| 'Int', |
| 'Long' |
| ], |
| } |
| |
| all_types = type_map['floating_point'] + type_map['integral'] |
| type_map['all'] = all_types |
| |
| all_backends = ['CPU', 'CUDA', 'SparseCPU', 'SparseCUDA'] |
| default_backends = ['CPU', 'CUDA'] |
| |
| |
| def process_types_and_backends(option): |
| # if specific pairs were not listed, then enumerate them |
| # based on the backend and type attributes |
| # if backend or type is not defined, it is assumed to be all of them |
| if 'backend_type_pairs' not in option: |
| backends = option.get('backends', default_backends) |
| if option.get('aten_sparse', False): |
| backends = all_backends |
| |
| types = option.get('types', all_types) |
| |
| pairs = [[p, t] for p in backends for t in types] |
| else: |
| pairs = option['backend_type_pairs'] |
| |
| # expand type alias (integral, floating_point, all) |
| def expand(pair): |
| p, t = pair |
| assert(p in all_backends) |
| if t in type_map: |
| return [(p, tt) for tt in type_map[t]] |
| assert(t in all_types) |
| return [(p, t)] |
| pairs = set(p for pair in pairs for p in expand(pair)) |
| |
| # disable CUDA Half if there is a Sparse argument |
| for arg in option.get('arguments', []): |
| if arg['type'] == 'THSTensor*': |
| pairs.discard(('CUDA', 'Half')) |
| |
| # special case remove Half for cpu unless it is explicitly enabled, |
| if not option.get('cpu_half', False): |
| pairs.discard(('CPU', 'Half')) |
| |
| # sort the result for easy reading |
| option['backend_type_pairs'] = sorted([p for p in pairs]) |
| |
| |
| def exclude(declaration): |
| return 'only_register' in declaration or declaration.get('python_name') == 'ndimension' |
| |
| |
| def add_variants(option): |
| option.setdefault('variants', ['method']) |
| |
| # if we have 'output' arguments, generate a variant where |
| # we mark oututs as allocate = True, and where the method variant |
| # is disabled... |
| |
| |
| def handle_outputs_taken_as_arguments(options): |
| new_options = [] |
| |
| def is_nullable(arg): |
| return (arg['type'] in {'THIntegerTensor*', 'THTensor*'} and |
| arg.get('default', '') in {'NULL', 'nullptr'}) |
| |
| for option in options: |
| for arg in option['arguments']: |
| # mark arguments which can be null |
| if is_nullable(arg): |
| arg['is_nullable'] = True |
| |
| if any('output' in arg for arg in option['arguments']): |
| allocate_option = deepcopy(option) |
| # the allocating option needs to be marked |
| for arg in allocate_option['arguments']: |
| if 'output' in arg: |
| arg['allocate'] = True |
| |
| # the original option, which takes arguments for the results, |
| # is no longer a method, and has _out added to indicte it takes |
| # output arguments |
| if 'function' in option['variants']: |
| if 'method' in option['variants']: |
| option['variants'].remove('method') |
| option['api_name'] += '_out' |
| new_options.append(option) |
| |
| new_options.append(allocate_option) |
| else: |
| new_options.append(option) |
| return new_options |
| |
| |
| def sanitize_return(option): |
| ret = option['return'] |
| m = re.match('argument (\d+(,\d+)*)', ret) |
| if m is not None: |
| arguments = [int(x) for x in m.group(1).split(',')] |
| option['return'] = {'kind': 'arguments', 'arguments': arguments} |
| elif ret == 'self': |
| option['return'] = {'kind': 'arguments', 'arguments': []} |
| for i, x in enumerate(option['arguments']): |
| if x['name'] == 'self': |
| option['return']['arguments'].append(i) |
| break |
| else: |
| option['return'] = {'kind': 'type', 'type': option['return']} |
| |
| |
| def set_mode(option): |
| option['mode'] = option.get('mode', 'TH') |
| |
| # To enable 0-dim support in TH operations |
| # we find all places where a single Scalar replaced with a Tensor |
| # as an argument is still a valid function |
| # we then mark the tensor variant with a key zero_dim_dispatch_when_scalar: name |
| # where 'name' is the name of the argument that should be a scalar |
| # during dispatch, if that argument is marked internally as holding a scalar |
| # then the method will dispatch to that function. |
| |
| |
| def discover_zero_dim_tensor_operations(declaration): |
| def exclude(arg): |
| return arg.get('ignore_check') |
| |
| def signature(option, i=None, value=None): |
| elements = [TYPE_FORMAL_GENERIC.get(arg['type'], arg['type']) |
| if i is None or j != i else value |
| for j, arg in enumerate(option['arguments']) |
| if not exclude(arg)] |
| return '#'.join(elements) |
| signature_to_option = {signature(option): option |
| for option in declaration['options']} |
| |
| for option in declaration['options']: |
| for i, arg in enumerate(option['arguments']): |
| if arg['type'] == 'real': |
| signature_of_tensor_version = signature(option, i, 'Tensor &') |
| if signature_of_tensor_version in signature_to_option: |
| tensor_version = \ |
| signature_to_option[signature_of_tensor_version] |
| names = [arg['name'] for arg in tensor_version['arguments'] |
| if not exclude(arg)] |
| tensor_version['zero_dim_dispatch_when_scalar'] = names[i] |
| # print("FOUND "+str(i) ) |
| # print("Scalar Version ===== ") |
| # print(yaml.dump(option)) |
| # print("Tensor Version ===== ") |
| # print(yaml.dump(tensor_version)) |
| # print("SHARED "+names[i]) |
| |
| |
| def discover_sparse_tensor_operations(declaration): |
| def exclude(arg): |
| return arg.get('ignore_check') |
| |
| def signature(option, i=None, value=None): |
| elements = [TYPE_FORMAL_GENERIC.get(arg['type'], arg['type']) |
| if i is None or j != i else value |
| for j, arg in enumerate(option['arguments']) |
| if not exclude(arg)] |
| return '#'.join(elements) |
| |
| name = declaration['name'] |
| if name == 'add' or name == 'add_': |
| signature_to_option = {signature(option): option |
| for option in declaration['options']} |
| |
| for option in declaration['options']: |
| for i, arg in enumerate(option['arguments']): |
| if (arg['type'] == 'THSTensor*' and |
| option.get('aten_dense_sparse', False)): |
| signature_of_tensor_version = signature( |
| option, i, 'Tensor &') |
| if signature_of_tensor_version in signature_to_option: |
| tensor_version = \ |
| signature_to_option[signature_of_tensor_version] |
| raw_args = len(tensor_version['arguments']) |
| names = [arg['name'] for arg in tensor_version['arguments'] |
| if not exclude(arg)] |
| filtered_args = len(names) |
| tensor_version['when_sparse_dispatch'] = names[i - |
| (raw_args - filtered_args)] |
| |
| |
| def run(declarations): |
| declarations = [d for d in declarations if not exclude(d)] |
| for declaration in declarations: |
| common_with_cwrap.set_declaration_defaults(declaration) |
| declaration['options'] = [deepcopy(o) for o in declaration['options']] |
| declaration['options'] = common_with_cwrap.filter_unique_options( |
| declaration['options'], |
| allow_kwarg=False, |
| type_to_signature=TYPE_FORMAL_GENERIC, |
| remove_self=True) |
| common_with_cwrap.sort_by_number_of_options(declaration) |
| discover_zero_dim_tensor_operations(declaration) |
| discover_sparse_tensor_operations(declaration) |
| |
| for option in declaration['options']: |
| set_mode(option) |
| sanitize_return(option) |
| process_types_and_backends(option) |
| add_variants(option) |
| declaration['options'] = handle_outputs_taken_as_arguments( |
| declaration['options']) |
| return declarations |
