[llvm] [1/2][AMDGPU] Fixed crash due to virtual register defs not dominating uses (PR #198472)

[via llvm-commits]

================
@@ -476,6 +492,69 @@ void AMDGPURewriteAGPRCopyMFMAImpl::collectSpillIndexUses(
   }
 }
 
+bool AMDGPURewriteAGPRCopyMFMAImpl::insertImplicitDefsForLiveSegments(
+    LiveInterval &StackLI, int Slot, ArrayRef<MachineInstr *> SpillStores,
+    Register NewVReg) const {
+  const unsigned ScanLimit = ImplicitDefScanLimit;
+
+  SmallPtrSet<MachineInstr *, 4> StoreSet(SpillStores.begin(),
+                                          SpillStores.end());
+
+  SmallVector<MachineInstr *, 2> InsertedDefs;
+
+  for (const LiveRange::Segment &Seg : StackLI) {
+    MachineInstr *MI = LIS.getInstructionFromIndex(Seg.start);
+    // Fall back to LIS query in case the segment start does not correspond
+    // to an instruction.
+    MachineBasicBlock *MBB =
+        MI ? MI->getParent() : LIS.getMBBFromIndex(Seg.start);
+
+    // Bail-out: we could not resolve any MBB for this segment start.
+    // Roll back any defs already inserted for earlier segments -- the
+    // caller must skip unspilling this slot entirely.
+    if (!MBB) {
+      for (MachineInstr *Def : InsertedDefs) {
+        LIS.RemoveMachineInstrFromMaps(*Def);
+        Def->eraseFromParent();
+      }
+      return false;
+    }
+
+    // Scan forward from the segment start looking for a spill store to this
+    // slot. LiveStacks liveness for a store can start a few instructions
+    // before the store itself, so we allow a small window controlled by
+    // ImplicitDefScanLimit. Redundant IMPLICIT_DEFs (e.g. because of a small
+    // scan limit) are harmless, since they will be cleaned up by downstream
+    // transformations.
+    MachineBasicBlock::iterator It = MI ? MI->getIterator() : MBB->begin();
+    bool FoundStore = false;
+    for (unsigned I = 0; I < ScanLimit && It != MBB->end(); ++I, ++It) {
+      if (StoreSet.count(&*It)) {
+        FoundStore = true;
+        break;
+      }
+    }
----------------
ruiling wrote:

I just realized that a live-segment might span across multiple blocks (sorry I missed that). So actually we need going backward through the CFG at block level. We need to start from each end-position of the live-segment, we get the block that contains that position, then check if the live-range of the stack slot was alive at the block start. If it is alive, go checking at the terminator of the predecessors. If the stack slot is not live there, insert implicit_def. otherwise repeat the check going backward the CFG. We may need some logic to avoid repeated traversal. Sounds reasonable to you? (I think in real case, the non-dominating case comes from phi-node, maybe we don't even need to handle case that the live-segment start from the middle of a block but does not have a stack store.)

https://github.com/llvm/llvm-project/pull/198472