Improve Windows Compatibility(for lib/TH)

torch/lib/TH/THAllocator.c

@@ -39,7 +39,11 @@ struct THMapAllocatorContext_ {
   char *filename; /* file name */
   int flags;
   ptrdiff_t size; /* mapped size */
+#ifdef _WIN32
+  HANDLE handle;
+#else
   int fd;
+#endif
 };
 
 #define TH_ALLOC_ALIGNMENT 64
@@ -68,17 +72,25 @@ THMapAllocatorContext *THMapAllocatorContext_new(const char *filename, int flags
   }
   ctx->flags = flags;
   ctx->size = 0;
+#ifdef _WIN32
+  ctx->handle = INVALID_HANDLE_VALUE;
+#else
   ctx->fd = -1;
+#endif
 
   return ctx;
 }
 
 THMapAllocatorContext *THMapAllocatorContext_newWithFd(const char *filename, int fd, int flags)
 {
+#ifdef _WIN32
+  THError("THMapAllocatorContext_newWithFd is unsupported on Windows");
+#else
   THMapAllocatorContext *ctx = THMapAllocatorContext_new(filename, flags);
   ctx->fd = fd;
 
   return ctx;
+#endif
 }
 
 char * THMapAllocatorContext_filename(THMapAllocatorContext *ctx)
@@ -88,7 +100,11 @@ char * THMapAllocatorContext_filename(THMapAllocatorContext *ctx)
 
 int THMapAllocatorContext_fd(THMapAllocatorContext *ctx)
 {
+#ifdef _WIN32
+  THError("THMapAllocatorContext_fd is unsupported on Windows");
+#else
   return ctx->fd;
+#endif
 }
 
 ptrdiff_t THMapAllocatorContext_size(THMapAllocatorContext *ctx)
