Return namedtuples from torch.* function with multiple return arguments for C++ operators (#15429)

Summary:
Partially fixes: https://github.com/pytorch/pytorch/issues/394

Implementation detail:

Codegen is modified to generate codes that looks like below:
```C++
static PyObject * THPVariable_svd(PyObject* self_, PyObject* args, PyObject* kwargs)
{
  HANDLE_TH_ERRORS
  static PythonArgParser parser({
    "svd(Tensor input, bool some=True, bool compute_uv=True, *, TensorList[3] out=None)",
  }, /*traceable=*/true);

  ParsedArgs<6> parsed_args;
  auto r = parser.parse(args, kwargs, parsed_args);
  static PyStructSequence_Field fields0[] = {
    {"U", ""}, {"S", ""}, {"V", ""}, {nullptr}
  };
  static PyStructSequence_Desc desc0 = {
    "torch.return_types.svd_out", nullptr,
    fields0, 3
  };
  static PyTypeObject type0;
  static bool namedtuple_type_initialized0 = false;
  if (!namedtuple_type_initialized0) {
    PyStructSequence_InitType(&type0, &desc0);
    namedtuple_type_initialized0 = true;
  }
  static PyStructSequence_Field fields1[] = {
    {"U", ""}, {"S", ""}, {"V", ""}, {nullptr}
  };
  static PyStructSequence_Desc desc1 = {
    "torch.return_types.svd", nullptr,
    fields1, 3
  };
  static PyTypeObject type1;
  static bool namedtuple_type_initialized1 = false;
  if (!namedtuple_type_initialized1) {
    PyStructSequence_InitType(&type1, &desc1);
    namedtuple_type_initialized1 = true;
  }
  if (r.idx == 0) {
    if (r.isNone(3)) {
      return wrap(&type1, dispatch_svd(r.tensor(0), r.toBool(1), r.toBool(2)));
    } else {
      auto results = r.tensorlist_n<3>(3);
      return wrap(&type0, dispatch_svd(r.tensor(0), r.toBool(1), r.toBool(2), results[0], results[1], results[2]));
    }
  }
  Py_RETURN_NONE;
  END_HANDLE_TH_ERRORS
}
```
Types are defined as static member of `THPVariable_${op_name}` functions, and initialized at the first time the function is called.

When parsing function prototypes in `native_functions.yaml`, the parser will set the specified name as `field_name` when see things like `-> (Tensor t1, ...)`. These field names will be the field names of namedtuple. The class of namedtuples will be named `torch.return_types.${op_name}`.

In some python 2, `PyStructSequence` is not a subtype of tuple, so we have to create some functions to check if an object is a tuple or namedtuple for compatibility issue.

Operators in `native_functions.yaml` are changed such that only `max` and `svd` are generated as namedtuple. Tests are added for these two operators to see if the return value works as expected. Docs for these two ops are also updated to explicitly mention the return value is a namedtuple. More ops will be added in later PRs.

There is some issue with Windows build of linker unable to resolve `PyStructSequence_UnnamedField`, and some workaround is added to deal with this case.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/15429

Differential Revision: D13709678

Pulled By: ezyang

fbshipit-source-id: 23a511c9436977098afc49374e9a748b6e30bccf

3 files changed, 65 insertions, 10 deletions

diff --git a/tools/autograd/derivatives.yaml b/tools/autograd/derivatives.yaml
index 9ae87a7475..621496d8be 100644
--- a/tools/autograd/derivatives.yaml
+++ b/tools/autograd/derivatives.yaml
@@ -479,7 +479,7 @@
   self: zeros_like(self.expand(at::infer_size(self.sizes(), mask.sizes()))).masked_scatter_(mask, grad)
 
 - name: max(Tensor self, int64_t dim, bool keepdim)
-  self: index_select_backward(grad, dim, result1, self.sizes(), keepdim)
+  self: index_select_backward(grad, dim, indices, self.sizes(), keepdim)
 
 - name: max(Tensor self)
   self: select_equals_backward(grad, self, result)
@@ -791,7 +791,7 @@
   self: sum_backward(grad, self.sizes(), dim, keepdim).to(self.type().scalarType())
 
 - name: svd(Tensor self, bool some, bool compute_uv)
-  self: svd_backward(grads, self, some, compute_uv, result0, result1, result2)
+  self: svd_backward(grads, self, some, compute_uv, U, S, V)
 
 - name: symeig(Tensor self, bool eigenvectors, bool upper)
   self: symeig_backward(grads, self, eigenvectors, upper, result0, result1)
@@ -1067,10 +1067,10 @@
   self: adaptive_avg_pool3d_backward(grad, self)
 
 - name: adaptive_max_pool2d(Tensor self, IntList output_size)
-  self: adaptive_max_pool2d_backward(grad, self, indices)
+  self: adaptive_max_pool2d_backward(grad, self, result1)
 
 - name: adaptive_max_pool3d(Tensor self, IntList output_size)
-  self: adaptive_max_pool3d_backward(grad, self, indices)
+  self: adaptive_max_pool3d_backward(grad, self, result1)
 
 - name: avg_pool2d(Tensor self, IntList kernel_size, IntList stride, IntList padding, bool ceil_mode, bool count_include_pad)
   self: avg_pool2d_backward(grad, self, kernel_size, stride, padding, ceil_mode, count_include_pad)
@@ -1079,16 +1079,16 @@
   self: avg_pool3d_backward(grad, self, kernel_size, stride, padding, ceil_mode, count_include_pad)
 
 - name: fractional_max_pool2d(Tensor self, IntList kernel_size, IntList output_size, Tensor random_samples)
-  self: fractional_max_pool2d_backward(grad, self, kernel_size, output_size, indices)
+  self: fractional_max_pool2d_backward(grad, self, kernel_size, output_size, result1)
 
 - name: fractional_max_pool3d(Tensor self, IntList kernel_size, IntList output_size, Tensor random_samples)
   self: fractional_max_pool3d_backward(grad, self, kernel_size, output_size, indices)
 
 - name: max_pool2d_with_indices(Tensor self, IntList kernel_size, IntList stride, IntList padding, IntList dilation, bool ceil_mode)
-  self: max_pool2d_with_indices_backward(grad, self, kernel_size, stride, padding, dilation, ceil_mode, indices)
+  self: max_pool2d_with_indices_backward(grad, self, kernel_size, stride, padding, dilation, ceil_mode, result1)
 
 - name: max_pool3d_with_indices(Tensor self, IntList kernel_size, IntList stride, IntList padding, IntList dilation, bool ceil_mode)
-  self: max_pool3d_with_indices_backward(grad, self, kernel_size, stride, padding, dilation, ceil_mode, indices)
+  self: max_pool3d_with_indices_backward(grad, self, kernel_size, stride, padding, dilation, ceil_mode, result1)
 
 - name: max_unpool2d(Tensor self, Tensor indices, IntList output_size)
   self: max_unpool2d_backward(grad, self, indices, output_size)
diff --git a/tools/autograd/gen_autograd.py b/tools/autograd/gen_autograd.py
index 64c1285530..5789fa3937 100644
--- a/tools/autograd/gen_autograd.py
+++ b/tools/autograd/gen_autograd.py
@@ -223,8 +223,10 @@ def main():
                         help='path to Declarations.yaml')
     parser.add_argument('out', metavar='OUT',
                         help='path to output directory')
