[WebAssembly] Optionally compress relocation sites in the code section

[sbc100]

Summary

This change adds the ability for lld to remove LEB padding from
code section. This effectively shrinks the size of the resulting binary
in proportion to the number of code relocations.

Since there will be a performance cost this is currently only active
for -O2. Some toolchains may instead want to perform this compression
as a post linker step (for example running a binary through binaryen will
automatically compress these values).

I imagine we might want to make this the default in the future and
force people to opt out with -O0.

[sbc100]

• cleanup

[sbc100]

• cleanup

[sbc100]

• cleanup

[sbc100]

I will need to update this patch to handle FUNCTION_OFFSET relocations once we land https://reviews.llvm.org/D45118

[sbc100]

Add comment

[sbc100]

• remove some asserts

[sbc100]

rebase

[sbc100]

• no need to new flag

[sbc100]

• add comment

[sbc100]

• revert

[ruiu]

Overall looking good, but I think I have a concern about the lack of comments. Could you expand comments to make the .cpp file more self-explanatory so that a first-time reader who's not very familiar with wasm can understand the code by reading it?

[ncw]

This looks like a nice size optimisation, great.

[sbc100]

• add comments

[sbc100]

OK to land now @ruiu ?

[sbc100]

Planning on landing this today unless there are any more comments.

[ruiu]

I'm sorry, I missed your ping.

wasm/Driver.cpp

294

I think 0 should be default. For example, if you don't pass -O to gcc or clang, they emit a non-optimized code.

wasm/InputChunks.cpp

149

When a buffer is given, it is usually the first argument. It is better to follow the convention for readability.

197

Can you also mention that this function only computes the length of the compressed section and doesn't actually write anything?

[sbc100]

• feedback

[ruiu]

LGTM

test/wasm/compress-relocs.ll

+; RUN: llc -filetype=obj %p/Inputs/call-indirect.ll -o %t2.o
+; RUN: llc -filetype=obj %s -o %t.o
+; RUN: wasm-ld --check-signatures -o %t.wasm %t2.o %t.o
+; RUN: obj2yaml %t.wasm | FileCheck %s
+
+; RUN: wasm-ld --check-signatures --compress-relocs -o %t-compressed.wasm %t2.o %t.o
+; RUN: obj2yaml %t-compressed.wasm | FileCheck %s -check-prefix=COMPRESS
+
+target triple = "wasm32-unknown-unknown-wasm"
+
+define i32 @foo() {
+entry:
+  ret i32 2
+}
+
+define void @_start() local_unnamed_addr {
+entry:
+  ret void
+}
+
+; CHECK:    Body:            4100280284888080002100410028028088808000118080808000001A2000118180808000001A0B
+; COMPRESS: Body:            41002802840821004100280280081100001A20001101001A0B

