[Attributor] Enable heap-to-stack of any size

[wsmoses]

Summary

Enable Attributor's heap-to-stack to lower unbounded allocations given a max size of -1

[jdoerfert]

LG, two nits.

llvm/lib/Transforms/IPO/AttributorAttributes.cpp

5173

Can you put these 4 lines in a helper lambda, we now replicate it a lot.

llvm/test/Transforms/Attributor/lowerheap.ll

36

can you move this in the heap2stack test case

[wsmoses]

Fix constant cast bug and test calloc

[jdoerfert]

LG, one nit comment below, address as makes sense.

llvm/lib/Transforms/IPO/AttributorAttributes.cpp

5012

If MaxHeapToStackSize is not -1, we can mark it nuw at least because we checked for overflow.
Should we just call the constructor, going through a builder seems kinda weird, except if we add a constant folder. Also, add a name, "h2s.calloc.size" or something.

llvm/lib/Transforms/IPO/AttributorAttributes.cpp

for (Instruction *MallocCall : MallocCalls) {
         A.deleteAfterManifest(*FreeCall);
         HasChanged = ChangeStatus::CHANGED;
       }
 
       LLVM_DEBUG(dbgs() << "H2S: Removing malloc call: " << *MallocCall
                         << "\n");
 
       Align Alignment;
-      Constant *Size;
+      Value *Size;
       if (isCallocLikeFn(MallocCall, TLI)) {
-        auto *Num = cast<ConstantInt>(MallocCall->getOperand(0));
-        auto *SizeT = cast<ConstantInt>(MallocCall->getOperand(1));
-        APInt TotalSize = SizeT->getValue() * Num->getValue();
-        Size =
-            ConstantInt::get(MallocCall->getOperand(0)->getType(), TotalSize);
+        auto *Num = MallocCall->getOperand(0);
+        auto *SizeT = MallocCall->getOperand(1);
+        IRBuilder<> B(MallocCall);
+        Size = B.CreateMul(Num, SizeT, "h2s.calloc.size");
       } else if (isAlignedAllocLikeFn(MallocCall, TLI)) {
-        Size = cast<ConstantInt>(MallocCall->getOperand(1));
+        Size = MallocCall->getOperand(1);
         Alignment = MaybeAlign(cast<ConstantInt>(MallocCall->getOperand(0))
                                    ->getValue()
                                    .getZExtValue())
                         .valueOrOne();
       } else {
-        Size = cast<ConstantInt>(MallocCall->getOperand(0));
+        Size = MallocCall->getOperand(0);
       }
 
       unsigned AS = cast<PointerType>(MallocCall->getType())->getAddressSpace();
       Instruction *AI =
           new AllocaInst(Type::getInt8Ty(F->getContext()), AS, Size, Alignment,
                          "", MallocCall->getNextNode());
 
       if (AI->getType() != MallocCall->getType())
