FastISel: take care to update iterators when removing instructions.

[t.p.northover]

Summary

We keep a few iterators into the basic block we're selecting while performing FastISel. Usually this is fine, but occasionally code wants to remove already-emitted instructions. When this happens we have to be careful to update those iterators so they're not pointint at dangling memory.

[t.p.northover]

Ping.

[aschwaighofer]

LGTM

llvm/lib/CodeGen/SelectionDAG/FastISel.cpp

while (FuncInfo.InsertPt != FuncInfo.MBB->end() &&
     ++FuncInfo.InsertPt;
 }
 
 void FastISel::removeDeadCode(MachineBasicBlock::iterator I,
                               MachineBasicBlock::iterator E) {
   assert(I.isValid() && E.isValid() && std::distance(I, E) > 0 &&
          "Invalid iterator!");
   while (I != E) {
+    if (LastFlushPoint == I)
+      LastFlushPoint = E;
+    if (SavedInsertPt == I)
+      SavedInsertPt = E;
+    if (EmitStartPt == I)
+      EmitStartPt = E.isValid() ? &*E : nullptr;
+    if (LastLocalValue == I)
+      LastLocalValue = E.isValid() ? &*E : nullptr;
+
     MachineInstr *Dead = &*I;
     ++I;
     Dead->eraseFromParent();
     ++NumFastIselDead;
   }
   recomputeInsertPt();
 }
 

llvm/lib/Target/AArch64/AArch64FastISel.cpp

if (IntExtVal) {
     // could select it. Emit a copy to subreg if necessary. FastISel will remove
     // it when it selects the integer extend.
     unsigned Reg = lookUpRegForValue(IntExtVal);
     auto *MI = MRI.getUniqueVRegDef(Reg);
     if (!MI) {
       if (RetVT == MVT::i64 && VT <= MVT::i32) {
         if (WantZExt) {
           // Delete the last emitted instruction from emitLoad (SUBREG_TO_REG).
-          std::prev(FuncInfo.InsertPt)->eraseFromParent();
-          ResultReg = std::prev(FuncInfo.InsertPt)->getOperand(0).getReg();
+          MachineBasicBlock::iterator I(std::prev(FuncInfo.InsertPt));
+          ResultReg = std::prev(I)->getOperand(0).getReg();
+          removeDeadCode(I, std::next(I));
         } else
           ResultReg = fastEmitInst_extractsubreg(MVT::i32, ResultReg,
                                                  /*IsKill=*/true,
                                                  AArch64::sub_32);
       }
       updateValueMap(I, ResultReg);
       return true;
     }
 
     // The integer extend has already been emitted - delete all the instructions
     // that have been emitted by the integer extend lowering code and use the
     // result from the load instruction directly.
     while (MI) {
       Reg = 0;
       for (auto &Opnd : MI->uses()) {
         if (Opnd.isReg()) {
           Reg = Opnd.getReg();
           break;
         }
       }
-      MI->eraseFromParent();
+      MachineBasicBlock::iterator I(MI);
+      removeDeadCode(I, std::next(I));
       MI = nullptr;
       if (Reg)
         MI = MRI.getUniqueVRegDef(Reg);
     }
     updateValueMap(IntExtVal, ResultReg);
     return true;
   }
 
