Skip to content

Commit

Permalink
THAlloc trigger GC when heap size exceeds soft max
Browse files Browse the repository at this point in the history
  • Loading branch information
adamlerer committed Jul 20, 2015
1 parent 65e5544 commit 2fbcdf2
Showing 2 changed files with 65 additions and 4 deletions.
67 changes: 63 additions & 4 deletions lib/TH/THGeneral.c
Original file line number Diff line number Diff line change
@@ -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);
}

2 changes: 2 additions & 0 deletions lib/TH/THGeneral.h.in
Original file line number Diff line number Diff line change
@@ -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__)

0 comments on commit 2fbcdf2

Please sign in to comment.