WIP: [llvm][MemoryBuiltins] Add initialization category to getInitialValueOfAllocation

[jmciver]

Summary

Part of a sequence of exploratory patches using poison semantics for
uninitialized memory.

The enumeration InitializationCategory will be used to provide backwards
compatibility to existing optimizations where further information is required to
handle poison and freeze poison semantics.

Callers are refactored at this time to only use the returned constant. Future
commits will enable the use of the InitializationCategory.

llvm/include/llvm/Analysis/MemoryBuiltins.h

 /// one Value* (operand to the allocation) with another. This is useful when
 /// doing abstract interpretation.
 std::optional<APInt> getAllocSize(
     const CallBase *CB, const TargetLibraryInfo *TLI,
     function_ref<const Value *(const Value *)> Mapper = [](const Value *V) {
       return V;
     });
 
+bool loadHasFreezeBits(const LoadInst *const LI);
+
+/// Provide categorical information concerning constant selection.
+enum class InitializationCategory {
+  /// This enumerator will be returned when the load is determined to be of an
+  /// address that was allocated using:
+  ///
+  ///   * calloc or compliant new derivative
+  ///   * malloc in the absence of !freeze_bits
+  ///
+  /// In the first case the returned constant will be zero. This is in keeping
+  /// with current behavior where initialized AllocFnKind::Zeroed results in a
+  /// zero constant of arbitrary type. In the second case the return constant
+  /// will be poison.
+  Constant,
+  /// This enumerator will be returned when the load accompanied by !freeze_bits
+  /// and classification of the allocator such that memory is determined to be
+  /// in an uninitialized state. The returned constant will be null of arbitrary
+  /// type.
+  ///
+  /// The MemoryBuiltins API is used by clients that remove loads and stores.
+  /// The reason for the null of arbitrary type is it is differentiated
+  /// from zero. If zero were used optimizations like dead store elimination
+  /// (DSE) would remove the store 0.
+  FreezePoison,
+  /// Returned when getInitialValueOfAllocation is unable to determine the
+  /// classification of the allocation operations. No constant object is
+  /// produced and thus the returned pointer will be a nullptr. This provides
+  /// backward compatibility with the current API.
+  Unknown
+};
+
 /// If this is a call to an allocation function that initializes memory to a
-/// fixed value, return said value in the requested type. If this is a call to
-/// alloca instruction the returned value is undef. Otherwise, return nullptr.
-Constant *getInitialValueOfAllocation(const Value *V,
-                                      const TargetLibraryInfo *TLI,
-                                      Type *Ty);
+/// fixed value, return said value in the requested type. Otherwise, return
+/// nullptr.
+///
+/// This function also serves to determine the return value of a load
+/// instruction directly using allocation without an intermediate store
+/// instruction(s). This mode is enabled by passing a load institution pointer.
+std::pair<InitializationCategory, Constant *>
+getInitialValueOfAllocation(const Value *V, const TargetLibraryInfo *TLI,
+                            Type *Ty, const LoadInst *LI = nullptr);
 
 /// If a function is part of an allocation family (e.g.
 /// malloc/realloc/calloc/free), return the identifier for its family
 /// of functions.
 std::optional<StringRef> getAllocationFamily(const Value *I,
                                              const TargetLibraryInfo *TLI);
 
 //===----------------------------------------------------------------------===//

llvm/lib/Analysis/MemoryBuiltins.cpp

llvm::getAllocSize(const CallBase *CB, const TargetLibraryInfo *TLI,
 
   bool Overflow;
   Size = Size.umul_ov(NumElems, Overflow);
   if (Overflow)
     return std::nullopt;
   return Size;
 }
 
