[AArch64] Remove ToBeRemoved from AArch64MIPeepholeOpt

[dmgreen]

Summary

The ToBeRemoved is used to remove any MachineInstructions that are no longer needed, making sure we don't invalidate the iterator that is currently in use by erasing the instruction straight away. This makes issues for keeping the code in SSA from though, where subsequent transforms that require SSA form may have been broken by previous.

If, instead, we use make_early_inc_range the iteration issue shouldn't be present, so long as we do not remove the subsequent instruction in the peephole optimizations. That way the code between transforms is kept in SSA form, meaning hopefully less things that can go wrong.

[jaykang10]

Thanks for this patch @dmgreen!
As you mentioned, the ToBeRemoved was to keep the iterator safe. I think the make_early_inc_range is better and it should be ok.
LGTM.

[dmgreen]

Thanks - lets hope this doesn't cause further issues. The MIPeephole optimizations seem to be very easy to get wrong in small and subtle ways.

My testing hasn't shown anything, but please yell if anything comes up.

llvm/lib/Target/AArch64/AArch64MIPeepholeOpt.cpp

struct AArch64MIPeepholeOpt : public MachineFunctionPass {
   ///
   /// 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(unsigned PosOpc, unsigned NegOpc, MachineInstr &MI,
-                   SmallSetVector<MachineInstr *, 8> &ToBeRemoved);
+  bool visitADDSUB(unsigned PosOpc, unsigned NegOpc, MachineInstr &MI);
   template <typename T>
-  bool visitADDSSUBS(OpcodePair PosOpcs, OpcodePair NegOpcs, MachineInstr &MI,
-                     SmallSetVector<MachineInstr *, 8> &ToBeRemoved);
+  bool visitADDSSUBS(OpcodePair PosOpcs, OpcodePair NegOpcs, MachineInstr &MI);
 
   template <typename T>
