[EarlyCSE] Insert unreachable for !noundef undef load

This extends llvm#96639 for mem2reg to earlycse. If we try to replace
a load with !noundef metadata by undef, insert a non-terminator
unreachable.

I've added a handleLoadOfUndef() helper for this purpose, which is
shared between mem2reg and earlycse (and in the future also gvn
etc). This also provides a place to experiment with replacing
uninit loads with freeze poison.

Author: nikic

llvm/include/llvm/Transforms/Utils/Local.h

@@ -428,6 +428,11 @@ void copyMetadataForLoad(LoadInst &Dest, const LoadInst &Source);
 /// its original position.
 void patchReplacementInstruction(Instruction *I, Value *Repl);
 
+/// Handle undef values when replacing a load \p I with the loaded value
+/// \p Repl. If \p Repl is undef, this will introduce a non-terminator
+/// unreachable for !noundef loads.
+Value *handleLoadOfUndef(Instruction &I, Value *Repl);
+
 // Replace each use of 'From' with 'To', if that use does not belong to basic
 // block where 'From' is defined. Returns the number of replacements made.
 unsigned replaceNonLocalUsesWith(Instruction *From, Value *To);

llvm/lib/Transforms/Scalar/EarlyCSE.cpp

@@ -1585,6 +1585,7 @@ bool EarlyCSE::processNode(DomTreeNode *Node) {
           LLVM_DEBUG(dbgs() << "Skipping due to debug counter\n");
           continue;
         }
+        Op = handleLoadOfUndef(Inst, Op);
         if (InVal.IsLoad)
           if (auto *I = dyn_cast<Instruction>(Op))
             combineMetadataForCSE(I, &Inst, false);

llvm/lib/Transforms/Utils/Local.cpp

@@ -3450,6 +3450,19 @@ void llvm::patchReplacementInstruction(Instruction *I, Value *Repl) {
   combineMetadataForCSE(ReplInst, I, false);
 }
 
