 #include "THCGeneral.h"
 
 /*
-  Base xcorr2 routine: 3D input, 3D output, 4D kernel
-  All chunks of data should be contiguous
+ * Description:
+ *   This code provides convolutions and xcorrelations that are API compatible with
+ *   the ones in THLabConv.
+ *
+ * History:
+ *   July 21, 2011, 11:21PM  -  Clement Farabet  -  Creation, based conv2d routine
  */
-__global__ void validXCorr2ptr(float *input, float *kernel, float *output,
-                               long input_n, long input_h, long input_w, 
-                               long kernel_n, long kernel_h, long kernel_w,
-                               long stride_h, long stride_w)
+
+#define CUDA_SHARED_MEM_SIZE (4*1024-32) // this is given by nVidia: max shared mem per block
+
+/*
+ * Description:
+ *   base conv2D routine: 3D input, 3D output, 4D kernel
+ *
+ *   - all chunks of data should be contiguous
+ *   - the swapkernel flag can be used to generate a conv2 instead of xcorr2
+ *   - the templated kernel size is useful to generate code that's 2x faster
+ *     but can be set to 0 to allow arbitrary kernel sizes
+ *
+ * to do:
+ *   + the most crucial thing is to make the threads make contiguous
+ *     memory accesses wrt each other, to enable memory reads to coalesce.
+ *     this should give an extra 4 to 8x speedup
+ */
+template <bool swapkernel, int T_kernel_h, int T_kernel_w>
+  __global__ void conv2generic(float *input, float *kernel, float *output,
+                               int input_n, int input_h, int input_w,
+                               int kernel_n, int kernel_h, int kernel_w,
+                               int stride_h, int stride_w,
+                               int patch_h, int patch_w)
 {
   // output dimensions
-  long output_h = (input_h - kernel_h) / stride_h + 1;
-  long output_w = (input_w - kernel_w) / stride_w + 1;
-  long output_n = kernel_n / input_n;
-
-  // pre-load B_Y pixels from B_Y*filtersPerThread filters
-  //__shared__ float shFilters[B_Y*channelCache][B_Y*filtersPerThread];
-
-  // pre-load B_Y pixels from B_X*imgsPerThread images
-  //__shared__ float shImages[B_Y*channelCache][B_X*imgsPerThread];
+  int output_h = (input_h - kernel_h) / stride_h + 1;
+  int output_w = (input_w - kernel_w) / stride_w + 1;
+
+  // xcorr or conv
+  int koffset = swapkernel ? kernel_w*kernel_h-1 : 0;
 
   // generate offsets according to block/thread ids
-  long yy_start = threadIdx.y * 4; // 4 lines per thread
-  long yy_end = yy_start + 4; if (yy_end > output_h) yy_end = output_h;
-  long xx_start = 0;
-  long xx_end = output_w;
-
-  // convolution loop
-  long oo, ii, xx, yy;
-  for(oo = 0; oo < output_n; oo++) {
-    for(ii = 0; ii < input_n; ii++) {
-      for(yy = yy_start; yy < yy_end; yy++) {
-        for(xx = xx_start; xx < xx_end; xx++) {
-          // Dot product in two dimensions... (between input image and the mask)
-          float *input_p = input + ii*input_h*input_w + yy*stride_h*input_w + xx*stride_w;
-          float *output_p = output + oo*output_h*output_w + yy*output_w + xx;
-          float *kernel_p = kernel + oo*output_n + ii;
-          float sum = 0;
-          long kx, ky;
-          for(ky = 0; ky < kernel_h; ky++) {
-            for(kx = 0; kx < kernel_w; kx++) {
-              sum += input_p[kx]*kernel_p[kx];
+  int xx_start = threadIdx.x * patch_w;
+  int xx_end = xx_start + patch_w; if (xx_end > output_w) xx_end = output_w;
+  int yy_start = threadIdx.y * patch_h;
+  int yy_end = yy_start + patch_h; if (yy_end > output_h) yy_end = output_h;
+  int oo_start = blockIdx.x;
+  int oo_end = oo_start+1;
+  int ii_start = 0;
+  int ii_end = input_n;
+
+  // iterators
+  int oo, ii, xx, yy, kx, ky, kk;
+
+  // do the kernels fit in shared mem ?
+  if (input_n*kernel_w*kernel_h <= CUDA_SHARED_MEM_SIZE) {
+
+    // put the kernel in shared memory
+    __shared__ float shared_kernel[CUDA_SHARED_MEM_SIZE];
+
+    // first thread of each block does the copy
+    if ((threadIdx.x == 0) && (threadIdx.y == 0)) {
+      for (kk = 0; kk < kernel_w*kernel_h*input_n; kk++) {
+        shared_kernel[kk] = kernel[input_n*kernel_w*kernel_h*blockIdx.x + kk];
+      }
+    }
+
+    // sync threads
+    __syncthreads();
+
+    // templated kernel size
+    if ((T_kernel_w > 0) && (T_kernel_h > 0)) {
+      // unrolled convolution loop
+      for(oo = oo_start; oo < oo_end; oo++) {
+        for(ii = ii_start; ii < ii_end; ii++) {
+          for(yy = yy_start; yy < yy_end; yy++) {
+            for(xx = xx_start; xx < xx_end; xx++) {
+              // Dot product in two dimensions... (between input image and the mask)
+              float *input_p = input + ii*input_h*input_w + yy*stride_h*input_w + xx*stride_w;
+              float *output_p = output + oo*output_h*output_w + yy*output_w + xx;
+              float *kernel_p = shared_kernel + ii * kernel_w * kernel_h + koffset;
+              float sum = 0;
+              if (swapkernel) {
+#pragma unroll
+                for(ky = 0; ky < T_kernel_h; ky++) {
+#pragma unroll
+                  for(kx = 0; kx < T_kernel_w; kx++) {
+                    sum += input_p[kx]*(*kernel_p--);
+                  }
+                  input_p += input_w;
+                }
+              } else {
+#pragma unroll
+                for(ky = 0; ky < T_kernel_h; ky++) {
+#pragma unroll
+                  for(kx = 0; kx < T_kernel_w; kx++) {
+                    sum += input_p[kx]*(*kernel_p++);
+                  }
+                  input_p += input_w;
+                }
+              }
+              *output_p += sum;
             }
-            input_p += input_w; /* next input line */
-            kernel_p += kernel_w; /* next mask line */
           }
-          *output_p += sum;
         }
       }
+    } else {
+      // default convolution loop
+      for(oo = oo_start; oo < oo_end; oo++) {
+        for(ii = ii_start; ii < ii_end; ii++) {
+          for(yy = yy_start; yy < yy_end; yy++) {
+            for(xx = xx_start; xx < xx_end; xx++) {
+              // Dot product in two dimensions... (between input image and the mask)
+              float *input_p = input + ii*input_h*input_w + yy*stride_h*input_w + xx*stride_w;
+              float *output_p = output + oo*output_h*output_w + yy*output_w + xx;
+              float *kernel_p = shared_kernel + ii * kernel_w * kernel_h + koffset;
+              float sum = 0;
+              for(ky = 0; ky < kernel_h; ky++) {
+                for(kx = 0; kx < kernel_w; kx++) {
+                  if (swapkernel) sum += input_p[kx]*(*kernel_p--);
+                  else sum += input_p[kx]*(*kernel_p++);
+                }
+                input_p += input_w;
+              }
+              *output_p += sum;
+            }
+          }
+        }
+      }
+    }
+
+  } else { // not enough shared mem for kernels, simply stream them
+
+    // convolution loop
+    for(oo = oo_start; oo < oo_end; oo++) {
+      for(ii = ii_start; ii < ii_end; ii++) {
+        for(yy = yy_start; yy < yy_end; yy++) {
+          for(xx = xx_start; xx < xx_end; xx++) {
+            // Dot product in two dimensions... (between input image and the mask)
+            float *input_p = input + ii*input_h*input_w + yy*stride_h*input_w + xx*stride_w;
+            float *output_p = output + oo*output_h*output_w + yy*output_w + xx;
+            float *kernel_p = kernel + (oo * input_n + ii) * kernel_w * kernel_h + koffset;
+            float sum = 0;
+            for(ky = 0; ky < kernel_h; ky++) {
+              for(kx = 0; kx < kernel_w; kx++) {
+                if (swapkernel) sum += input_p[kx]*(*kernel_p--);
+                else sum += input_p[kx]*(*kernel_p++);
+              }
+              input_p += input_w;
+            }
+            *output_p += sum;
+          }
+        }
+      }
     }
   }
 }
 
 /*
-  3D input, 4D kernel, 3D output
-  matrix vector product like
-  y <- Ax + beta*y
-*/
+ * API-compatible with THRealTensor_conv2Dmv
+ * 3D input, 4D kernel, 3D output
+ * matrix vector product like: y <- Ax + beta*y
+ */
 TH_API void THCudaTensor_conv2Dmv(THCudaTensor *output, float beta, THCudaTensor *input,
                                   THCudaTensor *kernel, long srow, long scol, const char *type)
 {

  float *weight_data = THCudaTensor_data(kernel);
  float *output_data = THCudaTensor_data(output);

  // auto compute nb of blocks and threads
  dim3 blocks(nOutputPlane);
  dim3 threads(1, nOutputRows / 4);
  if ((nOutputRows % 4) != 0) threads.y++;

  // convolution: input with kernel, 4 modes (full, valid, xcorr2 or conv2)
  if (type[0] == 'f') {

  } else { // 'v'

    if (type[1] == 'x') {
      if ((nKernelCols == 3) && (nKernelRows == 3))
        conv2generic <false, 3, 3> <<<blocks, threads>>> (input_data, weight_data, output_data,
                                                          nInputPlane, nInputRows, nInputCols,
                                                          nOutputPlane*nInputPlane, nKernelRows, nKernelCols,
                                                          srow, scol, patch_h, patch_w);
      if ((nKernelCols == 5) && (nKernelRows == 5))
        conv2generic <false, 5, 5> <<<blocks, threads>>> (input_data, weight_data, output_data,
                                                          nInputPlane, nInputRows, nInputCols,
                                                          nOutputPlane*nInputPlane, nKernelRows, nKernelCols,
                                                          srow, scol, patch_h, patch_w);
      if ((nKernelCols == 7) && (nKernelRows == 7))
        conv2generic <false, 7, 7> <<<blocks, threads>>> (input_data, weight_data, output_data,
                                                          nInputPlane, nInputRows, nInputCols,
                                                          nOutputPlane*nInputPlane, nKernelRows, nKernelCols,
                                                          srow, scol, patch_h, patch_w);
      if ((nKernelCols == 9) && (nKernelRows == 9))
        conv2generic <false, 9, 9> <<<blocks, threads>>> (input_data, weight_data, output_data,
                                                          nInputPlane, nInputRows, nInputCols,
                                                          nOutputPlane*nInputPlane, nKernelRows, nKernelCols,
                                                          srow, scol, patch_h, patch_w);
      if ((nKernelCols == 11) && (nKernelRows == 11))
        conv2generic <false, 11, 11> <<<blocks, threads>>> (input_data, weight_data, output_data,
                                                            nInputPlane, nInputRows, nInputCols,
                                                            nOutputPlane*nInputPlane, nKernelRows, nKernelCols,
                                                            srow, scol, patch_h, patch_w);
      if ((nKernelCols == 13) && (nKernelRows == 13))
        conv2generic <false, 13, 13> <<<blocks, threads>>> (input_data, weight_data, output_data,
                                                            nInputPlane, nInputRows, nInputCols,
                                                            nOutputPlane*nInputPlane, nKernelRows, nKernelCols,
                                                            srow, scol, patch_h, patch_w);
      else
        conv2generic <false, 0 , 0> <<<blocks, threads>>> (input_data, weight_data, output_data,
                                                           nInputPlane, nInputRows, nInputCols,
                                                           nOutputPlane*nInputPlane, nKernelRows, nKernelCols,
                                                           srow, scol, patch_h, patch_w);
    } else { // 'c'
      if ((nKernelCols == 3) && (nKernelRows == 3))
        conv2generic <true, 3, 3> <<<blocks, threads>>> (input_data, weight_data, output_data,
                                                         nInputPlane, nInputRows, nInputCols,
                                                         nOutputPlane*nInputPlane, nKernelRows, nKernelCols,
                                                         srow, scol, patch_h, patch_w);
      if ((nKernelCols == 5) && (nKernelRows == 5))
        conv2generic <true, 5, 5> <<<blocks, threads>>> (input_data, weight_data, output_data,
                                                         nInputPlane, nInputRows, nInputCols,
                                                         nOutputPlane*nInputPlane, nKernelRows, nKernelCols,
                                                         srow, scol, patch_h, patch_w);
      if ((nKernelCols == 7) && (nKernelRows == 7))
        conv2generic <true, 7, 7> <<<blocks, threads>>> (input_data, weight_data, output_data,
                                                         nInputPlane, nInputRows, nInputCols,
                                                         nOutputPlane*nInputPlane, nKernelRows, nKernelCols,
                                                         srow, scol, patch_h, patch_w);
      if ((nKernelCols == 9) && (nKernelRows == 9))
        conv2generic <true, 9, 9> <<<blocks, threads>>> (input_data, weight_data, output_data,
                                                         nInputPlane, nInputRows, nInputCols,
                                                         nOutputPlane*nInputPlane, nKernelRows, nKernelCols,
                                                         srow, scol, patch_h, patch_w);
      if ((nKernelCols == 11) && (nKernelRows == 11))
        conv2generic <true, 11, 11> <<<blocks, threads>>> (input_data, weight_data, output_data,
                                                           nInputPlane, nInputRows, nInputCols,
                                                           nOutputPlane*nInputPlane, nKernelRows, nKernelCols,
                                                           srow, scol, patch_h, patch_w);
      if ((nKernelCols == 13) && (nKernelRows == 13))
        conv2generic <true, 13, 13> <<<blocks, threads>>> (input_data, weight_data, output_data,
                                                           nInputPlane, nInputRows, nInputCols,
                                                           nOutputPlane*nInputPlane, nKernelRows, nKernelCols,
                                                           srow, scol, patch_h, patch_w);
      else
        conv2generic <true, 0 , 0> <<<blocks, threads>>> (input_data, weight_data, output_data,
                                                          nInputPlane, nInputRows, nInputCols,
                                                          nOutputPlane*nInputPlane, nKernelRows, nKernelCols,
                                                          srow, scol, patch_h, patch_w);
    }

  }
