cpython/Objects/mimalloc/init.c

/* ----------------------------------------------------------------------------
Copyright (c) 2018-2022, Microsoft Research, Daan Leijen
This is free software; you can redistribute it and/or modify it under the
terms of the MIT license. A copy of the license can be found in the file
"LICENSE" at the root of this distribution.
-----------------------------------------------------------------------------*/
#include "mimalloc.h"
#include "mimalloc/internal.h"
#include "mimalloc/prim.h"

#include <string.h>  // memcpy, memset
#include <stdlib.h>  // atexit


// Empty page used to initialize the small free pages array
const mi_page_t _mi_page_empty;

#define MI_PAGE_EMPTY() …

#if (MI_SMALL_WSIZE_MAX==128)
#if (MI_PADDING>0) && (MI_INTPTR_SIZE >= 8)
#define MI_SMALL_PAGES_EMPTY …
#elif (MI_PADDING>0)
#define MI_SMALL_PAGES_EMPTY …
#else
#define MI_SMALL_PAGES_EMPTY …
#endif
#else
#error "define right initialization sizes corresponding to MI_SMALL_WSIZE_MAX"
#endif

// Empty page queues for every bin
#define QNULL(sz) …
#define MI_PAGE_QUEUES_EMPTY …

#define MI_STAT_COUNT_NULL() …

// Empty statistics
#if MI_STAT>1
#define MI_STAT_COUNT_END_NULL …
#else
#define MI_STAT_COUNT_END_NULL() …
#endif

#define MI_STATS_NULL …


// Empty slice span queues for every bin
#define SQNULL(sz) …
#define MI_SEGMENT_SPAN_QUEUES_EMPTY …


// --------------------------------------------------------
// Statically allocate an empty heap as the initial
// thread local value for the default heap,
// and statically allocate the backing heap for the main
// thread so it can function without doing any allocation
// itself (as accessing a thread local for the first time
// may lead to allocation itself on some platforms)
// --------------------------------------------------------

mi_decl_cache_align const mi_heap_t _mi_heap_empty = …;

#define tld_empty_stats …
#define tld_empty_os …

mi_decl_cache_align static const mi_tld_t tld_empty = …;

mi_threadid_t _mi_thread_id(void) mi_attr_noexcept { … }

// the thread-local default heap for allocation
mi_decl_thread mi_heap_t* _mi_heap_default = …;

extern mi_heap_t _mi_heap_main;

static mi_tld_t tld_main = …;

mi_heap_t _mi_heap_main = …;

bool _mi_process_is_initialized = …;  // set to `true` in `mi_process_init`.

mi_stats_t _mi_stats_main = …;


static void mi_heap_main_init(void) { … }

mi_heap_t* _mi_heap_main_get(void) { … }


/* -----------------------------------------------------------
  Initialization and freeing of the thread local heaps
----------------------------------------------------------- */

// note: in x64 in release build `sizeof(mi_thread_data_t)` is under 4KiB (= OS page size).
mi_thread_data_t;


// Thread meta-data is allocated directly from the OS. For
// some programs that do not use thread pools and allocate and
// destroy many OS threads, this may causes too much overhead
// per thread so we maintain a small cache of recently freed metadata.

#define TD_CACHE_SIZE …
static _Atomic(mi_thread_data_t*) td_cache[TD_CACHE_SIZE];

static mi_thread_data_t* mi_thread_data_zalloc(void) { … }

static void mi_thread_data_free( mi_thread_data_t* tdfree ) { … }

void _mi_thread_data_collect(void) { … }

// Initialize the thread local default heap, called from `mi_thread_init`
static bool _mi_heap_init(void) { … }

void _mi_tld_init(mi_tld_t* tld, mi_heap_t* bheap) { … }

// Free the thread local default heap (called from `mi_thread_done`)
static bool _mi_heap_done(mi_heap_t* heap) { … }



// --------------------------------------------------------
// Try to run `mi_thread_done()` automatically so any memory
// owned by the thread but not yet released can be abandoned
// and re-owned by another thread.
//
// 1. windows dynamic library:
//     call from DllMain on DLL_THREAD_DETACH
// 2. windows static library:
//     use `FlsAlloc` to call a destructor when the thread is done
// 3. unix, pthreads:
//     use a pthread key to call a destructor when a pthread is done
//
// In the last two cases we also need to call `mi_process_init`
// to set up the thread local keys.
// --------------------------------------------------------

// Set up handlers so `mi_thread_done` is called automatically
static void mi_process_setup_auto_thread_done(void) { … }


bool _mi_is_main_thread(void) { … }

static _Atomic(size_t) thread_count = …;

size_t  _mi_current_thread_count(void) { … }

// This is called from the `mi_malloc_generic`
void mi_thread_init(void) mi_attr_noexcept
{ … }

