[AIEX] Fix invalid remat after super-reg rewrite/unbundle

llvm/include/llvm/CodeGen/TargetSubtargetInfo.h

@@ -4,7 +4,7 @@
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
-// Modifications (c) Copyright 2023-2025 Advanced Micro Devices, Inc. or its
+// Modifications (c) Copyright 2023-2026 Advanced Micro Devices, Inc. or its
 // affiliates
 //
 //===----------------------------------------------------------------------===//
@@ -21,11 +21,13 @@
 #include "llvm/ADT/StringRef.h"
 #include "llvm/CodeGen/MacroFusion.h"
 #include "llvm/CodeGen/PBQPRAConstraint.h"
+#include "llvm/CodeGen/Register.h"
 #include "llvm/CodeGen/SchedulerRegistry.h"
 #include "llvm/IR/GlobalValue.h"
 #include "llvm/MC/MCSubtargetInfo.h"
 #include "llvm/Support/CodeGen.h"
 #include <memory>
+#include <optional>
 #include <vector>
 
 namespace llvm {
@@ -141,6 +143,24 @@ class TargetSubtargetInfo : public MCSubtargetInfo {
     return nullptr;
   }
 
+  /// Optional target hook used by InlineSpiller to recover a "logical group
+  /// original" for stack-slot sharing when a target pass has deliberately
+  /// severed the VirtRegMap split-from chain (via clearSplitFromReg) for
+  /// correctness reasons (e.g. to stop SplitKit::defFromParent from
+  /// rematerializing through a stale ancestor LiveInterval).
+  ///
+  /// If the target returns a valid Register \p R, InlineSpiller will use
+  /// \p R as the "Original" for stack-slot sharing (HoistSpillHelper /
+  /// MergeableSpills) instead of VRM.getOriginal(VirtReg). The returned
+  /// register must still have a valid LiveInterval; otherwise the override
+  /// is ignored.
+  ///
+  /// Default: no override.
+  virtual std::optional<Register>
+  getSpillGroupOriginal(const MachineFunction &MF, Register VirtReg) const {
+    return std::nullopt;
+  }
+
   /// Resolve a SchedClass at runtime, where SchedClass identifies an
   /// MCSchedClassDesc with the isVariant property. This may return the ID of
   /// another variant SchedClass, but repeated invocation must quickly terminate

llvm/include/llvm/CodeGen/VirtRegMap.h

@@ -4,7 +4,7 @@
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
-// Modifications (c) Copyright 2023-2024 Advanced Micro Devices, Inc. or its
+// Modifications (c) Copyright 2023-2026 Advanced Micro Devices, Inc. or its
 // affiliates
 //
 //===----------------------------------------------------------------------===//
@@ -176,6 +176,14 @@ class VirtRegMap {
     }
   }
 
