Teach the AArch64 backend to materialize immediates using a pair of ORR-immediate instructions.

[resistor]

Summary

Credit to czwarich for figuring out the algorithm to test for this.

[fhahn]

Credit to czwarich for figuring out the algorithm to test for this.

Would be good to add a description.

[HsiangKai]

llvm/test/CodeGen/AArch64/arm64-movi.ll

321–322

Remove the FIXME.

390–391

Ditto.

[resistor]

Refactor and add more code comments. Update test comments.

Updating D140952: Teach the AArch64 backend to materialize immediates using a pair of ORR-immediate

instructions.

[resistor]

In D140952#4025370, @fhahn wrote:

Credit to czwarich for figuring out the algorithm to test for this.

Would be good to add a description.

I have attempted to document the approach very verbosely in the code. Are there specific parts where more detail would be helpful?

[resistor]

Ping

[resistor]

Ping

[dmgreen]

Hello. Feel free to add reviewers to a patch, otherwise people might not see it or think it is still WIP. The patch looks like a nice addition.

Can you clang-format the patch? I don't believe that ifs and returns should be on the same line.

There is another use of expandMOVImm in the MachineCombiner. Is that OK because it only generated single instructions? It should possibly just be generating a MOVi32imm/MOVi64imm instead.

llvm/lib/Target/AArch64/AArch64ExpandImm.cpp

308

Remove can

345

"as an AND of an OR"? or a maybe explain the inverse of the immediate thing. I can see why this would work, but find it difficult to follow - it would be good to add a longer comment explaining how it works.

llvm/lib/Target/AArch64/AArch64ExpandPseudoInsts.cpp

169

Can this happen?

[resistor]

Update for review comments

Updating D140952: Teach the AArch64 backend to materialize immediates using a pair of ORR-immediate

instructions.

llvm/lib/Target/AArch64/AArch64ExpandImm.cpp

345

The ORR is a red herring. The initial ORR is just a move-immediate, which the following AND then uses as its LHS.

llvm/lib/Target/AArch64/AArch64ExpandPseudoInsts.cpp

169

I'm not aware that anything creates it today, but I thought it was better to have the API of be orthogonal, so as to minimize surprise for any future folks touching this area.

[efriedma]

FYI, there's an incomplete patch D78981 which has misc two-instruction patterns for immediates, if you're looking in the area. I found when I was looking at it that the more obscure patterns trigger quite rarely in practice in the codebases I was looking at, though, so I never got around to completing it.

[dmgreen]

Thanks. There are a couple of suggestions below, otherwise LGTM.

llvm/lib/Target/AArch64/AArch64ExpandImm.cpp

345

Yeah. Its not really "as a pair of ANDs of logical immediates". It might be better to say " as an AND of a pair of logical immediates"

llvm/lib/Target/AArch64/AArch64ExpandPseudoInsts.cpp

169

An AND of 0 would just be a way of creating 0, and an AND would be a strange way of creating it. The Or with a WZR is the canonical form of a mov, but the And sounds like it would make less sense and might be a bug if it did come up. Perhaps change it to an assert that I->Op1 != 0 instead.

[resistor]

llvm/lib/Target/AArch64/AArch64ExpandPseudoInsts.cpp

169

That's not what's happening here. The 0 isn't in the immediate operand, but in the register number. It's a signaling mechanism that is used to tell the AArch64ExpandPseudos pass what to do.

[dmgreen]

llvm/lib/Target/AArch64/AArch64ExpandPseudoInsts.cpp

169

But XZR is zero, so this is doing 0 & imm, which is always zero. So the result will always just be an awkward way of writing movi 0. It's not a big deal either way, but if we want a zero we should be using the more canonical ORRWrs that processors will recognize and optimize for. But if anything gets here it is likely a bug in AArch64_IMM::expandMOVImm.

llvm/lib/Target/AArch64/AArch64ExpandImm.cpp

static bool trySequenceOfOnes(uint64_t UImm,
   // Create the second MOVK instruction.
   Insn.push_back({ AArch64::MOVKXi, getChunk(UImm, SecondMovkIdx),
              AArch64_AM::getShifterImm(AArch64_AM::LSL,
                                              SecondMovkIdx * 16) });
 
   return true;
 }
 