void mi_thread_done(void) mi_attr_noexcept { … }

void _mi_thread_done(mi_heap_t* heap)
{ … }

void _mi_heap_set_default_direct(mi_heap_t* heap)  { … }


// --------------------------------------------------------
// Run functions on process init/done, and thread init/done
// --------------------------------------------------------
static void mi_cdecl mi_process_done(void);

static bool os_preloading = …;    // true until this module is initialized
static bool mi_redirected = …;   // true if malloc redirects to mi_malloc

// Returns true if this module has not been initialized; Don't use C runtime routines until it returns false.
bool mi_decl_noinline _mi_preloading(void) { … }

mi_decl_nodiscard bool mi_is_redirected(void) mi_attr_noexcept { … }

// Communicate with the redirection module on Windows
#if defined(_WIN32) && defined(MI_SHARED_LIB) && !defined(MI_WIN_NOREDIRECT)
#ifdef __cplusplus
extern "C" {
#endif
mi_decl_export void _mi_redirect_entry(DWORD reason) {
  // called on redirection; careful as this may be called before DllMain
  if (reason == DLL_PROCESS_ATTACH) {
    mi_redirected = true;
  }
  else if (reason == DLL_PROCESS_DETACH) {
    mi_redirected = false;
  }
  else if (reason == DLL_THREAD_DETACH) {
    mi_thread_done();
  }
}
__declspec(dllimport) bool mi_cdecl mi_allocator_init(const char** message);
__declspec(dllimport) void mi_cdecl mi_allocator_done(void);
#ifdef __cplusplus
}
#endif
#else
static bool mi_allocator_init(const char** message) { … }
static void mi_allocator_done(void) { … }
#endif

// Called once by the process loader
static void mi_process_load(void) { … }

#if defined(_WIN32) && (defined(_M_IX86) || defined(_M_X64))
#include <intrin.h>
mi_decl_cache_align bool _mi_cpu_has_fsrm = false;

static void mi_detect_cpu_features(void) {
  // FSRM for fast rep movsb support (AMD Zen3+ (~2020) or Intel Ice Lake+ (~2017))
  int32_t cpu_info[4];
  __cpuid(cpu_info, 7);
  _mi_cpu_has_fsrm = ((cpu_info[3] & (1 << 4)) != 0); // bit 4 of EDX : see <https://en.wikipedia.org/wiki/CPUID#EAX=7,_ECX=0:_Extended_Features>
}
#else
static void mi_detect_cpu_features(void) { … }
#endif

// Initialize the process; called by thread_init or the process loader
void mi_process_init(void) mi_attr_noexcept { … }

// Called when the process is done (through `at_exit`)
static void mi_cdecl mi_process_done(void) { … }



#if defined(_WIN32) && defined(MI_SHARED_LIB)
  // Windows DLL: easy to hook into process_init and thread_done
  __declspec(dllexport) BOOL WINAPI DllMain(HINSTANCE inst, DWORD reason, LPVOID reserved) {
    MI_UNUSED(reserved);
    MI_UNUSED(inst);
    if (reason==DLL_PROCESS_ATTACH) {
      mi_process_load();
    }
    else if (reason==DLL_PROCESS_DETACH) {
      mi_process_done();
    }
    else if (reason==DLL_THREAD_DETACH) {
      if (!mi_is_redirected()) {
        mi_thread_done();
      }
    }
    return TRUE;
  }

#elif defined(_MSC_VER)
  // MSVC: use data section magic for static libraries
  // See <https://www.codeguru.com/cpp/misc/misc/applicationcontrol/article.php/c6945/Running-Code-Before-and-After-Main.htm>
  static int _mi_process_init(void) {
    mi_process_load();
    return 0;
  }
  typedef int(*_mi_crt_callback_t)(void);
  #if defined(_M_X64) || defined(_M_ARM64)
    __pragma(comment(linker, "/include:" "_mi_msvc_initu"))
    #pragma section(".CRT$XIU", long, read)
  #else
    __pragma(comment(linker, "/include:" "__mi_msvc_initu"))
  #endif
  #pragma data_seg(".CRT$XIU")
  mi_decl_externc _mi_crt_callback_t _mi_msvc_initu[] = { &_mi_process_init };
  #pragma data_seg()

#elif defined(__cplusplus)
  // C++: use static initialization to detect process start
  static bool _mi_process_init(void) {
    mi_process_load();
    return (_mi_heap_main.thread_id != 0);
  }
  static bool mi_initialized = _mi_process_init();

#elif defined(__GNUC__) || defined(__clang__)
  // GCC,Clang: use the constructor attribute
  static void __attribute__((constructor)) _mi_process_init(void) { … }

#else
#pragma message("define a way to call mi_process_load on your platform")
#endif