[lld][WebAssembly] Allow first thread to use static TLS region

[sbc100]

Summary

It turns out we were already allocating static address space for TLS
data along with the non-TLS static data, but this space was going
unused/ignored.

With this change, we include the TLS segment in __wasm_init_memory
(which does the work of loading the passive segments into memory when a
module is first loaded). We also set the __tls_base global to point
to the start of this segment.

This means that the runtime can use this static copy of the TLS data for
the first/primary thread if it chooses, rather than doing a runtime
allocation prior to calling __wasm_init_tls.

Practically speaking, this will allow emscripten to avoid dynamic
allocation of TLS region on the main thread.

Part of work towards https://github.com/emscripten-core/emscripten/issues/16948

[sbc100]

• tweaks

[tlively]

lld/test/wasm/tls.s

109 ↗

(On Diff #431099)

Where does this number come from?

lld/wasm/Writer.cpp

280

Why these changes?

973–975

Why did this change?

1212–1218

Is there only ever one TLS segment? It feels weird to have this in a loop.

1272–1273

Can you elaborate on this comment? I don't understand why the condition should be removed. Surely we need to continue to avoid doing data.drop on TLS segments?

[sbc100]

• simplify

[sbc100]

OK I simplified and refactored and re-wrote the description. Maybe it will be more clear what is going on here now?

lld/wasm/Writer.cpp

280

P

280

Previously __tls_base was not set at all (just set to zero) in the shared memory case, and we relied on the embedder always calling __wasm_tls_init to set it.

After this change the embedder only needs to call __wasm_tls_init on secondary threads.

In the primary thread we initialize __tls_basehere to the offset of the segment from __memory_base, then we call
__wasm_tls_init as part of __wasm_memory_init.. which adjusts it accordingly.

[tlively]

Thanks, this makes more sense now. LGTM with fixes and extra information in that comment.

Also, in the third paragraph of the description:

"This means that the runtime and" => "This means that the runtime can"

lld/wasm/Writer.cpp

1195

[sbc100]

• fix test

lld/test/wasm/data-segments.ll

 ; PASSIVE-PIC-NEXT:        Body:            0B
 ; PASSIVE-PIC-NEXT:      - Index:           1
 ; PASSIVE-PIC-NEXT:        Locals:          []
 ; PASSIVE-PIC-NEXT:        Body:            {{.*}}
 ; PASSIVE-PIC-NEXT:      - Index:           2
 ; PASSIVE-PIC-NEXT:        Locals:
 ; PASSIVE32-PIC-NEXT:          - Type:            I32
 ; PASSIVE64-PIC-NEXT:          - Type:            I64
-; PASSIVE-PIC-NEXT:              Count:           1
+; PASSIVE-PIC-NEXT:              Count:           2
 ; PASSIVE-PIC-NEXT:        Body:            {{.*}}
 ; PASSIVE-PIC-NEXT:      - Index:           3
 ; PASSIVE-PIC-NEXT:        Locals:          []
 ; PASSIVE-PIC-NEXT:        Body:            0B
 ; PASSIVE-PIC-NEXT:  - Type:            DATA
 ; PASSIVE-PIC-NEXT:    Segments:
 ; PASSIVE-PIC-NEXT:      - SectionOffset:   3
 ; PASSIVE-PIC-NEXT:        InitFlags:       1

 
 ; DIS-NEXT:             i32.const       0
 ; DIS-NEXT:             i32.const       1
 ; DIS-NEXT:             i32.atomic.rmw.cmpxchg  0
 ; DIS-NEXT:             br_table        {0, 1, 2}      # 1:     down to label1
 ; DIS-NEXT:                                            # 2:     down to label0
 ; DIS-NEXT:            end
 
+; NOPIC-DIS-NEXT:      [[PTR]].const   1024
+; NOPIC-DIS-NEXT:      [[PTR]].const   1024
+; NOPIC-DIS-NEXT:      global.set      1
+; PIC-DIS-NEXT:        [[PTR]].const   0
+; PIC-DIS-NEXT:        global.get      1
+; PIC-DIS-NEXT:        [[PTR]].add
+; PIC-DIS-NEXT:        local.tee       1
+; PIC-DIS-NEXT:        global.set      {{\d*}}
+; PIC-DIS-NEXT:        local.get       1
+; DIS-NEXT:            i32.const       0
+; DIS-NEXT:            i32.const       4
+; DIS-NEXT:            memory.init  0, 0
+
 ; NOPIC-DIS-NEXT:      [[PTR]].const   1028
 ; PIC-DIS-NEXT:        [[PTR]].const   4
 ; PIC-DIS-NEXT:        global.get      1
 ; PIC-DIS-NEXT:        [[PTR]].add
 
 ; DIS-NEXT:            i32.const       0
 ; DIS-NEXT:            i32.const       13
 ; DIS-NEXT:            memory.init     1, 0

lld/test/wasm/tls_init_symbols.s

 # CHECK:       - Type:            CUSTOM
 # CHECK-NEXT:    Name:            name
 # CHECK-NEXT:    FunctionNames:
 # CHECK-NEXT:      - Index:           0
 # CHECK-NEXT:        Name:            __wasm_call_ctors
 # CHECK-NEXT:      - Index:           1
 # CHECK-NEXT:        Name:            __wasm_init_tls
 # CHECK-NEXT:      - Index:           2
-# CHECK-NEXT:        Name:            __wasm_apply_global_tls_relocs
+# CHECK-NEXT:        Name:            __wasm_init_memory
 # CHECK-NEXT:      - Index:           3
+# CHECK-NEXT:        Name:            __wasm_apply_global_tls_relocs
+# CHECK-NEXT:      - Index:           4
 # CHECK-NEXT:        Name:            _start
 
 # DIS:       <__wasm_init_tls>:
 # DIS:        local.get 0
 # DIS-NEXT:   global.set  1
 # DIS-NEXT:   local.get 0
 # DIS-NEXT:   i32.const 0
 # DIS-NEXT:   i32.const 4
 # DIS-NEXT:   memory.init 0, 0
-# DIS-NEXT:   call  2
+# DIS-NEXT:   call  3
 # DIS-NEXT:   end
 
 # DIS:      <__wasm_apply_global_tls_relocs>:
 # DIS:        global.get  1
 # DIS-NEXT:   i32.const 0
 # DIS-NEXT:   i32.add
 # DIS-NEXT:   global.set  4
 # DIS-NEXT:   end
 
 # DIS:      <_start>:
 # DIS:        global.get  4
 # DIS-NEXT:   end
 

lld/wasm/Writer.cpp

void Writer::writeSections() {
   uint8_t *buf = buffer->getBufferStart();
   parallelForEach(outputSections, [buf](OutputSection *s) {
     assert(s->isNeeded());
     s->writeTo(buf);
   });
 }
 
 static void setGlobalPtr(DefinedGlobal *g, uint64_t memoryPtr) {
+  LLVM_DEBUG(dbgs() << "setGlobalPtr " << g->getName() << " -> " << memoryPtr << "\n");
   g->global->setPointerValue(memoryPtr);
 }
 
 // Fix the memory layout of the output binary.  This assigns memory offsets
 // to each of the input data sections as well as the explicit stack region.
 // The default memory layout is as follows, from low to high.
 //
 //  - initialized data (starting at Config->globalBase)
void Writer::layoutMemory() {
   for (OutputSegment *seg : segments) {
     out.dylinkSec->memAlign = std::max(out.dylinkSec->memAlign, seg->alignment);
     memoryPtr = alignTo(memoryPtr, 1ULL << seg->alignment);
     seg->startVA = memoryPtr;
     log(formatv("mem: {0,-15} offset={1,-8} size={2,-8} align={3}", seg->name,
                 memoryPtr, seg->size, seg->alignment));
 
     if (!config->relocatable && seg->isTLS()) {
-      if (config->sharedMemory) {
+      if (WasmSym::tlsSize) {
         auto *tlsSize = cast<DefinedGlobal>(WasmSym::tlsSize);
         setGlobalPtr(tlsSize, seg->size);
-
+      }
+      if (WasmSym::tlsAlign) {
         auto *tlsAlign = cast<DefinedGlobal>(WasmSym::tlsAlign);
         setGlobalPtr(tlsAlign, int64_t{1} << seg->alignment);
-      } else if (WasmSym::tlsBase) {
+      }
+      if (!config->sharedMemory && WasmSym::tlsBase) {
         auto *tlsBase = cast<DefinedGlobal>(WasmSym::tlsBase);
         setGlobalPtr(tlsBase, memoryPtr);
       }
     }
 
     memoryPtr += seg->size;
   }
 
std::string functionBody;
     writeUleb128(os, bodyContent.size(), "function size");
     os << bodyContent;
   }
   ArrayRef<uint8_t> body = arrayRefFromStringRef(saver().save(functionBody));
   cast<SyntheticFunction>(func->function)->setBody(body);
 }
 
 bool Writer::needsPassiveInitialization(const OutputSegment *segment) {
-  // TLS segments are initialized separately
-  if (segment->isTLS())
-    return false;
   // If bulk memory features is supported then we can perform bss initialization
   // (via memory.fill) during `__wasm_init_memory`.
   if (config->importMemory && !segment->requiredInBinary())
     return true;
   return segment->initFlags & WASM_DATA_SEGMENT_IS_PASSIVE;
 }
 
 bool Writer::hasPassiveInitializedSegments() {
void Writer::createSyntheticInitFunctions() {
   // dropped in __wasm_init_memory, which is registered as the start function
   // We also initialize bss segments (using memory.fill) as part of this
   // function.
   if (hasPassiveInitializedSegments()) {
     WasmSym::initMemory = symtab->addSyntheticFunction(
         "__wasm_init_memory", WASM_SYMBOL_VISIBILITY_HIDDEN,
         make<SyntheticFunction>(nullSignature, "__wasm_init_memory"));
     WasmSym::initMemory->markLive();
+    if (config->sharedMemory) {
+      // This global is assigned during  __wasm_init_memory in the shared memory
+      // case.
+      WasmSym::tlsBase->markLive();
+    }
   }
 
   if (config->sharedMemory && out.globalSec->needsTLSRelocations()) {
     WasmSym::applyGlobalTLSRelocs = symtab->addSyntheticFunction(
         "__wasm_apply_global_tls_relocs", WASM_SYMBOL_VISIBILITY_HIDDEN,
         make<SyntheticFunction>(nullSignature,
                                 "__wasm_apply_global_tls_relocs"));
     WasmSym::applyGlobalTLSRelocs->markLive();
auto writeGetFlagAddress = [&]() {
         writePtrConst(os, flagAddress, is64, "flag address");
       }
     };
 
     if (config->sharedMemory) {
       // With PIC code we cache the flag address in local 0
       if (config->isPic) {
         writeUleb128(os, 1, "num local decls");
-        writeUleb128(os, 1, "local count");
+        writeUleb128(os, 2, "local count");
         writeU8(os, is64 ? WASM_TYPE_I64 : WASM_TYPE_I32, "address type");
         writeU8(os, WASM_OPCODE_GLOBAL_GET, "GLOBAL_GET");
         writeUleb128(os, WasmSym::memoryBase->getGlobalIndex(), "memory_base");
         writePtrConst(os, flagAddress, is64, "flag address");
         writeU8(os, is64 ? WASM_OPCODE_I64_ADD : WASM_OPCODE_I32_ADD, "add");
         writeU8(os, WASM_OPCODE_LOCAL_SET, "local.set");
         writeUleb128(os, 0, "local 0");
       } else {
for (const OutputSegment *s : segments) {
         // For passive BSS segments we can simple issue a memory.fill(0).
         // For non-BSS segments we do a memory.init.  Both these
         // instructions take as thier first argument the destination
         // address.
         writePtrConst(os, s->startVA, is64, "destination address");
         if (config->isPic) {
           writeU8(os, WASM_OPCODE_GLOBAL_GET, "GLOBAL_GET");
           writeUleb128(os, WasmSym::memoryBase->getGlobalIndex(),
-                       "memory_base");
+                       "__memory_base");
           writeU8(os, is64 ? WASM_OPCODE_I64_ADD : WASM_OPCODE_I32_ADD,
                   "i32.add");
         }
+
+        // When we initialize the TLS segment we also set the `__tls_base`
+        // global.  This allows the runtime to use this static copy of the
+        // TLS data for the first/main thread.
+        if (config->sharedMemory && s->isTLS()) {
+          if (config->isPic) {
+            // Cache the result of the addionion in local 0
+            writeU8(os, WASM_OPCODE_LOCAL_TEE, "local.tee");
+            writeUleb128(os, 1, "local 1");
+          } else {
+            writePtrConst(os, s->startVA, is64, "destination address");
+          }
+          writeU8(os, WASM_OPCODE_GLOBAL_SET, "GLOBAL_SET");
+          writeUleb128(os, WasmSym::tlsBase->getGlobalIndex(),
+                       "__tls_base");
+          if (config->isPic) {
+            writeU8(os, WASM_OPCODE_LOCAL_GET, "local.tee");
+            writeUleb128(os, 1, "local 1");
+          }
+        }
+
         if (s->isBss) {
           writeI32Const(os, 0, "fill value");
           writeI32Const(os, s->size, "memory region size");
           writeU8(os, WASM_OPCODE_MISC_PREFIX, "bulk-memory prefix");
           writeUleb128(os, WASM_OPCODE_MEMORY_FILL, "memory.fill");
           writeU8(os, 0, "memory index immediate");
         } else {
           writeI32Const(os, 0, "source segment offset");
           writeI32Const(os, s->size, "memory region size");
           writeU8(os, WASM_OPCODE_MISC_PREFIX, "bulk-memory prefix");
           writeUleb128(os, WASM_OPCODE_MEMORY_INIT, "memory.init");
           writeUleb128(os, s->index, "segment index immediate");
           writeU8(os, 0, "memory index immediate");
         }
       }
     }
 
if (config->sharedMemory) {
       writeU8(os, WASM_OPCODE_DROP, "drop");
 
       // Unconditionally drop passive data segments
       writeU8(os, WASM_OPCODE_END, "end $drop");
     }
 
     for (const OutputSegment *s : segments) {
       if (needsPassiveInitialization(s) && !s->isBss) {
+        // The TLS region should not be dropped since its is needed
+        // during the intiailizing of each thread (__wasm_init_tls).
+        if (config->sharedMemory && s->isTLS())
+          continue;
         // data.drop instruction
         writeU8(os, WASM_OPCODE_MISC_PREFIX, "bulk-memory prefix");
         writeUleb128(os, WASM_OPCODE_DATA_DROP, "data.drop");
         writeUleb128(os, s->index, "segment index immediate");
       }
     }
 
     // End the function
std::string bodyContent;
 
     writeUleb128(os, 0, "num locals");
     if (tlsSeg) {
       writeU8(os, WASM_OPCODE_LOCAL_GET, "local.get");
       writeUleb128(os, 0, "local index");
 
       writeU8(os, WASM_OPCODE_GLOBAL_SET, "global.set");
       writeUleb128(os, WasmSym::tlsBase->getGlobalIndex(), "global index");
-      WasmSym::tlsBase->markLive();
 
       // FIXME(wvo): this local needs to be I64 in wasm64, or we need an extend op.
       writeU8(os, WASM_OPCODE_LOCAL_GET, "local.get");
       writeUleb128(os, 0, "local index");
 
       writeI32Const(os, 0, "segment offset");
 
       writeI32Const(os, tlsSeg->size, "memory region size");
if (!config->relocatable) {
     if (!config->relocatable && !config->isPic &&
         !WasmSym::callCtors->isUsedInRegularObj &&
         !WasmSym::callCtors->isExported()) {
       log("-- createCommandExportWrappers");
       createCommandExportWrappers();
     }
   }
 
-  if (WasmSym::initTLS && WasmSym::initTLS->isLive())
+  if (WasmSym::initTLS && WasmSym::initTLS->isLive()) {
+    log("-- createInitTLSFunction");
     createInitTLSFunction();
+  }
 
   if (errorCount())
     return;
 
   log("-- calculateTypes");
   calculateTypes();
   log("-- calculateExports");
   calculateExports();

llvm/include/llvm/BinaryFormat/Wasm.h

 };
 
 // Opcodes used in initializer expressions.
 enum : unsigned {
   WASM_OPCODE_END = 0x0b,
   WASM_OPCODE_CALL = 0x10,
   WASM_OPCODE_LOCAL_GET = 0x20,
   WASM_OPCODE_LOCAL_SET = 0x21,
+  WASM_OPCODE_LOCAL_TEE = 0x22,
   WASM_OPCODE_GLOBAL_GET = 0x23,
   WASM_OPCODE_GLOBAL_SET = 0x24,
   WASM_OPCODE_I32_STORE = 0x36,
   WASM_OPCODE_I64_STORE = 0x37,
   WASM_OPCODE_I32_CONST = 0x41,
   WASM_OPCODE_I64_CONST = 0x42,
   WASM_OPCODE_F32_CONST = 0x43,
   WASM_OPCODE_F64_CONST = 0x44,