preprocess_declarations.py

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',
        'Bool',
    ],
    'quantized': [
        'QInt8',
        'QUInt8',
        'QInt32',
    ]
}

all_types = type_map['floating_point'] + type_map['integral'] + type_map['quantized']
type_map['all'] = all_types

all_backends = ['CPU', 'CUDA', 'SparseCPU', 'SparseCUDA', 'MkldnnCPU', 'QuantizedCPU']
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_types' not in option:
        backends = option.get('backends', default_backends)
        if isinstance(option.get('type_method_definition_dispatch'), dict):
            backends = option.get('type_method_definition_dispatch').keys()
        backends = set(backends)

        backend_types = {}
        for backend in backends:
            if backend == 'QuantizedCPU':
                backend_types[backend] = type_map['quantized']
            else:
                backend_types[backend] = option.get('types', all_types)
    else:
        backend_types = option['backend_types']

    # expand type alias (integral, floating_point, all)
    def expand(types):
        ret = []
        for t in types:
            if t in type_map:
                ret.extend(type_map[t])
            else:
                assert(t in all_types)
                ret.append(t)
        return ret

    for backend in backend_types.keys():
        assert(backend in all_backends)
        backend_types[backend] = set(expand(backend_types[backend]))

    # special case remove Half for cpu unless it is explicitly enabled
    if not option.get('cpu_half', False):
        if 'CPU' in backend_types:
            backend_types['CPU'].discard('Half')

    # special cases remove bool for cpu and cuda unless it is explicitly enabled
    if not option.get('cpu_bool', False):
        if 'CPU' in backend_types:
            backend_types['CPU'].discard('Bool')

    if not option.get('cuda_bool', False):
        if 'CUDA' in backend_types:
            backend_types['CUDA'].discard('Bool')

    # sort the result for easy reading
    for backend in backend_types.keys():
        backend_types[backend] = sorted([type for type in backend_types[backend]])
    option['backend_types'] = backend_types


def exclude(declaration):
    return 'only_register' in declaration or declaration.get('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 {None, 'NULL', 'nullptr'})

    def should_generate_out_variant(option):
        if 'function' in option['variants'] and option['mode'] != 'native':
            # don't generate _out variants for in-place functions
            return re.search('(^__i|[^_]_$)', option['api_name']) is None
        return False

    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 should_generate_out_variant(option):
                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(r'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)


def is_extended_method(option):
    if 'method' in option['variants']:
        return False
    else:
        return True


def run(declarations):
    declarations = [d for d in declarations if not exclude(d)]
    non_extended_methods = set()
    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)
            if option['mode'] != 'native':
                sanitize_return(option)
            process_types_and_backends(option)
            add_variants(option)
            if not is_extended_method(option):
                non_extended_methods.add(option['api_name'])
        declaration['options'] = handle_outputs_taken_as_arguments(
            declaration['options'])
    # We (very unfortunately) have overloaded virtual methods. Because
    # of C++'s rules, we cannot move one overload without doing some
    # extra work to make sure that overload in a superclass and an
    # overload in a subclass resolve together. I've chosen to resolve
    # this problem simply by moving ALL overloads of a method which
    # occurs in Tensor to Type.  This is why we have to first compute
    # which methods *names* go on type, and then move ALL overloads
    # of this name to Type.
    for declaration in declarations:
        for option in declaration['options']:
            option['extended_method'] = option['api_name'] not in non_extended_methods
    return declarations