+static uint64_t GetRunOfOnesStartingAt(uint64_t V, uint64_t StartPosition) {
+  uint64_t NumOnes = llvm::countTrailingOnes(V >> StartPosition);
+
+  uint64_t UnshiftedOnes;
+  if (NumOnes == 64) {
+    UnshiftedOnes = ~0ULL;
+  } else {
+    UnshiftedOnes = (1ULL << NumOnes) - 1;
+  }
+  return UnshiftedOnes << StartPosition;
+}
+
+static uint64_t rotl(uint64_t n, uint64_t d) {
+  if (d == 0)
+    return n;
+  return (n << d) | (n >> (64 - d));
+}
+
+static uint64_t rotr(uint64_t n, uint64_t d) {
+  if (d == 0)
+    return n;
+  return (n >> d) | (n << (64 - d));
+}
+
+static uint64_t MaximallyReplicateSubImmediate(uint64_t V, uint64_t Subset) {
+  uint64_t Result = Subset;
+
+  // 64, 32, 16, 8, 4, 2
+  for (uint64_t i = 0; i < 6; ++i) {
+    uint64_t Rotation = 1 << (6 - i);
+    uint64_t Closure = Result | rotl(Result, Rotation);
+    if (Closure != (Closure & V)) {
+      break;
+    }
+    Result = Closure;
+  }
+
+  return Result;
+}
+
+// Find the logical immediate that covers the most bits in RemainingBits,
+// allowing for additional bits to be set that were set in OriginalBits.
+static uint64_t maximalLogicalImmWithin(uint64_t RemainingBits,
+                                        uint64_t OriginalBits) {
+  // Find the first set bit.
+  uint32_t Position = llvm::countTrailingZeros(RemainingBits);
+
+  // Get the first run of set bits.
+  uint64_t FirstRun = GetRunOfOnesStartingAt(OriginalBits, Position);
+
+  // Replicate the run as many times as possible, as long as the bits are set in
+  // RemainingBits.
+  uint64_t MaximalImm = MaximallyReplicateSubImmediate(OriginalBits, FirstRun);
+
+  return MaximalImm;
+}
+
+static std::optional<std::pair<uint64_t, uint64_t>>
+decomposeIntoOrrOfLogicalImmediates(uint64_t UImm) {
+  if (UImm == 0 || ~UImm == 0)
+    return std::nullopt;
+
+  // Make sure we don't have a run of ones split around the rotation boundary.
+  uint32_t InitialTrailingOnes = llvm::countTrailingOnes(UImm);
+  uint64_t RotatedBits = rotr(UImm, InitialTrailingOnes);
+
+  // Find the largest logical immediate that fits within the full immediate.
+  uint64_t MaximalImm1 = maximalLogicalImmWithin(RotatedBits, RotatedBits);
+
+  // Remove all bits that are set by this mask.
+  uint64_t RemainingBits = RotatedBits & ~MaximalImm1;
+
+  // Find the largest logical immediate covering the remaining bits, allowing
+  // for additional bits to be set that were also set in the original immediate.
+  uint64_t MaximalImm2 = maximalLogicalImmWithin(RemainingBits, RotatedBits);
+
+  // If any bits still haven't been covered, then give up.
+  if (RemainingBits & ~MaximalImm2)
+    return std::nullopt;
+
+  // Make sure to un-rotate the immediates.
+  return std::make_pair(rotl(MaximalImm1, InitialTrailingOnes),
+                        rotl(MaximalImm2, InitialTrailingOnes));
+}
+
+// Attempt to expand an immediate as the ORR of a pair of logical immediates.
+static bool tryOrrOfLogicalImmediates(uint64_t UImm,
+                                      SmallVectorImpl<ImmInsnModel> &Insn) {
+  auto MaybeDecomposition = decomposeIntoOrrOfLogicalImmediates(UImm);
+  if (MaybeDecomposition == std::nullopt)
+    return false;
+  uint64_t Imm1 = MaybeDecomposition->first;
+  uint64_t Imm2 = MaybeDecomposition->second;
+
+  uint64_t Encoding1, Encoding2;
+  bool Imm1Success = AArch64_AM::processLogicalImmediate(Imm1, 64, Encoding1);
+  bool Imm2Success = AArch64_AM::processLogicalImmediate(Imm2, 64, Encoding2);
+
+  if (Imm1Success && Imm2Success) {
+    // Create the ORR-immediate instructions.
+    Insn.push_back({AArch64::ORRXri, 0, Encoding1});
+    Insn.push_back({AArch64::ORRXri, 1, Encoding2});
+    return true;
+  }
+
+  return false;
+}
+
+// Attempt to expand an immediate as the AND of a pair of logical immediates.
+// This is done by applying DeMorgan's law, under which logical immediates
+// are closed.
+static bool tryAndOfLogicalImmediates(uint64_t UImm,
+                                      SmallVectorImpl<ImmInsnModel> &Insn) {
+  // Apply DeMorgan's law to turn this into an ORR problem.
+  auto MaybeDecomposition = decomposeIntoOrrOfLogicalImmediates(~UImm);
+  if (MaybeDecomposition == std::nullopt)
+    return false;
+  uint64_t Imm1 = MaybeDecomposition->first;
+  uint64_t Imm2 = MaybeDecomposition->second;
+
+  uint64_t Encoding1, Encoding2;
+  bool Imm1Success = AArch64_AM::processLogicalImmediate(~Imm1, 64, Encoding1);
+  bool Imm2Success = AArch64_AM::processLogicalImmediate(~Imm2, 64, Encoding2);
+
+  if (Imm1Success && Imm2Success) {
+    // Materialize Imm1, the LHS of the AND
+    Insn.push_back({AArch64::ORRXri, 0, Encoding1});
+    // AND Imm1 with Imm2
+    Insn.push_back({AArch64::ANDXri, 1, Encoding2});
+    return true;
+  }
+
+  return false;
+}
+
 /// \brief Expand a MOVi32imm or MOVi64imm pseudo instruction to a
 /// MOVZ or MOVN of width BitSize followed by up to 3 MOVK instructions.
 static inline void expandMOVImmSimple(uint64_t Imm, unsigned BitSize,
               unsigned OneChunks, unsigned ZeroChunks,
               SmallVectorImpl<ImmInsnModel> &Insn) {
   const unsigned Mask = 0xFFFF;
 
   // Use a MOVZ or MOVN instruction to set the high bits, followed by one or
if (AArch64_AM::processLogicalImmediate(ZeroChunk, BitSize, Encoding) ||
       // Create the MOVK instruction.
       const unsigned Imm16 = getChunk(UImm, Shift / 16);
       Insn.push_back({ AArch64::MOVKXi, Imm16,
            AArch64_AM::getShifterImm(AArch64_AM::LSL, Shift) });
       return;
     }
   }
 
