[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/Inputs/branch-hints.ll

@@ -0,0 +1,29 @@
+target triple = "wasm32-unknown-unknown"
+
+define i32 @test0(i32 %a) {
+entry:
+  %cmp0 = icmp eq i32 %a, 0
+  ; This metadata hints that the true branch is overwhelmingly likely.
+  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/test/wasm/code-metadata-branch-hints.ll

@@ -0,0 +1,37 @@
+; RUN: llc -filetype=obj %s -o %t1.o -mattr=+branch-hinting
+; RUN: llc -filetype=obj %S/Inputs/branch-hints.ll -o %t2.o -mattr=+branch-hinting
+; RUN: wasm-ld --export-all -o %t.wasm %t1.o %t2.o
+; RUN: obj2yaml %t.wasm | FileCheck %s
+
+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}
+
+; CHECK:        - Type:            CUSTOM
+; CHECK-NEXT:     Name:            metadata.code.branch_hint
+; CHECK-NEXT:     Payload:         '84808080008180808000010501008280808000010501018380808000020701000E0101848080800001050101'
+
+; CHECK:        - Type:            CUSTOM
+; CHECK:          Name:            target_features
+; CHECK-NEXT:     Features:
+; CHECK:            - Prefix:          USED
+; CHECK-NEXT:         Name:            branch-hinting

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/lib/Target/WebAssembly/MCTargetDesc/WebAssemblyAsmBackend.cpp

@@ -21,6 +21,7 @@
 #include "llvm/MC/MCSubtargetInfo.h"
 #include "llvm/MC/MCSymbol.h"
 #include "llvm/MC/MCWasmObjectWriter.h"
+#include "llvm/Support/Casting.h"
 #include "llvm/Support/raw_ostream.h"
 
 using namespace llvm;
@@ -44,6 +45,9 @@ class WebAssemblyAsmBackend final : public MCAsmBackend {
   std::unique_ptr<MCObjectTargetWriter>
   createObjectTargetWriter() const override;
 
+  std::pair<bool, bool> relaxLEB128(const MCAssembler &Asm, MCLEBFragment &LF,
+                                    int64_t &Value) const override;
+
   bool writeNopData(raw_ostream &OS, uint64_t Count,
                     const MCSubtargetInfo *STI) const override;
 };
@@ -70,6 +74,27 @@ WebAssemblyAsmBackend::getFixupKindInfo(MCFixupKind Kind) const {
   return Infos[Kind - FirstTargetFixupKind];
 }
 
+std::pair<bool, bool>
+WebAssemblyAsmBackend::relaxLEB128(const MCAssembler &assembler,
+                                   MCLEBFragment &LF, int64_t &Value) const {
+  const MCExpr &Expr = LF.getValue();
+  if (Expr.getKind() == MCExpr::ExprKind::SymbolRef) {
+    const MCSymbolRefExpr &SymExpr = llvm::cast<MCSymbolRefExpr>(Expr);
+    if (static_cast<WebAssembly::Specifier>(SymExpr.getSpecifier()) ==
+        WebAssembly::S_DEBUG_REF) {
+      Value = assembler.getSymbolOffset(SymExpr.getSymbol());
+      return std::make_pair(true, false);
+    }
+  }
+  // currently, this is only used for leb128 encoded function indices
+  // that require relocations
+  LF.getFixups().push_back(
+      MCFixup::create(0, &Expr, WebAssembly::fixup_uleb128_i32, Expr.getLoc()));
+  // ensure that the stored placeholder is large enough to hold any 32-bit val
+  Value = UINT32_MAX;
+  return std::make_pair(true, false);
+}
+
 bool WebAssemblyAsmBackend::writeNopData(raw_ostream &OS, uint64_t Count,
                                          const MCSubtargetInfo *STI) const {
   for (uint64_t I = 0; I < Count; ++I)

llvm/lib/Target/WebAssembly/MCTargetDesc/WebAssemblyMCAsmInfo.h

@@ -37,6 +37,7 @@ enum Specifier {
   S_TBREL,     // Table index relative to __table_base
   S_TLSREL,    // Memory address relative to __tls_base
   S_TYPEINDEX, // Reference to a symbol's type (signature)
+  S_DEBUG_REF, // Marker placed for generation of metadata.code.* section
 };
 }
 } // end namespace llvm

llvm/lib/Target/WebAssembly/MCTargetDesc/WebAssemblyWasmObjectWriter.cpp

@@ -93,10 +93,13 @@ unsigned WebAssemblyWasmObjectWriter::getRelocType(
   case WebAssembly::S_None:
     break;
   case WebAssembly::S_FUNCINDEX:
+    if (static_cast<unsigned>(Fixup.getKind()) ==
+        WebAssembly::fixup_uleb128_i32)
+      return wasm::R_WASM_FUNCTION_INDEX_LEB;
     return wasm::R_WASM_FUNCTION_INDEX_I32;
   }
 
-  switch (unsigned(Fixup.getKind())) {
+  switch (static_cast<unsigned>(Fixup.getKind())) {
   case WebAssembly::fixup_sleb128_i32:
     if (SymA.isFunction())
       return wasm::R_WASM_TABLE_INDEX_SLEB;

llvm/lib/Target/WebAssembly/WebAssembly.td

@@ -90,6 +90,10 @@ def FeatureWideArithmetic :
       SubtargetFeature<"wide-arithmetic", "HasWideArithmetic", "true",
                        "Enable wide-arithmetic instructions">;
 
+def FeatureBranchHinting :
+      SubtargetFeature<"branch-hinting", "HasBranchHinting", "true",
+                       "Enable branch hints for branch instructions">;
+
 //===----------------------------------------------------------------------===//
 // Architectures.
 //===----------------------------------------------------------------------===//
@@ -142,7 +146,7 @@ def : ProcessorModel<"bleeding-edge", NoSchedModel,
                        FeatureMultivalue, FeatureMutableGlobals,
                        FeatureNontrappingFPToInt, FeatureRelaxedSIMD,
                        FeatureReferenceTypes, FeatureSIMD128, FeatureSignExt,
-                       FeatureTailCall]>;
+                       FeatureTailCall, FeatureBranchHinting]>;
 
 //===----------------------------------------------------------------------===//
 // Target Declaration