[SelectOpti] Restrict load sinking

[apostolakis]

Summary

This is a follow-up to D133777, which resolved a use-after-free case but
did not cover all possible memory bugs due to misplacement of loads.

In short, the overall problem was that sinked loads could be moved after
state-modifying instructions leading to memory bugs.

The solution is to restrict load sinking unless it is found to be sound.
i) Within a basic block (to-be-sinked load and select-user are in the same BB),
loads can be sinked only if there is no intervening state-modifying instruction.
This is a conservative approach to avoid resorting to alias analysis to detect
potential memory overlap.
ii) Across basic blocks, sinking of loads is avoided. This is because going over
multiple basic blocks looking for memory conflicts could be computationally
expensive and also unlikely to allow loads to sink. Further, experiments showed
that not sinking these loads has a slight positive performance effect.
Maybe for some of these loads, having some separation allows enough time
for the load to be executed in time for its user. This is not the case for
floating point operations that benefit more from sinking.

The solution in D133777 was essentially undone in this patch,
since the latter is a complete solution to the observed problem.

Overall, the performance impact of this patch is minimal.
Tested on two internal Google workloads with instrPGO.
Search application showed <0.05% perf difference,
while the database one showed a slight improvement,
but not statistically significant.

[apostolakis]

Typo in comment

[davidxl]

llvm/test/CodeGen/X86/select-optimize.ll

247

why is this unsafe to sink?

[apostolakis]

llvm/test/CodeGen/X86/select-optimize.ll

247

it is not unsafe to sink this load but the current version of the transformation does not check for sink-safety if such loads from other basic blocks and thus they should not be sunk.
See patch description for why cross-basic-block load moves were avoided, even though there might be cases such as this one where sinking would have been safe.

[davidxl]

llvm/lib/CodeGen/SelectOptimize.cpp

713

In most of the cases, the load and sinkpt are not in the same bb right -- as the case in the test case I commented.

[apostolakis]

llvm/lib/CodeGen/SelectOptimize.cpp

713

Actually, most of the sinkable loads are in the same basic block.
For a search workload, only 2.8% of the sinkable loads were in a different BB, while the rest (i.e., the vast majority, 97.2%) were in the same BB as the sinkpt.
The intuition here is that sinkable loads, as defined in this optimization, are loads that have a single-use (only used for the purpose of computing the select operand). Thus, it is expected that they will be close-by in the same BB as the select, i.e., their use.
If they are not in the same BB, there might be a good reason for that. Either there is a state-modifying and aliasing instruction in-between that prevents them for moving, or due to some other optimization (e.g., for both reported bugs, the loaded stack objects were freed early).

Note also that for the search workload, out of all the sinkable loads in the same BB, only 3 load-sinks (0.06% of sinkable loads) were conservatively skipped due to the check for potential (although non-aliasing in these cased) state-modifying instructions in line 717. So, even for same-BB loads, the impact of this transformation seems minimal. This also explains how rare it was for an actual memory bug to occur.

These statistics also justify the small perf impact of this patch

[apostolakis]

llvm/lib/CodeGen/SelectOptimize.cpp

713

Oh I now realized that your comment was on this check within isSafeToSinkLoad, and not about skipping cross-BB load sinks in line 760, although the provided statistics should provide more clarity overall.

Regarding this specific check: isSafeToSinkLoad function is only called for cases where both LoadI and SinkPt are in the same BB (see callsite in line 762). Additionally, the iteration here will only work properly for intra-BB searches, so to prevent any misuse it returns conservatively false.

Actually, re-reading this check I could have just called this function and avoid the if-statement in line 760

[davidxl]

Ok -- my confusion was the naming of the SinkPt variable -- it should actually be named 'SI'. I through the sinkpt is the actual bb after select is expanded into branch.

[apostolakis]

In D133999#3796300, @davidxl wrote:

Ok -- my confusion was the naming of the SinkPt variable -- it should actually be named 'SI'. I through the sinkpt is the actual bb after select is expanded into branch.

