[lld][WebAssembly] Take advantage of extended const expression when available

[sbc100]

Summary

In particular, when extended-const is available and we building PIC
code we no longer need to combine output segments into a single segment
that can be initialized at __memory_base. Instead each segment
can encode its offset from __memory_base in its initializer. e.g.

(i32.add (global.get __memory_base) (i32.const offset)

[sbc100]

• cleanup

[dschuff]

Would there still be an advantage to combining the segments? e.g. would the library load faster if there were only one?

[sbc100]

In D121420#3379498, @dschuff wrote:

Would there still be an advantage to combining the segments? e.g. would the library load faster if there were only one?

We already combine segments of the same type/prefix unless --no-merge-data-segments is passed... so in the normal case there are just a fixed number or possible output segments: .data .bss .tdata .tbss, plus any user-defined segments. This later-phase combining was always just a stop gap to handler the lack of extended const expressions. See the description of https://reviews.llvm.org/D96453 for why we delay the merging in this case.

[sbc100]

In D121420#3379513, @sbc100 wrote:

In D121420#3379498, @dschuff wrote:

Would there still be an advantage to combining the segments? e.g. would the library load faster if there were only one?

We already combine segments of the same type/prefix unless --no-merge-data-segments is passed... so in the normal case there are just a fixed number or possible output segments: .data .bss .tdata .tbss, plus any user-defined segments. This later-phase combining was always just a stop gap to handler the lack of extended const expressions. See the description of https://reviews.llvm.org/D96453 for why we delay the merging in this case.

Even if we decided that this second level combining was a good idea more generally that isn't really related to this PR and could be an option for all build (unrelated to PIC or sharedMemory, or extendedConst). But i think keeping the high level sections separate by default probably makes sense, given how few there are in most programs. Also, binaryen can always combine them if folks want to micro-optimize.

[sbc100]

• rebase

[dschuff]

Ah, right I forgot we already combined the default segments. That's probably fine as a default behavior.

lld/test/wasm/pie.ll

96

150

I don't see a RUN line for DISASSEM-EXTENDED-CONST?

lld/test/wasm/tls-non-shared-memory.s

144

lld/wasm/SyntheticSections.h

277

Does this comment need updating?

lld/wasm/Writer.cpp

930–932

[sbc100]

• feedback

[sbc100]

lld/wasm/SyntheticSections.h

277

I think the comment is still accurate. What has changed is that these globals might not need relocation.. but this comment doesn't touch on the relocation part.

[sbc100]

• rebase

[dschuff]

lld/test/wasm/pie.ll

119

maybe write what instructions these encoding correspond to?

lld/wasm/Writer.cpp

930

1567

1569

[sbc100]

• feedback

lld/test/wasm/pie.ll

 ; CHECK-NEXT:         Name:            __wasm_start
 
 ; DISASSEM:       <__wasm_start>:
 ; DISASSEM-EMPTY:
 ; DISASSEM-NEXT:   call 3
 ; DISASSEM-NEXT:   call 2
 ; DISASSEM-NEXT:   end
 
+; Run the same test with extended-const support.  When this is available
+; we don't need __wasm_apply_global_relocs and instead rely on the add
+; instruction in the InitExpr.  We also, therefore, do not need these globals
+; to be mutable.
+
+; RUN: llc -relocation-model=pic -mattr=+extended-const,+mutable-globals,+atomics,+bulk-memory -filetype=obj %s -o %t.extended.o
+; RUN: wasm-ld --no-gc-sections --allow-undefined --experimental-pic -pie -o %t.extended.wasm %t.extended.o
+; RUN: obj2yaml %t.extended.wasm | FileCheck %s --check-prefix=EXTENDED-CONST
+
+; EXTENDED-CONST-NOT: __wasm_apply_global_relocs
+
+; EXTENDED-CONST:       - Type:            GLOBAL
+; EXTENDED-CONST-NEXT:    Globals:
+; EXTENDED-CONST-NEXT:      - Index:           4
+; EXTENDED-CONST-NEXT:        Type:            I32
+; EXTENDED-CONST-NEXT:        Mutable:         false
+; EXTENDED-CONST-NEXT:        InitExpr:
+; EXTENDED-CONST-NEXT:          Opcode:        GLOBAL_GET
+; EXTENDED-CONST-NEXT:          Index:         1
+; EXTENDED-CONST-NEXT:      - Index:           5
+; EXTENDED-CONST-NEXT:        Type:            I32
+; EXTENDED-CONST-NEXT:        Mutable:         false
+; EXTENDED-CONST-NEXT:        InitExpr:
+; EXTENDED-CONST-NEXT:          Extended:        true
+; EXTENDED-CONST-NEXT:          Body:            230141046A0B
+; EXTENDED-CONST-NEXT:      - Index:           6
+; EXTENDED-CONST-NEXT:        Type:            I32
+; EXTENDED-CONST-NEXT:        Mutable:         false
+; EXTENDED-CONST-NEXT:        InitExpr:
+; EXTENDED-CONST-NEXT:          Extended:        true
+; EXTENDED-CONST-NEXT:          Body:            2301410C6A0B
+
+;      EXTENDED-CONST:  - Type:            START
+; EXTENDED-CONST-NEXT:    StartFunction:   2
+
+;      EXTENDED-CONST:    FunctionNames:
+; EXTENDED-CONST-NEXT:      - Index:           0
+; EXTENDED-CONST-NEXT:        Name:            external_func
+; EXTENDED-CONST-NEXT:      - Index:           1
+; EXTENDED-CONST-NEXT:        Name:            __wasm_call_ctors
+; EXTENDED-CONST-NEXT:      - Index:           2
+; EXTENDED-CONST-NEXT:        Name:            __wasm_apply_data_relocs
+
 ; Run the same test with threading support.  In this mode
 ; we expect __wasm_init_memory and __wasm_apply_data_relocs
 ; to be generated along with __wasm_start as the start
 ; function.
 
 ; RUN: llc -relocation-model=pic -mattr=+mutable-globals,+atomics,+bulk-memory -filetype=obj %s -o %t.shmem.o
 ; RUN: wasm-ld --no-gc-sections --shared-memory --allow-undefined --experimental-pic -pie -o %t.shmem.wasm %t.shmem.o
 ; RUN: obj2yaml %t.shmem.wasm | FileCheck %s --check-prefix=SHMEM
-; RUN: llvm-objdump --disassemble-symbols=__wasm_start --no-show-raw-insn --no-leading-addr %t.shmem.wasm | FileCheck %s --check-prefixes DISASSEM-SHMEM
+; RUN: llvm-objdump --disassemble-symbols=__wasm_start --no-show-raw-insn --no-leading-addr %t.shmem.wasm | FileCheck %s --check-prefix DISASSEM-SHMEM
 
 ; SHMEM:         - Type:            START
 ; SHMEM-NEXT:      StartFunction:   6
 
 ; DISASSEM-SHMEM:       <__wasm_start>:
 ; DISASSEM-SHMEM-EMPTY:
 ; DISASSEM-SHMEM-NEXT:   call 5
 ; DISASSEM-SHMEM-NEXT:   call 3
 ; DISASSEM-SHMEM-NEXT:   end
 
 ; SHMEM:         FunctionNames:
 ; SHMEM-NEXT:      - Index:           0

 ; SHMEM-NEXT:      - Index:           6
 ; SHMEM-NEXT:        Name:            __wasm_start
 ; SHMEM-NEXT:      - Index:           7
 ; SHMEM-NEXT:        Name:            foo
 ; SHMEM-NEXT:      - Index:           8
 ; SHMEM-NEXT:        Name:            get_data_address
 ; SHMEM-NEXT:      - Index:           9
 ; SHMEM-NEXT:        Name:            _start
-

lld/test/wasm/tls-non-shared-memory.s

 # RUN: obj2yaml %t.wasm | FileCheck %s
 
 # RUN: wasm-ld --experimental-pic -shared -o %t.so %t.o
 # RUN: obj2yaml %t.so | FileCheck %s --check-prefixes=SHARED,PIC
 
 # RUN: wasm-ld --experimental-pic --no-gc-sections --no-entry -pie -o %t-pie.wasm %t.o
 # RUN: obj2yaml %t-pie.wasm | FileCheck %s --check-prefixes=PIE,PIC
 
+# RUN: wasm-ld --experimental-pic --features=atomics,bulk-memory,extended-const --no-gc-sections --no-entry -pie -o %t-extended-const.wasm %t.o
+# RUN: obj2yaml %t-extended-const.wasm | FileCheck %s --check-prefixes=EXT-CONST
+
 #      CHECK:   - Type:            GLOBAL
 # __stack_pointer
 # CHECK-NEXT:     Globals:
 # CHECK-NEXT:       - Index:           0
 # CHECK-NEXT:         Type:            I32
 # CHECK-NEXT:         Mutable:         true
 # CHECK-NEXT:         InitExpr:
 # CHECK-NEXT:           Opcode:          I32_CONST

 # PIC-NEXT:    Segments:
 # PIC-NEXT:      - SectionOffset:   6
 # PIC-NEXT:        InitFlags:       0
 # PIC-NEXT:        Offset:
 # PIC-NEXT:          Opcode:          GLOBAL_GET
 # PIC-NEXT:          Index:           {{\d*}}
 # PIC-NEXT:        Content:         2B0000002A000000
 # PIC-NEXT:  - Type:            CUSTOM
+
+# Unless we have extended-const, in which case the merging is not needed.
+# The first segment is placed directly at `__memory_base` and the second
+# one is offset from `__memory_base` using `i32.add` and a constant.
+
+#      EXT-CONST:  - Type:            DATA
+# EXT-CONST-NEXT:    Segments:
+# EXT-CONST-NEXT:      - SectionOffset:   6
+# EXT-CONST-NEXT:        InitFlags:       0
+# EXT-CONST-NEXT:        Offset:
+# EXT-CONST-NEXT:          Opcode:          GLOBAL_GET
+# EXT-CONST-NEXT:          Index:           1
+# EXT-CONST-NEXT:        Content:         2B000000
+# EXT-CONST-NEXT:      - SectionOffset:   18
+# EXT-CONST-NEXT:        InitFlags:       0
+# EXT-CONST-NEXT:        Offset:
+# EXT-CONST-NEXT:          Extended:        true
+# This instruction sequence decodes to:
+# (global.get[0x23] 0x1 i32.const[0x41] 0x04 i32.add[0x6A] end[0x0b])
+# EXT-CONST-NEXT:          Body:            230141046A0B
+# EXT-CONST-NEXT:        Content:         2A000000

lld/wasm/Config.h

struct Configuration {
   bool compressRelocations;
   bool demangle;
   bool disableVerify;
   bool experimentalPic;
   bool emitRelocs;
   bool exportAll;
   bool exportDynamic;
   bool exportTable;
+  bool extendedConst;
   bool growableTable;
   bool gcSections;
   bool importMemory;
   bool sharedMemory;
   bool importTable;
   bool importUndefined;
   llvm::Optional<bool> is64;
   bool mergeDataSegments;

lld/wasm/OutputSections.cpp

unsigned segmentCount = std::count_if(
       [](OutputSegment *segment) { return segment->requiredInBinary(); });
 #ifndef NDEBUG
   unsigned activeCount = std::count_if(
       segments.begin(), segments.end(), [](OutputSegment *segment) {
         return (segment->initFlags & WASM_DATA_SEGMENT_IS_PASSIVE) == 0;
       });
 #endif
 
-  assert((config->sharedMemory || !config->isPic || activeCount <= 1) &&
+  assert((config->sharedMemory || !config->isPic || config->extendedConst ||
+          activeCount <= 1) &&
          "output segments should have been combined by now");
 
   writeUleb128(os, segmentCount, "data segment count");
   os.flush();
   bodySize = dataSectionHeader.size();
+  bool is64 = config->is64.getValueOr(false);
 
   for (OutputSegment *segment : segments) {
     if (!segment->requiredInBinary())
       continue;
     raw_string_ostream os(segment->header);
     writeUleb128(os, segment->initFlags, "init flags");
     if (segment->initFlags & WASM_DATA_SEGMENT_HAS_MEMINDEX)
       writeUleb128(os, 0, "memory index");
     if ((segment->initFlags & WASM_DATA_SEGMENT_IS_PASSIVE) == 0) {
+      if (config->isPic && config->extendedConst) {
+        writeU8(os, WASM_OPCODE_GLOBAL_GET, "global get");
+        writeUleb128(os, WasmSym::memoryBase->getGlobalIndex(),
+                     "literal (global index)");
+        if (segment->startVA) {
+          writePtrConst(os, segment->startVA, is64, "offset");
+          writeU8(os, is64 ? WASM_OPCODE_I64_ADD : WASM_OPCODE_I32_ADD, "add");
+        }
+        writeU8(os, WASM_OPCODE_END, "opcode:end");
+      } else {
         WasmInitExpr initExpr;
         initExpr.Extended = false;
         if (config->isPic) {
+          assert(segment->startVA == 0);
           initExpr.Inst.Opcode = WASM_OPCODE_GLOBAL_GET;
           initExpr.Inst.Value.Global = WasmSym::memoryBase->getGlobalIndex();
         } else {
-        initExpr = intConst(segment->startVA, config->is64.getValueOr(false));
+          initExpr = intConst(segment->startVA, is64);
         }
         writeInitExpr(os, initExpr);
       }
+    }
     writeUleb128(os, segment->size, "segment size");
     os.flush();
 
     segment->sectionOffset = bodySize;
     bodySize += segment->header.size() + segment->size;
     log("Data segment: size=" + Twine(segment->size) + ", startVA=" +
         Twine::utohexstr(segment->startVA) + ", name=" + segment->name);
 

lld/wasm/SyntheticSections.h

uint32_t numGlobals() const {
     return inputGlobals.size() + dataAddressGlobals.size() +
            internalGotSymbols.size();
   }
   bool isNeeded() const override { return numGlobals() > 0; }
   void assignIndexes() override;
   void writeBody() override;
   void addGlobal(InputGlobal *global);
 
   // Add an internal GOT entry global that corresponds to the given symbol.
   // Normally GOT entries are imported and assigned by the external dynamic
   // linker.  However, when linking PIC code statically or when linking with
   // -Bsymbolic we can internalize GOT entries by declaring globals the hold
   // symbol addresses.
   //
   // For the static linking case these internal globals can be completely
   // eliminated by a post-link optimizer such as wasm-opt.
   //
   // TODO(sbc): Another approach to optimizing these away could be to use
   // specific relocation types combined with linker relaxation which could
   // transform a `global.get` to an `i32.const`.
   void addInternalGOTEntry(Symbol *sym);
   bool needsRelocations() {
+    if (config->extendedConst)
+      return false;
     return llvm::find_if(internalGotSymbols, [=](Symbol *sym) {
              return !sym->isTLS();
            }) != internalGotSymbols.end();
   }
   bool needsTLSRelocations() {
     return llvm::find_if(internalGotSymbols, [=](Symbol *sym) {
              return sym->isTLS();
            }) != internalGotSymbols.end();

lld/wasm/SyntheticSections.cpp

void GlobalSection::addInternalGOTEntry(Symbol *sym) {
   if (auto *F = dyn_cast<FunctionSymbol>(sym)) {
     ensureIndirectFunctionTable();
     out.elemSec->addEntry(F);
   }
   internalGotSymbols.push_back(sym);
 }
 
 void GlobalSection::generateRelocationCode(raw_ostream &os, bool TLS) const {
+  assert(!config->extendedConst);
   bool is64 = config->is64.getValueOr(false);
   unsigned opcode_ptr_const = is64 ? WASM_OPCODE_I64_CONST
                                    : WASM_OPCODE_I32_CONST;
   unsigned opcode_ptr_add = is64 ? WASM_OPCODE_I64_ADD
                                  : WASM_OPCODE_I32_ADD;
 
   for (const Symbol *sym : internalGotSymbols) {
     if (TLS != sym->isTLS())
for (InputGlobal *g : inputGlobals) {
     writeGlobalType(os, g->getType());
     writeInitExpr(os, g->getInitExpr());
   }
   bool is64 = config->is64.getValueOr(false);
   uint8_t itype = is64 ? WASM_TYPE_I64 : WASM_TYPE_I32;
   for (const Symbol *sym : internalGotSymbols) {
     bool mutable_ = false;
     if (!sym->isStub) {
-      // In the case of dynamic linking, these global must to be mutable since
-      // they get updated to the correct runtime value during
-      // `__wasm_apply_global_relocs`.
-      if (config->isPic && !sym->isTLS())
+      // In the case of dynamic linking, unless we have 'extended-const'
+      // available, these global must to be mutable since they get updated to
+      // the correct runtime value during `__wasm_apply_global_relocs`.
+      if (!config->extendedConst && config->isPic && !sym->isTLS())
         mutable_ = true;
       // With multi-theadeding any TLS globals must be mutable since they get
       // set during `__wasm_apply_global_tls_relocs`
       if (config->sharedMemory && sym->isTLS())
         mutable_ = true;
     }
     WasmGlobalType type{itype, mutable_};
+    writeGlobalType(os, type);
+
+    if (config->extendedConst && config->isPic && !sym->isTLS() &&
+        isa<DefinedData>(sym)) {
+      // We can use an extended init expression to add a constant
+      // offset of __memory_base.
+      auto *d = cast<DefinedData>(sym);
+      writeU8(os, WASM_OPCODE_GLOBAL_GET, "global get");
+      writeUleb128(os, WasmSym::memoryBase->getGlobalIndex(),
+                   "literal (global index)");
+      if (d->getVA()) {
+        writePtrConst(os, d->getVA(), is64, "offset");
+        writeU8(os, is64 ? WASM_OPCODE_I64_ADD : WASM_OPCODE_I32_ADD, "add");
+      }
+      writeU8(os, WASM_OPCODE_END, "opcode:end");
+    } else {
       WasmInitExpr initExpr;
       if (auto *d = dyn_cast<DefinedData>(sym))
         initExpr = intConst(d->getVA(), is64);
       else if (auto *f = dyn_cast<FunctionSymbol>(sym))
         initExpr = intConst(f->isStub ? 0 : f->getTableIndex(), is64);
       else {
         assert(isa<UndefinedData>(sym));
         initExpr = intConst(0, is64);
       }
-    writeGlobalType(os, type);
       writeInitExpr(os, initExpr);
     }
+  }
   for (const DefinedData *sym : dataAddressGlobals) {
     WasmGlobalType type{itype, false};
     writeGlobalType(os, type);
     writeInitExpr(os, intConst(sym->getVA(), is64));
   }
 }
 
 void GlobalSection::addGlobal(InputGlobal *global) {

lld/wasm/Writer.cpp

void Writer::populateTargetFeatures() {
 
   // Only infer used features if user did not specify features
   bool inferFeatures = !config->features.hasValue();
 
   if (!inferFeatures) {
     auto &explicitFeatures = config->features.getValue();
     allowed.insert(explicitFeatures.begin(), explicitFeatures.end());
     if (!config->checkFeatures)
-      return;
+      goto done;
   }
 
   // Find the sets of used, required, and disallowed features
   for (ObjFile *file : symtab->objectFiles) {
     StringRef fileName(file->getName());
     for (auto &feature : file->getWasmObj()->getTargetFeatures()) {
       switch (feature.Prefix) {
       case WASM_FEATURE_PREFIX_USED:
for (ObjFile *file : symtab->objectFiles) {
     tlsUsed = tlsUsed || llvm::any_of(file->segments, isTLS);
   }
 
   if (inferFeatures)
     for (const auto &key : used.keys())
       allowed.insert(std::string(key));
 
   if (!config->checkFeatures)
-    return;
+    goto done;
 
   if (config->sharedMemory) {
     if (disallowed.count("shared-mem"))
       error("--shared-memory is disallowed by " + disallowed["shared-mem"] +
             " because it was not compiled with 'atomics' or 'bulk-memory' "
             "features.");
 
     for (auto feature : {"atomics", "bulk-memory"})
for (ObjFile *file : symtab->objectFiles) {
     for (const auto &feature : required.keys()) {
       if (!objectFeatures.count(std::string(feature)))
         error(Twine("Missing target feature '") + feature + "' in " + fileName +
               ", required by " + required[feature] +
               ". Use --no-check-features to suppress.");
     }
   }
 
+done:
   // Normally we don't include bss segments in the binary.  In particular if
   // memory is not being imported then we can assume its zero initialized.
   // In the case the memory is imported, we and we can use the memory.fill
   // instrction than we can also avoid inluding the segments.
   if (config->importMemory && !allowed.count("bulk-memory"))
     config->emitBssSegments = true;
+
+  if (allowed.count("extended-const"))
+    config->extendedConst = true;
+
+  for (auto &feature : allowed)
+    log("Allowed feature: " + feature);
 }
 
 void Writer::checkImportExportTargetFeatures() {
   if (config->relocatable || !config->checkFeatures)
     return;
 
   if (out.targetFeaturesSec->features.count("mutable-globals") == 0) {
     for (const Symbol *sym : out.importSec->importedSymbols) {
void Writer::createOutputSegments() {
   for (OutputSegment *seg : segments)
     seg->finalizeInputSegments();
 }
 
 void Writer::combineOutputSegments() {
   // With PIC code we currently only support a single active data segment since
   // we only have a single __memory_base to use as our base address.  This pass
   // combines all data segments into a single .data segment.
-  // This restructions can be relaxed once we have extended constant
-  // expressions available:
-  // https://github.com/WebAssembly/extended-const
+  // This restruction does not apply when the extended const extension is
+  // available: https://github.com/WebAssembly/extended-const
+  assert(!config->extendedConst);
   assert(config->isPic && !config->sharedMemory);
   if (segments.size() <= 1)
     return;
   OutputSegment *combined = make<OutputSegment>(".data");
   combined->startVA = segments[0]->startVA;
   for (OutputSegment *s : segments) {
     bool first = true;
     for (InputChunk *inSeg : s->inputSegments) {
for (auto &name : config->requiredExports) {
     if (!sym || !sym->isDefined()) {
       if (config->unresolvedSymbols == UnresolvedPolicy::ReportError)
         error(Twine("symbol exported via --export not found: ") + name);
       if (config->unresolvedSymbols == UnresolvedPolicy::Warn)
         warn(Twine("symbol exported via --export not found: ") + name);
     }
   }
 
-  if (config->isPic && !config->sharedMemory) {
+  log("-- populateTargetFeatures");
+  populateTargetFeatures();
+
+  // When outputing PIC code each segment lives at at fixes offset from the
+  // `__memory_base` import.  Unless we support the extended const epression we
+  // can't do addition inside the constant expression, so we much combine the
+  // segment into a single one that can live at `__memory_base`.
+  if (config->isPic && !config->extendedConst && !config->sharedMemory) {
     // In shared memory mode all data segments are passive and initialized
     // via __wasm_init_memory.
     log("-- combineOutputSegments");
     combineOutputSegments();
   }
 
   log("-- createSyntheticSectionsPostLayout");
   createSyntheticSectionsPostLayout();
   log("-- populateProducers");
   populateProducers();
   log("-- calculateImports");
   calculateImports();
   log("-- scanRelocations");
   scanRelocations();
   log("-- finalizeIndirectFunctionTable");
   finalizeIndirectFunctionTable();
-  log("-- populateTargetFeatures");
-  populateTargetFeatures();
   log("-- createSyntheticInitFunctions");
   createSyntheticInitFunctions();
   log("-- assignIndexes");
   assignIndexes();
   log("-- calculateInitFunctions");
   calculateInitFunctions();
 
   if (!config->relocatable) {