[MemorySSA] Pass DT to the upward iterator for proper PhiTranslation.

[asbirlea]

Summary

A valid DominatorTree is needed to do PhiTranslation.
Before this patch, a MemoryUse could be optimized to an access outside a loop, while the address it loads from is modified in the loop.
This can lead to a miscompile.

[george.burgess.iv]

oof. thanks for this!

lgtm once a few nits are addressed.

llvm/include/llvm/Analysis/MemorySSA.h

726

nit: unless style around auto has changed, please spell out const DominatorTree &

(dunno offhand if const is possible here necessary, but ...)

secondary nit: since MSSA is a pretty big class with a fair few users, can we call this getDomTree() instead?

llvm/include/llvm/Analysis/MemorySSA.h

public:
   MemoryUseOrDef *getMemoryAccess(const Instruction *I) const {
     return cast_or_null<MemoryUseOrDef>(ValueToMemoryAccess.lookup(I));
   }
 
   MemoryPhi *getMemoryAccess(const BasicBlock *BB) const {
     return cast_or_null<MemoryPhi>(ValueToMemoryAccess.lookup(cast<Value>(BB)));
   }
 
+  DominatorTree &getDomTree() const { return *DT; }
+
   void dump() const;
   void print(raw_ostream &) const;
 
   /// Return true if \p MA represents the live on entry value
   ///
   /// Loads and stores from pointer arguments and other global values may be
   /// defined by memory operations that do not occur in the current function, so
   /// they may be live on entry to the function. MemorySSA represents such

 /// memory location through phi nodes for the user.
 class upward_defs_iterator
     : public iterator_facade_base<upward_defs_iterator,
                                   std::forward_iterator_tag,
                                   const MemoryAccessPair> {
   using BaseT = upward_defs_iterator::iterator_facade_base;
 
 public:
-  upward_defs_iterator(const MemoryAccessPair &Info)
+  upward_defs_iterator(const MemoryAccessPair &Info, DominatorTree *DT)
       : DefIterator(Info.first), Location(Info.second),
-        OriginalAccess(Info.first) {
+        OriginalAccess(Info.first), DT(DT) {
     CurrentPair.first = nullptr;
 
     WalkingPhi = Info.first && isa<MemoryPhi>(Info.first);
     fillInCurrentPair();
   }
 
   upward_defs_iterator() { CurrentPair.first = nullptr; }
 

 private:
   void fillInCurrentPair() {
     CurrentPair.first = *DefIterator;
     if (WalkingPhi && Location.Ptr) {
       PHITransAddr Translator(
           const_cast<Value *>(Location.Ptr),
           OriginalAccess->getBlock()->getModule()->getDataLayout(), nullptr);
       if (!Translator.PHITranslateValue(OriginalAccess->getBlock(),
-                                        DefIterator.getPhiArgBlock(), nullptr,
-                                        false))
+                                        DefIterator.getPhiArgBlock(), DT,
+                                        false)) {
         if (Translator.getAddr() != Location.Ptr) {
           CurrentPair.second = Location.getWithNewPtr(Translator.getAddr());
           return;
         }
       }
+    }
     CurrentPair.second = Location;
   }
 
   MemoryAccessPair CurrentPair;
   memoryaccess_def_iterator DefIterator;
   MemoryLocation Location;
   MemoryAccess *OriginalAccess = nullptr;
   bool WalkingPhi = false;
+  DominatorTree *DT = nullptr;
 };
 
-inline upward_defs_iterator upward_defs_begin(const MemoryAccessPair &Pair) {
-  return upward_defs_iterator(Pair);
+inline upward_defs_iterator upward_defs_begin(const MemoryAccessPair &Pair,
+                                              DominatorTree &DT) {
+  return upward_defs_iterator(Pair, &DT);
 }
 
 inline upward_defs_iterator upward_defs_end() { return upward_defs_iterator(); }
 
 inline iterator_range<upward_defs_iterator>
-upward_defs(const MemoryAccessPair &Pair) {
-  return make_range(upward_defs_begin(Pair), upward_defs_end());
+upward_defs(const MemoryAccessPair &Pair, DominatorTree &DT) {
+  return make_range(upward_defs_begin(Pair, DT), upward_defs_end());
 }
 
 /// Walks the defining accesses of MemoryDefs. Stops after we hit something that
 /// has no defining use (e.g. a MemoryPhi or liveOnEntry). Note that, when
 /// comparing against a null def_chain_iterator, this will compare equal only
 /// after walking said Phi/liveOnEntry.
 ///
 /// The UseOptimizedChain flag specifies whether to walk the clobbering

llvm/lib/Analysis/MemorySSA.cpp

for (const auto *MA : def_chain(MAP.first)) {
         (void)MU;
         assert (MU == Start &&
                 "Can only find use in def chain if Start is a use");
         continue;
       }
 
       assert(isa<MemoryPhi>(MA));
       Worklist.append(
-          upward_defs_begin({const_cast<MemoryAccess *>(MA), MAP.second}),
+          upward_defs_begin({const_cast<MemoryAccess *>(MA), MAP.second},
+                            MSSA.getDomTree()),
           upward_defs_end());
     }
   }
 
   // If the verify is done following an optimization, it's possible that
   // ClobberAt was a conservative clobbering, that we can now infer is not a
   // true clobbering access. Don't fail the verify if that's the case.
   // We do have accesses that claim they're optimized, but could be optimized
