Fix integer overflow in remainder operator (#5906)

vedanuj · soumith · commit 08b1324ec260 · 2018-03-20T22:05:34.000-04:00
* Fix integer overflow in remainder

* Fix remainder operator in CUDA

* Add tests for remainder integer overflow

* Add has_different_sign static function
diff --git a/aten/src/TH/generic/THTensorMath.c b/aten/src/TH/generic/THTensorMath.c
@@ -1052,6 +1052,10 @@ void THTensor_(fmod)(THTensor *r_, THTensor *t, real value)
   }
 }
 
+static inline bool has_different_sign(real a, real b) {
+  return (a < 0) != (b < 0);
+}
+
 void THTensor_(remainder)(THTensor *r_, THTensor *t, real value)
 {
   THTensor_(resizeAs)(r_, t);
@@ -1070,7 +1074,7 @@ void THTensor_(remainder)(THTensor *r_, THTensor *t, real value)
 #else
       // There is no NAN for integers
       rp[i] = tp[i] % value;
-      if (rp[i] * value < 0)
+      if (has_different_sign(rp[i], value))
         rp[i] += value;
 #endif
     }
@@ -1085,7 +1089,7 @@ void THTensor_(remainder)(THTensor *r_, THTensor *t, real value)
 #else
       // There is no NAN for integers
       TH_TENSOR_APPLY2_OMP(r_Size, r_Contig, tContig, real, r_, real, t, *r__data = *t_data % value;
-                                        if (*r__data * value < 0) *r__data += value;);
+                                        if (has_different_sign(*r__data, value)) *r__data += value;);
 #endif
     }
 #else
@@ -1098,7 +1102,7 @@ void THTensor_(remainder)(THTensor *r_, THTensor *t, real value)
 #else
     // There is no NAN for integers
     TH_TENSOR_APPLY2(real, r_, real, t, *r__data = *t_data % value;
-                                          if (*r__data * value < 0) *r__data += value;);
+                                          if (has_different_sign(*r__data, value)) *r__data += value;);
 #endif
   }
 }
diff --git a/aten/src/THC/THCNumerics.cuh b/aten/src/THC/THCNumerics.cuh
@@ -28,6 +28,11 @@ static inline __host__ __device__ scalar_t powi(scalar_t a, scalar_t b) {
   return result;
 }
 