-Constant *llvm::getInitialValueOfAllocation(const Value *V,
-                                            const TargetLibraryInfo *TLI,
-                                            Type *Ty) {
+bool llvm::loadHasFreezeBits(const LoadInst *const LI) {
+  return LI && LI->hasMetadata(LLVMContext::MD_freeze_bits);
+}
+
+static std::pair<InitializationCategory, Constant *>
+valueUsingLoadInstAndFreezeBits(const LoadInst *const LI, Type *const Ty) {
+  if (loadHasFreezeBits(LI))
+    return {InitializationCategory::FreezePoison, nullptr};
+  else
+    return {InitializationCategory::Constant, PoisonValue::get(Ty)};
+}
+
+std::pair<InitializationCategory, Constant *>
+llvm::getInitialValueOfAllocation(const Value *V, const TargetLibraryInfo *TLI,
+                                  Type *Ty, const LoadInst *LI) {
   if (isa<AllocaInst>(V))
-    return UndefValue::get(Ty);
+    return valueUsingLoadInstAndFreezeBits(LI, Ty);
 
   auto *Alloc = dyn_cast<CallBase>(V);
   if (!Alloc)
-    return nullptr;
+    return {InitializationCategory::Unknown, nullptr};
 
   // malloc are uninitialized (undef)
   if (getAllocationData(Alloc, MallocOrOpNewLike, TLI).has_value())
-    return UndefValue::get(Ty);
+    return valueUsingLoadInstAndFreezeBits(LI, Ty);
 
   AllocFnKind AK = getAllocFnKind(Alloc);
   if ((AK & AllocFnKind::Uninitialized) != AllocFnKind::Unknown)
-    return UndefValue::get(Ty);
+    return valueUsingLoadInstAndFreezeBits(LI, Ty);
   if ((AK & AllocFnKind::Zeroed) != AllocFnKind::Unknown)
-    return Constant::getNullValue(Ty);
+    return {InitializationCategory::Constant, Constant::getNullValue(Ty)};
 
-  return nullptr;
+  return {InitializationCategory::Unknown, nullptr};
 }
 
 struct FreeFnsTy {
   unsigned NumParams;
   // Name of default allocator function to group malloc/free calls by family
   MallocFamily Family;
 };
 

llvm/lib/Target/AMDGPU/AMDGPUPromoteAlloca.cpp

if (!Info.SrcIndex || !Info.DestIndex)
           Inst, "mem transfer inst is missing constant src and/or dst index");
   }
 
   LLVM_DEBUG(dbgs() << "  Converting alloca to vector " << *AllocaTy << " -> "
                     << *VectorTy << '\n');
   const unsigned VecStoreSize = DL->getTypeStoreSize(VectorTy);
 
   // Alloca is uninitialized memory. Imitate that by making the first value
-  // undef.
+  // poison.
   SSAUpdater Updater;
   Updater.Initialize(VectorTy, "promotealloca");
   Updater.AddAvailableValue(
       Alloca.getParent(),
-      getInitialValueOfAllocation(&Alloca, nullptr, VectorTy));
+      getInitialValueOfAllocation(&Alloca, nullptr, VectorTy).second);
 
   // First handle the initial worklist.
   SmallVector<LoadInst *, 4> DeferredLoads;
   forEachWorkListItem(WorkList, [&](Instruction *I) {
     BasicBlock *BB = I->getParent();
     // On the first pass, we only take values that are trivially known, i.e.
     // where AddAvailableValue was already called in this block.
     Value *Result = promoteAllocaUserToVector(

llvm/lib/Transforms/IPO/Attributor.cpp

 }
 
 Constant *
 AA::getInitialValueForObj(Attributor &A, const AbstractAttribute &QueryingAA,
                           Value &Obj, Type &Ty, const TargetLibraryInfo *TLI,
                           const DataLayout &DL, AA::RangeTy *RangePtr) {
   if (isa<AllocaInst>(Obj))
     return UndefValue::get(&Ty);
-  if (Constant *Init = getInitialValueOfAllocation(&Obj, TLI, &Ty))
+  if (Constant *Init = getInitialValueOfAllocation(&Obj, TLI, &Ty).second)
     return Init;
   auto *GV = dyn_cast<GlobalVariable>(&Obj);
   if (!GV)
     return nullptr;
 
   bool UsedAssumedInformation = false;
   Constant *Initializer = nullptr;
   if (A.hasGlobalVariableSimplificationCallback(*GV)) {

llvm/lib/Transforms/IPO/AttributorAttributes.cpp

auto AllocationIdentifierCB = [&](Instruction &I) {
         DeallocationInfos[CB] = new (A.Allocator) DeallocationInfo{CB, FreedOp};
         return true;
       }
       // To do heap to stack, we need to know that the allocation itself is
       // removable once uses are rewritten, and that we can initialize the
       // alloca to the same pattern as the original allocation result.
       if (isRemovableAlloc(CB, TLI)) {
         auto *I8Ty = Type::getInt8Ty(CB->getParent()->getContext());
-        if (nullptr != getInitialValueOfAllocation(CB, TLI, I8Ty)) {
+        if (nullptr != getInitialValueOfAllocation(CB, TLI, I8Ty).second) {
           AllocationInfo *AI = new (A.Allocator) AllocationInfo{CB};
           AllocationInfos[CB] = AI;
           if (TLI)
             TLI->getLibFunc(*CB, AI->LibraryFunctionId);
         }
       }
       return true;
     };
for (auto &It : AllocationInfos) {
           new AllocaInst(Type::getInt8Ty(F->getContext()), AS, Size, Alignment,
                          AI.CB->getName() + ".h2s", IP);
 
       if (Alloca->getType() != AI.CB->getType())
         Alloca = BitCastInst::CreatePointerBitCastOrAddrSpaceCast(
             Alloca, AI.CB->getType(), "malloc_cast", AI.CB);
 
       auto *I8Ty = Type::getInt8Ty(F->getContext());
-      auto *InitVal = getInitialValueOfAllocation(AI.CB, TLI, I8Ty);
+      auto *InitVal = getInitialValueOfAllocation(AI.CB, TLI, I8Ty).second;
       assert(InitVal &&
              "Must be able to materialize initial memory state of allocation");
 
       A.changeAfterManifest(IRPosition::inst(*AI.CB), *Alloca);
 
       if (auto *II = dyn_cast<InvokeInst>(AI.CB)) {
         auto *NBB = II->getNormalDest();
         BranchInst::Create(NBB, AI.CB->getParent());

llvm/lib/Transforms/IPO/GlobalOpt.cpp

static bool tryToOptimizeStoreOfAllocationToGlobal(GlobalVariable *GV,
                                                    CallInst *CI,
                                                    const DataLayout &DL,
                                                    TargetLibraryInfo *TLI) {
   if (!isRemovableAlloc(CI, TLI))
     // Must be able to remove the call when we get done..
     return false;
 
   Type *Int8Ty = Type::getInt8Ty(CI->getFunction()->getContext());
-  Constant *InitVal = getInitialValueOfAllocation(CI, TLI, Int8Ty);
+  Constant *InitVal = getInitialValueOfAllocation(CI, TLI, Int8Ty).second;
   if (!InitVal)
     // Must be able to emit a memset for initialization
     return false;
 
   uint64_t AllocSize;
   if (!getObjectSize(CI, AllocSize, DL, TLI, ObjectSizeOpts()))
     return false;
 

llvm/lib/Transforms/Scalar/DeadStoreElimination.cpp

bool storeIsNoop(MemoryDef *Def, const Value *DefUO) {
     else
       return false;
 
     if (!isRemovable(DefI))
       return false;
 
     if (StoredConstant) {
       Constant *InitC =
-          getInitialValueOfAllocation(DefUO, &TLI, StoredConstant->getType());
+          getInitialValueOfAllocation(DefUO, &TLI, StoredConstant->getType())
+              .second;
       // If the clobbering access is LiveOnEntry, no instructions between them
       // can modify the memory location.
       if (InitC && InitC == StoredConstant)
         return MSSA.isLiveOnEntryDef(
             MSSA.getSkipSelfWalker()->getClobberingMemoryAccess(Def, BatchAA));
     }
 
     if (!Store)

llvm/lib/Transforms/Scalar/GVN.cpp

GVNPass::AnalyzeLoadAvailability(LoadInst *Load, MemDepResult DepInfo,
 
   // Loading immediately after lifetime begin -> undef.
   if (isLifetimeStart(DepInst))
     return AvailableValue::get(UndefValue::get(Load->getType()));
 
   // In addition to allocator function calls this includes loading the alloca ->
   // undef.
   if (Constant *InitVal =
-          getInitialValueOfAllocation(DepInst, TLI, Load->getType()))
+          getInitialValueOfAllocation(DepInst, TLI, Load->getType()).second)
     return AvailableValue::get(InitVal);
 
   if (StoreInst *S = dyn_cast<StoreInst>(DepInst)) {
     // Reject loads and stores that are to the same address but are of
     // different types if we have to. If the stored value is convertable to
     // the loaded value, we can reuse it.
     if (!canCoerceMustAliasedValueToLoad(S->getValueOperand(), Load->getType(),
                                          DL))

llvm/lib/Transforms/Scalar/NewGVN.cpp

NewGVN::performSymbolicLoadCoercion(Type *LoadType, Value *LoadPtr,
   if (auto *II = dyn_cast<IntrinsicInst>(DepInst)) {
     if (II->getIntrinsicID() == Intrinsic::lifetime_start)
       return createConstantExpression(UndefValue::get(LoadType));
   }
   // If this load really doesn't depend on anything, then we must be loading an
   // undef value.  This can happen when loading for a fresh allocation with no
   // intervening stores, for example.  Note that this is only true in the case
   // that the result of the allocation is pointer equal to the load ptr.
-  else if (auto *InitVal = getInitialValueOfAllocation(DepInst, TLI, LoadType))
+  else if (auto *InitVal = getInitialValueOfAllocation(DepInst, TLI, LoadType).second)
     return createConstantExpression(InitVal);
 
   return nullptr;
 }
 
 const Expression *NewGVN::performSymbolicLoadEvaluation(Instruction *I) const {
   auto *LI = cast<LoadInst>(I);
 

llvm/lib/Transforms/Utils/PromoteMemoryToRegister.cpp

for (User *U : make_early_inc_range(AI->users())) {
     StoresByIndexTy::iterator I = llvm::lower_bound(
         StoresByIndex,
         std::make_pair(LoadIdx, static_cast<StoreInst *>(nullptr)),
         less_first());
     Value *ReplVal;
     if (I == StoresByIndex.begin()) {
       if (StoresByIndex.empty())
         // If there are no stores, the load takes the undef value.
-        ReplVal = getInitialValueOfAllocation(AI, nullptr, LI->getType());
+        ReplVal = getInitialValueOfAllocation(AI, nullptr, LI->getType()).second;
       else
         // There is no store before this load, bail out (load may be affected
         // by the following stores - see main comment).
         return false;
     } else {
       // Otherwise, there was a store before this load, the load takes its
       // value.
       ReplVal = std::prev(I)->second->getOperand(0);