Support jemalloc

couchbasedeps · Feb 5, 2025 · f87b8f2 · f87b8f2
1 parent 5ef8b19
commit f87b8f2
Show file tree

Hide file tree

Showing 29 changed files with 225 additions and 20 deletions.
diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -46,6 +46,8 @@ option(NO_WARMUP "Do not run a benchmark on each startup just to find the best l
 
 option(FIXED_LIBNAME "Use a non-versioned name for the library and no symbolic linking to variant names" OFF)
 
+option(USE_JEMALLOC "Use jemalloc globally across openBLAS" OFF)
+
 set(LIBNAMEPREFIX "" CACHE STRING  "Add a prefix to the openblas part of the library name" )
 set(LIBNAMESUFFIX "" CACHE STRING  "Add a suffix after the openblas part of the library name" )
 
@@ -105,6 +107,15 @@ message(WARNING "CMake support is experimental. It does not yet support all buil
 include("${PROJECT_SOURCE_DIR}/cmake/utils.cmake")
 include("${PROJECT_SOURCE_DIR}/cmake/system.cmake")
 
+if (USE_JEMALLOC)
+  # Allow the user to specify the path to the jemalloc library by specifying
+  # Jemalloc_ROOT as a CMake variable or environment variable.
+  cmake_policy(SET CMP0074 NEW)
+  add_definitions(-DUSE_JEMALLOC)
+  set(CMAKE_MODULE_PATH "${CMAKE_CURRENT_SOURCE_DIR}/cmake" ${CMAKE_MODULE_PATH})
+  find_package(Jemalloc REQUIRED)
+endif()
+
 set(OpenBLAS_LIBNAME ${LIBNAMEPREFIX}openblas${LIBNAMESUFFIX}${SUFFIX64_UNDERSCORE})
 
 set(BLASDIRS interface driver/level2 driver/level3 driver/others)
@@ -219,10 +230,16 @@ endif ()
 # add objects to the openblas lib
 if(NOT NO_LAPACK)
   add_library(LAPACK_OVERRIDES OBJECT ${LA_SOURCES})
+  if (USE_JEMALLOC)
+    target_link_libraries(LAPACK_OVERRIDES Jemalloc::jemalloc)
+  endif()
   list(APPEND TARGET_OBJS "$<TARGET_OBJECTS:LAPACK_OVERRIDES>")
 endif()
 if(NOT NO_LAPACKE)
   add_library(LAPACKE OBJECT ${LAPACKE_SOURCES})
+  if (USE_JEMALLOC)
+    target_link_libraries(LAPACKE Jemalloc::jemalloc)
+  endif()
   list(APPEND TARGET_OBJS "$<TARGET_OBJECTS:LAPACKE>")
 endif()
 #if(BUILD_RELAPACK)
@@ -233,11 +250,17 @@ set(OpenBLAS_LIBS "")
 if(BUILD_STATIC_LIBS)
   add_library(${OpenBLAS_LIBNAME}_static STATIC ${TARGET_OBJS} ${OpenBLAS_DEF_FILE})
   target_include_directories(${OpenBLAS_LIBNAME}_static INTERFACE $<INSTALL_INTERFACE:include/openblas${SUFFIX64}>)
+  if (USE_JEMALLOC)
+    target_link_libraries(${OpenBLAS_LIBNAME}_static Jemalloc::jemalloc)
+  endif()
   list(APPEND OpenBLAS_LIBS ${OpenBLAS_LIBNAME}_static)
 endif()
 if(BUILD_SHARED_LIBS)
   add_library(${OpenBLAS_LIBNAME}_shared SHARED ${TARGET_OBJS} ${OpenBLAS_DEF_FILE})
   target_include_directories(${OpenBLAS_LIBNAME}_shared INTERFACE $<INSTALL_INTERFACE:include/openblas${SUFFIX64}>)
+  if (USE_JEMALLOC)
+    target_link_libraries(${OpenBLAS_LIBNAME}_shared Jemalloc::jemalloc)
+  endif()
   list(APPEND OpenBLAS_LIBS ${OpenBLAS_LIBNAME}_shared)
 endif()
 if(BUILD_STATIC_LIBS)
@@ -510,6 +533,9 @@ if (BUILD_BENCHMARKS)
             (NOT ${target_name} STREQUAL "benchmark_max_COMPLEX") AND (NOT ${target_name} STREQUAL "benchmark_max_COMPLEX_DOUBLE") AND
             (NOT ${target_name} STREQUAL "benchmark_min_COMPLEX") AND (NOT ${target_name} STREQUAL "benchmark_min_COMPLEX_DOUBLE"))
           add_executable(${target_name} ${source})
+          if (USE_JEMALLOC)
+            target_link_libraries(${target_name} Jemalloc::jemalloc)
+          endif()
           target_include_directories(${target_name} PRIVATE ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_BINARY_DIR})
 	    target_link_libraries(${target_name} ${OpenBLAS_LIBNAME} )
 	  #  target_link_libraries(${target_name} ${OpenBLAS_LIBNAME} OpenMP::OpenMP_C)

