[llvm] [AMDGPU] Create local KnownBits in case DenseMap gets invalidated (PR #111568)

[Janek van Oirschot via llvm-commits]

https://github.com/JanekvO created https://github.com/llvm/llvm-project/pull/111568

KnownBits retrieved from DenseMap may invalidate if insertion requires a (re)growth.

Fixes https://github.com/llvm/llvm-project/issues/110930

>From e241240e9d0d58d15fded9b1e9d6b5c3f0318569 Mon Sep 17 00:00:00 2001
From: Janek van Oirschot <janek.vanoirschot at amd.com>
Date: Tue, 8 Oct 2024 18:34:32 +0100
Subject: [PATCH] [AMDGPU] Create local KnownBits in case DenseMap gets
 invalidated

---
 .../AMDGPU/MCTargetDesc/AMDGPUMCExpr.cpp      |   8 +-
 .../AMDGPU/mcexpr-knownbits-assign-crash.ll   | 120 ++++++++++++++++++
 2 files changed, 126 insertions(+), 2 deletions(-)
 create mode 100644 llvm/test/CodeGen/AMDGPU/mcexpr-knownbits-assign-crash.ll

diff --git a/llvm/lib/Target/AMDGPU/MCTargetDesc/AMDGPUMCExpr.cpp b/llvm/lib/Target/AMDGPU/MCTargetDesc/AMDGPUMCExpr.cpp
index 4fbd7d0f889457..9f5f4399679376 100644
--- a/llvm/lib/Target/AMDGPU/MCTargetDesc/AMDGPUMCExpr.cpp
+++ b/llvm/lib/Target/AMDGPU/MCTargetDesc/AMDGPUMCExpr.cpp
@@ -537,8 +537,12 @@ static void knownBitsMapHelper(const MCExpr *Expr, KnownBitsMap &KBM,
 
     // Variable value retrieval is not for actual use but only for knownbits
     // analysis.
-    knownBitsMapHelper(Sym.getVariableValue(/*SetUsed=*/false), KBM, Depth + 1);
-    KBM[Expr] = KBM[Sym.getVariableValue(/*SetUsed=*/false)];
+    const MCExpr *SymVal = Sym.getVariableValue(/*setUsed=*/false);
+    knownBitsMapHelper(SymVal, KBM, Depth + 1);
+
+    // Explicitly copy-construct so that there exists a local KnownBits in case
+    // KBM[SymVal] gets invalidated after a potential growth through KBM[Expr].
+    KBM[Expr] = KnownBits(KBM[SymVal]);
     return;
   }
   case MCExpr::ExprKind::Unary: {
diff --git a/llvm/test/CodeGen/AMDGPU/mcexpr-knownbits-assign-crash.ll b/llvm/test/CodeGen/AMDGPU/mcexpr-knownbits-assign-crash.ll
new file mode 100644
index 00000000000000..830996d2e24023
--- /dev/null
+++ b/llvm/test/CodeGen/AMDGPU/mcexpr-knownbits-assign-crash.ll
@@ -0,0 +1,120 @@
+; REQUIRES: asserts
+; RUN: llc -mtriple=amdgcn-amd-amdhsa -mcpu=gfx1030 -o - %s 2>&1 | FileCheck %s
+
+; CHECK-NOT: Assertion `BitWidth == RHS.BitWidth && "Bit widths must be same for comparison"' failed
+
+define void @RunTic() {
+  %call5.i1 = call i32 @G_CheckDemoStatus()
+  tail call void @D_AdvanceDemo()
+  call void @G_Ticker()
+  ret void
+}
+
+define void @G_Ticker() {
+  call void @G_DoReborn()
+  tail call void @F_Ticker()
+  tail call void @AM_Stop()
+  tail call void @F_StartFinale()
+  tail call void @D_AdvanceDemo()
+  %call.i.i449 = call i32 @R_FlatNumForName()
+  %call9.i.i = call i32 @R_TextureNumForName()
+  %call.i306 = tail call ptr @P_TempSaveGameFile()
+  %call1.i307 = call ptr @P_SaveGameFile()
+  call void (...) @I_Error()
+  ret void
+}
+
+define void @G_DoReborn() {
+  call void @P_RemoveMobj()
+  call void @P_SpawnMobj()
+  call void @P_SpawnPlayer()
+  call void (...) @I_Error()
+  ret void
+}
+
+define void @AM_Stop() {
+  ret void
+}
+
+define void @D_AdvanceDemo() {
+  ret void
+}
+
+define void @F_StartFinale() {
+  ret void
+}
+
+define void @F_Ticker() {
+  ret void
+}
+
+define void @G_PlayerReborn() {
+  ret void
+}
+
+define i32 @G_CheckDemoStatus() {
+  tail call void @I_Quit()
+  tail call void @D_AdvanceDemo()
+  call void (...) @I_Error()
+  ret i32 0
+}
+
+define void @HU_Start() {
+  ret void
+}
+
+define void @I_Quit() {
+  %fptr = load ptr, ptr null, align 8
+  tail call void %fptr()
+  ret void
+}
+
+define void @P_SetThingPosition() {
+  ret void
+}
+
+define void @P_RemoveMobj() {
+  ret void
+}
+
+define void @P_SpawnMobj() {
+  ret void
+}
+
+define void @P_SpawnPlayer() {
+  call void @G_PlayerReborn()
+  call void @P_SetThingPosition()
+  call void @P_SetupPsprites(ptr addrspace(1) null)
+  tail call void @HU_Start()
+  ret void
+}
+
+define void @P_SetupPsprites(ptr addrspace(1) %i) {
+  %fptr = load ptr, ptr addrspace(1) %i, align 8
+  tail call void %fptr()
+  ret void
+}
+
+define ptr @P_TempSaveGameFile() {
+  ret ptr null
+}
+
+define ptr @P_SaveGameFile() {
+  ret ptr null
+}
+
+define i32 @R_TextureNumForName() {
+  %ret = call i32 @R_FlatNumForName()
+  ret i32 0
+}
+
+define void @I_Error(...) {
+  %fptr = load ptr, ptr null, align 8
+  call void %fptr()
+  ret void
+}
+
+define i32 @R_FlatNumForName() {
+  call void (...) @I_Error()
+  unreachable
+}