Repository: Theano/Theano

Issue:
Pointers for writing ad-hoc operators involving 4D tensors (conv-like) on the GPU
So, I need to implement the following operation. Assume that we have a standard 2D convolution of a batch of images `I` of size `(N, C, W, H)` with a kernel weights `W` of size `(O, C, D, D)` (will assume kernels are square). Now let's write the convolution in the following form:
```
# This will be `(N*T) x (C*D1*D2)`
i2 = image2neigbs(I)
# This will be `(C*D1*D2)xO`
w2 = W.dimshuffle(1, 2, 3, 0).reshape((-1, O))
# This will be `(N*T)xO`
out = dot(i2, w2)
```
So far so good, however I need to calculate
```
# This will be (C*D1*D2)x(C*D1*D2)
X = dot(i2.T, i2)
```
I've tested the naive way to do it as wit the `image2neigbs` and here is the profiling:
```
<% time> <sum %> <apply time> <time per call> <type> <#call> <#apply> <Class name>
  83.3%    83.3%      39.646s       6.61e-01s     C       60      12   theano.tensor.nnet.neighbours.Images2Neibs
   6.4%    89.8%       3.068s       9.59e-03s     C      320      64   theano.gpuarray.basic_ops.GpuFromHost
   4.2%    94.0%       1.999s       1.11e-02s     C      180      36   theano.gpuarray.blas.GpuDot22
   1.4%    95.4%       0.660s       2.75e-03s     C      240      48   theano.gpuarray.basic_ops.HostFromGpu
   0.9%    96.3%       0.428s       3.89e-03s     C      110      22   theano.gpuarray.dnn.GpuDnnConvGradI
   0.8%    97.1%       0.377s       5.89e-04s     C      640     128   theano.gpuarray.basic_ops.GpuReshape
   0.6%    97.7%       0.298s       1.19e-02s     C       25       5   theano.gpuarray.blas.GpuGemm
   0.4%    98.1%       0.188s       1.71e-03s     C      110      22   theano.gpuarray.basic_ops.GpuJoin
   0.4%    98.5%       0.176s       9.78e-04s     C      180      36   theano.gpuarray.dnn.GpuDnnConv
   0.4%    98.8%       0.172s       1.11e-04s     C     1550     310   theano.gpuarray.elemwise.GpuElemwise
   0.3%    99.2%       0.157s       6.03e-04s     C      260      52   theano.gpuarray.elemwise.GpuCAReduceCuda
   0.2%    99.4%       0.113s       8.38e-04s     Py     135      27   theano.gpuarray.linalg.GpuCusolverSolve
   0.2%    99.6%       0.103s       2.89e-04s     C      355      71   theano.gpuarray.dnn.GpuDnnReduction
   0.2%    99.8%       0.102s       1.69e-03s     C       60      12   theano.gpuarray.dnn.GpuDnnConvGradW
   0.0%    99.9%       0.021s       8.65e-05s     C      245      49   theano.gpuarray.basic_ops.GpuAlloc
   0.0%    99.9%       0.020s       5.51e-05s     C      365      73   theano.gpuarray.basic_ops.GpuAllocEmpty
   0.0%    99.9%       0.012s       6.07e-04s     C       20       4   theano.gpuarray.dnn.GpuDnnPoolGrad
   0.0%    99.9%       0.005s       5.53e-06s     C      940     188   theano.gpuarray.basic_ops.GpuContiguous
   0.0%   100.0%       0.005s       2.12e-05s     C      230      46   theano.gpuarray.subtensor.GpuIncSubtensor
   0.0%   100.0%       0.004s       3.84e-04s     C       10       2   theano.gpuarray.pool.GpuDownsampleFactorMaxGradGrad
   ... (remaining 18 Classes account for   0.04%(0.02s) of the runtime)
```
As can be seen the `Images2Neibs` takes 83% of the compute time. At this stage I want to implement the whole operation as a proper GPU operator. However, I was wondering if anyone could point me out how to actually optimize this (e.g. if there are any code for convolutions which are not cuDNN) and how to manuver the 4D tensor in cuda (how to pass the sizes of the array to the funciton etc...).
I know this is noones job to do but any help would be most appriciated.