diff --git a/cmake/FindJemalloc.cmake b/cmake/FindJemalloc.cmake
@@ -0,0 +1,57 @@
+# - Try to find jemalloc
+# Once done this will define a target Jemalloc::jemalloc which will include all
+# required definitions and libraries.
+
+if (NOT FindJemalloc_included)
+
+  # First try finding Jemalloc using pkg-config.
+  find_package(PkgConfig QUIET)
+  if (PKG_CONFIG_FOUND)
+    pkg_check_modules(PC_JEMALLOC QUIET IMPORTED_TARGET GLOBAL jemalloc)
+
+    if (PC_JEMALLOC_FOUND)
+      # Ensure it created the PkgConfig target
+      if(NOT TARGET PkgConfig::PC_JEMALLOC)
+        message(FATAL_ERROR
+          "Found Jemalloc via pkg-config, but it did not create the PkgConfig::PC_JEMALLOC target.")
+      endif()
+
+      # Make the discovered target available as Jemalloc::jemalloc
+      add_library(Jemalloc::jemalloc ALIAS PkgConfig::PC_JEMALLOC)
+      set(_jemalloc_found TRUE)
+    endif()
+  endif()
+
+  # If that didn't find it, try finding Jemalloc using CMake's config
+  # mode.
+  if(NOT _jemalloc_found)
+    find_package(Jemalloc CONFIG)
+
+    if(Jemalloc_FOUND)
+      if(NOT TARGET Jemalloc::jemalloc)
+        # Ensure this found package created the standard target.
+        message(FATAL_ERROR
+          "Found Jemalloc, but it did not create the Jemalloc::jemalloc target.")
+      endif()
+      set(_jemalloc_found TRUE)
+    endif()
+  endif()
+
+  # Determine the version of the found jemalloc.
+  get_target_property(_jemalloc_include_dir Jemalloc::jemalloc INTERFACE_INCLUDE_DIRECTORIES)
+  set(_version_regex "^#define[ \t]+JEMALLOC_VERSION[ \t]+\"([^\"]+)\".*")
+  file(STRINGS "${_jemalloc_include_dir}/jemalloc/jemalloc.h"
+    JEMALLOC_VERSION REGEX "${_version_regex}")
+  string(REGEX REPLACE "${_version_regex}" "\\1"
+    JEMALLOC_VERSION "${JEMALLOC_VERSION}")
+  unset(_version_regex)
+
+  # handle the QUIET and REQUIRED arguments, verify version (if
+  # necessary), report found status, and set JEMALLOC_FOUND to TRUE.
+  include(FindPackageHandleStandardArgs)
+  find_package_handle_standard_args(
+    Jemalloc VERSION_VAR JEMALLOC_VERSION
+    REQUIRED_VARS _jemalloc_found)
+  set(FindJemalloc_included TRUE)
+
+endif (NOT FindJemalloc_included)
diff --git a/common.h b/common.h
@@ -82,6 +82,10 @@ extern "C" {
 #include <stdio.h>
 #include <stdlib.h>
 #include <string.h>
+#ifdef USE_JEMALLOC
+#define JEMALLOC_MANGLE
+#include <jemalloc/jemalloc.h>
+#endif
 
 #if !defined(_MSC_VER)
 #include <unistd.h>

diff --git a/cpp_thread_test/CMakeLists.txt b/cpp_thread_test/CMakeLists.txt
@@ -8,14 +8,20 @@ set(CMAKE_CXX_FLAGS "${CMAKE_C_FLAGS} -DADD${BU} -DCBLAS")
 if (USE_OPENMP)
 if (CPP_THREAD_SAFETY_TEST)
   message(STATUS building thread safety test)
-  add_executable(dgemm_thread_safety dgemm_thread_safety.cpp) 
+  add_executable(dgemm_thread_safety dgemm_thread_safety.cpp)
+  if (USE_JEMALLOC)
+    target_link_libraries(dgemm_thread_safety Jemalloc::jemalloc)
+  endif()
   target_link_libraries(dgemm_thread_safety ${OpenBLAS_LIBNAME})
   add_test( dgemm_thread_safety ${CMAKE_CURRENT_BINARY_DIR}/dgemm_thread_safety)
 endif()
 
 
 if (CPP_THREAD_SAFETY_TEST OR CPP_THREAD_SAFETY_GEMV)
-  add_executable(dgemv_thread_safety dgemv_thread_safety.cpp) 
+  add_executable(dgemv_thread_safety dgemv_thread_safety.cpp)
+  if (USE_JEMALLOC)
+    target_link_libraries(dgemv_thread_safety Jemalloc::jemalloc)
+  endif()
   target_link_libraries(dgemv_thread_safety ${OpenBLAS_LIBNAME})
   add_test(dgemv_thread_safety ${CMAKE_CURRENT_BINARY_DIR}/dgemv_thread_safety)
 endif()

diff --git a/ctest/CMakeLists.txt b/ctest/CMakeLists.txt
@@ -39,6 +39,9 @@ else()
     c_${float_char}blat1c.c
     c_${float_char}blas1.c)
 endif()
