1 files changed, 51 insertions, 6 deletions

diff --git a/generic/Square.c b/generic/Square.c
index 409055d..97baee3 100644
--- a/generic/Square.c
+++ b/generic/Square.c
@@ -6,12 +6,32 @@ static int nn_(Square_updateOutput)(lua_State *L)
 {
   THTensor *input = luaT_checkudata(L, 2, torch_(Tensor_id));
   THTensor *output = luaT_getfieldcheckudata(L, 1, "output", torch_(Tensor_id));
-  
-  THTensor_(resizeAs)(output, input);
 
-  TH_TENSOR_APPLY2(real, output, real, input,	\
-		   *output_data = *input_data * *input_data;);
+  THTensor_(resizeAs)(output, input);
+  
+  if (input->nDimension == 1 || !THTensor_(isContiguous)(input) || !THTensor_(isContiguous)(output))
+  {
+    TH_TENSOR_APPLY2(real, output, real, input,   \
+         *output_data = (*input_data) * (*input_data););
+  }
+  else
+  {
+    real* output_data = THTensor_(data)(output);
+    real* input_data  = THTensor_(data)(input);
+    long k;
 
+#pragma omp parallel for private(k)
+    for (k = 0; k < input->size[0]; k++)
+    {
+      real* ptr_output = output_data + k*input->stride[0];
+      real* ptr_input  = input_data  + k*input->stride[0];
+      long i;
+      for (i = 0; i < input->stride[0]; i++)
+      {
+        ptr_output[i] = ptr_input[i]*ptr_input[i];
+      }
+    }
+  }
   return 1;
 }
 
@@ -23,9 +43,34 @@ static int nn_(Square_updateGradInput)(lua_State *L)
 
   THTensor_(resizeAs)(gradInput, input);
 
-  TH_TENSOR_APPLY3(real, gradInput, real, gradOutput, real, input, \
-		   *gradInput_data = 2.0 * (*gradOutput_data) * (*input_data););
+  if (input->nDimension == 1 || 
+      !THTensor_(isContiguous)(input) || 
+      !THTensor_(isContiguous)(gradOutput) ||
+      !THTensor_(isContiguous)(gradInput))
+  {
+    TH_TENSOR_APPLY3(real, gradInput, real, gradOutput, real, input,  \
+         *gradInput_data  = (*gradOutput_data) * (*input_data););
+  }
+  else
+  {
+    real* gradOutput_data = THTensor_(data)(gradOutput);
+    real* gradInput_data  = THTensor_(data)(gradInput);
+    real* input_data     = THTensor_(data)(input);
+    long k;
 
+#pragma omp parallel for private(k)
+    for (k = 0; k < input->size[0]; k++)
+    {
+      real* ptr_gradOutput = gradOutput_data + k*input->stride[0];
+      real* ptr_gradInput  = gradInput_data  + k*input->stride[0];
+      real* ptr_input     = input_data     + k*input->stride[0];
+      long i;
+      for (i = 0; i < input->stride[0]; i++)
+      {
+        ptr_gradInput[i] = 2.0 * ptr_gradOutput[i] * ptr_input[i];
+      }
+    }
+  }
   return 1;
 }