TensorLib/Aten --> changes required in pytorch

Author: killeent

setup.py

@@ -173,10 +173,12 @@ def run(self):
         from tools.cwrap.plugins.WrapDim import WrapDim
         from tools.cwrap.plugins.AssertNDim import AssertNDim
         from tools.cwrap.plugins.Broadcast import Broadcast
+        from tools.cwrap.plugins.ProcessorSpecificPlugin import ProcessorSpecificPlugin
         thp_plugin = THPPlugin()
         cwrap('torch/csrc/generic/TensorMethods.cwrap', plugins=[
-            BoolOption(), thp_plugin, AutoGPU(condition='IS_CUDA'),
-            ArgcountSortPlugin(), KwargsPlugin(), AssertNDim(), WrapDim(), Broadcast()
+            ProcessorSpecificPlugin(), BoolOption(), thp_plugin,
+            AutoGPU(condition='IS_CUDA'), ArgcountSortPlugin(), KwargsPlugin(),
+            AssertNDim(), WrapDim(), Broadcast()
         ])
         cwrap('torch/csrc/cudnn/cuDNN.cwrap', plugins=[
             CuDNNPlugin(), NullableArguments()

tools/cwrap/cwrap.py

@@ -4,6 +4,7 @@
 from copy import deepcopy
 from .plugins import ArgcountChecker, OptionalArguments, ArgumentReferences, \
     BeforeAfterCall, ConstantArguments, ReturnArguments, GILRelease
+from ..shared import cwrap_common
 
 
 class cwrap(object):
@@ -76,7 +77,7 @@ def wrap_declarations(self, declarations):
             elif line == ']]':
                 in_declaration = False
                 declaration = yaml.load('\n'.join(declaration_lines))
-                self.set_declaration_defaults(declaration)
+                cwrap_common.set_declaration_defaults(declaration)
 
                 # Pass declaration in a list - maybe some plugins want to add
                 # multiple wrappers
@@ -104,24 +105,6 @@ def wrap_declarations(self, declarations):
 
         return '\n'.join(output)
 
-    def set_declaration_defaults(self, declaration):
-        declaration.setdefault('arguments', [])
-        declaration.setdefault('return', 'void')
-        if 'cname' not in declaration:
-            declaration['cname'] = declaration['name']
-        # Simulate multiple dispatch, even if it's not necessary
-        if 'options' not in declaration:
-            declaration['options'] = [{'arguments': declaration['arguments']}]
-            del declaration['arguments']
-        # Parse arguments (some of them can be strings)
-        for option in declaration['options']:
-            option['arguments'] = self.parse_arguments(option['arguments'])
-        # Propagate defaults from declaration to options
-        for option in declaration['options']:
-            for k, v in declaration.items():
-                if k != 'name' and k != 'options':
-                    option.setdefault(k, v)
-
     def parse_arguments(self, args):
         new_args = []
         for arg in args:

tools/cwrap/plugins/ArgcountSortPlugin.py

@@ -1,4 +1,6 @@
+import os
 from . import CWrapPlugin
+from ...shared import cwrap_common
 
 
 class ArgcountSortPlugin(CWrapPlugin):
@@ -7,8 +9,7 @@ def __init__(self, descending=True):
         self.descending = descending
 
     def process_declarations(self, declarations):
-        def num_checked_args(option):
-            return sum(map(lambda a: not a.get('ignore_check', False), option['arguments']))
         for declaration in declarations:
-            declaration['options'].sort(key=num_checked_args, reverse=self.descending)
+            cwrap_common.sort_by_number_of_options(declaration,
+                                                   self.descending)
         return declarations

tools/cwrap/plugins/OptionalArguments.py

@@ -1,58 +1,18 @@
+import os
 from copy import deepcopy
 from . import CWrapPlugin
 from itertools import product
+from ...shared import cwrap_common
 
 
 class OptionalArguments(CWrapPlugin):
 
     def process_declarations(self, declarations):
-        new_options = []
         for declaration in declarations:
-            for option in declaration['options']:
-                optional_args = []
-                for i, arg in enumerate(option['arguments']):
-                    if 'default' in arg:
-                        optional_args.append(i)
-                for permutation in product((True, False), repeat=len(optional_args)):
-                    option_copy = deepcopy(option)
-                    for i, bit in zip(optional_args, permutation):
-                        arg = option_copy['arguments'][i]
-                        if not bit:
-                            arg['type'] = 'CONSTANT'
-                            arg['ignore_check'] = True
-                            # PyYAML interprets NULL as None...
-                            arg['name'] = 'NULL' if arg['default'] is None else arg['default']
-                    new_options.append(option_copy)
-            declaration['options'] = self.filter_unique_options(new_options)
-        return declarations
+            cwrap_common.enumerate_options_due_to_default(
+                declaration,
+                allow_kwarg=True,
+                type_to_signature={},
+                remove_self=False)
 
-    def filter_unique_options(self, options):
-        def signature(option, kwarg_only_count):
-            if kwarg_only_count == 0:
-                kwarg_only_count = None
-            else:
-                kwarg_only_count = -kwarg_only_count
-            arg_signature = '#'.join(
-                arg['type']
-                for arg in option['arguments'][:kwarg_only_count]
-                if not arg.get('ignore_check'))
-            if kwarg_only_count is None:
-                return arg_signature
-            kwarg_only_signature = '#'.join(
-                arg['name'] + '#' + arg['type']
-                for arg in option['arguments'][kwarg_only_count:]
-                if not arg.get('ignore_check'))
-            return arg_signature + "#-#" + kwarg_only_signature
-        seen_signatures = set()
-        unique = []
-        for option in options:
-            for num_kwarg_only in range(0, len(option['arguments']) + 1):
-                sig = signature(option, num_kwarg_only)
-                if sig not in seen_signatures:
-                    if num_kwarg_only > 0:
-                        for arg in option['arguments'][-num_kwarg_only:]:
-                            arg['kwarg_only'] = True
-                    unique.append(option)
-                    seen_signatures.add(sig)
-                    break
-        return unique
+        return declarations

tools/cwrap/plugins/ProcessorSpecificPlugin.py

@@ -0,0 +1,90 @@
+from copy import deepcopy
+from . import CWrapPlugin
+import yaml
+
+
+class ProcessorSpecificPlugin(CWrapPlugin):
+
+    def process_declarations(self, declarations):
+        # In order to move Torch's random functions into the same cwrap
+        # declaration, we need to be able to handle the fact that on the CPU
+        # these functions take a generator argument, while on the GPU, they
+        # do not. As such, we would like to split those declarations at cwrap
+        # runtime into two separate declarations, one for the CPU (unchanged),
+        # and one for the GPU (with the generator argument removed).
+        #
+        # For example, the declaration arguments:
+        # arguments:
+        #   - THTensor* self
+        #   - arg: THGenerator* generator
+        #     default: THPDefaultGenerator->cdata
+        #     kwarg_only: True
+        #
+        # Would have the generator argument removed when generating for the GPU
+        # backend.
+
+        def arg_contains_generator(arg):
+            return (arg['type'] == 'THGenerator*' or (arg.get('default', None)
+                    is not None and 'THPDefaultGenerator' in
+                    str(arg.get('default', ""))))
+
+        def split_candidate(declaration):
+            # First, check and see if it is a declaration for both CPU/GPU
+            if all([proc in declaration['backends'] for
+                    proc in ['CPU', 'CUDA']]):
+                for option in declaration['options']:
+                    for argument in option['arguments']:
+                        if arg_contains_generator(argument):
+                            return True
+
+            return False
+
+        def can_we_handle_the_split(declaration):
+            # hook into here if the split cannot happen for some reason
+            return True
+
+        def generator_split(declaration):
+            # the split must make two changes: 1. remove the generator argument
+            # for the GPU, and 2. assign the correct backends/types to the
+            # split declaration
+            dec_cpu = declaration
+            dec_gpu = deepcopy(declaration)
+
+            # Remove GPU backend and types from dec_cpu
+            dec_cpu['backends'].remove('CUDA')
+            if dec_cpu.get('backend_type_pairs', False):
+                dec_cpu['backend_type_pairs'] = (
+                    [pair for pair in dec_cpu['backend_type_pairs'] if
+                     pair[1] == 'CPU'])
+            # also need to reach into options
+            for option in dec_cpu['options']:
+                option['backends'].remove('CUDA')
+
+            # Remove CPU backend and types from dec_gpu
+            dec_gpu['backends'].remove('CPU')
+            if dec_gpu.get('backend_type_pairs', False):
+                dec_gpu['backend_type_pairs'] = (
+                    [pair for pair in dec_gpu['backend_type_pairs'] if
+                     pair[1] == 'CUDA'])
+            # also need to reach into options
+            for option in dec_gpu['options']:
+                option['backends'].remove('CPU')
+
+            # Remove generator arguments from dec_gpu options
+            for option in dec_gpu['options']:
+                option['arguments'] = (
+                    [arg for arg in option['arguments'] if
+                     not arg_contains_generator(arg)])
+
+            return [dec_cpu, dec_gpu]
+
+        decs = []
+        for declaration in declarations:
+            if split_candidate(declaration):
+                assert(can_we_handle_the_split(declaration))
+                newdecs = generator_split(declaration)
+                decs.extend(newdecs)
+            else:
+                decs.append(declaration)
+
+        return decs

tools/cwrap/plugins/THPPlugin.py

@@ -334,8 +334,92 @@ def has_output_args(declaration):
                        for option in declaration['options']
                        for arg in option['arguments'])
 
+        def backends_types_to_defined_if_string(declaration):
+            # A declaration has two fields: 'backend', which stores a list of
+            # backends (currently 'cpu' and 'cuda') the declaration applies
+            # to, and 'types', which stores a list of real types the
+            # declaration applies to. In PyTorch, when a function is only
+            # supported by a subset of types, we wrap it in macro definition
+            # checks.
+            #
+            # Previously, we manually required the cwrap declaration to
+            # specify for which backend/type combinations a function was
+            # defined for. Now, we explicitly list the types and backends for
+            # a declaration, if it should only be supported for a specific
+            # subset of types, backends, or type-backend pairs.
+
+            types = declaration.get('types', [])
+            backends = declaration['backends']
+            all_backends = ['CPU', 'CUDA']
+
+            def get_defined_string(backend, real):
+                if backend == 'CUDA':
+                    if real == 'all':
+                        return "IS_CUDA"
+                    else:
+                        return 'CUDA_{0}'.format(real.upper())
+                else:
+                    if real == 'all':
+                        return "!IS_CUDA"
+                    else:
+                        return 'defined(TH_REAL_IS_{0})'.format(real.upper())
+
+            def expand_composite_type(p, t):
+                if t == 'floating_point':
+                    result = ['double', 'float']
+                    if p == 'CUDA':
+                        result.append('half')
+                elif t == 'integral':
+                    result = ['byte', 'char', 'short', 'int', 'long']
+                else:
+                    result = [t]
+                return result
+
+            defineds = []
+
+            # The logic below does not handle corner cases well. We allow the
+            # declaration to have a field 'backend_type_pairs' that stores a
+            # dictionary from type --> backend representing allowed
+            # combinations. Let's use these first.
+            for pair in declaration.get('backend_type_pairs', []):
+                p, t = pair
+                defineds.extend([get_defined_string(p, et) for et in
+                                expand_composite_type(p, t)])
+
+            # In the base case, types is empty and backends contains both
+            # 'CPU' and 'CUDA' --> this means we support all types, and our
+            # string should be empty, or simply the list of explict type
+            # backend pairs
+            if (len(types) == 0 and all([proc in backends for proc in
+                                         all_backends])):
+                return " || ".join(defineds)
+
+            # Case 2: types is empty, but only one backend type is specified
+            if len(types) == 0 and len(backends) == 1:
+                defineds.append('IS_CUDA' if backends[0] == 'CUDA' else
+                                "!IS_CUDA")
+                return " || ".join(defineds)
+
+            # Else, we loop overall all of the backend, type pairs and add
+            # them
+            for p in backends:
+                for t in types:
+                    defineds.extend([get_defined_string(p, et) for et in
+                                    expand_composite_type(p, t)])
+
+            return " || ".join(defineds)
+
         for declaration in declarations:
             # Disable all methods for THHalfTensor, unless cpu_half is True
+
+            dfstr = backends_types_to_defined_if_string(declaration)
+            if len(dfstr) > 0:
+                # for now, need to check for distributed defined if as well
+                if 'defined_if' in declaration:
+                    declaration['defined_if'] += ' && (' + dfstr + ')'
+                else:
+                    declaration['defined_if'] = dfstr
+
             if not declaration.get('cpu_half', False):
                 defined_if = '!defined(TH_REAL_IS_HALF)'
                 if 'defined_if' in declaration:
@@ -362,11 +446,13 @@ def has_output_args(declaration):
                 if option.get('sparse', False):
                     defined_if = option.get('defined_if', '')
                     option['defined_if'] = '!IS_DISTRIBUTED' + (' && ' if defined_if else '') + defined_if
-            if declaration.get('with_stateless', False) or declaration.get('only_stateless', False):
+
+            variants = declaration.get('variants', ['method'])
+            if 'function' in variants:
                 stateless_declaration = self.make_stateless(declaration)
                 new_declarations.append(stateless_declaration)
                 self.stateless_declarations.append(stateless_declaration)
-            if declaration.get('only_stateless', False):
+            if 'method' not in variants:
                 continue
 
             self.declarations.append(declaration)
@@ -379,9 +465,13 @@ def has_output_args(declaration):
 
         register_only = [d for d in declarations if d.get('only_register', False)]
         declarations = [d for d in declarations
-                        if (not d.get('only_stateless', False)) and (not d.get('only_register', False))]
-        self.declarations.extend(filter(lambda x: not x.get('only_stateless', False), register_only))
-        self.stateless_declarations.extend(filter(lambda x: x.get('only_stateless', False), register_only))
+                        if (('method' in d.get('variants', ['method'])) and
+                            (not d.get('only_register', False)))]
+        self.declarations.extend(filter(lambda x: 'method' in x.get('variants',
+                                 ['method']), register_only))
+        self.stateless_declarations.extend(filter(lambda x: 'method' not in
+                                           x.get('variants', ['method']),
+                                           register_only))
 
         self.process_docstrings()