diff --git a/llvm/include/llvm/Analysis/MemoryBuiltins.h b/llvm/include/llvm/Analysis/MemoryBuiltins.h --- a/llvm/include/llvm/Analysis/MemoryBuiltins.h +++ b/llvm/include/llvm/Analysis/MemoryBuiltins.h @@ -119,6 +119,12 @@ /// Gets the alignment argument for an aligned_alloc-like function Value *getAllocAlignment(const CallBase *V, const TargetLibraryInfo *TLI); +/// Returns the size arguments for a allocation function. If there is only one +/// size argument, the second pointer is null. getObjectSize is preferred if +/// only constant sizes are needed +std::pair getAllocSizeArgs(const CallBase *V, + const TargetLibraryInfo *TLI); + /// If this allocation function initializes memory to a fixed value, return /// said value in the requested type. Otherwise, return nullptr. Constant *getInitialValueOfAllocation(const CallBase *Alloc, diff --git a/llvm/lib/Analysis/MemoryBuiltins.cpp b/llvm/lib/Analysis/MemoryBuiltins.cpp --- a/llvm/lib/Analysis/MemoryBuiltins.cpp +++ b/llvm/lib/Analysis/MemoryBuiltins.cpp @@ -309,6 +309,24 @@ return V->getOperand(FnData->AlignParam); } +std::pair +llvm::getAllocSizeArgs(const CallBase *V, const TargetLibraryInfo *TLI) { + assert((isMallocOrCallocLikeFn(V, TLI) || isReallocLikeFn(V, TLI)) && + "getAllocSizeArgs can only be called on an allocation function that " + "isn't like strdup"); + const Optional FnData = getAllocationSize(V, TLI); + assert(FnData.hasValue()); + if (FnData->FstParam < 0) { + return std::pair(nullptr, nullptr); + } + if (FnData->SndParam < 0) { + return std::pair(V->getOperand(FnData->FstParam), + nullptr); + } + return std::pair(V->getOperand(FnData->FstParam), + V->getOperand(FnData->SndParam)); +} + Constant *llvm::getInitialValueOfAllocation(const CallBase *Alloc, const TargetLibraryInfo *TLI, Type *Ty) { diff --git a/llvm/lib/Transforms/IPO/AttributorAttributes.cpp b/llvm/lib/Transforms/IPO/AttributorAttributes.cpp --- a/llvm/lib/Transforms/IPO/AttributorAttributes.cpp +++ b/llvm/lib/Transforms/IPO/AttributorAttributes.cpp @@ -5771,13 +5771,6 @@ /// The call that allocates the memory. CallBase *const CB; - /// The kind of allocation. - const enum class AllocationKind { - MALLOC, - CALLOC, - ALIGNED_ALLOC, - } Kind; - /// The library function id for the allocation. LibFunc LibraryFunctionId = NotLibFunc; @@ -5834,18 +5827,16 @@ DeallocationInfos[CB] = new (A.Allocator) DeallocationInfo{CB}; return true; } - bool IsMalloc = isMallocLikeFn(CB, TLI); - bool IsAlignedAllocLike = !IsMalloc && isAlignedAllocLikeFn(CB, TLI); - bool IsCalloc = - !IsMalloc && !IsAlignedAllocLike && isCallocLikeFn(CB, TLI); - if (!IsMalloc && !IsAlignedAllocLike && !IsCalloc) + + if (!(isMallocOrCallocLikeFn(CB, TLI) && isAllocRemovable(CB, TLI))) return true; - auto Kind = - IsMalloc ? AllocationInfo::AllocationKind::MALLOC - : (IsCalloc ? AllocationInfo::AllocationKind::CALLOC - : AllocationInfo::AllocationKind::ALIGNED_ALLOC); - AllocationInfo *AI = new (A.Allocator) AllocationInfo{CB, Kind}; + auto *AllocValue = getInitialValueOfAllocation(CB, TLI, CB->getType()); + if (!AllocValue) { + return true; // Allocation has unknown but defined initial value + } + + AllocationInfo *AI = new (A.Allocator) AllocationInfo{CB}; AllocationInfos[CB] = AI; TLI->getLibFunc(*CB, AI->LibraryFunctionId); return true; @@ -5943,23 +5934,21 @@ Optional SizeAPI = getSize(A, *this, AI); if (SizeAPI.hasValue()) { Size = ConstantInt::get(AI.CB->getContext(), *SizeAPI); - } else if (AI.Kind == AllocationInfo::AllocationKind::CALLOC) { - auto *Num = AI.CB->getOperand(0); - auto *SizeT = AI.CB->getOperand(1); - IRBuilder<> B(AI.CB); - Size = B.CreateMul(Num, SizeT, "h2s.calloc.size"); - } else if (AI.Kind == AllocationInfo::AllocationKind::ALIGNED_ALLOC) { - Size = AI.CB->getOperand(1); } else { - Size = AI.CB->getOperand(0); + auto Args = getAllocSizeArgs(AI.CB, TLI); + if (!Args.second) { + Size = Args.first; + } else { + IRBuilder<> B(AI.CB); + Size = B.CreateMul(Args.first, Args.second, "h2s.calloc.size"); + } } Align Alignment(1); if (MaybeAlign RetAlign = AI.CB->getRetAlign()) Alignment = max(Alignment, RetAlign); - if (AI.Kind == AllocationInfo::AllocationKind::ALIGNED_ALLOC) { - Optional AlignmentAPI = - getAPInt(A, *this, *AI.CB->getArgOperand(0)); + if (Value *Align = getAllocAlignment(AI.CB, TLI)) { + Optional AlignmentAPI = getAPInt(A, *this, *Align); assert(AlignmentAPI.hasValue() && "Expected an alignment during manifest!"); Alignment = @@ -5985,15 +5974,19 @@ A.deleteAfterManifest(*AI.CB); } - // Zero out the allocated memory if it was a calloc. - if (AI.Kind == AllocationInfo::AllocationKind::CALLOC) { + // Set the initial value of the memory + auto *AllocValue = + getInitialValueOfAllocation(AI.CB, TLI, AI.CB->getType()); + assert(AllocValue); + if (!isa(AllocValue)) { auto *BI = new BitCastInst(Alloca, AI.CB->getType(), "calloc_bc", Alloca->getNextNode()); Value *Ops[] = { - BI, ConstantInt::get(F->getContext(), APInt(8, 0, false)), Size, + BI, AllocValue, Size, ConstantInt::get(Type::getInt1Ty(F->getContext()), false)}; - Type *Tys[] = {BI->getType(), AI.CB->getOperand(0)->getType()}; + Type *Tys[] = {BI->getType(), + getAllocSizeArgs(AI.CB, TLI).first->getType()}; Module *M = F->getParent(); Function *Fn = Intrinsic::getDeclaration(M, Intrinsic::memset, Tys); CallInst::Create(Fn, Ops, "", BI->getNextNode()); @@ -6018,20 +6011,24 @@ Optional getSize(Attributor &A, const AbstractAttribute &AA, AllocationInfo &AI) { + Function *F = getAnchorScope(); + const auto *TLI = A.getInfoCache().getTargetLibraryInfoForFunction(*F); + + // Only if the alignment is also constant we return a size. + Value *Align; + if ((Align = getAllocAlignment(AI.CB, TLI)) && + !getAPInt(A, AA, *Align).hasValue()) { + return llvm::None; + } - if (AI.Kind == AllocationInfo::AllocationKind::MALLOC) - return getAPInt(A, AA, *AI.CB->getArgOperand(0)); + auto Args = getAllocSizeArgs(AI.CB, TLI); - if (AI.Kind == AllocationInfo::AllocationKind::ALIGNED_ALLOC) - // Only if the alignment is also constant we return a size. - return getAPInt(A, AA, *AI.CB->getArgOperand(0)).hasValue() - ? getAPInt(A, AA, *AI.CB->getArgOperand(1)) - : llvm::None; + if (!Args.second) { + return getAPInt(A, AA, *Args.first); + } - assert(AI.Kind == AllocationInfo::AllocationKind::CALLOC && - "Expected only callocs are left"); - Optional Num = getAPInt(A, AA, *AI.CB->getArgOperand(0)); - Optional Size = getAPInt(A, AA, *AI.CB->getArgOperand(1)); + Optional Num = getAPInt(A, AA, *Args.first); + Optional Size = getAPInt(A, AA, *Args.second); if (!Num.hasValue() || !Size.hasValue()) return llvm::None; bool Overflow = false; @@ -6053,6 +6050,7 @@ ChangeStatus AAHeapToStackFunction::updateImpl(Attributor &A) { ChangeStatus Changed = ChangeStatus::UNCHANGED; const Function *F = getAnchorScope(); + const auto *TLI = A.getInfoCache().getTargetLibraryInfoForFunction(*F); const auto &LivenessAA = A.getAAFor(*this, IRPosition::function(*F), DepClassTy::NONE); @@ -6268,8 +6266,8 @@ continue; if (MaxHeapToStackSize == -1) { - if (AI.Kind == AllocationInfo::AllocationKind::ALIGNED_ALLOC) - if (!getAPInt(A, *this, *AI.CB->getArgOperand(0)).hasValue()) { + if (Value *Align = getAllocAlignment(AI.CB, TLI)) + if (!getAPInt(A, *this, *Align).hasValue()) { LLVM_DEBUG(dbgs() << "[H2S] Unknown allocation alignment: " << *AI.CB << "\n"); AI.Status = AllocationInfo::INVALID;