Merge pull request #8874 from vit-stepanovs/cmake-jemalloc-support

[CMake] Experimental Jemalloc support.

tensorflow/contrib/cmake/CMakeLists.txt

@@ -22,6 +22,7 @@ option(tensorflow_ENABLE_GPU "Enable GPU support" OFF)
 option(tensorflow_ENABLE_SSL_SUPPORT "Enable boringssl support" OFF)
 option(tensorflow_ENABLE_GRPC_SUPPORT "Enable gRPC support" ON)
 option(tensorflow_ENABLE_HDFS_SUPPORT "Enable HDFS support" OFF)
+option(tensorflow_ENABLE_JEMALLOC_SUPPORT "Enable jemalloc support" OFF)
 option(tensorflow_BUILD_CC_EXAMPLE "Build the C++ tutorial example" ON)
 option(tensorflow_BUILD_PYTHON_BINDINGS "Build the Python bindings" ON)
 option(tensorflow_BUILD_ALL_KERNELS "Build all OpKernels" ON)
@@ -94,6 +95,10 @@ if (tensorflow_WIN_CPU_SIMD_OPTIONS)
   endif()
 endif()
 
+if (tensorflow_ENABLE_JEMALLOC_SUPPORT)
+  add_definitions(-DTENSORFLOW_USE_JEMALLOC -DJEMALLOC_EXPORT=)
+endif()
+
 # External dependencies
 include(zlib)
 include(gif)
@@ -161,6 +166,12 @@ if(tensorflow_ENABLE_GRPC_SUPPORT)
   list(APPEND tensorflow_EXTERNAL_DEPENDENCIES grpc)
   include_directories(${GRPC_INCLUDE_DIRS})
 endif()
+if(tensorflow_ENABLE_JEMALLOC_SUPPORT)
+  include(jemalloc)
+  list(APPEND tensorflow_EXTERNAL_LIBRARIES ${jemalloc_STATIC_LIBRARIES})
+  list(APPEND tensorflow_EXTERNAL_DEPENDENCIES jemalloc)
+  include_directories(${jemalloc_INCLUDE_DIRS})
+endif()
 if(WIN32)
   list(APPEND tensorflow_EXTERNAL_LIBRARIES wsock32 ws2_32 shlwapi)
 endif()

tensorflow/contrib/cmake/external/jemalloc.cmake

@@ -0,0 +1,33 @@
+include (ExternalProject)
+
+set(jemalloc_INCLUDE_DIRS ${CMAKE_CURRENT_BINARY_DIR}/jemalloc/src/jemalloc/include)
+set(jemalloc_URL https://github.com/jemalloc/jemalloc-cmake/archive/jemalloc-cmake.4.3.1.tar.gz)
+set(jemalloc_HASH SHA256=f9be9a05fe906deb5c1c8ca818071a7d2e27d66fd87f5ba9a7bf3750bcedeaf0)
+set(jemalloc_BUILD ${CMAKE_CURRENT_BINARY_DIR}/jemalloc/src/jemalloc)
+
+if (WIN32)
+    set(jemalloc_INCLUDE_DIRS
+        ${jemalloc_INCLUDE_DIRS} 
+        ${CMAKE_CURRENT_BINARY_DIR}/jemalloc/src/jemalloc/include/msvc_compat
+    )
+    set(jemalloc_ADDITIONAL_CMAKE_OPTIONS -A x64)
+    set(jemalloc_STATIC_LIBRARIES ${jemalloc_BUILD}/Release/jemalloc.lib)
+else()
+    set(jemalloc_STATIC_LIBRARIES ${jemalloc_BUILD}/Release/jemalloc.a)
+endif()
+
+ExternalProject_Add(jemalloc
+    PREFIX jemalloc
+    URL ${jemalloc_URL}
+    URL_HASH ${jemalloc_HASH}
+    DOWNLOAD_DIR "${DOWNLOAD_LOCATION}"
+    BUILD_IN_SOURCE 1
+    CONFIGURE_COMMAND ${CMAKE_COMMAND}
+        -DCMAKE_BUILD_TYPE:STRING=Release
+        -DCMAKE_VERBOSE_MAKEFILE:BOOL=OFF
+        -Dwith-jemalloc-prefix:STRING=jemalloc_
+        -Dwithout-export:BOOL=ON
+        ${jemalloc_ADDITIONAL_CMAKE_OPTIONS}
+    BUILD_COMMAND ${CMAKE_COMMAND} --build . --config Release --target jemalloc
+    INSTALL_COMMAND ${CMAKE_COMMAND} -E echo "Skipping install step."
+)

tensorflow/core/platform/windows/port.cc

@@ -13,6 +13,10 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
+#ifdef TENSORFLOW_USE_JEMALLOC
+#include "jemalloc/jemalloc.h"
+#endif
+
 #include <stdio.h>
 #include <stdlib.h>
 #include <string.h>
@@ -53,16 +57,56 @@ int NumSchedulableCPUs() {
 }
 
 void* AlignedMalloc(size_t size, int minimum_alignment) {
-  return _aligned_malloc(size, minimum_alignment);
+#ifdef TENSORFLOW_USE_JEMALLOC
+    void* ptr = NULL;
+    // posix_memalign requires that the requested alignment be at least
+    // sizeof(void*). In this case, fall back on malloc which should return
+    // memory aligned to at least the size of a pointer.
+    const int required_alignment = sizeof(void*);
+    if (minimum_alignment < required_alignment) return Malloc(size);
+    int err = jemalloc_posix_memalign(&ptr, minimum_alignment, size);
+    if (err != 0) {
+        return NULL;
+    }
+    else {
+        return ptr;
+    }
+#else
+    return _aligned_malloc(size, minimum_alignment);
+#endif
 }
 
-void AlignedFree(void* aligned_memory) { _aligned_free(aligned_memory); }
+void AlignedFree(void* aligned_memory) {
+#ifdef TENSORFLOW_USE_JEMALLOC
+  jemalloc_free(aligned_memory);
+#else
+  _aligned_free(aligned_memory);
+#endif
+}
 
-void* Malloc(size_t size) { return ::malloc(size); }
+void* Malloc(size_t size) {
+#ifdef TENSORFLOW_USE_JEMALLOC
+  return jemalloc_malloc(size);
+#else
+  return malloc(size);
+#endif
+}
 
-void* Realloc(void* ptr, size_t size) { return ::realloc(ptr, size); }
+void* Realloc(void* ptr, size_t size) {
+#ifdef TENSORFLOW_USE_JEMALLOC
+  return jemalloc_realloc(ptr, size);
+#else
+  return realloc(ptr, size);
+#endif
+}
 
-void Free(void* ptr) { ::free(ptr); }
+void Free(void* ptr) {
+#ifdef TENSORFLOW_USE_JEMALLOC
+  return jemalloc_free(ptr);
+#else
+  return free(ptr);
+#endif
+}
 
 void MallocExtension_ReleaseToSystem(std::size_t num_bytes) {
   // No-op.