[AMDGPU] VGPR instruction placement should be aware of exec mask prol…

llvm/include/llvm/CodeGen/TargetInstrInfo.h

@@ -2166,6 +2166,12 @@ class LLVM_ABI TargetInstrInfo : public MCInstrInfo {
     return false;
   }
 
+  /// Get the insertion point for a live range split copy at the beginning of
+  /// a basic block. This is used by the register allocator's live range
+  /// splitting logic.
+  virtual MachineBasicBlock::iterator
+  getExecAwareInsertPoint(MachineBasicBlock &MBB, Register Reg) const;
+
   /// Allows targets to use appropriate copy instruction while spilitting live
   /// range of a register in register allocation.
   virtual unsigned getLiveRangeSplitOpcode(Register Reg,

llvm/lib/CodeGen/SplitKit.cpp

@@ -842,8 +842,8 @@ SlotIndex SplitEditor::leaveIntvAtTop(MachineBasicBlock &MBB) {
 
   unsigned RegIdx = 0;
   Register Reg = LIS.getInterval(Edit->get(RegIdx)).reg();
-  VNInfo *VNI = defFromParent(RegIdx, ParentVNI, Start, MBB,
-                              MBB.SkipPHIsLabelsAndDebug(MBB.begin(), Reg));
+  MachineBasicBlock::iterator InsertPt = TII.getExecAwareInsertPoint(MBB, Reg);
+  VNInfo *VNI = defFromParent(RegIdx, ParentVNI, Start, MBB, InsertPt);
   RegAssign.insert(Start, VNI->def, OpenIdx);
   LLVM_DEBUG(dump());
   return VNI->def;

llvm/lib/CodeGen/TargetInstrInfo.cpp

@@ -2222,3 +2222,9 @@ bool TargetInstrInfo::isGlobalMemoryObject(const MachineInstr *MI) const {
   return MI->isCall() || MI->hasUnmodeledSideEffects() ||
          (MI->hasOrderedMemoryRef() && !MI->isDereferenceableInvariantLoad());
 }
+
+MachineBasicBlock::iterator
+TargetInstrInfo::getExecAwareInsertPoint(MachineBasicBlock &MBB,
+                                         Register Reg) const {
+  return MBB.SkipPHIsLabelsAndDebug(MBB.begin(), Reg);
+}

llvm/lib/Target/AMDGPU/SIInstrInfo.cpp

@@ -9903,6 +9903,44 @@ bool SIInstrInfo::isBasicBlockPrologue(const MachineInstr &MI,
            MI.modifiesRegister(AMDGPU::EXEC, &RI)));
 }
 
+MachineBasicBlock::iterator
+SIInstrInfo::getExecAwareInsertPoint(MachineBasicBlock &MBB,
+                                     Register Reg) const {
+  if (!Reg)
+    return MBB.SkipPHIsLabelsAndDebug(MBB.begin());
+
+  // Instructions using SGPRS are independent of exec mask: skip.
+  const MachineRegisterInfo &MRI = MBB.getParent()->getRegInfo();
+  const TargetRegisterClass *RC = Reg.isPhysical() ? RI.getPhysRegBaseClass(Reg)
+                                                   : MRI.getRegClassOrNull(Reg);
+  if (RC && RI.isSGPRClass(RC))
+    return MBB.SkipPHIsLabelsAndDebug(MBB.begin(), Reg);
+
+  // For VGPRs, find a safe insert point after exec has been restored.
+  // A standard basic block prologue pattern is:
+  //   BB:
+  //     <PHIs>
+  //     <labels>
+  //     <SGPR spills/restores>
+  //     <scalar copies - live range splits for SGPRs>
+  //     s_or_b32 exec_lo, exec_lo, s_saved
+
+  // Look for exec restore instruction: it could be after the first SGPR copy.
+  MachineBasicBlock::iterator I = MBB.SkipPHIsLabelsAndDebug(MBB.begin(), Reg);
+  MachineBasicBlock::iterator E = MBB.end();
+
+  for (MachineBasicBlock::iterator It = I; It != E; ++It) {
+    if (!It->isTerminator() && It->modifiesRegister(AMDGPU::EXEC, &RI))
+      return std::next(It);
+
+    if (It->mayLoadOrStore() && !isSGPRSpill(It->getOpcode()) &&
+        !isWWMRegSpillOpcode(It->getOpcode()))
+      break;
+  }
+
+  return I;
+}
+
 MachineInstrBuilder
 SIInstrInfo::getAddNoCarry(MachineBasicBlock &MBB,
                            MachineBasicBlock::iterator I,

llvm/lib/Target/AMDGPU/SIInstrInfo.h

@@ -1579,6 +1579,9 @@ class SIInstrInfo final : public AMDGPUGenInstrInfo {
   bool isBasicBlockPrologue(const MachineInstr &MI,
                             Register Reg = Register()) const override;
 
+  MachineBasicBlock::iterator
+  getExecAwareInsertPoint(MachineBasicBlock &MBB, Register Reg) const override;
+
   MachineInstr *createPHIDestinationCopy(MachineBasicBlock &MBB,
                                          MachineBasicBlock::iterator InsPt,
                                          const DebugLoc &DL, Register Src,