Tweak TLS init/deinit (#4551)

* Delete the entire TLS in one go
* Put TLS allocation separately
* Keep TLS in a single array

kotlin-native/backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/llvm/ContextUtils.kt

@@ -515,7 +515,6 @@ internal class Llvm(val context: Context, val llvmModule: LLVMModuleRef) {
     val throwExceptionFunction = importRtFunction("ThrowException")
     val appendToInitalizersTail = importRtFunction("AppendToInitializersTail")
     val addTLSRecord = importRtFunction("AddTLSRecord")
-    val clearTLSRecord = importRtFunction("ClearTLSRecord")
     val lookupTLS = importRtFunction("LookupTLS")
     val initRuntimeIfNeeded = importRtFunction("Kotlin_initRuntimeIfNeeded")
     val mutationCheck = importRtFunction("MutationCheck")

kotlin-native/backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/llvm/IrToBitcode.kt

@@ -15,10 +15,8 @@ import org.jetbrains.kotlin.backend.konan.*
 import org.jetbrains.kotlin.backend.konan.descriptors.*
 import org.jetbrains.kotlin.backend.konan.ir.*
 import org.jetbrains.kotlin.backend.konan.llvm.coverage.LLVMCoverageInstrumentation
-import org.jetbrains.kotlin.backend.konan.serialization.KonanIrModuleFragmentImpl
 import org.jetbrains.kotlin.builtins.UnsignedType
 import org.jetbrains.kotlin.descriptors.Modality
-import org.jetbrains.kotlin.descriptors.ModuleDescriptor
 import org.jetbrains.kotlin.ir.IrElement
 import org.jetbrains.kotlin.ir.IrStatement
 import org.jetbrains.kotlin.ir.UNDEFINED_OFFSET
@@ -37,7 +35,6 @@ import org.jetbrains.kotlin.library.KotlinLibrary
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.resolve.descriptorUtil.classId
-import org.jetbrains.kotlin.resolve.descriptorUtil.module
 
 internal enum class FieldStorageKind {
     GLOBAL, // In the old memory model these are only accessible from the "main" thread.
@@ -385,9 +382,10 @@ internal class CodeGeneratorVisitor(val context: Context, val lifetimes: Map<IrE
 
     //-------------------------------------------------------------------------//
 
-    val INIT_GLOBALS = 0
+    // Must be synchronized with Runtime.cpp
-    val INIT_THREAD_LOCAL_GLOBALS = 1
+    val ALLOC_THREAD_LOCAL_GLOBALS = 0
-    val DEINIT_THREAD_LOCAL_GLOBALS = 2
+    val INIT_GLOBALS = 1
+    val INIT_THREAD_LOCAL_GLOBALS = 2
     val DEINIT_GLOBALS = 3
 
     private fun createInitBody(): LLVMValueRef {
@@ -397,7 +395,7 @@ internal class CodeGeneratorVisitor(val context: Context, val lifetimes: Map<IrE
             using(FunctionScope(initFunction, "init_body", it)) {
                 val bbInit = basicBlock("init", null)
                 val bbLocalInit = basicBlock("local_init", null)
-                val bbLocalDeinit = basicBlock("local_deinit", null)
+                val bbLocalAlloc = basicBlock("local_alloc", null)
                 val bbGlobalDeinit = basicBlock("global_deinit", null)
                 val bbDefault = basicBlock("default", null) {
                     unreachable()
@@ -406,18 +404,12 @@ internal class CodeGeneratorVisitor(val context: Context, val lifetimes: Map<IrE
                 switch(LLVMGetParam(initFunction, 0)!!,
                         listOf(Int32(INIT_GLOBALS).llvm                to bbInit,
                                Int32(INIT_THREAD_LOCAL_GLOBALS).llvm   to bbLocalInit,
-                               Int32(DEINIT_THREAD_LOCAL_GLOBALS).llvm to bbLocalDeinit,
+                               Int32(ALLOC_THREAD_LOCAL_GLOBALS).llvm  to bbLocalAlloc,
                                Int32(DEINIT_GLOBALS).llvm              to bbGlobalDeinit),
                         bbDefault)
 
                 // Globals initalizers may contain accesses to objects, so visit them first.
                 appendingTo(bbInit) {
-                    // Bit clumsy, global init may need access to TLS, thus it has to be ready to that point.
-                    if (context.llvm.tlsCount > 0) {
-                        val memory = LLVMGetParam(initFunction, 1)!!
-                        call(context.llvm.addTLSRecord, listOf(memory, context.llvm.tlsKey,
-                                Int32(context.llvm.tlsCount).llvm))
-                    }
                     context.llvm.fileInitializers
                             .forEach { irField ->
                                 if (irField.initializer?.expression !is IrConst<*>?) {
@@ -436,11 +428,6 @@ internal class CodeGeneratorVisitor(val context: Context, val lifetimes: Map<IrE
                 }
 
                 appendingTo(bbLocalInit) {
-                    if (context.llvm.tlsCount > 0) {
-                        val memory = LLVMGetParam(initFunction, 1)!!
-                        call(context.llvm.addTLSRecord, listOf(memory, context.llvm.tlsKey,
-                                Int32(context.llvm.tlsCount).llvm))
-                    }
                     context.llvm.fileInitializers
                             .forEach { irField ->
                                 if (irField.initializer != null && irField.storageKind == FieldStorageKind.THREAD_LOCAL) {
@@ -454,10 +441,11 @@ internal class CodeGeneratorVisitor(val context: Context, val lifetimes: Map<IrE
                     ret(null)
                 }
 
-                appendingTo(bbLocalDeinit) {
+                appendingTo(bbLocalAlloc) {
                     if (context.llvm.tlsCount > 0) {
                         val memory = LLVMGetParam(initFunction, 1)!!
-                        call(context.llvm.clearTLSRecord, listOf(memory, context.llvm.tlsKey))
+                        call(context.llvm.addTLSRecord, listOf(memory, context.llvm.tlsKey,
+                                Int32(context.llvm.tlsCount).llvm))
                     }
                     ret(null)
                 }

kotlin-native/runtime/src/legacymm/cpp/Memory.cpp

@@ -131,7 +131,6 @@ typedef KStdUnorderedSet<KRef> KRefSet;
 typedef KStdUnorderedMap<KRef, KInt> KRefIntMap;
 typedef KStdDeque<KRef> KRefDeque;
 typedef KStdDeque<KRefList> KRefListDeque;
-typedef KStdUnorderedMap<void**, std::pair<KRef*,int>> KThreadLocalStorageMap;
 
 // A little hack that allows to enable -O2 optimizations
 // Prevents clang from replacing FrameOverlay struct
@@ -584,15 +583,78 @@ private:
   }
 };
 
+namespace {
+
+class ThreadLocalStorage {
+public:
+    using Key = void**;
+
+    void Init() noexcept { map_ = konanConstructInstance<Map>(); }
+
+    void Deinit() noexcept {
+        RuntimeAssert(map_->size() == 0, "Must be already cleared");
+        konanDestructInstance(map_);
+    }
+
+    void Add(Key key, int size) noexcept {
+        RuntimeAssert(storage_ == nullptr, "Storage must not be committed");
+        auto it = map_->find(key);
+        if (it != map_->end()) {
+            RuntimeAssert(it->second.second == size, "Attempt to add TLS record with the same key and different size");
+            return;
+        }
+        map_->emplace(key, std::make_pair(size_, size));
+        size_ += size;
+    }
+
+    void Commit() noexcept {
+        RuntimeAssert(storage_ == nullptr, "Cannot commit storage twice");
+        storage_ = reinterpret_cast<KRef*>(konanAllocMemory(size_ * sizeof(KRef)));
+    }
+
+    void Clear() noexcept {
+        RuntimeAssert(storage_ != nullptr, "Storage must be committed");
+        for (int i = 0; i < size_; ++i) {
+            UpdateHeapRef(storage_ + i, nullptr);
+        }
+        konanFreeMemory(storage_);
+        map_->clear();
+    }
+
+    KRef* Lookup(Key key, int index) noexcept {
+        RuntimeAssert(storage_ != nullptr, "Storage must be committed");
+        // In many cases there is only one module, so this is one element cache.
+        if (lastKey_ == key) {
+            return storage_ + lastOffset_ + index;
+        }
+        auto it = map_->find(key);
+        RuntimeAssert(it != map_->end(), "Must be there");
+        int offset = it->second.first;
+        RuntimeAssert(offset + index < size_, "Out of bound in TLS access");
+        lastKey_ = key;
+        lastOffset_ = offset;
+        return storage_ + offset + index;
+    }
+
+private:
+    using Map = KStdUnorderedMap<Key, std::pair<int, int>>;
+
+    Map* map_ = nullptr;
+    KRef* storage_ = nullptr;
+    int size_ = 0;
+    int lastOffset_ = 0;
+    Key lastKey_ = nullptr;
+};
+
+} // namespace
+
 struct MemoryState {
 #if TRACE_MEMORY
   // Set of all containers.
   ContainerHeaderSet* containers;
 #endif
 
-  KThreadLocalStorageMap* tlsMap;
+  ThreadLocalStorage tls;
-  KRef* tlsMapLastStart;
-  void* tlsMapLastKey;
 
 #if USE_GC
   // Finalizer queue - linked list of containers scheduled for finalization.
@@ -1995,7 +2057,7 @@ MemoryState* initMemory(bool firstRuntime) {
   memoryState->allocSinceLastGcThreshold = kMaxGcAllocThreshold;
   memoryState->gcErgonomics = true;
 #endif
-  memoryState->tlsMap = konanConstructInstance<KThreadLocalStorageMap>();
+  memoryState->tls.Init();
   memoryState->foreignRefManager = ForeignRefManager::create();
   bool firstMemoryState = atomicAdd(&aliveMemoryStatesCount, 1) == 1;
   switch (Kotlin_getDestroyRuntimeMode()) {
@@ -2051,8 +2113,7 @@ void deinitMemory(MemoryState* memoryState, bool destroyRuntime) {
   konanDestructInstance(memoryState->toFree);
   konanDestructInstance(memoryState->roots);
   konanDestructInstance(memoryState->toRelease);
-  RuntimeAssert(memoryState->tlsMap->size() == 0, "Must be already cleared");
+  memoryState->tls.Deinit();
-  konanDestructInstance(memoryState->tlsMap);
   RuntimeAssert(memoryState->finalizerQueue == nullptr, "Finalizer queue must be empty");
   RuntimeAssert(memoryState->finalizerQueueSize == 0, "Finalizer queue must be empty");
 #endif // USE_GC
@@ -3535,44 +3596,19 @@ void Kotlin_Any_share(ObjHeader* obj) {
 }
 
 RUNTIME_NOTHROW void AddTLSRecord(MemoryState* memory, void** key, int size) {
-  auto* tlsMap = memory->tlsMap;
+    memory->tls.Add(key, size);
-  auto it = tlsMap->find(key);
-  if (it != tlsMap->end()) {
-    RuntimeAssert(it->second.second == size, "Size must be consistent");
-    return;
-  }
-  KRef* start = reinterpret_cast<KRef*>(konanAllocMemory(size * sizeof(KRef)));
-  tlsMap->emplace(key, std::make_pair(start, size));
 }
 
-RUNTIME_NOTHROW void ClearTLSRecord(MemoryState* memory, void** key) {
+RUNTIME_NOTHROW void CommitTLSStorage(MemoryState* memory) {
-  auto* tlsMap = memory->tlsMap;
+    memory->tls.Commit();
-  auto it = tlsMap->find(key);
+}
-  if (it != tlsMap->end()) {
+
-    KRef* start = it->second.first;
+RUNTIME_NOTHROW void ClearTLS(MemoryState* memory) {
-    int count = it->second.second;
+    memory->tls.Clear();
-    for (int i = 0; i < count; i++) {
-      UpdateHeapRef(start + i, nullptr);
-    }
-    konanFreeMemory(start);
-    tlsMap->erase(it);
-  }
 }
 
 RUNTIME_NOTHROW KRef* LookupTLS(void** key, int index) {
-  auto* state = memoryState;
+    return memoryState->tls.Lookup(key, index);
-  auto* tlsMap = state->tlsMap;
-  // In many cases there is only one module, so this one element cache.
-  if (state->tlsMapLastKey == key) {
-    return state->tlsMapLastStart + index;
-  }
-  auto it = tlsMap->find(key);
-  RuntimeAssert(it != tlsMap->end(), "Must be there");
-  RuntimeAssert(index < it->second.second, "Out of bound in TLS access");
-  KRef* start = it->second.first;
-  state->tlsMapLastKey = key;
-  state->tlsMapLastStart = start;
-  return start + index;
 }
 
 

kotlin-native/runtime/src/main/cpp/CompilerExport.cpp

@@ -35,7 +35,6 @@ void EnsureDeclarationsEmitted() {
     ensureUsed(EnterFrame);
     ensureUsed(LeaveFrame);
     ensureUsed(AddTLSRecord);
-    ensureUsed(ClearTLSRecord);
     ensureUsed(LookupTLS);
     ensureUsed(MutationCheck);
     ensureUsed(CheckLifetimesConstraint);

kotlin-native/runtime/src/main/cpp/Memory.h

@@ -225,8 +225,10 @@ void FreezeSubgraph(ObjHeader* obj);
 void EnsureNeverFrozen(ObjHeader* obj);
 // Add TLS object storage, called by the generated code.
 void AddTLSRecord(MemoryState* memory, void** key, int size) RUNTIME_NOTHROW;
-// Clear TLS object storage, called by the generated code.
+// Allocate storage for TLS. `AddTLSRecord` cannot be called after this.
-void ClearTLSRecord(MemoryState* memory, void** key) RUNTIME_NOTHROW;
+void CommitTLSStorage(MemoryState* memory) RUNTIME_NOTHROW;
+// Clear TLS object storage.
+void ClearTLS(MemoryState* memory) RUNTIME_NOTHROW;
 // Lookup element in TLS object storage.
 ObjHeader** LookupTLS(void** key, int index) RUNTIME_NOTHROW;
 

kotlin-native/runtime/src/main/cpp/Runtime.cpp

@@ -55,10 +55,11 @@ struct RuntimeState {
     RuntimeStatus status = RuntimeStatus::kUninitialized;
 };
 
+// Must be synchronized with IrToBitcode.kt
 enum {
-  INIT_GLOBALS = 0,
+  ALLOC_THREAD_LOCAL_GLOBALS = 0,
-  INIT_THREAD_LOCAL_GLOBALS = 1,
+  INIT_GLOBALS = 1,
-  DEINIT_THREAD_LOCAL_GLOBALS = 2,
+  INIT_THREAD_LOCAL_GLOBALS = 2,
   DEINIT_GLOBALS = 3
 };
 
@@ -120,6 +121,8 @@ RuntimeState* initRuntime() {
           result->worker = WorkerInit(true);
   }
 
+  InitOrDeinitGlobalVariables(ALLOC_THREAD_LOCAL_GLOBALS, result->memoryState);
+  CommitTLSStorage(result->memoryState);
   // Keep global variables in state as well.
   if (firstRuntime) {
     konan::consoleInit();
@@ -148,7 +151,7 @@ void deinitRuntime(RuntimeState* state, bool destroyRuntime) {
       // Nothing to do.
       break;
   }
-  InitOrDeinitGlobalVariables(DEINIT_THREAD_LOCAL_GLOBALS, state->memoryState);
+  ClearTLS(state->memoryState);
   if (destroyRuntime)
     InitOrDeinitGlobalVariables(DEINIT_GLOBALS, state->memoryState);
   auto workerId = GetWorkerId(state->worker);

kotlin-native/runtime/src/mm/cpp/Stubs.cpp

@@ -166,7 +166,11 @@ RUNTIME_NOTHROW void AddTLSRecord(MemoryState* memory, void** key, int size) {
     RuntimeCheck(false, "Unimplemented");
 }
 
-RUNTIME_NOTHROW void ClearTLSRecord(MemoryState* memory, void** key) {
+RUNTIME_NOTHROW void CommitTLSStorage(MemoryState* memory) {
+    RuntimeCheck(false, "Unimplemented");
+}
+
+RUNTIME_NOTHROW void ClearTLS(MemoryState* memory) {
     RuntimeCheck(false, "Unimplemented");
 }