[lld][WebAssembly] Add suppport for PIC + passive data initialization

[sbc100]

Summary

This change improves our support for shared memory to include
PIC executables (and shared libraries).

To handle this case the linker-generated __wasm_init_memory
function (that only exists in shared memory builds) must be
capable of loading memory segements at non-const offsets based
on the runtime value of __memory_base.

[sbc100]

Still need to add testing..

[dschuff]

lld/wasm/Writer.cpp

880

should this be conditionalized for wasm64?
alternatively, throw a nice error if we don't want to implement it right now.

[tlively]

Looks good so far!

lld/wasm/Writer.cpp

934

960–961

This repeated code chunk could be extracted into a small lambda to reduce repetition.

[sbc100]

• rebase + add test

[sbc100]

feedback

[tlively]

lld/wasm/Writer.cpp

940–947

I think this will make the logic clearer.

[sbc100]

• fix start

lld/test/lit.cfg.py

 
 # Configuration file for the 'lit' test runner.
 
 # name: The name of this test suite.
 config.name = 'lld'
 
 # testFormat: The test format to use to interpret tests.
 #
-# For now we require '&&' between commands, until they get globally killed and
-# the test runner updated.
+# For now we require '&&' between commands, until they get globally killed and the test runner updated.
 config.test_format = lit.formats.ShTest(not llvm_config.use_lit_shell)
 
 # suffixes: A list of file extensions to treat as test files.
 config.suffixes = ['.ll', '.s', '.test', '.yaml', '.objtxt']
 
 # excludes: A list of directories to exclude from the testsuite. The 'Inputs'
 # subdirectories contain auxiliary inputs for various tests in their parent
 # directories.

lld/test/wasm/data-segments.ll

 ; RUN: llc --mtriple=wasm32-unknown-unknown -filetype=obj %s -o %t.atomics.o -mattr=+atomics
 ; RUN: llc --mtriple=wasm32-unknown-unknown -filetype=obj %s -o %t.bulk-mem.o -mattr=+bulk-memory
 ; RUN: llc --mtriple=wasm32-unknown-unknown -filetype=obj %s -o %t.atomics.bulk-mem.o -mattr=+atomics,+bulk-memory
 ; RUN: llc --mtriple=wasm64-unknown-unknown -filetype=obj %s -o %t.atomics.bulk-mem64.o -mattr=+atomics,+bulk-memory
+; RUN: llc --mtriple=wasm32-unknown-unknown -filetype=obj %s -o %t.atomics.bulk-mem.pic.o -relocation-model=pic -mattr=+atomics,+bulk-memory,+mutable-globals
+; RUN: llc --mtriple=wasm64-unknown-unknown -filetype=obj %s -o %t.atomics.bulk-mem.pic-mem64.o -relocation-model=pic -mattr=+atomics,+bulk-memory,+mutable-globals
 
 ; atomics, shared memory => error
 ; RUN: not wasm-ld -no-gc-sections --no-entry --shared-memory --max-memory=131072 %t.atomics.o -o %t.atomics.wasm 2>&1 | FileCheck %s --check-prefix ERROR
 
 ; bulk memory, unshared memory => active segments
 ; RUN: wasm-ld -no-gc-sections --no-entry %t.bulk-mem.o -o %t.bulk-mem.wasm
 ; RUN: obj2yaml %t.bulk-mem.wasm | FileCheck %s --check-prefix ACTIVE
 
 ; atomics, bulk memory, shared memory => passive segments
 ; RUN: wasm-ld -no-gc-sections --no-entry --shared-memory --max-memory=131072 %t.atomics.bulk-mem.o -o %t.atomics.bulk-mem.wasm
 ; RUN: obj2yaml %t.atomics.bulk-mem.wasm | FileCheck %s --check-prefixes PASSIVE,PASSIVE32
 
 ; Also test with wasm64
 ; RUN: wasm-ld -mwasm64 -no-gc-sections --no-entry --shared-memory --max-memory=131072 %t.atomics.bulk-mem64.o -o %t.atomics.bulk-mem64.wasm
 ; RUN: obj2yaml %t.atomics.bulk-mem64.wasm | FileCheck %s --check-prefixes PASSIVE,PASSIVE64
 
 ; Also test in combination with PIC/pie
