[MachineOutliner] Teach outliner to set live-ins

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.

Differential Revision: https://reviews.llvm.org/D78605

lib/CodeGen/MachineOutliner.cpp

@@ -1178,12 +1178,35 @@ MachineFunction *MachineOutliner::createOutlinedFunction(
     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)) {

lib/Target/AArch64/AArch64InstrInfo.cpp

@@ -6450,7 +6450,8 @@ void AArch64InstrInfo::buildOutlinedFrame(
     IsLeafFunction = false;
 
     // LR has to be a live in so that we can save it.
-    MBB.addLiveIn(AArch64::LR);
+    if (!MBB.isLiveIn(AArch64::LR))
+      MBB.addLiveIn(AArch64::LR);
 
     MachineBasicBlock::iterator It = MBB.begin();
     MachineBasicBlock::iterator Et = MBB.end();
@@ -6529,8 +6530,13 @@ void AArch64InstrInfo::buildOutlinedFrame(
   }
 
   // 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,

lib/Target/RISCV/RISCVInstrInfo.cpp

@@ -778,6 +778,8 @@ void RISCVInstrInfo::buildOutlinedFrame(
     }
   }
 
+  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)

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

@@ -55,7 +55,8 @@ body:             |
     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
@@ -63,4 +64,4 @@ body:             |
 # 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

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

@@ -196,9 +196,11 @@ body:             |
 # 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