[LICM] Make capture check more precise

[nikic]

Summary

During store promotion, we check whether the pointer was captured to exclude potential reads from other threads. However, we're only interested in captures before or inside the loop. Check this using PointerMayBeCapturedBefore against the loop header.

I'm not seeing any compile-time impact on CTMark for this change.

[reames]

LGTM w/one suggested comment.

llvm/lib/Transforms/Scalar/LICM.cpp

1957

I'd suggest inlining this helper into it's sole caller and simply updating the existing block comment. (Optional.)

[nikic]

llvm/lib/Transforms/Scalar/LICM.cpp

1957

There's two callers though :) One in isKnownNonEscaping, and the other in promoteLoopAccessesToScalars.

[reames]

llvm/lib/Transforms/Scalar/LICM.cpp

1957

Oops, yep.

Though why I wrote the second caller that way as opposed to calling the same helper is not immediately obvious. Cleanup opportunity?

p.s. LGTM still stands.

llvm/lib/Transforms/Scalar/LICM.cpp

void instructionDeleted(Instruction *I) const override {
     SafetyInfo.removeInstruction(I);
     if (AST)
       AST->deleteValue(I);
     if (MSSAU)
       MSSAU->removeMemoryAccess(I);
   }
 };
 
+bool isNotCapturedBeforeOrInLoop(const Value *V, const Loop *L,
+                                 DominatorTree *DT) {
+  // We can perform the captured-before check against any instruction in the
+  // loop header, as the loop header is reachable from any instruction inside
+  // the loop.
+  // TODO: ReturnCaptures=true shouldn't be necessary here.
+  return !PointerMayBeCapturedBefore(V, /* ReturnCaptures */ true,
+                                     /* StoreCaptures */ true,
+                                     L->getHeader()->getTerminator(), DT);
+}
 
 /// Return true iff we can prove that a caller of this function can not inspect
 /// the contents of the provided object in a well defined program.
-bool isKnownNonEscaping(Value *Object, const TargetLibraryInfo *TLI) {
+bool isKnownNonEscaping(Value *Object, const Loop *L,
+                        const TargetLibraryInfo *TLI, DominatorTree *DT) {
   if (isa<AllocaInst>(Object))
     // Since the alloca goes out of scope, we know the caller can't retain a
     // reference to it and be well defined.  Thus, we don't need to check for
     // capture.
     return true;
 
   // For all other objects we need to know that the caller can't possibly
   // have gotten a reference to the object.  There are two components of
   // that:
   //   1) Object can't be escaped by this function.  This is what
   //      PointerMayBeCaptured checks.
   //   2) Object can't have been captured at definition site.  For this, we
   //      need to know the return value is noalias.  At the moment, we use a
   //      weaker condition and handle only AllocLikeFunctions (which are
   //      known to be noalias).  TODO
   return isAllocLikeFn(Object, TLI) &&
-    !PointerMayBeCaptured(Object, true, true);
+         isNotCapturedBeforeOrInLoop(Object, L, DT);
 }
 
 } // namespace
 
 /// Try to promote memory values to scalars by sinking stores out of the
 /// loop and moving loads to before the loop.  We do this by looping over
 /// the stores in the loop, looking for stores to Must pointers which are
 /// loop invariant.