Discussion:
There are GPU conv2d implementations that are not from cuDNN, see [here](https://github.com/Theano/Theano/blob/c2e14ce1649f0dc1a56de7288cc0f9dda2fb60bc/theano/gpuarray/corr_gemm.c) for instance.
I'm not particularly familiar with that code, but I think you could re-use the im2col / col2im kernels to the the equivalent of img2neighbs.
In fact, it may be a good idea to expose them as Ops that could replace img2neighbs/neighb2img if that is actually true.
Also, I thought there was a GPU implementation of img2neighbs, but your profile seems to keep it on CPU, I'm not sure why.
Oh, shoot, I did not notice that the operator is not the CPU. Even stranger that I'm using `assert_no_cpu_op = raise`. Are there any documentation on what exactly is `im2col` doing as it sounds very close to the `image2neibs`, or at least examples of how to call it so I can try to understand it.
im2col was copied from caffe via torch as part of the CorrMM implementation, if I remember correctly, while image2neibs was developed in parallel in Theano directly.
@f0k or @gvtulder may be able to give you more info.
The high-level explanations of im2col I've seen (https://petewarden.com/2015/04/20/why-gemm-is-at-the-heart-of-deep-learning/, https://arxiv.org/abs/1410.0759) make it seem like it does something close to image2neibs, but I cannot find anything more precise for the code itself.

`assert_no_cpu_op` is not really reliable, since it only detects gpu->cpu->gpu transfers, and in that case the CPU Op was directly on an input.
> The high-level explanations of im2col I've seen (https://petewarden.com/2015/04/20/why-gemm-is-at-the-heart-of-deep-learning/, https://arxiv.org/abs/1410.0759) make it seem like it does something close to image2neibs, but I cannot find anything more precise for the code itself.

"How GEMM works for Convolutions" in the first link explains it pretty well. The catch is that within the convolution implementation, we loop over the batch to call im2col separately for each input example, to minimize the amount of temporary memory needed (where "we" means the developers of caffe the implementation is copied from).

> In fact, it may be a good idea to expose them as Ops that could replace img2neighbs/neighb2img if that is actually true.

Yes, this would be useful, if GPU memory is plentiful. I never used img2neighbs, but I guess it does the same, possibly with a different memory layout. im2col has a lot of bells and whistles (padding, strides, dilation) and is probably even more powerful than img2neighbs.
@f0k Thanks for the pointers. I understand how GEMM works for Convolutions, but generally speaking, there are 2 ways to do it - via expanding the input, which is what I think both im2col and image2neibs are doing, and via expanding the kernel, which results I think in a Toplitz matrix multiplied with the input.

>  The catch is that within the convolution implementation, we loop over the batch to call im2col separately for each input example, to minimize the amount of temporary memory needed

So, this part I do not understand fully. Are you referring to the actual `im2col` as an Op where you output the resulting tensor or are you referring to how it is used "on the fly" in the convolution implementation to produce the output I guess for a single image and then perform GEMM for that image while looping over the batch? Also, image2neighbs also has padding, modes and strides I don't think it has dilation. I think the two operators are potentially doing the same, however, if in reality the `im2col` might be implemented only for a single example such that it is used internally in the convolution as you said, which I was hoping to understand.
@f0k is referring to the "on the fly" implementations which are internal to the CorrMM/GpuCorrMM ops.

Those ops loops over the batch dimension of the images to do the convolutions and so the im2col and col2im functions just apply to one image. I don't think it would be hard to modify these codes to take out the matrix multiplies and just convert them into im2col/col2im ops.
https://github.com/Theano/Theano/blob/ebaef5af9a2888013d188ca441b563b0f48c6191/theano/tensor/nnet/corr_gemm.c#L258
https://github.com/Theano/Theano/blob/ebaef5af9a2888013d188ca441b563b0f48c6191/theano/gpuarray/corr_gemm.c#L499

@JesseLivezey Thanks for the clarification. 

To my understanding, this is what the `image2neighbs` is doing, however not "on the fly", but static for the whole batch. Unfortunately, in my example:
```
# This will be `(N*T) x (C*D1*D2)`
i2 = image2neigbs(I)
# This will be (C*D1*D2)x(C*D1*D2)
Covar = T.dot(i2.T,i2)
```
I don't think the "on the fly" strategy can be applied as you need to basically do the outer product of the `im2col(x)` and `im2col(y)` for all possible combination of `x` and `y`, where these are images in the batch. Unless of course there is some clever trick which I'm missing. 
@botev it does look like there is a gpu version of images2neibs
https://github.com/Theano/Theano/blob/5df0cfd8a6cb700301d92520d9018f20b32baca5/theano/gpuarray/neighbours.py
so it's probably easiest to first figure out why this isn't being optimized onto the gpu.

I see this optimization exists
https://github.com/Theano/Theano/blob/5df0cfd8a6cb700301d92520d9018f20b32baca5/theano/gpuarray/neighbours.py#L583
@JesseLivezey Yeah I fixed this here: https://github.com/Theano/Theano/pull/6163. I'm not sure why but the Op lifter seems to need to be in `opt.py`

Resolving patch:
diff --git a/theano/gpuarray/neighbours.py b/theano/gpuarray/neighbours.py
--- a/theano/gpuarray/neighbours.py
+++ b/theano/gpuarray/neighbours.py
@@ -12,7 +12,6 @@
 
 from .basic_ops import (as_gpuarray_variable, GpuKernelBase, Kernel,
                         infer_context_name)
-from .opt import register_opt2, op_lifter, register_opt
 from .type import GpuArrayType, gpu_context_type
 
 
@@ -578,11 +577,3 @@ def c_code(self, node, name, inp, out, sub):
     def perform(self, node, inp, out, params):
         # Disable the perform method from the CPU version
         Op.perform(self, node, inp, out, params)
-
-
-@register_opt('fast_compile')
-@op_lifter([Images2Neibs])
-@register_opt2([Images2Neibs], 'fast_compile')
-def local_gpua_images2neibs(op, context_name, inputs, outputs):
-    if op.mode in ['valid', 'half', 'full', 'ignore_borders', 'wrap_centered']:
-        return GpuImages2Neibs(op.mode)
diff --git a/theano/gpuarray/opt.py b/theano/gpuarray/opt.py
--- a/theano/gpuarray/opt.py
+++ b/theano/gpuarray/opt.py
@@ -33,6 +33,7 @@
                                               AbstractConv3d,
                                               AbstractConv3d_gradWeights,
                                               AbstractConv3d_gradInputs)