wasm/Config.h

 #include "llvm/BinaryFormat/Wasm.h"
 
 namespace lld {
 namespace wasm {
 
 struct Configuration {
   bool AllowUndefined;
   bool CheckSignatures;
+  bool CompressRelocTargets;
   bool Demangle;
   bool ExportTable;
   bool GcSections;
   bool ImportMemory;
   bool ImportTable;
   bool PrintGcSections;
   bool Relocatable;
   bool StripAll;

wasm/Driver.cpp

void LinkerDriver::link(ArrayRef<const char *> ArgsArr) {
   V.push_back("wasm-ld (LLVM option parsing)");
   for (auto *Arg : Args.filtered(OPT_mllvm))
     V.push_back(Arg->getValue());
   cl::ParseCommandLineOptions(V.size(), V.data());
 
   errorHandler().ErrorLimit = args::getInteger(Args, OPT_error_limit, 20);
 
   Config->AllowUndefined = Args.hasArg(OPT_allow_undefined);
+  Config->CompressRelocTargets =
+      Args.hasFlag(OPT_compress_relocs, OPT_no_compress_relocs, false);
   Config->CheckSignatures =
       Args.hasFlag(OPT_check_signatures, OPT_no_check_signatures, false);
   Config->Demangle = Args.hasFlag(OPT_demangle, OPT_no_demangle, true);
   Config->Entry = getEntry(Args, Args.hasArg(OPT_relocatable) ? "" : "_start");
   Config->ExportTable = Args.hasArg(OPT_export_table);
   Config->ImportMemory = Args.hasArg(OPT_import_memory);
   Config->ImportTable = Args.hasArg(OPT_import_table);
   Config->OutputFile = Args.getLastArgValue(OPT_o);
   Config->Relocatable = Args.hasArg(OPT_relocatable);
   Config->GcSections =
       Args.hasFlag(OPT_gc_sections, OPT_no_gc_sections, !Config->Relocatable);
   Config->PrintGcSections =
       Args.hasFlag(OPT_print_gc_sections, OPT_no_print_gc_sections, false);
   Config->SearchPaths = args::getStrings(Args, OPT_L);
   Config->StripAll = Args.hasArg(OPT_strip_all);
   Config->StripDebug = Args.hasArg(OPT_strip_debug);
   Config->StackFirst = Args.hasArg(OPT_stack_first);
   errorHandler().Verbose = Args.hasArg(OPT_verbose);
   ThreadsEnabled = Args.hasFlag(OPT_threads, OPT_no_threads, true);
 
   Config->InitialMemory = args::getInteger(Args, OPT_initial_memory, 0);
   Config->GlobalBase = args::getInteger(Args, OPT_global_base, 1024);
   Config->MaxMemory = args::getInteger(Args, OPT_max_memory, 0);
   Config->ZStackSize =
       args::getZOptionValue(Args, OPT_z, "stack-size", WasmPageSize);
 
+  // We can't compress relocation targets when writing a relocatable file.
+  if (Config->Relocatable)
+    Config->CompressRelocTargets = false;
+
   if (auto *Arg = Args.getLastArg(OPT_allow_undefined_file))
     readImportFile(Arg->getValue());
 
   if (!Args.hasArg(OPT_INPUT)) {
     error("no input files");
     return;
   }
 

wasm/InputChunks.h

 class OutputSegment;
 
 class InputChunk {
 public:
   enum Kind { DataSegment, Function, SyntheticFunction, Section };
 
   Kind kind() const { return SectionKind; }
 
-  uint32_t getSize() const { return data().size(); }
+  virtual uint32_t getSize() const { return data().size(); }
 
   void copyRelocations(const WasmSection &Section);
 
-  void writeTo(uint8_t *SectionStart) const;
+  virtual void writeTo(uint8_t *SectionStart) const;
 
   ArrayRef<WasmRelocation> getRelocations() const { return Relocations; }
 
   virtual StringRef getName() const = 0;
   virtual StringRef getDebugName() const = 0;
   virtual uint32_t getComdat() const = 0;
   StringRef getComdatName() const;
 
public:
   unsigned Live : 1;
 
 protected:
   InputChunk(ObjFile *F, Kind K)
       : File(F), Live(!Config->GcSections), SectionKind(K) {}
   virtual ~InputChunk() = default;
   virtual ArrayRef<uint8_t> data() const = 0;
   virtual uint32_t getInputSectionOffset() const = 0;
+  virtual uint32_t getInputSize() const { return getSize(); };
 
   std::vector<WasmRelocation> Relocations;
   Kind SectionKind;
 };
 
 // Represents a WebAssembly data segment which can be included as part of
 // an output data segments.  Note that in WebAssembly, unlike ELF and other
 // formats, used the term "data segment" to refer to the continous regions of
public:
   InputFunction(const WasmSignature &S, const WasmFunction *Func, ObjFile *F)
       : InputChunk(F, InputChunk::Function), Signature(S), Function(Func) {}
 
   static bool classof(const InputChunk *C) {
     return C->kind() == InputChunk::Function ||
            C->kind() == InputChunk::SyntheticFunction;
   }
 
+  void writeTo(uint8_t *SectionStart) const override;
   StringRef getName() const override { return Function->SymbolName; }
   StringRef getDebugName() const override { return Function->DebugName; }
   uint32_t getComdat() const override { return Function->Comdat; }
+  uint32_t getSize() const override {
+    return Config->CompressRelocTargets && File ? CompressedSize : data().size();
+  }
   uint32_t getFunctionIndex() const { return FunctionIndex.getValue(); }
   bool hasFunctionIndex() const { return FunctionIndex.hasValue(); }
   void setFunctionIndex(uint32_t Index);
   uint32_t getTableIndex() const { return TableIndex.getValue(); }
   bool hasTableIndex() const { return TableIndex.hasValue(); }
   void setTableIndex(uint32_t Index);
+  void calculateSize();
 
   const WasmSignature &Signature;
 
 protected:
+  void compressRelocations();
+
   ArrayRef<uint8_t> data() const override {
+    assert(!Config->CompressRelocTargets);
     return File->CodeSection->Content.slice(getInputSectionOffset(),
                                             Function->Size);
   }
+
+  uint32_t getInputSize() const override { return Function->Size; }
+
   uint32_t getInputSectionOffset() const override {
     return Function->CodeSectionOffset;
   }
 
   const WasmFunction *Function;
   llvm::Optional<uint32_t> FunctionIndex;
   llvm::Optional<uint32_t> TableIndex;
+
+  std::vector<uint32_t> RelocationDeltas;
+  unsigned CompressedFuncSize = 0;
+  unsigned CompressedSize = 0;
 };
 
 class SyntheticFunction : public InputFunction {
 public:
   SyntheticFunction(const WasmSignature &S, StringRef Name,
                     StringRef DebugName = {})
       : InputFunction(S, nullptr, nullptr), Name(Name), DebugName(DebugName) {
     SectionKind = InputChunk::SyntheticFunction;

wasm/InputChunks.cpp

if (Index == UINT32_MAX)
     return StringRef();
   return File->getWasmObj()->linkingData().Comdats[Index];
 }
 
 void InputChunk::copyRelocations(const WasmSection &Section) {
   if (Section.Relocations.empty())
     return;
   size_t Start = getInputSectionOffset();
-  size_t Size = getSize();
+  size_t Size = getInputSize();
   for (const WasmRelocation &R : Section.Relocations)
     if (R.Offset >= Start && R.Offset < Start + Size)
       Relocations.push_back(R);
 }
 
 // Copy this input chunk to an mmap'ed output file and apply relocations.
 void InputChunk::writeTo(uint8_t *Buf) const {
   // Copy contents

 }
 
 void InputFunction::setTableIndex(uint32_t Index) {
   DEBUG(dbgs() << "InputFunction::setTableIndex: " << getName() << " -> "
                << Index << "\n");
   assert(!hasTableIndex());
   TableIndex = Index;
 }
+
+static unsigned getRelocWidthPadded(const WasmRelocation &Rel) {
+  switch (Rel.Type) {
+  case R_WEBASSEMBLY_TYPE_INDEX_LEB:
+  case R_WEBASSEMBLY_FUNCTION_INDEX_LEB:
+  case R_WEBASSEMBLY_GLOBAL_INDEX_LEB:
+  case R_WEBASSEMBLY_MEMORY_ADDR_LEB:
+  case R_WEBASSEMBLY_TABLE_INDEX_SLEB:
+  case R_WEBASSEMBLY_MEMORY_ADDR_SLEB:
+    return 5;
+  case R_WEBASSEMBLY_TABLE_INDEX_I32:
+  case R_WEBASSEMBLY_MEMORY_ADDR_I32:
+    return 4;
+  default:
+    llvm_unreachable("unknown relocation type");
+  }
+}
+
+static unsigned writeCompressedReloc(const WasmRelocation &Rel, uint8_t *Buf,
+                                     uint32_t Value) {
+  switch (Rel.Type) {
+  case R_WEBASSEMBLY_TYPE_INDEX_LEB:
+  case R_WEBASSEMBLY_FUNCTION_INDEX_LEB:
+  case R_WEBASSEMBLY_GLOBAL_INDEX_LEB:
+  case R_WEBASSEMBLY_MEMORY_ADDR_LEB:
+    return encodeULEB128(Value, Buf);
+  case R_WEBASSEMBLY_TABLE_INDEX_SLEB:
+  case R_WEBASSEMBLY_MEMORY_ADDR_SLEB:
+    return encodeSLEB128(static_cast<int32_t>(Value), Buf);
+  case R_WEBASSEMBLY_TABLE_INDEX_I32:
+  case R_WEBASSEMBLY_MEMORY_ADDR_I32:
+    return 4;
+  default:
+    llvm_unreachable("unknown relocation type");
+  }
+}
+
+static unsigned getRelocWidth(const WasmRelocation &Rel, uint32_t Value) {
+  uint8_t Target[5];
+  return writeCompressedReloc(Rel, Target, Value);
+}
+
+void InputFunction::calculateSize() {
+  if (!File || !Config->CompressRelocTargets)
+    return;
+
+  DEBUG(dbgs() << "calculateSize: " << getName() "\n");
+
+  const uint8_t *SecStart = File->CodeSection->Content.data();
+  const uint8_t *FuncStart = SecStart + getInputSectionOffset();
+  uint32_t FunctionSizeLength;
+  uint32_t ExistingFuncSize = decodeULEB128(FuncStart, &FunctionSizeLength);
+
+  uint32_t Start = getInputSectionOffset();
+  uint32_t End = Start + Function->Size;
+
+  uint32_t LastRelocEnd = Start + FunctionSizeLength;
+  for (WasmRelocation &Rel : Relocations) {
+    DEBUG(dbgs() << "  region: " << (Rel.Offset - LastRelocEnd) << "\n");
+    assert(Rel.Offset >= LastRelocEnd);
+    CompressedFuncSize += Rel.Offset - LastRelocEnd;
+    CompressedFuncSize += getRelocWidth(Rel, File->calcNewValue(Rel));
+    LastRelocEnd = Rel.Offset + getRelocWidthPadded(Rel);
+    assert(LastRelocEnd <= End);
+  }
+  DEBUG(dbgs() << "  final region: " << (End - LastRelocEnd) << "\n");
+  CompressedFuncSize += End - LastRelocEnd;
+
+  // Now we know how long the resulting function is we can add the encoding
+  // of its length
+  uint8_t Buf[5];
+  CompressedSize = CompressedFuncSize + encodeULEB128(CompressedFuncSize, Buf);
+
+  DEBUG(dbgs() << "  calculateSize orig: " << Function->Size << "\n");
+  DEBUG(dbgs() << "  calculateSize  new: " << CompressedSize << "\n");
+}
+
+void InputFunction::writeTo(uint8_t *Buf) const {
+  if (!File || !Config->CompressRelocTargets)
+    return InputChunk::writeTo(Buf);
+
+  Buf += OutputOffset;
+  uint8_t *Orig = Buf;
+
+  const uint8_t *SecStart = File->CodeSection->Content.data();
+  const uint8_t *FuncStart = SecStart + getInputSectionOffset();
+  const uint8_t *End = FuncStart + Function->Size;
+  uint32_t Count;
+  decodeULEB128(Buf, &Count);
+  FuncStart += Count;
+
+  DEBUG(dbgs() << "write func: " << getName() << "\n");
+  Buf += encodeULEB128(CompressedFuncSize, Buf);
+  const uint8_t *LastRelocEnd = FuncStart;
+  for (const WasmRelocation &Rel : Relocations) {
+    unsigned ChunkSize = (SecStart + Rel.Offset) - LastRelocEnd;
+    DEBUG(dbgs() << "  write chunk: " << ChunkSize << "\n");
+    memcpy(Buf, LastRelocEnd, ChunkSize);
+    Buf += ChunkSize;
+    Buf += writeCompressedReloc(Rel, Buf, File->calcNewValue(Rel));
+    LastRelocEnd = SecStart + Rel.Offset + getRelocWidthPadded(Rel);
+  }
+
+  assert(LastRelocEnd <= End);
+  unsigned ChunkSize = End - LastRelocEnd;
+  DEBUG(dbgs() << "  write final chunk: " << ChunkSize << "\n");
+  memcpy(Buf, LastRelocEnd, ChunkSize);
+  DEBUG(dbgs() << "  total: " << (Buf + ChunkSize - Orig) << "\n");
+}

wasm/Options.td

 
 // The follow flags are shared with the ELF linker
 def color_diagnostics: F<"color-diagnostics">,
   HelpText<"Use colors in diagnostics">;
 
 def color_diagnostics_eq: J<"color-diagnostics=">,
   HelpText<"Use colors in diagnostics">;
 
+defm compress_relocs: B<"compress-relocs",
+    "Compress relocation targets that contain padding",
+    "Do not compress relocation targets that contain padding">;
+
 defm demangle: B<"demangle",
     "Demangle symbol names",
     "Do not demangle symbol names">;
 
 def entry: S<"entry">, MetaVarName<"<entry>">,
   HelpText<"Name of entry point symbol">;
 
 def error_limit: J<"error-limit=">,

wasm/OutputSections.cpp

 
 CodeSection::CodeSection(ArrayRef<InputFunction *> Functions)
     : OutputSection(WASM_SEC_CODE), Functions(Functions) {
   assert(Functions.size() > 0);
 
   raw_string_ostream OS(CodeSectionHeader);
   writeUleb128(OS, Functions.size(), "function count");
   OS.flush();
-  BodySize = CodeSectionHeader.size();
 
-  for (InputChunk *Func : Functions) {
+  BodySize = CodeSectionHeader.size();
+  for (InputFunction *Func : Functions) {
     Func->OutputOffset = BodySize;
+    Func->calculateSize();
     BodySize += Func->getSize();
   }
 
   createHeader(BodySize);
 }
 
 void CodeSection::writeTo(uint8_t *Buf) {
   log("writing " + toString(*this));
void CodeSection::writeTo(uint8_t *Buf) {
 
   // Write code section bodies
   parallelForEach(Functions,
                   [&](const InputChunk *Chunk) { Chunk->writeTo(Buf); });
 }
 
 uint32_t CodeSection::numRelocations() const {
   uint32_t Count = 0;
-  for (const InputChunk *Func : Functions)
-    Count += Func->NumRelocations();
+  for (const InputChunk *C : Functions)
+    Count += C->NumRelocations();
   return Count;
 }
 
 void CodeSection::writeRelocations(raw_ostream &OS) const {
   for (const InputChunk *C : Functions)
     C->writeRelocations(OS);
 }
 

wasm/Writer.cpp

void Writer::createSections() {
   createDataSection();
   createCustomSections();
 
   // Custom sections
   if (Config->Relocatable) {
     createLinkingSection();
     createRelocSections();
   }
+
   if (!Config->StripDebug && !Config->StripAll)
     createNameSection();
 
   for (OutputSection *S : OutputSections) {
     S->setOffset(FileSize);
     S->finalizeContents();
     FileSize += S->getSize();
   }