[JITLink][LoongArch] Add R_LARCH_{B16,B21} relocations support (#121096)

Author: zhaoqi5

llvm/include/llvm/ExecutionEngine/JITLink/loongarch.h

@@ -41,6 +41,50 @@ enum EdgeKind_loongarch : Edge::Kind {
   ///
   Pointer32,
 
+  /// A 16-bit PC-relative branch.
+  ///
+  /// Represents a PC-relative branch to a target within +/-128Kb. The target
+  /// must be 4-byte aligned.
+  ///
+  /// Fixup expression:
+  ///   Fixup <- (Target - Fixup + Addend) >> 2 : int16
+  ///
+  /// Notes:
+  ///   The '16' in the name refers to the number operand bits and follows the
+  /// naming convention used by the corresponding ELF relocations. Since the low
+  /// two bits must be zero (because of the 4-byte alignment of the target) the
+  /// operand is effectively a signed 18-bit number.
+  ///
+  /// Errors:
+  ///   - The result of the unshifted part of the fixup expression must be
+  ///     4-byte aligned otherwise an alignment error will be returned.
+  ///   - The result of the fixup expression must fit into an int16 otherwise an
+  ///     out-of-range error will be returned.
+  ///
+  Branch16PCRel,
+
+  /// A 21-bit PC-relative branch.
+  ///
+  /// Represents a PC-relative branch to a target within +/-4Mb. The Target must
+  /// be 4-byte aligned.
+  ///
+  /// Fixup expression:
+  ///   Fixup <- (Target - Fixup + Addend) >> 2 : int21
+  ///
+  /// Notes:
+  ///   The '21' in the name refers to the number operand bits and follows the
+  /// naming convention used by the corresponding ELF relocations. Since the low
+  /// two bits must be zero (because of the 4-byte alignment of the target) the
+  /// operand is effectively a signed 23-bit number.
+  ///
+  /// Errors:
+  ///   - The result of the unshifted part of the fixup expression must be
+  ///     4-byte aligned otherwise an alignment error will be returned.
+  ///   - The result of the fixup expression must fit into an int21 otherwise an
+  ///     out-of-range error will be returned.
+  ///
+  Branch21PCRel,
+
   /// A 26-bit PC-relative branch.
   ///
   /// Represents a PC-relative call or branch to a target within +/-128Mb. The
@@ -213,6 +257,37 @@ inline Error applyFixup(LinkGraph &G, Block &B, const Edge &E) {
     *(ulittle32_t *)FixupPtr = Value;
     break;
   }
+  case Branch16PCRel: {
+    int64_t Value = TargetAddress - FixupAddress + Addend;
+
+    if (!isInt<18>(Value))
+      return makeTargetOutOfRangeError(G, B, E);
+
+    if (!isShiftedInt<16, 2>(Value))
+      return makeAlignmentError(orc::ExecutorAddr(FixupAddress), Value, 4, E);
+
+    uint32_t RawInstr = *(little32_t *)FixupPtr;
+    uint32_t Imm = static_cast<uint32_t>(Value >> 2);
+    uint32_t Imm15_0 = extractBits(Imm, /*Hi=*/15, /*Lo=*/0) << 10;
+    *(little32_t *)FixupPtr = RawInstr | Imm15_0;
+    break;
+  }
+  case Branch21PCRel: {
+    int64_t Value = TargetAddress - FixupAddress + Addend;
+
+    if (!isInt<23>(Value))
+      return makeTargetOutOfRangeError(G, B, E);
+
+    if (!isShiftedInt<21, 2>(Value))
+      return makeAlignmentError(orc::ExecutorAddr(FixupAddress), Value, 4, E);
+
+    uint32_t RawInstr = *(little32_t *)FixupPtr;
+    uint32_t Imm = static_cast<uint32_t>(Value >> 2);
+    uint32_t Imm15_0 = extractBits(Imm, /*Hi=*/15, /*Lo=*/0) << 10;
+    uint32_t Imm20_16 = extractBits(Imm, /*Hi=*/20, /*Lo=*/16);
+    *(little32_t *)FixupPtr = RawInstr | Imm15_0 | Imm20_16;
+    break;
+  }
   case Branch26PCRel: {
     int64_t Value = TargetAddress - FixupAddress + Addend;
 

llvm/lib/ExecutionEngine/JITLink/ELF_loongarch.cpp

@@ -58,6 +58,10 @@ class ELFLinkGraphBuilder_loongarch : public ELFLinkGraphBuilder<ELFT> {
       return Pointer32;
     case ELF::R_LARCH_32_PCREL:
       return Delta32;
+    case ELF::R_LARCH_B16:
+      return Branch16PCRel;
+    case ELF::R_LARCH_B21:
+      return Branch21PCRel;
     case ELF::R_LARCH_B26:
       return Branch26PCRel;
     case ELF::R_LARCH_PCALA_HI20:

llvm/lib/ExecutionEngine/JITLink/loongarch.cpp

@@ -44,6 +44,8 @@ const char *getEdgeKindName(Edge::Kind K) {
     KIND_NAME_CASE(Delta32)
     KIND_NAME_CASE(NegDelta32)
     KIND_NAME_CASE(Delta64)
+    KIND_NAME_CASE(Branch16PCRel)
+    KIND_NAME_CASE(Branch21PCRel)
     KIND_NAME_CASE(Branch26PCRel)
     KIND_NAME_CASE(Page20)
     KIND_NAME_CASE(PageOffset12)

llvm/test/ExecutionEngine/JITLink/LoongArch/ELF_loongarch32_relocations.s

@@ -103,6 +103,30 @@ test_gotoffset12_external:
     ld.w $a0, $a0, %got_pc_lo12(external_data)
     .size test_gotoffset12_external, .-test_gotoffset12_external
 
+## Check R_LARCH_B16 relocation for compare and branch instructions.
+
+# jitlink-check: decode_operand(test_br16, 2)[17:0] = \
+# jitlink-check:   (test_br16_target - test_br16)[17:0]
+    .globl test_br16, test_br16_target
+    .p2align 2
+test_br16:
+    beq $t1, $t2, %b16(test_br16_target)
+    .skip (1 << 16)
+test_br16_target:
+    .size test_br16, .-test_br16
+
+## Check R_LARCH_B21 relocation for compare and branch instructions.
+
+# jitlink-check: decode_operand(test_br21, 1)[22:0] = \
+# jitlink-check:   (test_br21_target - test_br21)[22:0]
+    .globl test_br21, test_br21_target
+    .p2align 2
+test_br21:
+    beqz $t1, %b21(test_br21_target)
+    .skip (1 << 21)
+test_br21_target:
+    .size test_br21, .-test_br21
+
 
     .globl named_data
     .p2align 4

llvm/test/ExecutionEngine/JITLink/LoongArch/ELF_loongarch64_relocations.s

@@ -116,7 +116,6 @@ test_gotoffset12_external:
     ld.d $a0, $a0, %got_pc_lo12(external_data)
     .size test_gotoffset12_external, .-test_gotoffset12_external
 
-
 ## Check R_LARCH_CALL36 relocation of a local function call.
 
 # jitlink-check: decode_operand(local_func_call36, 1)[19:0] = \
@@ -130,6 +129,31 @@ local_func_call36:
     jirl $ra, $ra, 0
     .size local_func_call36, .-local_func_call36
 
+## Check R_LARCH_B16 relocation for compare and branch instructions.
+
+# jitlink-check: decode_operand(test_br16, 2)[17:0] = \
+# jitlink-check:   (test_br16_target - test_br16)[17:0]
+    .globl test_br16, test_br16_target
+    .p2align 2
+test_br16:
+    beq $t1, $t2, %b16(test_br16_target)
+    .skip (1 << 16)
+test_br16_target:
+    .size test_br16, .-test_br16
+
+## Check R_LARCH_B21 relocation for compare and branch instructions.
+
+# jitlink-check: decode_operand(test_br21, 1)[22:0] = \
+# jitlink-check:   (test_br21_target - test_br21)[22:0]
+    .globl test_br21, test_br21_target
+    .p2align 2
+test_br21:
+    beqz $t1, %b21(test_br21_target)
+    .skip (1 << 21)
+test_br21_target:
+    .size test_br21, .-test_br21
+
+
     .globl named_data
     .p2align 4
     .type named_data,@object