Preserve nonnull metadata on Loads through SROA & mem2reg.

[luqmana]

Summary

https://llvm.org/bugs/show_bug.cgi?id=31142 :

SROA was dropping the nonnull metadata on loads from allocas that got optimized out. This patch simply preserves nonnull metadata on loads through SROA and mem2reg.

[chandlerc]

I don't get it, this patch looks like it is *removing* SROA's existing logic to propagate metadata? And it doesn't add any new test cases, just deletes an existing test?

[luqmana]

Whoops, that was my mistake. The diff was inverted and should be corrected now.

[davide]

Are you always guaranteed that the range doesn't contain zero?

[luqmana]

In D27114#605657, @davide wrote:

Are you always guaranteed that the range doesn't contain zero?

Is it valid for the range to contain 0 if the load was already marked as nonnull?

[arielb1]

Could you also preserve !range metadata? It behaves similarly to nonzero, and inttoptr/ptrtoint interchanges them.

[efriedma]

lib/Transforms/Utils/PromoteMemoryToRegister.cpp

405

This isn't a legal transform: you're assigning "nonnull" to a completely unrelated instruction (which could have other users) just because it happens to be a LoadInst. I guess you could get away with this if you can prove there aren't any other uses of the value, though.

[luqmana]

lib/Transforms/Utils/PromoteMemoryToRegister.cpp

405

Hmmm, I could just add a check for it being the only use but is it illegal? Just walking through this code:

• We have an alloca with a single store (OnlyStore)
• The value stored was from a Load (ReplVal)
• Now, we have a Load (LI) from this alloca that's dominated by this store and it has !nonnull.

Adding !nonnull to ReplVal seems reasonable since (assuming) the !nonnull on LI was valid, we know that value must be non null. So, can we not just propagate it back to ReplVal?

[efriedma]

lib/Transforms/Utils/PromoteMemoryToRegister.cpp

405

(assuming) the !nonnull on LI was valid, we know that value must be non null

!nonnull basically means "if this load would return null, we can treat it as poison". So you can't propagate backwards unless you can prove that the load will actually execute (LI post-dominates ReplVal), and that the loaded value will be used in a way that causes undefined behavior (a branch or call which uses LI post-dominates ReplVal). In practice, this is essentially impossible to check in LLVM.

[arielb1]

lib/Transforms/Utils/PromoteMemoryToRegister.cpp

405

I am quite sure I was told that !nonnull causes UB rather than poison. Is there some real documentation for that?

[efriedma]

I am quite sure I was told that !nonnull causes UB rather than poison. Is there some real documentation for that?

Err, no, you're probably right; I haven't really worked with it, and LangRef doesn't say anything, so "must" defaults to undefined behavior.

You still can't propagate backwards without proving post-dominance.

[luqmana]

In D27114#626971, @efriedma wrote:

I am quite sure I was told that !nonnull causes UB rather than poison. Is there some real documentation for that?

Err, no, you're probably right; I haven't really worked with it, and LangRef doesn't say anything, so "must" defaults to undefined behavior.

You still can't propagate backwards without proving post-dominance.

That should hold in this case no? Lines 361-381 assert that LI is dominated by the Store and the Store trivially dominates ReplVal?

[efriedma]

Lines 361-381 assert that LI is dominated by the Store and the Store trivially dominates ReplVal?

Post-dominates, not dominates... they're sort of similar, but post-dominance involves following the control flow edges the other way. A compiler textbook should cover this.

[luqmana]

Address comments from review.

• Preserve nonnull metadata on Loads through SROA & mem2reg.
• Pass PostDominatorTree to PromoteMemoryToRegister.
• Assert Post-Dominance of Load from alloca before propogating metadata.

[luqmana]

In D27114#627071, @efriedma wrote:

Lines 361-381 assert that LI is dominated by the Store and the Store trivially dominates ReplVal?

Post-dominates, not dominates... they're sort of similar, but post-dominance involves following the control flow edges the other way. A compiler textbook should cover this.

Ah yes, missed that detail. Updated to also consider post-dominance.

[luqmana]

Remove unnecessary changes.

[efriedma]

