Support other shift amounts.

rj-jesus · rj-jesus · commit 0a2928939cd0 · 2025-10-21T09:06:56.000-07:00
diff --git a/llvm/lib/Target/AArch64/AArch64ExpandImm.cpp b/llvm/lib/Target/AArch64/AArch64ExpandImm.cpp
@@ -239,15 +239,18 @@ static bool trySequenceOfOnes(uint64_t UImm,
   return true;
 }
 
-// Attempt to expand 64-bit immediate values whose negated upper half match
-// the lower half (for example, 0x1234'5678'edcb'a987).
-// Immediates of this form can generally be expanded via a sequence of
-// MOVN+MOVK to expand the lower half, followed by an EOR to shift and negate
-// the result to the upper half, e.g.:
+// Attempt to expand 64-bit immediate values that consist of shifted negated
+// components such as 0x1234'5678'edcb'a987, where the upper half is the
+// negation of the lower half. Immediates of this form can generally be
+// expanded via a sequence of MOVN+MOVK to expand the lower half, followed by
+// an EOR or EON to shift and negate the result to the upper half, for example:
 //   mov  x0, #-22137          // =0xffffffffffffa987
 //   movk x0, #60875, lsl #16  // =0xffffffffedcba987
 //   eor  x0, x0, x0, lsl #32  // =0xffffffffedcba987 ^ 0xedcba98700000000
 //                                =0x12345678edcba987.
+// The logic extends to other shift amounts in the range [17, 48) (outside that
+// range we get runs of ones/zeros that are optimised separately).
+//
 // When the lower half contains a 16-bit chunk of ones, such as
 // 0x0000'5678'ffff'a987, the intermediate MOVK is redundant.
 // Similarly, when it contains a 16-bit chunk of zeros, such as
@@ -261,35 +264,44 @@ static bool trySequenceOfOnes(uint64_t UImm,
 // compared to the default expansion based on MOV and MOVKs.
 static bool tryCopyWithNegation(uint64_t Imm, bool AllowThreeSequence,
                                 SmallVectorImpl<ImmInsnModel> &Insn) {
-  // We need the negation of the upper half of Imm to match the lower half.
   // Degenerate cases where Imm is a run of ones should be handled separately.
-  if ((~Imm >> 32) != (Imm & 0xffffffffULL) || llvm::isShiftedMask_64(Imm))
+  if (!Imm || llvm::isShiftedMask_64(Imm))
     return false;
 
   const unsigned Mask = 0xffff;
-  unsigned Opc = AArch64::EORXrs;
 
-  // If we have a chunk of all zeros in the lower half, we can save a MOVK by
-  // materialising the negated immediate and negating the result with an EON.
-  if ((Imm & Mask) == 0 || ((Imm >> 16) & Mask) == 0) {
-    Opc = AArch64::EONXrs;
-    Imm = ~Imm;
-  }
+  auto tryExpansion = [&](unsigned Opc, uint64_t C, unsigned N) {
+    assert((C >> 32) == 0xffffffffULL && "Invalid immediate");
+    const unsigned Imm0 = C & Mask;
+    const unsigned Imm16 = (C >> 16) & Mask;
+    if (Imm0 != Mask && Imm16 != Mask && !AllowThreeSequence)
+      return false;
+
+    if (Imm0 != Mask) {
+      Insn.push_back({AArch64::MOVNXi, Imm0 ^ Mask, 0});
+      if (Imm16 != Mask)
+        Insn.push_back({AArch64::MOVKXi, Imm16, 16});
+    } else {
+      Insn.push_back({AArch64::MOVNXi, Imm16 ^ Mask, 16});
+    }
 
-  unsigned Imm0 = Imm & Mask;
-  unsigned Imm16 = (Imm >> 16) & Mask;
-  if (Imm0 != Mask && Imm16 != Mask && !AllowThreeSequence)
-    return false;
-  if (Imm0 != Mask) {
-    Insn.push_back({AArch64::MOVNXi, Imm0 ^ Mask, 0});
-    if (Imm16 != Mask)
-      Insn.push_back({AArch64::MOVKXi, Imm16, 16});
-  } else {
-    Insn.push_back({AArch64::MOVNXi, Imm16 ^ Mask, 16});
+    Insn.push_back({Opc, 0, N});
+    return true;
+  };
+
+  for (unsigned N = 17; N < 48; ++N) {
+    // Attempt EOR.
+    uint64_t C = 0xffffffff00000000ULL | (Imm ^ (Imm << N));
+    if ((C ^ (C << N)) == Imm && tryExpansion(AArch64::EORXrs, C, N))
+      return true;
+
+    // Attempt EON.
+    C = 0xffffffff00000000ULL | (Imm ^ ~(~Imm << N));
+    if ((C ^ ~(C << N)) == Imm && tryExpansion(AArch64::EONXrs, C, N))
+      return true;
   }
 
-  Insn.push_back({Opc, 0, 32});
-  return true;
+  return false;
 }
 
 static uint64_t GetRunOfOnesStartingAt(uint64_t V, uint64_t StartPosition) {
@@ -698,7 +710,7 @@ void AArch64_IMM::expandMOVImm(uint64_t Imm, unsigned BitSize,
   if (BitSize == 64 && trySequenceOfOnes(UImm, Insn))
     return;
 
-  // Attempt to use a sequence of MOVN+MOVK+EOR (shifted register).
+  // Attempt to use a sequence of MOVN+MOVK+EOR/EON (shifted register).
   if (tryCopyWithNegation(Imm, /*AllowThreeSequence=*/true, Insn))
     return;
 
diff --git a/llvm/test/CodeGen/AArch64/arm64-movi.ll b/llvm/test/CodeGen/AArch64/arm64-movi.ll
@@ -557,9 +557,8 @@ define i64 @orr_32_eor_64() nounwind {
 define i64 @movn_0_eon_lsl_17() {
 ; CHECK-LABEL: movn_0_eon_lsl_17:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    mov x0, #-4370 // =0xffffffffffffeeee
-; CHECK-NEXT:    movk x0, #8738, lsl #16
-; CHECK-NEXT:    movk x0, #65534, lsl #32
+; CHECK-NEXT:    mov x0, #-61167 // =0xffffffffffff1111
+; CHECK-NEXT:    eon x0, x0, x0, lsl #17
 ; CHECK-NEXT:    ret
   ret i64 u0xfffffffe2222eeee
 }
@@ -576,19 +575,17 @@ define i64 @movn_0_eon_lsl_32() {
 define i64 @movn_0_eon_lsl_47() {
 ; CHECK-LABEL: movn_0_eon_lsl_47:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    mov x0, #61166 // =0xeeee
-; CHECK-NEXT:    movk x0, #32768, lsl #32
-; CHECK-NEXT:    movk x0, #34952, lsl #48
+; CHECK-NEXT:    mov x0, #-61167 // =0xffffffffffff1111
+; CHECK-NEXT:    eon x0, x0, x0, lsl #47
 ; CHECK-NEXT:    ret
   ret i64 u0x888880000000eeee
 }
 
 define i64 @movn_1_eon_lsl_17() {
 ; CHECK-LABEL: movn_1_eon_lsl_17:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    mov x0, #286261248 // =0x11100000
-; CHECK-NEXT:    movk x0, #8739, lsl #32
-; CHECK-NEXT:    movk x0, #65534, lsl #48
+; CHECK-NEXT:    mov x0, #-4008574977 // =0xffffffff1111ffff
+; CHECK-NEXT:    eon x0, x0, x0, lsl #17
 ; CHECK-NEXT:    ret
   ret i64 u0xfffe222311100000
 }
@@ -605,19 +602,17 @@ define i64 @movn_1_eon_lsl_32() {
 define i64 @movn_1_eon_lsl_46() {
 ; CHECK-LABEL: movn_1_eon_lsl_46:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    mov x0, #4008509440 // =0xeeed0000
-; CHECK-NEXT:    movk x0, #49152, lsl #32
-; CHECK-NEXT:    movk x0, #49151, lsl #48
+; CHECK-NEXT:    mov x0, #-4008509441 // =0xffffffff1112ffff
+; CHECK-NEXT:    eon x0, x0, x0, lsl #46
 ; CHECK-NEXT:    ret
   ret i64 u0xbfffc000eeed0000
 }
 
 define i64 @movn_0_eor_lsl_17() {
 ; CHECK-LABEL: movn_0_eor_lsl_17:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    mov x0, #4369 // =0x1111
-; CHECK-NEXT:    movk x0, #56797, lsl #16
-; CHECK-NEXT:    movk x0, #1, lsl #32
+; CHECK-NEXT:    mov x0, #-61167 // =0xffffffffffff1111
+; CHECK-NEXT:    eor x0, x0, x0, lsl #17
 ; CHECK-NEXT:    ret
   ret i64 u0x00000001dddd1111
 }
@@ -635,18 +630,16 @@ define i64 @movn_0_eor_lsl_47() {
 ; CHECK-LABEL: movn_0_eor_lsl_47:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    mov x0, #-61167 // =0xffffffffffff1111
-; CHECK-NEXT:    movk x0, #32767, lsl #32
-; CHECK-NEXT:    movk x0, #30583, lsl #48
+; CHECK-NEXT:    eor x0, x0, x0, lsl #47
 ; CHECK-NEXT:    ret
   ret i64 u0x77777fffffff1111
 }
 
 define i64 @movn_1_eor_lsl_17() {
 ; CHECK-LABEL: movn_1_eor_lsl_17:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    mov x0, #-286261249 // =0xffffffffeeefffff
-; CHECK-NEXT:    movk x0, #56796, lsl #32
-; CHECK-NEXT:    movk x0, #1, lsl #48
+; CHECK-NEXT:    mov x0, #-4008574977 // =0xffffffff1111ffff
+; CHECK-NEXT:    eor x0, x0, x0, lsl #17
 ; CHECK-NEXT:    ret
   ret i64 u0x0001dddceeefffff
 }
@@ -663,42 +656,38 @@ define i64 @movn_1_eor_lsl_32() {
 define i64 @movn_1_eor_lsl_46() {
 ; CHECK-LABEL: movn_1_eor_lsl_46:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    mov x0, #-4008509441
-; CHECK-NEXT:    movk x0, #16383, lsl #32
-; CHECK-NEXT:    movk x0, #16384, lsl #48
+; CHECK-NEXT:    mov x0, #-4008509441 // =0xffffffff1112ffff
+; CHECK-NEXT:    eor x0, x0, x0, lsl #46
 ; CHECK-NEXT:    ret
   ret i64 u0x40003fff1112ffff
 }
 
 define i64 @movn_movk_eon_lsl_17() {
 ; CHECK-LABEL: movn_movk_eon_lsl_17:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    mov x0, #43399 // =0xa987
-; CHECK-NEXT:    movk x0, #16699, lsl #16
-; CHECK-NEXT:    movk x0, #9320, lsl #32
-; CHECK-NEXT:    movk x0, #65534, lsl #48
+; CHECK-NEXT:    mov x0, #-43400 // =0xffffffffffff5678
+; CHECK-NEXT:    movk x0, #4660, lsl #16
+; CHECK-NEXT:    eon x0, x0, x0, lsl #17
 ; CHECK-NEXT:    ret
   ret i64 u0xfffe2468413ba987
 }
 
 define i64 @movn_movk_eon_lsl_47() {
 ; CHECK-LABEL: movn_movk_eon_lsl_47:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    mov x0, #61166 // =0xeeee
-; CHECK-NEXT:    movk x0, #21554, lsl #16
-; CHECK-NEXT:    movk x0, #32768, lsl #32
-; CHECK-NEXT:    movk x0, #34952, lsl #48
+; CHECK-NEXT:    mov x0, #-61167 // =0xffffffffffff1111
+; CHECK-NEXT:    movk x0, #43981, lsl #16
+; CHECK-NEXT:    eon x0, x0, x0, lsl #47
 ; CHECK-NEXT:    ret
   ret i64 u0x888880005432eeee
 }
 
 define i64 @movn_movk_eor_lsl_17() {
 ; CHECK-LABEL: movn_movk_eor_lsl_17:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    mov x0, #22136 // =0x5678
-; CHECK-NEXT:    movk x0, #48836, lsl #16
-; CHECK-NEXT:    movk x0, #56215, lsl #32
-; CHECK-NEXT:    movk x0, #1, lsl #48
+; CHECK-NEXT:    mov x0, #-43400 // =0xffffffffffff5678
+; CHECK-NEXT:    movk x0, #4660, lsl #16
+; CHECK-NEXT:    eor x0, x0, x0, lsl #17
 ; CHECK-NEXT:    ret
   ret i64 u0x0001db97bec45678
 }
@@ -716,10 +705,9 @@ define i64 @movn_movk_eor_lsl_32() {
 define i64 @movn_movk_eor_lsl_47() {
 ; CHECK-LABEL: movn_movk_eor_lsl_47:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    mov x0, #4369 // =0x1111
+; CHECK-NEXT:    mov x0, #-61167 // =0xffffffffffff1111
 ; CHECK-NEXT:    movk x0, #43981, lsl #16
-; CHECK-NEXT:    movk x0, #32767, lsl #32
-; CHECK-NEXT:    movk x0, #30583, lsl #48
+; CHECK-NEXT:    eor x0, x0, x0, lsl #47
 ; CHECK-NEXT:    ret
   ret i64 u0x77777fffabcd1111
 }
diff --git a/llvm/test/CodeGen/AArch64/srem-seteq-illegal-types.ll b/llvm/test/CodeGen/AArch64/srem-seteq-illegal-types.ll
@@ -61,14 +61,13 @@ define <3 x i1> @test_srem_vec(<3 x i33> %X) nounwind {
 ; CHECK-NEXT:    sbfx x8, x0, #0, #33
 ; CHECK-NEXT:    sbfx x10, x1, #0, #33
 ; CHECK-NEXT:    movk x9, #29127, lsl #16
-; CHECK-NEXT:    mov x13, #7281 // =0x1c71
+; CHECK-NEXT:    mov x13, #-7282 // =0xffffffffffffe38e
 ; CHECK-NEXT:    sbfx x12, x2, #0, #33
 ; CHECK-NEXT:    movk x9, #50972, lsl #32
-; CHECK-NEXT:    movk x13, #29127, lsl #16
+; CHECK-NEXT:    movk x13, #36408, lsl #16
 ; CHECK-NEXT:    movk x9, #7281, lsl #48
-; CHECK-NEXT:    movk x13, #50972, lsl #32
+; CHECK-NEXT:    eon x13, x13, x13, lsl #33
 ; CHECK-NEXT:    smulh x11, x8, x9
-; CHECK-NEXT:    movk x13, #7281, lsl #48
 ; CHECK-NEXT:    smulh x9, x10, x9
 ; CHECK-NEXT:    smulh x13, x12, x13
 ; CHECK-NEXT:    add x11, x11, x11, lsr #63
