feat: using jemalloc?

barretenberg/cpp/CMakeLists.txt

@@ -33,6 +33,7 @@ option(OMP_MULTITHREADING "Enable OMP multi-threading" OFF)
 option(FUZZING "Build ONLY fuzzing harnesses" OFF)
 option(DISABLE_TBB "Intel Thread Building Blocks" ON)
 option(COVERAGE "Enable collecting coverage from tests" OFF)
+option(ENABLE_JEMALLOC "Use the sophisticated, thread-optimized fragmentation-friendly jemalloc" ON)
 option(ENABLE_ASAN "Address sanitizer for debugging tricky memory corruption" OFF)
 option(ENABLE_HEAVY_TESTS "Enable heavy tests when collecting coverage" OFF)
 # Note: Must do 'sudo apt-get install libdw-dev' or equivalent
@@ -158,6 +159,7 @@ include(GNUInstallDirs)
 include(cmake/arch.cmake)
 include(cmake/threading.cmake)
 include(cmake/tracy.cmake)
+include(cmake/jemalloc.cmake)
 include(cmake/gtest.cmake)
 include(cmake/benchmark.cmake)
 include(cmake/module.cmake)
@@ -175,6 +177,13 @@ if(DISABLE_AZTEC_VM)
 endif()
 include(cmake/backward-cpp.cmake)
 
+if(ENABLE_JEMALLOC)
+    add_compile_options(-DENABLE_JEMALLOC=1)
+    SET(JEMALLOC_LIBS jemalloc)
+else()
+    SET(JEMALLOC_LIBS)
+endif()
+
 if (WASM)
     set(DISABLE_AZTEC_VM ON)
 endif()

barretenberg/cpp/cmake/jemalloc.cmake

@@ -0,0 +1,23 @@
+# Based on https://github.com/apple/foundationdb/blob/0946f495797408b28d02b52b1c0a7c4b6a401736/cmake/Jemalloc.cmake
+
+add_library(jemalloc INTERFACE)
+add_library(im_jemalloc_pic STATIC IMPORTED)
+add_library(im_jemalloc STATIC IMPORTED)
+include(ExternalProject)
+set(JEMALLOC_DIR "${CMAKE_BINARY_DIR}/jemalloc")
+ExternalProject_add(Jemalloc_project
+  URL "https://github.com/jemalloc/jemalloc/releases/download/5.3.0/jemalloc-5.3.0.tar.bz2"
+  URL_HASH SHA256=2db82d1e7119df3e71b7640219b6dfe84789bc0537983c3b7ac4f7189aecfeaa
+  BUILD_BYPRODUCTS "${JEMALLOC_DIR}/include/jemalloc/jemalloc.h"
+  "${JEMALLOC_DIR}/lib/libjemalloc.a"
+  "${JEMALLOC_DIR}/lib/libjemalloc_pic.a"
+  CONFIGURE_COMMAND ./configure --prefix=${JEMALLOC_DIR} --enable-static
+  BUILD_IN_SOURCE ON
+  BUILD_COMMAND make
+  INSTALL_DIR "${JEMALLOC_DIR}"
+  INSTALL_COMMAND make install)
+add_dependencies(im_jemalloc Jemalloc_project)
+add_dependencies(im_jemalloc_pic Jemalloc_project)
+set_target_properties(im_jemalloc_pic PROPERTIES IMPORTED_LOCATION "${JEMALLOC_DIR}/lib/libjemalloc_pic.a")
+set_target_properties(im_jemalloc PROPERTIES IMPORTED_LOCATION "${JEMALLOC_DIR}/lib/libjemalloc.a")
+target_link_libraries(jemalloc INTERFACE im_jemalloc_pic im_jemalloc)

barretenberg/cpp/cmake/module.cmake

@@ -54,6 +54,7 @@ function(barretenberg_module MODULE_NAME)
             PUBLIC
             ${ARGN}
             ${TRACY_LIBS}
+            ${JEMALLOC_LIBS}
             ${TBB_IMPORTED_TARGETS}
         )
 
@@ -91,6 +92,7 @@ function(barretenberg_module MODULE_NAME)
             ${MODULE_NAME}_test_objects
             PRIVATE
             ${TRACY_LIBS}
+            ${JEMALLOC_LIBS}
             GTest::gtest
             GTest::gtest_main
             GTest::gmock_main
@@ -150,6 +152,7 @@ function(barretenberg_module MODULE_NAME)
             GTest::gtest_main
             GTest::gmock_main
             ${TRACY_LIBS}
+            ${JEMALLOC_LIBS}
             ${TBB_IMPORTED_TARGETS}
         )
 
@@ -242,6 +245,7 @@ function(barretenberg_module MODULE_NAME)
                 PRIVATE
                 benchmark::benchmark
                 ${TRACY_LIBS}
+                ${JEMALLOC_LIBS}
                 ${TBB_IMPORTED_TARGETS}
             )
 
@@ -258,6 +262,7 @@ function(barretenberg_module MODULE_NAME)
                 ${ARGN}
                 benchmark::benchmark
                 ${TRACY_LIBS}
+                ${JEMALLOC_LIBS}
                 ${TBB_IMPORTED_TARGETS}
             )
             if(CHECK_CIRCUIT_STACKTRACES)

barretenberg/cpp/src/barretenberg/bb/CMakeLists.txt

@@ -13,6 +13,7 @@ if (NOT(FUZZING) AND NOT(WASM))
         env
         circuit_checker
         ${TRACY_LIBS}
+        ${JEMALLOC_LIBS}
         libdeflate::libdeflate_static
     )
     if(CHECK_CIRCUIT_STACKTRACES)