@@ -105,15 +121,49 @@ void THMapAllocatorContext_free(THMapAllocatorContext *ctx)
 
 static void *_map_alloc(void* ctx_, ptrdiff_t size)
 {
-  if (size == 0) {
+  if (size == 0)
     return NULL;
-  }
 
   THMapAllocatorContext *ctx = ctx_;
   void *data = NULL;
 
 #ifdef _WIN32
+  if (ctx->flags & TH_ALLOCATOR_MAPPED_SHAREDMEM)
+  {
+    char *filename;
+    LARGE_INTEGER hfilesz;
+
+    if (ctx->filename[0] == '/')
+      filename = ctx->filename + 1;
+    else
+      filename = ctx->filename;
+
+    hfilesz.QuadPart = size;
+
+    if (ctx->flags & TH_ALLOCATOR_MAPPED_EXCLUSIVE)
+    {
+      ctx->handle = CreateFileMapping(INVALID_HANDLE_VALUE, NULL, PAGE_READWRITE, hfilesz.HighPart, hfilesz.LowPart, filename);
+    }
+    else if (ctx->flags & TH_ALLOCATOR_MAPPED_NOCREATE)
+    {
+      ctx->handle = OpenFileMapping(FILE_MAP_ALL_ACCESS, FALSE, filename);
+    }
+    else
+    {
+      THError("Excpected either TH_ALLOCATOR_MAPPED_EXCLUSIVE or TH_ALLOCATOR_MAPPED_NOCREATE");
+    }
+
+    if (ctx->handle == NULL)
+      THError("Couldn't open shared file mapping: <%s>, error code: <%d>", filename, GetLastError());
+
+    ctx->size = size;
+    data = MapViewOfFile(ctx->handle, FILE_MAP_ALL_ACCESS, 0, 0, size);
+    if (!data)
+      THError("Couldn't map view of shared file <%s>, error code: <%d>", filename, GetLastError());
+  }
+  else
   {
+
     HANDLE hfile;
     HANDLE hmfile;
     LARGE_INTEGER hfilesz;
@@ -143,9 +193,7 @@ static void *_map_alloc(void* ctx_, ptrdiff_t size)
     }
 
     if (GetFileSizeEx(hfile, &hfilesz) == 0)
-    {
       THError("could not get file size: <%s>; error code: <%d>", ctx->filename, GetLastError());
-    }
 
     if(size > 0)
     {
@@ -342,6 +390,8 @@ static void THMapAllocator_free(void* ctx_, void* data) {
   THMapAllocatorContext *ctx = ctx_;
 
 #ifdef _WIN32
+  if ((ctx->flags & TH_ALLOCATOR_MAPPED_KEEPFD) || (ctx->flags & TH_ALLOCATOR_MAPPED_SHAREDMEM))
+    CloseHandle(ctx->handle);
   if(UnmapViewOfFile(data) == 0)
     THError("could not unmap the shared memory file");
 #else /* _WIN32 */

torch/lib/TH/THAllocator.h

@@ -22,7 +22,7 @@ typedef struct THAllocator {
 /* default malloc/free allocator. malloc and realloc raise an error (using
  * THError) on allocation failure.
  */
-extern THAllocator THDefaultAllocator;
+TH_API THAllocator THDefaultAllocator;
 
 /* file map allocator
  */
@@ -37,7 +37,7 @@ TH_API void THMapAllocatorContext_free(THMapAllocatorContext *ctx);
 TH_API void THRefcountedMapAllocator_incref(THMapAllocatorContext *ctx, void *data);
 TH_API int THRefcountedMapAllocator_decref(THMapAllocatorContext *ctx, void *data);
 
-extern THAllocator THMapAllocator;
-extern THAllocator THRefcountedMapAllocator;
+TH_API THAllocator THMapAllocator;
+TH_API THAllocator THRefcountedMapAllocator;
 
 #endif

torch/lib/TH/THAtomic.c

@@ -19,75 +19,73 @@
 static pthread_mutex_t ptm = PTHREAD_MUTEX_INITIALIZER;
 #endif
 
-void THAtomicSet(int volatile *a, int newvalue)
+void THAtomicSet(int32_t volatile *a, int32_t newvalue)
 {
 #if defined(USE_C11_ATOMICS)
   atomic_store(a, newvalue);
 #elif defined(USE_MSC_ATOMICS)
-  assert(sizeof(int) == sizeof(long));
-  _InterlockedExchange((long*)a, newvalue);
+  assert(sizeof(int) == sizeof(int32_t));
+  _InterlockedExchange((int32_t*)a, newvalue);
 #elif defined(USE_GCC_ATOMICS)
   __sync_lock_test_and_set(a, newvalue);
 #else
-  int oldvalue;
+  int32_t oldvalue;
   do {
     oldvalue = *a;
   } while (!THAtomicCompareAndSwap(a, oldvalue, newvalue));
 #endif
 }
 
-int THAtomicGet(int volatile *a)
+int THAtomicGet(int32_t volatile *a)
 {
 #if defined(USE_C11_ATOMICS)
   return atomic_load(a);
 #else
-  int value;
+  int32_t value;
   do {
     value = *a;
   } while (!THAtomicCompareAndSwap(a, value, value));
   return value;
 #endif
 }
 
-int THAtomicAdd(int volatile *a, int value)
+int THAtomicAdd(int32_t volatile *a, int32_t value)
 {
 #if defined(USE_C11_ATOMICS)
   return atomic_fetch_add(a, value);
 #elif defined(USE_MSC_ATOMICS)
-  assert(sizeof(int) == sizeof(long));
-  return _InterlockedExchangeAdd((long*)a, value);
+  return _InterlockedExchangeAdd((int32_t*)a, value);
 #elif defined(USE_GCC_ATOMICS)
   return __sync_fetch_and_add(a, value);
 #else
-  int oldvalue;
+  int32_t oldvalue;
   do {
     oldvalue = *a;
   } while (!THAtomicCompareAndSwap(a, oldvalue, (oldvalue + value)));
   return oldvalue;
 #endif
 }
 
-void THAtomicIncrementRef(int volatile *a)
+void THAtomicIncrementRef(int32_t volatile *a)
 {
   THAtomicAdd(a, 1);
 }
 
-int THAtomicDecrementRef(int volatile *a)
+int THAtomicDecrementRef(int32_t volatile *a)
 {
   return (THAtomicAdd(a, -1) == 1);
 }
 
-int THAtomicCompareAndSwap(int volatile *a, int oldvalue, int newvalue)
+int THAtomicCompareAndSwap(int32_t volatile *a, int32_t oldvalue, int32_t newvalue)
 {
 #if defined(USE_C11_ATOMICS)
   return atomic_compare_exchange_strong(a, &oldvalue, newvalue);
 #elif defined(USE_MSC_ATOMICS)
-  assert(sizeof(int) == sizeof(long));
-  return (_InterlockedCompareExchange((long*)a, (long)newvalue, (long)oldvalue) == (long)oldvalue);
+  return (_InterlockedCompareExchange((int32_t*)a, (int32_t)newvalue, (int32_t)oldvalue) == (int32_t)oldvalue);
 #elif defined(USE_GCC_ATOMICS)
   return __sync_bool_compare_and_swap(a, oldvalue, newvalue);
 #elif defined(USE_PTHREAD_ATOMICS)
-  int ret = 0;
+  int32_t ret = 0;
   pthread_mutex_lock(&ptm);
   if(*a == oldvalue) {
     *a = newvalue;
@@ -106,62 +104,62 @@ int THAtomicCompareAndSwap(int volatile *a, int oldvalue, int newvalue)
 #endif
 }
 
-void THAtomicSetLong(long volatile *a, long newvalue)
+void THAtomicSetLong(int64_t volatile *a, int64_t newvalue)
 {
 #if defined(USE_C11_ATOMICS)
   atomic_store(a, newvalue);
 #elif defined(USE_MSC_ATOMICS)
-  _InterlockedExchange(a, newvalue);
+  _InterlockedExchange64(a, newvalue);
 #elif defined(USE_GCC_ATOMICS)
   __sync_lock_test_and_set(a, newvalue);
 #else
-  long oldvalue;
+  int64_t oldvalue;
   do {
     oldvalue = *a;
   } while (!THAtomicCompareAndSwapLong(a, oldvalue, newvalue));
 #endif
 }
 
-long THAtomicGetLong(long volatile *a)
+int64_t THAtomicGetLong(int64_t volatile *a)
 {
 #if defined(USE_C11_ATOMICS)
   return atomic_load(a);
 #else
-  long value;
+  int64_t value;
   do {
     value = *a;
   } while (!THAtomicCompareAndSwapLong(a, value, value));
   return value;
 #endif
 }
 
-long THAtomicAddLong(long volatile *a, long value)
+int64_t THAtomicAddLong(int64_t volatile *a, int64_t value)
 {
 #if defined(USE_C11_ATOMICS)
   return atomic_fetch_add(a, value);
 #elif defined(USE_MSC_ATOMICS)
-  return _InterlockedExchangeAdd(a, value);
+  return _InterlockedExchangeAdd64(a, value);
 #elif defined(USE_GCC_ATOMICS)
   return __sync_fetch_and_add(a, value);
 #else
-  long oldvalue;
+  int64_t oldvalue;
   do {
     oldvalue = *a;
   } while (!THAtomicCompareAndSwapLong(a, oldvalue, (oldvalue + value)));
   return oldvalue;
 #endif
 }
 
-long THAtomicCompareAndSwapLong(long volatile *a, long oldvalue, long newvalue)
+int64_t THAtomicCompareAndSwapLong(int64_t volatile *a, int64_t oldvalue, int64_t 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);
+  return (_InterlockedCompareExchange64(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;
+  int64_t ret = 0;
   pthread_mutex_lock(&ptm);
   if(*a == oldvalue) {
     *a = newvalue;

torch/lib/TH/THAtomic.h

@@ -21,25 +21,25 @@
 /*
  * *a = newvalue
 */
-TH_API void THAtomicSet(int volatile *a, int newvalue);
+TH_API void THAtomicSet(int32_t volatile *a, int32_t newvalue);
 
 /*
  * return *a
 */
-TH_API int THAtomicGet(int volatile *a);
+TH_API int32_t THAtomicGet(int32_t volatile *a);
 
 /*
  * *a += value,
  * return previous *a
 */
-TH_API int THAtomicAdd(int volatile *a, int value);
+TH_API int32_t THAtomicAdd(int32_t volatile *a, int32_t value);
 
 /*
  * check if (*a == oldvalue)
  * if true: set *a to newvalue, return 1
  * if false: return 0
 */
-TH_API int THAtomicCompareAndSwap(int volatile *a, int oldvalue, int newvalue);
+TH_API int32_t THAtomicCompareAndSwap(int32_t volatile *a, int32_t oldvalue, int32_t newvalue);
 
 
 /******************************************************************************
@@ -49,13 +49,13 @@ TH_API int THAtomicCompareAndSwap(int volatile *a, int oldvalue, int newvalue);
 /*
  * *a++
 */
-TH_API void THAtomicIncrementRef(int volatile *a);
+TH_API void THAtomicIncrementRef(int32_t volatile *a);
 
 /*
  * *a--,
  * return 1 if *a == 0 after the operation, 0 otherwise
 */
-TH_API int THAtomicDecrementRef(int volatile *a);
+TH_API int32_t THAtomicDecrementRef(int32_t volatile *a);
 
 
 
@@ -66,25 +66,25 @@ TH_API int THAtomicDecrementRef(int volatile *a);
 /*
  * *a = newvalue
 */
-TH_API void THAtomicSetLong(long volatile *a, long newvalue);
+TH_API void THAtomicSetLong(int64_t volatile *a, int64_t newvalue);
 
 /*
  * return *a
 */
-TH_API long THAtomicGetLong(long volatile *a);
+TH_API int64_t THAtomicGetLong(int64_t volatile *a);
 
 /*
  * *a += value,
  * return previous *a
 */
-TH_API long THAtomicAddLong(long volatile *a, long value);
+TH_API int64_t THAtomicAddLong(int64_t volatile *a, int64_t 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);
+TH_API int64_t THAtomicCompareAndSwapLong(int64_t volatile *a, int64_t oldvalue, int64_t newvalue);