Allow GpuCAReduce do unary elemwise operation on the input.

theano/sandbox/cuda/basic_ops.py

@@ -503,31 +503,49 @@ class GpuCAReduce(GpuOp):
     GPUs are not especially well-suited to reduction operations so it is
     quite possible that the GPU might be slower for some cases.
 
+    pre_scalar_op: if present, must be a scalar op with only 1
+    input. We will execute it on the input value before reduction.
+
     """
 
-    def __init__(self, reduce_mask, scalar_op):
+    def __init__(self, reduce_mask, scalar_op, pre_scalar_op=None):
         self.reduce_mask = tuple(reduce_mask)
         self.scalar_op = scalar_op
         # used to make sure that calls to scalar op
         # have unique name arguments
         self._n_scalar_op_calls = 0
+        self.pre_scalar_op = pre_scalar_op
+        if pre_scalar_op:
+            assert pre_scalar_op.nin == 1
 
     def __eq__(self, other):
         return (type(self) == type(other) and
                 self.reduce_mask == other.reduce_mask and
-                self.scalar_op == other.scalar_op)
+                self.scalar_op == other.scalar_op and
+                self.pre_scalar_op == other.pre_scalar_op)
 
     def __hash__(self):
         return (hash(type(self)) ^
                 hash(self.reduce_mask) ^
-                hash(type(self.scalar_op)))
+                hash(type(self.scalar_op)) ^
+                hash(type(self.pre_scalar_op)))
 
     def __str__(self):
-        return "GpuCAReduce{%s}{%s}" % (
+        pre = ""
+        if self.pre_scalar_op:
+            pre = "pre=%s,red=" % str(self.pre_scalar_op)
+        return "GpuCAReduce{%s%s}{%s}" % (
+                pre,
                 str(self.scalar_op),
                 ','.join(str(i) for i in self.reduce_mask)
                 )
 
+    def __setstate__(self, d):
+        self.__dict__.update(d)
+        # For unpickling of old ops.
+        if not hasattr(self, "pre_scalar_op"):
+            self.pre_scalar_op = None
+
     def make_node(self, x):
         if (x.type.ndim != len(self.reduce_mask)):
             raise TypeError("x must have rank %i" % len(self.reduce_mask))
@@ -889,15 +907,33 @@ class GpuCAReduce(GpuOp):
         else:
             assert isinstance(self.scalar_op, (scal.Maximum,
                                                scal.Minimum))
+            if self.pre_scalar_op:
+                #dtype = node.inputs[0].dtype
+                dtype = 'float32'
+
+                dummy_var = scal.Scalar(dtype=dtype)()
+
+                dummy_node = self.pre_scalar_op.make_node(dummy_var)
+
+                dummy_name = 'assign_init_pre_scalar_op' + str(self._n_scalar_op_calls)
+                self._n_scalar_op_calls += 1
+                t = self.pre_scalar_op.c_code(dummy_node, dummy_name,
+                                              (first_item,), ("",), {})
+                assert t.startswith(' = ')
+                first_item = t[3:]
+                if first_item[-1] == ';':
+                    first_item = first_item[:-1]
+
             return first_item
 
-    def _assign_reduce(self, node, name, left, right, sub):
+    def _assign_reduce(self, node, name, left, right, sub, pre):
         """
             node: the node argument to this op's c_code
             name: the name argument to this op's c_code
             left: a C code string identifying an lvalue
             right: a C code string identifying an expression
             sub: the sub argument to this op's c_code
+            pre: If True, we will add the pre_scalar_op.c_code
 
             returns C code to reduce left and right, assigning the
             result to left."""
@@ -913,7 +949,17 @@ class GpuCAReduce(GpuOp):
 
         dummy_name = name + '_scalar_op' + str(self._n_scalar_op_calls)
         self._n_scalar_op_calls += 1
-
+        if pre and self.pre_scalar_op:
+            assert left == "myresult"
+            dummy_node = self.pre_scalar_op.make_node(dummy_left)
+            dummy_name = name + '_scalar_op' + str(self._n_scalar_op_calls)
+            self._n_scalar_op_calls += 1
+            t = self.pre_scalar_op.c_code(dummy_node, dummy_name,
+                                          (right,), ("",), sub)
+            assert t.startswith(' = ')
+            right = t[3:]
+            if right[-1] == ';':
+                right = right[:-1]
         return self.scalar_op.c_code(dummy_node, dummy_name, (left, right),
                                      (left,), sub)
 
@@ -939,7 +985,8 @@ class GpuCAReduce(GpuOp):
         {
             int idx = threadNum - (threadCount >> 1) * 2;"""
 
