[llvm] [InstCombine] Avoid Allocating Arrays Too Large For the Target (PR #70980)

[Qiongsi Wu via llvm-commits]

https://github.com/qiongsiwu created https://github.com/llvm/llvm-project/pull/70980

The current logic in `simplifyAllocaArraySize` may result in an array to large for a target. For example, `%1 = alloca i32, i32 -1` may result in an array type of `[4294967295 x i32]`. Such an array can cause a crash in the code generator. This PR adds logic to check the size of the array type, and sets the array size to 0 if the array is too big for a target. This behaviour is consistent with what LLVM does when an `alloca`'s first operand is too large for the target. 

>From 2ac39f713f58f01095707e907a1db50cf58d87cf Mon Sep 17 00:00:00 2001
From: Qiongsi Wu <qwu at ibm.com>
Date: Wed, 1 Nov 2023 16:18:15 -0400
Subject: [PATCH] Stop instcombine from alloca-ing arrays that are too large

---
 .../InstCombineLoadStoreAlloca.cpp            | 15 ++++++++++---
 .../test/Transforms/InstCombine/alloca-big.ll | 22 ++++++++++++++++++-
 2 files changed, 33 insertions(+), 4 deletions(-)

diff --git a/llvm/lib/Transforms/InstCombine/InstCombineLoadStoreAlloca.cpp b/llvm/lib/Transforms/InstCombine/InstCombineLoadStoreAlloca.cpp
index b72b68c68d985bd..8d90fc9b63ca049 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineLoadStoreAlloca.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineLoadStoreAlloca.cpp
@@ -210,14 +210,23 @@ static Instruction *simplifyAllocaArraySize(InstCombinerImpl &IC,
   if (const ConstantInt *C = dyn_cast<ConstantInt>(AI.getArraySize())) {
     if (C->getValue().getActiveBits() <= 64) {
       Type *NewTy = ArrayType::get(AI.getAllocatedType(), C->getZExtValue());
+
+      // Make sure we do not create an array type larger than pointers on the
+      // target can index.
+      unsigned MaxArrSizeBitWidth =
+          IC.getDataLayout().getPointerTypeSizeInBits(AI.getType());
+      APInt ArrayAllocSize(64, IC.getDataLayout().getTypeAllocSize(NewTy));
+      if (ArrayAllocSize.getActiveBits() > MaxArrSizeBitWidth)
+        NewTy = ArrayType::get(AI.getAllocatedType(), 0);
+
       AllocaInst *New = IC.Builder.CreateAlloca(NewTy, AI.getAddressSpace(),
                                                 nullptr, AI.getName());
       New->setAlignment(AI.getAlign());
 
       replaceAllDbgUsesWith(AI, *New, *New, DT);
 
-      // Scan to the end of the allocation instructions, to skip over a block of
-      // allocas if possible...also skip interleaved debug info
+      // Scan to the end of the allocation instructions, to skip over a block
+      // of allocas if possible...also skip interleaved debug info
       //
       BasicBlock::iterator It(New);
       while (isa<AllocaInst>(*It) || isa<DbgInfoIntrinsic>(*It))
@@ -243,7 +252,7 @@ static Instruction *simplifyAllocaArraySize(InstCombinerImpl &IC,
     return IC.replaceInstUsesWith(AI, Constant::getNullValue(AI.getType()));
 
   // Ensure that the alloca array size argument has type equal to the offset
-  // size of the alloca() pointer, which, in the tyical case, is intptr_t,
+  // size of the alloca() pointer, which, in the typical case, is intptr_t,
   // so that any casting is exposed early.
   Type *PtrIdxTy = IC.getDataLayout().getIndexType(AI.getType());
   if (AI.getArraySize()->getType() != PtrIdxTy) {
diff --git a/llvm/test/Transforms/InstCombine/alloca-big.ll b/llvm/test/Transforms/InstCombine/alloca-big.ll
index 6925f1ba988dc1e..2ccfb5afa67493d 100644
--- a/llvm/test/Transforms/InstCombine/alloca-big.ll
+++ b/llvm/test/Transforms/InstCombine/alloca-big.ll
@@ -1,5 +1,6 @@
 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
-; RUN: opt < %s -passes=instcombine -S | FileCheck %s
+; RUN: opt < %s -passes=instcombine -data-layout="p:64:64" -S | FileCheck %s
+; RUN: opt < %s -passes=instcombine -data-layout="p:32:32" -S | FileCheck --check-prefix=CHECK32 %s
 
 ; OSS-Fuzz: https://bugs.chromium.org/p/oss-fuzz/issues/detail?id=5223
 define void @test_bigalloc(ptr %dst) {
@@ -7,8 +8,27 @@ define void @test_bigalloc(ptr %dst) {
 ; CHECK-NEXT:    [[TMP1:%.*]] = alloca [18446744069414584320 x i8], align 1
 ; CHECK-NEXT:    store ptr [[TMP1]], ptr [[DST:%.*]], align 8
 ; CHECK-NEXT:    ret void
+; CHECK32-LABEL: @test_bigalloc(
+; CHECK32-NEXT:    [[TMP1:%.*]] = alloca [0 x i8], align 1
+; CHECK32-NEXT:    store ptr [[TMP1]], ptr [[DST:%.*]], align 4
+; CHECK32-NEXT:    ret void
 ;
   %1 = alloca i8, i864 -4294967296
   store ptr %1, ptr %dst
   ret void
 }
+
+define void @test_negalloc(ptr %dst) {
+; CHECK-LABEL: @test_negalloc(
+; CHECK-NEXT:    [[TMP1:%.*]] = alloca [4294967295 x i32], align 4
+; CHECK-NEXT:    store ptr [[TMP1]], ptr [[DST:%.*]], align 8
+; CHECK-NEXT:    ret void
+; CHECK32-LABEL: @test_negalloc(
+; CHECK32-NEXT:    [[TMP1:%.*]] = alloca [0 x i32], align 4
+; CHECK32-NEXT:    store ptr [[TMP1]], ptr [[DST:%.*]], align 4
+; CHECK32-NEXT:    ret void
+;
+  %1 = alloca i32, i32 -1
+  store ptr %1, ptr %dst
+  ret void
+}