2 files changed, 71 insertions, 20 deletions

diff --git a/llvm/lib/Transforms/Scalar/LICM.cpp b/llvm/lib/Transforms/Scalar/LICM.cpp
index 0e7f534c32a..0ac53ba92cb 100644
--- a/llvm/lib/Transforms/Scalar/LICM.cpp
+++ b/llvm/lib/Transforms/Scalar/LICM.cpp
@@ -154,10 +154,7 @@ namespace {
     /// Simple Analysis hook. Delete loop L from alias set map.
     void deleteAnalysisLoop(Loop *L) override;
 
-    /// Returns an owning pointer to an alias set which incorporates aliasing
-    /// info from all subloops of L, but does not include instructions in L
-    /// itself.
-    AliasSetTracker *collectAliasInfoFromSubLoops(Loop *L);
+    AliasSetTracker *collectAliasInfoForLoop(Loop *L);
   };
 }
 
@@ -188,22 +185,13 @@ bool LICM::runOnLoop(Loop *L, LPPassManager &LPM) {
 
   assert(L->isLCSSAForm(*DT) && "Loop is not in LCSSA form.");
 
-  CurAST = collectAliasInfoFromSubLoops(L);
+  CurAST = collectAliasInfoForLoop(L);
 
   CurLoop = L;
 
   // Get the preheader block to move instructions into...
   Preheader = L->getLoopPreheader();
 
-  // Loop over the body of this loop, looking for calls, invokes, and stores.
-  // Because subloops have already been incorporated into AST, we skip blocks in
-  // subloops.
-  //
-  for (BasicBlock *BB : L->blocks()) {
-    if (LI->getLoopFor(BB) == L)        // Ignore blocks in subloops.
-      CurAST->add(*BB);                 // Incorporate the specified basic block
-  }
-
   // Compute loop safety information.
   LICMSafetyInfo SafetyInfo;
   computeLICMSafetyInfo(&SafetyInfo, CurLoop);
@@ -1039,13 +1027,20 @@ bool llvm::promoteLoopAccessesToScalars(AliasSet &AS,
 }
 
 /// Returns an owning pointer to an alias set which incorporates aliasing info
-/// from all subloops of L, but does not include instructions in L itself.
-///
-AliasSetTracker *LICM::collectAliasInfoFromSubLoops(Loop *L) {
+/// from L and all subloops of L.
+AliasSetTracker *LICM::collectAliasInfoForLoop(Loop *L) {
   AliasSetTracker *CurAST = nullptr;
+  SmallVector<Loop *, 4> RecomputeLoops;
   for (Loop *InnerL : L->getSubLoops()) {
-    AliasSetTracker *InnerAST = LoopToAliasSetMap[InnerL];
-    assert(InnerAST && "Where is my AST?");
+    auto MapI = LoopToAliasSetMap.find(InnerL);
+    // If the AST for this inner loop is missing it may have been merged into
+    // some other loop's AST and then that loop unrolled, and so we need to
+    // recompute it.
+    if (MapI == LoopToAliasSetMap.end()) {
+      RecomputeLoops.push_back(InnerL);
+      continue;
+    }
+    AliasSetTracker *InnerAST = MapI->second;
 
     if (CurAST != nullptr) {
       // What if InnerLoop was modified by other passes ?
@@ -1057,10 +1052,27 @@ AliasSetTracker *LICM::collectAliasInfoFromSubLoops(Loop *L) {
     } else {
       CurAST = InnerAST;
     }
-    LoopToAliasSetMap.erase(InnerL);
+    LoopToAliasSetMap.erase(MapI);
   }
   if (CurAST == nullptr)
     CurAST = new AliasSetTracker(*AA);
+
+  auto mergeLoop = [&](Loop *L) {
+    // Loop over the body of this loop, looking for calls, invokes, and stores.
+    // Because subloops have already been incorporated into AST, we skip blocks
+    // in subloops.
+    for (BasicBlock *BB : L->blocks())
+      if (LI->getLoopFor(BB) == L) // Ignore blocks in subloops.
+        CurAST->add(*BB);          // Incorporate the specified basic block
+  };
+
+  // Add everything from the sub loops that are no longer directly available.
+  for (Loop *InnerL : RecomputeLoops)
+    mergeLoop(InnerL);
+
+  // And merge in this loop.
+  mergeLoop(L);
+
   return CurAST;
 }
 
diff --git a/llvm/test/Transforms/LICM/alias-set-tracker-loss.ll b/llvm/test/Transforms/LICM/alias-set-tracker-loss.ll
new file mode 100644
index 00000000000..378d908f698
--- /dev/null
+++ b/llvm/test/Transforms/LICM/alias-set-tracker-loss.ll
@@ -0,0 +1,39 @@
+; RUN: opt -S -licm -loop-unroll < %s
+;
+; This test contains a carefully rotated set of three nested loops. The middle
+; loop can be unrolled leaving one copy of the inner loop inside the outer
+; loop. Because of how LICM works, when this middle loop is unrolled and
+; removed, its alias set tracker is destroyed and no longer available when LICM
+; runs on the outer loop.
+
+define void @f() {
+entry:
+  br label %l1
+
+l2.l1.loopexit_crit_edge:
+  br label %l1.loopexit
+
+l1.loopexit:
+  br label %l1.backedge
+
+l1:
+  br i1 undef, label %l1.backedge, label %l2.preheader
+
+l1.backedge:
+  br label %l1
+
+l2.preheader:
+  br i1 true, label %l1.loopexit, label %l3.preheader.lr.ph
+
+l3.preheader.lr.ph:
+  br label %l3.preheader
+
+l2.loopexit:
+  br i1 true, label %l2.l1.loopexit_crit_edge, label %l3.preheader
+
+l3.preheader:
+  br label %l3
+
+l3:
+  br i1 true, label %l3, label %l2.loopexit
+}