barretenberg/cpp/src/barretenberg/common/jemalloc_helper.h

@@ -0,0 +1,84 @@
+
+// Aztec-specific: We include this file from
+// https://github.com/tabokie/rocksdb/blob/54dd1848405edeaf63a96432ec11471a133a7633/port/jemalloc_helper.h
+// to help with clang issues with linking jemalloc. File replicated below.
+
+//  Copyright (c) 2011-present, Facebook, Inc.  All rights reserved.
+//  This source code is licensed under both the GPLv2 (found in the
+//  COPYING file in the root directory) and Apache 2.0 License
+//  (found in the LICENSE.Apache file in the root directory).
+
+#pragma once
+
+#if defined(__clang__)
+// glibc's `posix_memalign()` declaration specifies `throw()` while clang's
+// declaration does not. There is a hack in clang to make its re-declaration
+// compatible with glibc's if they are declared consecutively. That hack breaks
+// if yet another `posix_memalign()` declaration comes between glibc's and
+// clang's declarations. Include "mm_malloc.h" here ensures glibc's and clang's
+// declarations both come before "jemalloc.h"'s `posix_memalign()` declaration.
+//
+// This problem could also be avoided if "jemalloc.h"'s `posix_memalign()`
+// declaration did not specify `throw()` when built with clang.
+#include <mm_malloc.h>
+#endif
+
+#ifdef ROCKSDB_JEMALLOC
+#ifdef __FreeBSD__
+#include <malloc_np.h>
+#else
+#define JEMALLOC_MANGLE
+#include <jemalloc/jemalloc.h>
+#endif
+
+#ifndef JEMALLOC_CXX_THROW
+#define JEMALLOC_CXX_THROW
+#endif
+
+#if defined(OS_WIN) && defined(_MSC_VER)
+
+// MSVC does not have weak symbol support. As long as ROCKSDB_JEMALLOC is defined,
+// Jemalloc memory allocator is used.
+static inline bool HasJemalloc()
+{
+    return true;
+}
+
+#else
+
+// Declare non-standard jemalloc APIs as weak symbols. We can null-check these
+// symbols to detect whether jemalloc is linked with the binary.
+extern "C" JEMALLOC_ALLOCATOR JEMALLOC_RESTRICT_RETURN void JEMALLOC_NOTHROW* mallocx(size_t, int) JEMALLOC_ATTR(malloc)
+    JEMALLOC_ALLOC_SIZE(1) __attribute__((__weak__));
+extern "C" JEMALLOC_ALLOCATOR JEMALLOC_RESTRICT_RETURN void JEMALLOC_NOTHROW* rallocx(void*, size_t, int)
+    JEMALLOC_ALLOC_SIZE(2) __attribute__((__weak__));
+extern "C" size_t JEMALLOC_NOTHROW xallocx(void*, size_t, size_t, int) __attribute__((__weak__));
+extern "C" size_t JEMALLOC_NOTHROW sallocx(const void*, int) JEMALLOC_ATTR(pure) __attribute__((__weak__));
+extern "C" void JEMALLOC_NOTHROW dallocx(void*, int) __attribute__((__weak__));
+extern "C" void JEMALLOC_NOTHROW sdallocx(void*, size_t, int) __attribute__((__weak__));
+extern "C" size_t JEMALLOC_NOTHROW nallocx(size_t, int) JEMALLOC_ATTR(pure) __attribute__((__weak__));
+extern "C" int JEMALLOC_NOTHROW mallctl(const char*, void*, size_t*, void*, size_t) __attribute__((__weak__));
+extern "C" int JEMALLOC_NOTHROW mallctlnametomib(const char*, size_t*, size_t*) __attribute__((__weak__));
+extern "C" int JEMALLOC_NOTHROW mallctlbymib(const size_t*, size_t, void*, size_t*, void*, size_t)
+    __attribute__((__weak__));
+extern "C" void JEMALLOC_NOTHROW malloc_stats_print(void (*)(void*, const char*), void*, const char*)
+    __attribute__((__weak__));
+extern "C" size_t JEMALLOC_NOTHROW malloc_usable_size(JEMALLOC_USABLE_SIZE_CONST void*) JEMALLOC_CXX_THROW
+    __attribute__((__weak__));
+
+// Check if Jemalloc is linked with the binary. Note the main program might be
+// using a different memory allocator even this method return true.
+// It is loosely based on folly::usingJEMalloc(), minus the check that actually
+// allocate memory and see if it is through jemalloc, to handle the dlopen()
+// case:
+// https://github.com/facebook/folly/blob/76cf8b5841fb33137cfbf8b224f0226437c855bc/folly/memory/Malloc.h#L147
+static inline bool HasJemalloc()
+{
+    return mallocx != nullptr && rallocx != nullptr && xallocx != nullptr && sallocx != nullptr && dallocx != nullptr &&
+           sdallocx != nullptr && nallocx != nullptr && mallctl != nullptr && mallctlnametomib != nullptr &&
+           mallctlbymib != nullptr && malloc_stats_print != nullptr && malloc_usable_size != nullptr;
+}
+
+#endif
+
+#endif // ROCKSDB_JEMALLOC

barretenberg/cpp/src/barretenberg/common/mem.hpp

@@ -6,6 +6,10 @@
 #include <cstdlib>
 #include <memory>
 
+#ifdef ENABLE_JEMALLOC
+#include "jemalloc_helper.h"
+#endif
+
 // This can be altered to capture stack traces, though more expensive
 // so wrap TracyAlloc or TracyAllocS. We disable these if gates are being tracked
 // Gates are hackishly tracked as if they were memory, for the sweet sweet memory