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 @@ -117,6 +117,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 @@ -310,6 +310,22 @@ return V->getOperand(FnData->AlignParam); } +std::pair +llvm::getAllocSizeArgs(const CallBase *V, const TargetLibraryInfo *TLI) { + const Optional FnData = getAllocationSize(V, TLI); + assert(FnData.hasValue() && + "getAllocSizeArgs can only be called on an allocation function"); + 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 @@ -5773,9 +5773,8 @@ /// The kind of allocation. const enum class AllocationKind { - MALLOC, - CALLOC, - ALIGNED_ALLOC, + AllocZeros, + AllocUndef, } Kind; /// The library function id for the allocation. @@ -5834,16 +5833,19 @@ 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)) return true; - auto Kind = - IsMalloc ? AllocationInfo::AllocationKind::MALLOC - : (IsCalloc ? AllocationInfo::AllocationKind::CALLOC - : AllocationInfo::AllocationKind::ALIGNED_ALLOC); + + auto *AllocValue = getInitialValueOfAllocation(CB, TLI, CB->getType()); + AllocationInfo::AllocationKind Kind; + if (AllocValue->isNullValue()) { + Kind = AllocationInfo::AllocationKind::AllocZeros; + } else if (AllocValue->containsUndefOrPoisonElement()) { + Kind = AllocationInfo::AllocationKind::AllocUndef; + } else { + return true; // Allocation has unknown but defined initial value + } AllocationInfo *AI = new (A.Allocator) AllocationInfo{CB, Kind}; AllocationInfos[CB] = AI; @@ -5943,27 +5945,27 @@ 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)); - assert(AlignmentAPI.hasValue() && - "Expected an alignment during manifest!"); - Alignment = - max(Alignment, MaybeAlign(AlignmentAPI.getValue().getZExtValue())); + if (Value *Align = getAllocAlignment(AI.CB, TLI)) { + Optional AlignmentAPI = getAPInt(A, *this, *Align); + if (AlignmentAPI.hasValue()) { + Alignment = max(Alignment, + MaybeAlign(AlignmentAPI.getValue().getZExtValue())); + } else { + Alignment = max(Alignment, MaybeAlign(llvm::Value::MaximumAlignment)); + } } unsigned AS = cast(AI.CB->getType())->getAddressSpace(); @@ -5986,14 +5988,15 @@ } // Zero out the allocated memory if it was a calloc. - if (AI.Kind == AllocationInfo::AllocationKind::CALLOC) { + if (AI.Kind == AllocationInfo::AllocationKind::AllocZeros) { auto *BI = new BitCastInst(Alloca, AI.CB->getType(), "calloc_bc", Alloca->getNextNode()); Value *Ops[] = { BI, ConstantInt::get(F->getContext(), APInt(8, 0, false)), 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,25 +6021,14 @@ Optional getSize(Attributor &A, const AbstractAttribute &AA, AllocationInfo &AI) { - - if (AI.Kind == AllocationInfo::AllocationKind::MALLOC) - return getAPInt(A, AA, *AI.CB->getArgOperand(0)); - - 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; - - 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)); - if (!Num.hasValue() || !Size.hasValue()) - return llvm::None; - bool Overflow = false; - Size = Size.getValue().umul_ov(Num.getValue(), Overflow); - return Overflow ? llvm::None : Size; + Function *F = getAnchorScope(); + const auto *TLI = A.getInfoCache().getTargetLibraryInfoForFunction(*F); + const auto DL = A.getDataLayout(); + uint64_t Size; + if (getObjectSize(AI.CB, Size, DL, TLI) && Size > 0) { + return APInt(64, Size); + } + return llvm::None; } /// Collection of all malloc-like calls in a function with associated @@ -6053,6 +6045,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 +6261,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;