+from theano.tensor.nnet.neighbours import Images2Neibs
 import theano.tensor.nlinalg as nlinalg
 import theano.tensor.signal.pool as pool
 import theano.tensor.slinalg as slinalg
@@ -76,6 +77,7 @@
 from .linalg import (GpuCusolverSolve, MATRIX_STRUCTURES_SOLVE, GpuCholesky,
                      cusolver_available, GpuMagmaMatrixInverse, gpu_svd,
                      GpuMagmaCholesky, gpu_qr, GpuMagmaEigh)
+from .neighbours import GpuImages2Neibs
 
 _logger = logging.getLogger("theano.gpuarray.opt")
 
@@ -2086,6 +2088,14 @@ def local_gpu_maxandargmax(op, context_name, inputs, outputs):
     return op
 
 
+@register_opt('fast_compile')
+@op_lifter([Images2Neibs])
+@register_opt2([Images2Neibs], 'fast_compile')
+def local_gpua_images2neibs(op, context_name, inputs, outputs):
+    if op.mode in ['valid', 'half', 'full', 'ignore_borders', 'wrap_centered']:
+        return GpuImages2Neibs(op.mode)
+
+
 # solve
 @register_opt('fast_compile')
 @op_lifter([slinalg.Solve])
diff --git a/theano/tensor/slinalg.py b/theano/tensor/slinalg.py
--- a/theano/tensor/slinalg.py
+++ b/theano/tensor/slinalg.py
@@ -325,6 +325,9 @@ def grad(self, inputs, output_gradients):
 """Optimized implementation of :func:`theano.tensor.slinalg.solve` when A is lower triangular."""
 solve_upper_triangular = Solve(A_structure='upper_triangular', lower=False)
 """Optimized implementation of :func:`theano.tensor.slinalg.solve` when A is upper triangular."""
+# symmetric solves
+solve_symmetric = Solve(A_structure='symmetric')
+"""Optimized implementation of :func:`theano.tensor.slinalg.solve` when A is symmetric."""
 
 # TODO: Optimizations to replace multiplication by matrix inverse
 #      with solve() Op (still unwritten)


Test patch: