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

[Shilei Tian via llvm-commits]

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

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.

>From ce763231946f5f9c8103d3d13c5b94245e35b57b 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] [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..03477e0c95523
--- /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 undef, half undef, half undef, half undef, half undef, half undef, half undef, half undef, half undef, half undef, half undef, half undef, half undef, half undef, half undef>
+}