Merge pull request #2291 from sebastien-j/issue_1763_2

theano/gof/cutils.py

@@ -13,8 +13,7 @@ if os.path.exists(os.path.join(config.compiledir, 'cutils_ext.so')):
     os.remove(os.path.join(config.compiledir, 'cutils_ext.so'))
 
 
-def compile_cutils():
-    """Do just the compilation of cutils_ext"""
+def compile_cutils_code():
 
     types = ['npy_' + t for t in ['int8', 'int16', 'int32', 'int64', 'int128',
         'int256', 'uint8', 'uint16', 'uint32', 'uint64', 'uint128', 'uint256',
@@ -26,7 +25,7 @@ def compile_cutils():
 
     inplace_map_template = """
     #if defined(%(typen)s)
-    static void %(type)s_inplace_add(PyArrayMapIterObject *mit, PyArrayIterObject *it)
+    static void %(type)s_inplace_add(PyArrayMapIterObject *mit, PyArrayIterObject *it, int inc_or_set)
     {
         int index = mit->size;
         while (index--) {
@@ -39,10 +38,10 @@ def compile_cutils():
     #endif
     """
 
-    floatadd = "((%(type)s*)mit->dataptr)[0] = ((%(type)s*)mit->dataptr)[0] + ((%(type)s*)it->dataptr)[0];"
+    floatadd = "((%(type)s*)mit->dataptr)[0] = inc_or_set * ((%(type)s*)mit->dataptr)[0] + ((%(type)s*)it->dataptr)[0];"
     complexadd = """
-    ((%(type)s*)mit->dataptr)[0].real = ((%(type)s*)mit->dataptr)[0].real + ((%(type)s*)it->dataptr)[0].real;
-    ((%(type)s*)mit->dataptr)[0].imag = ((%(type)s*)mit->dataptr)[0].imag + ((%(type)s*)it->dataptr)[0].imag;
+    ((%(type)s*)mit->dataptr)[0].real = inc_or_set * ((%(type)s*)mit->dataptr)[0].real + ((%(type)s*)it->dataptr)[0].real;
+    ((%(type)s*)mit->dataptr)[0].imag = inc_or_set * ((%(type)s*)mit->dataptr)[0].imag + ((%(type)s*)it->dataptr)[0].imag;
     """
 
     fns = ''.join([inplace_map_template % {'type': t, 'typen': t.upper(),
@@ -72,37 +71,13 @@ def compile_cutils():
             "-1000};")
 
     code = ("""
-        #include <Python.h>
-        #include "numpy/arrayobject.h"
-
-        extern "C"{
-        static PyObject *
-        run_cthunk(PyObject *self, PyObject *args)
-        {
-          PyObject *py_cthunk = NULL;
-          if(!PyArg_ParseTuple(args,"O",&py_cthunk))
-            return NULL;
-
-          if (!PyCObject_Check(py_cthunk)) {
-            PyErr_SetString(PyExc_ValueError,
-                           "Argument to run_cthunk must be a PyCObject.");
-            return NULL;
-          }
-          void * ptr_addr = PyCObject_AsVoidPtr(py_cthunk);
-          int (*fn)(void*) = (int (*)(void*))(ptr_addr);
-          void* it = PyCObject_GetDesc(py_cthunk);
-          int failure = fn(it);
-
-          return Py_BuildValue("i", failure);
-        }
-
         #if NPY_API_VERSION >= 0x00000008
-        typedef void (*inplace_map_binop)(PyArrayMapIterObject *, PyArrayIterObject *);
+        typedef void (*inplace_map_binop)(PyArrayMapIterObject *, PyArrayIterObject *, int inc_or_set);
         """ + fns + fn_array + type_number_array +
 
 """
 static int
-map_increment(PyArrayMapIterObject *mit, PyObject *op, inplace_map_binop add_inplace)
+map_increment(PyArrayMapIterObject *mit, PyObject *op, inplace_map_binop add_inplace, int inc_or_set)
 {
     PyArrayObject *arr = NULL;
     PyArrayIterObject *it;
@@ -130,7 +105,7 @@ map_increment(PyArrayMapIterObject *mit, PyObject *op, inplace_map_binop add_inp
         return -1;
     }
 
-    (*add_inplace)(mit, it);
+    (*add_inplace)(mit, it, inc_or_set);
 
     Py_DECREF(arr);
     Py_DECREF(it);
@@ -142,19 +117,20 @@ static PyObject *
 inplace_increment(PyObject *dummy, PyObject *args)
 {
     PyObject *arg_a = NULL, *index=NULL, *inc=NULL;
+    int inc_or_set = 1;
     PyArrayObject *a;
     inplace_map_binop add_inplace = NULL;
     int type_number = -1;
-    int i =0;
+    int i = 0;
     PyArrayMapIterObject * mit;
 
-    if (!PyArg_ParseTuple(args, "OOO", &arg_a, &index,
-            &inc)) {
+    if (!PyArg_ParseTuple(args, "OOO|i", &arg_a, &index,
+            &inc, &inc_or_set)) {
         return NULL;
     }
     if (!PyArray_Check(arg_a)) {
-         PyErr_SetString(PyExc_ValueError, "needs an ndarray as first argument");
-         return NULL;
+        PyErr_SetString(PyExc_ValueError, "needs an ndarray as first argument");
+        return NULL;
     }
 
     a = (PyArrayObject *) arg_a;
@@ -187,7 +163,7 @@ inplace_increment(PyObject *dummy, PyObject *args)
     if (mit == NULL) {
         goto fail;
     }
-    if (map_increment(mit, inc, add_inplace) != 0) {
+    if (map_increment(mit, inc, add_inplace, inc_or_set) != 0) {
         goto fail;
     }
 
@@ -202,9 +178,40 @@ fail:
     return NULL;
 }
         #endif
+""")
 
+    return code
 
-        static PyMethodDef CutilsExtMethods[] = {
+def compile_cutils():
+    """Do just the compilation of cutils_ext"""
+    code = ("""
+        #include <Python.h>
+        #include "numpy/arrayobject.h"
+
+        extern "C"{
+        static PyObject *
+        run_cthunk(PyObject *self, PyObject *args)
+        {
+          PyObject *py_cthunk = NULL;
+          if(!PyArg_ParseTuple(args,"O",&py_cthunk))
+            return NULL;
+
+          if (!PyCObject_Check(py_cthunk)) {
+            PyErr_SetString(PyExc_ValueError,
+                           "Argument to run_cthunk must be a PyCObject.");
+            return NULL;
+          }
+          void * ptr_addr = PyCObject_AsVoidPtr(py_cthunk);
+          int (*fn)(void*) = (int (*)(void*))(ptr_addr);
+          void* it = PyCObject_GetDesc(py_cthunk);
+          int failure = fn(it);
+
+          return Py_BuildValue("i", failure);
+         }""")
+
+    code += compile_cutils_code()
+
+    code += ("""static PyMethodDef CutilsExtMethods[] = {
             {"run_cthunk",  run_cthunk, METH_VARARGS|METH_KEYWORDS,
              "Run a theano cthunk."},
             #if NPY_API_VERSION >= 0x00000008
@@ -214,7 +221,6 @@ fail:
             #endif
             {NULL, NULL, 0, NULL}        /* Sentinel */
         };""")
-
     if PY3:
         # This is not the most efficient code, but it is written this way to
         # highlight the changes needed to make 2.x code compile under python 3.

theano/tensor/subtensor.py

 from copy import copy
 from itertools import izip
+import os
 import sys
 from textwrap import dedent
 import warnings
@@ -1823,6 +1824,68 @@ class AdvancedIncSubtensor1(Op):
 
         return Apply(self, [x_, y_, ilist_], [x_.type()])
 
+    def copy_of_x(self, x):
+        """
+            :param x: a string giving the name of a C variable
+                pointing to an array
+
+            :return: C code expression to make a copy of x
+
+            Base class uses PyArrayObject *, subclasses may override for
+            different types of arrays.
+        """
+        # Parameters of PyArrary_FromAny are:
+        # array
+        # dtype: we pass NULL to say any dtype is acceptable, so the existing
+        #        dtype will be copied
+        # min_depth: we pass 0 to have this parameter ignored
+        # max_depth: we pass 0 to have this parameter ignored
+        # requirements: here we pass NPY_ARRAY_ENSURECOPY to force a copy
+        # context: this is almost always NULL, I'm not sure what it's used for
+        return """(PyArrayObject*)PyArray_FromAny(py_%(x)s, NULL, 0, 0,
+                NPY_ARRAY_ENSURECOPY, NULL)""" % locals()
+
+    def c_support_code(self):
+        from theano.gof.cutils import compile_cutils_code
+        return compile_cutils_code()
+
+    def c_code(self, node, name, input_names, output_names, sub):
+        numpy_ver = [int(n) for n in numpy.__version__.split('.')[:2]]
+        if bool(numpy_ver < [1, 8]):
+            raise NotImplementedError
+        x, y, idx = input_names
+        out = output_names[0]
+        fail = sub['fail']
+        inc_or_set = 1 - self.set_instead_of_inc
+        if self.inplace:  # convert bool to int
+            inplace = 1
+        else:
+            inplace = 0
+        copy_of_x = self.copy_of_x(x)
+
+        return """
+        if (%(inplace)s)
+        {
+            if (%(x)s != %(out)s)
+            {
+                Py_XDECREF(%(out)s);
+                Py_INCREF(%(x)s);
+                %(out)s = %(x)s;
+            }
+        }
+        else
+        {
+            Py_XDECREF(%(out)s);
+            %(out)s = %(copy_of_x)s;
+        }
+        PyObject *arglist = Py_BuildValue("OOOi",%(out)s, %(idx)s, %(y)s, %(inc_or_set)d);
+        inplace_increment(NULL, arglist);
+        Py_XDECREF(arglist);
+        """ % locals()
+
+    def c_code_cache_version(self):
+        return (1,)
+
     def perform(self, node, inp, out_):
         # TODO opt to make this inplace
         x, y, idx = inp
@@ -1832,6 +1895,7 @@ class AdvancedIncSubtensor1(Op):
         # In Numpy, x[idx] += y doesn't work if the same index is present
         # many times: it does it only once. Is it a bug? In any case, for
         # this reason we implement our own 'inc' iteration.
+
         if self.set_instead_of_inc:
             x[idx] = y
         else: