[SROA] propagate !range metadata when moving loads

[arielb1]

Summary

This tries to propagate !range metadata to a pre-existing load when a load is optimized out. This is done instead of adding an assume because converting loads to and from assumes creates a lot of IR.

This requires D37215 in order to work, but I've put it in a separate patch because it affects a different section of the code (and also for easier bisection if this introduces a regression).

[hfinkel]

lib/Transforms/Utils/Local.cpp

1893 ↗

(On Diff #112906)

return here.

lib/Transforms/Utils/PromoteMemoryToRegister.cpp

315 ↗

(On Diff #112906)

Huh?

318 ↗

(On Diff #112906)

Brace on previous line.

328 ↗

(On Diff #112906)

Two space indent.

329 ↗

(On Diff #112906)

We should refactor this (calling this function here seems odd - the load is not an assume). I think that you're trying to reuse the logic here for the dominance checking, which is fine, but you don't need the isEphemeralValueOf check). Please feel free to break isValidAssumeForContext into two parts, and reuse here the first part.

330 ↗

(On Diff #112906)

Odd indentation.

[arielb1]

address review comments

[hfinkel]

This looks good, but, have you checked for compile-time regressions? I'm concerned here because:

We've now made addAssumptionsFromMetadata linear in the number of instructions between the instruction and its replacement. That should be documented with the function declaration (at least). As the replacement always dominates the original instruction, when they're in the same BB, we'll always be doing that linear walk. As you're making the walk go through normal loads, I can imagine this becoming quadratic in some inputs (unless there's something about the way that SROA uses this that prevents that, and if so, we need to document that too).

If this can become quadratic, I think we'll want to maintain some kind of BB-region analysis, there we partition each BB into regions based on guaranteed execution, and then use that instead of walking.

include/llvm/Analysis/ValueTracking.h

375 ↗

(On Diff #113699)

reached and executed -> later executed

lib/Transforms/Utils/PromoteMemoryToRegister.cpp

314 ↗

(On Diff #113699)

Brace on previous line.

316 ↗

(On Diff #113699)

Do you need the llvm:: here? If not, please remove it.

[arielb1]

Don't use ValueTracking's isGuaranteedToBeExecuted to avoid potential O(n^2). This is much uglier but should be O(n log n) in all cases.

[arielb1]

style

[hfinkel]

lib/Transforms/Utils/PromoteMemoryToRegister.cpp

168 ↗

(On Diff #113761)

... guaranteed to execute once the first one executes.

209 ↗

(On Diff #113761)

How is that renumbering handled correctly?

Also, there's some assumption here that we don't insert new instructions that would break apart an existing internal? (I'm guessing not, as that seems like it would be a semantics-changing transformation, and we don't do that here, but we should document the assumptions here).

232 ↗

(On Diff #113761)

Line too long (and you should just use auto here).

372 ↗

(On Diff #113761)

Odd indentation.

420 ↗

(On Diff #113761)

Odd indentation.

475–476 ↗

(On Diff #113761)

Odd indentation.

578 ↗

(On Diff #113761)

Odd indentation.

998 ↗

(On Diff #113761)

Odd indentation.

[arielb1]

address review comments

[hfinkel]

You can benchmarks on the ValueTracking change? How did those turn out?

lib/Transforms/Utils/PromoteMemoryToRegister.cpp

172 ↗

(On Diff #113778)

Variable names start with a capital letter.

264 ↗

(On Diff #113778)

Start comments with a capital letter.

270 ↗

(On Diff #113778)

Local variable names start with a capital letter.

280 ↗

(On Diff #113778)

If we that -> If we find that

[arielb1]

address review comments

[hfinkel]

In D37216#861566, @hfinkel wrote:

You can benchmarks on the ValueTracking change? How did those turn out?

Great. This LGTM.

[arielb1]

This *still* hasn't committed. What is this waiting for?

[arielb1]

Ok. I was just going over my todo list for in-progress items.

[arielb1]

Has this still not landed in this month? Did this become blocked on anything?

[lebedev.ri]

In D37216#910089, @arielb1 wrote:

Has this still not landed in this month? Did this become blocked on anything?

Once the differential is accepted, you are supposed to commit it. There is no auto-commit.
If you don't have commit rights yet, you should request them.
Or you can ask others to commit it for you.

[nagisa]

Would be nice if somebody committed this. None of us have the commit bit.

[davide]

In D37216#935153, @nagisa wrote:

Would be nice if somebody committed this. None of us have the commit bit.

Please rebase the patch on ToT and I'll commit for you.

[arielb1]

Rebased and ready to land

[davide]

rL319096, cheers.

[davide]

Ariel, we had to revert this because it broke some bots. Feel free to fix the assertion and update the patch

[nox]

In D37216#937116, @davide wrote:

Ariel, we had to revert this because it broke some bots. Feel free to fix the assertion and update the patch

What was the failing assertion?

llvm/trunk/lib/Transforms/Scalar/SROA.cpp

if (VecTy) {
       if (LI.isVolatile())
         NewLI->setAtomic(LI.getOrdering(), LI.getSyncScopeID());
 
       // Any !nonnull metadata or !range metadata on the old load is also valid
       // on the new load. This is even true in some cases even when the loads
       // are different types, for example by mapping !nonnull metadata to
       // !range metadata by modeling the null pointer constant converted to the
       // integer type.
-      // FIXME: Add support for range metadata here. Currently the utilities
-      // for this don't propagate range metadata in trivial cases from one
-      // integer load to another, don't handle non-addrspace-0 null pointers
-      // correctly, and don't have any support for mapping ranges as the
-      // integer type becomes winder or narrower.
       if (MDNode *N = LI.getMetadata(LLVMContext::MD_nonnull))
         copyNonnullMetadata(LI, N, *NewLI);
-
-      // Try to preserve nonnull metadata
+      if (MDNode *N = LI.getMetadata(LLVMContext::MD_range))
+        copyRangeMetadata(DL, LI, N, *NewLI);
       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()) {
for (;;) {
       auto AS = LI->getPointerAddressSpace();
       auto *PartPtrTy = PartTy->getPointerTo(AS);
       LoadInst *PLoad = IRB.CreateAlignedLoad(
           getAdjustedPtr(IRB, DL, BasePtr,
                          APInt(DL.getPointerSizeInBits(AS), PartOffset),
                          PartPtrTy, BasePtr->getName() + "."),
           getAdjustedAlignment(LI, PartOffset, DL), /*IsVolatile*/ false,
           LI->getName());
-      PLoad->copyMetadata(*LI, LLVMContext::MD_mem_parallel_loop_access); 
+      PLoad->copyMetadata(*LI, LLVMContext::MD_mem_parallel_loop_access);
 
       // Append this load onto the list of split loads so we can find it later
       // to rewrite the stores.
       SplitLoads.push_back(PLoad);
 
       // Now build a new slice for the alloca.
       NewSlices.push_back(
           Slice(BaseOffset + PartOffset, BaseOffset + PartOffset + PartSize,

llvm/trunk/lib/Transforms/Utils/Local.cpp

void llvm::copyNonnullMetadata(const LoadInst &OldLI, MDNode *N,
   auto *NonNullInt = ConstantExpr::getAdd(NullInt, ConstantInt::get(ITy, 1));
   NewLI.setMetadata(LLVMContext::MD_range,
                     MDB.createRange(NonNullInt, NullInt));
 }
 
 void llvm::copyRangeMetadata(const DataLayout &DL, const LoadInst &OldLI,
                              MDNode *N, LoadInst &NewLI) {
   auto *NewTy = NewLI.getType();
+  auto *OldTy = OldLI.getType();
 
-  // Give up unless it is converted to a pointer where there is a single very
-  // valuable mapping we can do reliably.
-  // FIXME: It would be nice to propagate this in more ways, but the type
-  // conversions make it hard.
-  if (!NewTy->isPointerTy())
+  if (DL.getTypeStoreSizeInBits(NewTy) == DL.getTypeSizeInBits(OldTy) &&
+      NewTy->isIntegerTy()) {
+    // An integer with the same number of bits - give it the range
+    // metadata!.
+    NewLI.setMetadata(LLVMContext::MD_range, N);
     return;
+  }
 
+  if (NewTy->isPointerTy()) {
+    // Try to convert the !range metadata to !nonnull metadata on the
+    // new pointer.
     unsigned BitWidth = DL.getTypeSizeInBits(NewTy);
     if (!getConstantRangeFromMetadata(*N).contains(APInt(BitWidth, 0))) {
       MDNode *NN = MDNode::get(OldLI.getContext(), None);
       NewLI.setMetadata(LLVMContext::MD_nonnull, NN);
     }
   }
+}
 
 namespace {
 
 /// A potential constituent of a bitreverse or bswap expression. See
 /// collectBitParts for a fuller explanation.
 struct BitPart {
   BitPart(Value *P, unsigned BW) : Provider(P) {
     Provenance.resize(BW);

llvm/trunk/lib/Transforms/Utils/PromoteMemoryToRegister.cpp

 // 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/ADT/ArrayRef.h"
 #include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/Optional.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/Statistic.h"
 #include "llvm/ADT/TinyPtrVector.h"
 #include "llvm/ADT/Twine.h"
 #include "llvm/Analysis/AssumptionCache.h"
 #include "llvm/Analysis/InstructionSimplify.h"

 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/Module.h"
 #include "llvm/IR/Type.h"
 #include "llvm/IR/User.h"
 #include "llvm/Support/Casting.h"
 #include "llvm/Transforms/Utils/Local.h"
 #include "llvm/Transforms/Utils/PromoteMemToReg.h"
 #include <algorithm>
-#include <cassert>
+
 #include <iterator>
 #include <utility>
 #include <vector>
 
 using namespace llvm;
 
 #define DEBUG_TYPE "mem2reg"
 
public:
   RenamePassData(BasicBlock *B, BasicBlock *P, ValVector V)
       : BB(B), Pred(P), Values(std::move(V)) {}
 
   BasicBlock *BB;
   BasicBlock *Pred;
   ValVector Values;
 };
 
+/// \brief Semi-open interval of instructions that are guaranteed to
+/// all execute if the first one does.
+class GuaranteedExecutionRange {
+public:
+  unsigned Start;
+  unsigned End;
+
+  GuaranteedExecutionRange(unsigned S, unsigned E): Start(S), End(E) {}
+};
+
 /// \brief This assigns and keeps a per-bb relative ordering of load/store
 /// instructions in the block that directly load or store an alloca.
 ///
 /// This functionality is important because it avoids scanning large basic
 /// blocks multiple times when promoting many allocas in the same block.
 class LargeBlockInfo {
   /// \brief For each instruction that we track, keep the index of the
   /// instruction.
   ///
   /// The index starts out as the number of the instruction from the start of
   /// the block.
   DenseMap<const Instruction *, unsigned> InstNumbers;
 
+  /// \brief For each basic block we track, keep track of the intervals
+  /// of instruction numbers of instructions that transfer control
+  /// to their successors, for propagating metadata.
+  DenseMap<const BasicBlock *, Optional<SmallVector<GuaranteedExecutionRange, 4>>>
+    GuaranteedExecutionIntervals;
+
 public:
 
-  /// This code only looks at accesses to allocas.
+  /// This code looks for stores to allocas, and for loads both for
+  /// allocas and for transferring metadata.
   static bool isInterestingInstruction(const Instruction *I) {
-    return (isa<LoadInst>(I) && isa<AllocaInst>(I->getOperand(0))) ||
+    return isa<LoadInst>(I) ||
            (isa<StoreInst>(I) && isa<AllocaInst>(I->getOperand(1)));
   }
 
+  /// Compute the GuaranteedExecutionIntervals for a given BB.
+  ///
+  /// This is valid and remains valid as long as each interesting
+  /// instruction (see isInterestingInstruction) that
+  ///   A) existed when this LBI was cleared
+  ///   B) has not been deleted (deleting interesting instructions is fine)
+  /// are run in the same program executions and in the same order
+  /// as when this LBI was cleared.
+  ///
+  /// Because `PromoteMemoryToRegister` does not move memory loads at
+  /// all, this assumption is satisfied in this pass.
+  SmallVector<GuaranteedExecutionRange, 4> computeGEI(const BasicBlock *BB) {
+    SmallVector<GuaranteedExecutionRange, 4> GuaranteedExecutionIntervals;
+
+    unsigned InstNo = 0;
+    bool InRange = false;
+    unsigned FirstInstInRange = 0;
+    for (const Instruction &BBI : *BB) {
+      if (isGuaranteedToTransferExecutionToSuccessor(&BBI)) {
+        if (!InRange && isInterestingInstruction(&BBI)) {
+          InRange = true;
+          FirstInstInRange = InstNo;
+        }
+      } else {
+        if (InRange) {
+          assert(FirstInstInRange < InstNo && "Can't push an empty range here.");
+          GuaranteedExecutionIntervals.emplace_back(FirstInstInRange, InstNo);
+        }
+        InRange = false;
+      }
+
+      if (isInterestingInstruction(&BBI)) {
+        auto It = InstNumbers.find(&BBI);
+        assert(It != InstNumbers.end() &&
+               InstNo <= It->second &&
+               "missing number for interesting instruction");
+        InstNo = It->second + 1;
+      }
+    }
+
+    if (InRange) {
+      assert(FirstInstInRange < InstNo && "Can't push an empty range here.");
+      GuaranteedExecutionIntervals.emplace_back(FirstInstInRange, InstNo);
+    }
+
+    return GuaranteedExecutionIntervals;
+  }
+
+  /// Return true if, when CxtI executes, it is guaranteed that either
+  /// I had executed already or that I is guaranteed to be later executed.
+  ///
+  /// The useful property this guarantees is that if I exhibits undefined
+  /// behavior under some circumstances, then the whole program will exhibit
+  /// undefined behavior at CxtI.
+  bool isGuaranteedToBeExecuted(const Instruction *CxtI, const Instruction *I) {
+    const BasicBlock *BB = CxtI->getParent();
+
+    if (BB != I->getParent()) {
+      // Instructions in different basic blocks, so control flow
+      // can diverge between them (we could track this with
+      // postdoms, but we don't bother).
+      return false;
+    }
+
+    unsigned Index1 = getInstructionIndex(CxtI);
+    unsigned Index2 = getInstructionIndex(I);
+
+    auto& BBGEI = GuaranteedExecutionIntervals[BB];
+    if (!BBGEI.hasValue()) {
+      BBGEI.emplace(computeGEI(BB));
+    }
+
+    // We want to check whether I and CxtI are in the same range. To do that,
+    // we notice that CxtI can only be in the first range R where
+    // CxtI.end < R.end. If we find that range using binary search,
+    // we can check whether I and CxtI are both in it.
+    GuaranteedExecutionRange Bound(Index1, Index1);
+    auto R = std::upper_bound(
+      BBGEI->begin(), BBGEI->end(), Bound,
+      [](GuaranteedExecutionRange I_, GuaranteedExecutionRange R) {
+        return I_.End < R.End;
+      });
+
+    return R != BBGEI->end() &&
+      R->Start <= Index1 && Index1 < R->End &&
+      R->Start <= Index2 && Index2 < R->End;
+  }
+
   /// Get or calculate the index of the specified instruction.
   unsigned getInstructionIndex(const Instruction *I) {
     assert(isInterestingInstruction(I) &&
            "Not a load/store to/from an alloca?");
 
     // If we already have this instruction number, return it.
     DenseMap<const Instruction *, unsigned>::iterator It = InstNumbers.find(I);
     if (It != InstNumbers.end())
       return It->second;
 
     // Scan the whole block to get the instruction.  This accumulates
     // information for every interesting instruction in the block, in order to
     // avoid gratuitus rescans.
     const BasicBlock *BB = I->getParent();
     unsigned InstNo = 0;
+    GuaranteedExecutionIntervals.erase(BB);
     for (const Instruction &BBI : *BB)
       if (isInterestingInstruction(&BBI))
         InstNumbers[&BBI] = InstNo++;
+
     It = InstNumbers.find(I);
 
     assert(It != InstNumbers.end() && "Didn't insert instruction?");
     return It->second;
   }
 
   void deleteValue(const Instruction *I) { InstNumbers.erase(I); }
 
-  void clear() { InstNumbers.clear(); }
+  void clear() {
+    InstNumbers.clear();
+    GuaranteedExecutionIntervals.clear();
+  }
 };
 
 struct PromoteMem2Reg {
   /// The alloca instructions being promoted.
   std::vector<AllocaInst *> Allocas;
 
   DominatorTree &DT;
   DIBuilder DIB;
unsigned getNumPreds(const BasicBlock *BB) {
     return NP - 1;
   }
 
   void ComputeLiveInBlocks(AllocaInst *AI, AllocaInfo &Info,
                            const SmallPtrSetImpl<BasicBlock *> &DefBlocks,
                            SmallPtrSetImpl<BasicBlock *> &LiveInBlocks);
   void RenamePass(BasicBlock *BB, BasicBlock *Pred,
                   RenamePassData::ValVector &IncVals,
+                  LargeBlockInfo &LBI,
                   std::vector<RenamePassData> &Worklist);
   bool QueuePhiNode(BasicBlock *BB, unsigned AllocaIdx, unsigned &Version);
 };
 
 } // end 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 addAssumptionsFromMetadata(LoadInst *LI,
+                                       Value *ReplVal,
+                                       DominatorTree &DT,
+                                       const DataLayout &DL,
+              LargeBlockInfo &LBI,
+                                       AssumptionCache *AC)
+{
+  if (LI->getMetadata(LLVMContext::MD_nonnull) &&
+      !isKnownNonZero(ReplVal, DL, 0, AC, LI, &DT)) {
+    addAssumeNonNull(AC, LI);
+  }
+
+  if (auto *N = LI->getMetadata(LLVMContext::MD_range)) {
+    // Range metadata is harder to use as an assumption,
+    // so don't try to add one, but *do* try to copy
+    // the metadata to a load in the same BB.
+    if (LoadInst *NewLI = dyn_cast<LoadInst>(ReplVal)) {
+      DEBUG(dbgs() << "trying to move !range metadata from" <<
+        *LI << " to" << *NewLI << "\n");
+      if (LBI.isGuaranteedToBeExecuted(LI, NewLI)) {
+        copyRangeMetadata(DL, *LI, N, *NewLI);
+      }
+    }
+  }
+}
+
 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))
for (auto UI = AI->user_begin(), E = AI->user_end(); UI != E;) {
     // 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) &&
-        !isKnownNonZero(ReplVal, DL, 0, AC, LI, &DT))
-      addAssumeNonNull(AC, LI);
+    addAssumptionsFromMetadata(LI, ReplVal, DT, DL, LBI, AC);
 
     LI->replaceAllUsesWith(ReplVal);
     LI->eraseFromParent();
     LBI.deleteValue(LI);
   }
 
   // Finally, after the scan, check to see if the store is all that is left.
   if (!Info.UsingBlocks.empty())
if (I == StoresByIndex.begin()) {
         // 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);
+      addAssumptionsFromMetadata(LI, ReplVal, DT, DL, LBI, AC);
 
       LI->replaceAllUsesWith(ReplVal);
     }
 
     LI->eraseFromParent();
     LBI.deleteValue(LI);
   }
 
for (unsigned AllocaNum = 0; AllocaNum != Allocas.size(); ++AllocaNum) {
 
     unsigned CurrentVersion = 0;
     for (BasicBlock *BB : PHIBlocks)
       QueuePhiNode(BB, AllocaNum, CurrentVersion);
   }
 
   if (Allocas.empty())
     return; // All of the allocas must have been trivial!
-
   LBI.clear();
 
   // Set the incoming values for the basic block to be null values for all of
   // the alloca's.  We do this in case there is a load of a value that has not
   // been stored yet.  In this case, it will get this null value.
   RenamePassData::ValVector Values(Allocas.size());
   for (unsigned i = 0, e = Allocas.size(); i != e; ++i)
     Values[i] = UndefValue::get(Allocas[i]->getAllocatedType());
 
   // Walks all basic blocks in the function performing the SSA rename algorithm
   // and inserting the phi nodes we marked as necessary
   std::vector<RenamePassData> RenamePassWorkList;
   RenamePassWorkList.emplace_back(&F.front(), nullptr, std::move(Values));
   do {
     RenamePassData RPD = std::move(RenamePassWorkList.back());
     RenamePassWorkList.pop_back();
     // RenamePass may add new worklist entries.
-    RenamePass(RPD.BB, RPD.Pred, RPD.Values, RenamePassWorkList);
+    RenamePass(RPD.BB, RPD.Pred, RPD.Values, LBI, RenamePassWorkList);
   } while (!RenamePassWorkList.empty());
 
+  LBI.clear();
   // The renamer uses the Visited set to avoid infinite loops.  Clear it now.
   Visited.clear();
 
   // Remove the allocas themselves from the function.
   for (Instruction *A : Allocas) {
     // If there are any uses of the alloca instructions left, they must be in
     // unreachable basic blocks that were not processed by walking the dominator
     // tree. Just delete the users now.

 
 /// \brief Recursively traverse the CFG of the function, renaming loads and
 /// stores to the allocas which we are promoting.
 ///
 /// IncomingVals indicates what value each Alloca contains on exit from the
 /// predecessor block Pred.
 void PromoteMem2Reg::RenamePass(BasicBlock *BB, BasicBlock *Pred,
                                 RenamePassData::ValVector &IncomingVals,
+       LargeBlockInfo &LBI,
                                 std::vector<RenamePassData> &Worklist) {
 NextIteration:
   // If we are inserting any phi nodes into this BB, they will already be in the
   // block.
   if (PHINode *APN = dyn_cast<PHINode>(BB->begin())) {
     // If we have PHI nodes to update, compute the number of edges from Pred to
     // BB.
     if (PhiToAllocaMap.count(APN)) {
if (LoadInst *LI = dyn_cast<LoadInst>(I)) {
       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) &&
-          !isKnownNonZero(V, SQ.DL, 0, AC, LI, &DT))
-        addAssumeNonNull(AC, LI);
+      addAssumptionsFromMetadata(LI, V, DT, SQ.DL, LBI, AC);
 
       // Anything using the load now uses the current value.
       LI->replaceAllUsesWith(V);
       BB->getInstList().erase(LI);
+      LBI.deleteValue(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;
 
       // what value were we writing?
       IncomingVals[ai->second] = SI->getOperand(0);
       // Record debuginfo for the store before removing it.
       for (DbgInfoIntrinsic *DII : AllocaDbgDeclares[ai->second])
         ConvertDebugDeclareToDebugValue(DII, SI, DIB);
       BB->getInstList().erase(SI);
+      LBI.deleteValue(SI);
     }
   }
 
   // 'Recurse' to our successors.
   succ_iterator I = succ_begin(BB), E = succ_end(BB);
   if (I == E)
     return;
 

llvm/trunk/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.
 
+%pair = type { i64, [0 x i64], [1 x i64] }
+
 declare void @llvm.memcpy.p0i8.p0i8.i64(i8* nocapture writeonly, i8* nocapture readonly, i64, i32, i1)
 
 ; Check that we do basic propagation of nonnull when rewriting.
 define i8* @propagate_nonnull(i32* %v) {
 ; CHECK-LABEL: define i8* @propagate_nonnull(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    %[[A:.*]] = alloca i8*
 ; CHECK-NEXT:    %[[V_CAST:.*]] = bitcast i32* %v to i8*
entry:
   %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
 }
 
+; Make sure we propagate the !range attribute when we expand loads.
+define i64 @propagate_range(%pair* dereferenceable(16)) {
+; CHECK-LABEL: define i64 @propagate_range(
+; CHECK-NEXT:  start:
+; CHECK-NEXT:      %[[SROA_IDX:.*]] = getelementptr inbounds %pair
+; CHECK-NEXT:      %[[RESULT:.*]] = load i64, i64* %[[SROA_IDX]], align 8, !range !1
+; CHECK:           ret i64 %[[RESULT]]
+start:
+  %a = alloca %pair
+  %1 = bitcast %pair* %0 to i8*
+  %2 = bitcast %pair* %a to i8*
+  call void @llvm.memcpy.p0i8.p0i8.i64(i8* %2, i8* %1, i64 16, i32 8, i1 false)
+  %3 = getelementptr inbounds %pair, %pair* %a, i32 0, i32 0
+  %4 = load i64, i64* %3, !range !1
+  ret i64 %4
+}
+
 ; Make sure we properly handle the !nonnull attribute when we convert
 ; a pointer load to an integer load.
 ; FIXME: While this doesn't do anythnig actively harmful today, it really
 ; should propagate the !nonnull metadata to range metadata. The irony is, it
 ; *does* initially, but then we lose that !range metadata before we finish
 ; SROA.
 define i8* @propagate_nonnull_to_int() {
 ; CHECK-LABEL: define i8* @propagate_nonnull_to_int(
entry:
   %a.gep1.cast = bitcast i8** %a.gep1 to i64*
   store i64 42, i64* %a.gep1.cast
   store i64 0, i64* %a.gep0.cast
   %load = load i8*, i8** %a.gep1, align 8, !nonnull !0
   ret i8* %load
 }
 
 !0 = !{}
+!1 = !{i64 0, i64 2}