Merge pull request #2219 from yosefe/topic/ucm-dynamic-mmap-thresh

UCM/MALLOC: Add dynamic mmap threshold

src/ucm/malloc/malloc_hook.c

@@ -55,6 +55,8 @@
 #define UCM_OPERATOR_VEC_NEW_SYMBOL    "_Znam"
 #define UCM_OPERATOR_VEC_DELETE_SYMBOL "_ZdaPv"
 
+/* Maximal size for mmap threshold - 32mb */
+#define UCM_DEFAULT_MMAP_THRESHOLD_MAX (4ul * 1024 * 1024 * sizeof(long))
 
 /* Take out 12 LSB's, since they are the page-offset on most systems */
 #define ucm_mmap_addr_hash(_addr) \
@@ -77,7 +79,10 @@ typedef struct ucm_malloc_hook_state {
     pthread_mutex_t       install_mutex; /* Protect hooks installation */
     int                   install_state; /* State of hook installation */
     int                   installed_events; /* Which events are working */
+    int                   mmap_thresh_set; /* mmap threshold set by user */
+    int                   trim_thresh_set; /* trim threshold set by user */
     int                   hook_called; /* Our malloc hook was called */
+    size_t                max_freed_size; /* Maximal size released so far */
     size_t                (*usable_size)(void*); /* function pointer to get usable size */
 
     ucm_release_func_t    free; /* function pointer to release memory */
@@ -111,7 +116,10 @@ static ucm_malloc_hook_state_t ucm_malloc_hook_state = {
     .install_mutex    = PTHREAD_MUTEX_INITIALIZER,
     .install_state    = 0,
     .installed_events = 0,
+    .mmap_thresh_set  = 0,
+    .trim_thresh_set  = 0,
     .hook_called      = 0,
+    .max_freed_size   = 0,
     .usable_size      = malloc_usable_size,
     .free             = free,
     .heap_start       = (void*)-1,
@@ -122,6 +130,7 @@ static ucm_malloc_hook_state_t ucm_malloc_hook_state = {
     .num_env_strs     = 0
 };
 
+int ucm_dlmallopt_get(int); /* implemented in ptmalloc */
 
 static void ucm_malloc_mmaped_ptr_add(void *ptr)
 {
@@ -226,10 +235,36 @@ static void *ucm_malloc_impl(size_t size, const char *debug_name)
     return ptr;
 }
 
+static void ucm_malloc_adjust_thresholds(size_t size)
+{
+    int mmap_thresh;
+
+    if (size > ucm_malloc_hook_state.max_freed_size) {
+        if (ucm_global_config.enable_dynamic_mmap_thresh &&
+            !ucm_malloc_hook_state.trim_thresh_set &&
+            !ucm_malloc_hook_state.mmap_thresh_set) {
+            /* new mmap threshold is increased to the size of released block,
+             * new trim threshold is twice that size.
+             */
+            mmap_thresh = ucs_min(ucs_max(ucm_dlmallopt_get(M_MMAP_THRESHOLD), size),
+                                  UCM_DEFAULT_MMAP_THRESHOLD_MAX);
+            ucm_trace("adjust mmap threshold to %d", mmap_thresh);
+            ucm_dlmallopt(M_MMAP_THRESHOLD, mmap_thresh);
+            ucm_dlmallopt(M_TRIM_THRESHOLD, mmap_thresh * 2);
+        }
+
+        /* avoid adjusting the threshold for every released block, do it only
+         * if the size is larger than ever before.
+         */
+        ucm_malloc_hook_state.max_freed_size = size;
+    }
+}
+
 static inline void ucm_mem_free(void *ptr, size_t size)
 {
     VALGRIND_FREELIKE_BLOCK(ptr, 0);
     VALGRIND_MAKE_MEM_UNDEFINED(ptr, size); /* Make memory accessible to ptmalloc3 */
+    ucm_malloc_adjust_thresholds(size);
     ucm_dlfree(ptr);
 }
 
@@ -562,6 +597,24 @@ static void ucm_malloc_install_symbols(ucm_reloc_patch_t *patches)
     }
 }
 