-  bool visitAND(unsigned Opc, MachineInstr &MI,
-                SmallSetVector<MachineInstr *, 8> &ToBeRemoved);
-  bool visitORR(MachineInstr &MI,
-                SmallSetVector<MachineInstr *, 8> &ToBeRemoved);
+  bool visitAND(unsigned Opc, MachineInstr &MI);
+  bool visitORR(MachineInstr &MI);
   bool runOnMachineFunction(MachineFunction &MF) override;
 
   StringRef getPassName() const override {
     return "AArch64 MI Peephole Optimization pass";
   }
 
   void getAnalysisUsage(AnalysisUsage &AU) const override {
     AU.setPreservesCFG();
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(
-    unsigned Opc, MachineInstr &MI,
-    SmallSetVector<MachineInstr *, 8> &ToBeRemoved) {
+    unsigned Opc, MachineInstr &MI) {
   // 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.
 
   return splitTwoPartImm<T>(
-      MI, ToBeRemoved,
+      MI,
       [Opc](T Imm, unsigned RegSize, T &Imm0, T &Imm1) -> Optional<OpcodePair> {
         if (splitBitmaskImm(Imm, RegSize, Imm0, Imm1))
           return std::make_pair(Opc, Opc);
         return None;
       },
       [&TII = TII](MachineInstr &MI, OpcodePair Opcode, unsigned Imm0,
                    unsigned Imm1, Register SrcReg, Register NewTmpReg,
                    Register NewDstReg) {
         DebugLoc DL = MI.getDebugLoc();
         MachineBasicBlock *MBB = MI.getParent();
         BuildMI(*MBB, MI, DL, TII->get(Opcode.first), NewTmpReg)
             .addReg(SrcReg)
             .addImm(Imm0);
         BuildMI(*MBB, MI, DL, TII->get(Opcode.second), NewDstReg)
             .addReg(NewTmpReg)
             .addImm(Imm1);
       });
 }
 
-bool AArch64MIPeepholeOpt::visitORR(
-    MachineInstr &MI, SmallSetVector<MachineInstr *, 8> &ToBeRemoved) {
+bool AArch64MIPeepholeOpt::visitORR(MachineInstr &MI) {
   // 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)>;
   if (MI.getOperand(3).getImm() != 0)
     return false;
 
   if (MI.getOperand(1).getReg() != AArch64::WZR)
if (SrcMI->getOperand(1).getSubReg() == AArch64::ssub) {
       CpySrc = MRI->createVirtualRegister(&AArch64::FPR32RegClass);
       BuildMI(*SrcMI->getParent(), SrcMI, SrcMI->getDebugLoc(),
               TII->get(TargetOpcode::COPY), CpySrc)
           .add(SrcMI->getOperand(1));
     }
     BuildMI(*SrcMI->getParent(), SrcMI, SrcMI->getDebugLoc(),
             TII->get(AArch64::FMOVSWr), SrcMI->getOperand(0).getReg())
         .addReg(CpySrc);
-    ToBeRemoved.insert(SrcMI);
+    SrcMI->eraseFromParent();
   }
   else if (SrcMI->getOpcode() <= TargetOpcode::GENERIC_OP_END)
     return false;
 
   Register DefReg = MI.getOperand(0).getReg();
   Register SrcReg = MI.getOperand(2).getReg();
   MRI->replaceRegWith(DefReg, SrcReg);
   MRI->clearKillFlags(SrcReg);
-  // replaceRegWith changes MI's definition register. Keep it for SSA form until
-  // deleting MI.
-  MI.getOperand(0).setReg(DefReg);
-  ToBeRemoved.insert(&MI);
-
   LLVM_DEBUG(dbgs() << "Removed: " << MI << "\n");
+  MI.eraseFromParent();
 
   return true;
 }
 
 template <typename T>
 static bool splitAddSubImm(T Imm, unsigned RegSize, T &Imm0, T &Imm1) {
   // The immediate must be in the form of ((imm0 << 12) + imm1), in which both
   // imm0 and imm1 are non-zero 12-bit unsigned int.
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(
-    unsigned PosOpc, unsigned NegOpc, MachineInstr &MI,
-    SmallSetVector<MachineInstr *, 8> &ToBeRemoved) {
+    unsigned PosOpc, unsigned NegOpc, MachineInstr &MI) {
   // 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.
 
   return splitTwoPartImm<T>(
-      MI, ToBeRemoved,
+      MI,
       [PosOpc, NegOpc](T Imm, unsigned RegSize, T &Imm0,
                        T &Imm1) -> Optional<OpcodePair> {
         if (splitAddSubImm(Imm, RegSize, Imm0, Imm1))
           return std::make_pair(PosOpc, PosOpc);
         if (splitAddSubImm(-Imm, RegSize, Imm0, Imm1))
           return std::make_pair(NegOpc, NegOpc);
         return None;
       },
return splitTwoPartImm<T>(
             .addReg(NewTmpReg)
             .addImm(Imm1)
             .addImm(0);
       });
 }
 
 template <typename T>
 bool AArch64MIPeepholeOpt::visitADDSSUBS(
-    OpcodePair PosOpcs, OpcodePair NegOpcs, MachineInstr &MI,
-    SmallSetVector<MachineInstr *, 8> &ToBeRemoved) {
+    OpcodePair PosOpcs, OpcodePair NegOpcs, MachineInstr &MI) {
   // Try the same transformation as ADDSUB but with additional requirement
   // that the condition code usages are only for Equal and Not Equal
   return splitTwoPartImm<T>(
-      MI, ToBeRemoved,
+      MI,
       [PosOpcs, NegOpcs, &MI, &TRI = TRI, &MRI = MRI](
           T Imm, unsigned RegSize, T &Imm0, T &Imm1) -> Optional<OpcodePair> {
         OpcodePair OP;
         if (splitAddSubImm(Imm, RegSize, Imm0, Imm1))
           OP = PosOpcs;
         else if (splitAddSubImm(-Imm, RegSize, Imm0, Imm1))
           OP = NegOpcs;
         else
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,
+    MachineInstr &MI,
     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))
bool AArch64MIPeepholeOpt::splitTwoPartImm(
   // replaceRegWith changes MI's definition register. Keep it for SSA form until
   // deleting MI. Only if we made a new destination register.
   if (DstReg != NewDstReg) {
     MRI->replaceRegWith(DstReg, NewDstReg);
     MI.getOperand(0).setReg(DstReg);
   }
 
   // Record the MIs need to be removed.
-  ToBeRemoved.insert(&MI);
+  MI.eraseFromParent();
   if (SubregToRegMI)
-    ToBeRemoved.insert(SubregToRegMI);
-  ToBeRemoved.insert(MovMI);
+    SubregToRegMI->eraseFromParent();
+  MovMI->eraseFromParent();
 
   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!");
 
   bool Changed = false;
-  SmallSetVector<MachineInstr *, 8> ToBeRemoved;
 
   for (MachineBasicBlock &MBB : MF) {
-    for (MachineInstr &MI : MBB) {
+    for (MachineInstr &MI : make_early_inc_range(MBB)) {
       switch (MI.getOpcode()) {
       default:
         break;
       case AArch64::ANDWrr:
-        Changed = visitAND<uint32_t>(AArch64::ANDWri, MI, ToBeRemoved);
+        Changed = visitAND<uint32_t>(AArch64::ANDWri, MI);
         break;
       case AArch64::ANDXrr:
-        Changed = visitAND<uint64_t>(AArch64::ANDXri, MI, ToBeRemoved);
+        Changed = visitAND<uint64_t>(AArch64::ANDXri, MI);
         break;
       case AArch64::ORRWrs:
-        Changed = visitORR(MI, ToBeRemoved);
+        Changed = visitORR(MI);
         break;
       case AArch64::ADDWrr:
-        Changed = visitADDSUB<uint32_t>(AArch64::ADDWri, AArch64::SUBWri, MI,
-                                        ToBeRemoved);
+        Changed = visitADDSUB<uint32_t>(AArch64::ADDWri, AArch64::SUBWri, MI);
         break;
       case AArch64::SUBWrr:
-        Changed = visitADDSUB<uint32_t>(AArch64::SUBWri, AArch64::ADDWri, MI,
-                                        ToBeRemoved);
+        Changed = visitADDSUB<uint32_t>(AArch64::SUBWri, AArch64::ADDWri, MI);
         break;
       case AArch64::ADDXrr:
-        Changed = visitADDSUB<uint64_t>(AArch64::ADDXri, AArch64::SUBXri, MI,
-                                        ToBeRemoved);
+        Changed = visitADDSUB<uint64_t>(AArch64::ADDXri, AArch64::SUBXri, MI);
         break;
       case AArch64::SUBXrr:
-        Changed = visitADDSUB<uint64_t>(AArch64::SUBXri, AArch64::ADDXri, MI,
-                                        ToBeRemoved);
+        Changed = visitADDSUB<uint64_t>(AArch64::SUBXri, AArch64::ADDXri, MI);
         break;
       case AArch64::ADDSWrr:
         Changed = visitADDSSUBS<uint32_t>({AArch64::ADDWri, AArch64::ADDSWri},
                                           {AArch64::SUBWri, AArch64::SUBSWri},
-                                          MI, ToBeRemoved);
+                                          MI);
         break;
       case AArch64::SUBSWrr:
         Changed = visitADDSSUBS<uint32_t>({AArch64::SUBWri, AArch64::SUBSWri},
                                           {AArch64::ADDWri, AArch64::ADDSWri},
-                                          MI, ToBeRemoved);
+                                          MI);
         break;
       case AArch64::ADDSXrr:
         Changed = visitADDSSUBS<uint64_t>({AArch64::ADDXri, AArch64::ADDSXri},
                                           {AArch64::SUBXri, AArch64::SUBSXri},
-                                          MI, ToBeRemoved);
+                                          MI);
         break;
       case AArch64::SUBSXrr:
         Changed = visitADDSSUBS<uint64_t>({AArch64::SUBXri, AArch64::SUBSXri},
                                           {AArch64::ADDXri, AArch64::ADDSXri},
-                                          MI, ToBeRemoved);
+                                          MI);
         break;
       }
     }
   }
 
-  for (MachineInstr *MI : ToBeRemoved)
-    MI->eraseFromParent();
-
   return Changed;
 }
 
 FunctionPass *llvm::createAArch64MIPeepholeOptPass() {
   return new AArch64MIPeepholeOpt();
 }

