[llvm] [RegAlloc] Fix use-after-free in `RegAllocBase::cleanupFailedVReg` (PR #151435)

[Shilei Tian via llvm-commits]

https://github.com/shiltian updated https://github.com/llvm/llvm-project/pull/151435

>From 4f1c8bc26102c057ad383d4d7cc9c5067f8d63e7 Mon Sep 17 00:00:00 2001
From: Shilei Tian <i at tianshilei.me>
Date: Wed, 30 Jul 2025 22:03:34 -0400
Subject: [PATCH 1/2] [RegAlloc] Fix use-after-free in
 `RegAllocBase::cleanupFailedVReg`

Since #128400 already mentions it's not clear about the necessity of removing intervals from regunits, this PR avoids the issue by simply skipping that step.

Fixes SWDEV-527146.
---
 llvm/lib/CodeGen/RegAllocBase.cpp                 |  4 +---
 .../use-after-free-after-cleanup-failed-vreg.ll   | 15 +++++++++++++++
 2 files changed, 16 insertions(+), 3 deletions(-)
 create mode 100644 llvm/test/CodeGen/AMDGPU/use-after-free-after-cleanup-failed-vreg.ll

diff --git a/llvm/lib/CodeGen/RegAllocBase.cpp b/llvm/lib/CodeGen/RegAllocBase.cpp
index 69b92917399fd..2400a1feea26e 100644
--- a/llvm/lib/CodeGen/RegAllocBase.cpp
+++ b/llvm/lib/CodeGen/RegAllocBase.cpp
@@ -178,10 +178,8 @@ void RegAllocBase::cleanupFailedVReg(Register FailedReg, MCRegister PhysReg,
     for (MCRegAliasIterator Aliases(PhysReg, TRI, true); Aliases.isValid();
          ++Aliases) {
       for (MachineOperand &MO : MRI->reg_operands(*Aliases)) {
-        if (MO.readsReg()) {
+        if (MO.readsReg())
           MO.setIsUndef(true);
-          LIS->removeAllRegUnitsForPhysReg(MO.getReg());
-        }
       }
     }
   }
diff --git a/llvm/test/CodeGen/AMDGPU/use-after-free-after-cleanup-failed-vreg.ll b/llvm/test/CodeGen/AMDGPU/use-after-free-after-cleanup-failed-vreg.ll
new file mode 100644
index 0000000000000..e5c239b178366
--- /dev/null
+++ b/llvm/test/CodeGen/AMDGPU/use-after-free-after-cleanup-failed-vreg.ll
@@ -0,0 +1,15 @@
+; RUN: not llc -mcpu=gfx1100 -mtriple=amdgcn-amd-amdhsa -stress-regalloc=4 %s 2>&1 | FileCheck %s
+
+; CHECK: ran out of registers during register allocation in function 'f'
+
+define <16 x half> @f(i1 %LGV2, <16 x half> %0) {
+BB:
+  br i1 %LGV2, label %SW_C3, label %SW_C
+
+SW_C:                                             ; preds = %BB
+  %B1 = fmul <16 x half> %0, zeroinitializer
+  ret <16 x half> %B1
+
+SW_C3:                                            ; preds = %BB
+  ret <16 x half> <half 0xH0000, half poison, half poison, half poison, half poison, half poison, half poison, half poison, half poison, half poison, half poison, half poison, half poison, half poison, half poison, half poison>
+}

>From 8a254cdc87cdaaba7541cfaaa451ea31a1de8e50 Mon Sep 17 00:00:00 2001
From: Shilei Tian <i at tianshilei.me>
Date: Wed, 30 Jul 2025 22:37:57 -0400
Subject: [PATCH 2/2] fix comments

---
 .../AMDGPU/use-after-free-after-cleanup-failed-vreg.ll        | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/llvm/test/CodeGen/AMDGPU/use-after-free-after-cleanup-failed-vreg.ll b/llvm/test/CodeGen/AMDGPU/use-after-free-after-cleanup-failed-vreg.ll
index e5c239b178366..2eb6bf7020418 100644
--- a/llvm/test/CodeGen/AMDGPU/use-after-free-after-cleanup-failed-vreg.ll
+++ b/llvm/test/CodeGen/AMDGPU/use-after-free-after-cleanup-failed-vreg.ll
@@ -1,6 +1,6 @@
-; RUN: not llc -mcpu=gfx1100 -mtriple=amdgcn-amd-amdhsa -stress-regalloc=4 %s 2>&1 | FileCheck %s
+; RUN: not llc -mcpu=gfx1100 -mtriple=amdgcn-amd-amdhsa -stress-regalloc=4 -filetype=null -verify-machineinstrs %s 2>&1 | FileCheck %s
 
-; CHECK: ran out of registers during register allocation in function 'f'
+; CHECK: error: <unknown>:0:0: ran out of registers during register allocation in function 'f'
 
 define <16 x half> @f(i1 %LGV2, <16 x half> %0) {
 BB: