[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

@@ -5276,6 +5276,8 @@ def mtail_call : Flag<["-"], "mtail-call">, Group<m_wasm_Features_Group>;
 def mno_tail_call : Flag<["-"], "mno-tail-call">, Group<m_wasm_Features_Group>;
 def mwide_arithmetic : Flag<["-"], "mwide-arithmetic">, Group<m_wasm_Features_Group>;
 def mno_wide_arithmetic : Flag<["-"], "mno-wide-arithmetic">, 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 mexec_model_EQ : Joined<["-"], "mexec-model=">, Group<m_wasm_Features_Driver_Group>,
                      Values<"command,reactor">,
                      HelpText<"Execution model (WebAssembly only)">,

clang/lib/Basic/Targets/WebAssembly.cpp

@@ -69,6 +69,7 @@ bool WebAssemblyTargetInfo::hasFeature(StringRef Feature) const {
       .Case("simd128", SIMDLevel >= SIMD128)
       .Case("tail-call", HasTailCall)
       .Case("wide-arithmetic", HasWideArithmetic)
+      .Case("branch-hinting", HasBranchHinting)
       .Default(false);
 }
 
@@ -116,6 +117,8 @@ void WebAssemblyTargetInfo::getTargetDefines(const LangOptions &Opts,
     Builder.defineMacro("__wasm_tail_call__");
   if (HasWideArithmetic)
     Builder.defineMacro("__wasm_wide_arithmetic__");
+  if (HasBranchHinting)
+    Builder.defineMacro("__wasm_branch_hinting__");
 
   Builder.defineMacro("__GCC_HAVE_SYNC_COMPARE_AND_SWAP_1");
   Builder.defineMacro("__GCC_HAVE_SYNC_COMPARE_AND_SWAP_2");
@@ -194,6 +197,7 @@ bool WebAssemblyTargetInfo::initFeatureMap(
     Features["multimemory"] = true;
     Features["tail-call"] = true;
     Features["wide-arithmetic"] = true;
+    Features["branch-hinting"] = true;
     setSIMDLevel(Features, RelaxedSIMD, true);
   };
   if (CPU == "generic") {
@@ -347,6 +351,14 @@ bool WebAssemblyTargetInfo::handleTargetFeatures(
       HasWideArithmetic = false;
       continue;
     }
+    if (Feature == "+branch-hinting") {
+      HasBranchHinting = true;
+      continue;
+    }
+    if (Feature == "-branch-hinting") {
+      HasBranchHinting = false;
+      continue;
+    }
 
     Diags.Report(diag::err_opt_not_valid_with_opt)
         << Feature << "-target-feature";

clang/lib/Basic/Targets/WebAssembly.h

@@ -72,6 +72,7 @@ class LLVM_LIBRARY_VISIBILITY WebAssemblyTargetInfo : public TargetInfo {
   bool HasSignExt = false;
   bool HasTailCall = false;
   bool HasWideArithmetic = false;
+  bool HasBranchHinting = false;
 
   std::string ABI;
 

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

@@ -0,0 +1,116 @@
+# RUN: rm -rf %t; split-file %s %t
+; RUN: llc -filetype=obj %t/f1.ll -o %t/f1.o -mattr=+branch-hinting
+; RUN: llc -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 -filetype=obj %t/f1.ll -o %t/f1.o -mattr=-branch-hinting
+; RUN: llc -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/Writer.cpp

@@ -165,14 +165,17 @@ void Writer::createCustomSections() {
   for (auto &pair : customSectionMapping) {
     StringRef name = pair.first;
     LLVM_DEBUG(dbgs() << "createCustomSection: " << name << "\n");
-
-    OutputSection *sec = make<CustomSection>(std::string(name), pair.second);
+    OutputSection *Sec;
+    if (name == "metadata.code.branch_hint")
+      Sec = make<CodeMetaDataSection>(std::string(name), pair.second);
+    else
+      Sec = make<CustomSection>(std::string(name), pair.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);
   }
 }
 

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;