auto MallocCallocCheck = [&](Instruction &I) {
     bool IsAlignedAllocLike = isAlignedAllocLikeFn(&I, TLI);
     bool IsCalloc = !IsMalloc && isCallocLikeFn(&I, TLI);
     if (!IsMalloc && !IsAlignedAllocLike && !IsCalloc) {
       BadMallocCalls.insert(&I);
       return true;
     }
 
     if (IsMalloc) {
+      if (MaxHeapToStackSize == -1) {
+        if (UsesCheck(I) || FreeCheck(I)) {
+          MallocCalls.insert(&I);
+          return true;
+        }
+      }
       if (auto *Size = dyn_cast<ConstantInt>(I.getOperand(0)))
         if (Size->getValue().ule(MaxHeapToStackSize))
           if (UsesCheck(I) || FreeCheck(I)) {
             MallocCalls.insert(&I);
             return true;
           }
     } else if (IsAlignedAllocLike && isa<ConstantInt>(I.getOperand(0))) {
+      if (MaxHeapToStackSize == -1) {
+        if (UsesCheck(I) || FreeCheck(I)) {
+          MallocCalls.insert(&I);
+          return true;
+        }
+      }
       // Only if the alignment and sizes are constant.
       if (auto *Size = dyn_cast<ConstantInt>(I.getOperand(1)))
         if (Size->getValue().ule(MaxHeapToStackSize))
           if (UsesCheck(I) || FreeCheck(I)) {
             MallocCalls.insert(&I);
             return true;
           }
     } else if (IsCalloc) {
+      if (MaxHeapToStackSize == -1) {
+        if (UsesCheck(I) || FreeCheck(I)) {
+          MallocCalls.insert(&I);
+          return true;
+        }
+      }
       bool Overflow = false;
       if (auto *Num = dyn_cast<ConstantInt>(I.getOperand(0)))
         if (auto *Size = dyn_cast<ConstantInt>(I.getOperand(1)))
           if ((Size->getValue().umul_ov(Num->getValue(), Overflow))
                   .ule(MaxHeapToStackSize))
             if (!Overflow && (UsesCheck(I) || FreeCheck(I))) {
               MallocCalls.insert(&I);
               return true;

llvm/test/Transforms/Attributor/lowerheap.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --function-signature --check-attributes
+; RUN: opt -max-heap-to-stack-size=-1 -attributor -enable-new-pm=0 -attributor-manifest-internal  -attributor-max-iterations-verify -attributor-annotate-decl-cs -attributor-max-iterations=1 -S < %s | FileCheck %s --check-prefixes=CHECK,NOT_CGSCC_NPM,NOT_CGSCC_OPM,NOT_TUNIT_NPM,IS__TUNIT____,IS________OPM,IS__TUNIT_OPM
+; RUN: opt -max-heap-to-stack-size=-1 -aa-pipeline=basic-aa -passes=attributor -attributor-manifest-internal  -attributor-max-iterations-verify -attributor-annotate-decl-cs -attributor-max-iterations=1 -S < %s | FileCheck %s --check-prefixes=CHECK,NOT_CGSCC_OPM,NOT_CGSCC_NPM,NOT_TUNIT_OPM,IS__TUNIT____,IS________NPM,IS__TUNIT_NPM
+; RUN: opt -max-heap-to-stack-size=-1 -attributor-cgscc -enable-new-pm=0 -attributor-manifest-internal  -attributor-annotate-decl-cs -S < %s | FileCheck %s --check-prefixes=CHECK,NOT_TUNIT_NPM,NOT_TUNIT_OPM,NOT_CGSCC_NPM,IS__CGSCC____,IS________OPM,IS__CGSCC_OPM
+; RUN: opt -max-heap-to-stack-size=-1 -aa-pipeline=basic-aa -passes=attributor-cgscc -attributor-manifest-internal  -attributor-annotate-decl-cs -S < %s | FileCheck %s --check-prefixes=CHECK,NOT_TUNIT_NPM,NOT_TUNIT_OPM,NOT_CGSCC_OPM,IS__CGSCC____,IS________NPM,IS__CGSCC_NPM
+
+declare i64 @subfn(i8*) #0
+
+declare noalias i8* @malloc(i64)
+declare noalias i8* @calloc(i64, i64)
+declare void @free(i8*)
+
+define i64 @f(i64 %len) {
+; IS________OPM-LABEL: define {{[^@]+}}@f
+; IS________OPM-SAME: (i64 [[LEN:%.*]]) {
+; IS________OPM-NEXT:  entry:
+; IS________OPM-NEXT:    [[MEM:%.*]] = call noalias i8* @malloc(i64 [[LEN]])
+; IS________OPM-NEXT:    [[RES:%.*]] = call i64 @subfn(i8* [[MEM]]) [[ATTR1:#.*]]
+; IS________OPM-NEXT:    call void @free(i8* [[MEM]])
+; IS________OPM-NEXT:    ret i64 [[RES]]
+;
+; IS________NPM-LABEL: define {{[^@]+}}@f
+; IS________NPM-SAME: (i64 [[LEN:%.*]]) {
+; IS________NPM-NEXT:  entry:
+; IS________NPM-NEXT:    [[TMP0:%.*]] = alloca i8, i64 [[LEN]], align 1
+; IS________NPM-NEXT:    [[RES:%.*]] = call i64 @subfn(i8* [[TMP0]]) [[ATTR2:#.*]]
+; IS________NPM-NEXT:    ret i64 [[RES]]
+;
+entry:
+  %mem = call i8* @malloc(i64 %len)
+  %res = call i64 @subfn(i8* %mem)
+  call void @free(i8* %mem)
+  ret i64 %res
+}
+
+
+define i64 @g(i64 %len) {
+; IS________OPM-LABEL: define {{[^@]+}}@g
+; IS________OPM-SAME: (i64 [[LEN:%.*]]) {
+; IS________OPM-NEXT:  entry:
+; IS________OPM-NEXT:    [[MEM:%.*]] = call noalias i8* @calloc(i64 [[LEN]], i64 noundef 8)
+; IS________OPM-NEXT:    [[RES:%.*]] = call i64 @subfn(i8* [[MEM]]) [[ATTR1]]
+; IS________OPM-NEXT:    call void @free(i8* [[MEM]])
+; IS________OPM-NEXT:    ret i64 [[RES]]
+;
+; IS________NPM-LABEL: define {{[^@]+}}@g
+; IS________NPM-SAME: (i64 [[LEN:%.*]]) {
+; IS________NPM-NEXT:  entry:
+; IS________NPM-NEXT:    [[TMP0:%.*]] = mul i64 [[LEN]], 8
+; IS________NPM-NEXT:    [[TMP1:%.*]] = alloca i8, i64 [[TMP0]], align 1
+; IS________NPM-NEXT:    [[CALLOC_BC:%.*]] = bitcast i8* [[TMP1]] to i8*
+; IS________NPM-NEXT:    call void @llvm.memset.p0i8.i64(i8* [[CALLOC_BC]], i8 0, i64 [[TMP0]], i1 false)
+; IS________NPM-NEXT:    [[RES:%.*]] = call i64 @subfn(i8* [[TMP1]]) [[ATTR2]]
+; IS________NPM-NEXT:    ret i64 [[RES]]
+;
+entry:
+  %mem = call i8* @calloc(i64 %len, i64 8)
+  %res = call i64 @subfn(i8* %mem)
+  call void @free(i8* %mem)
+  ret i64 %res
+}
+
+attributes #0 = { nounwind willreturn }