[JITLink][RISC-V] Support R_RISCV_SET_ULEB128/R_RISCV_SUB_ULEB128 (#153778)

Support bolt instrument the elf binary which has exception handling table.

Fixes #153775

Author: zengdage

lld/ELF/InputSection.cpp

@@ -19,6 +19,7 @@
 #include "llvm/Support/Compiler.h"
 #include "llvm/Support/Compression.h"
 #include "llvm/Support/Endian.h"
+#include "llvm/Support/LEB128.h"
 #include "llvm/Support/xxhash.h"
 #include <algorithm>
 #include <optional>

lld/ELF/Target.h

@@ -326,15 +326,6 @@ inline void write64(Ctx &ctx, void *p, uint64_t v) {
   llvm::support::endian::write64(p, v, ctx.arg.endianness);
 }
 
-// Overwrite a ULEB128 value and keep the original length.
-inline uint64_t overwriteULEB128(uint8_t *bufLoc, uint64_t val) {
-  while (*bufLoc & 0x80) {
-    *bufLoc++ = 0x80 | (val & 0x7f);
-    val >>= 7;
-  }
-  *bufLoc = val;
-  return val;
-}
 } // namespace elf
 } // namespace lld
 

llvm/include/llvm/ExecutionEngine/JITLink/riscv.h

@@ -221,6 +221,18 @@ enum EdgeKind_riscv : Edge::Kind {
   /// Fixup expression:
   ///   Fixup <- Fixup - Target + Addend
   NegDelta32,
+
+  /// Set ULEB128-encoded value
+  ///
+  /// Fixup expression:
+  ///   Fixup <- Target + Addend
+  R_RISCV_SET_ULEB128,
+
+  /// Subtract from ULEB128-encoded value
+  ///
+  /// Fixup expression:
+  ///   Fixup <- V - Target - Addend
+  R_RISCV_SUB_ULEB128,
 };
 
 /// Returns a string name for the given riscv edge. For debugging purposes

llvm/include/llvm/Support/LEB128.h

@@ -221,6 +221,16 @@ inline uint64_t decodeULEB128AndIncUnsafe(const uint8_t *&p) {
   return decodeULEB128AndInc(p, nullptr);
 }
 
+/// Overwrite a ULEB128 value and keep the original length.
+inline uint64_t overwriteULEB128(uint8_t *bufLoc, uint64_t val) {
+  while (*bufLoc & 0x80) {
+    *bufLoc++ = 0x80 | (val & 0x7f);
+    val >>= 7;
+  }
+  *bufLoc = val;
+  return val;
+}
+
 enum class LEB128Sign { Unsigned, Signed };
 
 template <LEB128Sign Sign, typename T, typename U = char,

llvm/lib/ExecutionEngine/JITLink/ELF_riscv.cpp

@@ -154,18 +154,22 @@ class ELFJITLinker_riscv : public JITLinker<ELFJITLinker_riscv> {
                      std::unique_ptr<LinkGraph> G, PassConfiguration PassConfig)
       : JITLinker(std::move(Ctx), std::move(G), std::move(PassConfig)) {
     JITLinkerBase::getPassConfig().PostAllocationPasses.push_back(
-        [this](LinkGraph &G) { return gatherRISCVPCRelHi20(G); });
+        [this](LinkGraph &G) { return gatherRISCVPairs(G); });
   }
 
 private:
   DenseMap<std::pair<const Block *, orc::ExecutorAddrDiff>, const Edge *>
       RelHi20;
+  DenseMap<std::pair<const Block *, orc::ExecutorAddrDiff>, const Edge *>
+      SetULEB128;
 
-  Error gatherRISCVPCRelHi20(LinkGraph &G) {
+  Error gatherRISCVPairs(LinkGraph &G) {
     for (Block *B : G.blocks())
       for (Edge &E : B->edges())
         if (E.getKind() == R_RISCV_PCREL_HI20)
           RelHi20[{B, E.getOffset()}] = &E;
+        else if (E.getKind() == R_RISCV_SET_ULEB128)
+          SetULEB128[{B, E.getOffset()}] = &E;
 
     return Error::success();
   }
@@ -189,6 +193,20 @@ class ELFJITLinker_riscv : public JITLinker<ELFJITLinker_riscv> {
                                     "for LO12 PCREL relocation type");
   }
 
