[JITLink][AArch64] Implement MoveWide16 generic edge.

[sunho]

Summary

Implements MoveWide16 generic edge kind that can be used to patch MOVZ/MOVK (imm16) instructions.

[lhames]

Is it possible to add a # jitlink-check testcase for this?

[lhames]

Oh -- test cases are in https://reviews.llvm.org/D127585.

Ok -- LGTM. Thanks @sunho!

[lhames]

Oh -- getEdgeKindName is missing a case for the new relocation. Could you please add that?

[sunho]

Oops. I completely missed it. I've updated the code.

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

 
 enum EdgeKind_aarch64 : Edge::Kind {
   Branch26 = Edge::FirstRelocation,
   Pointer32,
   Pointer64,
   Pointer64Anon,
   Page21,
   PageOffset12,
+  MoveWide16,
   GOTPage21,
   GOTPageOffset12,
   TLVPage21,
   TLVPageOffset12,
   PointerToGOT,
   PairedAddend,
   LDRLiteral19,
   Delta32,
if (ImplicitShift == 0)
         ImplicitShift = 4;
 
     return ImplicitShift;
   }
 
   return 0;
 }
 
+// Returns whether the Instr is MOVK/MOVZ (imm16) with a zero immediate field
+inline bool isMoveWideImm16(uint32_t Instr) {
+  constexpr uint32_t MoveWideImm16Mask = 0x5f9fffe0;
+  return (Instr & MoveWideImm16Mask) == 0x52800000;
+}
+
+// Returns the amount the address operand of MOVK/MOVZ (imm16)
+// should be shifted right by.
+//
+// The shift value is specfied in the assembly as LSL #<shift>.
+inline unsigned getMoveWide16Shift(uint32_t Instr) {
+  if (isMoveWideImm16(Instr)) {
+    uint32_t ImplicitShift = (Instr >> 21) & 0b11;
+    return ImplicitShift << 4;
+  }
+
+  return 0;
+}
+
 /// Apply fixup expression for edge to block content.
 inline Error applyFixup(LinkGraph &G, Block &B, const Edge &E) {
   using namespace support;
 
   char *BlockWorkingMem = B.getAlreadyMutableContent().data();
   char *FixupPtr = BlockWorkingMem + E.getOffset();
   orc::ExecutorAddr FixupAddress = B.getAddress() + E.getOffset();
 
case PageOffset12: {
     if (TargetOffset & ((1 << ImmShift) - 1))
       return make_error<JITLinkError>("PAGEOFF12 target is not aligned");
 
     uint32_t EncodedImm = (TargetOffset >> ImmShift) << 10;
     uint32_t FixedInstr = RawInstr | EncodedImm;
     *(ulittle32_t *)FixupPtr = FixedInstr;
     break;
   }
+  case MoveWide16: {
+    uint64_t TargetOffset =
+        (E.getTarget().getAddress() + E.getAddend()).getValue();
+
+    uint32_t RawInstr = *(ulittle32_t *)FixupPtr;
+    assert(isMoveWideImm16(RawInstr) &&
+           "RawInstr isn't a MOVK/MOVZ instruction");
+
+    unsigned ImmShift = getMoveWide16Shift(RawInstr);
+    uint32_t Imm = (TargetOffset >> ImmShift) & 0xffff;
+    uint32_t FixedInstr = RawInstr | (Imm << 5);
+    *(ulittle32_t *)FixupPtr = FixedInstr;
+    break;
+  }
   case LDRLiteral19: {
     assert((FixupAddress.getValue() & 0x3) == 0 && "LDR is not 32-bit aligned");
     assert(E.getAddend() == 0 && "LDRLiteral19 with non-zero addend");
     uint32_t RawInstr = *(ulittle32_t *)FixupPtr;
     assert(RawInstr == 0x58000010 && "RawInstr isn't a 64-bit LDR literal");
     int64_t Delta = E.getTarget().getAddress() - FixupAddress;
     if (Delta & 0x3)
       return make_error<JITLinkError>("LDR literal target is not 32-bit "

llvm/lib/ExecutionEngine/JITLink/aarch64.cpp

const char *getEdgeKindName(Edge::Kind R) {
   case Pointer64:
     return "Pointer64";
   case Pointer64Anon:
     return "Pointer64Anon";
   case Page21:
     return "Page21";
   case PageOffset12:
     return "PageOffset12";
+  case MoveWide16:
+    return "MoveWide16";
   case GOTPage21:
     return "GOTPage21";
   case GOTPageOffset12:
     return "GOTPageOffset12";
   case TLVPage21:
     return "TLVPage21";
   case TLVPageOffset12:
     return "TLVPageOffset12";