diff --git a/llvm/lib/Transforms/IPO/Attributor.cpp b/llvm/lib/Transforms/IPO/Attributor.cpp --- a/llvm/lib/Transforms/IPO/Attributor.cpp +++ b/llvm/lib/Transforms/IPO/Attributor.cpp @@ -5030,6 +5030,7 @@ LLVM_DEBUG(dbgs() << "H2S: Removing malloc call: " << *MallocCall << "\n"); + MaybeAlign Alignment; Constant *Size; if (isCallocLikeFn(MallocCall, TLI)) { auto *Num = cast(MallocCall->getOperand(0)); @@ -5037,13 +5038,19 @@ APInt TotalSize = SizeT->getValue() * Num->getValue(); Size = ConstantInt::get(MallocCall->getOperand(0)->getType(), TotalSize); + } else if (isAlignedAllocLikeFn(MallocCall, TLI)) { + Size = cast(MallocCall->getOperand(1)); + Alignment = MaybeAlign(cast(MallocCall->getOperand(0)) + ->getValue() + .getZExtValue()); } else { Size = cast(MallocCall->getOperand(0)); } unsigned AS = cast(MallocCall->getType())->getAddressSpace(); - Instruction *AI = new AllocaInst(Type::getInt8Ty(F->getContext()), AS, - Size, "", MallocCall->getNextNode()); + Instruction *AI = + new AllocaInst(Type::getInt8Ty(F->getContext()), AS, Size, Alignment, + "", MallocCall->getNextNode()); if (AI->getType() != MallocCall->getType()) AI = new BitCastInst(AI, MallocCall->getType(), "malloc_bc", @@ -5175,8 +5182,9 @@ return true; bool IsMalloc = isMallocLikeFn(&I, TLI); + bool IsAlignedAllocLike = isAlignedAllocLikeFn(&I, TLI); bool IsCalloc = !IsMalloc && isCallocLikeFn(&I, TLI); - if (!IsMalloc && !IsCalloc) { + if (!IsMalloc && !IsAlignedAllocLike && !IsCalloc) { BadMallocCalls.insert(&I); return true; } @@ -5188,6 +5196,14 @@ MallocCalls.insert(&I); return true; } + } else if (IsAlignedAllocLike && isa(I.getOperand(0))) { + // Only if the alignment and sizes are constant. + if (auto *Size = dyn_cast(I.getOperand(1))) + if (Size->getValue().ule(MaxHeapToStackSize)) + if (UsesCheck(I) || FreeCheck(I)) { + MallocCalls.insert(&I); + return true; + } } else if (IsCalloc) { bool Overflow = false; if (auto *Num = dyn_cast(I.getOperand(0))) @@ -5219,8 +5235,9 @@ /// See AbstractAttribute::trackStatistics(). void trackStatistics() const override { - STATS_DECL(MallocCalls, Function, - "Number of malloc calls converted to allocas"); + STATS_DECL( + MallocCalls, Function, + "Number of malloc/calloc/aligned_alloc calls converted to allocas"); for (auto *C : MallocCalls) if (!BadMallocCalls.count(C)) ++BUILD_STAT_NAME(MallocCalls, Function); diff --git a/llvm/test/Transforms/Attributor/heap_to_stack.ll b/llvm/test/Transforms/Attributor/heap_to_stack.ll --- a/llvm/test/Transforms/Attributor/heap_to_stack.ll +++ b/llvm/test/Transforms/Attributor/heap_to_stack.ll @@ -76,6 +76,26 @@ ret void } +declare noalias i8* @aligned_alloc(i64, i64) + +define void @test3b(i8* %p) { + %1 = tail call noalias i8* @aligned_alloc(i64 32, i64 128) + ; CHECK: %1 = alloca i8, i64 128, align 32 + ; CHECK-NEXT: tail call void @nofree_arg_only + tail call void @nofree_arg_only(i8* %1, i8* %p) + ; CHECK-NOT: @free(i8* %1) + tail call void @free(i8* %1) + ret void +} + +; leave alone non-constant alignments. +define void @test3c(i64 %alignment) { + %1 = tail call noalias i8* @aligned_alloc(i64 %alignment, i64 128) + ; CHECK: tail call noalias i8* @aligned_alloc + tail call void @free(i8* %1) + ret void +} + declare noalias i8* @calloc(i64, i64) define void @test0() { @@ -90,7 +110,7 @@ ret void } -; TEST 4 +; TEST 4 define void @test4() { %1 = tail call noalias i8* @malloc(i64 4) ; CHECK: %1 = alloca i8, i64 4 @@ -219,7 +239,7 @@ ret i32 %3 } -; TEST 11 +; TEST 11 define void @test11() { %1 = tail call noalias i8* @malloc(i64 4)