bool llvm::promoteLoopAccessesToScalars(
   bool IsKnownThreadLocalObject = false;
   if (SafetyInfo->anyBlockMayThrow()) {
     // If a loop can throw, we have to insert a store along each unwind edge.
     // That said, we can't actually make the unwind edge explicit. Therefore,
     // we have to prove that the store is dead along the unwind edge.  We do
     // this by proving that the caller can't have a reference to the object
     // after return and thus can't possibly load from the object.
     Value *Object = getUnderlyingObject(SomePtr);
-    if (!isKnownNonEscaping(Object, TLI))
+    if (!isKnownNonEscaping(Object, CurLoop, TLI, DT))
       return false;
     // Subtlety: Alloca's aren't visible to callers, but *are* potentially
     // visible to other threads if captured and used during their lifetimes.
     IsKnownThreadLocalObject = !isa<AllocaInst>(Object);
   }
 
   // Check that all of the pointers in the alias set have the same type.  We
   // cannot (yet) promote a memory location that is loaded and stored in
bool llvm::promoteLoopAccessesToScalars(
   // memory model.
   if (!SafeToInsertStore) {
     if (IsKnownThreadLocalObject)
       SafeToInsertStore = true;
     else {
       Value *Object = getUnderlyingObject(SomePtr);
       SafeToInsertStore =
           (isAllocLikeFn(Object, TLI) || isa<AllocaInst>(Object)) &&
-          !PointerMayBeCaptured(Object, true, true);
+          isNotCapturedBeforeOrInLoop(Object, CurLoop, DT);
     }
   }
 
   // If we've still failed to prove we can sink the store, give up.
   if (!SafeToInsertStore)
     return false;
 
   // Otherwise, this is safe to promote, lets do it!

llvm/test/Transforms/LICM/promote-capture.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt -S -passes='loop-mssa(licm)' < %s | FileCheck %s
 
 declare i1 @cond(i32 %v) readnone
 declare void @capture(i32* %p) readnone
 
 define void @test_captured_after_loop(i32 %len) {
 ; CHECK-LABEL: @test_captured_after_loop(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[COUNT:%.*]] = alloca i32, align 4
 ; CHECK-NEXT:    store i32 0, i32* [[COUNT]], align 4
+; CHECK-NEXT:    [[COUNT_PROMOTED:%.*]] = load i32, i32* [[COUNT]], align 4
 ; CHECK-NEXT:    br label [[LOOP:%.*]]
 ; CHECK:       loop:
-; CHECK-NEXT:    [[I:%.*]] = phi i32 [ 0, [[ENTRY:%.*]] ], [ [[I_NEXT:%.*]], [[LATCH:%.*]] ]
+; CHECK-NEXT:    [[C_INC2:%.*]] = phi i32 [ [[COUNT_PROMOTED]], [[ENTRY:%.*]] ], [ [[C_INC1:%.*]], [[LATCH:%.*]] ]
+; CHECK-NEXT:    [[I:%.*]] = phi i32 [ 0, [[ENTRY]] ], [ [[I_NEXT:%.*]], [[LATCH]] ]
 ; CHECK-NEXT:    [[COND:%.*]] = call i1 @cond(i32 [[I]])
 ; CHECK-NEXT:    br i1 [[COND]], label [[IF:%.*]], label [[LATCH]]
 ; CHECK:       if:
-; CHECK-NEXT:    [[C:%.*]] = load i32, i32* [[COUNT]], align 4
-; CHECK-NEXT:    [[C_INC:%.*]] = add i32 [[C]], 1
-; CHECK-NEXT:    store i32 [[C_INC]], i32* [[COUNT]], align 4
+; CHECK-NEXT:    [[C_INC:%.*]] = add i32 [[C_INC2]], 1
 ; CHECK-NEXT:    br label [[LATCH]]
 ; CHECK:       latch:
+; CHECK-NEXT:    [[C_INC1]] = phi i32 [ [[C_INC]], [[IF]] ], [ [[C_INC2]], [[LOOP]] ]
 ; CHECK-NEXT:    [[I_NEXT]] = add nuw i32 [[I]], 1
 ; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i32 [[I_NEXT]], [[LEN:%.*]]
 ; CHECK-NEXT:    br i1 [[CMP]], label [[EXIT:%.*]], label [[LOOP]]
 ; CHECK:       exit:
+; CHECK-NEXT:    [[C_INC1_LCSSA:%.*]] = phi i32 [ [[C_INC1]], [[LATCH]] ]
+; CHECK-NEXT:    store i32 [[C_INC1_LCSSA]], i32* [[COUNT]], align 4
 ; CHECK-NEXT:    call void @capture(i32* [[COUNT]])
 ; CHECK-NEXT:    ret void
 ;
 entry:
   %count = alloca i32
   store i32 0, i32* %count
   br label %loop