Don't have malloc-free pairs that cross DLL boundaries. (#17302)

Summary:
See https://blogs.msdn.microsoft.com/oldnewthing/20060915-04/?p=29723
for more background on this requirement on Windows.

Fixes #17239.

Signed-off-by: Edward Z. Yang <ezyang@fb.com>

cc xkszltl peterjc123
Pull Request resolved: https://github.com/pytorch/pytorch/pull/17302

Differential Revision: D14150067

Pulled By: ezyang

fbshipit-source-id: 9dc16ca781ff17515b8df1bb55492477e7843d4c

3 files changed, 13 insertions, 14 deletions

diff --git a/aten/src/TH/THGeneral.cpp b/aten/src/TH/THGeneral.cpp
index e2d6942dcf..8df8ff5dce 100644
--- a/aten/src/TH/THGeneral.cpp
+++ b/aten/src/TH/THGeneral.cpp
@@ -196,7 +196,7 @@ void* THRealloc(void *ptr, ptrdiff_t size)
 
 void THFree(void *ptr)
 {
-  free(ptr);
+  c10::free_cpu(ptr);
 }
 
 double THLog10(const double x)
diff --git a/c10/core/CPUAllocator.cpp b/c10/core/CPUAllocator.cpp
index 41ded8ca27..e418cb8b0d 100644
--- a/c10/core/CPUAllocator.cpp
+++ b/c10/core/CPUAllocator.cpp
@@ -76,6 +76,14 @@ void* alloc_cpu(size_t nbytes) {
   return data;
 }
 
+void free_cpu(void* data) {
+#ifdef _MSC_VER
+  _aligned_free(data);
+#else
+  free(data);
+#endif
+}
+
 // A virtual struct that is used to report C10's memory allocation and
 // deallocation status
 class C10_API MemoryAllocationReporter {
@@ -99,32 +107,22 @@ struct C10_API DefaultCPUAllocator final : at::Allocator {
       getMemoryAllocationReporter().New(data, nbytes);
       return {data, data, &ReportAndDelete, at::Device(at::DeviceType::CPU)};
     }
-    return {data, data, &Delete, at::Device(at::DeviceType::CPU)};
+    return {data, data, &free_cpu, at::Device(at::DeviceType::CPU)};
   }
 
-#ifdef _MSC_VER
-  static void Delete(void* data) {
-    _aligned_free(data);
-  }
-#else
-  static void Delete(void* data) {
-    free(data);
-  }
-#endif
-
   static void ReportAndDelete(void* ptr) {
     if (!ptr) {
       return;
     }
     getMemoryAllocationReporter().Delete(ptr);
-    Delete(ptr);
+    free_cpu(ptr);
   }
 
   at::DeleterFnPtr raw_deleter() const override {
     if (FLAGS_caffe2_report_cpu_memory_usage) {
       return &ReportAndDelete;
     }
-    return &Delete;
+    return &free_cpu;
   }
 
  protected:
diff --git a/c10/core/CPUAllocator.h b/c10/core/CPUAllocator.h
index 66e940e4fa..17440df2ec 100644
--- a/c10/core/CPUAllocator.h
+++ b/c10/core/CPUAllocator.h
@@ -28,6 +28,7 @@ C10_API void NoDelete(void*);
 C10_API void memset_junk(void* data, size_t num);
 
 C10_API void* alloc_cpu(size_t nbytes);
+C10_API void free_cpu(void* data);
 
 // Get the CPU Alloctor.
 C10_API at::Allocator* GetCPUAllocator();