[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

@@ -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)