Fix size computation of array allocation in inline cost analysis

[eraman]

Summary

There are two bugs in computing size allocated by an array alloca:

• For primitive types, it computes the number of bits and there is a FIXME
• For aggregate types, the computed size is always 0.

The fix is to use DL.getTypeAllocSize(Ty) as it is done for static alloca

[eli.friedman]

While you're here, could you fix this to handle overflow correctly?

declare void @h(i32*)

define void @f(i128 %x) {
   %y = alloca i32, i128 %x
   call void @h(i32* %y)
   ret void
}

define void @g() {
  call void @f(i128 100000000000000000000000000000000000)
  ret void
}

[eraman]

Handle overflow in computing stack size.

[eli.friedman]

lib/Analysis/InlineCost.cpp

347

The addition can overflow as well.

Maybe use SaturatingMultiplyAdd from MathExtras.h instead of APInt?

[eraman]

Use SaturatingAdd and SaturatingMultiplyAdd to update AllocatedSize in visitAlloca

[eli.friedman]

LGTM.

lib/Analysis/InlineCost.cpp

 }
 
 bool CallAnalyzer::visitAlloca(AllocaInst &I) {
   // Check whether inlining will turn a dynamic alloca into a static
   // alloca and handle that case.
   if (I.isArrayAllocation()) {
     Constant *Size = SimplifiedValues.lookup(I.getArraySize());
     if (auto *AllocSize = dyn_cast_or_null<ConstantInt>(Size)) {
+      const DataLayout &DL = F.getParent()->getDataLayout();
       Type *Ty = I.getAllocatedType();
-      // FIXME: This can't be right. AllocatedSize is in *bytes*.
-      AllocatedSize += Ty->getPrimitiveSizeInBits() * AllocSize->getZExtValue();
+      AllocatedSize = SaturatingMultiplyAdd(
+          AllocSize->getLimitedValue(), DL.getTypeAllocSize(Ty), AllocatedSize);
       return Base::visitAlloca(I);
     }
   }
 
   // Accumulate the allocated size.
   if (I.isStaticAlloca()) {
     const DataLayout &DL = F.getParent()->getDataLayout();
     Type *Ty = I.getAllocatedType();
-    AllocatedSize += DL.getTypeAllocSize(Ty);
+    AllocatedSize = SaturatingAdd(DL.getTypeAllocSize(Ty), AllocatedSize);
   }
 
   // We will happily inline static alloca instructions.
   if (I.isStaticAlloca())
     return Base::visitAlloca(I);
 
   // FIXME: This is overly conservative. Dynamic allocas are inefficient for
   // a variety of reasons, and so we would like to not inline them into

test/Transforms/Inline/array-alloca.ll

+; RUN: opt -inline -S < %s | FileCheck %s
+%struct.A = type { i32 }
+
+define void @callee1(i32 %M) {
+entry:
+  %vla = alloca i32, i32 %M, align 16
+  ret void
+}
+
+define void @callee2(i32 %M) {
+entry:
+  %vla = alloca %struct.A, i32 %M, align 16
+  ret void
+}
+
+define void @callee3(i128 %M) {
+entry:
+  %vla = alloca i32, i128 %M, align 16
+  ret void
+}
+
+; CHECK-LABEL: @caller
+define void @caller() #0 {
+entry:
+  call void @caller()
+; CHECK-NOT: call void @callee1
+  call void @callee1(i32 256)
+; CHECK: call void @callee2
+  call void @callee2(i32 4096)
+; CHECK: call void @callee3
+; This is to test that there is no overflow in computing allocated size
+; call void @callee3(i128 0x8000000000000000);
+  call void @callee3(i128 9223372036854775808);
+  ret void
+}
+