[LLVM][WebAssembly] Implement branch hinting proposal

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,151 @@
+# RUN: rm -rf %t; split-file %s %t
+
+;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
+;; with branch hints
+; RUN: llc -mcpu=mvp -filetype=obj %t/f1.ll -o %t/f1.o -mattr=+branch-hinting
+; RUN: llc -mcpu=mvp -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
+
+; 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
+
+;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
+;; without branch hints
+; RUN: llc -mcpu=mvp -filetype=obj %t/f1.ll -o %t/f1.o -mattr=-branch-hinting
+; RUN: llc -mcpu=mvp -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
+
+; NCHECK-NOT:         Name:            metadata.code.branch_hint
+; NCHECK-NOT:         Name:            branch-hinting
+
+;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
+;; with branch hints, but only the _start function is not removed by lld (no --export-all)
+; RUN: llc -mcpu=mvp -filetype=obj %t/f1.ll -o %t/f1.o -mattr=+branch-hinting
+; RUN: llc -mcpu=mvp -filetype=obj %t/f2.ll -o %t/f2.o -mattr=+branch-hinting
+; RUN: wasm-ld -o %t.wasm %t/f2.o %t/f1.o
+; RUN: obj2yaml %t.wasm | FileCheck --check-prefixes=RCHECK %s
+
+; RCHECK:          - Type:            CUSTOM
+; RCHECK:            Name:            metadata.code.branch_hint
+; RCHECK-NEXT:       Entries:
+; RCHECK-NEXT:         - FuncIdx:         0
+; RCHECK-NEXT:           Hints:
+; RCHECK-NEXT:             - Offset:          5
+; RCHECK-NEXT:               Size:            1
+; RCHECK-NEXT:               Data:            UNLIKELY
+; RCHECK-NEXT:    - Type:            CODE
+
+; RCHECK:         - Type:            CUSTOM
+; RCHECK-NEXT:      Name:            name
+; RCHECK-NEXT:      FunctionNames:
+; RCHECK-NEXT:        - Index:           0
+; RCHECK-NEXT:          Name:            _start
+
+; RCHECK:        - Type:            CUSTOM
+; RCHECK:          Name:            target_features
+; RCHECK-NEXT:     Features:
+; RCHECK-NEXT:       - Prefix:          USED
+; RCHECK-NEXT:         Name:            branch-hinting
+;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
+
+#--- f1.ll
+target triple = "wasm32-unknown-unknown"
+
+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/InputChunks.cpp

@@ -206,7 +206,9 @@ uint64_t InputChunk::getTombstone() const {
   if (const auto *s = dyn_cast<InputSection>(this)) {
     return s->tombstoneValue;
   }
-
+  if (const auto *s = dyn_cast<CodeMetaDataInputSection>(this)) {
+    return s->tombstoneValue;
+  }
   return 0;
 }
 
@@ -548,10 +550,96 @@ uint64_t InputSection::getTombstoneForSection(StringRef name) {
   // 0 is a valid function index.
   if (name.starts_with("llvm.func_attr."))
     return UINT64_C(-1);
-  // Returning 0 means there is no tombstone value for this section, and relocation
-  // will just use the addend.
+  if (name.starts_with("metadata.code."))
+    return UINT64_C(-1);
+  // Returning 0 means there is no tombstone value for this section, and
+  // relocation will just use the addend.
   return 0;
 }
 
+void CodeMetaDataInputSection::finalizeContents() {
+  RelocatedData.resize(data().size());
+  memcpy(RelocatedData.data(), data().data(), data().size());
+  relocate(RelocatedData.data());
+
+  const uint8_t *Buf = RelocatedData.data();
+  const uint8_t *End = Buf + RelocatedData.size();
+
+  const char *Success = nullptr;
+  unsigned NumFunctionsEncodingSize;
+  const unsigned NumFunctions =
+      decodeULEB128(Buf, &NumFunctionsEncodingSize, End, &Success);
+  if (Success != nullptr) {
+    fatal("Failed to decode number of functions in " + name + ": " + Success);
+  }
+  Buf += NumFunctionsEncodingSize;
+
+  for (unsigned i = 0; i < NumFunctions; ++i) {
+    const uint8_t *FunctionStartOffset = Buf;
+
+    unsigned FuncIdxEncodingSize;
+    const unsigned FuncIndex =
+        decodeULEB128(Buf, &FuncIdxEncodingSize, End, &Success);
+    if (Success != nullptr) {
+      fatal("Failed to decode function index in " + name + ": " + Success);
+    }
+    Buf += FuncIdxEncodingSize;
+
+    unsigned NumHintsEncodingSize;
+    const unsigned NumHints =
+        decodeULEB128(Buf, &NumHintsEncodingSize, End, &Success);
+    if (Success != nullptr) {
+      fatal("Failed to decode hint size in " + name + ": " + Success);
+    }
+    Buf += NumHintsEncodingSize;
+    for (unsigned j = 0; j < NumHints; ++j) {
+      unsigned OffsetEncodingSize;
+      decodeULEB128(Buf, &OffsetEncodingSize, End, &Success);
+      if (Success != nullptr) {
+        fatal("Failed to decode hint opcode in " + name + ": " + Success);
+      }
+      Buf += OffsetEncodingSize;
+
+      unsigned HintSizeEncodingSize;
+      const unsigned HintSize =
+          decodeULEB128(Buf, &HintSizeEncodingSize, End, &Success);
+      if (Success != nullptr) {
+        fatal("Failed to decode hint operand in " + name + ": " + Success);
+      }
+      Buf += HintSizeEncodingSize;
+      Buf += HintSize;
+    }
+    // skip entries for removed functions
+    // this is the whole reason we need to parse the hints here again
+    if (FuncIndex == static_cast<uint32_t>(getTombstone())) {
+      continue;
+    }
+    ArrayRef FunctionData(FunctionStartOffset, Buf - FunctionStartOffset);
+    Hints.emplace_back(FunctionData);
+  }
+  assert(Buf == End && "CodeMetaDataInputSection: not all data was consumed");
+}
+
+void CodeMetaDataInputSection::writeTo(uint8_t *Buf) const {
+  // write sorted hints to output buffer
+  // hints are implicitly sorted, since function indices are assigned by
+  // section order
+  for (const auto &ref : Hints) {
+    memcpy(Buf + outSecOff, ref.data(), ref.size());
+    Buf += ref.size();
+  }
+}
+
+uint32_t CodeMetaDataInputSection::getSize() const {
+  uint32_t size = 0;
+  for (const auto &ref : Hints) {
+    size += ref.size();
+  }
+  return size;
+}
+
+uint32_t CodeMetaDataInputSection::getNumFuncHints() const {
+  return Hints.size();
+}
 } // namespace wasm
 } // namespace lld

lld/wasm/InputChunks.h

@@ -45,6 +45,7 @@ class InputChunk {
     Function,
     SyntheticFunction,
     Section,
+    CodeMetadata,
   };
 
   StringRef name;
@@ -359,6 +360,26 @@ class InputSection : public InputChunk {
   const WasmSection &section;
 };
 
+class CodeMetaDataInputSection : public InputSection {
+public:
+  CodeMetaDataInputSection(const WasmSection &s, ObjFile *f, uint32_t alignment)
+      : InputSection(s, f, alignment) {
+    sectionKind = CodeMetadata;
+  }
+
+  static bool classof(const InputChunk *c) {
+    return c->kind() == InputChunk::CodeMetadata;
+  }
+  void finalizeContents();
+  void writeTo(uint8_t *Buf) const;
+  uint32_t getSize() const;
+  uint32_t getNumFuncHints() const;
+
+private:
+  // func_idx -> generic hints array
+  SmallVector<ArrayRef<uint8_t>> Hints;
+  std::vector<uint8_t> RelocatedData;
+};
 } // namespace wasm
 
 std::string toString(const wasm::InputChunk *);

lld/wasm/InputFiles.cpp

@@ -545,6 +545,8 @@ void ObjFile::parse(bool ignoreComdats) {
       uint32_t alignment = getCustomSectionAlignment(section);
       if (shouldMerge(section))
         customSec = make<MergeInputChunk>(section, this, alignment);
+      else if (section.Name == "metadata.code.branch_hint")
+        customSec = make<CodeMetaDataInputSection>(section, this, alignment);
       else
         customSec = make<InputSection>(section, this, alignment);
       customSec->discarded = isExcludedByComdat(customSec);