[AArch64] Fixes ADD/SUB opt bug and abstracts shared behavior in MIPeepholeOpt for ADD, SUB, and AND.

[red1bluelost]

Summary

This fixes a bug where (SUBREG_TO_REG 0 (MOVi32imm <negative-number>) sub_32)
would generate invalid code since the top 32-bits were not zeroed when inspecting the
immediate value. A new test was added for this case.

Change to abstract shared behavior in MIPeepholeOpt. Both
visitAND and visitADDSUB attempt to split an RR instruction with an immediate
operand into two RI instructions with the immediate split.

The differing behavior lies in how the immediate is split into two pieces and
how the new instructions are built. The rest of the behavior (adding new VRegs,
checking for the MOVImm, constraining reg classes, removing old intructions)
are shared between the operations.

The new helper function splitTwoPartImm implements the shared behavior and
delegates differing behavior to two function objects passed by the caller.
One function object splits the immediate into two values and returns the
opcode to use if it is a valid split. The other function object builds
the new instructions.

I felt this abstraction would help since I believe it will help reduce the
code repetition when adding new instructions of the pattern, such as
SUBS for this conditional optimization.

Tested it locally by running check all with compiler-rt, mlir, clang-tools-extra,
flang, llvm, and clang enabled.

[dmgreen]

llvm/lib/Target/AArch64/AArch64MIPeepholeOpt.cpp

354–355

Do we have test case for this case with an ADD?

[red1bluelost]

Adds test case for edge case that happens when a MOVi32imm <negative-number> into SUBREG_TO_REG occurs.

[red1bluelost]

llvm/lib/Target/AArch64/AArch64MIPeepholeOpt.cpp

354–355

Turns out this was a very subtle but that I didn't catch in D117429. See example here: https://llvm.godbolt.org/z/Gnq1x8rh1

I added a test for this. We definitely want to have the Imm &= 0xFFFF'FFFF since assigning to T Imm will sign extend the value, not zero extend.

I guess in this case the revision is no longer NFC.

[dmgreen]

OK nice. I see. LGTM then

[red1bluelost]

Could you commit on my behalf when you have the chance. Same info: Micah Weston - [e-mail]
Thanks!

llvm/lib/Target/AArch64/AArch64MIPeepholeOpt.cpp

 //    ORRWrs, we can remove the ORRWrs because the upper 32 bits of the source
 //    operand are set to zero.
 //
 //===----------------------------------------------------------------------===//
 
 #include "AArch64ExpandImm.h"
 #include "AArch64InstrInfo.h"
 #include "MCTargetDesc/AArch64AddressingModes.h"