+  // Attempt to use a sequence of two ORR-immediate instructions.
+  if (tryOrrOfLogicalImmediates(Imm, Insn))
+    return;
+
+  // Attempt to use a sequence of ORR-immediate followed by AND-immediate.
+  if (tryAndOfLogicalImmediates(Imm, Insn))
+    return;
+
   // FIXME: Add more two-instruction sequences.
 
   // Three instruction sequences.
   //
   // Prefer MOVZ/MOVN followed by two MOVK; it's more readable, and possibly
   // the fastest sequence with fast literal generation. (If neither MOVK is
   // part of a fast literal generation pair, it could be slower than the
   // four-instruction sequence, but we won't worry about that for now.)

llvm/lib/Target/AArch64/AArch64ExpandPseudoInsts.cpp

bool AArch64ExpandPseudo::expandMOVImm(MachineBasicBlock &MBB,
   for (auto I = Insn.begin(), E = Insn.end(); I != E; ++I) {
     bool LastItem = std::next(I) == E;
     switch (I->Opcode)
     {
     default: llvm_unreachable("unhandled!"); break;
 
     case AArch64::ORRWri:
     case AArch64::ORRXri:
+      if (I->Op1 == 0) {
         MIBS.push_back(BuildMI(MBB, MBBI, MI.getDebugLoc(), TII->get(I->Opcode))
                            .add(MI.getOperand(0))
                            .addReg(BitSize == 32 ? AArch64::WZR : AArch64::XZR)
                            .addImm(I->Op2));
+      } else {
+        Register DstReg = MI.getOperand(0).getReg();
+        bool DstIsDead = MI.getOperand(0).isDead();
+        MIBS.push_back(
+            BuildMI(MBB, MBBI, MI.getDebugLoc(), TII->get(I->Opcode))
+                .addReg(DstReg, RegState::Define |
+                                    getDeadRegState(DstIsDead && LastItem) |
+                                    RenamableState)
+                .addReg(DstReg)
+                .addImm(I->Op2));
+      }
+      break;
+    case AArch64::ANDXri:
+      if (I->Op1 == 0) {
+        MIBS.push_back(BuildMI(MBB, MBBI, MI.getDebugLoc(), TII->get(I->Opcode))
+                           .add(MI.getOperand(0))
+                           .addReg(BitSize == 32 ? AArch64::WZR : AArch64::XZR)
+                           .addImm(I->Op2));
+      } else {
+        Register DstReg = MI.getOperand(0).getReg();
+        bool DstIsDead = MI.getOperand(0).isDead();
+        MIBS.push_back(
+            BuildMI(MBB, MBBI, MI.getDebugLoc(), TII->get(I->Opcode))
+                .addReg(DstReg, RegState::Define |
+                                    getDeadRegState(DstIsDead && LastItem) |
+                                    RenamableState)
+                .addReg(DstReg)
+                .addImm(I->Op2));
+      }
       break;
     case AArch64::MOVNWi:
     case AArch64::MOVNXi:
     case AArch64::MOVZWi:
     case AArch64::MOVZXi: {
       bool DstIsDead = MI.getOperand(0).isDead();
       MIBS.push_back(BuildMI(MBB, MBBI, MI.getDebugLoc(), TII->get(I->Opcode))
         .addReg(DstReg, RegState::Define |

llvm/test/CodeGen/AArch64/arm64-movi.ll

 ; CHECK-LABEL: movz_skip1_movk:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    mov x0, #4660
 ; CHECK-NEXT:    movk x0, #34388, lsl #32
 ; CHECK-NEXT:    ret
   ret i64 147695335379508
 }
 
-; FIXME: Prefer "mov w0, #2863311530; lsl x0, x0, #4"
-; or "mov x0, #-6148914691236517206; and x0, x0, #45812984480"
 define i64 @orr_lsl_pattern() nounwind {
 ; CHECK-LABEL: orr_lsl_pattern:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    mov x0, #43680
-; CHECK-NEXT:    movk x0, #43690, lsl #16
-; CHECK-NEXT:    movk x0, #10, lsl #32
+; CHECK-NEXT:    mov x0, #-6148914691236517206
+; CHECK-NEXT:    and x0, x0, #0x1fffffffe0
 ; CHECK-NEXT:    ret
   ret i64 45812984480
 }
 
 ; FIXME: prefer "mov x0, #-16639; lsl x0, x0, #24"
 define i64 @mvn_lsl_pattern() nounwind {
 ; CHECK-LABEL: mvn_lsl_pattern:
 ; CHECK:       // %bb.0:

 define i64 @orr_movk15() nounwind {
 ; CHECK-LABEL: orr_movk15:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    mov x0, #549755813887
 ; CHECK-NEXT:    movk x0, #63487, lsl #16
 ; CHECK-NEXT:    ret
   ret i64 549621596159
 }
 
-; FIXME: prefer "mov x0, #2147483646; orr x0, x0, #36028659580010496"
 define i64 @orr_movk16() nounwind {
 ; CHECK-LABEL: orr_movk16:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    mov x0, #36028659580010496
-; CHECK-NEXT:    movk x0, #65534
-; CHECK-NEXT:    movk x0, #32767, lsl #16
+; CHECK-NEXT:    mov x0, #2147483646
+; CHECK-NEXT:    orr x0, x0, #0x7fffe0007fffe0
 ; CHECK-NEXT:    ret
   ret i64 36028661727494142
 }
 
 define i64 @orr_movk17() nounwind {
 ; CHECK-LABEL: orr_movk17:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    mov x0, #-1099511627776
 ; CHECK-NEXT:    movk x0, #65280, lsl #16
 ; CHECK-NEXT:    ret
   ret i64 -1095233437696
 }
 
 define i64 @orr_movk18() nounwind {
 ; CHECK-LABEL: orr_movk18:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    mov x0, #137438887936
 ; CHECK-NEXT:    movk x0, #65473
 ; CHECK-NEXT:    ret
   ret i64 137438953409
 }
 
-; FIXME: prefer "mov x0, #72340172838076673; and x0, x0, #2199023255296"
 define i64 @orr_and() nounwind {
 ; CHECK-LABEL: orr_and:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    mov x0, #256
-; CHECK-NEXT:    movk x0, #257, lsl #16
-; CHECK-NEXT:    movk x0, #257, lsl #32
+; CHECK-NEXT:    mov x0, #72340172838076673
+; CHECK-NEXT:    and x0, x0, #0xffffffffff00
 ; CHECK-NEXT:    ret
   ret i64 1103823438080
 }
 
 ; FIXME: prefer "mov w0, #-1431655766; movk x0, #9, lsl #32"
 define i64 @movn_movk() nounwind {
 ; CHECK-LABEL: movn_movk:
 ; CHECK:       // %bb.0:

 ; CHECK-LABEL: movn_eor:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    mov x0, #3689348814741910323
 ; CHECK-NEXT:    movk x0, #52428
 ; CHECK-NEXT:    movk x0, #8455, lsl #16
 ; CHECK-NEXT:    ret
   ret i64 3689348814437076172
 }
 
-; FIXME: prefer "mov x0, #536866816; orr x0, x0, #0x3fff800000000000"
 define i64 @orr_orr_64() nounwind {
 ; CHECK-LABEL: orr_orr_64:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    mov x0, #4611545280939032576
-; CHECK-NEXT:    movk x0, #61440
-; CHECK-NEXT:    movk x0, #8191, lsl #16
+; CHECK-NEXT:    mov x0, #536866816
+; CHECK-NEXT:    orr x0, x0, #0x3fff800000000000
 ; CHECK-NEXT:    ret
   ret i64 4611545281475899392
 }
 
-; FIXME: prefer "mov x0, #558551907040256; orr x0, x0, #0x1000100010001000"
 define i64 @orr_orr_32() nounwind {
 ; CHECK-LABEL: orr_orr_32:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    mov x0, #-287953294993589248
-; CHECK-NEXT:    movk x0, #7169, lsl #16
-; CHECK-NEXT:    movk x0, #7169, lsl #48
+; CHECK-NEXT:    mov x0, #558551907040256
+; CHECK-NEXT:    orr x0, x0, #0x1c001c001c001c00
 ; CHECK-NEXT:    ret
   ret i64 2018171185438784512
 }
 
-; FIXME: prefer "mov x0, #281479271743489; orr x0, x0, #0x1000100010001000"
 define i64 @orr_orr_16() nounwind {
 ; CHECK-LABEL: orr_orr_16:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    mov x0, #4097
-; CHECK-NEXT:    movk x0, #4097, lsl #16
-; CHECK-NEXT:    movk x0, #4097, lsl #32
-; CHECK-NEXT:    movk x0, #4097, lsl #48
+; CHECK-NEXT:    mov x0, #1152939097061330944
+; CHECK-NEXT:    orr x0, x0, #0x1000100010001
 ; CHECK-NEXT:    ret
   ret i64 1153220576333074433
 }
 
-; FIXME: prefer "mov x0, #144680345676153346; orr x0, x0, #0x1818181818181818"
 define i64 @orr_orr_8() nounwind {
 ; CHECK-LABEL: orr_orr_8:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    mov x0, #6682
-; CHECK-NEXT:    movk x0, #6682, lsl #16
-; CHECK-NEXT:    movk x0, #6682, lsl #32
-; CHECK-NEXT:    movk x0, #6682, lsl #48
+; CHECK-NEXT:    mov x0, #144680345676153346
+; CHECK-NEXT:    orr x0, x0, #0x1818181818181818
 ; CHECK-NEXT:    ret
   ret i64 1880844493789993498
 }
 
-; FIXME: prefer "mov x0, #-6148914691236517206; orr x0, x0, #0x0FFFFF0000000000"
 define i64 @orr_64_orr_8() nounwind {
 ; CHECK-LABEL: orr_64_orr_8:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    mov x0, #-6148914691236517206
-; CHECK-NEXT:    movk x0, #65450, lsl #32
-; CHECK-NEXT:    movk x0, #45055, lsl #48
+; CHECK-NEXT:    orr x0, x0, #0xfffff0000000000
 ; CHECK-NEXT:    ret
   ret i64 -5764607889538110806
 }