[llvm] [CodeGen] Handling Oversized Alloca Types under 32 bit Mode to Avoid Code Generator Crash (PR #71472)

[Qiongsi Wu via llvm-commits]

https://github.com/qiongsiwu updated https://github.com/llvm/llvm-project/pull/71472

>From 2e98ef4542f04bebad5fc02ce4491710399bcabd Mon Sep 17 00:00:00 2001
From: Qiongsi Wu <qwu at ibm.com>
Date: Mon, 6 Nov 2023 19:17:55 -0500
Subject: [PATCH 1/2] Fix code generator crash when it sees an oversized alloca
 type

---
 .../SelectionDAG/SelectionDAGBuilder.cpp      |  7 +--
 llvm/test/CodeGen/PowerPC/alloca-neg-size.ll  | 46 +++++++++++++++++++
 2 files changed, 50 insertions(+), 3 deletions(-)
 create mode 100644 llvm/test/CodeGen/PowerPC/alloca-neg-size.ll

diff --git a/llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp b/llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp
index aab0d5c5a348bfe..d5ffaf28ca2d499 100644
--- a/llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp
@@ -4138,9 +4138,10 @@ void SelectionDAGBuilder::visitAlloca(const AllocaInst &I) {
                                           APInt(IntPtr.getScalarSizeInBits(),
                                                 TySize.getKnownMinValue())));
   else
-    AllocSize =
-        DAG.getNode(ISD::MUL, dl, IntPtr, AllocSize,
-                    DAG.getConstant(TySize.getFixedValue(), dl, IntPtr));
+    AllocSize = DAG.getNode(ISD::MUL, dl, IntPtr, AllocSize,
+                            DAG.getConstant(APInt(IntPtr.getScalarSizeInBits(),
+                                                  TySize.getFixedValue()),
+                                            dl, IntPtr));
 
   // Handle alignment.  If the requested alignment is less than or equal to
   // the stack alignment, ignore it.  If the size is greater than or equal to
diff --git a/llvm/test/CodeGen/PowerPC/alloca-neg-size.ll b/llvm/test/CodeGen/PowerPC/alloca-neg-size.ll
new file mode 100644
index 000000000000000..ba22c0a71294b8d
--- /dev/null
+++ b/llvm/test/CodeGen/PowerPC/alloca-neg-size.ll
@@ -0,0 +1,46 @@
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 3
+; The instcombine pass can turn
+;     alloca i32, i32 -1
+; to
+;     alloca [4294967295 x i32], i32 1
+; because it zero extends the NumElements to unit64_t.
+; The zero extension can lead to oversized arrays on a 32 bit system.
+; Alloca-ing an array of size bigger than half of the address space
+; is most likely an undefined behaviour, but the code generator
+; should not crash in such situations.
+; RUN: llc < %s -mtriple=powerpc-ibm-aix-xcoff | FileCheck %s
+define void @test_negalloc(ptr %dst, i32 %cond) {
+; CHECK-LABEL: test_negalloc:
+; CHECK:       # %bb.0: # %entry
+; CHECK-NEXT:    stw 31, -4(1)
+; CHECK-NEXT:    stwu 1, -80(1)
+; CHECK-NEXT:    cmplwi 4, 0
+; CHECK-NEXT:    mr 31, 1
+; CHECK-NEXT:    beq 0, L..BB0_2
+; CHECK-NEXT:  # %bb.1: # %if.then
+; CHECK-NEXT:    li 4, 0
+; CHECK-NEXT:    addi 5, 31, 80
+; CHECK-NEXT:    stwux 5, 1, 4
+; CHECK-NEXT:    addi 4, 1, 32
+; CHECK-NEXT:    b L..BB0_3
+; CHECK-NEXT:  L..BB0_2:
+; CHECK-NEXT:    addi 4, 31, 44
+; CHECK-NEXT:  L..BB0_3: # %if.end
+; CHECK-NEXT:    stw 4, 0(3)
+; CHECK-NEXT:    lwz 1, 0(1)
+; CHECK-NEXT:    lwz 31, -4(1)
+; CHECK-NEXT:    blr
+entry:
+  %0 = alloca [8 x i32], i32 1, align 4
+  %tobool = icmp ne i32 %cond, 0
+  br i1 %tobool, label %if.then, label %if.end
+
+if.then:
+  %vla1 = alloca [4294967295 x i32], i32 1, align 4
+  br label %if.end
+
+if.end:
+  %arr = phi ptr [%0, %entry], [%vla1, %if.then]
+  store ptr %arr, ptr %dst
+  ret void
+}

>From e978491b85ab5a31ed2ee6e19acdfcf326a4ba93 Mon Sep 17 00:00:00 2001
From: Qiongsi Wu <qwu at ibm.com>
Date: Tue, 7 Nov 2023 09:19:09 -0500
Subject: [PATCH 2/2] Address review comments in the test

---
 .../{alloca-neg-size.ll => alloca-oversized.ll}    | 14 +++++---------
 1 file changed, 5 insertions(+), 9 deletions(-)
 rename llvm/test/CodeGen/PowerPC/{alloca-neg-size.ll => alloca-oversized.ll} (72%)

diff --git a/llvm/test/CodeGen/PowerPC/alloca-neg-size.ll b/llvm/test/CodeGen/PowerPC/alloca-oversized.ll
similarity index 72%
rename from llvm/test/CodeGen/PowerPC/alloca-neg-size.ll
rename to llvm/test/CodeGen/PowerPC/alloca-oversized.ll
index ba22c0a71294b8d..e102131d397218e 100644
--- a/llvm/test/CodeGen/PowerPC/alloca-neg-size.ll
+++ b/llvm/test/CodeGen/PowerPC/alloca-oversized.ll
@@ -1,16 +1,12 @@
 ; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 3
-; The instcombine pass can turn
+;
 ;     alloca i32, i32 -1
-; to
+; and
 ;     alloca [4294967295 x i32], i32 1
-; because it zero extends the NumElements to unit64_t.
-; The zero extension can lead to oversized arrays on a 32 bit system.
-; Alloca-ing an array of size bigger than half of the address space
-; is most likely an undefined behaviour, but the code generator
-; should not crash in such situations.
+; are equivalent and the assembly code sequences generated are the same.
 ; RUN: llc < %s -mtriple=powerpc-ibm-aix-xcoff | FileCheck %s
-define void @test_negalloc(ptr %dst, i32 %cond) {
-; CHECK-LABEL: test_negalloc:
+define void @test_oversized(ptr %dst, i32 %cond) {
+; CHECK-LABEL: test_oversized:
 ; CHECK:       # %bb.0: # %entry
 ; CHECK-NEXT:    stw 31, -4(1)
 ; CHECK-NEXT:    stwu 1, -80(1)