llvm/test/CodeGen/AArch64/peephole-orr.mir

+# NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py
+# RUN: llc -run-pass=aarch64-mi-peephole-opt -o - -mtriple=aarch64-unknown-linux -verify-machineinstrs %s | FileCheck %s
+
+---
+name: copy_multiple_uses
+tracksRegLiveness: true
+body: |
+  ; CHECK-LABEL: name: copy_multiple_uses
+  ; CHECK: bb.0.entry:
+  ; CHECK-NEXT:   successors: %bb.1(0x80000000)
+  ; CHECK-NEXT:   liveins: $w0, $q0
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT:   [[COPY:%[0-9]+]]:fpr128 = COPY $q0
+  ; CHECK-NEXT:   [[COPY1:%[0-9]+]]:gpr32 = COPY $w0
+  ; CHECK-NEXT:   [[COPY2:%[0-9]+]]:gpr32sp = COPY $w0
+  ; CHECK-NEXT:   B %bb.1
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT: bb.1:
+  ; CHECK-NEXT:   successors: %bb.1(0x40000000), %bb.2(0x40000000)
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT:   [[SUBSWri:%[0-9]+]]:gpr32 = SUBSWri [[COPY2]], 1, 0, implicit-def $nzcv
+  ; CHECK-NEXT:   [[COPY3:%[0-9]+]]:fpr32 = COPY [[COPY]].ssub
+  ; CHECK-NEXT:   [[FMOVSWr:%[0-9]+]]:gpr32 = FMOVSWr [[COPY3]]
+  ; CHECK-NEXT:   [[ADDWrr:%[0-9]+]]:gpr32 = ADDWrr [[FMOVSWr]], [[SUBSWri]]
+  ; CHECK-NEXT:   Bcc 2, %bb.1, implicit $nzcv
+  ; CHECK-NEXT:   B %bb.2
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT: bb.2:
+  ; CHECK-NEXT:   $w0 = COPY [[ADDWrr]]
+  ; CHECK-NEXT:   RET_ReallyLR implicit $w0
+  bb.0.entry:
+    liveins: $w0, $q0
+    %0:fpr128 = COPY $q0
+    %1:gpr32 = COPY $w0
+    %6:gpr32sp = COPY $w0
+    B %bb.1
+
+  bb.1:
+    %7:gpr32 = SUBSWri %6, 1, 0, implicit-def $nzcv
+    %2:gpr32 = COPY %0.ssub:fpr128
+    %3:gpr32 = ORRWrs $wzr, %2:gpr32, 0
+    %5:gpr32 = ADDWrr %2:gpr32, %7:gpr32
+    Bcc 2, %bb.1, implicit $nzcv
+    B %bb.2
+
+  bb.2:
+    $w0 = COPY %5
+    RET_ReallyLR implicit $w0