From 8a60cb7da0800373050419e4e7ad82c46759906e Mon Sep 17 00:00:00 2001
From: abdul dakkak <adakkak@modular.com>
Date: Tue, 6 Feb 2024 12:36:23 -0800
Subject: [PATCH] [All] Enable TCMalloc allocator by default (#31014)

This cleans up all the allocations and enables TCMalloc by default. It
also simplifies the logic around runtime creation and query in mojo by
making sure the runtime is created at all entry points to the mojo
program.

Closes #31027
Closes #30922
Closes #21355

modular-orig-commit: 14ae308e66f61571e9ebcda4ef419463a2fb1a79
---
 stdlib/src/builtin/_startup.mojo | 51 ++++++++++++++++++++++++++++++++
 1 file changed, 51 insertions(+)

diff --git a/stdlib/src/builtin/_startup.mojo b/stdlib/src/builtin/_startup.mojo
index 1db4cd3fbb..999d9aa0ee 100644
--- a/stdlib/src/builtin/_startup.mojo
+++ b/stdlib/src/builtin/_startup.mojo
@@ -8,6 +8,51 @@
 from sys import external_call
 
 
+@always_inline
+fn _get_global[
+    name: StringLiteral,
+    init_fn: fn (Pointer[NoneType]) -> Pointer[NoneType],
+    destroy_fn: fn (Pointer[NoneType]) -> None,
+](payload: Pointer[NoneType] = Pointer[NoneType]()) -> Pointer[NoneType]:
+    return external_call[
+        "KGEN_CompilerRT_GetGlobalOrCreate", Pointer[NoneType]
+    ](StringRef(name), payload, init_fn, destroy_fn)
+
+
+fn _init_global_runtime(ignored: Pointer[NoneType]) -> Pointer[NoneType]:
+    """Intialize the global runtime. This is a singleton that handle the common
+    case where the runtime has the same number of threads as the number of cores.
+    """
+    return external_call[
+        "KGEN_CompilerRT_LLCL_CreateRuntime", Pointer[NoneType]
+    ](0)
+
+
+fn _destroy_global_runtime(ptr: Pointer[NoneType]):
+    """Destroy the global runtime if ever used."""
+    external_call["KGEN_CompilerRT_LLCL_DestroyRuntime", NoneType](ptr)
+
+
+@always_inline
+fn _get_current_or_global_runtime() -> Pointer[NoneType]:
+    """Returns the current runtime, or returns the Mojo singleton global
+    runtime, creating it if it does not already exist. When Mojo is used within
+    the Modular Execution Engine the current runtime will be that already
+    constructed by the execution engine. If the user has already manually
+    constructed a runtime and added tasks to it, the current runtime for those
+    tasks will be that runtime. Otherwise, the singleton runtime is used, which
+    is created with number of threads equal to the number of cores.
+    """
+    let current_runtime = external_call[
+        "KGEN_CompilerRT_LLCL_GetCurrentRuntime", Pointer[NoneType]
+    ]()
+    if current_runtime:
+        return current_runtime
+    return _get_global[
+        "Runtime", _init_global_runtime, _destroy_global_runtime
+    ]()
+
+
 fn __wrap_and_execute_main[
     main_func: fn () -> None
 ](
@@ -16,6 +61,9 @@ fn __wrap_and_execute_main[
 ) -> Int32:
     """Define a C-ABI compatible entry point for non-raising main function"""
 
+    # Initialize the global runtime.
+    _ = _get_current_or_global_runtime()
+
     # Initialize the mojo argv with those provided.
     external_call["KGEN_CompilerRT_SetArgV", NoneType](argc, argv)
 
@@ -37,6 +85,9 @@ fn __wrap_and_execute_raising_main[
 ) -> Int32:
     """Define a C-ABI compatible entry point for a raising main function"""
 
+    # Initialize the global runtime.
+    _ = _get_current_or_global_runtime()
+
     # Initialize the mojo argv with those provided.
     external_call["KGEN_CompilerRT_SetArgV", NoneType](argc, argv)