-        new_version += self._assign_reduce(node, name, 'buf[idx]','buf[threadNum]', sub)
+        new_version += self._assign_reduce(node, name, 'buf[idx]',
+                                           'buf[threadNum]', sub, False)
 
         new_version += """
         }
@@ -958,8 +1005,8 @@ class GpuCAReduce(GpuOp):
               float temp = buf[threadNum + halfPoint];
               """
 
-        new_version += self._assign_reduce(node, name,
-                                           'buf[threadNum]', 'temp', sub)
+        new_version += self._assign_reduce(node, name, 'buf[threadNum]',
+                                           'temp', sub, False)
 
         new_version += """
             }
@@ -989,8 +1036,8 @@ class GpuCAReduce(GpuOp):
             for (int i = threadNum + warpSize; i < threadCount; i += warpSize)
             {
                 """
-        current_version += self._assign_reduce(node, name,
-                                               'myresult', 'buf[i]', sub) + """
+        current_version += self._assign_reduce(node, name, 'myresult',
+                                               'buf[i]', sub, False) + """
             }
             buf[threadNum] = myresult;
         /*Comment this optimization as it don't work on Fermi GPU.
@@ -1002,7 +1049,7 @@ class GpuCAReduce(GpuOp):
             current_version += self._assign_reduce(node, name,
                                                    'buf[threadNum]',
                                                    'buf[threadNum+%d]' % num,
-                                                   sub)
+                                                   sub, False)
         current_version += """
                 if (threadNum == 0)
                 {
@@ -1019,7 +1066,7 @@ class GpuCAReduce(GpuOp):
             this_if = "if (threadNum + %d < threadCount) " % num + \
                 self._assign_reduce(node, name,
                                     'buf[threadNum]','buf[threadNum+%d]' % num,
-                                    sub)
+                                    sub, False)
             current_version += this_if
         current_version += """
                 if (threadNum == 0)
@@ -1037,7 +1084,8 @@ class GpuCAReduce(GpuOp):
     #Threads must be organized as: threadNum%nb_reduce correspond to the same sum
     #nb_reduce<=warpSize
     def _k_reduce_buf_multiple(self, z_pos, node, name, nb_reduce):
-        reduce_fct = self._assign_reduce(node, name, 'myresult', 'buf[i]', {})
+        reduce_fct = self._assign_reduce(node, name, 'myresult',
+                                         'buf[i]', {}, True)
         return """
         __syncthreads(); // some kernel do multiple reduction.
         buf[threadNum] = myresult;
@@ -1609,7 +1657,7 @@ class GpuCAReduce(GpuOp):
             reducebuf = self._k_reduce_buf('Z[0]', node, nodename, sub={})
             reduce_fct = self._assign_reduce(node, nodename, "myresult",
                                              "A[i0]",
-                                             {})
+                                             {}, True)
             reduce_init = self._assign_init("A[0]")
             print >> sio, """
             static __global__ void kernel_reduce_ccontig_%(nodename)s(
@@ -1640,7 +1688,7 @@ class GpuCAReduce(GpuOp):
             reducebuf = self._k_reduce_buf('Z[0]', node, nodename, sub={})
             reduce_fct = self._assign_reduce(node, nodename, "myresult",
                                              "A[i0 * sA0]",
-                                             {})
+                                             {}, True)
             reduce_init = self._assign_init("A[0]")
             print >> sio, """
             static __global__ void kernel_reduce_1_%(nodename)s(
@@ -1671,7 +1719,7 @@ class GpuCAReduce(GpuOp):
             reducebuf = self._k_reduce_buf('Z[0]', node, nodename, sub={})
             reduce_fct = self._assign_reduce(node, nodename, "myresult",
                                              "A[i0 * sA0 + i1 * sA1]",
-                                             {})
+                                             {}, True)
             reduce_init = self._assign_init("A[0]")
 
             print >> sio, """
@@ -1755,7 +1803,8 @@ class GpuCAReduce(GpuOp):
             reduce_fct = self._assign_reduce(
                 node, nodename, "myresult",
                 "A[i3 * sA3 + i2 * sA2 + i1 * sA1 + i0 * sA0]",
-                {})
+                {}, True)
+
             print >> sio, """
                 %(decl)s{
                     %(init)s
@@ -1783,7 +1832,7 @@ class GpuCAReduce(GpuOp):
                                            node, nodename, sub={})
             reduce_fct = self._assign_reduce(node, nodename, "myresult",
                                              "A[i0 * sA0 + i1 * sA1 + i2 * sA2]",
-                                             {})
+                                             {}, True)
             reduce_init = self._assign_init("A[i0 * sA0 + threadIdx.x * sA1 + i2 * sA2]")
             print >> sio, """
             static __global__ void kernel_reduce_010_%(nodename)s(
@@ -1822,7 +1871,7 @@ class GpuCAReduce(GpuOp):
         if self.reduce_mask == (0, 1, 0):
             reduce_fct = self._assign_reduce(node, nodename, "myresult",
                                              "X[a * sX0 + b * sX1 + c * sX2]",
-                                             {})
+                                             {}, True)
             reduce_init = self._assign_init("X[a * sX0 + 0 * sX1 + c * sX2]")
             print >> sio, """
             static __global__ void kernel_reduce_010_AD_%(nodename)s(
@@ -1882,7 +1931,7 @@ class GpuCAReduce(GpuOp):
                                                     'blockDim.x')
             reduce_fct = self._assign_reduce(node, nodename, "myresult",
                                              "A[i0 * sA0 + i1 * sA1 + i2 * sA2]",
-                                             {})
+                                             {}, True)
             reduce_init = self._assign_init("A[i0 * sA0 + 0 * sA1 + i2 * sA2]")
             print >> sio, """
             %(decl)s
@@ -1918,7 +1967,7 @@ class GpuCAReduce(GpuOp):
             reducebuf = self._k_reduce_buf('Z[blockIdx.x * sZ0]', node, nodename, sub = {})
             reduce_fct = self._assign_reduce(node, nodename, "myresult",
                                              "A[i0 * sA0 + i1 * sA1 + blockIdx.x * sA2]",
-                                             {})
+                                             {}, True)
             reduce_init = self._assign_init("A[blockIdx.x * sA2]")
             print >> sio, """
             static __global__ void kernel_reduce_110_%(nodename)s(
@@ -1959,7 +2008,7 @@ class GpuCAReduce(GpuOp):
             init = self._k_init(node, nodename)
             reduce_fct = self._assign_reduce(node, nodename, "myresult",
                                              "A[i0 * sA0 + i1 * sA1 + i2 * sA2]",
-                                             {})
+                                             {}, True)
             reduce_init = self._assign_init("A[i1 * sA1 + i2 * sA2]")
             print >> sio, """
             %(decl)s
@@ -1986,7 +2035,7 @@ class GpuCAReduce(GpuOp):
             init = self._k_init(node, nodename)
             reduce_fct = self._assign_reduce(node, nodename, "myresult",
                                              "A[i0 * sA0 + i1 * sA1 + i2 * sA2]",
-                                             {})
+                                             {}, True)
             reduce_init = self._assign_init("A[0]")
             print >> sio, """
             %(decl)s
@@ -2013,7 +2062,7 @@ class GpuCAReduce(GpuOp):
                                            node, nodename, sub={})
             reduce_fct = self._assign_reduce(node, nodename, "myresult",
                                              "A[i0 * sA0 + i1 * sA1 + i2 * sA2]",
-                                             {})
+                                             {}, True)
             reduce_init = self._assign_init("A[i0 * sA0 + i1 * sA1]")
             print >> sio, """
             static __global__ void kernel_reduce_001_%(nodename)s(
@@ -2056,7 +2105,7 @@ class GpuCAReduce(GpuOp):
             init = self._k_init(node, nodename)
             reduce_fct = self._assign_reduce(node, nodename, "myresult",
                                              "A[i0 * sA0 + i1 * sA1 + i2 * sA2 + i3 * sA3]",
-                                             {})
+                                             {}, True)
             reduce_init = self._assign_init("A[i0 * sA0 + i1 * sA1]")
             print >> sio, """
             %(decl)s
@@ -2089,7 +2138,7 @@ class GpuCAReduce(GpuOp):
             init = self._k_init(node, nodename)
             reduce_fct = self._assign_reduce(node, nodename, "myresult",
                                              "A[i0 * sA0 + i1 * sA1 + i2 * sA2 + i3 * sA3]",
-                                             {})
+                                             {}, True)
             reduce_init = self._assign_init("A[i0 * sA0 + i2 * sA2]")
             print >> sio, """
             %(decl)s
@@ -2120,7 +2169,7 @@ class GpuCAReduce(GpuOp):
             init = self._k_init(node, nodename)
             reduce_fct = self._assign_reduce(node, nodename, "myresult",
                                              "A[i0 * sA0 + i1 * sA1 + i2 * sA2 + i3 * sA3]",
-                                             {})
+                                             {}, True)
             reduce_init = self._assign_init("A[0]")
             print >> sio, """
             %(decl)s
@@ -2146,7 +2195,7 @@ class GpuCAReduce(GpuOp):
                                            node, nodename, sub={})
             reduce_fct = self._assign_reduce(node, nodename, "myresult",
                                              "A[i0 * sA0 + blockIdx.x * sA1 + i2 * sA2 + i3 * sA3]",
-                                             {})
+                                             {}, True)
             reduce_init = self._assign_init("A[blockIdx.x * sA1]")
             print >> sio, """
             static __global__ void kernel_reduce_1011_%(nodename)s(

theano/sandbox/cuda/opt.py

@@ -35,7 +35,7 @@ from theano.sandbox.cuda.blas import (GpuDownsampleFactorMax,
 from theano.sandbox.cuda.nnet import (
         GpuCrossentropySoftmaxArgmax1HotWithBias,
         GpuCrossentropySoftmax1HotWithBiasDx,
-        GpuSoftmax, GpuSoftmaxWithBias, GpuSqrSumAx0)
+        GpuSoftmax, GpuSoftmaxWithBias)
 from theano.sandbox.cuda.elemwise import SupportCodeError
 from theano.scalar.basic_scipy import Erfinv
 from theano.sandbox.cuda.elemwise import erfinv_gpu
@@ -685,17 +685,22 @@ def local_gpu_careduce(node):
     return False
 
 
-@register_opt()#"fast_compile")
+@register_opt("low_memory")
 @local_optimizer([GpuCAReduce])
-def local_gpu_sqr_sum_ax0(node):
+def local_gpu_elemwise_careduce(node):
     if (isinstance(node.op, GpuCAReduce) and
-        isinstance(node.op.scalar_op, theano.scalar.basic.Add) and
-        node.op.reduce_mask == (1, 0) and
+        node.op.pre_scalar_op is None and
         node.inputs[0].owner and
         isinstance(node.inputs[0].owner.op, GpuElemwise) and
-        isinstance(node.inputs[0].owner.op.scalar_op, theano.scalar.basic.Sqr)
+        # The Op support all scalar with 1 inputs.  We don't
+        # automatically add more case, as some like trigonometic
+        # operation with some reduction pattern will probably result
+        # to slow down.
+        isinstance(node.inputs[0].owner.op.scalar_op, scal.basic.Sqr)
         ):
-        return [GpuSqrSumAx0()(node.inputs[0].owner.inputs[0])]
+        op = node.op
+        inp = node.inputs[0].owner.inputs[0]
+        return [GpuCAReduce(op.reduce_mask, op.scalar_op, scal.basic.sqr)(inp)]
 
 
 @register_opt()

theano/sandbox/cuda/tests/test_basic_ops.py

@@ -60,6 +60,10 @@ def test_careduce():
     1110,1101,1011
 
     TODO: test with broadcast
+
+    We test with the pre_scalar_op sqr in all cases. This cover all
+    code, with and without it the pre_scalar_op.
+
     """
     for scalar_op, careduce_op in [
             (theano.scalar.mul, tensor.elemwise.CAReduceDtype),
@@ -132,7 +136,7 @@ def test_careduce():
             pat = tensor_pattern_to_gpu_pattern(shape, pattern)
 
             a = tensor.TensorType('float32', (False,) * len(shape))()
-            b = op(a)
+            b = op(a*a)
             val = numpy.random.rand(numpy.prod(shape)).reshape(shape)
     #        val = numpy.ones(shape)
     #        val = numpy.arange(numpy.prod(shape)).reshape(shape)
@@ -142,6 +146,10 @@ def test_careduce():
             assert tcn.GpuCAReduce in [x.op.__class__
                                        for x in f.maker.fgraph.toposort()], (
                                            scalar_op, shape, pattern)
+            if tcn.GpuElemwise in [x.op.__class__
+                                   for x in f.maker.fgraph.toposort()]:
+                assert tcn.GpuReshape in [x.op.__class__
+                                          for x in f.maker.fgraph.toposort()]
             assert op.__class__ in [x.op.__class__
                                     for x in f2.maker.fgraph.toposort()], (
                                            scalar_op, shape, pattern)
@@ -210,7 +218,7 @@ def test_careduce():
             dim_pattern[0] = 1
             dim_pattern[1] = 0
             a = a.dimshuffle(dim_pattern)
-            b = op(a)
+            b = op(a*a)
             val = numpy.random.rand(numpy.prod(shape)).reshape(shape)
     #        val = numpy.ones(shape)
     #        val = numpy.arange(numpy.prod(shape)).reshape(shape)
@@ -220,6 +228,8 @@ def test_careduce():
             assert tcn.GpuCAReduce in [x.op.__class__
                                        for x in f.maker.fgraph.toposort()], (
                                            scalar_op, shape, pattern)
+            assert tcn.GpuElemwise not in [x.op.__class__
+                                           for x in f.maker.fgraph.toposort()]
             assert op.__class__ in [x.op.__class__
                                     for x in f2.maker.fgraph.toposort()], (
                                            scalar_op, shape, pattern)
@@ -242,8 +252,8 @@ def test_careduce():
             shape = numpy.asarray(shape) * 2
             a = tensor.TensorType('float32', (False,) * len(shape))()
             a2 = tcn.CudaNdarrayType((False,) * len(shape))()
-            b = op(a)
-            b2 = op(a2)
+            b = op(a*a)
+            b2 = op(a2*a2)
             val = numpy.random.rand(numpy.prod(shape)).reshape(shape)
     #        val = numpy.ones(shape)
     #        val = numpy.arange(numpy.prod(shape)).reshape(shape)
@@ -266,6 +276,8 @@ def test_careduce():
             assert tcn.GpuCAReduce in [x.op.__class__
                                        for x in f2.maker.fgraph.toposort()], (
                                            scalar_op, shape, pattern)
+            assert tcn.GpuElemwise not in [x.op.__class__
+                                           for x in f.maker.fgraph.toposort()]
             assert op.__class__ in [x.op.__class__
                                     for x in f.maker.fgraph.toposort()], (
                                            scalar_op, shape, pattern)

theano/sandbox/cuda/tests/test_nnet.py

@@ -264,10 +264,24 @@ def test_sqr_sum_ax0():
         gout = f_gpu(data)
         assert numpy.allclose(out, gout), numpy.absolute(out - gout)
 
-    cmp(10, 15)
-    cmp(120000, 15)
-    cmp(15, 120000)
-    cmp(4000, 4000)
-    cmp(0, 15)
-    cmp(10, 0)
-    cmp(0, 0)
+    #cmp(10, 15)
+    #cmp(120000, 15)
+    #cmp(15, 120000)
+    #cmp(4000, 4000)
+    #cmp(0, 15)
+    #cmp(10, 0)
+    #cmp(0, 0)
+
+    m = mode_with_gpu.excluding("local_gpu_sqr_sum_ax0")
+    f_gpu2 = theano.function([x], z, mode=m)
+    n, m = 4000, 4000
+    data = numpy.arange(n * m, dtype='float32').reshape(n, m)
+    import time
+    t0 = time.time()
+    for i in range(1000):
+        f_gpu(data)
+    t1 = time.time()
+    for i in range(1000):
+        f_gpu2(data)
+    t2 = time.time()
+    print t1 - t0, t2 - t1