+#include "llvm/ADT/Optional.h"
 #include "llvm/ADT/SetVector.h"
 #include "llvm/CodeGen/MachineDominators.h"
 #include "llvm/CodeGen/MachineLoopInfo.h"
 
 using namespace llvm;
 
 #define DEBUG_TYPE "aarch64-mi-peephole-opt"
 
 namespace {
 
 struct AArch64MIPeepholeOpt : public MachineFunctionPass {
   static char ID;
 
   AArch64MIPeepholeOpt() : MachineFunctionPass(ID) {
     initializeAArch64MIPeepholeOptPass(*PassRegistry::getPassRegistry());
   }
 
   const AArch64InstrInfo *TII;
+  const AArch64RegisterInfo *TRI;
   MachineLoopInfo *MLI;
   MachineRegisterInfo *MRI;
 
+  template <typename T>
+  using SplitAndOpcFunc =
+      std::function<Optional<unsigned>(T, unsigned, T &, T &)>;
+  using BuildMIFunc =
+      std::function<void(MachineInstr &, unsigned, unsigned, unsigned, Register,
+                         Register, Register)>;
+
+  /// For instructions where an immediate operand could be split into two
+  /// separate immediate instructions, use the splitTwoPartImm two handle the
+  /// optimization.
+  ///
+  /// To implement, the following function types must be passed to
+  /// splitTwoPartImm. A SplitAndOpcFunc must be implemented that determines if
+  /// splitting the immediate is valid and returns the associated new opcode. A
+  /// BuildMIFunc must be implemented to build the two immediate instructions.
+  ///
+  /// Example Pattern (where IMM would require 2+ MOV instructions):
+  ///     %dst = <Instr>rr %src IMM [...]
+  /// becomes:
+  ///     %tmp = <Instr>ri %src (encode half IMM) [...]
+  ///     %dst = <Instr>ri %tmp (encode half IMM) [...]
+  template <typename T>
+  bool splitTwoPartImm(MachineInstr &MI,
+                       SmallSetVector<MachineInstr *, 8> &ToBeRemoved,
+                       SplitAndOpcFunc<T> SplitAndOpc, BuildMIFunc BuildInstr);
+
   bool checkMovImmInstr(MachineInstr &MI, MachineInstr *&MovMI,
                         MachineInstr *&SubregToRegMI);
 
   template <typename T>
-  bool visitADDSUB(MachineInstr &MI,
-                   SmallSetVector<MachineInstr *, 8> &ToBeRemoved, bool IsAdd);
+  bool visitADDSUB(unsigned PosOpc, unsigned NegOpc, MachineInstr &MI,
+                   SmallSetVector<MachineInstr *, 8> &ToBeRemoved);
   template <typename T>
-  bool visitAND(MachineInstr &MI,
+  bool visitAND(unsigned Opc, MachineInstr &MI,
                 SmallSetVector<MachineInstr *, 8> &ToBeRemoved);
   bool visitORR(MachineInstr &MI,
                 SmallSetVector<MachineInstr *, 8> &ToBeRemoved);
   bool runOnMachineFunction(MachineFunction &MF) override;
 
   StringRef getPassName() const override {
     return "AArch64 MI Peephole Optimization pass";
   }
static bool splitBitmaskImm(T Imm, unsigned RegSize, T &Imm1Enc, T &Imm2Enc) {
 
   Imm1Enc = AArch64_AM::encodeLogicalImmediate(NewImm1, RegSize);
   Imm2Enc = AArch64_AM::encodeLogicalImmediate(NewImm2, RegSize);
   return true;
 }
 
 template <typename T>
 bool AArch64MIPeepholeOpt::visitAND(
-    MachineInstr &MI, SmallSetVector<MachineInstr *, 8> &ToBeRemoved) {
+    unsigned Opc, MachineInstr &MI,
+    SmallSetVector<MachineInstr *, 8> &ToBeRemoved) {
   // Try below transformation.
   //
   // MOVi32imm + ANDWrr ==> ANDWri + ANDWri
   // MOVi64imm + ANDXrr ==> ANDXri + ANDXri
   //
   // The mov pseudo instruction could be expanded to multiple mov instructions
   // later. Let's try to split the constant operand of mov instruction into two
   // bitmask immediates. It makes only two AND instructions intead of multiple
   // mov + and instructions.
 
-  unsigned RegSize = sizeof(T) * 8;
-  assert((RegSize == 32 || RegSize == 64) &&
-         "Invalid RegSize for AND bitmask peephole optimization");
-
-  // Perform several essential checks against current MI.
-  MachineInstr *MovMI = nullptr, *SubregToRegMI = nullptr;
-  if (!checkMovImmInstr(MI, MovMI, SubregToRegMI))
-    return false;
-
-  // Split the bitmask immediate into two.
-  T UImm = static_cast<T>(MovMI->getOperand(1).getImm());
-  // For the 32 bit form of instruction, the upper 32 bits of the destination
-  // register are set to zero. If there is SUBREG_TO_REG, set the upper 32 bits
-  // of UImm to zero.
-  if (SubregToRegMI)
-    UImm &= 0xFFFFFFFF;
-  T Imm1Enc;
-  T Imm2Enc;
-  if (!splitBitmaskImm(UImm, RegSize, Imm1Enc, Imm2Enc))
-    return false;
-
-  // Create new AND MIs.
+  return splitTwoPartImm<T>(
+      MI, ToBeRemoved,
+      [Opc](T Imm, unsigned RegSize, T &Imm0, T &Imm1) -> Optional<unsigned> {
+        if (splitBitmaskImm(Imm, RegSize, Imm0, Imm1))
+          return Opc;
+        return None;
+      },
+      [&TII = TII](MachineInstr &MI, unsigned Opcode, unsigned Imm0,
+                   unsigned Imm1, Register SrcReg, Register NewTmpReg,
+                   Register NewDstReg) {
         DebugLoc DL = MI.getDebugLoc();
         MachineBasicBlock *MBB = MI.getParent();
-  const TargetRegisterClass *ANDImmRC =
-      (RegSize == 32) ? &AArch64::GPR32spRegClass : &AArch64::GPR64spRegClass;
-  Register DstReg = MI.getOperand(0).getReg();
-  Register SrcReg = MI.getOperand(1).getReg();
-  Register NewTmpReg = MRI->createVirtualRegister(ANDImmRC);
-  Register NewDstReg = MRI->createVirtualRegister(ANDImmRC);
-  unsigned Opcode = (RegSize == 32) ? AArch64::ANDWri : AArch64::ANDXri;
-
-  MRI->constrainRegClass(NewTmpReg, MRI->getRegClass(SrcReg));
         BuildMI(*MBB, MI, DL, TII->get(Opcode), NewTmpReg)
             .addReg(SrcReg)
-      .addImm(Imm1Enc);
-
-  MRI->constrainRegClass(NewDstReg, MRI->getRegClass(DstReg));
+            .addImm(Imm0);
         BuildMI(*MBB, MI, DL, TII->get(Opcode), NewDstReg)
             .addReg(NewTmpReg)
-      .addImm(Imm2Enc);
-
-  MRI->replaceRegWith(DstReg, NewDstReg);
-  // replaceRegWith changes MI's definition register. Keep it for SSA form until
-  // deleting MI.
-  MI.getOperand(0).setReg(DstReg);
-
-  ToBeRemoved.insert(&MI);
-  if (SubregToRegMI)
-    ToBeRemoved.insert(SubregToRegMI);
-  ToBeRemoved.insert(MovMI);
-
-  return true;
+            .addImm(Imm1);
+      });
 }
 
 bool AArch64MIPeepholeOpt::visitORR(
     MachineInstr &MI, SmallSetVector<MachineInstr *, 8> &ToBeRemoved) {
   // Check this ORR comes from below zero-extend pattern.
   //
   // def : Pat<(i64 (zext GPR32:$src)),
   //           (SUBREG_TO_REG (i32 0), (ORRWrs WZR, GPR32:$src, 0), sub_32)>;
static bool splitAddSubImm(T Imm, unsigned RegSize, T &Imm0, T &Imm1) {
   // Split Imm into (Imm0 << 12) + Imm1;
   Imm0 = (Imm >> 12) & 0xfff;
   Imm1 = Imm & 0xfff;
   return true;
 }
 
 template <typename T>
 bool AArch64MIPeepholeOpt::visitADDSUB(
-    MachineInstr &MI, SmallSetVector<MachineInstr *, 8> &ToBeRemoved,
-    bool IsAdd) {
+    unsigned PosOpc, unsigned NegOpc, MachineInstr &MI,
+    SmallSetVector<MachineInstr *, 8> &ToBeRemoved) {
   // Try below transformation.
   //
   // MOVi32imm + ADDWrr ==> ADDWri + ADDWri
   // MOVi64imm + ADDXrr ==> ADDXri + ADDXri
   //
   // MOVi32imm + SUBWrr ==> SUBWri + SUBWri
   // MOVi64imm + SUBXrr ==> SUBXri + SUBXri
   //
   // The mov pseudo instruction could be expanded to multiple mov instructions
   // later. Let's try to split the constant operand of mov instruction into two
   // legal add/sub immediates. It makes only two ADD/SUB instructions intead of
   // multiple `mov` + `and/sub` instructions.
 
-  unsigned RegSize = sizeof(T) * 8;
-  assert((RegSize == 32 || RegSize == 64) &&
-         "Invalid RegSize for legal add/sub immediate peephole optimization");
-
-  // Perform several essential checks against current MI.
-  MachineInstr *MovMI, *SubregToRegMI;
-  if (!checkMovImmInstr(MI, MovMI, SubregToRegMI))
-    return false;
-
-  // Split the immediate to Imm0 and Imm1, and calculate the Opcode.
-  T Imm = static_cast<T>(MovMI->getOperand(1).getImm()), Imm0, Imm1;
-  unsigned Opcode;
-  if (splitAddSubImm(Imm, RegSize, Imm0, Imm1)) {
-    if (IsAdd)
-      Opcode = RegSize == 32 ? AArch64::ADDWri : AArch64::ADDXri;
-    else
-      Opcode = RegSize == 32 ? AArch64::SUBWri : AArch64::SUBXri;
-  } else if (splitAddSubImm(-Imm, RegSize, Imm0, Imm1)) {
-    if (IsAdd)
-      Opcode = RegSize == 32 ? AArch64::SUBWri : AArch64::SUBXri;
-    else
-      Opcode = RegSize == 32 ? AArch64::ADDWri : AArch64::ADDXri;
-  } else {
-    return false;
-  }
-
-  // Create new ADD/SUB MIs.
+  return splitTwoPartImm<T>(
+      MI, ToBeRemoved,
+      [PosOpc, NegOpc](T Imm, unsigned RegSize, T &Imm0,
+                       T &Imm1) -> Optional<unsigned> {
+        if (splitAddSubImm(Imm, RegSize, Imm0, Imm1))
+          return PosOpc;
+        if (splitAddSubImm(-Imm, RegSize, Imm0, Imm1))
+          return NegOpc;
+        return None;
+      },
+      [&TII = TII](MachineInstr &MI, unsigned Opcode, unsigned Imm0,
+                   unsigned Imm1, Register SrcReg, Register NewTmpReg,
+                   Register NewDstReg) {
         DebugLoc DL = MI.getDebugLoc();
         MachineBasicBlock *MBB = MI.getParent();
-  const TargetRegisterClass *RC =
-      (RegSize == 32) ? &AArch64::GPR32spRegClass : &AArch64::GPR64spRegClass;
-  Register DstReg = MI.getOperand(0).getReg();
-  Register SrcReg = MI.getOperand(1).getReg();
-  Register NewTmpReg = MRI->createVirtualRegister(RC);
-  Register NewDstReg = MRI->createVirtualRegister(RC);
-
-  MRI->constrainRegClass(SrcReg, RC);
         BuildMI(*MBB, MI, DL, TII->get(Opcode), NewTmpReg)
             .addReg(SrcReg)
             .addImm(Imm0)
             .addImm(12);
-
-  MRI->constrainRegClass(NewDstReg, MRI->getRegClass(DstReg));
         BuildMI(*MBB, MI, DL, TII->get(Opcode), NewDstReg)
             .addReg(NewTmpReg)
             .addImm(Imm1)
             .addImm(0);
-
-  MRI->replaceRegWith(DstReg, NewDstReg);
-  // replaceRegWith changes MI's definition register. Keep it for SSA form until
-  // deleting MI.
-  MI.getOperand(0).setReg(DstReg);
-
-  // Record the MIs need to be removed.
-  ToBeRemoved.insert(&MI);
-  if (SubregToRegMI)
-    ToBeRemoved.insert(SubregToRegMI);
-  ToBeRemoved.insert(MovMI);
-
-  return true;
+      });
 }
 
 // Checks if the corresponding MOV immediate instruction is applicable for
 // this peephole optimization.
 bool AArch64MIPeepholeOpt::checkMovImmInstr(MachineInstr &MI,
                                             MachineInstr *&MovMI,
                                             MachineInstr *&SubregToRegMI) {
   // Check whether current MBB is in loop and the AND is loop invariant.
if (!MRI->hasOneUse(MovMI->getOperand(0).getReg()))
     return false;
   if (SubregToRegMI && !MRI->hasOneUse(SubregToRegMI->getOperand(0).getReg()))
     return false;
 
   // It is OK to perform this peephole optimization.
   return true;
 }
 
+template <typename T>
+bool AArch64MIPeepholeOpt::splitTwoPartImm(
+    MachineInstr &MI, SmallSetVector<MachineInstr *, 8> &ToBeRemoved,
+    SplitAndOpcFunc<T> SplitAndOpc, BuildMIFunc BuildInstr) {
+  unsigned RegSize = sizeof(T) * 8;
+  assert((RegSize == 32 || RegSize == 64) &&
+         "Invalid RegSize for legal immediate peephole optimization");
+
+  // Perform several essential checks against current MI.
+  MachineInstr *MovMI, *SubregToRegMI;
+  if (!checkMovImmInstr(MI, MovMI, SubregToRegMI))
+    return false;
+
+  // Split the immediate to Imm0 and Imm1, and calculate the Opcode.
+  T Imm = static_cast<T>(MovMI->getOperand(1).getImm()), Imm0, Imm1;
+  // For the 32 bit form of instruction, the upper 32 bits of the destination
+  // register are set to zero. If there is SUBREG_TO_REG, set the upper 32 bits
+  // of Imm to zero. This is essential if the Immediate value was a negative
+  // number since it was sign extended when we assign to the 64-bit Imm.
+  if (SubregToRegMI)
+    Imm &= 0xFFFFFFFF;
+  unsigned Opcode;
+  if (auto R = SplitAndOpc(Imm, RegSize, Imm0, Imm1))
+    Opcode = R.getValue();
+  else
+    return false;
+
+  // Create new ADD/SUB MIs.
+  MachineFunction *MF = MI.getMF();
+  const TargetRegisterClass *RC =
+      TII->getRegClass(TII->get(Opcode), 0, TRI, *MF);
+  const TargetRegisterClass *ORC =
+      TII->getRegClass(TII->get(Opcode), 1, TRI, *MF);
+  Register DstReg = MI.getOperand(0).getReg();
+  Register SrcReg = MI.getOperand(1).getReg();
+  Register NewTmpReg = MRI->createVirtualRegister(RC);
+  Register NewDstReg = MRI->createVirtualRegister(RC);
+
+  MRI->constrainRegClass(SrcReg, RC);
+  MRI->constrainRegClass(NewTmpReg, ORC);
+  MRI->constrainRegClass(NewDstReg, MRI->getRegClass(DstReg));
+
+  BuildInstr(MI, Opcode, Imm0, Imm1, SrcReg, NewTmpReg, NewDstReg);
+
+  MRI->replaceRegWith(DstReg, NewDstReg);
+  // replaceRegWith changes MI's definition register. Keep it for SSA form until
+  // deleting MI.
+  MI.getOperand(0).setReg(DstReg);
+
+  // Record the MIs need to be removed.
+  ToBeRemoved.insert(&MI);
+  if (SubregToRegMI)
+    ToBeRemoved.insert(SubregToRegMI);
+  ToBeRemoved.insert(MovMI);
+
+  return true;
+}
+
 bool AArch64MIPeepholeOpt::runOnMachineFunction(MachineFunction &MF) {
   if (skipFunction(MF.getFunction()))
     return false;
 
   TII = static_cast<const AArch64InstrInfo *>(MF.getSubtarget().getInstrInfo());
+  TRI = static_cast<const AArch64RegisterInfo *>(
+      MF.getSubtarget().getRegisterInfo());
   MLI = &getAnalysis<MachineLoopInfo>();
   MRI = &MF.getRegInfo();
 
-  assert (MRI->isSSA() && "Expected to be run on SSA form!");
+  assert(MRI->isSSA() && "Expected to be run on SSA form!");
 
   bool Changed = false;
   SmallSetVector<MachineInstr *, 8> ToBeRemoved;
 
   for (MachineBasicBlock &MBB : MF) {
     for (MachineInstr &MI : MBB) {
       switch (MI.getOpcode()) {
       default:
         break;
       case AArch64::ANDWrr:
-        Changed = visitAND<uint32_t>(MI, ToBeRemoved);
+        Changed = visitAND<uint32_t>(AArch64::ANDWri, MI, ToBeRemoved);
         break;
       case AArch64::ANDXrr:
-        Changed = visitAND<uint64_t>(MI, ToBeRemoved);
+        Changed = visitAND<uint64_t>(AArch64::ANDXri, MI, ToBeRemoved);
         break;
       case AArch64::ORRWrs:
         Changed = visitORR(MI, ToBeRemoved);
         break;
       case AArch64::ADDWrr:
-        Changed = visitADDSUB<uint32_t>(MI, ToBeRemoved, true);
+        Changed = visitADDSUB<uint32_t>(AArch64::ADDWri, AArch64::SUBWri, MI,
+                                        ToBeRemoved);
         break;
       case AArch64::SUBWrr:
-        Changed = visitADDSUB<uint32_t>(MI, ToBeRemoved, false);
+        Changed = visitADDSUB<uint32_t>(AArch64::SUBWri, AArch64::ADDWri, MI,
+                                        ToBeRemoved);
         break;
       case AArch64::ADDXrr:
-        Changed = visitADDSUB<uint64_t>(MI, ToBeRemoved, true);
+        Changed = visitADDSUB<uint64_t>(AArch64::ADDXri, AArch64::SUBXri, MI,
+                                        ToBeRemoved);
         break;
       case AArch64::SUBXrr:
-        Changed = visitADDSUB<uint64_t>(MI, ToBeRemoved, false);
+        Changed = visitADDSUB<uint64_t>(AArch64::SUBXri, AArch64::ADDXri, MI,
+                                        ToBeRemoved);
         break;
       }
     }
   }
 
   for (MachineInstr *MI : ToBeRemoved)
     MI->eraseFromParent();
 
   return Changed;
 }
 
 FunctionPass *llvm::createAArch64MIPeepholeOptPass() {
   return new AArch64MIPeepholeOpt();
 }

llvm/test/CodeGen/AArch64/addsub.ll

; CHECK-NEXT:    ret
   %a0 = call {i8, i1} @llvm.uadd.with.overflow.i8(i8 %nota, i8 %b)
   %e0 = extractvalue {i8, i1} %a0, 0
   %e1 = extractvalue {i8, i1} %a0, 1
   %res = add i8 %e0, 1
   store i8 %res, i8* %p
   ret i1 %e1
 }
 
+; This is a unique edge case that will generate the following MIR
+;   MOVi32imm -1000000
+;   SUBREG_TO_REG 0, killed %1, %subreg.sub_32
+; When using a 64-bit unsigned for the "-1000000" immediate, the code
+; must make sure to zero out the top 32 bits since SUBREG_TO_REG is
+; zero extending the value
+define i64 @addl_0x80000000(i64 %a) {
+; CHECK-LABEL: addl_0x80000000:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    mov w8, #48576
+; CHECK-NEXT:    movk w8, #65520, lsl #16
+; CHECK-NEXT:    add x0, x0, x8
+; CHECK-NEXT:    ret
+  %b = add i64 %a, 4293967296
+  ret i64 %b
+}
+
 ; TODO: adds/subs