+    parser.add_argument('autograd', metavar='AUTOGRAD',
+                        help='path to autograd directory')
     args = parser.parse_args()
-    gen_autograd(args.declarations, args.out)
+    gen_autograd(args.declarations, args.out, args.autograd)
 
 
 if __name__ == '__main__':
diff --git a/tools/autograd/gen_python_functions.py b/tools/autograd/gen_python_functions.py
index 4794b024bd..d837dd79be 100644
--- a/tools/autograd/gen_python_functions.py
+++ b/tools/autograd/gen_python_functions.py
@@ -53,6 +53,7 @@ static PyObject * ${pycname}(PyObject* self_, PyObject* args, PyObject* kwargs)
   ${unpack_self}
   ParsedArgs<${max_args}> parsed_args;
   auto r = parser.parse(args, kwargs, parsed_args);
+  ${declare_namedtuple_return_types}
   ${dispatch}
   Py_RETURN_NONE;
   END_HANDLE_TH_ERRORS
@@ -63,8 +64,9 @@ PY_VARIABLE_METHOD_NOARGS = CodeTemplate("""\
 static PyObject * ${pycname}(PyObject* self_, PyObject* args)
 {
   HANDLE_TH_ERRORS
+  ${declare_namedtuple_return_types}
   ${unpack_self}
-  return wrap(${dispatch_name}(${actuals}));
+  return wrap(${namedtuple_return_type}${dispatch_name}(${actuals}));
   END_HANDLE_TH_ERRORS
 }
 """)
@@ -100,7 +102,7 @@ PY_VARIABLE_SET_REQUIRES_GRAD = CodeTemplate("""\
 ${call_dispatch}.set_requires_grad(${requires_grad})""")
 
 PY_VARIABLE_WRAP = CodeTemplate("""\
-return wrap(${call_dispatch});""")
+return wrap(${namedtuple_return_type}${call_dispatch});""")
 
 PY_VARIABLE_DISPATCH = CodeTemplate("""\
 inline ${simple_return_type} ${dispatch_name}(${formal_args}) {
@@ -113,6 +115,22 @@ inline ${simple_return_type} ${dispatch_name}(${formal_args}) {
 PY_VARIABLE_METHOD_DEF = CodeTemplate("""\
 {"${name}", (PyCFunction)${pycname}, ${flags}, NULL},""")
 
+PY_RETURN_NAMEDTUPLE_DEF = CodeTemplate("""\
+static PyStructSequence_Field fields${namedtuple_type_index}[] = {
+  ${namedtuple_fields} {nullptr}
+};
+static PyStructSequence_Desc desc${namedtuple_type_index} = {
+  "torch.return_types.${name}", nullptr,
+  fields${namedtuple_type_index}, ${namedtuple_size}
+};
+static PyTypeObject type${namedtuple_type_index};
+static bool namedtuple_type_initialized${namedtuple_type_index} = false;
+if (!namedtuple_type_initialized${namedtuple_type_index}) {
+  PyStructSequence_InitType(&type${namedtuple_type_index}, &desc${namedtuple_type_index});
+  namedtuple_type_initialized${namedtuple_type_index} = true;
+}
+""")
+
 UNPACK_SELF = "auto& self = reinterpret_cast<THPVariable*>(self_)->cdata;"
 
 PYTHON_FUNCTION_SIGNATURE = CodeTemplate("""\
@@ -589,6 +607,32 @@ def create_python_bindings(python_functions, has_self, is_module=False):
             python_binding_arguments.append(requires_grad_arg)
         return python_binding_arguments
 
+    def emit_namedtuple_return_type_def(declaration, next_index):
+        returns = declaration['returns']
+        if len(returns) <= 1 or all(['field_name' not in x for x in returns]):
+            declaration['namedtuple_return_type'] = ''
+            return '', next_index
+        declaration['namedtuple_type_index'] = next_index
+        declaration['namedtuple_fields'] = ''
+        for x in returns:
+            # See Note [field_name versus name]
+            if 'field_name' not in x:
+                # When building on Windows, `PyStructSequence_UnnamedField` could not be
+                # resolved by the linker for some reason, which cause error in building:
+                #
+                # python_nn_functions.cpp.obj : error LNK2001: unresolved external symbol
+                # PyStructSequence_UnnamedField
+                #
+                # Thus, at this point in time, we do not support unnamed
+                # fields in namedtuple; you must either name all fields,
+                # or none of them.
+                raise ValueError("Unnamed field is not supported by codegen")
+            else:
+                declaration['namedtuple_fields'] += '{"' + x['field_name'] + '", ""}, '
+        declaration['namedtuple_size'] = len(returns)
+        declaration['namedtuple_return_type'] = '&type{}, '.format(next_index)
+        return PY_RETURN_NAMEDTUPLE_DEF.substitute(declaration), next_index + 1
+
     def process_function(name, declarations):
         for declaration in declarations:
             declaration['python_binding_arguments'] = get_python_binding_arguments(declaration)
@@ -601,11 +645,19 @@ def create_python_bindings(python_functions, has_self, is_module=False):
             'max_args': max(len(o['arguments']) + len(o['python_binding_arguments']) for o in declarations),
             'unpack_self': [],
             'dispatch': [],
+            'declare_namedtuple_return_types': '',
         }
 
         if has_self:
             env['unpack_self'] = [UNPACK_SELF]
 
+        # generate namedtuple type declare
+        next_index = 0
+        for declaration in declarations:
+            typedef, next_index = emit_namedtuple_return_type_def(declaration, next_index)
+            env['declare_namedtuple_return_types'] += typedef
+
+        # emit dispatch
         grouped = group_declarations(declarations)
         for i, dictionary in enumerate(grouped):
             signature = dictionary['signature']
@@ -629,6 +681,7 @@ def create_python_bindings(python_functions, has_self, is_module=False):
             tmpl = PY_VARIABLE_METHOD_NOARGS
             env['actuals'] = ['self']
             env['flags'] = 'METH_NOARGS'
+            env['namedtuple_return_type'] = declarations[0]['namedtuple_return_type']
         else:
             tmpl = PY_VARIABLE_METHOD_VARARGS
             env['flags'] = 'METH_VARARGS | METH_KEYWORDS'