walkToPhiOrClobber(DefPath &Desc, const MemoryAccess *StopAt = nullptr,
 
     assert(isa<MemoryPhi>(Desc.Last) &&
            "Ended at a non-clobber that's not a phi?");
     return {Desc.Last, false, MayAlias};
   }
 
   void addSearches(MemoryPhi *Phi, SmallVectorImpl<ListIndex> &PausedSearches,
                    ListIndex PriorNode) {
-    auto UpwardDefs = make_range(upward_defs_begin({Phi, Paths[PriorNode].Loc}),
-                                 upward_defs_end());
+    auto UpwardDefs = make_range(
+        upward_defs_begin({Phi, Paths[PriorNode].Loc}, DT), upward_defs_end());
     for (const MemoryAccessPair &P : UpwardDefs) {
       PausedSearches.push_back(Paths.size());
       Paths.emplace_back(P.second, P.first, PriorNode);
     }
   }
 
   /// Represents a search that terminated after finding a clobber. This clobber
   /// may or may not be present in the path of defs from LastNode..SearchStart,

llvm/test/Analysis/MemorySSA/phi-translation.ll

 ; NOLIMIT: MemoryUse(4) MayAlias
 ; NOLIMIT-NEXT: load i8, i8* %p1
 ; LIMIT: MemoryUse(3) MayAlias
 ; LIMIT-NEXT: load i8, i8* %p1
   load i8, i8* %p1
   br label %while.cond
 }
 
+; CHECK-LABEL: define i32 @use_not_optimized_due_to_backedge
+define i32 @use_not_optimized_due_to_backedge(i32* nocapture %m_i_strides, i32* nocapture readonly %eval_left_dims) {
+entry:
+; CHECK: 1 = MemoryDef(liveOnEntry)
+; CHECK_NEXT: store i32 1, i32* %m_i_strides, align 4
+  store i32 1, i32* %m_i_strides, align 4
+  br label %for.body
+
+for.cond.cleanup:                                 ; preds = %for.inc
+  ret i32 %m_i_size.1
+
+for.body:                                         ; preds = %entry, %for.inc
+; CHECK: 4 = MemoryPhi({entry,1},{for.inc,3})
+; CHECK-NEXT: %indvars.iv = phi i64 [ 0, %entry ], [ %indvars.iv.next, %for.inc ]
+  %indvars.iv = phi i64 [ 0, %entry ], [ %indvars.iv.next, %for.inc ]
+  %m_i_size.022 = phi i32 [ 1, %entry ], [ %m_i_size.1, %for.inc ]
+  %indvars.iv.next = add nuw nsw i64 %indvars.iv, 1
+  %cmp1 = icmp eq i64 %indvars.iv, 0
+  %arrayidx2 = getelementptr inbounds i32, i32* %m_i_strides, i64 %indvars.iv
+; CHECK: MemoryUse(4) MayAlias
+; CHECK-NEXT: %0 = load i32, i32* %arrayidx2, align 4
+  %0 = load i32, i32* %arrayidx2, align 4
+  %arrayidx4 = getelementptr inbounds i32, i32* %eval_left_dims, i64 %indvars.iv
+; CHECK: MemoryUse(4) MayAlias
+; CHECK-NEXT: %1 = load i32, i32* %arrayidx4, align 4
+  %1 = load i32, i32* %arrayidx4, align 4
+  %mul = mul nsw i32 %1, %0
+  br i1 %cmp1, label %if.then, label %for.inc
+
+if.then:                                          ; preds = %for.body
+  %arrayidx7 = getelementptr inbounds i32, i32* %m_i_strides, i64 %indvars.iv.next
+; CHECK: 2 = MemoryDef(4)
+; CHECK-NEXT: store i32 %mul, i32* %arrayidx7, align 4
+  store i32 %mul, i32* %arrayidx7, align 4
+  br label %for.inc
+
+for.inc:                                          ; preds = %for.body, %if.then
+; CHECK: 3 = MemoryPhi({for.body,4},{if.then,2})
+; CHECK-NEXT: %m_i_size.1 = phi i32 [ %m_i_size.022, %if.then ], [ %mul, %for.body ]
+  %m_i_size.1 = phi i32 [ %m_i_size.022, %if.then ], [ %mul, %for.body ]
+  br i1 %cmp1, label %for.body, label %for.cond.cleanup
+}