[lld] [RISCV][LLD] Zcmt RISC-V extension in lld (PR #183450)

[via llvm-commits]

================
@@ -1630,3 +1770,89 @@ void elf::mergeRISCVAttributesSections(Ctx &ctx) {
 }
 
 void elf::setRISCVTargetInfo(Ctx &ctx) { ctx.target.reset(new RISCV(ctx)); }
+
+TableJumpSection::TableJumpSection(Ctx &ctx)
+    : SyntheticSection(ctx, ".riscv.jvt", SHT_PROGBITS, SHF_ALLOC,
+                       /*alignment=*/64) {}
+
+int TableJumpSection::getCMJTEntryIndex(const Symbol *sym) const {
+  auto it = cmjtCandidates.find(sym);
+  return it != cmjtCandidates.end() ? it->second.index : -1;
+}
+
+int TableJumpSection::getCMJALTEntryIndex(const Symbol *sym) const {
+  auto it = cmjaltCandidates.find(sym);
+  return it != cmjaltCandidates.end() ? it->second.index : -1;
+}
+
+void TableJumpSection::addEntry(const Symbol *sym, int saved, bool isCMJT) {
+  auto &candidates = isCMJT ? cmjtCandidates : cmjaltCandidates;
+  candidates.try_emplace(sym, TableJumpEntry{0, -1}).first->second.saved +=
+      saved;
+}
+
+// Sort candidates by code size reduction (descending), truncate to maxSize,
+// drop entries whose reduction doesn't cover a table entry, and assign each
+// survivor a table index starting at baseIdx.
+static void
+selectEntries(Ctx &ctx,
+              llvm::MapVector<const Symbol *, TableJumpEntry> &candidates,
+              uint32_t maxSize, uint32_t baseIdx) {
+  SmallVector<std::pair<const Symbol *, TableJumpEntry>, 0> entries(
+      candidates.begin(), candidates.end());
+  // MapVector iteration is in insertion order, so stable_sort gives a
+  // deterministic ordering keyed on saved bytes descending.
+  llvm::stable_sort(entries, [](const auto &a, const auto &b) {
+    return a.second.saved > b.second.saved;
+  });
+  if (entries.size() > maxSize)
+    entries.resize(maxSize);
+  while (!entries.empty() &&
+         entries.back().second.saved < (int)ctx.arg.wordsize)
+    entries.pop_back();
+
+  candidates.clear();
+  for (auto [i, entry] : llvm::enumerate(entries))
+    candidates[entry.first] = {entry.second.saved, (int)(baseIdx + i)};
+}
+
+void TableJumpSection::finalizeContents() {
+  selectEntries(ctx, cmjtCandidates, maxCMJTEntrySize, /*baseIdx=*/0);
+  selectEntries(ctx, cmjaltCandidates, maxCMJALTEntrySize, startCMJALTEntryIdx);
+
+  int cmjtBenefit = 0, cmjaltBenefit = 0;
+  for (auto &[_, e] : cmjtCandidates)
+    cmjtBenefit += e.saved;
+  for (auto &[_, e] : cmjaltCandidates)
+    cmjaltBenefit += e.saved;
+
+  // cm.jalt forces padding the cm.jt region to startCMJALTEntryIdx entries.
+  // Drop cm.jalt if its benefit doesn't cover the padding plus its own entries.
+  size_t pad = startCMJALTEntryIdx - cmjtCandidates.size();
+  if (cmjaltBenefit < (int)((pad + cmjaltCandidates.size()) * ctx.arg.wordsize))
+    cmjaltCandidates.clear();
+
+  // Give up entirely if cm.jt doesn't reduce code size.
+  if (cmjtBenefit <= (int)(cmjtCandidates.size() * ctx.arg.wordsize)) {
+    Log(ctx) << "table jump relaxation didn't reduce code size";
+    cmjtCandidates.clear();
+    cmjaltCandidates.clear();
+  }
+}
+
+size_t TableJumpSection::getSize() const {
+  if (!cmjaltCandidates.empty())
+    return (startCMJALTEntryIdx + cmjaltCandidates.size()) * ctx.arg.wordsize;
+  return cmjtCandidates.size() * ctx.arg.wordsize;
+}
+
+void TableJumpSection::writeTo(uint8_t *buf) {
+  for (auto *cands : {&cmjtCandidates, &cmjaltCandidates})
+    for (auto &[sym, entry] : *cands) {
+      uint8_t *p = buf + entry.index * ctx.arg.wordsize;
+      if (ctx.arg.is64)
+        write64le(p, sym->getVA(ctx));
----------------
LukeZhuang wrote:

Done

https://github.com/llvm/llvm-project/pull/183450