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

[via llvm-commits]

Author: Qiongsi Wu
Date: 2023-11-14T10:52:51-05:00
New Revision: c8b11091e832503f5440e9128fbd61ffa9ff86fd

URL: https://github.com/llvm/llvm-project/commit/c8b11091e832503f5440e9128fbd61ffa9ff86fd
DIFF: https://github.com/llvm/llvm-project/commit/c8b11091e832503f5440e9128fbd61ffa9ff86fd.diff

LOG: [SelectionDAG] Handling Oversized Alloca Types under 32 bit Mode to Avoid Code Generator Crash (#71472)

Situations may arise leading to negative `NumElements` argument of an
`alloca` instruction. In this case the `NumElements` is treated as a
large unsigned value. 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.

Added: 
    llvm/test/CodeGen/AArch64/alloca-oversized.ll
    llvm/test/CodeGen/PowerPC/alloca-oversized.ll

Modified: 
    llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp b/llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp
index aab0d5c5a348bfe..f0021f9f0907d37 100644
--- a/llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp
@@ -4128,7 +4128,7 @@ void SelectionDAGBuilder::visitAlloca(const AllocaInst &I) {
 
   SDValue AllocSize = getValue(I.getArraySize());
 
-  EVT IntPtr = TLI.getPointerTy(DAG.getDataLayout(), I.getAddressSpace());
+  EVT IntPtr = TLI.getPointerTy(DL, I.getAddressSpace());
   if (AllocSize.getValueType() != IntPtr)
     AllocSize = DAG.getZExtOrTrunc(AllocSize, dl, IntPtr);
 
@@ -4137,10 +4137,12 @@ void SelectionDAGBuilder::visitAlloca(const AllocaInst &I) {
                             DAG.getVScale(dl, IntPtr,
                                           APInt(IntPtr.getScalarSizeInBits(),
                                                 TySize.getKnownMinValue())));
-  else
-    AllocSize =
-        DAG.getNode(ISD::MUL, dl, IntPtr, AllocSize,
-                    DAG.getConstant(TySize.getFixedValue(), dl, IntPtr));
+  else {
+    SDValue TySizeValue =
+        DAG.getConstant(TySize.getFixedValue(), dl, MVT::getIntegerVT(64));
+    AllocSize = DAG.getNode(ISD::MUL, dl, IntPtr, AllocSize,
+                            DAG.getZExtOrTrunc(TySizeValue, 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/AArch64/alloca-oversized.ll b/llvm/test/CodeGen/AArch64/alloca-oversized.ll
new file mode 100644
index 000000000000000..e57bbcdf998003b
--- /dev/null
+++ b/llvm/test/CodeGen/AArch64/alloca-oversized.ll
@@ -0,0 +1,36 @@
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 3
+; RUN: llc < %s -mtriple=aarch64-none-linux-gnu -global-isel | FileCheck %s
+define void @test_oversized(ptr %dst, i32 %cond) {
+; CHECK-LABEL: test_oversized:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    stp x29, x30, [sp, #-16]! // 16-byte Folded Spill
+; CHECK-NEXT:    mov x29, sp
+; CHECK-NEXT:    sub sp, sp, #32
+; CHECK-NEXT:    .cfi_def_cfa w29, 16
+; CHECK-NEXT:    .cfi_offset w30, -8
+; CHECK-NEXT:    .cfi_offset w29, -16
+; CHECK-NEXT:    mov x8, sp
+; CHECK-NEXT:    mov x9, #2305843009213693952 // =0x2000000000000000
+; CHECK-NEXT:    sub x8, x8, x9
+; CHECK-NEXT:    sub x9, x29, #32
+; CHECK-NEXT:    mov sp, x8
+; CHECK-NEXT:    cmp w1, #0
+; CHECK-NEXT:    csel x8, x9, x8, eq
+; CHECK-NEXT:    str x8, [x0]
+; CHECK-NEXT:    mov sp, x29
+; CHECK-NEXT:    ldp x29, x30, [sp], #16 // 16-byte Folded Reload
+; CHECK-NEXT:    ret
+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 [18446744073709551615 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
+}

diff  --git a/llvm/test/CodeGen/PowerPC/alloca-oversized.ll b/llvm/test/CodeGen/PowerPC/alloca-oversized.ll
new file mode 100644
index 000000000000000..89922c348385869
--- /dev/null
+++ b/llvm/test/CodeGen/PowerPC/alloca-oversized.ll
@@ -0,0 +1,37 @@
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 3
+; RUN: llc < %s -mtriple=powerpc-ibm-aix-xcoff | FileCheck %s
+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)
+; 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
+}