Oh I see. Naming it as sink-point was indeed confusing. I will change it.

[davidxl]

llvm/test/CodeGen/X86/select-optimize.ll

203

For sunkable loads, should the lifetime marker still be preserved?

[apostolakis]

llvm/test/CodeGen/X86/select-optimize.ll

203

The lifetime-marker intrinsic writes to memory and thus there will not be any sunk loads bypassing lifetime markers.
This is why this patch also covers the use-after-free resolved by D133777.

[davidxl]

llvm/test/CodeGen/X86/select-optimize.ll

203

Can you add another case with lifemarker and check if it is not sunk?

[apostolakis]

llvm/test/CodeGen/X86/select-optimize.ll

203

Sure.

[apostolakis]

Add test with lifetime-marker, rename SinktPt to SI to avoid
confusion, and remove redundant check for same/different BB.

[davidxl]

lgtm

llvm/lib/CodeGen/SelectOptimize.cpp

 #include "llvm/CodeGen/TargetPassConfig.h"
 #include "llvm/CodeGen/TargetSchedule.h"
 #include "llvm/CodeGen/TargetSubtargetInfo.h"
 #include "llvm/IR/BasicBlock.h"
 #include "llvm/IR/Dominators.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/Instruction.h"
-#include "llvm/IR/IntrinsicInst.h"
 #include "llvm/IR/ProfDataUtils.h"
 #include "llvm/InitializePasses.h"
 #include "llvm/Pass.h"
 #include "llvm/Support/ScaledNumber.h"
 #include "llvm/Target/TargetMachine.h"
 #include "llvm/Transforms/Utils/SizeOpts.h"
 #include <algorithm>
 #include <memory>
private:
   // operand's (if any) dependence slice of any of the selects of the given
   // group.
   bool hasExpensiveColdOperand(const SmallVector<SelectInst *, 2> &ASI);
 
   // For a given source instruction, collect its backwards dependence slice
   // consisting of instructions exclusively computed for producing the operands
   // of the source instruction.
   void getExclBackwardsSlice(Instruction *I, std::stack<Instruction *> &Slice,
-                             bool ForSinking = false);
+                             bool ForSinking = false,
+                             Instruction *SinkPt = nullptr);
 
   // Returns true if the condition of the select is highly predictable.
   bool isSelectHighlyPredictable(const SelectInst *SI);
 
   // Loop-level checks to determine if a non-predicated version (with branches)
   // of the given loop is more profitable than its predicated version.
   bool checkLoopHeuristics(const Loop *L, const CostInfo LoopDepth[2]);
 
