[llvm] [regalloc][LiveRegMatrix] Add validity check for LiveRegMatrix to prevent dangling pointers (PR #168553)

[Valery Pykhtin via llvm-commits]

https://github.com/vpykhtin created https://github.com/llvm/llvm-project/pull/168553

- Implemented `isValid()` method in LiveRegMatrix to verify that all LiveInterval pointers are valid.
- Added assertion in RegAllocBase's `postOptimization()` to ensure no dangling pointers exist in LiveRegMatrix after spilling.

This is for testing only, do not merge. It will be submitted as part of ongoing fix to prevent test failures.

>From 5cb8212764d4c3940c08def7bea8293787c26b98 Mon Sep 17 00:00:00 2001
From: Valery Pykhtin <valery.pykhtin at amd.com>
Date: Fri, 7 Nov 2025 14:49:59 +0000
Subject: [PATCH] [regalloc] Add validity check for LiveRegMatrix to prevent
 dangling pointers

- Implemented `isValid()` method in LiveRegMatrix to verify that all LiveInterval pointers are valid.
- Added assertion in RegAllocBase's `postOptimization()` to ensure no dangling pointers exist in LiveRegMatrix after spilling.
---
 llvm/include/llvm/CodeGen/LiveRegMatrix.h | 10 ++++++++
 llvm/lib/CodeGen/LiveRegMatrix.cpp        | 29 +++++++++++++++++++++++
 llvm/lib/CodeGen/RegAllocBase.cpp         |  4 ++++
 3 files changed, 43 insertions(+)

diff --git a/llvm/include/llvm/CodeGen/LiveRegMatrix.h b/llvm/include/llvm/CodeGen/LiveRegMatrix.h
index 35add577d071a..a4a7c0375b9ab 100644
--- a/llvm/include/llvm/CodeGen/LiveRegMatrix.h
+++ b/llvm/include/llvm/CodeGen/LiveRegMatrix.h
@@ -170,6 +170,16 @@ class LiveRegMatrix {
   }
 
   Register getOneVReg(unsigned PhysReg) const;
+
+  /// Verify that all LiveInterval pointers in the matrix are valid.
+  /// This checks that each LiveInterval referenced in LiveIntervalUnion
+  /// actually exists in LiveIntervals and is not a dangling pointer.
+  /// Returns true if the matrix is valid, false if dangling pointers are found.
+  /// This is primarily useful for debugging heap-use-after-free issues.
+  /// This method uses a lazy approach - it builds a set of valid LiveInterval
+  /// pointers on-demand and has zero runtime/memory overhead during normal
+  /// register allocation.
+  bool isValid() const;
 };
 
 class LiveRegMatrixWrapperLegacy : public MachineFunctionPass {
diff --git a/llvm/lib/CodeGen/LiveRegMatrix.cpp b/llvm/lib/CodeGen/LiveRegMatrix.cpp
index e7238008d2c69..0f9571ea81389 100644
--- a/llvm/lib/CodeGen/LiveRegMatrix.cpp
+++ b/llvm/lib/CodeGen/LiveRegMatrix.cpp
@@ -12,11 +12,13 @@
 
 #include "llvm/CodeGen/LiveRegMatrix.h"
 #include "RegisterCoalescer.h"
+#include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/Statistic.h"
 #include "llvm/CodeGen/LiveInterval.h"
 #include "llvm/CodeGen/LiveIntervalUnion.h"
 #include "llvm/CodeGen/LiveIntervals.h"
 #include "llvm/CodeGen/MachineFunction.h"
+#include "llvm/CodeGen/MachineRegisterInfo.h"
 #include "llvm/CodeGen/TargetRegisterInfo.h"
 #include "llvm/CodeGen/TargetSubtargetInfo.h"
 #include "llvm/CodeGen/VirtRegMap.h"
@@ -290,6 +292,33 @@ Register LiveRegMatrix::getOneVReg(unsigned PhysReg) const {
   return MCRegister::NoRegister;
 }
 
+bool LiveRegMatrix::isValid() const {
+  // Build set of all valid LiveInterval pointers from LiveIntervals.
+  DenseSet<LiveInterval *> ValidIntervals;
+  for (unsigned RegIdx = 0, NumRegs = VRM->getRegInfo().getNumVirtRegs();
+       RegIdx < NumRegs; ++RegIdx) {
+    Register VReg = Register::index2VirtReg(RegIdx);
+    // Only track assigned registers since unassigned ones won't be in Matrix
+    if (VRM->hasPhys(VReg) && LIS->hasInterval(VReg))
+      ValidIntervals.insert(&LIS->getInterval(VReg));
+  }
+
+  // Now scan all LiveIntervalUnions in the matrix and verify each pointer
+  unsigned NumDanglingPointers = 0;
+  for (unsigned I = 0, Size = Matrix.size(); I != Size; ++I) {
+    MCRegUnit Unit = static_cast<MCRegUnit>(I);
+    for (const LiveInterval *LI : Matrix[Unit]) {
+      if (!ValidIntervals.contains(LI)) {
+        ++NumDanglingPointers;
+        dbgs() << "ERROR: LiveInterval pointer is not found in LiveIntervals:\n"
+               << "  Register Unit: " << printRegUnit(Unit, TRI) << '\n'
+               << "  LiveInterval pointer: " << LI << '\n';
+      }
+    }
+  }
+  return NumDanglingPointers == 0;
+}
+
 AnalysisKey LiveRegMatrixAnalysis::Key;
 
 LiveRegMatrix LiveRegMatrixAnalysis::run(MachineFunction &MF,
diff --git a/llvm/lib/CodeGen/RegAllocBase.cpp b/llvm/lib/CodeGen/RegAllocBase.cpp
index 2400a1feea26e..9ba6007fc1cb9 100644
--- a/llvm/lib/CodeGen/RegAllocBase.cpp
+++ b/llvm/lib/CodeGen/RegAllocBase.cpp
@@ -155,6 +155,10 @@ void RegAllocBase::allocatePhysRegs() {
 
 void RegAllocBase::postOptimization() {
   spiller().postOptimization();
+
+  // Verify LiveRegMatrix after spilling (no dangling pointers).
+  assert(Matrix->isValid() && "LiveRegMatrix validation failed");
+
   for (auto *DeadInst : DeadRemats) {
     LIS->RemoveMachineInstrFromMaps(*DeadInst);
     DeadInst->eraseFromParent();