[AArch64] Support expression results as immediate values in mov

Summary:
This patch adds support of using the result of an expression as an
immediate value. For example,

0:
.skip 4
 1:
mov x0, 1b - 0b

is assembled to

mov x0, #4

Currently it does not support expressions requiring relocation unless
explicitly specified. This fixes PR#45781.

Reviewers: peter.smith, ostannard, efriedma

Reviewed By: efriedma

Subscribers: nickdesaulniers, llozano, manojgupta, efriedma, ostannard, kristof.beyls, hiraditya, danielkiss, llvm-commits

Tags: #llvm

Differential Revision: https://reviews.llvm.org/D80028

Author: jcai19

llvm/lib/Target/AArch64/AsmParser/AArch64AsmParser.cpp

@@ -970,11 +970,15 @@ class AArch64Operand : public MCParsedAsmOperand {
   bool isMOVZMovAlias() const {
     if (!isImm()) return false;
 
-    const MCConstantExpr *CE = dyn_cast<MCConstantExpr>(getImm());
-    if (!CE) return false;
-    uint64_t Value = CE->getValue();
+    const MCExpr *E = getImm();
+    if (const MCConstantExpr *CE = dyn_cast<MCConstantExpr>(E)) {
+      uint64_t Value = CE->getValue();
 
-    return AArch64_AM::isMOVZMovAlias(Value, Shift, RegWidth);
+      return AArch64_AM::isMOVZMovAlias(Value, Shift, RegWidth);
+    }
+    // Only supports the case of Shift being 0 if an expression is used as an
+    // operand
+    return !Shift && E;
   }
 
   template<int RegWidth, int Shift>
@@ -1774,9 +1778,13 @@ class AArch64Operand : public MCParsedAsmOperand {
   void addMOVZMovAliasOperands(MCInst &Inst, unsigned N) const {
     assert(N == 1 && "Invalid number of operands!");
 
-    const MCConstantExpr *CE = cast<MCConstantExpr>(getImm());
-    uint64_t Value = CE->getValue();
-    Inst.addOperand(MCOperand::createImm((Value >> Shift) & 0xffff));
+    const MCConstantExpr *CE = dyn_cast<MCConstantExpr>(getImm());
+    if (CE) {
+      uint64_t Value = CE->getValue();
+      Inst.addOperand(MCOperand::createImm((Value >> Shift) & 0xffff));
+    } else {
+      addExpr(Inst, getImm());
+    }
   }
 
   template<int Shift>

llvm/lib/Target/AArch64/MCTargetDesc/AArch64AsmBackend.cpp

@@ -243,11 +243,22 @@ static uint64_t adjustFixupValue(const MCFixup &Fixup, const MCValue &Target,
         static_cast<AArch64MCExpr::VariantKind>(Target.getRefKind());
     if (AArch64MCExpr::getSymbolLoc(RefKind) != AArch64MCExpr::VK_ABS &&
         AArch64MCExpr::getSymbolLoc(RefKind) != AArch64MCExpr::VK_SABS) {
-      // VK_GOTTPREL, VK_TPREL, VK_DTPREL are movw fixups, but they can't
-      // ever be resolved in the assembler.
-      Ctx.reportError(Fixup.getLoc(),
-                      "relocation for a thread-local variable points to an "
-                      "absolute symbol");
+      if (!RefKind) {
+        // The fixup is an expression
+        if (SignedValue > 0xFFFF || SignedValue < -0xFFFF)
+          Ctx.reportError(Fixup.getLoc(),
+                          "fixup value out of range [-0xFFFF, 0xFFFF]");
+
+        // Invert the negative immediate because it will feed into a MOVN.
+        if (SignedValue < 0)
+          SignedValue = ~SignedValue;
+        Value = static_cast<uint64_t>(SignedValue);
+      } else
+        // VK_GOTTPREL, VK_TPREL, VK_DTPREL are movw fixups, but they can't
+        // ever be resolved in the assembler.
+        Ctx.reportError(Fixup.getLoc(),
+                        "relocation for a thread-local variable points to an "
+                        "absolute symbol");
       return Value;
     }
 
@@ -440,8 +451,9 @@ void AArch64AsmBackend::applyFixup(const MCAssembler &Asm, const MCFixup &Fixup,
   // FIXME: getFixupKindInfo() and getFixupKindNumBytes() could be fixed to
   // handle this more cleanly. This may affect the output of -show-mc-encoding.
   AArch64MCExpr::VariantKind RefKind =
-    static_cast<AArch64MCExpr::VariantKind>(Target.getRefKind());
-  if (AArch64MCExpr::getSymbolLoc(RefKind) == AArch64MCExpr::VK_SABS) {
+      static_cast<AArch64MCExpr::VariantKind>(Target.getRefKind());
+  if (AArch64MCExpr::getSymbolLoc(RefKind) == AArch64MCExpr::VK_SABS ||
+      (!RefKind && Fixup.getTargetKind() == AArch64::fixup_aarch64_movw)) {
     // If the immediate is negative, generate MOVN else MOVZ.
     // (Bit 30 = 0) ==> MOVN, (Bit 30 = 1) ==> MOVZ.
     if (SignedValue < 0)

llvm/lib/Target/AArch64/MCTargetDesc/AArch64MCCodeEmitter.cpp

@@ -569,23 +569,24 @@ unsigned AArch64MCCodeEmitter::fixMOVZ(const MCInst &MI, unsigned EncodedValue,
   if (UImm16MO.isImm())
     return EncodedValue;
 
-  const AArch64MCExpr *A64E = cast<AArch64MCExpr>(UImm16MO.getExpr());
-  switch (A64E->getKind()) {
-  case AArch64MCExpr::VK_DTPREL_G2:
-  case AArch64MCExpr::VK_DTPREL_G1:
-  case AArch64MCExpr::VK_DTPREL_G0:
-  case AArch64MCExpr::VK_GOTTPREL_G1:
-  case AArch64MCExpr::VK_TPREL_G2:
-  case AArch64MCExpr::VK_TPREL_G1:
-  case AArch64MCExpr::VK_TPREL_G0:
-    return EncodedValue & ~(1u << 30);
-  default:
-    // Nothing to do for an unsigned fixup.
-    return EncodedValue;
+  const MCExpr *E = UImm16MO.getExpr();
+  if (const AArch64MCExpr *A64E = dyn_cast<AArch64MCExpr>(E)) {
+    switch (A64E->getKind()) {
+    case AArch64MCExpr::VK_DTPREL_G2:
+    case AArch64MCExpr::VK_DTPREL_G1:
+    case AArch64MCExpr::VK_DTPREL_G0:
+    case AArch64MCExpr::VK_GOTTPREL_G1:
+    case AArch64MCExpr::VK_TPREL_G2:
+    case AArch64MCExpr::VK_TPREL_G1:
+    case AArch64MCExpr::VK_TPREL_G0:
+      return EncodedValue & ~(1u << 30);
+    default:
+      // Nothing to do for an unsigned fixup.
+      return EncodedValue;
+    }
   }
 
-
-  return EncodedValue & ~(1u << 30);
+  return EncodedValue;
 }
 
 void AArch64MCCodeEmitter::encodeInstruction(const MCInst &MI, raw_ostream &OS,

llvm/test/MC/AArch64/basic-a64-diagnostics.s

@@ -167,9 +167,9 @@
         // MOV alias should not accept any fiddling
         mov x2, xsp, #123
         mov wsp, w27, #0xfff, lsl #12
-// CHECK-ERROR: error: expected compatible register or logical immediate
+// CHECK-ERROR: error: invalid operand for instruction
 // CHECK-ERROR-NEXT:         mov x2, xsp, #123
-// CHECK-ERROR-NEXT:                 ^
+// CHECK-ERROR-NEXT:                      ^
 // CHECK-ERROR-NEXT: error: invalid operand for instruction
 // CHECK-ERROR-NEXT:         mov wsp, w27, #0xfff, lsl #12
 // CHECK-ERROR-NEXT:                       ^

llvm/test/MC/AArch64/mov-expr-as-immediate.s

@@ -0,0 +1,16 @@
+// RUN: llvm-mc -triple aarch64-none-linux-gnu %s -filetype=obj -o %t | llvm-objdump --triple aarch64-none-linux-gnu -Dr %t | FileCheck %s
+
+0:
+.skip 4
+1:
+mov x0, 1b - 0b
+// CHECK: mov x0, #4
+mov x0, 0b - 1b
+// CHECK: mov x0, #-4
+mov x0, 0b - 0b
+// CHECK: mov x0, #0
+mov x0, 1b - 2 - 0b + 6
+// CHECK: mov x0, #8
+mov x0, #:abs_g0_s:1b
+// CHECK: mov x0, #0
+// CHECK-NEXT: R_AARCH64_MOVW_SABS_G0	.text+0x4

llvm/test/MC/AArch64/mov-expression-as-immediate.s

@@ -0,0 +1,9 @@
+// RUN: llvm-mc -triple aarch64-none-linux-gnu %s -filetype=obj -o %t | llvm-objdump -d %t | FileCheck %s
+
+0:
+.skip 4
+1:
+mov x0, 1b - 0b
+// CHECK: mov x0, #4
+mov x0, 0b - 1b
+// CHECK: mov x0, #-4

llvm/test/MC/AArch64/mov-unsupported-expr-as-immediate.s

@@ -0,0 +1,17 @@
+// RUN: not llvm-mc -triple aarch64-none-linux-gnu %s -filetype=obj -o /dev/null 2>&1 | FileCheck %s
+
+0:
+.skip 0x10000
+1:
+mov x0, 1b - 0b
+// CHECK: error: fixup value out of range
+// CHECK: mov x0, 1b - 0b
+// CHECK: ^
+mov x0, 0b - 1b
+// CHECK: error: fixup value out of range
+// CHECK: mov x0, 0b - 1b
+// CHECK: ^
+mov x0, 1b
+// CHECK: error: invalid fixup for movz/movk instruction
+// CHECK: mov x0, 1b
+// CHECK: ^