Add a UndefinedVal to SSAUpdater, to use instead of Poison

Author: davemgreen

llvm/include/llvm/Transforms/Utils/SSAUpdater.h

@@ -55,6 +55,10 @@ class SSAUpdater {
   /// the vector.
   SmallVectorImpl<PHINode *> *InsertedPHIs;
 
+  /// The Undefined value (usually poison) to use for loads with no available
+  /// values or blocks considered unreachable.
+  Value *UndefinedVal = nullptr;
+
 public:
   /// If InsertedPHIs is specified, it will be filled
   /// in with all PHI Nodes created by rewriting.
@@ -134,6 +138,13 @@ class SSAUpdater {
   /// inserted values.
   void RewriteUseAfterInsertions(Use &U);
 
+  /// Set the UndefinedVal used for for loads with no available values or blocks
+  /// considered unreachable.
+  void SetUndefinedVal(Value *V);
+
+  /// Get the UndefinedVal or Poison if it is unset.
+  Value *GetUndefinedVal(Type *Ty);
+
 private:
   Value *GetValueAtEndOfBlockInternal(BasicBlock *BB);
   void UpdateDebugValue(Instruction *I, DbgValueInst *DbgValue);

llvm/include/llvm/Transforms/Utils/SSAUpdaterImpl.h

@@ -81,11 +81,14 @@ class SSAUpdaterImpl {
 
   BBMapTy BBMap;
   BumpPtrAllocator Allocator;
+  std::optional<ValT> UndefinedVal;
 
 public:
   explicit SSAUpdaterImpl(UpdaterT *U, AvailableValsTy *A,
-                          SmallVectorImpl<PhiT *> *Ins) :
-    Updater(U), AvailableVals(A), InsertedPHIs(Ins) {}
+                          SmallVectorImpl<PhiT *> *Ins,
+                          std::optional<ValT> UndefinedVal = std::nullopt)
+      : Updater(U), AvailableVals(A), InsertedPHIs(Ins),
+        UndefinedVal(UndefinedVal) {}
 
   /// GetValue - Check to see if AvailableVals has an entry for the specified
   /// BB and if so, return it.  If not, construct SSA form by first
@@ -97,7 +100,7 @@ class SSAUpdaterImpl {
 
     // Special case: bail out if BB is unreachable.
     if (BlockList.size() == 0) {
-      ValT V = Traits::GetPoisonVal(BB, Updater);
+      ValT V = UndefinedVal ? *UndefinedVal : Traits::GetPoisonVal(BB, Updater);
       (*AvailableVals)[BB] = V;
       return V;
     }
@@ -253,7 +256,9 @@ class SSAUpdaterImpl {
 
           // Treat an unreachable predecessor as a definition with 'poison'.
           if (Pred->BlkNum == 0) {
-            Pred->AvailableVal = Traits::GetPoisonVal(Pred->BB, Updater);
+            Pred->AvailableVal = UndefinedVal
+                                     ? *UndefinedVal
+                                     : Traits::GetPoisonVal(Pred->BB, Updater);
             (*AvailableVals)[Pred->BB] = Pred->AvailableVal;
             Pred->DefBB = Pred;
             Pred->BlkNum = PseudoEntry->BlkNum;

llvm/lib/Transforms/Scalar/SROA.cpp

@@ -5479,7 +5479,7 @@ void SROA::clobberUse(Use &U) {
     }
 }
 
-// A basic LoadAndStorePromoter that does not remove store nodes.
+/// A basic LoadAndStorePromoter that does not remove store nodes.
 class BasicLoadAndStorePromoter : public LoadAndStorePromoter {
 public:
   BasicLoadAndStorePromoter(ArrayRef<const Instruction *> Insts, SSAUpdater &S)
@@ -5505,14 +5505,13 @@ bool SROA::propagateStoredValuesToLoads(AllocaInst &AI, AllocaSlices &AS) {
     while (PartitionEnd != AS.end() &&
            (PartitionEnd->beginOffset() < EndOffset ||
             PartitionEnd->endOffset() <= EndOffset)) {
-      EndOffset = std::max(EndOffset, PartitionEnd->endOffset());
       if (AllSameAndValid) {
         AllSameAndValid &= PartitionEnd->beginOffset() == BeginOffset &&
                            PartitionEnd->endOffset() == EndOffset;
         Instruction *User =
             cast<Instruction>(PartitionEnd->getUse()->getUser());
         if (isa<LoadInst>(User) || isa<StoreInst>(User)) {
-          // LoadAndStorePromoter requires all the types are the same.
+          // LoadAndStorePromoter requires all the types to be the same.
           Type *UserTy = getLoadStoreType(User);
           if (PartitionType && UserTy != PartitionType)
             AllSameAndValid = false;
@@ -5521,6 +5520,7 @@ bool SROA::propagateStoredValuesToLoads(AllocaInst &AI, AllocaSlices &AS) {
         } else if (!isAssumeLikeIntrinsic(User))
           AllSameAndValid = false;
       }
+      EndOffset = std::max(EndOffset, PartitionEnd->endOffset());
       ++PartitionEnd;
     }
 
@@ -5534,8 +5534,7 @@ bool SROA::propagateStoredValuesToLoads(AllocaInst &AI, AllocaSlices &AS) {
       BasicLoadAndStorePromoter Promoter(Insts, SSA);
       // Add a zero value at the point of the alloca, to prevent the SSA updater
       // replacing loads with poison which would not be valid for padded loads.
-      SSA.AddAvailableValue(AI.getParent(),
-                            Constant::getNullValue(PartitionType));
+      SSA.SetUndefinedVal(Constant::getNullValue(PartitionType));
       Promoter.run(Insts);
     }
 

llvm/lib/Transforms/Utils/SSAUpdater.cpp

@@ -135,9 +135,10 @@ Value *SSAUpdater::GetValueInMiddleOfBlock(BasicBlock *BB) {
     }
   }
 
-  // If there are no predecessors, just return poison.
+  // If there are no predecessors, just return poison / the undefined input
+  // value.
   if (PredValues.empty())
-    return PoisonValue::get(ProtoType);
+    return GetUndefinedVal(ProtoType);
 
   // Otherwise, if all the merged values are the same, just use it.
   if (SingularValue)
@@ -359,10 +360,20 @@ Value *SSAUpdater::GetValueAtEndOfBlockInternal(BasicBlock *BB) {
   if (Value *V = AvailableVals[BB])
     return V;
 
-  SSAUpdaterImpl<SSAUpdater> Impl(this, &AvailableVals, InsertedPHIs);
+  SSAUpdaterImpl<SSAUpdater> Impl(this, &AvailableVals, InsertedPHIs,
+                                  UndefinedVal ? std::optional(UndefinedVal)
+                                               : std::nullopt);
   return Impl.GetValue(BB);
 }
 
+void SSAUpdater::SetUndefinedVal(Value *V) { UndefinedVal = V; }
+
+Value *SSAUpdater::GetUndefinedVal(Type *Ty) {
+  if (!UndefinedVal)
+    UndefinedVal = PoisonValue::get(Ty);
+  return UndefinedVal;
+}
+
 //===----------------------------------------------------------------------===//
 // LoadAndStorePromoter Implementation
 //===----------------------------------------------------------------------===//
@@ -484,7 +495,8 @@ void LoadAndStorePromoter::run(const SmallVectorImpl<Instruction *> &Insts) {
     replaceLoadWithValue(ALoad, NewVal);
 
     // Avoid assertions in unreachable code.
-    if (NewVal == ALoad) NewVal = PoisonValue::get(NewVal->getType());
+    if (NewVal == ALoad)
+      NewVal = SSA.GetUndefinedVal(NewVal->getType());
     ALoad->replaceAllUsesWith(NewVal);
     ReplacedLoads[ALoad] = NewVal;
   }

llvm/test/Transforms/SROA/non-capturing-call-readonly.ll

@@ -838,7 +838,7 @@ define i8 @dont_transform_load_only() {
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[A:%.*]] = alloca i8, align 1
 ; CHECK-NEXT:    call void @byte_user_of_alloca(ptr [[A]])
-; CHECK-NEXT:    ret i8 poison
+; CHECK-NEXT:    ret i8 0
 ;
 entry:
   %a = alloca i8

llvm/test/Transforms/SROA/readonlynocapture.ll

@@ -150,7 +150,7 @@ define i32 @notdominating() {
 ; CHECK-NEXT:    store i32 0, ptr [[A]], align 4
 ; CHECK-NEXT:    store i32 1, ptr [[B]], align 4
 ; CHECK-NEXT:    call void @callee(ptr [[A]])
-; CHECK-NEXT:    ret i32 poison
+; CHECK-NEXT:    ret i32 0
 ;
   %a = alloca {i32, i32}
   %b = getelementptr i32, ptr %a, i32 1
@@ -249,7 +249,7 @@ define void @incompletestruct(i1 %b, i1 %c) {
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[LII:%.*]] = alloca [[STRUCT_LOADIMMEDIATEINFO:%.*]], align 4
 ; CHECK-NEXT:    call void @llvm.lifetime.start.p0(i64 4, ptr nonnull [[LII]])
-; CHECK-NEXT:    [[BF_CLEAR4:%.*]] = and i32 poison, -262144
+; CHECK-NEXT:    [[BF_CLEAR4:%.*]] = and i32 0, -262144
 ; CHECK-NEXT:    [[BF_SET5:%.*]] = select i1 [[B:%.*]], i32 196608, i32 131072
 ; CHECK-NEXT:    [[BF_SET12:%.*]] = or disjoint i32 [[BF_SET5]], [[BF_CLEAR4]]
 ; CHECK-NEXT:    store i32 [[BF_SET12]], ptr [[LII]], align 4