+  Expected<const Edge &> getRISCVSetULEB128(const Block &B,
+                                            const Edge &E) const {
+    using namespace riscv;
+    assert(E.getKind() == R_RISCV_SUB_ULEB128 &&
+           "Can only have pair relocation for R_RISCV_SUB_ULEB128");
+
+    auto It = SetULEB128.find({&B, E.getOffset()});
+    if (It != SetULEB128.end())
+      return *It->second;
+
+    return make_error<JITLinkError>(
+        "No RISCV_SET_ULEB128 relocation type be found");
+  }
+
   Error applyFixup(LinkGraph &G, Block &B, const Edge &E) const {
     using namespace riscv;
     using namespace llvm::support;
@@ -467,6 +485,21 @@ class ELFJITLinker_riscv : public JITLinker<ELFJITLinker_riscv> {
       *(little32_t *)FixupPtr = static_cast<uint32_t>(Value);
       break;
     }
+    case R_RISCV_SET_ULEB128:
+      break;
+    case R_RISCV_SUB_ULEB128: {
+      auto SetULEB128 = getRISCVSetULEB128(B, E);
+      if (!SetULEB128)
+        return SetULEB128.takeError();
+      uint64_t Value = SetULEB128->getTarget().getAddress() +
+                       SetULEB128->getAddend() - E.getTarget().getAddress() -
+                       E.getAddend();
+      if (overwriteULEB128(reinterpret_cast<uint8_t *>(FixupPtr), Value) >=
+          0x80)
+        return make_error<StringError>("ULEB128 value exceeds available space",
+                                       inconvertibleErrorCode());
+      break;
+    }
     }
     return Error::success();
   }
@@ -843,6 +876,10 @@ class ELFLinkGraphBuilder_riscv : public ELFLinkGraphBuilder<ELFT> {
       return EdgeKind_riscv::R_RISCV_32_PCREL;
     case ELF::R_RISCV_ALIGN:
       return EdgeKind_riscv::AlignRelaxable;
+    case ELF::R_RISCV_SET_ULEB128:
+      return EdgeKind_riscv::R_RISCV_SET_ULEB128;
+    case ELF::R_RISCV_SUB_ULEB128:
+      return EdgeKind_riscv::R_RISCV_SUB_ULEB128;
     }
 
     return make_error<JITLinkError>(

llvm/lib/ExecutionEngine/JITLink/riscv.cpp

@@ -84,6 +84,10 @@ const char *getEdgeKindName(Edge::Kind K) {
     return "AlignRelaxable";
   case NegDelta32:
     return "NegDelta32";
+  case R_RISCV_SET_ULEB128:
+    return "R_RISCV_SET_ULEB128";
+  case R_RISCV_SUB_ULEB128:
+    return "R_RISCV_SUB_ULEB128";
   }
   return getGenericEdgeKindName(K);
 }

llvm/test/ExecutionEngine/JITLink/RISCV/ELF_reloc_uleb128.s

@@ -0,0 +1,21 @@
+# RUN: rm -rf %t && mkdir -p %t
+# RUN: llvm-mc -triple=riscv64 -filetype=obj -o %t/riscv64_reloc_uleb128.o %s
+# RUN: llvm-mc -triple=riscv32 -filetype=obj -o %t/riscv32_reloc_uleb128.o %s
+# RUN: llvm-jitlink -noexec -check %s %t/riscv64_reloc_uleb128.o
+# RUN: llvm-jitlink -noexec -check %s %t/riscv32_reloc_uleb128.o
+
+# jitlink-check: *{4}(foo+8) = 0x180
+
+.global main
+main:
+  lw a0, foo
+
+.section ".text","",@progbits
+.type foo,@function
+foo:
+  nop
+  nop
+  .reloc ., R_RISCV_SET_ULEB128, foo+129
+  .reloc ., R_RISCV_SUB_ULEB128, foo+1
+  .uleb128 0x80
+  .size foo, 8