PostDominatorTree doesn't provide the kind of post-dominance you need for this transformation: LLVM basic blocks have implicit edges which exit the function early, and those edges aren't reflected in the tree. For example, execution doesn't actually continue after a call to exit(). (See also the llvm::isGuaranteedToTransferExecutionToSuccessor() utility.)

[luqmana]

In D27114#627264, @efriedma wrote:

PostDominatorTree doesn't provide the kind of post-dominance you need for this transformation: LLVM basic blocks have implicit edges which exit the function early, and those edges aren't reflected in the tree. For example, execution doesn't actually continue after a call to exit(). (See also the llvm::isGuaranteedToTransferExecutionToSuccessor() utility.)

Hmm, is it possible then to do this sort of analysis in LLVM in this context? Or is this optimization just not possible currently?

[efriedma]

I don't think there's any existing analysis which computes the sort of post-dominator tree you want.

Anyway, I'm not sure this is really what you want... have you experimented with inserting llvm.assume calls?

[luqmana]

In D27114#628254, @efriedma wrote:

I don't think there's any existing analysis which computes the sort of post-dominator tree you want.

Anyway, I'm not sure this is really what you want... have you experimented with inserting llvm.assume calls?

The problem wasn't that. We do have llvm.assume which InstCombine turns assume( (load addr) != null ) to a load addr !nonnull. Then, SROA/mem2reg end up eating the !nonnull metadata.

Some more details at https://github.com/rust-lang/rust/issues/37945 which was linked in the bug link in the description.

[efriedma]

The problem wasn't that. We do have llvm.assume which InstCombine turns assume( (load addr) != null ) to a load addr !nonnull. Then, SROA/mem2reg end up eating the !nonnull metadata.

You can do the reverse transform in mem2reg: load addr !nonnull -> assume( (load addr) != null ), then when the load gets erased the nonnull assumption is preserved.

[chandlerc]

In D27114#628901, @efriedma wrote:

The problem wasn't that. We do have llvm.assume which InstCombine turns assume( (load addr) != null ) to a load addr !nonnull. Then, SROA/mem2reg end up eating the !nonnull metadata.

You can do the reverse transform in mem2reg: load addr !nonnull -> assume( (load addr) != null ), then when the load gets erased the nonnull assumption is preserved.

Or we could stop instcombine from doing this in the first place.

I think we should canonicalize on the assume formulation as it is strictly more general than the nonnull metadata formulation.

[luqmana]

Use assume to preserve nonnull-ness of Load.

[luqmana]

In D27114#629295, @chandlerc wrote:

In D27114#628901, @efriedma wrote:

The problem wasn't that. We do have llvm.assume which InstCombine turns assume( (load addr) != null ) to a load addr !nonnull. Then, SROA/mem2reg end up eating the !nonnull metadata.

You can do the reverse transform in mem2reg: load addr !nonnull -> assume( (load addr) != null ), then when the load gets erased the nonnull assumption is preserved.

Or we could stop instcombine from doing this in the first place.

I think we should canonicalize on the assume formulation as it is strictly more general than the nonnull metadata formulation.

For this change I just implemented @efriedma's suggestion. I think changing how nonnull-ness is canonicalized is outside the scope of this.

[chandlerc]

Could you clean up the formatting to match LLVM coding standards? At the very least clang-format would help here (lots of lines over 80 columns, indent looks all wrong).

[luqmana]

Run clang-format.

[spatel]

lib/Transforms/Utils/PromoteMemoryToRegister.cpp

973–986

Please make a helper function so we don't have duplicated code.

[luqmana]

Move duplicate code to helper function.

[spatel]

test/Transforms/SROA/preserve-nonnull.ll

2

I don't think you want to include -instcombine here and create a dependency on its canonicalization of the assume to metadata - especially since that may change based on the comments in this review and PR31518 ( https://llvm.org/bugs/show_bug.cgi?id=31518 ).

You can use "utils/update_test_checks.py" to auto-generate exact CHECK lines for your test.

[luqmana]

Simplify test and remove instcombine dependency.

[spatel]

Thanks for the changes.

Someone else can correct me if I'm wrong, but I think we're delaying improvements in nonnull optimization until we have a clang hack to avoid disastrous optimizations of libc functions:
http://lists.llvm.org/pipermail/cfe-dev/2017-January/052066.html

[chandlerc]

In D27114#636348, @spatel wrote:

