native_parse.py

from __future__ import print_function
import re
import yaml
import pprint
import sys
import copy

try:
    # use faster C loader if available
    from yaml import CLoader as Loader
except ImportError:
    from yaml import Loader

# [temp translations]
# We're currently incrementally moving from the custom func schema to the
# JIT signature schema incrementally. This will reduce overall complexity
# and increase compliance between these components. So for now we do simple
# type translations to continue to emit the legacy func schema for further
# processing by downstream tools. This will helps us avoid having to prematurely
# change all downstream tools to detect these new types.
def type_argument_translations(arg):
    type_and_name = [a.strip() for a in arg.rsplit(' ', 1)]
    name = ''
    if len(type_and_name) > 1:
        name = type_and_name[1]
    t = type_and_name[0]
    name = name.split('=')
    default = None
    nullable = False
    size = None  # Only applies to int[\d+] and Tensor[\d+] arguments
    if len(name) > 1:
        default = name[1]
    name = name[0]

    match = re.match(r'(Tensor.*)\((.+)\)(.*)', t)
    annotation = None
    if match:
        t = match.group(1) + match.group(3)
        annotation = match.group(2)

    # XXX: is_nullable flag can only annotate entire type as optional type,
    # need to special case Generator? logic to make ? only available in jit
    # TODO: deprecate is_nullable global flag, and parse the type
    # to support annotating complicated types with optional annotation
    nullable = (t != 'Generator?' and '?' in t)

    # This enables "Generator? x = None and translates to legacy
    # "Generator* x = nullptr". See [temp translations].
    if t == 'Generator?' and default == 'None':
        t = 'Generator*'
        default = 'nullptr'
    # Enables Generator? by translating to legacy Generator*.
    elif t == "Generator?":
        t = 'Generator*'
    # Enables Tensor[] by translating to legacy TensorList.
    elif t == 'Tensor[]' or t == 'Tensor?[]':
        t = 'TensorList'
    # Enables int[] by translating to legacy IntArrayRef.
    elif t == 'int[]':
        t = 'IntArrayRef'
    # Enables int by translating to legacy int64_t.
    elif t == 'int':
        t = 'int64_t'
    elif t == 'int?':
        t = 'int64_t?'
    elif t == 'int64_t':
        raise RuntimeError("Please use int and not int64_t. "
                           "See [temp translations] for details.")
    elif t == 'int64_t?':
        raise RuntimeError("Please use int? and not int64_t?. "
                           "See [temp translations] for details.")
    elif t == 'Dimname[]?':
        t = 'DimnameList?'
    # Enables float by translating to legacy double.
    elif t == 'float':
        t = 'double'
    # Enables str by translating to legacy std::string.
    elif t == 'str':
        t = 'std::string'
    elif t == 'double':
        raise RuntimeError("Please use float and not double. "
                           "See [temp translations] for details.")
    # Enables int[x] by translating to legacy IntArrayRef[x]. See [temp translations]
    elif re.match(r'int\[(\d+)\]', t):
        match = re.match(r'int\[(\d+)\]', t)
        t = 'IntArrayRef'
        size = int(match.group(1))
    # Enables bool[x] by translating to legacy std::array<bool,x>. See [temp translations]
    elif re.match(r'bool\[(\d+)\]', t):
        match = re.match(r'bool\[(\d+)\]', t)
        t = 'std::array<bool,{}>'.format(match.group(1))
    elif re.match(r'std::array', t):
        raise RuntimeError("Please use array notation, e.g. bool[3] and not std::array."
                           "See [temp translations] for details.")
    # Enables Dimname[x] by translating to DimnameList[x]. See [temp translations]
    elif re.match(r'Dimname\[(\d+)\]', t):
        match = re.match(r'Dimname\[(\d+)\]', t)
        t = 'DimnameList'
        size = int(match.group(1))

    # Legacy type sanitization. TODO: Do we really need this?
    if t == 'Generator*':
        t = 'Generator *'

    if not default:
        pass
    # This enables Tensor? x=None and translates to legacy
    # "Tensor? x={}". See [temp translations].
    elif t.startswith('Tensor?') and default == 'None':
        default = "{}"
    elif default == 'True':
        default = True
    elif default == 'False':
        default = False
    elif default == 'true':
        raise RuntimeError("Please use True and not true. "
                           "See [temp translations] for details.")
    elif default == 'false':
        raise RuntimeError("Please use False and not false. "
                           "See [temp translations] for details.")
    # Enables default argument [] by translating to legacy {}.
    # See [temp translations]
    elif default == '[]':
        default = '{}'
    # Enables lists by translating to legacy {.*}.
    # See [temp translations]
    elif re.match(r'\[.*\]', default):
        default = "{" + default[1:-1] + "}"
    elif default == 'None':
        default = 'c10::nullopt'
    # The JIT signature schema uses Mean, but in particular C++ needs
    # the legacy Reduction::Mean. So we'll continue emiting that until
    # we change this at either a JIT schema or C++ level.
    elif default == 'Mean':
        default = 'Reduction::Mean'
    elif default == 'contiguous_format':
        default = 'MemoryFormat::Contiguous'
    elif default == 'per_tensor_affine':
        default = 'QScheme::PER_TENSOR_AFFINE'
    else:
        try:
            default = int(default)
        except ValueError:
            try:
                default = float(default)
            except ValueError:
                pass

    return t, name, default, nullable, size, annotation


def parse_arguments(args, func_variants, declaration, func_return):
    arguments = []
    kwarg_only = False

    if len(args.strip()) == 0:
        return arguments

    # TODO: Use a real parser here; this will get bamboozled
    # by signatures that contain things like std::array<bool, 2> (note the space)
    for arg_idx, arg in enumerate(args.split(', ')):
        type_and_name = [a.strip() for a in arg.rsplit(' ', 1)]
        if type_and_name == ['*']:
            assert not kwarg_only
            kwarg_only = True
            continue

        t, name, default, nullable, size, annotation = type_argument_translations(arg)

        argument_dict = {'type': t.rstrip('?'), 'name': name, 'is_nullable': nullable, 'annotation': annotation}
        if size:
            argument_dict['size'] = size
        if default is not None:
            argument_dict['default'] = default
        if kwarg_only:
            argument_dict['kwarg_only'] = True
        arguments.append(argument_dict)

    is_out_fn = False
    arguments_out = []
    arguments_other = []
    for argument in arguments:
        if argument['type'] == "Tensor" and \
                argument['annotation'] and \
                re.match(r'^(.*!)$', argument['annotation']) and \
                argument.get('kwarg_only'):
            argument['output'] = True
            argument['kwarg_only'] = False
            arguments_out.append(argument)
            is_out_fn = True
        else:
            arguments_other.append(argument)

    arguments = arguments_out + arguments_other

    name = declaration['name']
    if is_out_fn:
        declaration['name'] += "_out"

    # Reverse splat of TensorOptions
    # As we move towards the JIT function schema for native_functions.yaml we need to support
    # the expanded version of TensorOptions. For now we discover whether there are three
    # types and names of keyword arguments: "ScalarType dtype", "Layout layout" and "Device device"
    # Each, if set, must have default arguments set to long or float, strided and "cpu" respectively.
    # They must appear in this order and in this order only in order for us to be able to process them.
    # In the future we will get rid of this specific processing as downstream consumers start relying
    # less on the content of Declarations.yaml. If you want to support more than this you'll
    # potentially have to extend the JIT.

    supported_topt_arguments = [
        [
            {'name': 'dtype', 'type': 'ScalarType', 'is_nullable': False, 'annotation': None},
            {'name': 'layout', 'type': 'Layout', 'is_nullable': False, 'annotation': None},
            {'name': 'device', 'type': 'Device', 'is_nullable': False, 'annotation': None},
            {'name': 'pin_memory', 'type': 'bool', 'is_nullable': False, 'annotation': None, 'default': False},
        ]
    ]
    supported_topt_arguments.append(copy.deepcopy(supported_topt_arguments[0]))
    for arg in supported_topt_arguments[1]:
        arg.update({'kwarg_only': True})
    supported_topt_arguments.append(copy.deepcopy(supported_topt_arguments[1]))
    for arg in supported_topt_arguments[2]:
        arg.update({'default': 'c10::nullopt', 'is_nullable': True})
    # add explicit support for what is needed for tril_indices / triu_indices
    supported_topt_arguments.append(
        [
            {'name': 'dtype', 'type': 'ScalarType', 'annotation': None, 'kwarg_only': True,
             'default': 'long', 'is_nullable': True},
            {'name': 'layout', 'type': 'Layout', 'annotation': None, 'kwarg_only': True,
             'default': 'c10::nullopt', 'is_nullable': True},
            {'name': 'device', 'type': 'Device', 'annotation': None, 'kwarg_only': True,
             'default': 'c10::nullopt', 'is_nullable': True},
            {'name': 'pin_memory', 'type': 'bool', 'annotation': None, 'kwarg_only': True,
             'default': 'c10::nullopt', 'is_nullable': True},
        ]
    )

    corresponding_topts = [
        {'type': 'TensorOptions', 'name': 'options', 'is_nullable': False, 'annotation': None},
    ]
    corresponding_topts.append(corresponding_topts[0].copy())
    corresponding_topts[1]['kwarg_only'] = True
    corresponding_topts.append(corresponding_topts[1].copy())
    corresponding_topts[2]['default'] = '{}'
    corresponding_topts.append(
        {'type': 'TensorOptions', 'name': 'options', 'is_nullable': False, 'annotation': None,
         'kwarg_only': True, 'default': 'at::kLong'})

    def check_topt_representation(topt_representation):
        for idx, supported_topt in enumerate(supported_topt_arguments):
            matches = all(topt_representation[i] == topt for i, topt in enumerate(supported_topt))
            if matches:
                return corresponding_topts[idx]
        return None

    def is_tensor_option(argument):
        return argument['name'] in ['dtype', 'layout', 'device', 'pin_memory']

    new_arguments = []
    idx = 0
    while idx < len(arguments):
        argument = arguments[idx]
        number_of_arguments = len(supported_topt_arguments[0])
        if is_tensor_option(argument) and len(arguments) - idx >= number_of_arguments:
            topt_representation = []
            for i in range(number_of_arguments):
                argument = arguments[idx]
                if not is_tensor_option(argument):
                    break
                topt_representation.append(argument)
                idx += 1
            if len(topt_representation) == number_of_arguments:
                merged_argument = check_topt_representation(topt_representation)
                assert merged_argument, \
                    "Unsupported combination of TensorOptions {}, the only currently supported combinations are {}"\
                    .format(str(topt_representation), str(supported_topt_arguments))
                new_arguments.append(merged_argument)
            else:
                new_arguments += topt_representation
        else:
            new_arguments.append(argument)
            idx += 1

    arguments = new_arguments

    # Sanity checks

    # TODO: convention is that the ith-argument correspond to the i-th return, but it would
    # be better if we just named everything and matched by name.
    for arg_idx, argument in enumerate(arguments_out):
        assert argument['annotation'] == func_return[arg_idx]['annotation'], \
            "For func {} writeable keyword Tensor arguments need to have a matching return Tensor. Further, " \
            "the ith-argument needs to correspond to the i-th return.".format(name)

    assert len(arguments_out) <= len(func_return), "func {} must return at least as many Tensors " \
        "as can be passed as output.".format(name)

    if name.endswith('_out'):
        raise RuntimeError("Native function {} may not be suffixed with _out as we transition to a unified schema. "
                           "Otherwise you will cause confusion amongst consumers of native functions.".format(name))

    if is_out_fn and func_variants not in [[], 'function', ['function']]:
        raise RuntimeError("Native functions with output MUST be declared with only the function variant; "
                           "e.g., variants: function; otherwise you will tickle a Python argument binding bug "
                           "(which usually manifests itself as the result variable being undefined.) "
                           "The culprit was: {}".format(name))
    if not is_out_fn:
        assert len(arguments_out) == 0, "func {} is not marked as output yet contains output " \
            "keyword arguments".format(name)

    # TODO: Explicit checking for void is a hack and should disappear after a more
    # functionally complete implementation of Tensor aliases.
    if declaration['inplace'] and len(func_return) > 0 and func_return[0]['type'] != "void":
        found_self = False
        for arg_idx, argument in enumerate(arguments):
            if argument['name'] == "self":
                assert argument['annotation'] and argument['annotation'].endswith("!"), \
                    "Inplace function \"{}\" needs to annotate Tensor argument named self " \
                    "as mutable.".format(name)
                found_self = True
                assert argument['annotation'] == func_return[arg_idx]['annotation'], \
                    "Inplace function annotations of function {} need to match between " \
                    "input and correponding output.".format(name)
                assert argument['name'] == func_return[arg_idx]['name'] or \
                    argument['name'] == func_return[arg_idx]['name'] + "_return"
                assert argument['type'] == func_return[arg_idx]['type']
        assert found_self, "Inplace function \"{}\" needs Tensor argument named self.".format(name)

    return arguments


def parse_return_arguments(return_decl, inplace, func_decl):
    arguments = []
    # TODO: Use a real parser here; this will get bamboozled
    # by signatures that contain things like std::array<bool, 2> (note the space)
    if return_decl[0] == '(' and return_decl[-1] == ')':
        return_decl = return_decl[1:-1]
    multiple_args = len(return_decl.split(', ')) > 1

    for arg_idx, arg in enumerate(return_decl.split(', ')):
        t, name, default, nullable, size, annotation = type_argument_translations(arg)
        # name of arguments and name of return sometimes have collision
        # in this case, we rename the return name to <name>_return.
        return_name = name
        if name in func_decl['func'].split('->')[0]:
            return_name = name + "_return"
        argument_dict = {'type': t, 'name': return_name, 'annotation': annotation}
        if name:
            # See Note [field_name versus name]
            argument_dict['field_name'] = name
        else:
            if t == "Tensor" and inplace:
                assert annotation and annotation.endswith("!"), \
                    "Return Tensor of function \"{}\" flagged as inplace needs to be " \
                    "annotated as mutable".format(func_decl['func'])
                argument_dict['name'] = 'self'
            else:
                argument_dict['name'] = 'result' if not multiple_args else 'result' + str(arg_idx)
        argument_dict['output'] = True
        arguments.append(argument_dict)
    return arguments


