Fixed double memory accesses of several pointwise operations. (#5068)

yongjik · soumith · commit e90699b8626b · 2018-02-07T09:21:41.000-08:00
Because nvcc does not know that in/out pointers do not alias each other,
if we assign a value to *out and then use *in again, the kernel has to
emit a write to *out and then another read from *in.

(Affected kernels become marginally faster after the fix.)
diff --git a/torch/lib/THC/THCTensorMathPointwise.cuh b/torch/lib/THC/THCTensorMathPointwise.cuh
@@ -273,13 +273,13 @@ struct TensorPowOp {
     } else if (StaticExp == 2) {
       *out = THCNumerics<T>::mul(*in, *in);
     } else if (StaticExp == 3) {
-      *out = THCNumerics<T>::mul(*in, *in);
-      *out = THCNumerics<T>::mul(*out, *in);
+      T square = THCNumerics<T>::mul(*in, *in);
+      *out = THCNumerics<T>::mul(square, *in);
     } else if (StaticExp == -1) {
       *out = THCNumerics<T>::cinv(*in);
     } else if (StaticExp == -2) {
-      *out = THCNumerics<T>::mul(*in, *in);
-      *out = THCNumerics<T>::cinv(*out);
+      T square = THCNumerics<T>::mul(*in, *in);
+      *out = THCNumerics<T>::cinv(square);
     } else {
       *out = THCNumerics<T>::pow(*in, val);
     }
@@ -295,8 +295,8 @@ struct TensorPowOp {
     } else if (StaticExp == -1) {
       *v = THCNumerics<T>::cinv(*v);
     } else if (StaticExp == -2) {
-      *v = THCNumerics<T>::mul(*v, *v);
-      *v = THCNumerics<T>::cinv(*v);
+      T square = THCNumerics<T>::mul(*v, *v);
+      *v = THCNumerics<T>::cinv(square);
     } else {
       *v = THCNumerics<T>::pow(*v, val);
     }
@@ -402,17 +402,19 @@ struct TensorDivOp<half> {
 template <typename T>
 struct TensorCRemainderOp {
   __device__ __forceinline__ void operator()(T* out, T* in) {
-    *out =  *out % *in;
-    if ((*out * *in)<0){
-      *out += *in;
+    T val =  *out % *in;
+    if ((val * *in)<0){
+      val += *in;
     }
+    *out = val;
   }
 
   __device__ __forceinline__ void operator()(T* out, T* in1, T* in2) {
-    *out = *in1 % *in2;
-    if ((*out * *in2)<0){
-      *out += *in2;
+    T val = *in1 % *in2;
+    if ((val * *in2)<0){
+      val += *in2;
     }
+    *out = val;
   }
 };
 
@@ -548,20 +550,24 @@ struct TensorCrossOp {
   TensorCrossOp(int64_t sx, int64_t sy, int64_t so) : sx(sx), sy(sy), so(so) {}
 
   __device__ __forceinline__ void operator()(T* out, T* x, T*y) {
-    out[0 * so] = THCNumerics<T>::sub(
+    T val0 = THCNumerics<T>::sub(
         THCNumerics<T>::mul(x[1 * sx], y[2 * sy]),
         THCNumerics<T>::mul(x[2 * sx], y[1 * sy])
     );
 
-    out[1 * so] = THCNumerics<T>::sub(
+    T val1 = THCNumerics<T>::sub(
         THCNumerics<T>::mul(x[2 * sx], y[0 * sy]),
         THCNumerics<T>::mul(x[0 * sx], y[2 * sy])
     );
 
-    out[2 * so] = THCNumerics<T>::sub(
+    T val2 = THCNumerics<T>::sub(
         THCNumerics<T>::mul(x[0 * sx], y[1 * sy]),
         THCNumerics<T>::mul(x[1 * sx], y[0 * sy])
     );
+
+    out[0 * so] = val0;
+    out[1 * so] = val1;
+    out[2 * so] = val2;
   }
 
   const int64_t sx, sy, so;