+Value *llvm::handleLoadOfUndef(Instruction &I, Value *Repl) {
+  if (isa<UndefValue>(Repl) && I.hasMetadata(LLVMContext::MD_noundef)) {
+    // Insert non-terminator unreachable.
+    LLVMContext &Ctx = I.getContext();
+    new StoreInst(ConstantInt::getTrue(Ctx),
+                  PoisonValue::get(PointerType::getUnqual(Ctx)),
+                  /*isVolatile=*/false, Align(1), &I);
+    return PoisonValue::get(Repl->getType());
+  }
+
+  return Repl;
+}
+
 template <typename RootType, typename ShouldReplaceFn>
 static unsigned replaceDominatedUsesWith(Value *From, Value *To,
                                          const RootType &Root,

llvm/lib/Transforms/Utils/PromoteMemoryToRegister.cpp

@@ -453,21 +453,14 @@ static void addAssumeNonNull(AssumptionCache *AC, LoadInst *LI) {
 static void convertMetadataToAssumes(LoadInst *LI, Value *Val,
                                      const DataLayout &DL, AssumptionCache *AC,
                                      const DominatorTree *DT) {
-  if (isa<UndefValue>(Val) && LI->hasMetadata(LLVMContext::MD_noundef)) {
-    // Insert non-terminator unreachable.
-    LLVMContext &Ctx = LI->getContext();
-    new StoreInst(ConstantInt::getTrue(Ctx),
-                  PoisonValue::get(PointerType::getUnqual(Ctx)),
-                  /*isVolatile=*/false, Align(1), LI);
-    return;
-  }
-
   // If the load was marked as nonnull we don't want to lose that information
   // when we erase this Load. So we preserve it with an assume. As !nonnull
   // returns poison while assume violations are immediate undefined behavior,
   // we can only do this if the value is known non-poison.
+  // Skip this for undef values, because they have already been converted to
+  // non-terminator unreachable.
   if (AC && LI->getMetadata(LLVMContext::MD_nonnull) &&
-      LI->getMetadata(LLVMContext::MD_noundef) &&
+      LI->getMetadata(LLVMContext::MD_noundef) && !isa<UndefValue>(Val) &&
       !isKnownNonZero(Val, SimplifyQuery(DL, DT, AC, LI)))
     addAssumeNonNull(AC, LI);
 }
@@ -567,6 +560,7 @@ rewriteSingleStoreAlloca(AllocaInst *AI, AllocaInfo &Info, LargeBlockInfo &LBI,
     if (ReplVal == LI)
       ReplVal = PoisonValue::get(LI->getType());
 
+    ReplVal = handleLoadOfUndef(*LI, ReplVal);
     convertMetadataToAssumes(LI, ReplVal, DL, AC, &DT);
     LI->replaceAllUsesWith(ReplVal);
     LI->eraseFromParent();
@@ -676,6 +670,7 @@ promoteSingleBlockAlloca(AllocaInst *AI, const AllocaInfo &Info,
       ReplVal = std::prev(I)->second->getOperand(0);
     }
 
+    ReplVal = handleLoadOfUndef(*LI, ReplVal);
     convertMetadataToAssumes(LI, ReplVal, DL, AC, &DT);
 
     // If the replacement value is the load, this must occur in unreachable
@@ -1162,6 +1157,7 @@ void PromoteMem2Reg::RenamePass(BasicBlock *BB, BasicBlock *Pred,
         continue;
 
       Value *V = IncomingVals[AI->second];
+      V = handleLoadOfUndef(*LI, V);
       convertMetadataToAssumes(LI, V, SQ.DL, AC, &DT);
 
       // Anything using the load now uses the current value.

llvm/test/Transforms/EarlyCSE/flags.ll

@@ -106,7 +106,8 @@ define ptr @store_to_load_forward(ptr %p, ptr %p2) {
 define i32 @load_undef_noundef(ptr %p) {
 ; CHECK-LABEL: @load_undef_noundef(
 ; CHECK-NEXT:    store i32 undef, ptr [[P:%.*]], align 4
-; CHECK-NEXT:    ret i32 undef
+; CHECK-NEXT:    store i1 true, ptr poison, align 1
+; CHECK-NEXT:    ret i32 poison
 ;
   store i32 undef, ptr %p
   %v = load i32, ptr %p, !noundef !{}

llvm/test/Transforms/Mem2Reg/preserve-nonnull-load-metadata.ll

@@ -144,7 +144,7 @@ define ptr @no_store_single_load_noundef() {
 ; CHECK-LABEL: @no_store_single_load_noundef(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    store i1 true, ptr poison, align 1
-; CHECK-NEXT:    ret ptr undef
+; CHECK-NEXT:    ret ptr poison
 ;
 entry:
   %buf = alloca ptr
@@ -158,10 +158,10 @@ define ptr @no_store_multiple_loads_noundef(i1 %c) {
 ; CHECK-NEXT:    br i1 [[C:%.*]], label [[IF:%.*]], label [[ELSE:%.*]]
 ; CHECK:       if:
 ; CHECK-NEXT:    store i1 true, ptr poison, align 1
-; CHECK-NEXT:    ret ptr undef
+; CHECK-NEXT:    ret ptr poison
 ; CHECK:       else:
 ; CHECK-NEXT:    store i1 true, ptr poison, align 1
-; CHECK-NEXT:    ret ptr undef
+; CHECK-NEXT:    ret ptr poison
 ;
 entry:
   %buf = alloca ptr
@@ -180,7 +180,7 @@ define ptr @no_store_single_load_nonnull_noundef() {
 ; CHECK-LABEL: @no_store_single_load_nonnull_noundef(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    store i1 true, ptr poison, align 1
-; CHECK-NEXT:    ret ptr undef
+; CHECK-NEXT:    ret ptr poison
 ;
 entry:
   %buf = alloca ptr
@@ -194,10 +194,10 @@ define ptr @no_store_multiple_loads_nonnull_noundef(i1 %c) {
 ; CHECK-NEXT:    br i1 [[C:%.*]], label [[IF:%.*]], label [[ELSE:%.*]]
 ; CHECK:       if:
 ; CHECK-NEXT:    store i1 true, ptr poison, align 1
-; CHECK-NEXT:    ret ptr undef
+; CHECK-NEXT:    ret ptr poison
 ; CHECK:       else:
 ; CHECK-NEXT:    store i1 true, ptr poison, align 1
-; CHECK-NEXT:    ret ptr undef
+; CHECK-NEXT:    ret ptr poison
 ;
 entry:
   %buf = alloca ptr