Thanks for the changes.

Someone else can correct me if I'm wrong, but I think we're delaying improvements in nonnull optimization until we have a clang hack to avoid disastrous optimizations of libc functions:
http://lists.llvm.org/pipermail/cfe-dev/2017-January/052066.html

Yes, that thread seems largely in favor of getting something in place in Clang first.

If no one else works on a patch to Clang I may in a week or so, but sadly my queue is very full at the moment.

[hfinkel]

In D27114#636391, @chandlerc wrote:

In D27114#636348, @spatel wrote:

Thanks for the changes.

Someone else can correct me if I'm wrong, but I think we're delaying improvements in nonnull optimization until we have a clang hack to avoid disastrous optimizations of libc functions:
http://lists.llvm.org/pipermail/cfe-dev/2017-January/052066.html

Yes, that thread seems largely in favor of getting something in place in Clang first.

If no one else works on a patch to Clang I may in a week or so, but sadly my queue is very full at the moment.

I don't think that's relevant to this patch. Once we have the nonnull in metadata form, we should preserve it. The particular problem with optimizing based on nonnull function-argument annotations is propagating the nonnull annotation from a callsite back into the caller. For that, we should wait on associated Clang updates (so that we can avoid doing it for memcpy and friends based on header-file annotations). This change should proceed.

[chandlerc]

Bunch of little comments below.

I'd like to see two high level thinsg before this goes in though:

1. I want to make sure Eli is happy with this approach w.r.t. any semantic differences between !nonnull metadata and the assume. Haven't seen an update from him since the switch of formulation.

2. I'd like at least a test case and a FIXME around the case where SROA rewrites the load and store to be integer loads and stores. This will happen fairly often due to things like unions, memcpy and other "bag of bits" interpretations of memory operations. We should work to not destroy !nonnull in the process by translating it to an assume. I'm happy for this to be handled in a subsequent patch, but I'd like to at least document the issue for anyone who ends up working on it.

lib/Transforms/Utils/PromoteMemoryToRegister.cpp

41

clang-format's header sort is wrong here (because PromoteMemToReg != PromoteMemoryToRegister).

306–313

This deosn't look like the right formatting... was clang-format misconfigured?

976–977

Here as well.

test/Transforms/SROA/preserve-nonnull.ll

7–8

Do you need these?

10

I would at least add the @ and maybe the define to this so that random other thinsg containing this string don't match confusingly down the road.

16

it's good to not use unnamed values in tests as they make future edits to the test brittle.

[luqmana]

Fix some formatting and update test.

[luqmana]

In D27114#637102, @chandlerc wrote:

Bunch of little comments below.

I'd like to see two high level thinsg before this goes in though:

1. I want to make sure Eli is happy with this approach w.r.t. any semantic differences between !nonnull metadata and the assume. Haven't seen an update from him since the switch of formulation.

2. I'd like at least a test case and a FIXME around the case where SROA rewrites the load and store to be integer loads and stores. This will happen fairly often due to things like unions, memcpy and other "bag of bits" interpretations of memory operations. We should work to not destroy !nonnull in the process by translating it to an assume. I'm happy for this to be handled in a subsequent patch, but I'd like to at least document the issue for anyone who ends up working on it.

Updated the patch to address the little comments. As for the second point, I think that's better handled in a separate patch. If Eli is happy with the current approach I'd like to move forward with just this patch.

[efriedma]

I want to make sure Eli is happy with this approach w.r.t. any semantic differences between !nonnull metadata and the assume. Haven't seen an update from him since the switch of formulation.

Semantically, this seems consistent with existing code dealing with !nonnull. That said, creating calls to @llvm.assume has the potential to cause performance regressions; there are a bunch of places in the optimizer which don't handle them well, so it's probably worth checking that this change doesn't create more problems than it solves.

[luqmana]

In D27114#637151, @efriedma wrote:

I want to make sure Eli is happy with this approach w.r.t. any semantic differences between !nonnull metadata and the assume. Haven't seen an update from him since the switch of formulation.

Semantically, this seems consistent with existing code dealing with !nonnull. That said, creating calls to @llvm.assume has the potential to cause performance regressions; there are a bunch of places in the optimizer which don't handle them well, so it's probably worth checking that this change doesn't create more problems than it solves.

