Fix lua GC to support allocations across multiple threads (e.g. threads sharedserialize)

Author: adamlerer

THAtomic.c

@@ -101,3 +101,61 @@ int THAtomicCompareAndSwap(int volatile *a, int oldvalue, int newvalue)
  return 0;
 #endif
 }
+
+long THAtomicGetLong(long volatile *a)
+{
+#if defined(USE_C11_ATOMICS)
+ return atomic_load(a);
+#else
+ long value;
+ do {
+ value = *a;
+ } while (!THAtomicCompareAndSwapLong(a, value, value));
+ return value;
+#endif
+}
+
+long THAtomicAddLong(long volatile *a, long value)
+{
+#if defined(USE_C11_ATOMICS)
+ return atomic_fetch_add(a, value);
+#elif defined(USE_MSC_ATOMICS)
+ return _InterlockedExchangeAdd(a, value);
+#elif defined(USE_GCC_ATOMICS)
+ return __sync_fetch_and_add(a, value);
+#else
+ long oldvalue;
+ do {
+ oldvalue = *a;
+ } while (!THAtomicCompareAndSwapLong(a, oldvalue, (oldvalue + value)));
+ return oldvalue;
+#endif
+}
+
+long THAtomicCompareAndSwapLong(long volatile *a, long oldvalue, long newvalue)
+{
+#if defined(USE_C11_ATOMICS)
+ return atomic_compare_exchange_strong(a, &oldvalue, newvalue);
+#elif defined(USE_MSC_ATOMICS)
+ return (_InterlockedCompareExchange(a, newvalue, oldvalue) == oldvalue);
+#elif defined(USE_GCC_ATOMICS)
+ return __sync_bool_compare_and_swap(a, oldvalue, newvalue);
+#elif defined(USE_PTHREAD_ATOMICS)
+ long ret = 0;
+ pthread_mutex_lock(&ptm);
+ if(*a == oldvalue) {
+ *a = newvalue;
+ ret = 1;
+ }
+ pthread_mutex_unlock(&ptm);
+ return ret;
+#else
+#warning THAtomic is not thread safe
+ if(*a == oldvalue) {
+ *a = newvalue;
+ return 1;
+ }
+ else
+ return 0;
+#endif
+}

THAtomic.h

@@ -57,4 +57,28 @@ TH_API void THAtomicIncrementRef(int volatile *a);
 */
 TH_API int THAtomicDecrementRef(int volatile *a);
 
+
+
+/******************************************************************************
+ * functions for long type
+ ******************************************************************************/
+
+/*
+ * return *a
+*/
+TH_API long THAtomicGetLong(long volatile *a);
+
+/*
+ * *a += value,
+ * return previous *a
+*/
+TH_API long THAtomicAddLong(long volatile *a, long value);
+
+/*
+ * check if (*a == oldvalue)
+ * if true: set *a to newvalue, return 1
+ * if false: return 0
+*/
+TH_API long THAtomicCompareAndSwapLong(long volatile *a, long oldvalue, long newvalue);
+
 #endif

THGeneral.c

@@ -1,11 +1,11 @@
 #include "THGeneral.h"
+#include "THAtomic.h"
 
 #ifndef TH_HAVE_THREAD
 #define __thread
 #endif
 
-#if defined(TH_DISABLE_HEAP_TRACKING)
-#elif (defined(__unix) || defined(_WIN32))
+#if (defined(__unix) || defined(_WIN32))
 #include <malloc.h>
 #elif defined(__APPLE__)
 #include <malloc/malloc.h>
@@ -101,8 +101,10 @@ void THSetArgErrorHandler( void (*torchArgErrorHandlerFunction_)(int argNumber,
 
 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
+static long heapSize = 0;
+static __thread long heapSoftmax = 300000000; // 300MB, adjusted upward dynamically
+static const double heapSoftmaxGrowthThresh = 0.8; // grow softmax if >80% max after GC
+static const double heapSoftmaxGrowthFactor = 1.4; // grow softmax by 40%
 
 /* Optional hook for integrating with a garbage-collected frontend.
  *
@@ -121,9 +123,7 @@ void THSetGCHandler( void (*torchGCFunction_)(void *data), void *data )
 }
 
 static long getAllocSize(void *ptr) {
-#if defined(TH_DISABLE_HEAP_TRACKING)
- return 0;
-#elif defined(__unix)
+#if defined(__unix)
  return malloc_usable_size(ptr);
 #elif defined(__APPLE__)
  return malloc_size(ptr);
@@ -138,20 +138,29 @@ static long getAllocSize(void *ptr) {
  * (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) {
+static void maybeTriggerGC(long curHeapSize) {
+ if(torchGCFunction && curHeapSize > heapSoftmax ) {
  torchGCFunction(torchGCData);
- if(torchHeapSize > torchHeapSizeSoftMax * 0.8) {
- torchHeapSizeSoftMax = torchHeapSizeSoftMax * 1.4;
+ long newHeapSize = THAtomicGetLong(&heapSize);
+ if(newHeapSize > heapSoftmax * heapSoftmaxGrowthThresh) {
+ heapSoftmax = heapSoftmax * heapSoftmaxGrowthFactor;
  }
  }
 }
 
 // hooks into the TH heap tracking
 void THHeapUpdate(long size) {
- torchHeapSize += size;
- if (size > 0)
- maybeTriggerGC();
+ long newHeapSize = THAtomicAddLong(&heapSize, size) + size;
+
+# ifdef TH_CHECK_HEAP_UPDATE
+ if (newHeapSize < 0) {
+ THError("Torch heap size <0 ?");
+ }
+#endif
+
+ if (size > 0) {
+ maybeTriggerGC(newHeapSize);
+ }
 }
 
 static void* THAllocInternal(long size)