[LICM] Hoist stores of invariant values to invariant addresses out of loops

Teach LICM to hoist stores out of loops when the store writes to a location otherwise unused in the loop, writes a value which is invariant, and is guaranteed to execute if the loop is entered.

Worth noting is that this transformation is partially overlapping with the existing promotion transformation. Reasons this is worthwhile anyway include:
 * For multi-exit loops, this doesn't require duplication of the store.
 * It kicks in for case where we can't prove we exit through a normal exit (i.e. we may throw), but can prove the store executes before that possible side exit.

Differential Revision: https://reviews.llvm.org/D50925

llvm-svn: 340974

6 files changed, 135 insertions, 16 deletions

diff --git a/llvm/lib/Analysis/AliasSetTracker.cpp b/llvm/lib/Analysis/AliasSetTracker.cpp
index 377db8a569f..51a61e9e2d9 100644
--- a/llvm/lib/Analysis/AliasSetTracker.cpp
+++ b/llvm/lib/Analysis/AliasSetTracker.cpp
@@ -256,9 +256,22 @@ Instruction* AliasSet::getUniqueInstruction() {
   if (AliasAny)
     // May have collapses alias set
     return nullptr;
-  if (size() != 0)
-    // Can't track source of pointer, might be many instruction
-    return nullptr;
+  if (begin() != end()) {
+    if (!UnknownInsts.empty())
+      // Another instruction found
+      return nullptr;
+    if (std::next(begin()) != end())
+      // Another instruction found
+      return nullptr;
+    Value *Addr = begin()->getValue();
+    assert(!Addr->user_empty() &&
+           "where's the instruction which added this pointer?");
+    if (std::next(Addr->user_begin()) != Addr->user_end())
+      // Another instruction found -- this is really restrictive
+      // TODO: generalize!
+      return nullptr;
+    return cast<Instruction>(*(Addr->user_begin()));
+  }
   if (1 != UnknownInsts.size())
     return nullptr;
   return cast<Instruction>(UnknownInsts[0]);
diff --git a/llvm/lib/Transforms/Scalar/LICM.cpp b/llvm/lib/Transforms/Scalar/LICM.cpp
index 94fcdaf3c08..ff7cb3c26c9 100644
--- a/llvm/lib/Transforms/Scalar/LICM.cpp
+++ b/llvm/lib/Transforms/Scalar/LICM.cpp
@@ -416,7 +416,8 @@ bool llvm::sinkRegion(DomTreeNode *N, AliasAnalysis *AA, LoopInfo *LI,
       //
       bool FreeInLoop = false;
       if (isNotUsedOrFreeInLoop(I, CurLoop, SafetyInfo, TTI, FreeInLoop) &&
-          canSinkOrHoistInst(I, AA, DT, CurLoop, CurAST, true, ORE)) {
+          canSinkOrHoistInst(I, AA, DT, CurLoop, CurAST, true, ORE) &&
+          !I.mayHaveSideEffects()) {
         if (sink(I, LI, DT, CurLoop, SafetyInfo, ORE, FreeInLoop)) {
           if (!FreeInLoop) {
             ++II;
@@ -604,8 +605,8 @@ namespace {
 /// concerns such as aliasing or speculation safety.  
 bool isHoistableAndSinkableInst(Instruction &I) {
   // Only these instructions are hoistable/sinkable.
-  return (isa<LoadInst>(I) || isa<CallInst>(I) ||
-          isa<FenceInst>(I) ||
+  return (isa<LoadInst>(I) || isa<StoreInst>(I) ||
+          isa<CallInst>(I) || isa<FenceInst>(I) || 
           isa<BinaryOperator>(I) || isa<CastInst>(I) ||
           isa<SelectInst>(I) || isa<GetElementPtrInst>(I) ||
           isa<CmpInst>(I) || isa<InsertElementInst>(I) ||
@@ -695,6 +696,7 @@ bool llvm::canSinkOrHoistInst(Instruction &I, AAResults *AA, DominatorTree *DT,
       // it's arguments with arbitrary offsets.  If we can prove there are no
       // writes to this memory in the loop, we can hoist or sink.
       if (AliasAnalysis::onlyAccessesArgPointees(Behavior)) {
+        // TODO: expand to writeable arguments
         for (Value *Op : CI->arg_operands())
           if (Op->getType()->isPointerTy() &&
               pointerInvalidatedByLoop(
@@ -729,6 +731,24 @@ bool llvm::canSinkOrHoistInst(Instruction &I, AAResults *AA, DominatorTree *DT,
     (void)FI; //suppress unused variable warning
     assert(UniqueI == FI && "AS must contain FI");
     return true;
+  } else if (auto *SI = dyn_cast<StoreInst>(&I)) {
+    if (!SI->isUnordered())
+      return false; // Don't sink/hoist volatile or ordered atomic store!
+
+    // We can only hoist a store that we can prove writes a value which is not
+    // read or overwritten within the loop.  For those cases, we fallback to
+    // load store promotion instead.
+    auto &AS = CurAST->getAliasSetFor(MemoryLocation::get(SI));
+
+    if (AS.isRef() || !AS.isMustAlias())
+      // Quick exit test, handled by the full path below as well.
+      return false;
+    auto *UniqueI = AS.getUniqueInstruction();
+    if (!UniqueI)
+      // other memory op, give up
+      return false;
+    assert(UniqueI == SI && "AS must contain SI");
+    return true;
   }
 
   assert(!I.mayReadOrWriteMemory() && "unhandled aliasing");
diff --git a/llvm/test/Transforms/LICM/atomics.ll b/llvm/test/Transforms/LICM/atomics.ll
index 15c461aeca2..23dc60c590a 100644
--- a/llvm/test/Transforms/LICM/atomics.ll
+++ b/llvm/test/Transforms/LICM/atomics.ll
@@ -173,10 +173,11 @@ loop:
 end:
   ret i32 %vala
 ; CHECK-LABEL: define i32 @test7b(
+; CHECK-LABEL: entry:
+; CHECK: store i32 5, i32* %x
+; CHECK-LABEL: loop:
 ; CHECK: load atomic i32, i32* %y monotonic
-
 ; CHECK-LABEL: end:
-; CHECK: store i32 5, i32* %x
 ; CHECK: store atomic i32 %{{.+}}, i32* %z unordered, align 4
 }
 
diff --git a/llvm/test/Transforms/LICM/funclet.ll b/llvm/test/Transforms/LICM/funclet.ll
index 6b5f11507ed..36d2d49d9d4 100644
--- a/llvm/test/Transforms/LICM/funclet.ll
+++ b/llvm/test/Transforms/LICM/funclet.ll
@@ -97,9 +97,11 @@ else:                                             ; preds = %postinvoke
 }
 
 ; CHECK-LABEL: define void @test3(
-; CHECK:      catchswitch within none
+; CHECK-LABEL: forbody.preheader:
 ; CHECK:      store i32 1, i32* %bc, align 4
 ; CHECK:      store i32 2, i32* %bc2, align 4
+; CHECK:      catchswitch within none
+; CHECK-LABEL: forbody:
 
 declare void @may_throw()
 
diff --git a/llvm/test/Transforms/LICM/promote-order.ll b/llvm/test/Transforms/LICM/promote-order.ll
index b7e0b7c6c43..79aa7899b21 100644
--- a/llvm/test/Transforms/LICM/promote-order.ll
+++ b/llvm/test/Transforms/LICM/promote-order.ll
@@ -10,6 +10,11 @@ target triple = "x86_64-apple-macosx10.8.0"
 @p = external global i8*
 
 define i32* @_Z4doiti(i32 %n, float* %tmp1, i32* %tmp3) nounwind {
+; CHECK-LABEL: for.body.lr.ph:
+; CHECK: store float 1.000000e+00, float* %tmp1
+; CHECK-LABEL: for.cond.for.end_crit_edge:
+; CHECK: store i32 1, i32* %tmp3
+
 entry:
   %cmp1 = icmp slt i32 0, %n
   br i1 %cmp1, label %for.body.lr.ph, label %for.end
@@ -25,9 +30,6 @@ for.body:                                         ; preds = %for.body, %for.body
   %cmp = icmp slt i32 %inc, %n
   br i1 %cmp, label %for.body, label %for.cond.for.end_crit_edge
 
-; CHECK: for.cond.for.end_crit_edge:
-; CHECK: store float 1.000000e+00, float* %tmp1
-; CHECK: store i32 1, i32* %tmp3
 for.cond.for.end_crit_edge:                       ; preds = %for.body
   %split = phi i32* [ %tmp3, %for.body ]
   br label %for.end
diff --git a/llvm/test/Transforms/LICM/store-hoisting.ll b/llvm/test/Transforms/LICM/store-hoisting.ll
index 651fd9b2a71..11922fa77ee 100644
--- a/llvm/test/Transforms/LICM/store-hoisting.ll
+++ b/llvm/test/Transforms/LICM/store-hoisting.ll
@@ -3,8 +3,9 @@
 
 define void @test(i32* %loc) {
 ; CHECK-LABEL: @test
-; CHECK-LABEL: exit:
+; CHECK-LABEL: entry:
 ; CHECK: store i32 0, i32* %loc
+; CHECK-LABEL: loop:
 entry:
   br label %loop
 
@@ -21,10 +22,9 @@ exit:
 
 define void @test_multiexit(i32* %loc, i1 %earlycnd) {
 ; CHECK-LABEL: @test_multiexit
-; CHECK-LABEL: exit1:
-; CHECK: store i32 0, i32* %loc
-; CHECK-LABEL: exit2:
+; CHECK-LABEL: entry:
 ; CHECK: store i32 0, i32* %loc
+; CHECK-LABEL: loop:
 entry:
   br label %loop
 
@@ -44,6 +44,24 @@ exit2:
   ret void
 }
 
+define i32* @false_negative_2use(i32* %loc) {
+; CHECK-LABEL: @false_negative_2use
+; CHECK-LABEL: exit:
+; CHECK: store i32 0, i32* %loc
+entry:
+  br label %loop
+
+loop:
+  %iv = phi i32 [0, %entry], [%iv.next, %loop]
+  store i32 0, i32* %loc
+  %iv.next = add i32 %iv, 1
+  %cmp = icmp slt i32 %iv, 200
+  br i1 %cmp, label %loop, label %exit
+
+exit:
+  ret i32* %loc
+}
+
 define void @neg_lv_value(i32* %loc) {
 ; CHECK-LABEL: @neg_lv_value
 ; CHECK-LABEL: exit:
@@ -227,4 +245,67 @@ exit:
   ret void
 }
 
+declare void @maythrow() inaccessiblememonly
+
+define void @neg_early_exit(i32* %loc) {
+; CHECK-LABEL: @neg_early_exit
+; CHECK-LABEL: body:
+; CHECK: store i32 0, i32* %loc
+; CHECK-LABEL: exit:
+entry:
+  br label %loop
+
+loop:
+  %iv = phi i32 [0, %entry], [%iv.next, %body]
+  %is_null = icmp eq i32* %loc, null
+  br i1 %is_null, label %exit, label %body
+body:
+  call void @maythrow()
+  store i32 0, i32* %loc
+  %iv.next = add i32 %iv, 1
+  %cmp = icmp slt i32 %iv, 200
+  br i1 %cmp, label %loop, label %exit
+
+exit:
+  ret void
+}
+
+define void @neg_early_throw(i32* %loc) {
+; CHECK-LABEL: @neg_early_throw
+; CHECK-LABEL: loop:
+; CHECK: store i32 0, i32* %loc
+; CHECK-LABEL: exit:
+entry:
+  br label %loop
+
+loop:
+  %iv = phi i32 [0, %entry], [%iv.next, %loop]
+  call void @maythrow()
+  store i32 0, i32* %loc
+  %iv.next = add i32 %iv, 1
+  %cmp = icmp slt i32 %iv, 200
+  br i1 %cmp, label %loop, label %exit
+
+exit:
+  ret void
+}
+
+define void @test_late_throw(i32* %loc) {
+; CHECK-LABEL: @test_late_throw
+; CHECK-LABEL: entry:
+; CHECK: store i32 0, i32* %loc
+; CHECK-LABEL: loop:
+entry:
+  br label %loop
 
+loop:
+  %iv = phi i32 [0, %entry], [%iv.next, %loop]
+  store i32 0, i32* %loc
+  call void @maythrow()
+  %iv.next = add i32 %iv, 1
+  %cmp = icmp slt i32 %iv, 200
+  br i1 %cmp, label %loop, label %exit
+
+exit:
+  ret void
+}