def parse_native_yaml(path):
    with open(path, 'r') as f:
        return yaml.load(f, Loader=Loader)


def propagate_field_names(output_arguments, return_arguments):
    if output_arguments:
        for i, r in enumerate(return_arguments):
            if 'field_name' in r:
                output_arguments[i]['field_name'] = r['field_name']

def is_named_tensor_only(declaration):
    return any(['Dimname' in arg['type'] for arg in declaration['arguments']])


def run(paths):
    declarations = []
    for path in paths:
        for func in parse_native_yaml(path):
            declaration = {'mode': 'native'}
            try:
                declaration['schema_string'] = "aten::" + func['func']
                if '->' in func['func']:
                    func_decl, return_decl = [x.strip() for x in func['func'].split('->')]
                else:
                    raise Exception('Expected return declaration')
                fn_name, arguments = func_decl.split('(', 1)
                if '.' in fn_name:
                    fn_name, overload_name = fn_name.split('.', 1)
                else:
                    overload_name = ''
                assert arguments[-1] == ")", "Expecting closing ) for {}".format(func['func'])
                arguments = arguments[:-1]  # Expect closing )
                declaration['name'] = func.get('name', fn_name)
                declaration['operator_name'] = func.get('name', fn_name)
                declaration['overload_name'] = func.get('overload_name', overload_name)
                declaration['inplace'] = re.search('(^__i|[^_]_$)', fn_name) is not None
                return_arguments = parse_return_arguments(return_decl, declaration['inplace'], func)
                arguments = parse_arguments(arguments, func.get('variants', []), declaration, return_arguments)
                output_arguments = [x for x in arguments if x.get('output')]
                propagate_field_names(output_arguments, return_arguments)
                declaration['return'] = return_arguments if len(output_arguments) == 0 else output_arguments
                declaration['variants'] = func.get('variants', ['function'])
                declaration['requires_tensor'] = func.get('requires_tensor', False)
                declaration['matches_jit_signature'] = func.get('matches_jit_signature', True)
                declaration['cpu_half'] = func.get('cpu_half', False)
                declaration['cpu_bfloat16'] = func.get('cpu_bfloat16', False)
                declaration['cpu_bool'] = func.get('cpu_bool', False)
                declaration['cuda_bool'] = func.get('cuda_bool', False)
                declaration['deprecated'] = func.get('deprecated', False)
                declaration['device_guard'] = func.get('device_guard', True)
                declaration['supports_named_tensor'] = func.get('supports_named_tensor', False)
                declaration['use_c10_dispatcher'] = func.get('use_c10_dispatcher', 'no')
                assert declaration['use_c10_dispatcher'] in ['no', 'unboxed_only', 'full']
                declaration['category_override'] = func.get('category_override', '')
                declaration['arguments'] = func.get('arguments', arguments)
                declaration['type_method_definition_dispatch'] = func.get('dispatch', declaration['name'])
                declaration['python_module'] = func.get('python_module', '')
                declarations.append(declaration)
            except Exception as e:
                msg = '''Exception raised in processing function:
{func}
Generated partial declaration:
{decl}'''.format(func=pprint.pformat(func), decl=pprint.pformat(declaration))
                print(msg, file=sys.stderr)
                raise e

    return declarations