Fix the root cause for _aligned_free issues on row_offset_ allocation (…

…pytorch#277)

Summary:
Pull Request resolved: pytorch#277

Previously in pytorch#275, we provided a quick fix on the test case (as row_offset_buf is declared but not used in the test case).

The root cause should be the inappropriate allocation API (not Windows compatible) in the Packing routines in FBGEMM. It turned out when we pass `row_offset` as `nullptr` in the Packing routine (e.g., `PackAWithQuantRowOffset`), we will use `fbgemmAlignedAlloc` to allocate the buffer (On Windows using `_aligned_malloc`) and set the flag `rowOffsetAllocatedHere` to be `true`. However, when we deallocate the `row_offset` buffer in the end, we use the normal `free` instead of `fbgemmAlignedFree`  (On Windows using `_aligned_free`). This is the root cause for the issue.

Reviewed By: jspark1105

Differential Revision: D19665851

fbshipit-source-id: 3d393aaf6a165d30cdedf9eed39db9b1d162d93d

include/fbgemm/Fbgemm.h

@@ -729,7 +729,7 @@ class FBGEMM_API PackAWithIm2Col
 
   ~PackAWithIm2Col() {
     if (rowOffsetAllocatedHere) {
-      free(row_offset_);
+      fbgemmAlignedFree(row_offset_);
     }
   }
 
@@ -819,7 +819,7 @@ class FBGEMM_API PackAWithRowOffset final
 
   ~PackAWithRowOffset() {
     if (rowOffsetAllocatedHere) {
-      free(row_offset_);
+      fbgemmAlignedFree(row_offset_);
     }
   }
 
@@ -911,7 +911,7 @@ class FBGEMM_API PackAWithQuantRowOffset final
 
   ~PackAWithQuantRowOffset() {
     if (rowOffsetAllocatedHere) {
-      free(row_offset_);
+      fbgemmAlignedFree(row_offset_);
     }
   }
 

src/Utils.cc

@@ -407,17 +407,18 @@ void* fbgemmAlignedAlloc(
     size_t align,
     size_t size,
     bool raiseException /*=false*/) {
-  void* aligned_mem;
+  void* aligned_mem = nullptr;
+  int ret;
 #ifdef _MSC_VER
   aligned_mem = _aligned_malloc(size, align);
+  ret = 0;
 #else
-  if (posix_memalign(&aligned_mem, align, size)) {
-    if (raiseException) {
-      throw std::bad_alloc();
-    }
-    return nullptr;
-  }
+  ret = posix_memalign(&aligned_mem, align, size);
 #endif
+  // Throw std::bad_alloc in the case of memory allocation failure.
+  if (raiseException || ret || aligned_mem == nullptr) {
+    throw std::bad_alloc();
+  }
   return aligned_mem;
 }