Merge pull request torch#758 from apaszke/master

Various improvements to THMapAllocator
gentlegene · Sep 8, 2016 · e5ebac6 · e5ebac6
2 parents 01d741a + 67d602d
commit e5ebac6
Show file tree

Hide file tree

Showing 3 changed files with 55 additions and 23 deletions.
diff --git a/lib/TH/CMakeLists.txt b/lib/TH/CMakeLists.txt
@@ -70,28 +70,6 @@ IF (CORTEXA9_FOUND)
   SET(CMAKE_C_FLAGS "-mcpu=cortex-a9 ${CMAKE_C_FLAGS}")
 ENDIF (CORTEXA9_FOUND)
 
-IF(UNIX)
-  INCLUDE(CheckFunctionExists)
-  SET(CMAKE_EXTRA_INCLUDE_FILES "sys/mman.h")
-  CHECK_FUNCTION_EXISTS(mmap HAVE_MMAP)
-  IF(HAVE_MMAP)
-    ADD_DEFINITIONS(-DHAVE_MMAP=1)
-  ENDIF(HAVE_MMAP)
-  ADD_DEFINITIONS(-D_FILE_OFFSET_BITS=64)
-  CHECK_FUNCTION_EXISTS(shm_open HAVE_SHM_OPEN)
-  IF(HAVE_SHM_OPEN)
-    ADD_DEFINITIONS(-DHAVE_SHM_OPEN=1)
-  ENDIF(HAVE_SHM_OPEN)
-  CHECK_FUNCTION_EXISTS(shm_unlink HAVE_SHM_UNLINK)
-  IF(HAVE_SHM_UNLINK)
-    ADD_DEFINITIONS(-DHAVE_SHM_UNLINK=1)
-  ENDIF(HAVE_SHM_UNLINK)
-  CHECK_FUNCTION_EXISTS(malloc_usable_size HAVE_MALLOC_USABLE_SIZE)
-  IF(HAVE_MALLOC_USABLE_SIZE)
-    ADD_DEFINITIONS(-DHAVE_MALLOC_USABLE_SIZE=1)
-  ENDIF(HAVE_MALLOC_USABLE_SIZE)
-ENDIF(UNIX)
-
 FIND_PACKAGE(SSE)
 IF(C_SSE2_FOUND)
   SET(CMAKE_C_FLAGS "${C_SSE2_FLAGS} -DUSE_SSE2 ${CMAKE_C_FLAGS}")
@@ -224,9 +202,34 @@ IF (UNIX AND NOT APPLE)
    CHECK_LIBRARY_EXISTS(rt clock_gettime "time.h" NEED_LIBRT)
    IF(NEED_LIBRT)
      TARGET_LINK_LIBRARIES(TH rt)
+     SET(CMAKE_REQUIRED_LIBRARIES ${CMAKE_REQUIRED_LIBRARIES} rt)
    ENDIF(NEED_LIBRT)
 ENDIF(UNIX AND NOT APPLE)
 
+IF(UNIX)
+  INCLUDE(CheckFunctionExists)
+  SET(CMAKE_EXTRA_INCLUDE_FILES "sys/mman.h")
+  CHECK_FUNCTION_EXISTS(mmap HAVE_MMAP)
+  IF(HAVE_MMAP)
+    ADD_DEFINITIONS(-DHAVE_MMAP=1)
+  ENDIF(HAVE_MMAP)
+  ADD_DEFINITIONS(-D_FILE_OFFSET_BITS=64)
+  CHECK_FUNCTION_EXISTS(shm_open HAVE_SHM_OPEN)
+  IF(HAVE_SHM_OPEN)
+    ADD_DEFINITIONS(-DHAVE_SHM_OPEN=1)
+  ENDIF(HAVE_SHM_OPEN)
+  CHECK_FUNCTION_EXISTS(shm_unlink HAVE_SHM_UNLINK)
+  IF(HAVE_SHM_UNLINK)
+    ADD_DEFINITIONS(-DHAVE_SHM_UNLINK=1)
+  ENDIF(HAVE_SHM_UNLINK)
+  CHECK_FUNCTION_EXISTS(malloc_usable_size HAVE_MALLOC_USABLE_SIZE)
+  IF(HAVE_MALLOC_USABLE_SIZE)
+    ADD_DEFINITIONS(-DHAVE_MALLOC_USABLE_SIZE=1)
+  ENDIF(HAVE_MALLOC_USABLE_SIZE)
+ENDIF(UNIX)
+
+
+
 IF(NOT MSVC)
   TARGET_LINK_LIBRARIES(TH m)
 ENDIF(NOT MSVC)

diff --git a/lib/TH/THAllocator.c b/lib/TH/THAllocator.c
@@ -242,7 +242,7 @@ static void *_map_alloc(void* ctx_, long size)
         if((fd = shm_open(ctx->filename, flags, (mode_t)0600)) == -1)
           THError("unable to open shared memory object <%s> in read-write mode", ctx->filename);
 #else
-        THError("unable to open file <%s> in sharedmem mode, shm_open unavailable on this platform");
+        THError("unable to open file <%s> in sharedmem mode, shm_open unavailable on this platform", ctx->filename);
 #endif
       }
       else
@@ -278,11 +278,16 @@ static void *_map_alloc(void* ctx_, long size)
             close(fd);
             THError("unable to stretch file <%s> to the right size", ctx->filename);
           }
+/* on OS X write returns with errno 45 (Opperation not supported) when used
+ * with a file descriptor obtained via shm_open
+ */
+#ifndef __APPLE__
           if((write(fd, "", 1)) != 1) /* note that the string "" contains the '\0' byte ... */
           {
             close(fd);
             THError("unable to write to file <%s>", ctx->filename);
           }
+#endif
         }
         else
         {
@@ -462,6 +467,18 @@ static void THRefcountedMapAllocator_free(void* ctx_, void* data) {
   THMapAllocatorContext_free(ctx);
 }
 
+void THRefcountedMapAllocator_incref(THMapAllocatorContext *ctx, void *data)
+{
+  THMapInfo *map_info = (THMapInfo*)(((char*)data) - TH_ALLOC_ALIGNMENT);
+  THAtomicIncrementRef(&map_info->refcount);
+}
+
+int THRefcountedMapAllocator_decref(THMapAllocatorContext *ctx, void *data)
+{
+  THMapInfo *map_info = (THMapInfo*)(((char*)data) - TH_ALLOC_ALIGNMENT);
+  return THAtomicDecrementRef(&map_info->refcount);
+}
+
 #else
 
 static void * THRefcountedMapAllocator_alloc(void *ctx, long size) {
@@ -478,6 +495,16 @@ static void THRefcountedMapAllocator_free(void* ctx_, void* data) {
   THError("refcounted file mapping not supported on your system");
 }
 
+void THRefcountedMapAllocator_incref(THMapAllocatorContext *ctx, void *data)
+{
+  THError("refcounted file mapping not supported on your system");
+}
+
+int THRefcountedMapAllocator_decref(THMapAllocatorContext *ctx, void *data)
+{
+  THError("refcounted file mapping not supported on your system");
+}
+
 #endif
 
 THAllocator THMapAllocator = {

diff --git a/lib/TH/THAllocator.h b/lib/TH/THAllocator.h
@@ -34,6 +34,8 @@ TH_API char * THMapAllocatorContext_filename(THMapAllocatorContext *ctx);
 TH_API int THMapAllocatorContext_fd(THMapAllocatorContext *ctx);
 TH_API long THMapAllocatorContext_size(THMapAllocatorContext *ctx);
 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;