+  if (USE_JEMALLOC)
+    target_link_libraries(x${float_char}cblat1 Jemalloc::jemalloc)
+  endif()
   target_link_libraries(x${float_char}cblat1 ${OpenBLAS_LIBNAME})
   if (USE_OPENMP AND (${CMAKE_Fortran_COMPILER_ID} STREQUAL GNU) AND (${CMAKE_C_COMPILER_ID} STREQUAL Clang))
         string(REGEX REPLACE "-fopenmp" "" CMAKE_Fortran_FLAGS "${CMAKE_Fortran_FLAGS}")
@@ -68,6 +71,9 @@ else()
     c_xerbla.c
     constant.c)
 endif()
+  if (USE_JEMALLOC)
+    target_link_libraries(x${float_char}cblat2 Jemalloc::jemalloc)
+  endif()
   target_link_libraries(x${float_char}cblat2 ${OpenBLAS_LIBNAME})
   if (USE_OPENMP AND (${CMAKE_Fortran_COMPILER_ID} STREQUAL GNU) AND (${CMAKE_C_COMPILER_ID} STREQUAL Clang))
         string(REGEX REPLACE "-fopenmp" "" CMAKE_Fortran_FLAGS "${CMAKE_Fortran_FLAGS}")
@@ -88,6 +94,9 @@ if (NOT NOFORTRAN)
     auxiliary.c
     c_xerbla.c
     constant.c)
+  if (USE_JEMALLOC)
+    target_link_libraries(x${float_char}cblat3 Jemalloc::jemalloc)
+  endif()
   if (USE_GEMM3M)
   if ((${float_char} STREQUAL "c") OR (${float_char} STREQUAL "z"))
   add_executable(x${float_char}cblat3_3m
@@ -97,6 +106,9 @@ if (NOT NOFORTRAN)
     auxiliary.c
     c_xerbla.c
     constant.c)
+  if (USE_JEMALLOC)
+    target_link_libraries(x${float_char}cblat3_3m Jemalloc::jemalloc)
+  endif()
   endif()
   endif()
 else()
@@ -107,6 +119,9 @@ else()
     auxiliary.c
     c_xerbla.c
     constant.c)
+    if (USE_JEMALLOC)
+      target_link_libraries(x${float_char}cblat3 Jemalloc::jemalloc)
+    endif()
   if (USE_GEMM3M)
   if ((${float_char} STREQUAL "c") OR (${float_char} STREQUAL "z"))
   add_executable(x${float_char}cblat3_3m
@@ -116,6 +131,9 @@ else()
     auxiliary.c
     c_xerbla.c
     constant.c)
+  if (USE_JEMALLOC)
+    target_link_libraries(x${float_char}cblat3_3m Jemalloc::jemalloc)
+  endif()
   endif()
   endif()
 endif()

diff --git a/driver/level2/CMakeLists.txt b/driver/level2/CMakeLists.txt
@@ -62,7 +62,7 @@ GenerateCombinationObjects("${UL_SOURCES}" "LOWER" "U" "" 1 "" "" 3)
 # gbmv uses a lowercase n and t
 GenerateNamedObjects("gbmv_k.c" "" "gbmv_n" false "" "" "" 3)
 GenerateNamedObjects("gbmv_k.c" "TRANS" "gbmv_t" false "" "" "" 3)
-# c/zgbmv 
+# c/zgbmv
 GenerateNamedObjects("zgbmv_k.c" "CONJ" "gbmv_r" false "" "" "" 2)
 GenerateNamedObjects("zgbmv_k.c" "TRANS;CONJ" "gbmv_c" false "" "" "" 2)
 GenerateNamedObjects("zgbmv_k.c" "XCONJ" "gbmv_o" false "" "" "" 2)
