Safety check for negative alloc_cpu() attempt (#17071)

Summary:
Some legacy TH code was relying on alloc to throw when called with negative number!!! E.g. `torch.linspace(0, 1, -1)`. And it breaks ASAN build. I still believe alloc should receive size_t, but I added a safety enforce inside.

It should fix ASAN. I'll follow up with a proper fix for empty_cpu (which is probably the right place to do it) separately
Pull Request resolved: https://github.com/pytorch/pytorch/pull/17071

Differential Revision: D14074157

Pulled By: dzhulgakov

fbshipit-source-id: 3ed3bdb873e446edecb558e1df491310fd7179e3

diff --git a/c10/core/CPUAllocator.cpp b/c10/core/CPUAllocator.cpp
index 10cff34..41ded8c 100644 (file)
--- a/c10/core/CPUAllocator.cpp
+++ b/c10/core/CPUAllocator.cpp
@@ -41,6 +41,11 @@ void* alloc_cpu(size_t nbytes) {
   if (nbytes == 0) {
     return nullptr;
   }
+  // We might have clowny upstream code that tries to alloc a negative number
+  // of bytes. Let's catch it early.
+  CAFFE_ENFORCE(
+    ((ptrdiff_t)nbytes) >= 0,
+    "alloc_cpu() seems to have been called with negative number: ", nbytes);
 
   void* data;
 #ifdef __ANDROID__