[llvm] Fix movk on aarch64 with an .equ operand (PR #124428)

[via llvm-commits]

https://github.com/DaGenix updated https://github.com/llvm/llvm-project/pull/124428

>From 6b772fb84c6aba5ca170dad99966516ba776d09e Mon Sep 17 00:00:00 2001
From: Palmer Cox <p at lmercox.com>
Date: Sat, 25 Jan 2025 15:09:19 -0500
Subject: [PATCH 1/2] Fix movk on aarch64 with an .equ operand

Prior to 5da801386c2b820a4596fc6d8da6b5f4a6da94b4, this code worked:

    .equ    p4_low_b0, 0x0000
    movk    x1, p4_low_b0, lsl 16

That commit fixed a different bug, but accidentally broke the case where
the second operand to movk is not a literal.

In 442f066fc464e953b7783230e95ccf2a67ebfb38, a fix was applied to handle
the case where the second operand is a value like "(Val) >> 16".
However, that didn't appear to fix the test case in this commit. In this
commit, we extend the change to handle the case where the second operand
is a identifier defined by .equ.
---
 llvm/lib/Target/AArch64/AsmParser/AArch64AsmParser.cpp | 4 +++-
 llvm/test/MC/AArch64/arm64-basic-a64-instructions.s    | 3 +++
 2 files changed, 6 insertions(+), 1 deletion(-)

diff --git a/llvm/lib/Target/AArch64/AsmParser/AArch64AsmParser.cpp b/llvm/lib/Target/AArch64/AsmParser/AArch64AsmParser.cpp
index d3eda48f3276e9..27b052825d2133 100644
--- a/llvm/lib/Target/AArch64/AsmParser/AArch64AsmParser.cpp
+++ b/llvm/lib/Target/AArch64/AsmParser/AArch64AsmParser.cpp
@@ -5017,7 +5017,9 @@ bool AArch64AsmParser::parseOperand(OperandVector &Operands, bool isCondCode,
       return true;
     E = SMLoc::getFromPointer(getLoc().getPointer() - 1);
     Operands.push_back(AArch64Operand::CreateImm(IdVal, S, E, getContext()));
-    return false;
+
+    // Parse an optional shift/extend modifier.
+    return parseOptionalShiftExtend(getTok());
   }
   case AsmToken::Integer:
   case AsmToken::Real:
diff --git a/llvm/test/MC/AArch64/arm64-basic-a64-instructions.s b/llvm/test/MC/AArch64/arm64-basic-a64-instructions.s
index 2f58eadfc8462c..5af4140b3e4101 100644
--- a/llvm/test/MC/AArch64/arm64-basic-a64-instructions.s
+++ b/llvm/test/MC/AArch64/arm64-basic-a64-instructions.s
@@ -16,3 +16,6 @@
 // CHECK: crc32ch  w13, w17, w25           // encoding: [0x2d,0x56,0xd9,0x1a]
 // CHECK: crc32cw  wzr, w3, w5             // encoding: [0x7f,0x58,0xc5,0x1a]
 // CHECK: crc32cx  w18, w16, xzr           // encoding: [0x12,0x5e,0xdf,0x9a]
+
+        .equ	p4_low_b0, 0x0000
+        movk x1, p4_low_b0, lsl 16

>From e93b820cf6757a3fdcf2c8d689e684c99b117397 Mon Sep 17 00:00:00 2001
From: Palmer Cox <p at lmercox.com>
Date: Sat, 25 Jan 2025 20:30:27 -0500
Subject: [PATCH 2/2] Move test to basic-a64-instructions.s

---
 llvm/test/MC/AArch64/arm64-basic-a64-instructions.s | 3 ---
 llvm/test/MC/AArch64/basic-a64-instructions.s       | 5 +++++
 2 files changed, 5 insertions(+), 3 deletions(-)

diff --git a/llvm/test/MC/AArch64/arm64-basic-a64-instructions.s b/llvm/test/MC/AArch64/arm64-basic-a64-instructions.s
index 5af4140b3e4101..2f58eadfc8462c 100644
--- a/llvm/test/MC/AArch64/arm64-basic-a64-instructions.s
+++ b/llvm/test/MC/AArch64/arm64-basic-a64-instructions.s
@@ -16,6 +16,3 @@
 // CHECK: crc32ch  w13, w17, w25           // encoding: [0x2d,0x56,0xd9,0x1a]
 // CHECK: crc32cw  wzr, w3, w5             // encoding: [0x7f,0x58,0xc5,0x1a]
 // CHECK: crc32cx  w18, w16, xzr           // encoding: [0x12,0x5e,0xdf,0x9a]
-
-        .equ	p4_low_b0, 0x0000
-        movk x1, p4_low_b0, lsl 16
diff --git a/llvm/test/MC/AArch64/basic-a64-instructions.s b/llvm/test/MC/AArch64/basic-a64-instructions.s
index 0ae23d672e4a3e..14ac11f581a552 100644
--- a/llvm/test/MC/AArch64/basic-a64-instructions.s
+++ b/llvm/test/MC/AArch64/basic-a64-instructions.s
@@ -3347,6 +3347,11 @@ _func:
 // CHECK: mov      x2, #5299989643264             // encoding: [0x42,0x9a,0xc0,0xd2]
 // CHECK: movk     xzr, #{{4321|0x10e1}}, lsl #48 // encoding: [0x3f,0x1c,0xe2,0xf2]
 
+	.equ equvalue, 0x0001
+        movk x1, equvalue, lsl 16
+// CHECK: .set equvalue, 1
+// CHECK-NEXT: movk x1, #1, lsl #16 // encoding: [0x21,0x00,0xa0,0xf2]
+
         movz x2, #:abs_g0:sym
         movk w3, #:abs_g0_nc:sym