@@ -223,3 +223,6 @@ if (USE_THREAD)
 endif ()
 
 add_library(driver_level2 OBJECT ${OPENBLAS_SRC})
+if (USE_JEMALLOC)
+  target_link_libraries(driver_level2 Jemalloc::jemalloc)
+endif()
diff --git a/driver/level3/CMakeLists.txt b/driver/level3/CMakeLists.txt
@@ -171,3 +171,6 @@ endforeach ()
 #
 
 add_library(driver_level3 OBJECT ${OPENBLAS_SRC})
+if (USE_JEMALLOC)
+  target_link_libraries(driver_level3 Jemalloc::jemalloc)
+endif()
diff --git a/driver/others/CMakeLists.txt b/driver/others/CMakeLists.txt
@@ -54,9 +54,9 @@ if (DYNAMIC_ARCH)
     list(APPEND COMMON_SOURCES dynamic_power.c)
   elseif (RISCV64)
     list(APPEND COMMON_SOURCES dynamic_riscv64.c detect_riscv64.c)
-  else ()  
+  else ()
     list(APPEND COMMON_SOURCES dynamic.c)
-  endif ()  
+  endif ()
 else ()
   list(APPEND COMMON_SOURCES parameter.c)
 endif ()
@@ -86,3 +86,6 @@ endif ()
 #endif
 
 add_library(driver_others OBJECT ${OPENBLAS_SRC} ${MEMORY} ${SMP_SOURCES} ${COMMON_SOURCES})
