[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 created https://github.com/llvm/llvm-project/pull/71472

`instcombine` currently generates large arrays when the `NumElements` argument of an `alloca` instruction is negative. Such large arrays may cause the size constant to overflow during code generation under 32 bit mode, leading to a crash. This PR limits the constant's bit width to the width of the pointer on the target. With this fix, 
```
alloca i32, i32 -1
```
and
```
alloca [4294967295 x i32], i32 1
```
generates the exact same PowerPC assembly code under 32 bit mode. 

>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] 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
+}