Make SpatialMaxPooling API consistent

modified:   SpatialDilatedMaxPooling.lua
modified:   SpatialMaxPooling.lua
renamed:    lib/THNN/generic/SpatialMaxPooling.c -> lib/THNN/generic/SpatialDilatedMaxPooling.c
modified:   lib/THNN/generic/SpatialMaxPooling.c
modified:   lib/THNN/generic/THNN.h
modified:   lib/THNN/init.c

6 files changed, 350 insertions, 285 deletions

diff --git a/SpatialDilatedMaxPooling.lua b/SpatialDilatedMaxPooling.lua
index 929459c..2f0eba0 100644
--- a/SpatialDilatedMaxPooling.lua
+++ b/SpatialDilatedMaxPooling.lua
@@ -15,7 +15,7 @@ function SpatialDilatedMaxPooling:updateOutput(input)
    self.iheight = input:size(dims-1)
    self.iwidth = input:size(dims)
 
-   input.THNN.SpatialMaxPooling_updateOutput(
+   input.THNN.SpatialDilatedMaxPooling_updateOutput(
       input:cdata(),
       self.output:cdata(),
       self.indices:cdata(),
@@ -29,7 +29,7 @@ function SpatialDilatedMaxPooling:updateOutput(input)
 end
 
 function SpatialDilatedMaxPooling:updateGradInput(input, gradOutput)
-   input.THNN.SpatialMaxPooling_updateGradInput(
+   input.THNN.SpatialDilatedMaxPooling_updateGradInput(
       input:cdata(),
       gradOutput:cdata(),
       self.gradInput:cdata(),
diff --git a/SpatialMaxPooling.lua b/SpatialMaxPooling.lua
index c05a876..8475b13 100644
--- a/SpatialMaxPooling.lua
+++ b/SpatialMaxPooling.lua
@@ -46,7 +46,6 @@ function SpatialMaxPooling:updateOutput(input)
       self.kW, self.kH,
       self.dW, self.dH,
       self.padW, self.padH,
-      1, 1,
       self.ceil_mode
    )
    return self.output
@@ -61,7 +60,6 @@ function SpatialMaxPooling:updateGradInput(input, gradOutput)
       self.kW, self.kH,
       self.dW, self.dH,
       self.padW, self.padH,
-      1, 1,
       self.ceil_mode
    )
    return self.gradInput
diff --git a/lib/THNN/generic/SpatialDilatedMaxPooling.c b/lib/THNN/generic/SpatialDilatedMaxPooling.c
new file mode 100644
index 0000000..6500f49
--- /dev/null
+++ b/lib/THNN/generic/SpatialDilatedMaxPooling.c
@@ -0,0 +1,316 @@
+#ifndef TH_GENERIC_FILE
+#define TH_GENERIC_FILE "generic/SpatialDilatedMaxPooling.c"
+#else
+
+static void THNN_(SpatialDilatedMaxPooling_updateOutput_frame)(
+          real *input_p,
+          real *output_p,
+          real *ind_p,
+          long nslices,
+          long iwidth,
+          long iheight,
+          long owidth,
+          long oheight,
+          int kW,
+          int kH,
+          int dW,
+          int dH,
+          int padW,
+          int padH,
+          int dilationW,
+          int dilationH
+          )
+{
+  long k;
+#pragma omp parallel for private(k)
+  for (k = 0; k < nslices; k++)
+  {
+    /* loop over output */
+    long i, j;
+    real *ip = input_p   + k*iwidth*iheight;
+    for(i = 0; i < oheight; i++)
+    {
+      for(j = 0; j < owidth; j++)
+      {
+        long hstart = i * dH - padH;
+        long wstart = j * dW - padW;
+        long hend = fminf(hstart + (kH - 1) * dilationH + 1, iheight);
+        long wend = fminf(wstart + (kW - 1) * dilationW + 1, iwidth);
+        while(hstart < 0)
+          hstart += dilationH;
+        while(wstart < 0)
+          wstart += dilationW;
+
+        /* local pointers */
+        real *op = output_p  + k*owidth*oheight + i*owidth + j;
+        real *indp = ind_p   + k*owidth*oheight + i*owidth + j;
+
+        /* compute local max: */
+        long maxindex = -1;
+        real maxval = -THInf;
+        long tcntr = 0;
+        long x,y;
+        for(y = hstart; y < hend; y += dilationH)
+        {
+          for(x = wstart; x < wend; x += dilationW)
+          {
+            tcntr = y*iwidth + x;
+            real val = *(ip + tcntr);
+            if (val > maxval)
+            {
+              maxval = val;
+              maxindex = tcntr;
+            }
+          }
+        }
+
+        /* set output to local max */
+        *op = maxval;
+
+        /* store location of max */
+        *indp = maxindex + TH_INDEX_BASE;
+      }
+    }
+  }
+}
+
+void THNN_(SpatialDilatedMaxPooling_updateOutput)(
+          THNNState *state,
+          THTensor *input,
+          THTensor *output,
+          THTensor *indices,
+          int kW,
+          int kH,
+          int dW,
+          int dH,
+          int padW,
+          int padH,
+          int dilationW,
+          int dilationH,
+          bool ceil_mode)
+{
+  int dimw = 2;
+  int dimh = 1;
+  long nbatch = 1;
+  long nslices;
+  long iheight;
+  long iwidth;
+  long oheight;
+  long owidth;
+  real *input_data;
+  real *output_data;
+  real *indices_data;
+
+
+  THArgCheck(input->nDimension == 3 || input->nDimension == 4 , 2, "3D or 4D (batch mode) tensor expected");
+
+  if (input->nDimension == 4)
+  {
+    nbatch = input->size[0];
+    dimw++;
+    dimh++;
+  }
+  THArgCheck(input->size[dimw] >= kW - padW && input->size[dimh] >= kH - padH, 2, "input image smaller than kernel size");
+  THArgCheck(kW/2 >= padW && kH/2 >= padH, 2, "pad should be smaller than half of kernel size");
+  
+  /* sizes */
+  nslices = input->size[dimh-1];
+  iheight = input->size[dimh];
+  iwidth = input->size[dimw];
+  if (ceil_mode)
+  {
+    oheight = (long)(ceil((float)(iheight - (dilationH * (kH - 1) + 1) + 2*padH) / dH)) + 1;
+    owidth  = (long)(ceil((float)(iwidth  - (dilationW * (kW - 1) + 1) + 2*padW) / dW)) + 1;
+  }
+  else
+  {
+    oheight = (long)(floor((float)(iheight - (dilationH * (kH - 1) + 1) + 2*padH) / dH)) + 1;
+    owidth  = (long)(floor((float)(iwidth  - (dilationW * (kW - 1) + 1) + 2*padW) / dW)) + 1;
+  }
+
+  if (owidth < 1 || oheight < 1)
+    THError("Given input size: (%dx%dx%d). Calculated output size: (%dx%dx%d). Output size is too small",
+            nslices,iheight,iwidth,nslices,oheight,owidth);
+
+  if (padW || padH)
+  {
+    // ensure that the last pooling starts inside the image
+    if ((oheight - 1)*dH >= iheight + padH)
+      --oheight;
+    if ((owidth  - 1)*dW >= iwidth  + padW)
+      --owidth;
+  }
+
+  /* get contiguous input */
+  input = THTensor_(newContiguous)(input);
+
+  /* resize output */
+  if (input->nDimension == 3)
+  {
+    THTensor_(resize3d)(output, nslices, oheight, owidth);
+    /* indices will contain the locations for each output point */
+    THTensor_(resize3d)(indices,  nslices, oheight, owidth);
+
+    input_data = THTensor_(data)(input);
+    output_data = THTensor_(data)(output);
+    indices_data = THTensor_(data)(indices);
+
+    THNN_(SpatialDilatedMaxPooling_updateOutput_frame)(input_data, output_data,
+                                              indices_data,
+                                              nslices,
+                                              iwidth, iheight,
+                                              owidth, oheight,
+                                              kW, kH, dW, dH,
+                                              padW, padH,
+                                              dilationW, dilationH
+                                              );
+  }
+  else
+  {
+    long p;
+
+    THTensor_(resize4d)(output, nbatch, nslices, oheight, owidth);
+    /* indices will contain the locations for each output point */
+    THTensor_(resize4d)(indices, nbatch, nslices, oheight, owidth);
+
+    input_data = THTensor_(data)(input);
+    output_data = THTensor_(data)(output);
+    indices_data = THTensor_(data)(indices);
+
+#pragma omp parallel for private(p)
+    for (p = 0; p < nbatch; p++)
+    {
+      THNN_(SpatialDilatedMaxPooling_updateOutput_frame)(input_data+p*nslices*iwidth*iheight, output_data+p*nslices*owidth*oheight,
+                                                indices_data+p*nslices*owidth*oheight,
+                                                nslices,
+                                                iwidth, iheight,
+                                                owidth, oheight,
+                                                kW, kH, dW, dH,
+                                                padW, padH,
+                                                dilationW, dilationH
+                                                );
+    }
+  }
+
+  /* cleanup */
+  THTensor_(free)(input);
+}
+
+static void THNN_(SpatialDilatedMaxPooling_updateGradInput_frame)(
+          real *gradInput_p,
+          real *gradOutput_p,
+          real *ind_p,
+          long nslices,
+          long iwidth,
+          long iheight,
+          long owidth,
+          long oheight,
+          int dW,
+          int dH)
+{
+  long k;
+#pragma omp parallel for private(k)
+  for (k = 0; k < nslices; k++)
+  {
+    real *gradInput_p_k = gradInput_p + k*iwidth*iheight;
+    real *gradOutput_p_k = gradOutput_p + k*owidth*oheight;
+    real *ind_p_k = ind_p + k*owidth*oheight;
+
+    /* calculate max points */
+    long i, j;
+    for(i = 0; i < oheight; i++)
+    {
+      for(j = 0; j < owidth; j++)
+      {
+        /* retrieve position of max */
+        long maxp = ind_p_k[i*owidth + j] - TH_INDEX_BASE;
+        /* update gradient */
+        gradInput_p_k[maxp] += gradOutput_p_k[i*owidth + j];
+      }
+    }
+  }
+}
+
+void THNN_(SpatialDilatedMaxPooling_updateGradInput)(
+          THNNState *state,
+          THTensor *input,
+          THTensor *gradOutput,
+          THTensor *gradInput,
+          THTensor *indices,
+          int kW,
+          int kH,
+          int dW,
+          int dH,
+          int padW,
+          int padH,
+          int dilationW,
+          int dilationH,
+          bool ceil_mode)
+{
+  int dimw = 2;
+  int dimh = 1;
+  long nbatch = 1;
+  int nslices;
+  int iheight;
+  int iwidth;
+  int oheight;
+  int owidth;
+  real *gradInput_data;
+  real *gradOutput_data;
+  real *indices_data;
+
+  /* get contiguous gradOutput */
+  gradOutput = THTensor_(newContiguous)(gradOutput);
+
+  /* resize */
+  THTensor_(resizeAs)(gradInput, input);
+  THTensor_(zero)(gradInput);
+
+  if (input->nDimension == 4) {
+    nbatch = input->size[0];
+    dimw++;
+    dimh++;
+  }
+
+  /* sizes */
+  nslices = input->size[dimh-1];
+  iheight = input->size[dimh];
+  iwidth = input->size[dimw];
+  oheight = gradOutput->size[dimh];
+  owidth = gradOutput->size[dimw];
+
+  /* get raw pointers */
+  gradInput_data = THTensor_(data)(gradInput);
+  gradOutput_data = THTensor_(data)(gradOutput);
+  indices_data = THTensor_(data)(indices);
+
+  /* backprop */
+  if (input->nDimension == 3)
+  {
+    THNN_(SpatialDilatedMaxPooling_updateGradInput_frame)(gradInput_data, gradOutput_data,
+                                                 indices_data,
+                                                 nslices,
+                                                 iwidth, iheight,
+                                                 owidth, oheight,
+                                                 dW, dH);
+  }
+  else
+  {
+    long p;
+#pragma omp parallel for private(p)
+    for (p = 0; p < nbatch; p++)
+    {
+      THNN_(SpatialDilatedMaxPooling_updateGradInput_frame)(gradInput_data+p*nslices*iwidth*iheight, gradOutput_data+p*nslices*owidth*oheight,
+                                                   indices_data+p*nslices*owidth*oheight,
+                                                   nslices,
+                                                   iwidth, iheight,
+                                                   owidth, oheight,
+                                                   dW, dH);
+    }
+  }
+
+  /* cleanup */
+  THTensor_(free)(gradOutput);
+}
+
+#endif
diff --git a/lib/THNN/generic/SpatialMaxPooling.c b/lib/THNN/generic/SpatialMaxPooling.c
index 3daef1d..e0fafb1 100644
--- a/lib/THNN/generic/SpatialMaxPooling.c
+++ b/lib/THNN/generic/SpatialMaxPooling.c
@@ -2,78 +2,6 @@
 #define TH_GENERIC_FILE "generic/SpatialMaxPooling.c"
 #else
 
-static void THNN_(SpatialMaxPooling_updateOutput_frame)(
-          real *input_p,
-          real *output_p,
-          real *ind_p,
-          long nslices,
-          long iwidth,
-          long iheight,
-          long owidth,
-          long oheight,
-          int kW,
-          int kH,
-          int dW,
-          int dH,
-          int padW,
-          int padH,
-          int dilationW,
-          int dilationH
-          )
-{
-  long k;
-#pragma omp parallel for private(k)
-  for (k = 0; k < nslices; k++)
-  {
-    /* loop over output */
-    long i, j;
-    real *ip = input_p   + k*iwidth*iheight;
-    for(i = 0; i < oheight; i++)
-    {
-      for(j = 0; j < owidth; j++)
-      {
-        long hstart = i * dH - padH;
-        long wstart = j * dW - padW;
-        long hend = fminf(hstart + (kH - 1) * dilationH + 1, iheight);
-        long wend = fminf(wstart + (kW - 1) * dilationW + 1, iwidth);
-        while(hstart < 0)
-          hstart += dilationH;
-        while(wstart < 0)
-          wstart += dilationW;
-
-        /* local pointers */
-        real *op = output_p  + k*owidth*oheight + i*owidth + j;
-        real *indp = ind_p   + k*owidth*oheight + i*owidth + j;
-
-        /* compute local max: */
-        long maxindex = -1;
-        real maxval = -THInf;
-        long tcntr = 0;
-        long x,y;
-        for(y = hstart; y < hend; y += dilationH)
-        {
-          for(x = wstart; x < wend; x += dilationW)
-          {
-            tcntr = y*iwidth + x;
-            real val = *(ip + tcntr);
-            if (val > maxval)
-            {
-              maxval = val;
-              maxindex = tcntr;
-            }
-          }
-        }
-
-        /* set output to local max */
-        *op = maxval;
-
-        /* store location of max */
-        *indp = maxindex + TH_INDEX_BASE;
-      }
-    }
-  }
-}
-
 void THNN_(SpatialMaxPooling_updateOutput)(
           THNNState *state,
           THTensor *input,
@@ -85,150 +13,12 @@ void THNN_(SpatialMaxPooling_updateOutput)(
           int dH,
           int padW,
           int padH,
-          int dilationW,
-          int dilationH,
           bool ceil_mode)
 {
-  int dimw = 2;
-  int dimh = 1;
-  long nbatch = 1;
-  long nslices;
-  long iheight;
-  long iwidth;
-  long oheight;
-  long owidth;
-  real *input_data;
-  real *output_data;
-  real *indices_data;
-
-
-  THArgCheck(input->nDimension == 3 || input->nDimension == 4 , 2, "3D or 4D (batch mode) tensor expected");
-
-  if (input->nDimension == 4)
-  {
-    nbatch = input->size[0];
-    dimw++;
-    dimh++;
-  }
-  THArgCheck(input->size[dimw] >= kW - padW && input->size[dimh] >= kH - padH, 2, "input image smaller than kernel size");
-  THArgCheck(kW/2 >= padW && kH/2 >= padH, 2, "pad should be smaller than half of kernel size");
-  
-  /* sizes */
-  nslices = input->size[dimh-1];
-  iheight = input->size[dimh];
-  iwidth = input->size[dimw];
-  if (ceil_mode)
-  {
-    oheight = (long)(ceil((float)(iheight - (dilationH * (kH - 1) + 1) + 2*padH) / dH)) + 1;
-    owidth  = (long)(ceil((float)(iwidth  - (dilationW * (kW - 1) + 1) + 2*padW) / dW)) + 1;
-  }
-  else
-  {
-    oheight = (long)(floor((float)(iheight - (dilationH * (kH - 1) + 1) + 2*padH) / dH)) + 1;
-    owidth  = (long)(floor((float)(iwidth  - (dilationW * (kW - 1) + 1) + 2*padW) / dW)) + 1;
-  }
-
-  if (owidth < 1 || oheight < 1)
-    THError("Given input size: (%dx%dx%d). Calculated output size: (%dx%dx%d). Output size is too small",
-            nslices,iheight,iwidth,nslices,oheight,owidth);
-
-  if (padW || padH)
-  {
-    // ensure that the last pooling starts inside the image
-    if ((oheight - 1)*dH >= iheight + padH)
-      --oheight;
-    if ((owidth  - 1)*dW >= iwidth  + padW)
-      --owidth;
-  }
-
-  /* get contiguous input */
-  input = THTensor_(newContiguous)(input);
-
-  /* resize output */
-  if (input->nDimension == 3)
-  {
-    THTensor_(resize3d)(output, nslices, oheight, owidth);
-    /* indices will contain the locations for each output point */
-    THTensor_(resize3d)(indices,  nslices, oheight, owidth);
-
-    input_data = THTensor_(data)(input);
-    output_data = THTensor_(data)(output);
-    indices_data = THTensor_(data)(indices);
-
-    THNN_(SpatialMaxPooling_updateOutput_frame)(input_data, output_data,
-                                              indices_data,
-                                              nslices,
-                                              iwidth, iheight,
-                                              owidth, oheight,
-                                              kW, kH, dW, dH,
-                                              padW, padH,
-                                              dilationW, dilationH
-                                              );
-  }
-  else
-  {
-    long p;
-
-    THTensor_(resize4d)(output, nbatch, nslices, oheight, owidth);
-    /* indices will contain the locations for each output point */
-    THTensor_(resize4d)(indices, nbatch, nslices, oheight, owidth);
-
-    input_data = THTensor_(data)(input);
-    output_data = THTensor_(data)(output);
-    indices_data = THTensor_(data)(indices);
-
-#pragma omp parallel for private(p)
-    for (p = 0; p < nbatch; p++)
-    {
-      THNN_(SpatialMaxPooling_updateOutput_frame)(input_data+p*nslices*iwidth*iheight, output_data+p*nslices*owidth*oheight,
-                                                indices_data+p*nslices*owidth*oheight,
-                                                nslices,
-                                                iwidth, iheight,
-                                                owidth, oheight,
-                                                kW, kH, dW, dH,
-                                                padW, padH,
-                                                dilationW, dilationH
-                                                );
-    }
-  }
-
-  /* cleanup */
-  THTensor_(free)(input);
-}
-
-static void THNN_(SpatialMaxPooling_updateGradInput_frame)(
-          real *gradInput_p,
-          real *gradOutput_p,
-          real *ind_p,
-          long nslices,
-          long iwidth,
-          long iheight,
-          long owidth,
-          long oheight,
-          int dW,
-          int dH)
-{
-  long k;
-#pragma omp parallel for private(k)
-  for (k = 0; k < nslices; k++)
-  {
-    real *gradInput_p_k = gradInput_p + k*iwidth*iheight;
-    real *gradOutput_p_k = gradOutput_p + k*owidth*oheight;
-    real *ind_p_k = ind_p + k*owidth*oheight;
-
-    /* calculate max points */
-    long i, j;
-    for(i = 0; i < oheight; i++)
-    {
-      for(j = 0; j < owidth; j++)
-      {
-        /* retrieve position of max */
-        long maxp = ind_p_k[i*owidth + j] - TH_INDEX_BASE;
-        /* update gradient */
-        gradInput_p_k[maxp] += gradOutput_p_k[i*owidth + j];
-      }
-    }
-  }
+  THNN_(SpatialDilatedMaxPooling_updateOutput)(
+      state, input, output, indices,
+      kW, kH, dW, dH, padW, padH, 1, 1, ceil_mode
+    );
 }
 
 void THNN_(SpatialMaxPooling_updateGradInput)(
@@ -243,74 +33,12 @@ void THNN_(SpatialMaxPooling_updateGradInput)(
           int dH,
           int padW,
           int padH,
-          int dilationW,
-          int dilationH,
           bool ceil_mode)
 {
-  int dimw = 2;
-  int dimh = 1;
-  long nbatch = 1;
-  int nslices;
-  int iheight;
-  int iwidth;
-  int oheight;
-  int owidth;
-  real *gradInput_data;
-  real *gradOutput_data;
-  real *indices_data;
-
-  /* get contiguous gradOutput */
-  gradOutput = THTensor_(newContiguous)(gradOutput);
-
-  /* resize */
-  THTensor_(resizeAs)(gradInput, input);
-  THTensor_(zero)(gradInput);
-
-  if (input->nDimension == 4) {
-    nbatch = input->size[0];
-    dimw++;
-    dimh++;
-  }
-
-  /* sizes */
-  nslices = input->size[dimh-1];
-  iheight = input->size[dimh];
-  iwidth = input->size[dimw];
-  oheight = gradOutput->size[dimh];
-  owidth = gradOutput->size[dimw];
-
-  /* get raw pointers */
-  gradInput_data = THTensor_(data)(gradInput);
-  gradOutput_data = THTensor_(data)(gradOutput);
-  indices_data = THTensor_(data)(indices);
-
-  /* backprop */
-  if (input->nDimension == 3)
-  {
-    THNN_(SpatialMaxPooling_updateGradInput_frame)(gradInput_data, gradOutput_data,
-                                                 indices_data,
-                                                 nslices,
-                                                 iwidth, iheight,
-                                                 owidth, oheight,
-                                                 dW, dH);
-  }
-  else
-  {
-    long p;
-#pragma omp parallel for private(p)
-    for (p = 0; p < nbatch; p++)
-    {
-      THNN_(SpatialMaxPooling_updateGradInput_frame)(gradInput_data+p*nslices*iwidth*iheight, gradOutput_data+p*nslices*owidth*oheight,
-                                                   indices_data+p*nslices*owidth*oheight,
-                                                   nslices,
-                                                   iwidth, iheight,
-                                                   owidth, oheight,
-                                                   dW, dH);
-    }
-  }
-
-  /* cleanup */
-  THTensor_(free)(gradOutput);
+  THNN_(SpatialDilatedMaxPooling_updateGradInput)(
+      state, input, gradOutput, gradInput, indices,
+      kW, kH, dW, dH, padW, padH, 1, 1, ceil_mode
+    );
 }
 
 #endif
diff --git a/lib/THNN/generic/THNN.h b/lib/THNN/generic/THNN.h
index 5313dad..319ffc7 100644
--- a/lib/THNN/generic/THNN.h
+++ b/lib/THNN/generic/THNN.h
@@ -841,7 +841,6 @@ TH_API void THNN_(SpatialMaxPooling_updateOutput)(
           int kW, int kH,
           int dW, int dH,
           int padW, int padH,
-          int dilationW, int dilationH,
           bool ceil_mode);
 TH_API void THNN_(SpatialMaxPooling_updateGradInput)(
           THNNState *state,
@@ -852,6 +851,27 @@ TH_API void THNN_(SpatialMaxPooling_updateGradInput)(
           int kW, int kH,
           int dW, int dH,
           int padW, int padH,
+          bool ceil_mode);
+
+TH_API void THNN_(SpatialDilatedMaxPooling_updateOutput)(
+          THNNState *state,
+          THTensor *input,
+          THTensor *output,
+          THTensor *indices,
+          int kW, int kH,
+          int dW, int dH,
+          int padW, int padH,
+          int dilationW, int dilationH,
+          bool ceil_mode);
+TH_API void THNN_(SpatialDilatedMaxPooling_updateGradInput)(
+          THNNState *state,
+          THTensor *input,
+          THTensor *gradOutput,
+          THTensor *gradInput,
+          THTensor *indices,
+          int kW, int kH,
+          int dW, int dH,
+          int padW, int padH,
           int dilationW, int dilationH,
           bool ceil_mode);
 
diff --git a/lib/THNN/init.c b/lib/THNN/init.c
index 739706c..d26c509 100644
--- a/lib/THNN/init.c
+++ b/lib/THNN/init.c
@@ -139,6 +139,9 @@
 #include "generic/SpatialMaxPooling.c"
 #include "THGenerateFloatTypes.h"
 
+#include "generic/SpatialDilatedMaxPooling.c"
+#include "THGenerateFloatTypes.h"
+
 #include "generic/SpatialMaxUnpooling.c"
 #include "THGenerateFloatTypes.h"