-; RUN: llc --mtriple=wasm32-unknown-unknown -filetype=obj -relocation-model=pic %s -o %t.atomics.bulk-mem.pic.o -mattr=+atomics,+bulk-memory,+mutable-globals
 ; RUN: wasm-ld --experimental-pic -pie -no-gc-sections --no-entry --shared-memory --max-memory=131072 %t.atomics.bulk-mem.pic.o -o %t.pic.wasm
+; RUN: obj2yaml %t.pic.wasm | FileCheck %s --check-prefixes PASSIVE-PIC,PASSIVE32-PIC
+
+; Also test in combination with PIC/pie + wasm64
+; RUN: wasm-ld -mwasm64 --experimental-pic -pie -no-gc-sections --no-entry --shared-memory --max-memory=131072 %t.atomics.bulk-mem.pic-mem64.o -o %t.pic-mem64.wasm
+; RUN: obj2yaml %t.pic-mem64.wasm | FileCheck %s --check-prefixes PASSIVE-PIC,PASSIVE64-PIC
 
 @a = hidden global [6 x i8] c"hello\00", align 1
 @b = hidden global [8 x i8] c"goodbye\00", align 1
 @c = hidden global [10000 x i8] zeroinitializer, align 1
 @d = hidden global i32 42, align 4
 
 @e = private constant [9 x i8] c"constant\00", align 1
 @f = private constant i8 43, align 4

 ; PASSIVE-NEXT:    Name:            name
 ; PASSIVE-NEXT:    FunctionNames:
 ; PASSIVE-NEXT:      - Index:           0
 ; PASSIVE-NEXT:        Name:            __wasm_call_ctors
 ; PASSIVE-NEXT:      - Index:           1
 ; PASSIVE-NEXT:        Name:            __wasm_init_memory
 ; PASSIVE-NEXT:      - Index:           2
 ; PASSIVE-NEXT:        Name:            __wasm_init_tls
+
+;      PASSIVE-PIC:  - Type:            START
+; PASSIVE-PIC-NEXT:    StartFunction:   2
+; PASSIVE-PIC-NEXT:  - Type:            DATACOUNT
+; PASSIVE-PIC-NEXT:    Count:           1
+; PASSIVE-PIC-NEXT:  - Type:            CODE
+; PASSIVE-PIC-NEXT:    Functions:
+; PASSIVE-PIC-NEXT:      - Index:           0
+; PASSIVE-PIC-NEXT:        Locals:          []
+; PASSIVE-PIC-NEXT:        Body:            10010B
+; PASSIVE-PIC-NEXT:      - Index:           1
+; PASSIVE-PIC-NEXT:        Locals:          []
+; PASSIVE-PIC-NEXT:        Body:            0B
+; PASSIVE-PIC-NEXT:      - Index:           2
+; PASSIVE-PIC-NEXT:        Locals:
+; PASSIVE32-PIC-NEXT:          - Type:            I32
+; PASSIVE64-PIC-NEXT:          - Type:            I64
+; PASSIVE-PIC-NEXT:            Count:           1
+; PASSIVE32-PIC-NEXT:        Body:            230141B4CE006A2100200041004101FE480200044020004101427FFE0102001A05410023016A410041B1CE00FC08000020004102FE1702002000417FFE0002001A0BFC09000B
+; PASSIVE64-PIC-NEXT:        Body:            230142B4CE006A2100200041004101FE480200044020004101427FFE0102001A05420023016A410041B1CE00FC08000020004102FE1702002000417FFE0002001A0BFC09000B
+; 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:   4
+; PASSIVE-PIC-NEXT:        InitFlags:       1
+
+;      PASSIVE-PIC:  - Type:            CUSTOM
+; PASSIVE-PIC-NEXT:    Name:            name
+; PASSIVE-PIC-NEXT:    FunctionNames:
+; PASSIVE-PIC-NEXT:      - Index:           0
+; PASSIVE-PIC-NEXT:        Name:            __wasm_call_ctors
+; PASSIVE-PIC-NEXT:      - Index:           1
+; PASSIVE-PIC-NEXT:        Name:            __wasm_apply_relocs
+; PASSIVE-PIC-NEXT:      - Index:           2
+; PASSIVE-PIC-NEXT:        Name:            __wasm_init_memory
+; PASSIVE-PIC-NEXT:      - Index:           3
+; PASSIVE-PIC-NEXT:        Name:            __wasm_init_tls

