[JITLink][RISCV] Use hashmap to find PCREL_HI20 edge

[hahnjo]

As noted in issues #68594 and #73935, JITLink/RISCV/ELF_ehframe.s fails with libstdc++'s expensive checks because getRISCVPCRelHi20 calls std::equal_range on the edges which may not be ordered by their offset. Instead let ELFJITLinker_riscv build a hashmap of all edges with type R_RISCV_PCREL_HI20 that can be looked up in constant time.

Closes #73935

[llvmbot]

@llvm/pr-subscribers-backend-risc-v

Author: Jonas Hahnfeld (hahnjo)

Changes

As noted in issues #68594 and #73935, JITLink/RISCV/ELF_ehframe.s fails with libstdc++'s expensive checks because getRISCVPCRelHi20 calls std::equal_range on the edges which may not be ordered by their offset. Instead let ELFJITLinker_riscv build a hashmap of all edges with type R_RISCV_PCREL_HI20 that can be looked up in constant time.

Closes #73935

---

Full diff: https://github.com/llvm/llvm-project/pull/78849.diff

1 Files Affected:

• (modified) llvm/lib/ExecutionEngine/JITLink/ELF_riscv.cpp (+35-33)

diff --git a/llvm/lib/ExecutionEngine/JITLink/ELF_riscv.cpp b/llvm/lib/ExecutionEngine/JITLink/ELF_riscv.cpp
index d0701ba08bd919..627f186ca59144 100644
--- a/llvm/lib/ExecutionEngine/JITLink/ELF_riscv.cpp
+++ b/llvm/lib/ExecutionEngine/JITLink/ELF_riscv.cpp
@@ -133,38 +133,6 @@ const uint8_t
 namespace llvm {
 namespace jitlink {
 
-static Expected<const Edge &> getRISCVPCRelHi20(const Edge &E) {
-  using namespace riscv;
-  assert((E.getKind() == R_RISCV_PCREL_LO12_I ||
-          E.getKind() == R_RISCV_PCREL_LO12_S) &&
-         "Can only have high relocation for R_RISCV_PCREL_LO12_I or "
-         "R_RISCV_PCREL_LO12_S");
-
-  const Symbol &Sym = E.getTarget();
-  const Block &B = Sym.getBlock();
-  orc::ExecutorAddrDiff Offset = Sym.getOffset();
-
-  struct Comp {
-    bool operator()(const Edge &Lhs, orc::ExecutorAddrDiff Offset) {
-      return Lhs.getOffset() < Offset;
-    }
-    bool operator()(orc::ExecutorAddrDiff Offset, const Edge &Rhs) {
-      return Offset < Rhs.getOffset();
-    }
-  };
-
-  auto Bound =
-      std::equal_range(B.edges().begin(), B.edges().end(), Offset, Comp{});
-
-  for (auto It = Bound.first; It != Bound.second; ++It) {
-    if (It->getKind() == R_RISCV_PCREL_HI20)
-      return *It;
-  }
-
-  return make_error<JITLinkError>(
-      "No HI20 PCREL relocation type be found for LO12 PCREL relocation type");
-}
-
 static uint32_t extractBits(uint32_t Num, unsigned Low, unsigned Size) {
   return (Num & (((1ULL << Size) - 1) << Low)) >> Low;
 }
@@ -184,9 +152,43 @@ class ELFJITLinker_riscv : public JITLinker<ELFJITLinker_riscv> {
 public:
   ELFJITLinker_riscv(std::unique_ptr<JITLinkContext> Ctx,
                      std::unique_ptr<LinkGraph> G, PassConfiguration PassConfig)
-      : JITLinker(std::move(Ctx), std::move(G), std::move(PassConfig)) {}
+      : JITLinker(std::move(Ctx), std::move(G), std::move(PassConfig)) {
+    JITLinkerBase::getPassConfig().PostAllocationPasses.push_back(
+        [this](LinkGraph &G) { return gatherRISCVPCRelHi20(G); });
+  }
 
 private:
+  DenseMap<std::pair<const Block *, orc::ExecutorAddrDiff>, const Edge *>
+      RelHi20;
+
+  Error gatherRISCVPCRelHi20(LinkGraph &G) {
+    for (Block *B : G.blocks())
+      for (Edge &E : B->edges())
+        if (E.getKind() == R_RISCV_PCREL_HI20)
+          RelHi20[{B, E.getOffset()}] = &E;
+
+    return Error::success();
+  }
+
+  Expected<const Edge &> getRISCVPCRelHi20(const Edge &E) const {
+    using namespace riscv;
+    assert((E.getKind() == R_RISCV_PCREL_LO12_I ||
+            E.getKind() == R_RISCV_PCREL_LO12_S) &&
+           "Can only have high relocation for R_RISCV_PCREL_LO12_I or "
+           "R_RISCV_PCREL_LO12_S");
+
+    const Symbol &Sym = E.getTarget();
+    const Block &B = Sym.getBlock();
+    orc::ExecutorAddrDiff Offset = Sym.getOffset();
+
+    auto It = RelHi20.find({&B, Offset});
+    if (It != RelHi20.end())
+      return *It->second;
+
+    return make_error<JITLinkError>("No HI20 PCREL relocation type be found "
+                                    "for LO12 PCREL relocation type");
+  }
+
   Error applyFixup(LinkGraph &G, Block &B, const Edge &E) const {
     using namespace riscv;
     using namespace llvm::support;

[kosarev]

@mtvec Can you take a look, please?

[mtvec]

This looks good to me but please wait for @lhames's review.

[kosarev]

@lhames Ping.

[lhames]

Looks good to me.

In the future I think we'll want to have edges sorted by offset (perhaps for performance we'll only sort lazily when needed), so it might be worth adding a comment to revisit this if / when we do start sorting edges.

[hahnjo]

In the future I think we'll want to have edges sorted by offset (perhaps for performance we'll only sort lazily when needed), so it might be worth adding a comment to revisit this if / when we do start sorting edges.

Yes, I considered that. It's maybe a bit trickier because passes can add edges on the fly (?), but on the other hand, even sorted offsets doesn't get us to constant lookup times as we have with the maps...