THAlloc trigger GC when heap size exceeds soft max

Author: adamlerer

THGeneral.c

@@ -3,6 +3,14 @@
 #ifndef TH_HAVE_THREAD
 #define __thread
 #endif
+
+#if defined(TH_DISABLE_HEAP_TRACKING)
+#elif (defined(__unix) || defined(_WIN32))
+#include <malloc.h>
+#elif defined(__APPLE__)
+#include <malloc/malloc.h>
+#endif
+
 /* Torch Error Handling */
 static void defaultTorchErrorHandlerFunction(const char *msg, void *data)
 {
@@ -13,9 +21,6 @@ static void defaultTorchErrorHandlerFunction(const char *msg, void *data)
 static __thread void (*torchErrorHandlerFunction)(const char *msg, void *data) = defaultTorchErrorHandlerFunction;
 static __thread void *torchErrorHandlerData;
 
-static __thread void (*torchGCFunction)(void *data) = NULL;
-static __thread void *torchGCData;
-
 void _THError(const char *file, const int line, const char *fmt, ...)
 {
  char msg[2048];
@@ -94,13 +99,62 @@ void THSetArgErrorHandler( void (*torchArgErrorHandlerFunction_)(int argNumber,
  torchArgErrorHandlerData = data;
 }
 
+static __thread void (*torchGCFunction)(void *data) = NULL;
+static __thread void *torchGCData;
+static __thread long torchHeapSize = 0;
+static __thread long torchHeapSizeSoftMax = 300000000; // 300MB, adjusted upward dynamically
+
+/* Optional hook for integrating with a garbage-collected frontend.
+ *
+ * If torch is running with a garbage-collected frontend (e.g. Lua),
+ * the GC isn't aware of TH-allocated memory so may not know when it
+ * needs to run. These hooks trigger the GC to run in two cases:
+ *
+ * (1) When a memory allocation (malloc, realloc, ...) fails
+ * (2) When the total TH-allocated memory hits a dynamically-adjusted
+ * soft maximum.
+ */
 void THSetGCHandler( void (*torchGCFunction_)(void *data), void *data )
 {
  torchGCFunction = torchGCFunction_;
  torchGCData = data;
 }
 
-void* THAllocInternal(long size)
+static long getAllocSize(void *ptr) {
+#if defined(TH_DISABLE_HEAP_TRACKING)
+ return 0;
+#elif defined(__unix)
+ return malloc_usable_size(ptr);
+#elif defined(__APPLE__)
+ return malloc_size(ptr);
+#elif defined(_WIN32)
+ return _msize(ptr);
+#else
+ return 0;
+#endif
+}
+
+/* (1) if the torch-allocated heap size exceeds the soft max, run GC
+ * (2) if post-GC heap size exceeds 80% of the soft max, increase the
+ * soft max by 40%
+ */
+static void maybeTriggerGC() {
+ if(torchGCFunction && torchHeapSize > torchHeapSizeSoftMax) {
+ torchGCFunction(torchGCData);
+ if(torchHeapSize > torchHeapSizeSoftMax * 0.8) {
+ torchHeapSizeSoftMax = torchHeapSizeSoftMax * 1.4;
+ }
+ }
+}
+
+// hooks into the TH heap tracking
+void THHeapUpdate(long size) {
+ torchHeapSize += size;
+ if (size > 0)
+ maybeTriggerGC();
+}
+
+static void* THAllocInternal(long size)
 {
  void *ptr;
 
@@ -119,6 +173,8 @@ void* THAllocInternal(long size)
  {
  ptr = malloc(size);
  }
+
+ THHeapUpdate(getAllocSize(ptr));
  return ptr;
 }
 
@@ -159,12 +215,14 @@ void* THRealloc(void *ptr, long size)
  if(size < 0)
  THError("$ Torch: invalid memory size -- maybe an overflow?");
 
+ THHeapUpdate(-getAllocSize(ptr));
  void *newptr = realloc(ptr, size);
 
  if(!newptr && torchGCFunction) {
  torchGCFunction(torchGCData);
  newptr = realloc(ptr, size);
  }
+ THHeapUpdate(getAllocSize(newptr ? newptr : ptr));
 
  if(!newptr)
  THError("$ Torch: not enough memory: you tried to reallocate %dGB. Buy new RAM!", size/1073741824);
@@ -174,6 +232,7 @@ void* THRealloc(void *ptr, long size)
 
 void THFree(void *ptr)
 {
+ THHeapUpdate(-getAllocSize(ptr));
  free(ptr);
 }
 

THGeneral.h.in

@@ -53,6 +53,8 @@ TH_API void* THAlloc(long size);
 TH_API void* THRealloc(void *ptr, long size);
 TH_API void THFree(void *ptr);
 TH_API void THSetGCHandler( void (*torchGCHandlerFunction)(void *data), void *data );
+// this hook should only be called by custom allocator functions
+TH_API void THHeapUpdate(long size);
 
 #define THError(...) _THError(__FILE__, __LINE__, __VA_ARGS__)
 #define THArgCheck(...) _THArgCheck(__FILE__, __LINE__, __VA_ARGS__)