[llvm] 3d475df - [Mem2Reg] Consistently preserve nonnull assume for uninit load

[Nikita Popov via llvm-commits]

Author: Nikita Popov
Date: 2022-07-12T12:53:08+02:00
New Revision: 3d475dfeb990df5c112535ecc4ec6ffc8f18d4b1

URL: https://github.com/llvm/llvm-project/commit/3d475dfeb990df5c112535ecc4ec6ffc8f18d4b1
DIFF: https://github.com/llvm/llvm-project/commit/3d475dfeb990df5c112535ecc4ec6ffc8f18d4b1.diff

LOG: [Mem2Reg] Consistently preserve nonnull assume for uninit load

When performing a !nonnull load from uninitialized memory, we
should preserve the nonnull assume just like in all other cases.
We already do this correctly in the generic mem2reg code, but
don't handle this case when using the optimized single-block
implementation.

Make sure that the optimized implementation exhibits the same
behavior as the generic implementation.

Added: 
    

Modified: 
    llvm/lib/Transforms/Utils/PromoteMemoryToRegister.cpp
    llvm/test/Transforms/Mem2Reg/preserve-nonnull-load-metadata.ll

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Transforms/Utils/PromoteMemoryToRegister.cpp b/llvm/lib/Transforms/Utils/PromoteMemoryToRegister.cpp
index aff692b36288e..bec1db896efbe 100644
--- a/llvm/lib/Transforms/Utils/PromoteMemoryToRegister.cpp
+++ b/llvm/lib/Transforms/Utils/PromoteMemoryToRegister.cpp
@@ -488,31 +488,33 @@ static bool promoteSingleBlockAlloca(AllocaInst *AI, const AllocaInfo &Info,
         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.
-        LI->replaceAllUsesWith(UndefValue::get(LI->getType()));
+        ReplVal = UndefValue::get(LI->getType());
       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.
-      // Note, if the load was marked as nonnull we don't want to lose that
-      // information when we erase it. So we preserve it with an assume.
-      Value *ReplVal = std::prev(I)->second->getOperand(0);
-      if (AC && LI->getMetadata(LLVMContext::MD_nonnull) &&
-          !isKnownNonZero(ReplVal, DL, 0, AC, LI, &DT))
-        addAssumeNonNull(AC, LI);
+      // Otherwise, there was a store before this load, the load takes its
+      // value.
+      ReplVal = std::prev(I)->second->getOperand(0);
+    }
 
-      // If the replacement value is the load, this must occur in unreachable
-      // code.
-      if (ReplVal == LI)
-        ReplVal = PoisonValue::get(LI->getType());
+    // Note, if the load was marked as nonnull we don't want to lose that
+    // information when we erase it. So we preserve it with an assume.
+    if (AC && LI->getMetadata(LLVMContext::MD_nonnull) &&
+        !isKnownNonZero(ReplVal, DL, 0, AC, LI, &DT))
+      addAssumeNonNull(AC, LI);
 
-      LI->replaceAllUsesWith(ReplVal);
-    }
+    // If the replacement value is the load, this must occur in unreachable
+    // code.
+    if (ReplVal == LI)
+      ReplVal = PoisonValue::get(LI->getType());
 
+    LI->replaceAllUsesWith(ReplVal);
     LI->eraseFromParent();
     LBI.deleteValue(LI);
   }

diff  --git a/llvm/test/Transforms/Mem2Reg/preserve-nonnull-load-metadata.ll b/llvm/test/Transforms/Mem2Reg/preserve-nonnull-load-metadata.ll
index cabe06130292c..1eb50f6483d03 100644
--- a/llvm/test/Transforms/Mem2Reg/preserve-nonnull-load-metadata.ll
+++ b/llvm/test/Transforms/Mem2Reg/preserve-nonnull-load-metadata.ll
@@ -95,6 +95,8 @@ fin:
 define float* @no_store_single_load() {
 ; CHECK-LABEL: @no_store_single_load(
 ; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[TMP0:%.*]] = icmp ne float* undef, null
+; CHECK-NEXT:    call void @llvm.assume(i1 [[TMP0]])
 ; CHECK-NEXT:    ret float* undef
 ;
 entry: