[JITLink][AArch64] Implement R_AARCH64_TSTBR14

[jobnoorman]

Summary

This relocation is used for the 14-bit immediate in test and branch
instructions.

[lhames]

LGTM. Thanks @jobnoorman!

llvm/include/llvm/ExecutionEngine/JITLink/aarch64.h

enum EdgeKind_aarch64 : Edge::Kind {
   ///
   /// Errors:
   ///   - The result of the unshifted part of the fixup expression must be
   ///     32-bit aligned otherwise an alignment error will be returned.
   ///   - The result of the fixup expression must fit into an int26 otherwise an
   ///     out-of-range error will be returned.
   Branch26PCRel,
 
+  /// A 14-bit PC-relative test and branch.
+  ///
+  /// Represents a PC-relative test and branch to a target within +/-32Kb. The
+  /// target must be 32-bit aligned.
+  ///
+  /// Fixup expression:
+  ///   Fixup <- (Target - Fixup + Addend) >> 2 : int14
+  ///
+  /// Notes:
+  ///   The '14' in the name refers to the number operand bits and follows the
+  /// naming convention used by the corresponding ELF relocation.
+  /// Since the low two bits must be zero (because of the 32-bit alignment of
+  /// the target) the operand is effectively a signed 16-bit number.
+  ///
+  ///
+  /// Errors:
+  ///   - The result of the unshifted part of the fixup expression must be
+  ///     32-bit aligned otherwise an alignment error will be returned.
+  ///   - The result of the fixup expression must fit into an int14 otherwise an
+  ///     out-of-range error will be returned.
+  TestAndBranch14PCRel,
+
   /// A 19-bit PC-relative conditional branch.
   ///
   /// Represents a PC-relative conditional branch to a target within +/-1Mb. The
   /// target must be 32-bit aligned.
   ///
   /// Fixup expression:
   ///   Fixup <- (Target - Fixup + Addend) >> 2 : int19
   ///

 const char *getEdgeKindName(Edge::Kind K);
 
 // Returns whether the Instr is LD/ST (imm12)
 inline bool isLoadStoreImm12(uint32_t Instr) {
   constexpr uint32_t LoadStoreImm12Mask = 0x3b000000;
   return (Instr & LoadStoreImm12Mask) == 0x39000000;
 }
 
+inline bool isTestAndBranchImm14(uint32_t Instr) {
+  constexpr uint32_t TestAndBranchImm14Mask = 0x7e000000;
+  return (Instr & TestAndBranchImm14Mask) == 0x36000000;
+}
+
 inline bool isCondBranchImm19(uint32_t Instr) {
   constexpr uint32_t CondBranchImm19Mask = 0xfe000000;
   return (Instr & CondBranchImm19Mask) == 0x54000000;
 }
 
 inline bool isCompAndBranchImm19(uint32_t Instr) {
   constexpr uint32_t CompAndBranchImm19Mask = 0x7e000000;
   return (Instr & CompAndBranchImm19Mask) == 0x34000000;
case LDRLiteral19: {
     if (Delta < -(1 << 20) || Delta > ((1 << 20) - 1))
       return makeTargetOutOfRangeError(G, B, E);
 
     uint32_t EncodedImm = ((static_cast<uint32_t>(Delta) >> 2) & 0x7ffff) << 5;
     uint32_t FixedInstr = RawInstr | EncodedImm;
     *(ulittle32_t *)FixupPtr = FixedInstr;
     break;
   }
+  case TestAndBranch14PCRel: {
+    assert((FixupAddress.getValue() & 0x3) == 0 &&
+           "Test and branch is not 32-bit aligned");
+    uint32_t RawInstr = *(ulittle32_t *)FixupPtr;
+    assert(isTestAndBranchImm14(RawInstr) &&
+           "RawInstr is not a test and branch");
+    int64_t Delta = E.getTarget().getAddress() + E.getAddend() - FixupAddress;
+    if (Delta & 0x3)
+      return make_error<JITLinkError>(
+          "Test and branch literal target is not 32-bit aligned");
+    if (!isInt<16>(Delta))
+      return makeTargetOutOfRangeError(G, B, E);
+    uint32_t EncodedImm = ((static_cast<uint32_t>(Delta) >> 2) & 0x3fff) << 5;
+    uint32_t FixedInstr = RawInstr | EncodedImm;
+    *(ulittle32_t *)FixupPtr = FixedInstr;
+    break;
+  }
   case CondBranch19PCRel: {
     assert((FixupAddress.getValue() & 0x3) == 0 &&
            "Conditional branch is not 32-bit aligned");
     uint32_t RawInstr = *(ulittle32_t *)FixupPtr;
     assert((isCondBranchImm19(RawInstr) || isCompAndBranchImm19(RawInstr)) &&
            "RawInstr is not a conditional branch");
     int64_t Delta = E.getTarget().getAddress() + E.getAddend() - FixupAddress;
     if (Delta & 0x3)

llvm/lib/ExecutionEngine/JITLink/ELF_aarch64.cpp

enum ELFAArch64RelocationKind : Edge::Kind {
     ELFLdSt16Abs12,
     ELFLdSt32Abs12,
     ELFLdSt64Abs12,
     ELFLdSt128Abs12,
     ELFMovwAbsG0,
     ELFMovwAbsG1,
     ELFMovwAbsG2,
     ELFMovwAbsG3,
+    ELFTstBr14,
     ELFCondBr19,
     ELFAbs32,
     ELFAbs64,
     ELFPrel32,
     ELFPrel64,
     ELFAdrGOTPage21,
     ELFLd64GOTLo12,
     ELFTLSDescAdrPage21,
getRelocationKind(const uint32_t Type) {
     case ELF::R_AARCH64_MOVW_UABS_G0_NC:
       return ELFMovwAbsG0;
     case ELF::R_AARCH64_MOVW_UABS_G1_NC:
       return ELFMovwAbsG1;
     case ELF::R_AARCH64_MOVW_UABS_G2_NC:
       return ELFMovwAbsG2;
     case ELF::R_AARCH64_MOVW_UABS_G3:
       return ELFMovwAbsG3;
+    case ELF::R_AARCH64_TSTBR14:
+      return ELFTstBr14;
     case ELF::R_AARCH64_CONDBR19:
       return ELFCondBr19;
     case ELF::R_AARCH64_ABS32:
       return ELFAbs32;
     case ELF::R_AARCH64_ABS64:
       return ELFAbs64;
     case ELF::R_AARCH64_PREL32:
       return ELFPrel32;
case ELFMovwAbsG3: {
           aarch64::getMoveWide16Shift(Instr) != 48)
         return make_error<JITLinkError>(
             "R_AARCH64_MOVW_UABS_G3 target is not a "
             "MOVK/MOVZ (imm16, LSL #48) instruction");
 
       Kind = aarch64::MoveWide16;
       break;
     }
+    case ELFTstBr14: {
+      uint32_t Instr = *(const ulittle32_t *)FixupContent;
+      if (!aarch64::isTestAndBranchImm14(Instr))
+        return make_error<JITLinkError>("R_AARCH64_TSTBR14 target is not a "
+                                        "test and branch instruction");
+
+      Kind = aarch64::TestAndBranch14PCRel;
+      break;
+    }
     case ELFCondBr19: {
       uint32_t Instr = *(const ulittle32_t *)FixupContent;
       if (!aarch64::isCondBranchImm19(Instr) &&
           !aarch64::isCompAndBranchImm19(Instr))
         return make_error<JITLinkError>("R_AARCH64_CONDBR19 target is not a "
                                         "conditional branch instruction");
 
       Kind = aarch64::CondBranch19PCRel;

llvm/lib/ExecutionEngine/JITLink/aarch64.cpp

const char *getEdgeKindName(Edge::Kind R) {
   case NegDelta32:
     return "NegDelta32";
   case Branch26PCRel:
     return "Branch26PCRel";
   case MoveWide16:
     return "MoveWide16";
   case LDRLiteral19:
     return "LDRLiteral19";
+  case TestAndBranch14PCRel:
+    return "TestAndBranch14PCRel";
   case CondBranch19PCRel:
     return "CondBranch19PCRel";
   case Page21:
     return "Page21";
   case PageOffset12:
     return "PageOffset12";
   case RequestGOTAndTransformToPage21:
     return "RequestGOTAndTransformToPage21";

llvm/test/ExecutionEngine/JITLink/AArch64/ELF_relocations.s

test_adr_gotpage_external:
         .size test_adr_gotpage_external, .-test_adr_gotpage_external
 
         .globl  test_ld64_gotlo12_external
         .p2align  2
 test_ld64_gotlo12_external:
         ldr   x0, [x0, :got_lo12:external_data]
         .size test_ld64_gotlo12_external, .-test_ld64_gotlo12_external
 
+# Check R_AARCH64_TSTBR14 for tbz
+#
+# jitlink-check: decode_operand(test_tstbr14_tbz, 2) = \
+# jitlink-check:     (test_tstbr14_tbz_target - test_tstbr14_tbz)[16:2]
+        .globl test_tstbr14_tbz, test_tstbr14_tbz_target
+        .p2align 2
+test_tstbr14_tbz:
+        tbz x0, 0, test_tstbr14_tbz_target
+        .skip (1 << 14)
+test_tstbr14_tbz_target:
+        .size test_tstbr14_tbz, .-test_tstbr14_tbz
+
+# Check R_AARCH64_TSTBR14 for tbnz
+#
+# jitlink-check: decode_operand(test_tstbr14_tbnz, 2) = \
+# jitlink-check:     (test_tstbr14_tbnz_target - test_tstbr14_tbnz)[16:2]
+        .globl test_tstbr14_tbnz, test_tstbr14_tbnz_target
+        .p2align 2
+test_tstbr14_tbnz:
+        tbnz x0, 0, test_tstbr14_tbnz_target
+        .skip (1 << 14)
+test_tstbr14_tbnz_target:
+        .size test_tstbr14_tbnz, .-test_tstbr14_tbnz
+
 # Check R_AARCH64_CONDBR19 for compare and branch instructions
 #
 # jitlink-check: decode_operand(test_condbr19_cbz, 1) = \
 # jitlink-check:     (test_condbr19_cbz_target - test_condbr19_cbz)[21:2]
         .globl test_condbr19_cbz, test_condbr19_cbz_target
         .p2align 2
 test_condbr19_cbz:
         cbz x0, test_condbr19_cbz_target