[MachineOutliner] Teach outliner to set live-ins

[efriedma]

Summary

Preserving liveness can be useful even late in the pipeline, if we're doing substantial optimization work afterwards. (See, for example, D76065.) Teach MachineOutliner how to correctly set live-ins on the basic block in outlined functions.

[paquette]

LGTM

llvm/lib/CodeGen/MachineOutliner.cpp

for (auto I = FirstCand.front(), E = std::next(FirstCand.back()); I != E;
     }
     NewMI->dropMemRefs(MF);
 
     // Don't keep debug information for outlined instructions.
     NewMI->setDebugLoc(DebugLoc());
     MBB.insert(MBB.end(), NewMI);
   }
 
-  TII.buildOutlinedFrame(MBB, MF, OF);
-
-  // Outlined functions shouldn't preserve liveness.
-  MF.getProperties().reset(MachineFunctionProperties::Property::TracksLiveness);
+  // Set normal properties for a late MachineFunction.
+  MF.getProperties().reset(MachineFunctionProperties::Property::IsSSA);
+  MF.getProperties().set(MachineFunctionProperties::Property::NoPHIs);
+  MF.getProperties().set(MachineFunctionProperties::Property::NoVRegs);
+  MF.getProperties().set(MachineFunctionProperties::Property::TracksLiveness);
   MF.getRegInfo().freezeReservedRegs(MF);
 