+static int ucm_malloc_mallopt(int param_number, int value)
+{
+    int success;
+
+    success = ucm_dlmallopt(param_number, value);
+    if (success) {
+        switch (param_number) {
+        case M_TRIM_THRESHOLD:
+            ucm_malloc_hook_state.trim_thresh_set = 1;
+            break;
+        case M_MMAP_THRESHOLD:
+            ucm_malloc_hook_state.mmap_thresh_set = 1;
+            break;
+        }
+    }
+    return success;
+}
+
 static ucm_reloc_patch_t ucm_malloc_symbol_patches[] = {
     { "free", ucm_free },
     { "realloc", ucm_realloc },
@@ -579,7 +632,7 @@ static ucm_reloc_patch_t ucm_malloc_symbol_patches[] = {
 };
 
 static ucm_reloc_patch_t ucm_malloc_optional_symbol_patches[] = {
-    { "mallopt", ucm_dlmallopt },
+    { "mallopt", ucm_malloc_mallopt },
     { "mallinfo", ucm_dlmallinfo },
     { "malloc_stats", ucm_dlmalloc_stats },
     { "malloc_trim", ucm_dlmalloc_trim },
@@ -595,7 +648,7 @@ static void ucm_malloc_install_optional_symbols()
     }
 }
 
-static void ucm_malloc_install_mallopt()
+static void ucm_malloc_set_env_mallopt()
 {
     /* copy values of M_MMAP_THRESHOLD and M_TRIM_THRESHOLD
      * if they were overriden by the user
@@ -605,13 +658,13 @@ static void ucm_malloc_install_mallopt()
     p = getenv("MALLOC_TRIM_THRESHOLD_");
     if (p) {
         ucm_debug("set trim_thresh to %d", atoi(p));
-        ucm_dlmallopt(M_TRIM_THRESHOLD, atoi(p));
+        ucm_malloc_mallopt(M_TRIM_THRESHOLD, atoi(p));
     }
 
     p = getenv("MALLOC_MMAP_THRESHOLD_");
     if (p) {
         ucm_debug("set mmap_thresh to %d", atoi(p));
-        ucm_dlmallopt(M_MMAP_THRESHOLD, atoi(p));
+        ucm_malloc_mallopt(M_MMAP_THRESHOLD, atoi(p));
     }
 }
 
@@ -698,14 +751,22 @@ ucs_status_t ucm_malloc_install(int events)
 
 out_install_opt_syms:
     ucm_malloc_install_optional_symbols();
-    ucm_malloc_install_mallopt();
+    ucm_malloc_set_env_mallopt();
 out_succ:
     status = UCS_OK;
 out_unlock:
     pthread_mutex_unlock(&ucm_malloc_hook_state.install_mutex);
     return status;
 }
 
+void ucm_malloc_state_reset(int default_mmap_thresh, int default_trim_thresh)
+{
+    ucm_malloc_hook_state.max_freed_size = 0;
+    ucm_dlmallopt(M_MMAP_THRESHOLD, default_mmap_thresh);
+    ucm_dlmallopt(M_TRIM_THRESHOLD, default_trim_thresh);
+    ucm_malloc_set_env_mallopt();
+}
+
 UCS_STATIC_INIT {
     ucs_spinlock_init(&ucm_malloc_hook_state.lock);
     kh_init_inplace(mmap_ptrs, &ucm_malloc_hook_state.ptrs);

src/ucm/malloc/malloc_hook.h

@@ -11,4 +11,6 @@
 
 ucs_status_t ucm_malloc_install(int events);
 
+void ucm_malloc_state_reset(int default_mmap_thresh, int default_trim_thresh);
+
 #endif

src/ucm/util/ucm_config.c

@@ -9,6 +9,7 @@
 #include <ucm/util/log.h>
 #include <ucs/config/parser.h>
 #include <ucs/type/component.h>
+#include <ucs/sys/checker.h>
 #include <string.h>
 #include <stdlib.h>
 
@@ -20,18 +21,20 @@
 #define UCM_EN_MMAP_RELOC_VAR    "MMAP_RELOC"
 #define UCM_EN_MALLOC_HOOKS_VAR  "MALLOC_HOOKS"
 #define UCM_EN_MALLOC_RELOC_VAR  "MALLOC_RELOC"
+#define UCM_EN_DYNAMIC_MMAP_VAR  "DYNAMIC_MMAP_THRESH"
 #define UCM_EN_CUDA_HOOKS_VAR    "CUDA_HOOKS"
 
 
 ucm_config_t ucm_global_config = {
-    .log_level            = UCS_LOG_LEVEL_WARN,
-    .alloc_alignment      = 16,
-    .enable_events        = 1,
-    .enable_mmap_reloc    = 1,
-    .enable_malloc_hooks  = 1,
-    .enable_malloc_reloc  = 0,
+    .log_level                  = UCS_LOG_LEVEL_WARN,
+    .alloc_alignment            = 16,
+    .enable_events              = 1,
+    .enable_mmap_reloc          = 1,
+    .enable_malloc_hooks        = 1,
+    .enable_malloc_reloc        = 0,
+    .enable_dynamic_mmap_thresh = 1,
 #if HAVE_CUDA
-    .enable_cuda_hooks    = 1
+    .enable_cuda_hooks          = 1
 #endif
 };
 
@@ -111,6 +114,18 @@ void ucm_config_print(FILE *stream, ucs_config_print_flags_t print_flags)
                               print_flags);
     fprintf(stream, "%s%s=%s\n", UCM_ENV_PREFIX, UCM_EN_MALLOC_RELOC_VAR,
             ucm_config_bool_to_string(ucm_global_config.enable_malloc_reloc));
+
+
+    ucm_config_print_bool_doc(stream,
+                              "Enable dynamic mmap threshold: for every released block, the\n"
+                              "mmap threshold is adjusted upward to the size of the size of\n"
+                              "the block, and trim threshold is adjust to twice the size of\n"
+                              "the dynamic mmap threshold.",
+                              print_flags);
+    fprintf(stream, "%s%s=%s\n", UCM_ENV_PREFIX, UCM_EN_DYNAMIC_MMAP_VAR,
+            ucm_config_bool_to_string(ucm_global_config.enable_dynamic_mmap_thresh));
+
+
 #if HAVE_CUDA
     fprintf(stream, "%s%s=%s\n", UCM_ENV_PREFIX, UCM_EN_CUDA_HOOKS_VAR,
             ucm_config_bool_to_string(ucm_global_config.enable_cuda_hooks));
@@ -165,6 +180,8 @@ ucs_status_t ucm_config_modify(const char *name, const char *value)
         ucm_config_set_value_bool(value, &ucm_global_config.enable_malloc_hooks);
     } else if (!strcmp(name, UCM_EN_MALLOC_RELOC_VAR)) {
         ucm_config_set_value_bool(value, &ucm_global_config.enable_malloc_reloc);
+    } else if (!strcmp(name, UCM_EN_DYNAMIC_MMAP_VAR)) {
+        ucm_config_set_value_bool(value, &ucm_global_config.enable_dynamic_mmap_thresh);
 #if HAVE_CUDA
     } else if (!strcmp(name, UCM_EN_CUDA_HOOKS_VAR)) {
         ucm_config_set_value_bool(value, &ucm_global_config.enable_cuda_hooks);
@@ -189,10 +206,17 @@ static void ucm_config_set(const char *name)
 }
 
 UCS_STATIC_INIT {
+    if (RUNNING_ON_VALGRIND) {
+        /* Valgrind limits the size of brk() segments to 8mb, so must use mmap
+         * for large allocations.
+         */
+        ucm_global_config.enable_dynamic_mmap_thresh = 0;
+    }
     ucm_config_set(UCM_LOG_LEVEL_VAR);
     ucm_config_set(UCM_ALLOC_ALIGN_VAR);
     ucm_config_set(UCM_EN_EVENTS_VAR);
     ucm_config_set(UCM_EN_MMAP_RELOC_VAR);
     ucm_config_set(UCM_EN_MALLOC_HOOKS_VAR);
     ucm_config_set(UCM_EN_MALLOC_RELOC_VAR);
+    ucm_config_set(UCM_EN_DYNAMIC_MMAP_VAR);
 }

src/ucm/util/ucm_config.h

@@ -18,6 +18,7 @@ typedef struct ucm_config {
     int             enable_mmap_reloc;
     int             enable_malloc_hooks;
     int             enable_malloc_reloc;
+    int             enable_dynamic_mmap_thresh;
 #if HAVE_CUDA
     int             enable_cuda_hooks;
 #endif