+template <typename scalar_t>
+static inline __host__ __device__ bool has_different_sign(scalar_t a, scalar_t b) {
+  return (a < 0) != (b < 0);
+}
+
 template <>
 struct THCNumerics<uint8_t> {
   static inline __host__ __device__ uint8_t min() { return 0; }
diff --git a/aten/src/THC/THCTensorMathPairwise.cu b/aten/src/THC/THCTensorMathPairwise.cu
@@ -245,14 +245,14 @@ struct TensorRemainderOp {
   TensorRemainderOp(T v) : val(v) {}
   __device__ __forceinline__ void operator()(T* out, T* in) {
     *out = *in % val;
-    if ((*out * val) < 0){
+    if (has_different_sign<T>(*out, val)){
       *out += val;
     }
   }
 
   __device__ __forceinline__ void operator()(T* v) {
     *v = *v % val;
-    if ((*v * val) < 0){
+    if (has_different_sign<T>(*v, val)){
       *v += val;
     }
   }
diff --git a/test/test_cuda.py b/test/test_cuda.py
@@ -1434,6 +1434,9 @@ def test_min_max_inits(self):
     def test_int_pow(self):
         TestTorch._test_int_pow(self, lambda x: x.cuda())
 
+    def test_remainder_overflow(self):
+        TestTorch._test_remainder_overflow(self, dtype=torch.cuda.int64)
+
     def test_var(self):
         cpu_tensor = torch.randn(2, 3, 3)
         gpu_tensor = cpu_tensor.cuda()
diff --git a/test/test_torch.py b/test/test_torch.py
@@ -780,6 +780,19 @@ def test_remainder(self):
         long_res1 = long_m1.clone()
         long_res1.remainder_(long_qs.unsqueeze(0).expand_as(long_res1))
 
+    @staticmethod
+    def _test_remainder_overflow(self, dtype=torch.int64):
+        # Check Integer Overflows
+        x = torch.tensor(23500, dtype=dtype)
+        q = 392486996410368
+        self.assertEqual(x % q, x)
+        self.assertEqual(-x % q, q - x)
+        self.assertEqual(x % -q, x - q)
+        self.assertEqual(-x % -q, -x)
+
+    def test_remainder_overflow(self):
+        self._test_remainder_overflow(self, dtype=torch.int64)
+
     def test_mm(self):
         # helper function
         def matrixmultiply(mat1, mat2):

Original file line number	Diff line number	Diff line change
`@@ -1052,6 +1052,10 @@ void THTensor_(fmod)(THTensor r_, THTensor t, real value)`
`1052`	`1052`	`}`
`1053`	`1053`	`}`
`1054`	`1054`
	`1055`	`+static inline bool has_different_sign(real a, real b) {`
	`1056`	`+ return (a < 0) != (b < 0);`
	`1057`	`+}`
	`1058`	`+`
`1055`	`1059`	`void THTensor_(remainder)(THTensor r_, THTensor t, real value)`
`1056`	`1060`	`{`
`1057`	`1061`	`THTensor_(resizeAs)(r_, t);`
`@@ -1070,7 +1074,7 @@ void THTensor_(remainder)(THTensor r_, THTensor t, real value)`
`1070`	`1074`	`#else`
`1071`	`1075`	`// There is no NAN for integers`
`1072`	`1076`	`rp[i] = tp[i] % value;`
`1073`		`- if (rp[i] * value < 0)`
	`1077`	`+ if (has_different_sign(rp[i], value))`
`1074`	`1078`	`rp[i] += value;`
`1075`	`1079`	`#endif`
`1076`	`1080`	`}`
`@@ -1085,7 +1089,7 @@ void THTensor_(remainder)(THTensor r_, THTensor t, real value)`
`1085`	`1089`	`#else`
`1086`	`1090`	`// There is no NAN for integers`
`1087`	`1091`	`TH_TENSOR_APPLY2_OMP(r_Size, r_Contig, tContig, real, r_, real, t, r__data = t_data % value;`
`1088`		`- if (r__data value < 0) *r__data += value;);`
	`1092`	`+ if (has_different_sign(r__data, value)) r__data += value;);`
`1089`	`1093`	`#endif`
`1090`	`1094`	`}`
`1091`	`1095`	`#else`
`@@ -1098,7 +1102,7 @@ void THTensor_(remainder)(THTensor r_, THTensor t, real value)`
`1098`	`1102`	`#else`
`1099`	`1103`	`// There is no NAN for integers`
`1100`	`1104`	`TH_TENSOR_APPLY2(real, r_, real, t, r__data = t_data % value;`
`1101`		`- if (r__data value < 0) *r__data += value;);`
	`1105`	`+ if (has_different_sign(r__data, value)) r__data += value;);`
`1102`	`1106`	`#endif`
`1103`	`1107`	`}`
`1104`	`1108`	`}`
Original file line number	Diff line number	Diff line change
`@@ -245,14 +245,14 @@ struct TensorRemainderOp {`
`245`	`245`	`TensorRemainderOp(T v) : val(v) {}`
`246`	`246`	`__device__ __forceinline__ void operator()(T* out, T* in) {`
`247`	`247`	`out = in % val;`
`248`		`- if ((out val) < 0){`
	`248`	`+ if (has_different_sign<T>(*out, val)){`
`249`	`249`	`*out += val;`
`250`	`250`	`}`
`251`	`251`	`}`
`252`	`252`
`253`	`253`	`__device__ __forceinline__ void operator()(T* v) {`
`254`	`254`	`v = v % val;`
`255`		`- if ((v val) < 0){`
	`255`	`+ if (has_different_sign<T>(*v, val)){`
`256`	`256`	`*v += val;`
`257`	`257`	`}`
`258`	`258`	`}`