BuildMI(*FuncInfo.MBB, FuncInfo.InsertPt, DbgLoc,
         .addReg(Reg, getKillRegState(true))
         .addImm(AArch64::sub_32);
     Reg = Reg64;
   } else {
     assert((MI->getOpcode() == TargetOpcode::COPY &&
             MI->getOperand(1).getSubReg() == AArch64::sub_32) &&
            "Expected copy instruction");
     Reg = MI->getOperand(1).getReg();
-    MI->eraseFromParent();
+    MachineBasicBlock::iterator I(MI);
+    removeDeadCode(I, std::next(I));
   }
   updateValueMap(I, Reg);
   return true;
 }
 
 bool AArch64FastISel::selectIntExt(const Instruction *I) {
   assert((isa<ZExtInst>(I) || isa<SExtInst>(I)) &&
          "Unexpected integer extend instruction.");

llvm/lib/Target/ARM/ARMFastISel.cpp

bool ARMFastISel::tryToFoldLoadIntoMI(MachineInstr *MI, unsigned OpNo,
 
   // See if we can handle this address.
   Address Addr;
   if (!ARMComputeAddress(LI->getOperand(0), Addr)) return false;
 
   unsigned ResultReg = MI->getOperand(0).getReg();
   if (!ARMEmitLoad(VT, ResultReg, Addr, LI->getAlignment(), isZExt, false))
     return false;
-  MI->eraseFromParent();
+  MachineBasicBlock::iterator I(MI);
+  removeDeadCode(I, std::next(I));
   return true;
 }
 
 unsigned ARMFastISel::ARMLowerPICELF(const GlobalValue *GV,
                                      unsigned Align, MVT VT) {
   bool UseGOT_PREL = !TM.shouldAssumeDSOLocal(*GV->getParent(), GV);
 
   LLVMContext *Context = &MF->getFunction().getContext();

llvm/lib/Target/PowerPC/PPCFastISel.cpp

if (!PPCComputeAddress(LI->getOperand(0), Addr))
     return false;
 
   unsigned ResultReg = MI->getOperand(0).getReg();
 
   if (!PPCEmitLoad(VT, ResultReg, Addr, nullptr, IsZExt,
         PPCSubTarget->hasSPE() ? PPC::EVLDD : PPC::LFD))
     return false;
 
-  MI->eraseFromParent();
+  MachineBasicBlock::iterator I(MI);
+  removeDeadCode(I, std::next(I));
   return true;
 }
 
 // Attempt to lower call arguments in a faster way than done by
 // the selection DAG code.
 bool PPCFastISel::fastLowerArguments() {
   // Defer to normal argument lowering for now.  It's reasonably
   // efficient.  Consider doing something like ARM to handle the

llvm/lib/Target/X86/X86FastISel.cpp

for (MachineInstr::mop_iterator I = Result->operands_begin(),
     unsigned IndexReg = constrainOperandRegClass(Result->getDesc(),
                                                  MO.getReg(), OperandNo);
     if (IndexReg == MO.getReg())
       continue;
     MO.setReg(IndexReg);
   }
 
   Result->addMemOperand(*FuncInfo.MF, createMachineMemOperandFor(LI));
-  MI->eraseFromParent();
+  MachineBasicBlock::iterator I(MI);
+  removeDeadCode(I, std::next(I));
   return true;
 }
 
 unsigned X86FastISel::fastEmitInst_rrrr(unsigned MachineInstOpcode,
                                         const TargetRegisterClass *RC,
                                         unsigned Op0, bool Op0IsKill,
                                         unsigned Op1, bool Op1IsKill,
                                         unsigned Op2, bool Op2IsKill,

llvm/test/CodeGen/AArch64/fast-isel-erase.ll

+; RUN: llc -mtriple=arm64-apple-ios -o - %s -fast-isel=1 -O0 | FileCheck %s
+
+; The zext can be folded into the load and removed, but doing so can invalidate
+; pointers internal to FastISel and cause a crash so it must be done carefully.
+define i32 @test() {
+; CHECK-LABEL: test:
+; CHECK: ldrh
+; CHECK: bl _callee
+; CHECK-NOT: uxth
+
+entry:
+  store i32 undef, i32* undef, align 4
+  %t81 = load i16, i16* undef, align 2
+  call void @callee()
+  %t82 = zext i16 %t81 to i32
+  %t83 = shl i32 %t82, 16
+  %t84 = or i32 undef, %t83
+  br label %end
+
+end:
+  %val = phi i32 [%t84, %entry]
+  ret i32 %val
+}
+
+declare void @callee()