[LICM] Store promotion when memory is thread local

[reames]

Summary

This patch teaches LICM's implementation of store promotion to exploit the fact that the memory location being accessed might be provable thread local. The fact it's thread local weakens the requirements for where we can insert stores since no other thread can observe the write. This allows us perform store promotion even in cases where the store is not guaranteed to execute in the loop.

Two key assumption worth drawing out is that this assumes a) no-capture is strong enough to imply no-escape, and b) standard allocation functions like malloc, calloc, and operator new return values which can be assumed not to have previously escaped.

In future work, it would be nice to generalize this so that it works without directly seeing the allocation site. I believe that the nocapture return attribute should be suitable for this purpose, but haven't investigated carefully. It's also likely that we could support unescaped allocas with similar reasoning, but since SROA and Mem2Reg should destroy those, they're less interesting than they first might seem.

[sanjoy]

lib/Transforms/Scalar/LICM.cpp

928 ↗

(On Diff #46568)

Though I think I understand it now, I found the code somewhat
difficult to follow because of the overlap between
GuaranteedToExecute && CanSpeculateLoad. Is it possible to
instead track CanSpeculateLoad and CanSpeculateStore?

[reames]

lib/Transforms/Scalar/LICM.cpp

928 ↗

(On Diff #46568)

This would be a completely reasonable refactoring, but I'd strongly prefer to do that as a separate change.

[reames]

ping

[sanjoy]

Overall the logic lgtm. However, should this be disabled when optimizing for size?

lib/Transforms/Scalar/LICM.cpp

1012 ↗

(On Diff #46568)

So part of the confusion here for me is that CanSpeculateLoad looks like it is really tracking something else. Does it make sense to rename it to CanPromoteAliasSet?

[reames]

In D16783#367396, @sanjoy wrote:

Overall the logic lgtm. However, should this be disabled when optimizing for size?

Why would the thread local version be disabled when optimizing for size and not the existing code? That doesn't seem to make any sense.

[reames]

lib/Transforms/Scalar/LICM.cpp

1012 ↗

(On Diff #46568)

I really not getting the confusion here. Per the comment on the variable you're asking me to rename, it really is answering only the speculation legality question. In particular, that by itself is not enough to establish promotion is legal. That's the whole point.

[sanjoy]

In D16783#368247, @reames wrote:

In D16783#367396, @sanjoy wrote:

Overall the logic lgtm. However, should this be disabled when optimizing for size?

Why would the thread local version be disabled when optimizing for size and not the existing code? That doesn't seem to make any sense.

Given that you're potentially duplicating a store, if this change
makes the promotion fire a lot more often maybe that's bad for
-Os. But I have no hard numbers for this.

lib/Transforms/Scalar/LICM.cpp

1012 ↗

(On Diff #46568)

I really not getting the confusion here.

That just proves you're smarter than I am. :)

My confusion initially was that you're checking if the store location
is dereferenceable, but yet use that to decide if the load can be
speculated. But now I get that this is really about proving that
any location in the alias set can be speculatively loaded from,
since they're all MustAlias.

Just as a understanding check -- if you remove this clause completely,
that will just be a quality of implementation regression, not a
correctness issue, right?

[reames]

In D16783#368332, @sanjoy wrote:

In D16783#368247, @reames wrote:

In D16783#367396, @sanjoy wrote:

Overall the logic lgtm. However, should this be disabled when optimizing for size?

Why would the thread local version be disabled when optimizing for size and not the existing code? That doesn't seem to make any sense.

Given that you're potentially duplicating a store, if this change
makes the promotion fire a lot more often maybe that's bad for
-Os. But I have no hard numbers for this.

I'd argue that either the whole transform should be restricted in Os, or not. My change isn't a change to any of the profitability heuristics, just a legality one.

p.s. Once you've considered my points, can you give another explicit LGTM? I want to make sure that still stands after our follow on discussion.

lib/Transforms/Scalar/LICM.cpp

1012 ↗

(On Diff #46568)

Correct. Using the store to establish dereferenceability lets the transform kick in more often, but is not required for correctness.

I'll add a clarify comment about the mustalias bit. That's not obvious from the code and might help to prevent future confusion.

[sanjoy]

lgtm

llvm/trunk/include/llvm/Transforms/Utils/LoopUtils.h

 /// the stores in the loop, looking for stores to Must pointers which are
 /// loop invariant. It takes AliasSet, Loop exit blocks vector, loop exit blocks
 /// insertion point vector, PredIteratorCache, LoopInfo, DominatorTree, Loop,
 /// AliasSet information for all instructions of the loop and loop safety
 /// information as arguments. It returns changed status.
 bool promoteLoopAccessesToScalars(AliasSet &, SmallVectorImpl<BasicBlock*> &,
                                   SmallVectorImpl<Instruction*> &,
                                   PredIteratorCache &, LoopInfo *,
-                                  DominatorTree *, Loop *, AliasSetTracker *,
-                                  LICMSafetyInfo *);
+                                  DominatorTree *, const TargetLibraryInfo *,
+                                  Loop *, AliasSetTracker *, LICMSafetyInfo *);
 
 /// \brief Computes safety information for a loop
 /// checks loop body & header for the possibility of may throw
 /// exception, it takes LICMSafetyInfo and loop as argument.
 /// Updates safety information in LICMSafetyInfo argument.
 void computeLICMSafetyInfo(LICMSafetyInfo *, Loop *);
 
 /// \brief Returns the instructions that use values defined in the loop.

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

 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/Transforms/Scalar.h"
 #include "llvm/ADT/Statistic.h"
 #include "llvm/Analysis/AliasAnalysis.h"
 #include "llvm/Analysis/AliasSetTracker.h"
 #include "llvm/Analysis/BasicAliasAnalysis.h"
+#include "llvm/Analysis/CaptureTracking.h"
 #include "llvm/Analysis/ConstantFolding.h"
 #include "llvm/Analysis/GlobalsModRef.h"
 #include "llvm/Analysis/LoopInfo.h"
 #include "llvm/Analysis/LoopPass.h"
+#include "llvm/Analysis/MemoryBuiltins.h"
 #include "llvm/Analysis/ScalarEvolution.h"
 #include "llvm/Analysis/ScalarEvolutionAliasAnalysis.h"
 #include "llvm/Analysis/TargetLibraryInfo.h"
 #include "llvm/Analysis/ValueTracking.h"
 #include "llvm/IR/CFG.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/DataLayout.h"
 #include "llvm/IR/DerivedTypes.h"
bool LICM::runOnLoop(Loop *L, LPPassManager &LPM) {
   if (!DisablePromotion && (Preheader || L->hasDedicatedExits())) {
     SmallVector<BasicBlock *, 8> ExitBlocks;
     SmallVector<Instruction *, 8> InsertPts;
     PredIteratorCache PIC;
 
     // Loop over all of the alias sets in the tracker object.
     for (AliasSet &AS : *CurAST)
       Changed |= promoteLoopAccessesToScalars(AS, ExitBlocks, InsertPts,
-                                              PIC, LI, DT, CurLoop, 
+                                              PIC, LI, DT, TLI, CurLoop,
                                               CurAST, &SafetyInfo);
 
     // Once we have promoted values across the loop body we have to recursively
     // reform LCSSA as any nested loop may now have values defined within the
     // loop used in the outer loop.
     // FIXME: This is really heavy handed. It would be a bit better to use an
     // SSAUpdater strategy during promotion that was LCSSA aware and reformed
     // it as it went.

 /// 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(AliasSet &AS,
                                         SmallVectorImpl<BasicBlock*>&ExitBlocks,
                                         SmallVectorImpl<Instruction*>&InsertPts,
                                         PredIteratorCache &PIC, LoopInfo *LI, 
-                                        DominatorTree *DT, Loop *CurLoop, 
+                                        DominatorTree *DT,
+                                        const TargetLibraryInfo *TLI,
+                                        Loop *CurLoop,
                                         AliasSetTracker *CurAST, 
                                         LICMSafetyInfo * SafetyInfo) { 
   // Verify inputs.
   assert(LI != nullptr && DT != nullptr && 
          CurLoop != nullptr && CurAST != nullptr && 
          SafetyInfo != nullptr && 
          "Unexpected Input to promoteLoopAccessesToScalars");
 
bool llvm::promoteLoopAccessesToScalars(AliasSet &AS,
   //    for () { if (c) *P += 1; }
   //
   // into:
   //
   //    tmp = *P;  for () { if (c) tmp +=1; } *P = tmp;
   //
   // is not safe, because *P may only be valid to access if 'c' is true.
   //
+  // The safety property divides into two parts:
+  // 1) The memory may not be dereferenceable on entry to the loop.  In this
+  //    case, we can't insert the required load in the preheader.
+  // 2) The memory model does not allow us to insert a store along any dynamic
+  //    path which did not originally have one.
+  //
   // It is safe to promote P if all uses are direct load/stores and if at
   // least one is guaranteed to be executed.
   bool GuaranteedToExecute = false;
 
+  // It is also safe to promote P if we can prove that speculating a load into
+  // the preheader is safe (i.e. proving dereferenceability on all
+  // paths through the loop), and that the memory can be proven thread local
+  // (so that the memory model requirement doesn't apply.)  We first establish
+  // the former, and then run a capture analysis below to establish the later.
+  // We can use any access within the alias set to prove dereferenceability
+  // since they're all must alias.
+  bool CanSpeculateLoad = false;
+
   SmallVector<Instruction*, 64> LoopUses;
   SmallPtrSet<Value*, 4> PointerMustAliases;
 
   // We start with an alignment of one and try to find instructions that allow
   // us to prove better alignment.
   unsigned Alignment = 1;
   AAMDNodes AATags;
   bool HasDedicatedExits = CurLoop->hasDedicatedExits();
 
+  // Don't sink stores from loops without dedicated block exits. Exits
+  // containing indirect branches are not transformed by loop simplify,
+  // make sure we catch that. An additional load may be generated in the
+  // preheader for SSA updater, so also avoid sinking when no preheader
+  // is available.
+  if (!HasDedicatedExits || !Preheader)
+    return false;
+  
+  const DataLayout &MDL = Preheader->getModule()->getDataLayout();
+
   // 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
   // different sizes.  While we are at it, collect alignment and AA info.
   bool Changed = false;
   for (AliasSet::iterator ASI = AS.begin(), E = AS.end(); ASI != E; ++ASI) {
     Value *ASIV = ASI->getValue();
     PointerMustAliases.insert(ASIV);
 
for (User *U : ASIV->users()) {
         continue;
 
       // If there is an non-load/store instruction in the loop, we can't promote
       // it.
       if (const LoadInst *Load = dyn_cast<LoadInst>(UI)) {
         assert(!Load->isVolatile() && "AST broken");
         if (!Load->isSimple())
           return Changed;
+
+        if (!GuaranteedToExecute && !CanSpeculateLoad)
+          CanSpeculateLoad =
+            isSafeToExecuteUnconditionally(*Load, DT, TLI, CurLoop,
+                                           SafetyInfo,
+                                           Preheader->getTerminator());
       } else if (const StoreInst *Store = dyn_cast<StoreInst>(UI)) {
         // Stores *of* the pointer are not interesting, only stores *to* the
         // pointer.
         if (UI->getOperand(1) != ASIV)
           continue;
         assert(!Store->isVolatile() && "AST broken");
         if (!Store->isSimple())
           return Changed;
-        // Don't sink stores from loops without dedicated block exits. Exits
-        // containing indirect branches are not transformed by loop simplify,
-        // make sure we catch that. An additional load may be generated in the
-        // preheader for SSA updater, so also avoid sinking when no preheader
-        // is available.
-        if (!HasDedicatedExits || !Preheader)
-          return Changed;
 
         // Note that we only check GuaranteedToExecute inside the store case
         // so that we do not introduce stores where they did not exist before
         // (which would break the LLVM concurrency model).
 
         // If the alignment of this instruction allows us to specify a more
         // restrictive (and performant) alignment and if we are sure this
         // instruction will be executed, update the alignment.
         // Larger is better, with the exception of 0 being the best alignment.
         unsigned InstAlignment = Store->getAlignment();
         if ((InstAlignment > Alignment || InstAlignment == 0) && Alignment != 0)
           if (isGuaranteedToExecute(*UI, DT, CurLoop, SafetyInfo)) {
             GuaranteedToExecute = true;
             Alignment = InstAlignment;
           }
 
         if (!GuaranteedToExecute)
           GuaranteedToExecute = isGuaranteedToExecute(*UI, DT, 
                                                       CurLoop, SafetyInfo);
 
+
+        if (!GuaranteedToExecute && !CanSpeculateLoad) {
+          CanSpeculateLoad =
+            isDereferenceableAndAlignedPointer(Store->getPointerOperand(),
+                                               Store->getAlignment(), MDL,
+                                               Preheader->getTerminator(),
+                                               DT, TLI);
+        }
       } else
         return Changed; // Not a load or store.
 
       // Merge the AA tags.
       if (LoopUses.empty()) {
         // On the first load/store, just take its AA tags.
         UI->getAAMetadata(AATags);
       } else if (AATags) {
         UI->getAAMetadata(AATags, /* Merge = */ true);
       }
 
       LoopUses.push_back(UI);
     }
   }
 
-  // If there isn't a guaranteed-to-execute instruction, we can't promote.
-  if (!GuaranteedToExecute)
+  // Check legality per comment above. Otherwise, we can't promote.
+  bool PromotionIsLegal = GuaranteedToExecute;
+  if (!PromotionIsLegal && CanSpeculateLoad) {
+    // If this is a thread local location, then we can insert stores along
+    // paths which originally didn't have them without violating the memory
+    // model. 
+    Value *Object = GetUnderlyingObject(SomePtr, MDL);
+    PromotionIsLegal = isAllocLikeFn(Object, TLI) &&
+      !PointerMayBeCaptured(Object, true, true);
+  }
+  if (!PromotionIsLegal)
     return Changed;
 
   // Figure out the loop exits and their insertion points, if this is the
   // first promotion.
   if (ExitBlocks.empty()) {
     CurLoop->getUniqueExitBlocks(ExitBlocks);
     InsertPts.clear();
     InsertPts.reserve(ExitBlocks.size());

llvm/trunk/test/Transforms/LICM/promote-tls.ll

+; RUN: opt -tbaa -basicaa -licm -S < %s | FileCheck %s
+
+; If we can prove a local is thread local, we can insert stores during
+; promotion which wouldn't be legal otherwise.
+
+target datalayout = "e-p:64:64:64-i1:8:8-i8:8:8-i16:16:16-i32:32:32-i64:64:64-f32:32:32-f64:64:64-v64:64:64-v128:128:128-a0:0:64-s0:64:64-f80:128:128-n8:16:32:64-S128"
+target triple = "x86_64-linux-generic"
+
+@p = external global i8*
+
+declare i8* @malloc(i64)
+
+; Exercise the TLS case
+define i32* @test(i32 %n) {
+entry:
+  ;; ignore the required null check for simplicity
+  %mem = call dereferenceable(16) noalias i8* @malloc(i64 16)
+  %addr = bitcast i8* %mem to i32*
+  br label %for.body.lr.ph
+
+for.body.lr.ph:                                   ; preds = %entry
+  br label %for.header
+
+for.header:
+  %i.02 = phi i32 [ 0, %for.body.lr.ph ], [ %inc, %for.body ]
+  %old = load i32, i32* %addr, align 4
+  ; deliberate impossible to analyze branch
+  %guard = load volatile i8*, i8** @p
+  %exitcmp = icmp eq i8* %guard, null
+  br i1 %exitcmp, label %for.body, label %early-exit
+
+early-exit:
+; CHECK-LABEL: early-exit:
+; CHECK: store i32 %new1.lcssa, i32* %addr, align 1
+  ret i32* null
+
+for.body:
+  %new = add i32 %old, 1
+  store i32 %new, i32* %addr, align 4
+  %inc = add nsw i32 %i.02, 1
+  %cmp = icmp slt i32 %inc, %n
+  br i1 %cmp, label %for.header, label %for.cond.for.end_crit_edge
+
+for.cond.for.end_crit_edge:                       ; preds = %for.body
+; CHECK-LABEL: for.cond.for.end_crit_edge:
+; CHECK: store i32 %new.lcssa, i32* %addr, align 1
+  %split = phi i32* [ %addr, %for.body ]
+  ret i32* null
+}
+
+declare i8* @not_malloc(i64)
+
+; Negative test - not TLS
+define i32* @test_neg(i32 %n) {
+entry:
+  ;; ignore the required null check for simplicity
+  %mem = call dereferenceable(16) noalias i8* @not_malloc(i64 16)
+  %addr = bitcast i8* %mem to i32*
+  br label %for.body.lr.ph
+
+for.body.lr.ph:                                   ; preds = %entry
+  br label %for.header
+
+for.header:
+  %i.02 = phi i32 [ 0, %for.body.lr.ph ], [ %inc, %for.body ]
+  %old = load i32, i32* %addr, align 4
+  ; deliberate impossible to analyze branch
+  %guard = load volatile i8*, i8** @p
+  %exitcmp = icmp eq i8* %guard, null
+  br i1 %exitcmp, label %for.body, label %early-exit
+
+early-exit:
+; CHECK-LABEL: early-exit:
+; CHECK-NOT: store
+  ret i32* null
+
+for.body:
+; CHECK-LABEL: for.body:
+; CHECK: store i32 %new, i32* %addr, align 4
+  %new = add i32 %old, 1
+  store i32 %new, i32* %addr, align 4
+  %inc = add nsw i32 %i.02, 1
+  %cmp = icmp slt i32 %inc, %n
+  br i1 %cmp, label %for.header, label %for.cond.for.end_crit_edge
+
+for.cond.for.end_crit_edge:                       ; preds = %for.body
+; CHECK-LABEL: for.cond.for.end_crit_edge:
+; CHECK-NOT: store
+  %split = phi i32* [ %addr, %for.body ]
+  ret i32* null
+}
+
+; Negative test - can't speculate load since branch
+; may control alignment
+define i32* @test_neg2(i32 %n) {
+entry:
+  ;; ignore the required null check for simplicity
+  %mem = call dereferenceable(16) noalias i8* @malloc(i64 16)
+  %addr = bitcast i8* %mem to i32*
+  br label %for.body.lr.ph
+
+for.body.lr.ph:                                   ; preds = %entry
+  br label %for.header
+
+for.header:
+  %i.02 = phi i32 [ 0, %for.body.lr.ph ], [ %inc, %for.body ]
+  ; deliberate impossible to analyze branch
+  %guard = load volatile i8*, i8** @p
+  %exitcmp = icmp eq i8* %guard, null
+  br i1 %exitcmp, label %for.body, label %early-exit
+
+early-exit:
+; CHECK-LABEL: early-exit:
+; CHECK-NOT: store
+  ret i32* null
+
+for.body:
+; CHECK-LABEL: for.body:
+; CHECK: store i32 %new, i32* %addr, align 4
+  %old = load i32, i32* %addr, align 4
+  %new = add i32 %old, 1
+  store i32 %new, i32* %addr, align 4
+  %inc = add nsw i32 %i.02, 1
+  %cmp = icmp slt i32 %inc, %n
+  br i1 %cmp, label %for.header, label %for.cond.for.end_crit_edge
+
+for.cond.for.end_crit_edge:                       ; preds = %for.body
+; CHECK-LABEL: for.cond.for.end_crit_edge:
+; CHECK-NOT: store
+  %split = phi i32* [ %addr, %for.body ]
+  ret i32* null
+}
+