+  // Compute live-in set for outlined fn
+  const MachineRegisterInfo &MRI = MF.getRegInfo();
+  const TargetRegisterInfo &TRI = *MRI.getTargetRegisterInfo();
+  LivePhysRegs LiveIns(TRI);
+  for (auto &Cand : OF.Candidates) {
+    // Figure out live-ins at the first instruction.
+    MachineBasicBlock &OutlineBB = *Cand.front()->getParent();
+    LivePhysRegs CandLiveIns(TRI);
+    CandLiveIns.addLiveOuts(OutlineBB);
+    for (const MachineInstr &MI :
+         reverse(make_range(Cand.front(), OutlineBB.end())))
+      CandLiveIns.stepBackward(MI);
+
+    // The live-in set for the outlined function is the union of the live-ins
+    // from all the outlining points.
+    for (MCPhysReg Reg : make_range(CandLiveIns.begin(), CandLiveIns.end()))
+      LiveIns.addReg(Reg);
+  }
+  addLiveIns(MBB, LiveIns);
+
+  TII.buildOutlinedFrame(MBB, MF, OF);
+
   // If there's a DISubprogram associated with this outlined function, then
   // emit debug info for the outlined function.
   if (DISubprogram *SP = getSubprogramOrNull(OF)) {
     // We have a DISubprogram. Get its DICompileUnit.
     DICompileUnit *CU = SP->getUnit();
     DIBuilder DB(M, true, CU);
     DIFile *Unit = SP->getFile();
     Mangler Mg;

llvm/lib/Target/AArch64/AArch64InstrInfo.cpp

if (std::any_of(MBB.instr_begin(), MBB.instr_end(), IsNonTailCall)) {
     // stack.
     assert(OF.FrameConstructionID != MachineOutlinerDefault &&
            "Can only fix up stack references once");
     fixupPostOutline(MBB);
 
     IsLeafFunction = false;
 
     // LR has to be a live in so that we can save it.
+    if (!MBB.isLiveIn(AArch64::LR))
       MBB.addLiveIn(AArch64::LR);
 
     MachineBasicBlock::iterator It = MBB.begin();
     MachineBasicBlock::iterator Et = MBB.end();
 
     if (OF.FrameConstructionID == MachineOutlinerTailCall ||
         OF.FrameConstructionID == MachineOutlinerThunk)
       Et = std::prev(MBB.end());
 
void AArch64InstrInfo::buildOutlinedFrame(
   if (OF.FrameConstructionID == MachineOutlinerTailCall ||
       OF.FrameConstructionID == MachineOutlinerThunk) {
     signOutlinedFunction(MF, MBB, ShouldSignReturnAddr,
                          ShouldSignReturnAddrWithAKey);
     return;
   }
 
   // It's not a tail call, so we have to insert the return ourselves.
+
+  // LR has to be a live in so that we can return to it.
+  if (!MBB.isLiveIn(AArch64::LR))
+    MBB.addLiveIn(AArch64::LR);
+
   MachineInstr *ret = BuildMI(MF, DebugLoc(), get(AArch64::RET))
-                          .addReg(AArch64::LR, RegState::Undef);
+                          .addReg(AArch64::LR);
   MBB.insert(MBB.end(), ret);
 
   signOutlinedFunction(MF, MBB, ShouldSignReturnAddr,
                        ShouldSignReturnAddrWithAKey);
 
   FI->setOutliningStyle("Function");
 
   // Did we have to modify the stack by saving the link register?

llvm/lib/Target/RISCV/RISCVInstrInfo.cpp

for (; I != E; ++I) {
       if (I->isCFIInstruction()) {
         I->removeFromParent();
         Changed = true;
         break;
       }
     }
   }
 
+  MBB.addLiveIn(RISCV::X5);
+
   // Add in a return instruction to the end of the outlined frame.
   MBB.insert(MBB.end(), BuildMI(MF, DebugLoc(), get(RISCV::JALR))
       .addReg(RISCV::X0, RegState::Define)
       .addReg(RISCV::X5)
       .addImm(0));
 }
 
 MachineBasicBlock::iterator RISCVInstrInfo::insertOutlinedCall(

llvm/test/CodeGen/AArch64/machine-outliner-calls.mir

 name:            baz
 tracksRegLiveness: true
 body:             |
   bb.0:
     liveins: $w0, $lr, $w8
     RET undef $lr
 
 # CHECK: name:            OUTLINED_FUNCTION_0
-# CHECK-DAG: bb.0:
+# CHECK: bb.0:
+# CHECK: liveins: $x19, $x20, $x21, $x22, $x23, $x24, $x25, $x26, $x27, $x28, $d15, $d8, $d9, $d10, $d11, $d12, $d13, $d14, $lr
 # CHECK-DAG: frame-setup CFI_INSTRUCTION def_cfa_offset -16
 # CHECK-NEXT: frame-setup CFI_INSTRUCTION offset $w30, 16
 # CHECK-NEXT: early-clobber $sp = STRXpre $lr, $sp, -16
 # CHECK-NEXT: BL @baz, implicit-def dead $lr, implicit $sp
 # CHECK-NEXT: $w17 = ORRWri $wzr, 1
 # CHECK-NEXT: $w17 = ORRWri $wzr, 1
 # CHECK-NEXT: early-clobber $sp, $lr = LDRXpost $sp, 16
-# CHECK-NEXT: RET undef $lr
+# CHECK-NEXT: RET $lr

llvm/test/CodeGen/AArch64/machine-outliner-retaddr-sign-sp-mod.mir

 # CHECK-NOT:            BL @OUTLINED_FUNCTION_{{.*}}
 # CHECK:                frame-destroy AUTIASP implicit-def $lr, implicit killed $lr, implicit $sp
 # CHECK-NEXT:           RET undef $lr
 
 # Outlined function that contains only legal sp modifications
 # CHECK:          name:            [[OUTLINED_FUNC]]
 # CHECK:          body:             |
 # CHECK-NEXT:       bb.0:
+# CHECK-NEXT: liveins: $lr
+# CHECK-NEXT: {{^  $}}
 # CHECK-NEXT:         frame-setup PACIASP implicit-def $lr, implicit $lr, implicit $sp
 # CHECK-NEXT:         frame-setup CFI_INSTRUCTION negate_ra_sign_state
 # CHECK-NEXT:         $sp = frame-setup SUBXri $sp, 16, 0
 # CHECK:              $sp = frame-destroy ADDXri $sp, 16, 0
 # CHECK-NEXT:         frame-destroy AUTIASP implicit-def $lr, implicit $lr, implicit $sp
-# CHECK-NEXT:         RET undef $lr
+# CHECK-NEXT:         RET $lr