Merge pull request #1080 from torch/unpoolfix

fixing OpenMP longjmp bugs in *MaxUnpooling

2 files changed, 61 insertions, 60 deletions

diff --git a/lib/THNN/generic/SpatialMaxUnpooling.c b/lib/THNN/generic/SpatialMaxUnpooling.c
index 1b7b517..3205386 100644
--- a/lib/THNN/generic/SpatialMaxUnpooling.c
+++ b/lib/THNN/generic/SpatialMaxUnpooling.c
@@ -4,13 +4,13 @@
 
 static void THNN_(SpatialMaxUnpooling_updateOutput_frame)(real *input_p, real *output_p,
                                                       THIndex_t *ind_p,
-                                                      long nslices,
-                                                      long iwidth, long iheight,
-                                                      long owidth, long oheight)
+                                                      int nslices,
+                                                      int iwidth, int iheight,
+                                                      int owidth, int oheight)
 {
-  long k;
+  int k;
   int has_error = 0;
-  long error_index;
+  THIndex_t error_index;
 #pragma omp parallel for private(k)
   for (k = 0; k < nslices; k++)
   {
@@ -18,7 +18,8 @@ static void THNN_(SpatialMaxUnpooling_updateOutput_frame)(real *input_p, real *o
     real *input_p_k = input_p + k*iwidth*iheight;
     THIndex_t *ind_p_k = ind_p + k*iwidth*iheight;
 
-    long i, j, maxp;
+    int i, j;
+    THIndex_t maxp;
     for(i = 0; i < iheight; i++)
     {
       for(j = 0; j < iwidth; j++)
@@ -37,7 +38,7 @@ static void THNN_(SpatialMaxUnpooling_updateOutput_frame)(real *input_p, real *o
     }
   }
   if (has_error) {
-    THError("found an invalid max index %ld (output volumes are of size %ldx%ld)",
+    THError("found an invalid max index %ld (output volumes are of size %dx%d)",
         error_index, oheight, owidth);
   }
 }
@@ -98,7 +99,7 @@ void THNN_(SpatialMaxUnpooling_updateOutput)(
   }
   else
   {
-    long p;
+    int p;
 
     THTensor_(resize4d)(output, nbatch, nslices, oheight, owidth);
     THTensor_(zero)(output);
@@ -107,14 +108,15 @@ void THNN_(SpatialMaxUnpooling_updateOutput)(
     output_data = THTensor_(data)(output);
     indices_data = THIndexTensor_(data)(indices);
 
-#pragma omp parallel for private(p)
     for (p = 0; p < nbatch; p++)
     {
-      THNN_(SpatialMaxUnpooling_updateOutput_frame)(input_data+p*nslices*iwidth*iheight, output_data+p*nslices*owidth*oheight,
-                                                indices_data+p*nslices*iwidth*iheight,
-                                                nslices,
-                                                iwidth, iheight,
-                                                owidth, oheight);
+      THNN_(SpatialMaxUnpooling_updateOutput_frame)(
+						    input_data+p*nslices*iwidth*iheight,
+						    output_data+p*nslices*owidth*oheight,
+						    indices_data+p*nslices*iwidth*iheight,
+						    nslices,
+						    iwidth, iheight,
+						    owidth, oheight);
     }
   }
 
@@ -125,11 +127,11 @@ void THNN_(SpatialMaxUnpooling_updateOutput)(
 
 static void THNN_(SpatialMaxUnpooling_updateGradInput_frame)(real *gradInput_p, real *gradOutput_p,
                                                          THIndex_t *ind_p,
-                                                         long nslices,
-                                                         long iwidth, long iheight,
-                                                         long owidth, long oheight)
+                                                         int nslices,
+                                                         int iwidth, int iheight,
+                                                         int owidth, int oheight)
 {
-  long k;
+  int k;
 #pragma omp parallel for private(k)
   for (k = 0; k < nslices; k++)
   {
@@ -137,14 +139,15 @@ static void THNN_(SpatialMaxUnpooling_updateGradInput_frame)(real *gradInput_p,
     real *gradOutput_p_k = gradOutput_p + k*owidth*oheight;
     THIndex_t *ind_p_k = ind_p + k*iwidth*iheight;
 
-    long i, j, maxp;
+    int i, j;
+    THIndex_t maxp;
     for(i = 0; i < iheight; i++)
     {
       for(j = 0; j < iwidth; j++)
       {
         maxp = ind_p_k[i*iwidth + j] - TH_INDEX_BASE; /* retrieve position of max */
-        if(maxp<0 || maxp>=owidth*oheight){
-            THError("invalid max index %d, owidth= %d, oheight= %d",maxp,owidth,oheight);
+        if(maxp < 0 || maxp >= owidth * oheight) {
+            THError("invalid max index %ld, owidth= %d, oheight= %d", maxp, owidth, oheight);
         }
         gradInput_p_k[i*iwidth + j] = gradOutput_p_k[maxp]; /* update gradient */
       }
@@ -193,7 +196,7 @@ void THNN_(SpatialMaxUnpooling_updateGradInput)(
 
   if(owidth!=gradOutput->size[dimw] || oheight!=gradOutput->size[dimh]){
     THError("Inconsistent gradOutput size. oheight= %d, owidth= %d, gradOutput: %dx%d",
-	    oheight, owidth,gradOutput->size[dimh],gradOutput->size[dimw]);
+	    oheight, owidth, gradOutput->size[dimh], gradOutput->size[dimw]);
   }
 
   /* get raw pointers */
@@ -212,8 +215,7 @@ void THNN_(SpatialMaxUnpooling_updateGradInput)(
   }
   else
   {
-    long p;
-#pragma omp parallel for private(p)
+    int p;
     for (p = 0; p < nbatch; p++)
     {
       THNN_(SpatialMaxUnpooling_updateGradInput_frame)(gradInput_data+p*nslices*iwidth*iheight, gradOutput_data+p*nslices*owidth*oheight,
diff --git a/lib/THNN/generic/VolumetricMaxUnpooling.c b/lib/THNN/generic/VolumetricMaxUnpooling.c
index d08d251..d9d9e59 100644
--- a/lib/THNN/generic/VolumetricMaxUnpooling.c
+++ b/lib/THNN/generic/VolumetricMaxUnpooling.c
@@ -57,13 +57,13 @@ static void THNN_(VolumetricMaxUnpooling_updateOutput_frame)(
           real *input_p,
           real *output_p,
           THIndex_t *ind_p,
-          long nslices,
-          long iT,
-          long iW,
-          long iH,
-          long oT,
-          long oW,
-          long oH,
+          int nslices,
+          int iT,
+          int iW,
+          int iH,
+          int oT,
+          int oW,
+          int oH,
           int dT,
           int dW,
           int dH,
@@ -71,24 +71,23 @@ static void THNN_(VolumetricMaxUnpooling_updateOutput_frame)(
           int pW,
           int pH)
 {
-  long k;
+  int k;
   int has_error = 0;
-  long error_index;
+  THIndex_t error_index;
 #pragma omp parallel for private(k)
   for (k = 0; k < nslices; k++)
   {
-    long ti, i, j, maxz, maxy, maxx;
+    int ti, i, j, maxz, maxy, maxx;
     for (ti = 0; ti < iT; ti++)
     {
       for (i = 0; i < iH; i++)
       {
         for (j = 0; j < iW; j++)
         {
-          long start_t = ti * dT - pT;
-          long start_h = i * dH - pH;
-          long start_w = j * dW - pW;
+          int start_t = ti * dT - pT;
+          int start_h = i * dH - pH;
+          int start_w = j * dW - pW;
 
-          //real *output_p_k = output_p + k*oT*oW*oH + ti*oW*oH*dT + i*oW*dH + j*dW;
           real *input_p_k = input_p + k*iT*iW*iH + ti*iW*iH + i*iW + j;
           THIndex_t *ind_p_k = ind_p + k*iT*iW*iH + ti*iW*iH + i*iW + j;
 
@@ -96,8 +95,9 @@ static void THNN_(VolumetricMaxUnpooling_updateOutput_frame)(
           maxy = ((unsigned char*)(ind_p_k))[1];
           maxx = ((unsigned char*)(ind_p_k))[2];
 
-          size_t idx = k*oT*oW*oH + oH*oW*(start_t+maxz) + oW*(start_h+maxy) + (start_w+maxx);
-          if (start_t+maxz<0 || start_h+maxy<0 || start_w+maxx<0 || start_t+maxz>=oT || start_h+maxy>=oH || start_w+maxx>=oW)
+          THIndex_t idx = k*oT*oW*oH + oH*oW*(start_t+maxz) + oW*(start_h+maxy) + (start_w+maxx);
+          if (start_t+maxz<0 || start_h+maxy<0 || start_w+maxx<0 || start_t+maxz>=oT
+	      || start_h+maxy>=oH || start_w+maxx>=oW)
           {
 #pragma omp critical
             {
@@ -113,7 +113,7 @@ static void THNN_(VolumetricMaxUnpooling_updateOutput_frame)(
   }
   if (has_error) {
     THError(
-        "found an invalid max index %ld (output volumes are of size %ldx%ldx%ld)",
+        "found an invalid max index %ld (output volumes are of size %dx%dx%d)",
         error_index, oT, oH, oW
     );
   }
@@ -189,7 +189,7 @@ void THNN_(VolumetricMaxUnpooling_updateOutput)(
   }
   else
   {
-    long p;
+    int p;
 
     THTensor_(resize5d)(output, nbatch, nslices, oT, oH, oW);
     THTensor_(zero)(output);
@@ -198,7 +198,6 @@ void THNN_(VolumetricMaxUnpooling_updateOutput)(
     output_data = THTensor_(data)(output);
     indices_data = THIndexTensor_(data)(indices);
 
-#pragma omp parallel for private(p)
     for (p = 0; p < nbatch; p++)
     {
       THNN_(VolumetricMaxUnpooling_updateOutput_frame)(
@@ -223,13 +222,13 @@ static void THNN_(VolumetricMaxUnpooling_updateGradInput_frame)(
           real *gradInput_p,
           real *gradOutput_p,
           THIndex_t *ind_p,
-          long nslices,
-          long iT,
-          long iW,
-          long iH,
-          long oT,
-          long oW,
-          long oH,
+          int nslices,
+          int iT,
+          int iW,
+          int iH,
+          int oT,
+          int oW,
+          int oH,
           int dT,
           int dW,
           int dH,
@@ -237,37 +236,38 @@ static void THNN_(VolumetricMaxUnpooling_updateGradInput_frame)(
           int pW,
           int pH)
 {
-  long k;
+  int k;
 #pragma omp parallel for private(k)
   for (k = 0; k < nslices; k++)
   {
-    long ti, i, j, maxz, maxy, maxx;
+    int ti, i, j, maxz, maxy, maxx;
     for (ti = 0; ti < iT; ti++)
     {
       for (i = 0; i < iH; i++)
       {
         for (j = 0; j < iW; j++)
         {
-          long start_t = ti * dT - pT;
-          long start_h = i * dH - pH;
-          long start_w = j * dW - pW;
+          int start_t = ti * dT - pT;
+          int start_h = i * dH - pH;
+          int start_w = j * dW - pW;
 
           real *gradInput_p_k = gradInput_p + k*iT*iW*iH + ti*iW*iH + i*iW + j;
-          //real *gradOutput_p_k = gradOutput_p + k*oT*oW*oH + ti*oW*oH*dT + i*oW*dH + j*dW;
           THIndex_t *ind_p_k = ind_p + k*iT*iW*iH + ti*iW*iH + i*iW + j;
 
           maxz = ((unsigned char*)(ind_p_k))[0]; /* retrieve position of max */
           maxy = ((unsigned char*)(ind_p_k))[1];
           maxx = ((unsigned char*)(ind_p_k))[2];
 
-          if (start_t+maxz<0 || start_h+maxy<0 || start_w+maxx<0 || start_t+maxz>=oT || start_h+maxy>=oH || start_w+maxx>=oW)
+          if (start_t+maxz<0 || start_h+maxy<0 || start_w+maxx<0
+	      || start_t+maxz>=oT || start_h+maxy>=oH || start_w+maxx>=oW)
           {
             THError(
               "invalid max index z= %d, y= %d, x= %d, oT= %d, oW= %d, oH= %d",
               start_t+maxz, start_h+maxy, start_w+maxx, oT, oW, oH
             );
           }
-          *gradInput_p_k = gradOutput_p[k*oT*oW*oH + oH*oW*(start_t+maxz) + oW*(start_h+maxy) + (start_w+maxx)]; /* update gradient */
+          *gradInput_p_k = gradOutput_p[k*oT*oW*oH + oH*oW*(start_t+maxz)
+					+ oW*(start_h+maxy) + (start_w+maxx)]; /* update gradient */
         }
       }
     }
@@ -349,8 +349,7 @@ void THNN_(VolumetricMaxUnpooling_updateGradInput)(
   }
   else
   {
-    long p;
-#pragma omp parallel for private(p)
+    int p;
     for (p = 0; p < nbatch; p++)
     {
       THNN_(VolumetricMaxUnpooling_updateGradInput_frame)(