Are there any existing performance tests so I can determine if this might cause a regression?

[efriedma]

Well, you can run the LLVM testsuite (http://llvm.org/docs/lnt/quickstart.html)... but that's written entirely written in C and C++, and I don't think clang produces nonnull metadata at all. Maybe you have some Rust performance tests you could use?

[spatel]

In D27114#637151, @efriedma wrote:

Semantically, this seems consistent with existing code dealing with !nonnull. That said, creating calls to @llvm.assume has the potential to cause performance regressions; there are a bunch of places in the optimizer which don't handle them well, so it's probably worth checking that this change doesn't create more problems than it solves.

Apologies if this isn't directly related to this patch, but one thing I noticed and mentioned in D28337: InstCombine is not calling into InstSimplify (and hence value tracking) with a context instruction in many cases. And so computeKnownBitsFromAssume() is not triggered in those cases. If we fix that (assuming that the lack of context instruction optional parameter is an oversight), we should make better use of the assumption machinery.

[efriedma]

Apologies if this isn't directly related to this patch, but one thing I noticed and mentioned in D28337: InstCombine is not calling into InstSimplify (and hence value tracking) with a context instruction in many cases. And so computeKnownBitsFromAssume() is not triggered in those cases. If we fix that (assuming that the lack of context instruction optional parameter is an oversight), we should make better use of the assumption machinery.

I was more referring to the fact than m_OneUse etc. doesn't work the way you want it to in the presence of llvm.assume calls.

[luqmana]

In D27114#637187, @efriedma wrote:

Well, you can run the LLVM testsuite (http://llvm.org/docs/lnt/quickstart.html)... but that's written entirely written in C and C++, and I don't think clang produces nonnull metadata at all. Maybe you have some Rust performance tests you could use?

Hmmm, I recall us not using as much assume due to supposed LLVM slowdowns but don't have any on hand.

[luqmana]

So at least running the test suite I didn't come across any performance regressions. Could we land this and observe the buildbots or something for any regressions and back out as necessary?

[luqmana]

Ping

[luqmana]

Ping

[efriedma]

Chandler, are you planning to continue reviewing this?

[luqmana]

Ping

[hfinkel]

lib/Transforms/Utils/PromoteMemoryToRegister.cpp

408

Even if we're going to do this in general, we shouldn't do this when we can otherwise prove that the address was nonnull. Can you call isKnownNonNullAt and only add the intrinsic when we need it?

The same applies to other places where you call addAssumeNonNull.

[luqmana]

Don't add assume if the value is already known to be nonnull.

[luqmana]

Ping?

[luqmana]

Ping

[spatel]

I'm not familiar with this code, but I'd like to see this functionality, so any assistance in reviewing is appreciated.
Other than the inline comment, this is good now?

lib/Transforms/Scalar/SROA.cpp

2391–2404

Is the intent to extend this for some subset of metadata types beyond just MD_nonnull?

Either way, just make this one line for now?
NewLI->copyMetadata(LI, LLVMContext::MD_nonnull);

[luqmana]

Use copyMetadata method instead of manually copying metadata.

[luqmana]

Ping?

[spatel]

Hal's earlier comment said we don't need to hold this up for clang, but for reference that fix is proposed now:
D30806

@chandlerc / @davide / @efriedma / @hfinkel - can you continue your review? I don't know enough to approve.

[efriedma]

I was hoping Chandler would continue to review, but I guess I can do it.

The test here is really bare-bones; it doesn't really cover all the relevant codepaths. It doesn't check the effect of the isKnownNonNullAt, and it doesn't cover one of the codepaths in PromoteMemoryToRegister.cpp at all.

Otherwise looks fine per earlier discussion.

[chandlerc]

In D27114#704478, @efriedma wrote:

I was hoping Chandler would continue to review, but I guess I can do it.

FWIW, I'm happy to continue it or for you to continue it. I was planning to GC these threads when the clang patch actually landed, but if you can get to them sooner, so much the better.

The test here is really bare-bones; it doesn't really cover all the relevant codepaths. It doesn't check the effect of the isKnownNonNullAt, and it doesn't cover one of the codepaths in PromoteMemoryToRegister.cpp at all.

FWIW, I think this really needs to be addressed before it lands.

[luqmana]

I added another test for all the relevant codepaths in PromoteMemoryToRegister.
I also addressed one case I missed earlier: when the alloca in question has
multiple stores but all within the same basic block.

[efriedma]

LGTM.

Luqman, thanks for putting up with the long delay in reviewing this.

Do you have commit access, or should I commit this for you?

[luqmana]

In D27114#707685, @efriedma wrote:

LGTM.

Luqman, thanks for putting up with the long delay in reviewing this.

Do you have commit access, or should I commit this for you?

Thanks! No worries and I can commit it.

lib/Transforms/Scalar/SROA.cpp

if (VecTy) {
                NewEndOffset == NewAllocaEndOffset &&
                (canConvertValue(DL, NewAllocaTy, TargetTy) ||
                 (IsLoadPastEnd && NewAllocaTy->isIntegerTy() &&
                  TargetTy->isIntegerTy()))) {
       LoadInst *NewLI = IRB.CreateAlignedLoad(&NewAI, NewAI.getAlignment(),
                                               LI.isVolatile(), LI.getName());
       if (LI.isVolatile())
         NewLI->setAtomic(LI.getOrdering(), LI.getSynchScope());
+
+      // Try to preserve nonnull metadata
+      if (TargetTy->isPointerTy())
+        NewLI->copyMetadata(LI, LLVMContext::MD_nonnull);
       V = NewLI;
 
       // If this is an integer load past the end of the slice (which means the
       // bytes outside the slice are undef or this load is dead) just forcibly
       // fix the integer size with correct handling of endianness.
       if (auto *AITy = dyn_cast<IntegerType>(NewAllocaTy))
         if (auto *TITy = dyn_cast<IntegerType>(TargetTy))
           if (AITy->getBitWidth() < TITy->getBitWidth()) {
             V = IRB.CreateZExt(V, TITy, "load.ext");
             if (DL.isBigEndian())
               V = IRB.CreateShl(V, TITy->getBitWidth() - AITy->getBitWidth(),
                                 "endian_shift");
           }
     } else {
       Type *LTy = TargetTy->getPointerTo();
       LoadInst *NewLI = IRB.CreateAlignedLoad(getNewAllocaSlicePtr(IRB, LTy),
                                               getSliceAlign(TargetTy),
                                               LI.isVolatile(), LI.getName());
       if (LI.isVolatile())

lib/Transforms/Utils/PromoteMemoryToRegister.cpp

 // This file promotes memory references to be register references.  It promotes
 // alloca instructions which only have loads and stores as uses.  An alloca is
 // transformed by using iterated dominator frontiers to place PHI nodes, then
 // traversing the function in depth-first order to rewrite loads and stores as
 // appropriate.
 //
 //===----------------------------------------------------------------------===//
 
-#include "llvm/Transforms/Utils/PromoteMemToReg.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/Statistic.h"
 #include "llvm/Analysis/AliasSetTracker.h"
+#include "llvm/Analysis/AssumptionCache.h"
 #include "llvm/Analysis/InstructionSimplify.h"
 #include "llvm/Analysis/IteratedDominanceFrontier.h"
 #include "llvm/Analysis/ValueTracking.h"
 #include "llvm/IR/CFG.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/DIBuilder.h"
 #include "llvm/IR/DebugInfo.h"
 #include "llvm/IR/DerivedTypes.h"
 #include "llvm/IR/Dominators.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/IntrinsicInst.h"
 #include "llvm/IR/Metadata.h"
 #include "llvm/IR/Module.h"
 #include "llvm/Transforms/Utils/Local.h"
+#include "llvm/Transforms/Utils/PromoteMemToReg.h"
 #include <algorithm>
 using namespace llvm;
 
 #define DEBUG_TYPE "mem2reg"
 
 STATISTIC(NumLocalPromoted, "Number of alloca's promoted within one block");
 STATISTIC(NumSingleStore,   "Number of alloca's promoted with a single store");
 STATISTIC(NumDeadAlloca,    "Number of dead alloca's removed");
private:
   void RenamePass(BasicBlock *BB, BasicBlock *Pred,
                   RenamePassData::ValVector &IncVals,
                   std::vector<RenamePassData> &Worklist);
   bool QueuePhiNode(BasicBlock *BB, unsigned AllocaIdx, unsigned &Version);
 };
 
 } // end of anonymous namespace
 
+/// Given a LoadInst LI this adds assume(LI != null) after it.
+static void addAssumeNonNull(AssumptionCache *AC, LoadInst *LI) {
+  Function *AssumeIntrinsic =
+      Intrinsic::getDeclaration(LI->getModule(), Intrinsic::assume);
+  ICmpInst *LoadNotNull = new ICmpInst(ICmpInst::ICMP_NE, LI,
+                                       Constant::getNullValue(LI->getType()));
+  LoadNotNull->insertAfter(LI);
+  CallInst *CI = CallInst::Create(AssumeIntrinsic, {LoadNotNull});
+  CI->insertAfter(LoadNotNull);
+  AC->registerAssumption(CI);
+}
+
 static void removeLifetimeIntrinsicUsers(AllocaInst *AI) {
   // Knowing that this alloca is promotable, we know that it's safe to kill all
   // instructions except for load and store.
 
   for (auto UI = AI->user_begin(), UE = AI->user_end(); UI != UE;) {
     Instruction *I = cast<Instruction>(*UI);
     ++UI;
     if (isa<LoadInst>(I) || isa<StoreInst>(I))

 ///
 /// When there is only a single store, we can use the domtree to trivially
 /// replace all of the dominated loads with the stored value. Do so, and return
 /// true if this has successfully promoted the alloca entirely. If this returns
 /// false there were some loads which were not dominated by the single store
 /// and thus must be phi-ed with undef. We fall back to the standard alloca
 /// promotion algorithm in that case.
 static bool rewriteSingleStoreAlloca(AllocaInst *AI, AllocaInfo &Info,
-                                     LargeBlockInfo &LBI,
-                                     DominatorTree &DT,
-                                     AliasSetTracker *AST) {
+                                     LargeBlockInfo &LBI, DominatorTree &DT,
+                                     AliasSetTracker *AST,
+                                     AssumptionCache *AC) {
   StoreInst *OnlyStore = Info.OnlyStore;
   bool StoringGlobalVal = !isa<Instruction>(OnlyStore->getOperand(0));
   BasicBlock *StoreBB = OnlyStore->getParent();
   int StoreIndex = -1;
 
   // Clear out UsingBlocks.  We will reconstruct it here if needed.
   Info.UsingBlocks.clear();
 
for (auto UI = AI->user_begin(), E = AI->user_end(); UI != E;) {
     }
 
     // Otherwise, we *can* safely rewrite this load.
     Value *ReplVal = OnlyStore->getOperand(0);
     // If the replacement value is the load, this must occur in unreachable
     // code.
     if (ReplVal == LI)
       ReplVal = UndefValue::get(LI->getType());
+
+    // 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.
+    if (AC && LI->getMetadata(LLVMContext::MD_nonnull) &&
+        !llvm::isKnownNonNullAt(ReplVal, LI, &DT))
+      addAssumeNonNull(AC, LI);
+
     LI->replaceAllUsesWith(ReplVal);
     if (AST && LI->getType()->isPointerTy())
       AST->deleteValue(LI);
     LI->eraseFromParent();
     LBI.deleteValue(LI);
   }
 
   // Finally, after the scan, check to see if the store is all that is left.

 ///  for (...) {
 ///    int t = *A;
 ///    if (!first_iteration)
 ///      use(t);
 ///    *A = 42;
 ///  }
 static bool promoteSingleBlockAlloca(AllocaInst *AI, const AllocaInfo &Info,
                                      LargeBlockInfo &LBI,
-                                     AliasSetTracker *AST) {
+                                     AliasSetTracker *AST,
+                                     DominatorTree &DT,
+                                     AssumptionCache *AC) {
   // The trickiest case to handle is when we have large blocks. Because of this,
   // this code is optimized assuming that large blocks happen.  This does not
   // significantly pessimize the small block case.  This uses LargeBlockInfo to
   // make it efficient to get the index of various operations in the block.
 
   // Walk the use-def list of the alloca, getting the locations of all stores.
   typedef SmallVector<std::pair<unsigned, StoreInst *>, 64> StoresByIndexTy;
   StoresByIndexTy StoresByIndex;
for (auto UI = AI->user_begin(), E = AI->user_end(); UI != E;) {
     if (I == StoresByIndex.begin()) {
       if (StoresByIndex.empty())
         // If there are no stores, the load takes the undef value.
         LI->replaceAllUsesWith(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
+    } else {
       // Otherwise, there was a store before this load, the load takes its value.
-      LI->replaceAllUsesWith(std::prev(I)->second->getOperand(0));
+      // 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) &&
+          !llvm::isKnownNonNullAt(ReplVal, LI, &DT))
+        addAssumeNonNull(AC, LI);
+
+      LI->replaceAllUsesWith(ReplVal);
+    }
 
     if (AST && LI->getType()->isPointerTy())
       AST->deleteValue(LI);
     LI->eraseFromParent();
     LBI.deleteValue(LI);
   }
 
   // Remove the (now dead) stores and alloca.
for (unsigned AllocaNum = 0; AllocaNum != Allocas.size(); ++AllocaNum) {
 
     // Calculate the set of read and write-locations for each alloca.  This is
     // analogous to finding the 'uses' and 'definitions' of each variable.
     Info.AnalyzeAlloca(AI);
 
     // If there is only a single store to this value, replace any loads of
     // it that are directly dominated by the definition with the value stored.
     if (Info.DefiningBlocks.size() == 1) {
-      if (rewriteSingleStoreAlloca(AI, Info, LBI, DT, AST)) {
+      if (rewriteSingleStoreAlloca(AI, Info, LBI, DT, AST, AC)) {
         // The alloca has been processed, move on.
         RemoveFromAllocasList(AllocaNum);
         ++NumSingleStore;
         continue;
       }
     }
 
     // If the alloca is only read and written in one basic block, just perform a
     // linear sweep over the block to eliminate it.
     if (Info.OnlyUsedInOneBlock &&
-        promoteSingleBlockAlloca(AI, Info, LBI, AST)) {
+        promoteSingleBlockAlloca(AI, Info, LBI, AST, DT, AC)) {
       // The alloca has been processed, move on.
       RemoveFromAllocasList(AllocaNum);
       continue;
     }
 
     // If we haven't computed a numbering for the BB's in the function, do so
     // now.
     if (BBNumbers.empty()) {
if (LoadInst *LI = dyn_cast<LoadInst>(I)) {
         continue;
 
       DenseMap<AllocaInst *, unsigned>::iterator AI = AllocaLookup.find(Src);
       if (AI == AllocaLookup.end())
         continue;
 
       Value *V = IncomingVals[AI->second];
 
+      // 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.
+      if (AC && LI->getMetadata(LLVMContext::MD_nonnull) &&
+          !llvm::isKnownNonNullAt(V, LI, &DT))
+        addAssumeNonNull(AC, LI);
+
       // Anything using the load now uses the current value.
       LI->replaceAllUsesWith(V);
       if (AST && LI->getType()->isPointerTy())
         AST->deleteValue(LI);
       BB->getInstList().erase(LI);
     } else if (StoreInst *SI = dyn_cast<StoreInst>(I)) {
       // Delete this instruction and mark the name as the current holder of the
       // value
       AllocaInst *Dest = dyn_cast<AllocaInst>(SI->getPointerOperand());
       if (!Dest)
         continue;
 
       DenseMap<AllocaInst *, unsigned>::iterator ai = AllocaLookup.find(Dest);
       if (ai == AllocaLookup.end())
         continue;

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

+; RUN: opt < %s -mem2reg -S | FileCheck %s
+
+; This tests that mem2reg preserves the !nonnull metadata on loads
+; from allocas that get optimized out.
+
+; Check the case where the alloca in question has a single store.
+define float* @single_store(float** %arg) {
+; CHECK-LABEL: define float* @single_store
+; CHECK: %arg.load = load float*, float** %arg, align 8
+; CHECK: [[ASSUME:%(.*)]] = icmp ne float* %arg.load, null
+; CHECK: call void @llvm.assume(i1 {{.*}}[[ASSUME]])
+; CHECK: ret float* %arg.load
+entry:
+  %buf = alloca float*
+  %arg.load = load float*, float** %arg, align 8
+  store float* %arg.load, float** %buf, align 8
+  %buf.load = load float*, float **%buf, !nonnull !0
+  ret float* %buf.load
+}
+
+; Check the case where the alloca in question has more than one
+; store but still within one basic block.
+define float* @single_block(float** %arg) {
+; CHECK-LABEL: define float* @single_block
+; CHECK: %arg.load = load float*, float** %arg, align 8
+; CHECK: [[ASSUME:%(.*)]] = icmp ne float* %arg.load, null
+; CHECK: call void @llvm.assume(i1 {{.*}}[[ASSUME]])
+; CHECK: ret float* %arg.load
+entry:
+  %buf = alloca float*
+  %arg.load = load float*, float** %arg, align 8
+  store float* null, float** %buf, align 8
+  store float* %arg.load, float** %buf, align 8
+  %buf.load = load float*, float **%buf, !nonnull !0
+  ret float* %buf.load
+}
+
+; Check the case where the alloca in question has more than one
+; store and also reads ands writes in multiple blocks.
+define float* @multi_block(float** %arg) {
+; CHECK-LABEL: define float* @multi_block
+; CHECK-LABEL: entry:
+; CHECK: %arg.load = load float*, float** %arg, align 8
+; CHECK: br label %next
+; CHECK-LABEL: next:
+; CHECK: [[ASSUME:%(.*)]] = icmp ne float* %arg.load, null
+; CHECK: call void @llvm.assume(i1 {{.*}}[[ASSUME]])
+; CHECK: ret float* %arg.load
+entry:
+  %buf = alloca float*
+  %arg.load = load float*, float** %arg, align 8
+  store float* null, float** %buf, align 8
+  br label %next
+next:
+  store float* %arg.load, float** %buf, align 8
+  %buf.load = load float*, float** %buf, !nonnull !0
+  ret float* %buf.load
+}
+
+; Check that we don't add an assume if it's not
+; necessary i.e. the value is already implied to be nonnull
+define float* @no_assume(float** %arg) {
+; CHECK-LABEL: define float* @no_assume
+; CHECK-LABEL: entry:
+; CHECK: %arg.load = load float*, float** %arg, align 8
+; CHECK: %cn = icmp ne float* %arg.load, null
+; CHECK: br i1 %cn, label %next, label %fin
+; CHECK-LABEL: next:
+; CHECK-NOT: call void @llvm.assume
+; CHECK: ret float* %arg.load
+; CHECK-LABEL: fin:
+; CHECK: ret float* null
+entry:
+  %buf = alloca float*
+  %arg.load = load float*, float** %arg, align 8
+  %cn = icmp ne float* %arg.load, null
+  br i1 %cn, label %next, label %fin
+next:
+; At this point the above nonnull check ensures that
+; the value %arg.load is nonnull in this block and thus
+; we need not add the assume.
+  store float* %arg.load, float** %buf, align 8
+  %buf.load = load float*, float** %buf, !nonnull !0
+  ret float* %buf.load
+fin:
+  ret float* null
+}
+
+!0 = !{}

test/Transforms/SROA/preserve-nonnull.ll

+; RUN: opt < %s -sroa -S | FileCheck %s
+;
+; Make sure that SROA doesn't lose nonnull metadata
+; on loads from allocas that get optimized out.
+
+; CHECK-LABEL: define float* @yummy_nonnull
+; CHECK: [[RETURN:%(.*)]] = load float*, float** %arg, align 8
+; CHECK: [[ASSUME:%(.*)]] = icmp ne float* {{.*}}[[RETURN]], null
+; CHECK: call void @llvm.assume(i1 {{.*}}[[ASSUME]])
+; CHECK: ret float* {{.*}}[[RETURN]]
+
+define float* @yummy_nonnull(float** %arg) {
+entry-block:
+  %buf = alloca float*
+
+  %_arg_i8 = bitcast float** %arg to i8*
+  %_buf_i8 = bitcast float** %buf to i8*
+  call void @llvm.memcpy.p0i8.p0i8.i64(i8* %_buf_i8, i8* %_arg_i8, i64 8, i32 8, i1 false)
+
+  %ret = load float*, float** %buf, align 8, !nonnull !0
+  ret float* %ret
+}
+
+declare void @llvm.memcpy.p0i8.p0i8.i64(i8* nocapture writeonly, i8* nocapture readonly, i64, i32, i1)
+
+!0 = !{}