[LLVM][WebAssembly] Implement branch hinting proposal

This commit implements the WebAssembly branch hinting proposal, as detailed at https://webassembly.github.io/branch-hinting/metadata/code/binary.html. This proposal introduces a mechanism to convey branch likelihood information to the WebAssembly engine, allowing for more effective performance optimizations.

The proposal specifies a new custom section named `metadata.code.branch_hint`. This section can contain a sequence of hints, where each hint is a single byte that applies to a corresponding `br_if` or `if` instruction. The hint values are:
 - `0x00` (`unlikely`): The branch is unlikely to be taken.
 - `0x01` (`likely`): The branch is likely to be taken.

This implementation includes the following changes:
 - Addition of the "branch-hinting" feature (flag)
 - Collection of edge probabilities in CFGStackify pass
 - Outputting of `metadata.code.branch_hint` section in WebAssemblyAsmPrinter
 - Addition of the `WebAssembly::Specifier::S_DEBUG_REF` symbol ref specifier
 - Custom relaxation of leb128 fragments for storage of uleb128 encoded function indices and instruction offsets
 - Custom handling of code metadata sections in lld, required since the proposal requires code metadata sections to start with a combined count of function hints, followed by an ordered list of function hints.

This change is purely an optimization and does not alter the semantics of WebAssembly programs.

Author: Lukasdoe

clang/include/clang/Driver/Options.td

@@ -5244,6 +5244,8 @@ def mharden_sls_EQ : Joined<["-"], "mharden-sls=">, Group<m_Group>,
 
 def matomics : Flag<["-"], "matomics">, Group<m_wasm_Features_Group>;
 def mno_atomics : Flag<["-"], "mno-atomics">, Group<m_wasm_Features_Group>;
+def mbranch_hinting : Flag<["-"], "mbranch-hinting">, Group<m_wasm_Features_Group>;
+def mno_branch_hinting : Flag<["-"], "mno-branch-hinting">, Group<m_wasm_Features_Group>;
 def mbulk_memory : Flag<["-"], "mbulk-memory">, Group<m_wasm_Features_Group>;
 def mno_bulk_memory : Flag<["-"], "mno-bulk-memory">, Group<m_wasm_Features_Group>;
 def mbulk_memory_opt : Flag<["-"], "mbulk-memory-opt">, Group<m_wasm_Features_Group>;

clang/lib/Basic/Targets/WebAssembly.cpp

@@ -53,6 +53,7 @@ bool WebAssemblyTargetInfo::setABI(const std::string &Name) {
 bool WebAssemblyTargetInfo::hasFeature(StringRef Feature) const {
   return llvm::StringSwitch<bool>(Feature)
       .Case("atomics", HasAtomics)
+      .Case("branch-hinting", HasBranchHinting)
       .Case("bulk-memory", HasBulkMemory)
       .Case("bulk-memory-opt", HasBulkMemoryOpt)
       .Case("call-indirect-overlong", HasCallIndirectOverlong)
@@ -86,6 +87,8 @@ void WebAssemblyTargetInfo::getTargetDefines(const LangOptions &Opts,
   defineCPUMacros(Builder, "wasm", /*Tuning=*/false);
   if (HasAtomics)
     Builder.defineMacro("__wasm_atomics__");
+  if (HasBranchHinting)
+    Builder.defineMacro("__wasm_branch_hinting__");
   if (HasBulkMemory)
     Builder.defineMacro("__wasm_bulk_memory__");
   if (HasBulkMemoryOpt)
@@ -188,6 +191,7 @@ bool WebAssemblyTargetInfo::initFeatureMap(
   auto addBleedingEdgeFeatures = [&]() {
     addGenericFeatures();
     Features["atomics"] = true;
+    Features["branch-hinting"] = true;
     Features["exception-handling"] = true;
     Features["extended-const"] = true;
     Features["fp16"] = true;
@@ -218,6 +222,14 @@ bool WebAssemblyTargetInfo::handleTargetFeatures(
       HasAtomics = false;
       continue;
     }
+    if (Feature == "+branch-hinting") {
+      HasBranchHinting = true;
+      continue;
+    }
+    if (Feature == "-branch-hinting") {
+      HasBranchHinting = false;
+      continue;
+    }
     if (Feature == "+bulk-memory") {
       HasBulkMemory = true;
       continue;

clang/lib/Basic/Targets/WebAssembly.h

@@ -58,6 +58,7 @@ class LLVM_LIBRARY_VISIBILITY WebAssemblyTargetInfo : public TargetInfo {
   } SIMDLevel = NoSIMD;
 
   bool HasAtomics = false;
+  bool HasBranchHinting = false;
   bool HasBulkMemory = false;
   bool HasBulkMemoryOpt = false;
   bool HasCallIndirectOverlong = false;

lld/test/wasm/code-metadata-branch-hints.ll

@@ -0,0 +1,115 @@
+# RUN: rm -rf %t; split-file %s %t
+; RUN: llc -mcpu=mvp -mtriple=wasm32-unknown-unknown -filetype=obj %t/f1.ll -o %t/f1.o -mattr=+branch-hinting
+; RUN: llc -mcpu=mvp -mtriple=wasm32-unknown-unknown -filetype=obj %t/f2.ll -o %t/f2.o -mattr=+branch-hinting
+; RUN: wasm-ld --export-all -o %t.wasm %t/f2.o %t/f1.o
+; RUN: obj2yaml %t.wasm | FileCheck --check-prefixes=CHECK %s
+
+; RUN: llc -mcpu=mvp -mtriple=wasm32-unknown-unknown -filetype=obj %t/f1.ll -o %t/f1.o -mattr=-branch-hinting
+; RUN: llc -mcpu=mvp -mtriple=wasm32-unknown-unknown -filetype=obj %t/f2.ll -o %t/f2.o -mattr=-branch-hinting
+; RUN: wasm-ld --export-all -o %t.wasm %t/f2.o %t/f1.o
+; RUN: obj2yaml %t.wasm | FileCheck --check-prefixes=NCHECK %s
+
+; CHECK:          - Type:            CUSTOM
+; CHECK:            Name:            metadata.code.branch_hint
+; CHECK-NEXT:       Entries:
+; CHECK-NEXT:         - FuncIdx:         1
+; CHECK-NEXT:           Hints:
+; CHECK-NEXT:             - Offset:          7
+; CHECK-NEXT:               Size:            1
+; CHECK-NEXT:               Data:            UNLIKELY
+; CHECK-NEXT:             - Offset:          14
+; CHECK-NEXT:               Size:            1
+; CHECK-NEXT:               Data:            LIKELY
+; CHECK-NEXT:         - FuncIdx:         2
+; CHECK-NEXT:           Hints:
+; CHECK-NEXT:             - Offset:          5
+; CHECK-NEXT:               Size:            1
+; CHECK-NEXT:               Data:            LIKELY
+; CHECK-NEXT:         - FuncIdx:         3
+; CHECK-NEXT:           Hints:
+; CHECK-NEXT:             - Offset:          5
+; CHECK-NEXT:               Size:            1
+; CHECK-NEXT:               Data:            UNLIKELY
+; CHECK-NEXT:         - FuncIdx:         4
+; CHECK-NEXT:           Hints:
+; CHECK-NEXT:             - Offset:          5
+; CHECK-NEXT:               Size:            1
+; CHECK-NEXT:               Data:            LIKELY
+
+; 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:            test0
+; CHECK-NEXT:        - Index:           2
+; CHECK-NEXT:          Name:            test1
+; CHECK-NEXT:        - Index:           3
+; CHECK-NEXT:          Name:            _start
+; CHECK-NEXT:        - Index:           4
+; CHECK-NEXT:          Name:            test_func1
+
+; CHECK:        - Type:            CUSTOM
+; CHECK:          Name:            target_features
+; CHECK-NEXT:     Features:
+; CHECK-NEXT:       - Prefix:          USED
+; CHECK-NEXT:         Name:            branch-hinting
+
+
+; NCHECK-NOT:         Name:            metadata.code.branch_hint
+; NCHECK-NOT:         Name:            branch-hinting
+
+#--- f1.ll
+define i32 @_start(i32 %a) {
+entry:
+  %cmp = icmp eq i32 %a, 0
+  br i1 %cmp, label %if.then, label %if.else, !prof !0
+if.then:
+  ret i32 1
+if.else:
+  ret i32 2
+}
+
+define i32 @test_func1(i32 %a) {
+entry:
+  %cmp = icmp eq i32 %a, 0
+  br i1 %cmp, label %if.then, label %if.else, !prof !1
+if.then:
+  ret i32 1
+if.else:
+  ret i32 2
+}
+
+!0 = !{!"branch_weights", i32 2000, i32 1}
+!1 = !{!"branch_weights", i32 1, i32 2000}
+
+#--- f2.ll
+target triple = "wasm32-unknown-unknown"
+
+define i32 @test0(i32 %a) {
+entry:
+  %cmp0 = icmp eq i32 %a, 0
+  br i1 %cmp0, label %if.then, label %ret1, !prof !0
+if.then:
+  %cmp1 = icmp eq i32 %a, 1
+  br i1 %cmp1, label %ret1, label %ret2, !prof !1
+ret1:
+  ret i32 2
+ret2:
+  ret i32 1
+}
+
+define i32 @test1(i32 %a) {
+entry:
+  %cmp = icmp eq i32 %a, 0
+  br i1 %cmp, label %if.then, label %if.else, !prof !1
+if.then:
+  ret i32 1
+if.else:
+  ret i32 2
+}
+
+; the resulting branch hint is actually reversed, since llvm-br is turned into br_unless, inverting branch probs
+!0 = !{!"branch_weights", i32 2000, i32 1}
+!1 = !{!"branch_weights", i32 1, i32 2000}

lld/wasm/OutputSections.cpp

@@ -270,6 +270,62 @@ void CustomSection::writeTo(uint8_t *buf) {
     section->writeTo(buf);
 }
 
+void CodeMetaDataSection::writeTo(uint8_t *buf) {
+  log("writing " + toString(*this) + " offset=" + Twine(offset) +
+      " size=" + Twine(getSize()) + " chunks=" + Twine(inputSections.size()));
+
+  assert(offset);
+  buf += offset;
+
+  // Write section header
+  memcpy(buf, header.data(), header.size());
+  buf += header.size();
+  memcpy(buf, nameData.data(), nameData.size());
+  buf += nameData.size();
+
+  uint32_t TotalNumHints = 0;
+  for (const InputChunk *section :
+       make_range(inputSections.rbegin(), inputSections.rend())) {
+    section->writeTo(buf);
+    unsigned EncodingSize;
+    uint32_t NumHints =
+        decodeULEB128(buf + section->outSecOff, &EncodingSize, nullptr);
+    if (EncodingSize != 5) {
+      fatal("Unexpected encoding size for function hint vec size in " + name +
+            ": must be exactly 5 bytes.");
+    }
+    TotalNumHints += NumHints;
+  }
+  encodeULEB128(TotalNumHints, buf, 5);
+}
+
+void CodeMetaDataSection::finalizeContents() {
+  finalizeInputSections();
+
+  raw_string_ostream os(nameData);
+  encodeULEB128(name.size(), os);
+  os << name;
+
+  bool firstSection = true;
+  for (InputChunk *section : inputSections) {
+    assert(!section->discarded);
+    payloadSize = alignTo(payloadSize, section->alignment);
+    if (firstSection) {
+      section->outSecOff = payloadSize;
+      payloadSize += section->getSize();
+      firstSection = false;
+    } else {
+      // adjust output offset so that each section write overwrites exactly the
+      // subsequent section's function hint vector size (which deduplicates)
+      section->outSecOff = payloadSize - 5;
+      // payload size should not include the hint vector size, which is deduped
+      payloadSize += section->getSize() - 5;
+    }
+  }
+
+  createHeader(payloadSize + nameData.size());
+}
+
 uint32_t CustomSection::getNumRelocations() const {
   uint32_t count = 0;
   for (const InputChunk *inputSect : inputSections)

lld/wasm/OutputSections.h

@@ -132,6 +132,15 @@ class CustomSection : public OutputSection {
   std::string nameData;
 };
 
+class CodeMetaDataSection : public CustomSection {
+public:
+  CodeMetaDataSection(std::string name, ArrayRef<InputChunk *> inputSections)
+      : CustomSection(name, inputSections) {}
+
+  void writeTo(uint8_t *buf) override;
+  void finalizeContents() override;
+};
+
 } // namespace wasm
 } // namespace lld
 

lld/wasm/SyntheticSections.h

@@ -43,7 +43,8 @@ class SyntheticSection : public OutputSection {
 
   void writeTo(uint8_t *buf) override {
     assert(offset);
-    log("writing " + toString(*this));
+    log("writing " + toString(*this) + " offset=" + Twine(offset) +
+        " size=" + Twine(getSize()));
     memcpy(buf + offset, header.data(), header.size());
     memcpy(buf + offset + header.size(), body.data(), body.size());
   }

lld/wasm/Writer.cpp

@@ -94,6 +94,8 @@ class Writer {
   void addSections();
 
   void createCustomSections();
+  // branch hint custom section must be placed before code section
+  void createBranchHintSection();
   void createSyntheticSections();
   void createSyntheticSectionsPostLayout();
   void finalizeSections();
@@ -164,15 +166,34 @@ void Writer::createCustomSections() {
   log("createCustomSections");
   for (auto &pair : customSectionMapping) {
     StringRef name = pair.first;
+    if (name == "metadata.code.branch_hint") {
+      // Branch hint section is created separately.
+      continue;
+    }
     LLVM_DEBUG(dbgs() << "createCustomSection: " << name << "\n");
+    OutputSection *Sec = make<CustomSection>(std::string(name), pair.second);
+    if (ctx.arg.relocatable || ctx.arg.emitRelocs) {
+      auto *sym = make<OutputSectionSymbol>(Sec);
+      out.linkingSec->addToSymtab(sym);
+      Sec->sectionSym = sym;
+    }
+    addSection(Sec);
+  }
+}
 
-    OutputSection *sec = make<CustomSection>(std::string(name), pair.second);
+void Writer::createBranchHintSection() {
+  log("createBranchHintSection");
+  std::string SectionName = "metadata.code.branch_hint";
+  if (const auto &Ins = customSectionMapping.find(SectionName);
+      Ins != customSectionMapping.end()) {
+    OutputSection *Sec =
+        make<CodeMetaDataSection>(std::move(SectionName), Ins->second);
     if (ctx.arg.relocatable || ctx.arg.emitRelocs) {
-      auto *sym = make<OutputSectionSymbol>(sec);
+      auto *sym = make<OutputSectionSymbol>(Sec);
       out.linkingSec->addToSymtab(sym);
-      sec->sectionSym = sym;
+      Sec->sectionSym = sym;
     }
-    addSection(sec);
+    addSection(Sec);
   }
 }
 
@@ -544,6 +565,8 @@ void Writer::addSections() {
   addSection(out.elemSec);
   addSection(out.dataCountSec);
 
+  createBranchHintSection();
+
   addSection(make<CodeSection>(out.functionSec->inputFunctions));
   addSection(make<DataSection>(segments));
 

llvm/include/llvm/BinaryFormat/Wasm.h

@@ -517,6 +517,22 @@ struct WasmSignature {
   WasmSignature() = default;
 };
 
+template <typename T> struct WasmCodeMetadataItemEntry {
+  uint32_t Offset;
+  uint32_t Size;
+  T Data;
+};
+
+template <typename T> struct WasmCodeMetadataFuncEntry {
+  uint32_t FuncIdx;
+  std::vector<WasmCodeMetadataItemEntry<T>> Hints;
+};
+
+enum class WasmCodeMetadataBranchHint : uint8_t {
+  UNLIKELY = 0x0,
+  LIKELY = 0x1,
+};
+
 // Useful comparison operators
 inline bool operator==(const WasmSignature &LHS, const WasmSignature &RHS) {
   return LHS.State == RHS.State && LHS.Returns == RHS.Returns &&

llvm/include/llvm/Object/Wasm.h

@@ -142,6 +142,10 @@ class LLVM_ABI WasmObjectFile : public ObjectFile {
   ArrayRef<wasm::WasmFeatureEntry> getTargetFeatures() const {
     return TargetFeatures;
   }
+  ArrayRef<wasm::WasmCodeMetadataFuncEntry<wasm::WasmCodeMetadataBranchHint>>
+  getBranchHints() const {
+    return BranchHints;
+  }
   ArrayRef<wasm::WasmSignature> types() const { return Signatures; }
   ArrayRef<wasm::WasmImport> imports() const { return Imports; }
   ArrayRef<wasm::WasmTable> tables() const { return Tables; }
@@ -275,12 +279,16 @@ class LLVM_ABI WasmObjectFile : public ObjectFile {
   Error parseProducersSection(ReadContext &Ctx);
   Error parseTargetFeaturesSection(ReadContext &Ctx);
   Error parseRelocSection(StringRef Name, ReadContext &Ctx);
+  Error parseCodeMetadataSection(StringRef Name, ReadContext &Ctx);
+  Error parseBranchHintSection(ReadContext &Ctx);
 
   wasm::WasmObjectHeader Header;
   std::vector<WasmSection> Sections;
   wasm::WasmDylinkInfo DylinkInfo;
   wasm::WasmProducerInfo ProducerInfo;
   std::vector<wasm::WasmFeatureEntry> TargetFeatures;
+  std::vector<wasm::WasmCodeMetadataFuncEntry<wasm::WasmCodeMetadataBranchHint>>
+      BranchHints;
   std::vector<wasm::WasmSignature> Signatures;
   std::vector<wasm::WasmTable> Tables;
   std::vector<wasm::WasmLimits> Memories;