+  /// clearSplitFromReg - Remove the split-from mapping for virtReg,
+  /// making it its own original. This restores the register to the
+  /// same canonical state as a freshly created vreg (no split parent).
+  void clearSplitFromReg(Register virtReg) {
+    assert(virtReg.isVirtual());
+    Virt2SplitMap[virtReg] = Register();
+  }
+
   /// returns the live interval virtReg is split from.
   Register getPreSplitReg(Register virtReg) const {
     return Virt2SplitMap[virtReg];

llvm/lib/CodeGen/InlineSpiller.cpp

@@ -4,6 +4,9 @@
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
+// Modifications (c) Copyright 2026 Advanced Micro Devices, Inc. or its
+// affiliates
+//
 //===----------------------------------------------------------------------===//
 //
 // The inline spiller modifies the machine function directly instead of
@@ -1289,6 +1292,16 @@ void InlineSpiller::spill(LiveRangeEdit &edit) {
          "Trying to spill a stack slot.");
   // Share a stack slot among all descendants of Original.
   Original = VRM.getOriginal(edit.getReg());
+  // Allow the target to redirect this lookup. Some target passes deliberately
+  // sever the VirtRegMap split-from chain (clearSplitFromReg) for correctness
+  // (e.g. to stop SplitKit from rematerializing through a stale ancestor LI),
+  // but still want spills of the descendants to share a stack slot with the
+  // logical group's original. The hook returns that "logical group original"
+  // when it should be used here.
+  if (auto SyntheticOrig =
+          MF.getSubtarget().getSpillGroupOriginal(MF, Original))
+    if (LIS.hasInterval(*SyntheticOrig))
+      Original = *SyntheticOrig;
   StackSlot = VRM.getStackSlot(Original);
   StackInt = nullptr;
 

llvm/lib/Target/AIE/AIEBaseSubtarget.cpp

@@ -16,6 +16,7 @@
 #include "AIE.h"
 #include "AIEBaseRegisterInfo.h"
 #include "AIEInterBlockScheduling.h"
+#include "AIEMachineFunctionInfo.h"
 #include "AIEMachineScheduler.h"
 #include "AIEMaxLatencyFinder.h"
 #include "AIERegMemEventTracker.h"
@@ -158,6 +159,17 @@ const AIEBaseSubtarget &AIEBaseSubtarget::get(const MachineFunction &MF) {
   return static_cast<const AIEBaseSubtarget &>(MF.getSubtarget());
 }
 
+std::optional<Register>
+AIEBaseSubtarget::getSpillGroupOriginal(const MachineFunction &MF,
+                                        Register VirtReg) const {
+  // The MFI is created lazily; if for some reason the MF has not allocated one
+  // (e.g. very early in pipeline), there is nothing to look up.
+  const auto *MFI = MF.getInfo<AIEMachineFunctionInfo>();
+  if (!MFI)
+    return std::nullopt;
+  return MFI->getSpillGroupOriginal(VirtReg);
+}
+
 namespace {
 
 // Set latency and declare height/depth dirty if it changes

llvm/lib/Target/AIE/AIEBaseSubtarget.h

@@ -96,6 +96,13 @@ class AIEBaseSubtarget : public TargetSubtargetInfo {
 
   // All AIE targets need post scheduling for correct instruction timing
   bool forcePostRAScheduling() const override { return true; }
+
+  /// See TargetSubtargetInfo::getSpillGroupOriginal. Forwards to the side map
+  /// in AIEMachineFunctionInfo populated by AIE register-rewriter passes when
+  /// they sever the VRM split-from chain for correctness.
+  std::optional<Register>
+  getSpillGroupOriginal(const MachineFunction &MF,
+                        Register VirtReg) const override;
 };
 } // namespace llvm
 

llvm/lib/Target/AIE/AIEMachineFunctionInfo.h

@@ -4,7 +4,7 @@
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
-// (c) Copyright 2023-2024 Advanced Micro Devices, Inc. or its affiliates
+// (c) Copyright 2023-2026 Advanced Micro Devices, Inc. or its affiliates
 //
 //===----------------------------------------------------------------------===//
 //
@@ -15,9 +15,11 @@
 #ifndef LLVM_LIB_TARGET_AIE_AIEMACHINEFUNCTIONINFO_H
 #define LLVM_LIB_TARGET_AIE_AIEMACHINEFUNCTIONINFO_H
 
+#include "llvm/ADT/DenseMap.h"
 #include "llvm/CodeGen/MIRYamlMapping.h"
 #include "llvm/CodeGen/MachineFrameInfo.h"
 #include "llvm/CodeGen/MachineFunction.h"
+#include "llvm/CodeGen/Register.h"
 #include <optional>
 
 namespace llvm {
@@ -74,6 +76,17 @@ class AIEMachineFunctionInfo : public MachineFunctionInfo {
 
   const TileMemoryPSV TileMemory;
 
+  /// Side map: descendant vreg -> "logical group original" (the VRM Original
+  /// before the split-from chain was deliberately severed by an AIE
+  /// register-rewriter pass).
+  ///
+  /// The chain is severed for correctness so that SplitKit::defFromParent
+  /// no longer rematerializes through the now-stale ancestor LiveInterval.
+  /// However, InlineSpiller still wants to share a stack slot among all
+  /// descendants of that group. This map remembers the pre-severance original
+  /// so getSpillGroupOriginal() can answer that lookup for InlineSpiller.
+  DenseMap<Register, Register> SpillGroupOriginal;
+
 public:
   //  AIEMachineFunctionInfo() = default;
 
@@ -89,6 +102,24 @@ class AIEMachineFunctionInfo : public MachineFunctionInfo {
   unsigned getBytesInStackArgArea() const { return BytesInStackArgArea; }
   void setBytesInStackArgArea(unsigned bytes) { BytesInStackArgArea = bytes; }
 
+  /// Record that \p Descendant logically belongs to the spill group whose
+  /// "original" is \p OldOriginal. Called by AIE register-rewriter passes
+  /// just before they sever the VRM split-from chain for \p Descendant.
+  void recordSpillGroupOriginal(Register Descendant, Register OldOriginal) {
+    SpillGroupOriginal[Descendant] = OldOriginal;
+  }
+
+  /// Return the "logical group original" recorded for \p V, if any. Used by
+  /// the AIEBaseSubtarget override of
+  /// TargetSubtargetInfo::getSpillGroupOriginal which InlineSpiller consults
+  /// for stack-slot sharing.
+  std::optional<Register> getSpillGroupOriginal(Register V) const {
+    auto It = SpillGroupOriginal.find(V);
+    if (It == SpillGroupOriginal.end())
+      return std::nullopt;
+    return It->second;
+  }
+
   MachineFunctionInfo *
   clone(BumpPtrAllocator &Allocator, MachineFunction &DestMF,
         const DenseMap<MachineBasicBlock *, MachineBasicBlock *> &Src2DstMBB)

llvm/lib/Target/AIE/AIESuperRegRewriter.cpp

@@ -4,7 +4,7 @@
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
-// (c) Copyright 2023-2025 Advanced Micro Devices, Inc. or its affiliates
+// (c) Copyright 2023-2026 Advanced Micro Devices, Inc. or its affiliates
 //
 //===----------------------------------------------------------------------===//
 
@@ -114,6 +114,11 @@ bool AIESuperRegRewriter::runOnMachineFunction(MachineFunction &MF) {
     }
   }
 
+  // Snapshot Originals whose LI is about to go stale.
+  SmallSet<Register, 8> TaintedOriginals;
+  for (auto &[VReg, _] : AssignedPhysRegs)
+    TaintedOriginals.insert(VRM.getOriginal(VReg));
+
   // Re-write all the collected VRegs
   for (auto &[VReg, PhysRegAndSubRegs] : AssignedPhysRegs) {
     const Register PhysReg = PhysRegAndSubRegs.first;
@@ -122,6 +127,9 @@ bool AIESuperRegRewriter::runOnMachineFunction(MachineFunction &MF) {
                                       LRM, LIS, Indexes, DebugVars);
   }
 
+  // Prevent SplitKit from rematerializing through stale ancestor LIs.
+  AIESuperRegUtils::clearStaleSplitFromMappings(TaintedOriginals, MF, MRI, VRM);
+
   LLVM_DEBUG(VRM.dump());
   return !AssignedPhysRegs.empty();
 }

llvm/lib/Target/AIE/AIESuperRegUtils.cpp

@@ -10,6 +10,7 @@
 #include "AIESuperRegUtils.h"
 #include "AIEBaseInstrInfo.h"
 #include "AIEBaseRegisterInfo.h"
+#include "AIEMachineFunctionInfo.h"
 #include "llvm/CodeGen/LiveDebugVariables.h"
 #include "llvm/CodeGen/LiveIntervals.h"
 #include "llvm/CodeGen/LiveRegMatrix.h"
@@ -19,6 +20,7 @@
 #include "llvm/CodeGen/MachineRegisterInfo.h"
 #include "llvm/CodeGen/SlotIndexes.h"
 #include "llvm/CodeGen/TargetInstrInfo.h"
+#include "llvm/CodeGen/VirtRegMap.h"
 #include "llvm/Support/Debug.h"
 
 #define DEBUG_TYPE "aie-ra"
@@ -376,16 +378,6 @@ void rewriteSuperReg(Register Reg, std::optional<Register> AssignedPhysReg,
   // Step 4: Remove the original register's live interval
   LIS.removeInterval(Reg);
 
-  // Step 4b: Clear stale ancestor live intervals. The operand rewrite in
-  // step 3 modified instructions in-place (e.g., stripping sub-register
-  // indices). Any ancestor register in the VRM split chain still has VNInfos
-  // pointing to those instruction slots. If a later Greedy pass traces back
-  // via VRM.getOriginal(), it would find a stale instruction and could produce
-  // an invalid rematerialization. Clearing the ancestor interval prevents this.
-  Register Original = VRM.getOriginal(Reg);
-  if (Original != Reg && LIS.hasInterval(Original))
-    LIS.getInterval(Original).clear();
-
   // Step 5: Filter out empty subregisters
   markEffectiveEmptyCopiesDead(SubRegToVReg, MRI, TRI, LIS);
 
@@ -413,6 +405,42 @@ bool isRegUsedBy2DOr3DInstruction(const MachineRegisterInfo &MRI,
       });
 }
 
+void clearStaleSplitFromMappings(const SmallSet<Register, 8> &TaintedOriginals,
+                                 MachineFunction &MF, MachineRegisterInfo &MRI,
+                                 VirtRegMap &VRM) {
+  if (TaintedOriginals.empty())
+    return;
+
+  // Record the pre-severance "logical group original" in the target-side
+  // side map so that InlineSpiller (via TargetSubtargetInfo::
+  // getSpillGroupOriginal) can still merge sibling spills onto a shared
+  // stack slot after we cut the VRM split-from chain below.
+  auto *MFI = MF.getInfo<AIEMachineFunctionInfo>();
+
+  for (unsigned I = 0, E = MRI.getNumVirtRegs(); I != E; ++I) {
+    const Register V = Register::index2VirtReg(I);
+    if (MRI.reg_nodbg_empty(V))
+      continue;
+    const Register Orig = VRM.getPreSplitReg(V);
+    if (!Orig || !TaintedOriginals.count(Orig))
+      continue;
+
+    LLVM_DEBUG({
+      const TargetRegisterInfo *TRI = MRI.getTargetRegisterInfo();
+      dbgs() << "  Clearing stale split-from for " << printReg(V, TRI, 0, &MRI)
+             << " (was split from " << printReg(Orig, TRI, 0, &MRI)
+             << "); recorded for spill-group sharing\n";
+    });
+    // Remember the chain so InlineSpiller can still group V's spills with
+    // the rest of Orig's descendants on a shared stack slot.
+    if (MFI)
+      MFI->recordSpillGroupOriginal(V, Orig);
+    // Restore V to the canonical "no split parent" state so getOriginal(V)==V
+    // and SplitKit::defFromParent stops consulting the (stale) ancestor LI.
+    VRM.clearSplitFromReg(V);
+  }
+}
+
 void repairLiveIntervals(SmallSet<Register, 8> &RegistersToRepair,
                          VirtRegMap &VRM, LiveRegMatrix &LRM,
                          LiveIntervals &LIS) {

llvm/lib/Target/AIE/AIESuperRegUtils.h

@@ -19,6 +19,7 @@
 
 namespace llvm {
 class Register;
+class MachineFunction;
 class MachineRegisterInfo;
 struct AIEBaseRegisterInfo;
 class MachineInstr;
@@ -84,6 +85,29 @@ void repairLiveIntervals(SmallSet<Register, 8> &RegistersToRepair,
                          VirtRegMap &VRM, LiveRegMatrix &LRM,
                          LiveIntervals &LIS);
 
+/// Sever VRM split-from chain for descendants of \p TaintedOriginals so that
+/// SplitKit::defFromParent consults the descendant's own (repaired) LI, not
+/// the stale ancestor LI which may still hold VNs at slots whose MIs were
+/// rewritten/unbundled by an AIE register-rewriter pass. Each affected
+/// descendant is restored via VRM.clearSplitFromReg() to the canonical
+/// "no split parent" state of a freshly created vreg.
+///
+///   before:                    after:
+///   %0 (stale LI)              %0 (stale LI, ignored)
+///    | split-from               x  (chain cut)
+///   %35 ----.                  %35 (no split parent)
+///    | split-from
+///   %141..%144 (future split   greedy splits will use %35's LI
+///   would consult %0's LI)     instead of %0's
+///
+/// Before severing, the pre-severance Original is recorded in
+/// AIEMachineFunctionInfo's spill-group side map so that InlineSpiller (via
+/// TargetSubtargetInfo::getSpillGroupOriginal) can still merge sibling spills
+/// of these descendants onto a shared stack slot.
+void clearStaleSplitFromMappings(const SmallSet<Register, 8> &TaintedOriginals,
+                                 MachineFunction &MF, MachineRegisterInfo &MRI,
+                                 VirtRegMap &VRM);
+
 } // namespace llvm::AIESuperRegUtils
 
 #endif

llvm/lib/Target/AIE/AIEUnallocatedSuperRegRewriter.cpp

@@ -229,13 +229,23 @@ bool AIEUnallocatedSuperRegRewriter::runOnMachineFunction(MachineFunction &MF) {
     return false;
   }
 
+  // Snapshot Originals whose LI is about to go stale.
+  SmallSet<Register, 8> TaintedOriginals;
+  for (auto &P : Info.ExpandableRegs)
+    TaintedOriginals.insert(VRM.getOriginal(P.first));
+  for (auto &P : Info.RewritableRegs)
+    TaintedOriginals.insert(VRM.getOriginal(P.first));
+
   LLVM_DEBUG(dbgs() << "Expanding copy bundles...\n");
   expandCopyBundles(Info.ExpandableRegs, MRI, Indexes, LIS, VRM, LRM);
 
   LLVM_DEBUG(dbgs() << "Performing register rewrites...\n");
   rewriteCandidates(Info.RewritableRegs, MRI, TRI, VRM, LRM, LIS, Indexes,
                     DebugVars);
 
+  // Prevent SplitKit from rematerializing through stale ancestor LIs.
+  AIESuperRegUtils::clearStaleSplitFromMappings(TaintedOriginals, MF, MRI, VRM);
+
   LLVM_DEBUG(dbgs() << "Successfully rewrote " << Info.RewritableRegs.size()
                     << " register(s)\n");