[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

292

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

184

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

232

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 -o %t.wasm %t2.o %t.o
+; RUN: obj2yaml %t.wasm | FileCheck %s
+
+; RUN: wasm-ld -O2 -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/ADT/StringSet.h"
 #include "llvm/BinaryFormat/Wasm.h"
 
 namespace lld {
 namespace wasm {
 
 struct Configuration {
   bool AllowUndefined;
+  bool CompressRelocTargets;
   bool Demangle;
   bool ExportTable;
   bool GcSections;
   bool ImportMemory;
   bool ImportTable;
   bool PrintGcSections;
   bool Relocatable;
   bool StripAll;
   bool StripDebug;
   bool StackFirst;
   uint32_t GlobalBase;
   uint32_t InitialMemory;
   uint32_t MaxMemory;
+  uint32_t Optimize;
   uint32_t ZStackSize;
   llvm::StringRef Entry;
   llvm::StringRef OutputFile;
 
   llvm::StringSet<> AllowUndefinedSymbols;
   std::vector<llvm::StringRef> SearchPaths;
 };
 
 // The only instance of Configuration struct.
 extern Configuration *Config;
 
 } // namespace wasm
 } // namespace lld
 
 #endif

wasm/Driver.cpp

void LinkerDriver::link(ArrayRef<const char *> ArgsArr) {
   Config->AllowUndefined = Args.hasArg(OPT_allow_undefined);
   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);
   errorHandler().FatalWarnings =
       Args.hasFlag(OPT_fatal_warnings, OPT_no_fatal_warnings, false);
   Config->ImportMemory = Args.hasArg(OPT_import_memory);
   Config->ImportTable = Args.hasArg(OPT_import_table);
+  Config->Optimize = args::getInteger(Args, OPT_O, 1);
   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);
 
+  Config->CompressRelocTargets = Config->Optimize > 1 && !Config->Relocatable;
+
   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(); };
 
   // Verifies the existing data at relocation targets matches our expectations.
   // This is performed only debug builds as an extra sanity check.
   void verifyRelocTargets() const;
 
   std::vector<WasmRelocation> Relocations;
   Kind SectionKind;
 };
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; }
   const ArrayRef<uint8_t> getFunctionBody() const { return Function->Body; }
   uint32_t getFunctionInputOffset() const { return getInputSectionOffset(); }
+  uint32_t getSize() const override {
+    if (Config->CompressRelocTargets && File) {
+      assert(CompressedSize);
+      return CompressedSize;
+    }
+    return 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);
 
+  // The size of a given input function can depend on the values of the
+  // LEB relocations within it.  This finalizeContents method is called after
+  // all the symbol values have be calcualted but before getSize() is ever
+  // called.
+  void calculateSize();
+
   const WasmSignature &Signature;
 
 protected:
   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;
+  uint32_t CompressedFuncSize = 0;
+  uint32_t 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);
 }
 
 void InputChunk::verifyRelocTargets() const {
   for (const WasmRelocation &Rel : Relocations) {
     uint32_t ExistingValue;

 }
 
 void InputFunction::setTableIndex(uint32_t Index) {
   DEBUG(dbgs() << "InputFunction::setTableIndex: " << getName() << " -> "
                << Index << "\n");
   assert(!hasTableIndex());
   TableIndex = Index;
 }
+
+// Write a relocation value without padding and return the number of bytes
+// witten.
+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 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 getRelocWidth(const WasmRelocation &Rel, uint32_t Value) {
+  uint8_t Target[5];
+  return writeCompressedReloc(Rel, Target, Value);
+}
+
+// Relocations of type LEB and SLEB in the code section are padded to 5 bytes
+// so that a fast linker can blindly overwrite them without needing to worry
+// about the number of bytes needed to encode the values.
+// However, for optimal output the code section can be compressed to remove
+// the padding then outputting non-relocatable files.
+// In this case we need to perform a size calculation based on the value at each
+// relocation.  At best we end up saving 4 bytes for each relocation entry.
+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;
+  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");
+    CompressedFuncSize += Rel.Offset - LastRelocEnd;
+    CompressedFuncSize += getRelocWidth(Rel, File->calcNewValue(Rel));
+    LastRelocEnd = Rel.Offset + getRelocWidthPadded(Rel);
+  }
+  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");
+}
+
+// Override the default writeTo method so that we can (optionally) write the
+// compressed version of the function.
+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);
+  }
+
+  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

 def no_color_diagnostics: F<"no-color-diagnostics">,
   HelpText<"Do not use colors in diagnostics">;
 
 def no_fatal_warnings: F<"no-fatal-warnings">;
 
 def o: JoinedOrSeparate<["-"], "o">, MetaVarName<"<path>">,
   HelpText<"Path to file to write output">;
 
+def O: JoinedOrSeparate<["-"], "O">, HelpText<"Optimize output file size">;
+
 defm print_gc_sections: B<"print-gc-sections",
     "List removed unused sections",
     "Do not list removed unused sections">;
 
 def relocatable: F<"relocatable">, HelpText<"Create relocatable object file">;
 
 def strip_all: F<"strip-all">, HelpText<"Strip all symbols">;
 

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) {
+  for (InputFunction *Func : Functions) {
     Func->OutputOffset = BodySize;
+    Func->calculateSize();
     BodySize += Func->getSize();
   }
 
   createHeader(BodySize);
 }
 
 void CodeSection::writeTo(uint8_t *Buf) {
   log("writing " + toString(*this));