lld/wasm/Driver.cpp

if (config->isPic) {
     WasmSym::memoryBase->markLive();
     WasmSym::tableBase->markLive();
   } else {
     // For non-PIC code
     WasmSym::stackPointer = createGlobalVariable("__stack_pointer", true);
     WasmSym::stackPointer->markLive();
   }
 
-  if (config->sharedMemory && !config->shared) {
+  if (config->sharedMemory) {
     // Passive segments are used to avoid memory being reinitialized on each
     // thread's instantiation. These passive segments are initialized and
     // dropped in __wasm_init_memory, which is registered as the start function
     WasmSym::initMemory = symtab->addSyntheticFunction(
         "__wasm_init_memory", WASM_SYMBOL_VISIBILITY_HIDDEN,
         make<SyntheticFunction>(nullSignature, "__wasm_init_memory"));
     WasmSym::initMemoryFlag = symtab->addSyntheticDataSymbol(
         "__wasm_init_memory_flag", WASM_SYMBOL_VISIBILITY_HIDDEN);

lld/wasm/SyntheticSections.cpp

void ExportSection::writeBody() {
   raw_ostream &os = bodyOutputStream;
 
   writeUleb128(os, exports.size(), "export count");
   for (const WasmExport &export_ : exports)
     writeExport(os, export_);
 }
 
 bool StartSection::isNeeded() const {
-  return !config->relocatable && hasInitializedSegments && config->sharedMemory;
+  return WasmSym::initMemory != nullptr;
 }
 
 void StartSection::writeBody() {
   raw_ostream &os = bodyOutputStream;
   writeUleb128(os, WasmSym::initMemory->getFunctionIndex(), "function index");
 }
 
 void ElemSection::addEntry(FunctionSymbol *sym) {

lld/wasm/Writer.cpp

return std::find_if(segments.begin(), segments.end(),
                         return this->needsPassiveInitialization(s);
                       }) != segments.end();
 }
 
 void Writer::createInitMemoryFunction() {
   LLVM_DEBUG(dbgs() << "createInitMemoryFunction\n");
   assert(WasmSym::initMemoryFlag);
   uint64_t flagAddress = WasmSym::initMemoryFlag->getVirtualAddress();
+  bool is64 = config->is64.getValueOr(false);
   std::string bodyContent;
   {
     raw_string_ostream os(bodyContent);
+    // 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");
+      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, WASM_OPCODE_I32_ADD, "add");
+      writeU8(os, WASM_OPCODE_LOCAL_SET, "local.set");
+      writeUleb128(os, 0, "local 0");
+    } else {
       writeUleb128(os, 0, "num locals");
+    }
 
     if (hasPassiveInitializedSegments()) {
       // Initialize memory in a thread-safe manner. The thread that successfully
       // increments the flag from 0 to 1 is is responsible for performing the
       // memory initialization. Other threads go sleep on the flag until the
       // first thread finishing initializing memory, increments the flag to 2,
       // and wakes all the other threads. Once the flag has been set to 2,
       // subsequently started threads will skip the sleep. All threads
if (hasPassiveInitializedSegments()) {
       //      (i32.const $__init_memory_flag)
       //      (i32.const -1u)
       //     )
       //    )
       //   )
       //  )
       //  ( ... drop data segments ... )
       // )
+      //
+      // When we are building with PIC, calculate the flag location using:
+      //
+      //    (global.get $__memory_base)
+      //    (i32.const $__init_memory_flag)
+      //    (i32.const 1)
 
-      bool is64 = config->is64.getValueOr(false);
+      auto writeGetFlagAddress = [&]() {
+        if (config->isPic) {
+          writeU8(os, WASM_OPCODE_LOCAL_GET, "local.get");
+          writeUleb128(os, 0, "local 0");
+        } else {
+          writePtrConst(os, flagAddress, is64, "flag address");
+        }
+      };
 
       // Atomically check whether this is the main thread.
-      writePtrConst(os, flagAddress, is64, "flag address");
+      writeGetFlagAddress();
       writeI32Const(os, 0, "expected flag value");
       writeI32Const(os, 1, "flag value");
       writeU8(os, WASM_OPCODE_ATOMICS_PREFIX, "atomics prefix");
       writeUleb128(os, WASM_OPCODE_I32_RMW_CMPXCHG, "i32.atomic.rmw.cmpxchg");
       writeMemArg(os, 2, 0);
       writeU8(os, WASM_OPCODE_IF, "IF");
       writeU8(os, WASM_TYPE_NORESULT, "blocktype");
 
       // Did not increment 0, so wait for main thread to initialize memory
-      writePtrConst(os, flagAddress, is64, "flag address");
+      writeGetFlagAddress();
       writeI32Const(os, 1, "expected flag value");
       writeI64Const(os, -1, "timeout");
 
       writeU8(os, WASM_OPCODE_ATOMICS_PREFIX, "atomics prefix");
       writeUleb128(os, WASM_OPCODE_I32_ATOMIC_WAIT, "i32.atomic.wait");
       writeMemArg(os, 2, 0);
       writeU8(os, WASM_OPCODE_DROP, "drop");
 
       writeU8(os, WASM_OPCODE_ELSE, "ELSE");
 
       // Did increment 0, so conditionally initialize passive data segments
       for (const OutputSegment *s : segments) {
         if (needsPassiveInitialization(s)) {
           // 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");
+            writeU8(os, WASM_OPCODE_I32_ADD, "i32.add");
+          }
           // source segment offset
           writeI32Const(os, 0, "segment offset");
           // memory region size
           writeI32Const(os, s->size, "memory region size");
           // memory.init instruction
           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");
         }
       }
 
       // Set flag to 2 to mark end of initialization
-      writePtrConst(os, flagAddress, is64, "flag address");
+      writeGetFlagAddress();
       writeI32Const(os, 2, "flag value");
       writeU8(os, WASM_OPCODE_ATOMICS_PREFIX, "atomics prefix");
       writeUleb128(os, WASM_OPCODE_I32_ATOMIC_STORE, "i32.atomic.store");
       writeMemArg(os, 2, 0);
 
       // Notify any waiters that memory initialization is complete
-      writePtrConst(os, flagAddress, is64, "flag address");
+      writeGetFlagAddress();
       writeI32Const(os, -1, "number of waiters");
       writeU8(os, WASM_OPCODE_ATOMICS_PREFIX, "atomics prefix");
       writeUleb128(os, WASM_OPCODE_ATOMIC_NOTIFY, "atomic.notify");
       writeMemArg(os, 2, 0);
       writeU8(os, WASM_OPCODE_DROP, "drop");
 
       writeU8(os, WASM_OPCODE_END, "END");
 
std::string bodyContent;
     // End the function, returning the return values from the user's code.
     writeU8(os, WASM_OPCODE_END, "END");
   }
 
   createFunction(f, bodyContent);
 }
 
 void Writer::createInitTLSFunction() {
-  if (!WasmSym::initTLS->isLive())
-    return;
-
   std::string bodyContent;
   {
     raw_string_ostream os(bodyContent);
 
     OutputSegment *tlsSeg = nullptr;
     for (auto *seg : segments) {
       if (seg->name == ".tdata") {
         tlsSeg = seg;
if (!config->relocatable) {
     if (!config->relocatable && !config->isPic &&
         !WasmSym::callCtors->isUsedInRegularObj &&
         !WasmSym::callCtors->isExported()) {
       log("-- createCommandExportWrappers");
       createCommandExportWrappers();
     }
   }
 
-  if (!config->relocatable && config->sharedMemory && !config->shared)
+  if (WasmSym::initTLS && WasmSym::initTLS->isLive())
     createInitTLSFunction();
 
   if (errorCount())
     return;
 
   log("-- calculateTypes");
   calculateTypes();
   log("-- calculateExports");

llvm/include/llvm/BinaryFormat/Wasm.h

enum : unsigned {
   WASM_EXTERNAL_EVENT = 0x4,
 };
 
 // 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_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,