Fix potential UB when input is empty (pytorch#6242)

If the source and result tensors are empty, arr_in and arr_out may be null (and size will be 0). This previously called memcpy(null, null, 0), which is UB according to http://en.cppreference.com/w/cpp/string/byte/memcpy. Note that either one of these changes would be sufficient. (Detected by UBSan)
runtian-zhou · Apr 4, 2018 · fd580ce · fd580ce
1 parent 2f0bb19
commit fd580ce
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 4 deletions.
diff --git a/aten/src/ATen/native/UnaryOps.cpp b/aten/src/ATen/native/UnaryOps.cpp
@@ -28,7 +28,9 @@ Tensor& _ ## op ## _out_cuda(Tensor& result, const Tensor& self) {            \
 Tensor& _ ## op ## _out_cpu(Tensor& result, const Tensor& self) {             \
   if (result.is_contiguous() && self.is_contiguous()) {                       \
     result.resize_(self.sizes());                                             \
-    op ## Impl(result, self);                                                 \
+    if (result.numel() > 0) {                                                 \
+      op ## Impl(result, self);                                               \
+    }                                                                         \
     return result;                                                            \
   }                                                                           \
   return at::_ ## op ## _out(result, self);                                   \

diff --git a/aten/src/ATen/native/cpu/UnaryOpsKernel.cpp b/aten/src/ATen/native/cpu/UnaryOpsKernel.cpp
@@ -21,9 +21,11 @@ static void unary_kernel(scalar_t* arr_out, const scalar_t* arr_in, int64_t size
     value.store(arr_out + k);
   }
   auto leftover = size - k;
-  Vec a;
-  a.load_partial(arr_in + k, leftover);
-  func(a).store_partial(arr_out + k, leftover);
+  if (leftover > 0) {
+    Vec a;
+    a.load_partial(arr_in + k, leftover);
+    func(a).store_partial(arr_out + k, leftover);
+  }
 }
 
 template <class scalar_t, class F>