+if (USE_JEMALLOC)
+  target_link_libraries(driver_others Jemalloc::jemalloc)
+endif()
diff --git a/interface/CMakeLists.txt b/interface/CMakeLists.txt
@@ -28,9 +28,9 @@ set(BLAS1_MANGLED_SOURCES
 # these all have 'z' sources for complex versions
 set(BLAS2_SOURCES
   gemv.c ger.c
-  trsv.c trmv.c 
+  trsv.c trmv.c
   syr2.c gbmv.c
-  sbmv.c 
+  sbmv.c
   spr2.c
   tbsv.c tbmv.c
   tpsv.c tpmv.c
@@ -246,3 +246,6 @@ if ( BUILD_COMPLEX16 AND NOT  BUILD_DOUBLE)
 endif ()
 
 add_library(interface OBJECT ${OPENBLAS_SRC})
+if (USE_JEMALLOC)
+  target_link_libraries(interface Jemalloc::jemalloc)
+endif()
diff --git a/kernel/CMakeLists.txt b/kernel/CMakeLists.txt
@@ -1349,6 +1349,10 @@ endif ()
     set_target_properties(kernel${TSUFFIX} PROPERTIES COMPILE_FLAGS "${KERNEL_DEFINITIONS}")
     get_target_property(KERNEL_INCLUDE_DIRECTORIES kernel${TSUFFIX} INCLUDE_DIRECTORIES)
     set_target_properties(kernel${TSUFFIX} PROPERTIES INCLUDE_DIRECTORIES "${KERNEL_INCLUDE_DIRECTORIES};${TARGET_CONF_DIR}")
+    if (USE_JEMALLOC)
+      target_link_libraries(kernel${TSUFFIX} Jemalloc::jemalloc)
+      set_target_properties(kernel${TSUFFIX} PROPERTIES INCLUDE_DIRECTORIES "${JEMALLOC_INCLUDE_DIRS};${KERNEL_INCLUDE_DIRECTORIES};${TARGET_CONF_DIR}")
+    endif()
     if (USE_GEMM3M)
       target_compile_definitions(kernel${TSUFFIX} PRIVATE USE_GEMM3M)
     endif()

diff --git a/lapack-netlib/BLAS/SRC/CMakeLists.txt b/lapack-netlib/BLAS/SRC/CMakeLists.txt
@@ -103,4 +103,9 @@ set_target_properties(
   VERSION ${LAPACK_VERSION}
   SOVERSION ${LAPACK_MAJOR_VERSION}
   )
+
+if (USE_JEMALLOC)
+  target_link_libraries(${BLASLIB} Jemalloc::jemalloc)
+endif()
+
 lapack_install_library(${BLASLIB})
diff --git a/lapack-netlib/CBLAS/src/CMakeLists.txt b/lapack-netlib/CBLAS/src/CMakeLists.txt
@@ -126,3 +126,7 @@ target_include_directories(${CBLASLIB} PUBLIC
 )
 target_link_libraries(${CBLASLIB} PRIVATE ${BLAS_LIBRARIES})
 lapack_install_library(${CBLASLIB})
+
+if (USE_JEMALLOC)
+  target_link_libraries(${CBLASLIB} Jemalloc::jemalloc)
+endif()
diff --git a/lapack-netlib/LAPACKE/CMakeLists.txt b/lapack-netlib/LAPACKE/CMakeLists.txt
@@ -83,6 +83,11 @@ target_include_directories(${LAPACKELIB} PUBLIC
   $<BUILD_INTERFACE:${CMAKE_CURRENT_SOURCE_DIR}/include>
   $<INSTALL_INTERFACE:include>
 )
+
+if (USE_JEMALLOC)
+  target_link_libraries(${LAPACKELIB} Jemalloc::jemalloc)
+endif()
+
 if(WIN32 AND NOT UNIX)
   target_compile_definitions(${LAPACKELIB} PUBLIC HAVE_LAPACK_CONFIG_H LAPACK_COMPLEX_STRUCTURE)
   message(STATUS "Windows BUILD")

diff --git a/lapack-netlib/LAPACKE/include/lapack.h b/lapack-netlib/LAPACKE/include/lapack.h
@@ -11,6 +11,10 @@
 #include "lapacke_mangling.h"
 
 #include <stdlib.h>
+#ifdef USE_JEMALLOC
+#define JEMALLOC_MANGLE
+#include <jemalloc/jemalloc.h>
+#endif
 #include <stdarg.h>
 #include <inttypes.h>
 

diff --git a/lapack-netlib/SRC/CMakeLists.txt b/lapack-netlib/SRC/CMakeLists.txt
@@ -307,7 +307,7 @@ set(DLASRC
    dlaqtr.f dlar1v.f dlar2v.f iladlr.f iladlc.f
    dlarf.f  dlarfb.f dlarfb_gett.f dlarfg.f dlarfgp.f dlarft.f dlarfx.f dlarfy.f
    dlargv.f dlarrv.f dlartv.f
-   dlarz.f  dlarzb.f dlarzt.f dlaswp.f dlasy2.f 
+   dlarz.f  dlarzb.f dlarzt.f dlaswp.f dlasy2.f
    dlasyf.f dlasyf_rook.f dlasyf_rk.f dlasyf_aa.f
    dlatbs.f dlatdf.f dlatps.f dlatrd.f dlatrs.f dlatrz.f dlauu2.f
    dlauum.f dopgtr.f dopmtr.f dorg2l.f dorg2r.f
@@ -521,6 +521,10 @@ set_target_properties(
   SOVERSION ${LAPACK_MAJOR_VERSION}
   )
 
+if (USE_JEMALLOC)
+  target_link_libraries(${LAPACKLIB} Jemalloc::jemalloc)
+endif()
+
 if(USE_XBLAS)
   target_link_libraries(${LAPACKLIB} PRIVATE ${XBLAS_LIBRARY})
 endif()

diff --git a/lapack-netlib/SRC/ctrsyl3.c b/lapack-netlib/SRC/ctrsyl3.c
@@ -1,5 +1,9 @@
 #include <math.h>
 #include <stdlib.h>
+#ifdef USE_JEMALLOC
+#define JEMALLOC_MANGLE
+#include <jemalloc/jemalloc.h>
+#endif
 #include <string.h>
 #include <stdio.h>
 #include <complex.h>

diff --git a/lapack-netlib/SRC/dtrsyl3.c b/lapack-netlib/SRC/dtrsyl3.c
@@ -1,5 +1,7 @@
 #include <math.h>
 #include <stdlib.h>
+#define JEMALLOC_MANGLE
+#include <jemalloc/jemalloc.h> 
 #include <string.h>
 #include <stdio.h>
 #include <complex.h>

diff --git a/lapack-netlib/SRC/iparam2stage.c b/lapack-netlib/SRC/iparam2stage.c
@@ -1,5 +1,9 @@
 #include <math.h>
 #include <stdlib.h>
+#ifdef USE_JEMALLOC
+#define JEMALLOC_MANGLE
+#include <jemalloc/jemalloc.h>
+#endif
 #include <string.h>
 #include <stdio.h>
 #include <ctype.h>

diff --git a/lapack-netlib/SRC/ztrsyl3.c b/lapack-netlib/SRC/ztrsyl3.c
@@ -2,6 +2,10 @@
 #include <stdlib.h>
 #include <string.h>
 #include <stdio.h>
+#ifdef USE_JEMALLOC
+#define JEMALLOC_MANGLE
+#include <jemalloc/jemalloc.h>
+#endif
 #include <complex.h>
 #ifdef complex
 #undef complex