for (SelectGroup &ASI : ProfSIGroups) {
     SmallVector<std::stack<Instruction *>, 2> TrueSlices, FalseSlices;
     typedef std::stack<Instruction *>::size_type StackSizeType;
     StackSizeType maxTrueSliceLen = 0, maxFalseSliceLen = 0;
     for (SelectInst *SI : ASI) {
       // For each select, compute the sinkable dependence chains of the true and
       // false operands.
       if (auto *TI = dyn_cast<Instruction>(SI->getTrueValue())) {
         std::stack<Instruction *> TrueSlice;
-        getExclBackwardsSlice(TI, TrueSlice, true);
+        getExclBackwardsSlice(TI, TrueSlice, true, SI);
         maxTrueSliceLen = std::max(maxTrueSliceLen, TrueSlice.size());
         TrueSlices.push_back(TrueSlice);
       }
       if (auto *FI = dyn_cast<Instruction>(SI->getFalseValue())) {
         std::stack<Instruction *> FalseSlice;
-        getExclBackwardsSlice(FI, FalseSlice, true);
+        getExclBackwardsSlice(FI, FalseSlice, true, SI);
         maxFalseSliceLen = std::max(maxFalseSliceLen, FalseSlice.size());
         FalseSlices.push_back(FalseSlice);
       }
     }
     // In the case of multiple select instructions in the same group, the order
     // of non-dependent instructions (instructions of different dependence
     // slices) in the true/false blocks appears to affect performance.
     // Interleaving the slices seems to experimentally be the optimal approach.
if (TrueBlock == FalseBlock) {
       assert(TrueBlock == nullptr &&
              "Unexpected basic block transform while optimizing select");
 
       FalseBlock = BasicBlock::Create(SI->getContext(), "select.false",
                                       EndBlock->getParent(), EndBlock);
       auto *FalseBranch = BranchInst::Create(EndBlock, FalseBlock);
       FalseBranch->setDebugLoc(SI->getDebugLoc());
     }
-    // If there was sinking, move any lifetime-end intrinsic calls found in the
-    // StartBlock to the newly-created end block to ensure sound lifetime info
-    // after sinking (in case any use is sinked).
-    else {
-      SmallVector<Instruction *, 2> EndLifetimeCalls;
-      for (Instruction &I : *StartBlock) {
-        if (IntrinsicInst *II = dyn_cast<IntrinsicInst>(&I)) {
-          if (II->getIntrinsicID() == Intrinsic::lifetime_end) {
-            EndLifetimeCalls.push_back(&I);
-          }
-        }
-      }
-      for (auto *LC : EndLifetimeCalls) {
-        LC->moveBefore(&*EndBlock->getFirstInsertionPt());
-      }
-    }
 
     // Insert the real conditional branch based on the original condition.
     // If we did not create a new block for one of the 'true' or 'false' paths
     // of the condition, it means that side of the branch goes to the end block
     // directly and the path originates from the start block from the point of
     // view of the new PHI.
     BasicBlock *TT, *FT;
     if (TrueBlock == nullptr) {
if (ColdI) {
       if (AdjSliceCost >=
           ColdOperandMaxCostMultiplier * TargetTransformInfo::TCC_Expensive)
         return true;
     }
   }
   return false;
 }
 
+// Check if it is safe to move LoadI to the SinkPt.
+// Conservatively assume it is safe only if there is no instruction
+// modifying memory in-between the load and the sink point.
+static bool isSafeToSinkLoad(Instruction *LoadI, Instruction *SinkPt) {
+  // Ensure that the instructions are on the basic block.
+  if (LoadI->getParent() != SinkPt->getParent())
+    return false;
+  auto It = LoadI->getIterator();
+  while (&*It != SinkPt) {
+    if (It->mayWriteToMemory())
+      return false;
+    It++;
+  }
+  return true;
+}
+
 // For a given source instruction, collect its backwards dependence slice
 // consisting of instructions exclusively computed for the purpose of producing
 // the operands of the source instruction. As an approximation
 // (sufficiently-accurate in practice), we populate this set with the
 // instructions of the backwards dependence slice that only have one-use and
 // form an one-use chain that leads to the source instruction.
 void SelectOptimize::getExclBackwardsSlice(Instruction *I,
                                            std::stack<Instruction *> &Slice,
-                                           bool ForSinking) {
+                                           bool ForSinking,
+                                           Instruction *SinkPt) {
   SmallPtrSet<Instruction *, 2> Visited;
   std::queue<Instruction *> Worklist;
   Worklist.push(I);
   while (!Worklist.empty()) {
     Instruction *II = Worklist.front();
     Worklist.pop();
 
     // Avoid cycles.
     if (!Visited.insert(II).second)
       continue;
 
     if (!II->hasOneUse())
       continue;
 
     // Cannot soundly sink instructions with side-effects.
     // Terminator or phi instructions cannot be sunk.
     // Avoid sinking other select instructions (should be handled separetely).
     if (ForSinking && (II->isTerminator() || II->mayHaveSideEffects() ||
                        isa<SelectInst>(II) || isa<PHINode>(II)))
       continue;
 
+    // Avoid sinking loads in order not to skip state-modifying instructions,
+    // that may alias with the loaded address.
+    // Only allow sinking of loads within the same basic block that are
+    // conservatively proven to be safe.
+    if (ForSinking && II->mayReadFromMemory() && SinkPt) {
+      if (II->getParent() != SinkPt->getParent()) {
+        continue;
+      } else if (!isSafeToSinkLoad(II, SinkPt)) {
+        continue;
+      }
+    }
+
     // Avoid considering instructions with less frequency than the source
     // instruction (i.e., avoid colder code regions of the dependence slice).
     if (BFI->getBlockFreq(II->getParent()) < BFI->getBlockFreq(I->getParent()))
       continue;
 
     // Eligible one-use instruction added to the dependence slice.
     Slice.push(II);
 

llvm/test/CodeGen/X86/select-optimize.ll

 ; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[CMP:%.*]], i32 [[X:%.*]], i32 [[LOAD]], !prof [[PROF18]]
 ; CHECK-NEXT:    ret i32 [[SEL]]
 ;
   %load = load i32, ptr %a, align 8
   %sel = select i1 %cmp, i32 %x, i32 %load, !prof !17
   ret i32 %sel
 }
 
-; Cold value operand with load in its one-use dependence slice shoud result
+; Cold value operand with load in its one-use dependence slice should result
 ; into a branch with sinked dependence slice.
-; The lifetime-end intrinsics should be soundly preserved.
 define i32 @expensive_val_operand3(ptr nocapture %a, i32 %b, i32 %y, i1 %cmp) {
 ; CHECK-LABEL: @expensive_val_operand3(
 ; CHECK-NEXT:    [[SEL_FROZEN:%.*]] = freeze i1 [[CMP:%.*]]
 ; CHECK-NEXT:    br i1 [[SEL_FROZEN]], label [[SELECT_TRUE_SINK:%.*]], label [[SELECT_END:%.*]], !prof [[PROF18]]
 ; CHECK:       select.true.sink:
 ; CHECK-NEXT:    [[LOAD:%.*]] = load i32, ptr [[A:%.*]], align 8
 ; CHECK-NEXT:    [[X:%.*]] = add i32 [[LOAD]], [[B:%.*]]
 ; CHECK-NEXT:    br label [[SELECT_END]]
 ; CHECK:       select.end:
 ; CHECK-NEXT:    [[SEL:%.*]] = phi i32 [ [[X]], [[SELECT_TRUE_SINK]] ], [ [[Y:%.*]], [[TMP0:%.*]] ]
-; CHECK-NEXT:    call void @llvm.lifetime.end.p0(i64 2, ptr nonnull [[A]])
 ; CHECK-NEXT:    ret i32 [[SEL]]
 ;
   %load = load i32, ptr %a, align 8
-  call void @llvm.lifetime.end.p0(i64 2, ptr nonnull %a)
   %x = add i32 %load, %b
   %sel = select i1 %cmp, i32 %x, i32 %y, !prof !17
   ret i32 %sel
 }
 
-; Multiple uses of the load value operand.
-define i32 @expensive_val_operand4(i32 %a, ptr nocapture %b, i32 %x, i1 %cmp) {
+; Expensive cold value operand with unsafe-to-sink load (partial slice sinking).
+define i32 @expensive_val_operand4(ptr nocapture %a, i32 %b, i32 %y, i1 %cmp) {
 ; CHECK-LABEL: @expensive_val_operand4(
+; CHECK-NEXT:    [[LOAD:%.*]] = load i32, ptr [[A:%.*]], align 8
+; CHECK-NEXT:    call void @free(ptr [[A]])
+; CHECK-NEXT:    [[SEL_FROZEN:%.*]] = freeze i1 [[CMP:%.*]]
+; CHECK-NEXT:    br i1 [[SEL_FROZEN]], label [[SELECT_TRUE_SINK:%.*]], label [[SELECT_END:%.*]], !prof [[PROF18]]
+; CHECK:       select.true.sink:
+; CHECK-NEXT:    [[X:%.*]] = add i32 [[LOAD]], [[B:%.*]]
+; CHECK-NEXT:    br label [[SELECT_END]]
+; CHECK:       select.end:
+; CHECK-NEXT:    [[SEL:%.*]] = phi i32 [ [[X]], [[SELECT_TRUE_SINK]] ], [ [[Y:%.*]], [[TMP0:%.*]] ]
+; CHECK-NEXT:    ret i32 [[SEL]]
+;
+  %load = load i32, ptr %a, align 8
+  call void @free(ptr %a)
+  %x = add i32 %load, %b
+  %sel = select i1 %cmp, i32 %x, i32 %y, !prof !17
+  ret i32 %sel
+}
+
+; Expensive cold value operand with potentially-unsafe-to-sink load (located
+; in a different basic block and thus unchecked for sinkability).
+define i32 @expensive_val_operand5(ptr nocapture %a, i32 %b, i32 %y, i1 %cmp) {
+; CHECK-LABEL: @expensive_val_operand5(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[LOAD:%.*]] = load i32, ptr [[A:%.*]], align 8
+; CHECK-NEXT:    br label [[BB1:%.*]]
+; CHECK:       bb1:
+; CHECK-NEXT:    [[SEL_FROZEN:%.*]] = freeze i1 [[CMP:%.*]]
+; CHECK-NEXT:    br i1 [[SEL_FROZEN]], label [[SELECT_TRUE_SINK:%.*]], label [[SELECT_END:%.*]], !prof [[PROF18]]
+; CHECK:       select.true.sink:
+; CHECK-NEXT:    [[X:%.*]] = add i32 [[LOAD]], [[B:%.*]]
+; CHECK-NEXT:    br label [[SELECT_END]]
+; CHECK:       select.end:
+; CHECK-NEXT:    [[SEL:%.*]] = phi i32 [ [[X]], [[SELECT_TRUE_SINK]] ], [ [[Y:%.*]], [[BB1]] ]
+; CHECK-NEXT:    ret i32 [[SEL]]
+;
+entry:
+  %load = load i32, ptr %a, align 8
+  br label %bb1
+bb1:                                 ; preds = %entry
+  %x = add i32 %load, %b
+  %sel = select i1 %cmp, i32 %x, i32 %y, !prof !17
+  ret i32 %sel
+}
+
+; Multiple uses of the load value operand.
+define i32 @expensive_val_operand6(i32 %a, ptr nocapture %b, i32 %x, i1 %cmp) {
+; CHECK-LABEL: @expensive_val_operand6(
 ; CHECK-NEXT:    [[LOAD:%.*]] = load i32, ptr [[B:%.*]], align 4
 ; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[CMP:%.*]], i32 [[X:%.*]], i32 [[LOAD]]
 ; CHECK-NEXT:    [[ADD:%.*]] = add i32 [[SEL]], [[LOAD]]
 ; CHECK-NEXT:    ret i32 [[ADD]]
 ;
   %load = load i32, ptr %b, align 4
   %sel = select i1 %cmp, i32 %x, i32 %load
   %add = add i32 %sel, %load
for.body:                                         ; preds = %for.body.preheader, %for.body
   %indvars.iv.next = add nuw nsw i64 %indvars.iv, 1
   %exitcond.not = icmp eq i64 %indvars.iv.next, %wide.trip.count
   br i1 %exitcond.not, label %for.cond.cleanup.loopexit, label %for.body
 }
 
 ; Function Attrs: nounwind readnone speculatable willreturn
 declare void @llvm.dbg.value(metadata, metadata, metadata)
 
-; Function Attrs: argmemonly mustprogress nocallback nofree nosync nounwind willreturn
-declare void @llvm.lifetime.end.p0(i64 immarg, ptr nocapture)
+declare void @free(ptr nocapture)
 
 !llvm.module.flags = !{!0, !26, !27}
 !0 = !{i32 1, !"ProfileSummary", !1}
 !1 = !{!2, !3, !4, !5, !6, !7, !8, !9}
 !2 = !{!"ProfileFormat", !"InstrProf"}
 !3 = !{!"TotalCount", i64 10000}
 !4 = !{!"MaxCount", i64 10}
 !5 = !{!"MaxInternalCount", i64 1}