Repository: pytorch/pytorch
Pull request: Cwrap arg assign
My first `cwrap` PR, thorough review needed !

Extend cwrap stucture to assign the arguments to variables named `arg_*` before the function call.
This simplifies the `before_call` that can now access directly the arguments (the previous behaviour with `${arg0}` is still allowed right now for backward compatibility).
Also allows the `before_call` to modify the arguments!!

Also fixes #146 as the unpacking is done before.

Sample before:
```cpp
      PyThreadState *_save = NULL;
      try {
        long ndim = ((THPTensor*)self)->cdata->nDimension;
        THPUtils_assert(ndim == 2, "t_() expects a 2D tensor, but self is %ldD", ndim);
        
        Py_UNBLOCK_THREADS;
        THTensor_(transpose)(LIBRARY_STATE ((THPTensor*)self)->cdata, ((THPTensor*)self)->cdata, 0, 1);
        Py_BLOCK_THREADS;
        Py_INCREF(self);
        return (PyObject*)self;
      } catch (...) {
        if (_save) {
          Py_BLOCK_THREADS;
        }
        throw;
      }
```

Sample after:
```cpp
      THTensor* arg_self = ((THPTensor*)self)->cdata;
      
      PyThreadState *_save = NULL;
      try {
        long ndim = arg_self->nDimension;
        THPUtils_assert(ndim == 2, "t_() expects a 2D tensor, but self is %ldD", ndim);
        
        Py_UNBLOCK_THREADS;
        THTensor_(transpose)(LIBRARY_STATE arg_self, arg_self, 0, 1);
        Py_BLOCK_THREADS;
        Py_INCREF(self);
        return (PyObject*)self;
      } catch (...) {
        if (_save) {
          Py_BLOCK_THREADS;
        }
        throw;
      }
```

Problem statement:
THPPlugin calls Python API without GIL
THPPlugin calls into the Python API without holding the GIL. Most of the calls look benign (read-only access), but it probably should be cleaned up:

``` python
...
        Py_UNBLOCK_THREADS;
        THTensor_(set)(LIBRARY_STATE ((THPTensor*)self)->cdata, ((THPTensor*)PyTuple_GET_ITEM(args, 0))->cdata);
        Py_BLOCK_THREADS;
...
```

Hints from issue discussion:

Commit message: fix flake8

Patch to review:
diff --git a/tools/cwrap/cwrap.py b/tools/cwrap/cwrap.py
index a232ec8d57cc..c12d9eca1353 100644
--- a/tools/cwrap/cwrap.py
+++ b/tools/cwrap/cwrap.py
@@ -7,6 +7,8 @@
 
 
 class cwrap(object):
+    BASE_INDENT_SIZE = 6
+
     RETURN_WRAPPERS = {
         'void': Template('Py_RETURN_NONE;'),
         'long': Template('return PyLong_FromLong($result);'),
@@ -16,15 +18,19 @@ class cwrap(object):
 
     OPTION_TEMPLATE = Template("""
     ${els}if ($arg_check) {
+      $pre_arg_assign
+      $arg_assign
       $code
     """)
 
+    ARG_ASSIGN_TEMPLATE = Template("""${type} ${name} = ${unpack};""")
+
     OPTION_CODE_TEMPLATE = [
         '$call',
         '$return_result',
     ]
 
-    FUNCTION_CALL_TEMPLATE = Template("$capture_result$cname($arg_unpack);")
+    FUNCTION_CALL_TEMPLATE = Template("$capture_result$cname($call_arg);")
 
     DEFAULT_PLUGIN_CLASSES = [ArgcountChecker, ConstantArguments, OptionalArguments,
                               ArgumentReferences, BeforeAfterCall, ReturnArguments, GILRelease]
@@ -148,6 +154,9 @@ def get_return_wrapper(self, option):
     def get_wrapper_template(self, declaration):
         return self.search_plugins('get_wrapper_template', (declaration,), lambda _: None)
 
+    def get_formal_args(self, arguments):
+        return self.search_plugins('get_formal_args', (arguments,), lambda _: arguments)
+
     def get_arg_accessor(self, arg, option):
         def wrap_accessor(arg, _):
             if arg.get('idx') is None:
@@ -178,9 +187,44 @@ def map_selected_arguments(self, base_fn_name, plugin_fn_name, option, arguments
             res = tmpl.substitute(arg=accessor, idx=arg.get('idx'))
             for plugin in self.plugins:
                 res = getattr(plugin, plugin_fn_name)(res, arg, accessor)
+
             result.append(res)
         return result
 
+    def build_arg_assign(self, arguments, arg_unpack):
+        assignement = []
+        call_arg = []
+        # If type names needs to be changed
+        arguments = self.get_formal_args(arguments)
+        for arg, unpack in zip(arguments, arg_unpack):
+            if arg['type'] == 'CONSTANT':
+                call_arg.append(str(arg['name']))
+            else:
+                var_name = "arg_" + str(arg.get('formal_name', arg['name']))
+                res = self.ARG_ASSIGN_TEMPLATE.substitute(
+                    type=arg['type'],
+                    name=var_name,
+                    unpack=unpack)
+
+                if var_name not in call_arg:
+                    assignement.append(res)
+                call_arg.append(var_name)
+        return assignement, call_arg
+
+    def indent_code(self, code):
+        if code == '':
+            return code
+        code_lines = map(lambda s: s.strip(), code.split('\n'))
+        code = '\n'
+        depth = self.BASE_INDENT_SIZE
+        for line in code_lines:
+            depth -= line.count('}') * 2
+            code += ' ' * depth + line + '\n'
+            depth += line.count('{') * 2
+            depth += line.count('(') * 4
+            depth -= line.count(')') * 4
+        return code[:-1]
+
     def generate_option(self, option, is_first):
         checked_args = list(filter(
             lambda arg: 'ignore_check' not in arg or not arg['ignore_check'],
@@ -199,22 +243,29 @@ def generate_option(self, option, is_first):
         for plugin in self.plugins:
             arg_checks = plugin.process_all_checks(arg_checks, option)
 
-        # Generate unpacks
+        # Generate pre_arg assign
+        pre_arg_assign = []
+        for plugin in self.plugins:
+            pre_arg_assign = plugin.process_pre_arg_assign(pre_arg_assign, option)
+
+        # Generate arg assignment and call arguments
         arg_unpack = self.map_selected_arguments('get_type_unpack',
                                                  'process_single_unpack', option, option['arguments'])
-        arg_unpack = ', '.join(arg_unpack)
+        arg_assign, call_arg = self.build_arg_assign(option['arguments'], arg_unpack)
+
+        call_arg = ', '.join(call_arg)
         for plugin in self.plugins:
-            arg_unpack = plugin.process_all_unpacks(arg_unpack, option)
+            call_arg = plugin.process_all_unpacks(call_arg, option)
 
         # Generate call
         try:
             return_result = self.get_return_wrapper(option).substitute()
             call = self.FUNCTION_CALL_TEMPLATE.substitute(capture_result='',
-                                                          cname=option['cname'], arg_unpack=arg_unpack)
+                                                          cname=option['cname'], call_arg=call_arg)
         except KeyError:
             return_result = self.get_return_wrapper(option).substitute(result='__result')
             call = self.FUNCTION_CALL_TEMPLATE.substitute(capture_result=(option['return'] + ' __result = '),
-                                                          cname=option['cname'], arg_unpack=arg_unpack)
+                                                          cname=option['cname'], call_arg=call_arg)
 
         code_template = deepcopy(self.OPTION_CODE_TEMPLATE)
         for plugin in self.plugins:
@@ -222,19 +273,15 @@ def generate_option(self, option, is_first):
                                                                 option)
         code_template = Template('\n'.join(code_template))
         code = code_template.substitute(call=call, return_result=return_result)
-        code_lines = map(lambda s: s.strip(), code.split('\n'))
-        code = '\n'
-        depth = 6
-        for line in code_lines:
-            depth -= line.count('}') * 2
-            code += ' ' * depth + line + '\n'
-            depth += line.count('{') * 2
-            depth += line.count('(') * 4
-            depth -= line.count(')') * 4
+        code = self.indent_code(code)
+        pre_arg_assign = self.indent_code('\n'.join(pre_arg_assign))
+        arg_assign = self.indent_code('\n'.join(arg_assign))
 
         # Put everything together
         return self.OPTION_TEMPLATE.substitute(
             els=('} else ' if not is_first else ''),
             arg_check=arg_checks,
+            pre_arg_assign=pre_arg_assign,
+            arg_assign=arg_assign,
             code=code,
         )
diff --git a/tools/cwrap/plugins/AutoGPU.py b/tools/cwrap/plugins/AutoGPU.py
index d5574240a8a4..c8198cee0c92 100644
--- a/tools/cwrap/plugins/AutoGPU.py
+++ b/tools/cwrap/plugins/AutoGPU.py
@@ -15,7 +15,7 @@ def __init__(self, has_self=True, condition=None):
 #endif
 """
 
-    def process_option_code_template(self, template, option):
+    def process_pre_arg_assign(self, template, option):
         if not option.get('auto_gpu', True):
             return template
         call = 'THCPAutoGPU __autogpu_guard = THCPAutoGPU(args{});'.format(
diff --git a/tools/cwrap/plugins/BoolOption.py b/tools/cwrap/plugins/BoolOption.py
index c686aa4ffdd2..e243a1179eb5 100644
--- a/tools/cwrap/plugins/BoolOption.py
+++ b/tools/cwrap/plugins/BoolOption.py
@@ -9,11 +9,20 @@ class BoolOption(CWrapPlugin):
     def is_bool_option(self, arg):
         return arg['type'] == 'bool' and 'if_true' in arg and 'if_false' in arg
 
+    def process_declarations(self, declarations):
+        for declaration in declarations:
+            for option in declaration['options']:
+                for arg in option['arguments']:
+                    if self.is_bool_option(arg):
+                        arg['is_bool_option'] = True
+                        arg['type'] = 'const char*'
+        return declarations
+
     def get_type_check(self, arg, option):
-        if self.is_bool_option(arg):
+        if arg.get('is_bool_option', False):
             return Template('PyBool_Check($arg)')
 
     def get_type_unpack(self, arg, option):
-        if self.is_bool_option(arg):
+        if arg.get('is_bool_option', False):
             return Template(self.UNPACK_TEMPLATE.safe_substitute(
                 if_true=arg['if_true'], if_false=arg['if_false']))
diff --git a/tools/cwrap/plugins/CuDNNPlugin.py b/tools/cwrap/plugins/CuDNNPlugin.py
index 74f9f6b1f4fd..5944446d88d5 100644
--- a/tools/cwrap/plugins/CuDNNPlugin.py
+++ b/tools/cwrap/plugins/CuDNNPlugin.py
@@ -1,4 +1,5 @@
 from string import Template
+import copy
 from copy import deepcopy
 from . import CWrapPlugin
 from itertools import product
@@ -17,6 +18,10 @@ class CuDNNPlugin(CWrapPlugin):
         'double': Template('THPDoubleUtils_unpackReal($arg)'),
     }
 
+    INPUT_ARGUMENT_MAP = {
+        'THTensor*': 'THVoidTensor*',
+    }
+
     TYPE_CHECK = {
         'Convolution*': Template('THPWrapper_check($arg)'),
         'THTensor*': Template('(PyObject*)Py_TYPE($arg) == tensorClass'),
@@ -79,6 +84,16 @@ def get_type_unpack(self, arg, option):
     def get_type_check(self, arg, option):
         return self.TYPE_CHECK.get(arg['type'], None)
 
+    def get_formal_args(self, arguments):
+        formal_args = []
+        for arg in arguments:
+            arg = copy.copy(arg)
+            new_type = self.INPUT_ARGUMENT_MAP.get(arg['type'])
+            if new_type is not None:
+                arg['type'] = new_type
+            formal_args.append(arg)
+        return formal_args
+
     def get_wrapper_template(self, declaration):
         arg_desc = []
         for option in declaration['options']:
diff --git a/tools/cwrap/plugins/GenericNN.py b/tools/cwrap/plugins/GenericNN.py
index e5741f7dee61..e7883c63376c 100644
--- a/tools/cwrap/plugins/GenericNN.py
+++ b/tools/cwrap/plugins/GenericNN.py
@@ -64,6 +64,7 @@ class GenericNN(CWrapPlugin):
         'THTensor*': 'thpp::Tensor*',
         'THCTensor*': 'thpp::Tensor*',
         'THIndexTensor*': 'thpp::Tensor*',
+        'THCIndexTensor*': 'thpp::Tensor*',
         'THIndex_t': 'long',
         'accreal': 'double',
     }
@@ -98,11 +99,19 @@ def process_declarations(self, declarations):
     def get_arg_accessor(self, arg, option):
         return self.get_type_unpack(arg, option)
 
+    def process_pre_arg_assign(self, pre_arg_assign, option):
+        if option['backend'] == 'cunn':
+            # Enclose arg_assign with CUDA guard
+            pre_arg_assign.append('#ifdef WITH_CUDA')
+        return pre_arg_assign
+
     def process_option_code_template(self, template, option):
-        code = '// fill me in'
+        template = []
+        if option['backend'] == 'cunn':
+            template.append('#endif')
 
         def base_cast(arg, CReal, real):
-            name = arg['formal_name']
+            name = 'arg_' + arg['formal_name']
             type = arg['type']
             if type in self.REAL_TENSOR_TYPES:
                 return ('(TH{CReal}Tensor*){name}->cdata()'
@@ -120,7 +129,7 @@ def base_cast(arg, CReal, real):
         def cast(arg, CReal, real):
             expr = base_cast(arg, CReal, real)
             if arg.get('optional', False):
-                name = arg['formal_name']
+                name = 'arg_' + arg['formal_name']
                 return '{name} ? {expr} : NULL'.format(name=name, expr=expr)
             return expr
 
@@ -135,6 +144,7 @@ def cast(arg, CReal, real):
                 name=option['cname'],
                 float_args=',\n'.join(float_args),
                 double_args=',\n'.join(double_args))
+            template.append(code)
 
         elif option['backend'] == 'cunn':
             float_args = []
@@ -150,11 +160,13 @@ def cast(arg, CReal, real):
                 float_args=',\n'.join(float_args),
                 double_args=',\n'.join(double_args),
                 half_args=',\n'.join(half_args))
+            template.append(code)
 
-        return [code, '']
+        template.append('')
+        return template
 
     def get_type_unpack(self, arg, option):
-        return Template(arg['name'])
+        return Template(arg.get('formal_name', arg['name']))
 
     def get_type_check(self, arg, option):
         if option['backend'] == 'cunn':
diff --git a/tools/cwrap/plugins/THPPlugin.py b/tools/cwrap/plugins/THPPlugin.py
index 43d8d319a370..2186d6d22cd1 100644
--- a/tools/cwrap/plugins/THPPlugin.py
+++ b/tools/cwrap/plugins/THPPlugin.py
@@ -182,6 +182,7 @@ def _allocate(typename, tmpl, cuda_tmpl=None, sparse=False):
         'double': 'float',
         'accreal': '" RealStr "',
         'bool': 'bool',
+        'const char*': 'bool',  # Can come only from bool option.
     }
 
     OUT_INIT = """
@@ -385,6 +386,7 @@ def make_stateless(self, declaration):
         for option in declaration['options']:
             for arg in option['arguments']:
                 if arg['name'] == 'self':
+                    arg['formal_name'] = 'self'
                     arg['name'] = 'source'
         return declaration
 
@@ -458,7 +460,7 @@ def process_all_checks(self, code, option):
 
         return code
 
-    def process_option_code_template(self, template, option):
+    def process_pre_arg_assign(self, template, option):
         new_args = []
         for arg in option['arguments']:
             if not option.get('output_provided', True) and arg.get('output'):
diff --git a/tools/cwrap/plugins/__init__.py b/tools/cwrap/plugins/__init__.py
index f04b164f8c9f..9a38af92a959 100644
--- a/tools/cwrap/plugins/__init__.py
+++ b/tools/cwrap/plugins/__init__.py
@@ -16,6 +16,9 @@ def get_return_wrapper(self, option):
     def get_wrapper_template(self, declaration):
         pass
 
+    def get_formal_args(self, arguments):
+        pass
+
     def get_arg_accessor(self, arg, option):
         pass
 
@@ -46,6 +49,9 @@ def process_declarations(self, declarations):
     def process_option_code_template(self, template, option):
         return template
 
+    def process_pre_arg_assign(self, template, option):
+        return template
+
 
 from .StandaloneExtension import StandaloneExtension
 from .NullableArguments import NullableArguments
diff --git a/torch/csrc/generic/TensorMethods.cwrap b/torch/csrc/generic/TensorMethods.cwrap
index e1f6b5d8b205..f100e7f99833 100644
--- a/torch/csrc/generic/TensorMethods.cwrap
+++ b/torch/csrc/generic/TensorMethods.cwrap
@@ -52,9 +52,11 @@
 #endif
 
 #if IS_CUDA
+#define THBoolTensor THCudaByteTensor
 #define THPBoolTensor THCPByteTensor
 #define THPBoolTensorClass THCPByteTensorClass
 #else
+#define THBoolTensor THByteTensor
 #define THPBoolTensor THPByteTensor
 #define THPBoolTensorClass THPByteTensorClass
 #endif
@@ -96,6 +98,7 @@ typedef THLongStorage THStride;
 #undef THIndexTensor_
 #undef THPIndexTensor
 #undef THPIndexTensorClass
+#undef THBoolTensor
 #undef THPBoolTensor
 #undef THPBoolTensorClass
 #undef RealStr
diff --git a/torch/csrc/generic/methods/Tensor.cwrap b/torch/csrc/generic/methods/Tensor.cwrap
index 701d6b946809..27568b20f607 100644
--- a/torch/csrc/generic/methods/Tensor.cwrap
+++ b/torch/csrc/generic/methods/Tensor.cwrap
@@ -152,7 +152,7 @@ PyObject * THPTensor_(setIndex)(THPTensor *self, PyObject *args)
         - THTensor* self
         - CONSTANT NULL, 0, NULL, NULL
     - cname: setStorage
-      before_call: THLongStoragePtr __storage_size = THLongStorage_newWithSize1(THStorage_(size)(LIBRARY_STATE ((THPStorage*)$arg1)->cdata));
+      before_call: THLongStoragePtr __storage_size = THLongStorage_newWithSize1(THStorage_(size)(LIBRARY_STATE arg_storage));
       arguments:
         - THTensor* self
         - THStorage* storage
@@ -411,7 +411,7 @@ THTensor* THTensor_(transpose_neg)(THTensor *self, THTensor *src, int dim0, int
   cname: newTranspose
   return: THTensor*
   before_call: |
-    long ndim = ((THPTensor*)${arg0})->cdata->nDimension;
+    long ndim = arg_self->nDimension;
     THPUtils_assert(ndim == 2, "t() expects a 2D tensor, but self is %ldD", ndim);
   arguments:
     - THTensor* self
@@ -425,7 +425,7 @@ THTensor* THTensor_(transpose_neg)(THTensor *self, THTensor *src, int dim0, int
   auto_gpu: False
   return: self
   before_call: |
-    long ndim = ((THPTensor*)${arg0})->cdata->nDimension;
+    long ndim = arg_self->nDimension;
     THPUtils_assert(ndim == 2, "t_() expects a 2D tensor, but self is %ldD", ndim);
   arguments:
     - THTensor* self
@@ -655,8 +655,8 @@ THTensor* THTensor_(transpose_neg)(THTensor *self, THTensor *src, int dim0, int
   with_stateless: True
   return: argument 0
   before_call: |
-    THLongStoragePtr _size = THIndexTensor_(newSizeOf)(LIBRARY_STATE ((THPIndexTensor*)$arg3)->cdata);
-    THTensor_(resize)(LIBRARY_STATE ((THPTensor*)$arg0)->cdata, _size, NULL);
+    THLongStoragePtr _size = THIndexTensor_(newSizeOf)(LIBRARY_STATE arg_index);
+    THTensor_(resize)(LIBRARY_STATE arg_result, _size, NULL);
   arguments:
     - arg: THTensor* result
       output: True


Code review comments:
diff_hunk: @@ -178,9 +187,44 @@ def map_selected_arguments(self, base_fn_name, plugin_fn_name, option, arguments
             res = tmpl.substitute(arg=accessor, idx=arg.get('idx'))
             for plugin in self.plugins:
                 res = getattr(plugin, plugin_fn_name)(res, arg, accessor)
+
             result.append(res)
         return result
 
+    def build_arg_assign(self, arguments, arg_unpack):
+        assignement = []
+        call_arg = []
+        # If type names needs to be changed
+        arguments = self.get_formal_args(arguments)
path: tools/cwrap/cwrap.py
text: @user1:
get_assign_args? I think replacing all references to `formal` with `assign` makes this function easier to understand.

diff_hunk: @@ -199,42 +243,45 @@ def generate_option(self, option, is_first):
         for plugin in self.plugins:
             arg_checks = plugin.process_all_checks(arg_checks, option)
 
-        # Generate unpacks
+        # Generate pre_arg assign
+        pre_arg_assign = []
+        for plugin in self.plugins:
+            pre_arg_assign = plugin.process_pre_arg_assign(pre_arg_assign, option)
+
+        # Generate arg assignment and call arguments
         arg_unpack = self.map_selected_arguments('get_type_unpack',
                                                  'process_single_unpack', option, option['arguments'])
-        arg_unpack = ', '.join(arg_unpack)
+        arg_assign, call_arg = self.build_arg_assign(option['arguments'], arg_unpack)
+
+        call_arg = ', '.join(call_arg)
         for plugin in self.plugins:
-            arg_unpack = plugin.process_all_unpacks(arg_unpack, option)
+            call_arg = plugin.process_all_unpacks(call_arg, option)
path: tools/cwrap/cwrap.py
text: @user1:
I'd consider changing the name of that method. It's no longer working on the argument unpacking, but on the call args.

diff_hunk: @@ -178,9 +187,44 @@ def map_selected_arguments(self, base_fn_name, plugin_fn_name, option, arguments
             res = tmpl.substitute(arg=accessor, idx=arg.get('idx'))
             for plugin in self.plugins:
                 res = getattr(plugin, plugin_fn_name)(res, arg, accessor)
+
             result.append(res)
         return result
 
+    def build_arg_assign(self, arguments, arg_unpack):
path: tools/cwrap/cwrap.py
text: @user1:
build_option_args? It builds both the assignments and the call args.