[LoopDataPrefetch] Move prefetch to dominating position of two accesses

[jonpa]

Summary

I found that the LBM benchmark benefits significantly if enabling (stores) prefetching on the hot loop with '-mllvm -min-prefetch-stride=128 -mllvm -loop-prefetch-writes' on SystemZ. I then observed that the same memory addresses are accessed on different paths in different iterations. Currently, the LoopDataPrefetch pass emits one prefetch and then skips adding another one for the next close/identical access without taking this into consideration. I therefore tried the idea of making sure that the one prefetch actually dominates both accesses by moving it if necessary. For what I can see now on SystemZ at least, this patch gives yet another 7-8% of improvement on LBM (matching gcc)...

Is this generally a good idea? The prefetch distance heuristic may get slightly thrown off by this in some bigger loops, I guess.

[jonpa]

ping

[jonpa]

This change is included instead in https://reviews.llvm.org/D70228.

lib/Transforms/Scalar/LoopDataPrefetch.cpp

 
 STATISTIC(NumPrefetches, "Number of prefetches inserted");
 
 namespace {
 
 /// Loop prefetch implementation class.
 class LoopDataPrefetch {
 public:
-  LoopDataPrefetch(AssumptionCache *AC, LoopInfo *LI, ScalarEvolution *SE,
-                   const TargetTransformInfo *TTI,
+  LoopDataPrefetch(AssumptionCache *AC, DominatorTree *DT, LoopInfo *LI,
+                   ScalarEvolution *SE, const TargetTransformInfo *TTI,
                    OptimizationRemarkEmitter *ORE)
-      : AC(AC), LI(LI), SE(SE), TTI(TTI), ORE(ORE) {}
+      : AC(AC), DT(DT), LI(LI), SE(SE), TTI(TTI), ORE(ORE) {}
 
   bool run();
 
 private:
   bool runOnLoop(Loop *L);
 
   /// Check if the stride of the accesses is large enough to
   /// warrant a prefetch.
private:
 
   unsigned getMaxPrefetchIterationsAhead() {
     if (MaxPrefetchIterationsAhead.getNumOccurrences() > 0)
       return MaxPrefetchIterationsAhead;
     return TTI->getMaxPrefetchIterationsAhead();
   }
 
   AssumptionCache *AC;
+  DominatorTree *DT;
   LoopInfo *LI;
   ScalarEvolution *SE;
   const TargetTransformInfo *TTI;
   OptimizationRemarkEmitter *ORE;
 };
 
 /// Legacy class for inserting loop data prefetches.
 class LoopDataPrefetchLegacyPass : public FunctionPass {
 public:
   static char ID; // Pass ID, replacement for typeid
   LoopDataPrefetchLegacyPass() : FunctionPass(ID) {
     initializeLoopDataPrefetchLegacyPassPass(*PassRegistry::getPassRegistry());
   }
 
   void getAnalysisUsage(AnalysisUsage &AU) const override {
     AU.addRequired<AssumptionCacheTracker>();
+    AU.addRequired<DominatorTreeWrapperPass>();
     AU.addPreserved<DominatorTreeWrapperPass>();
     AU.addRequired<LoopInfoWrapperPass>();
     AU.addPreserved<LoopInfoWrapperPass>();
     AU.addRequired<OptimizationRemarkEmitterWrapperPass>();
     AU.addRequired<ScalarEvolutionWrapperPass>();
     AU.addPreserved<ScalarEvolutionWrapperPass>();
     AU.addRequired<TargetTransformInfoWrapperPass>();
   }
if (!ConstStride)
     return false;
 
   unsigned AbsStride = std::abs(ConstStride->getAPInt().getSExtValue());
   return TargetMinStride <= AbsStride;
 }
 
 PreservedAnalyses LoopDataPrefetchPass::run(Function &F,
                                             FunctionAnalysisManager &AM) {
+  DominatorTree *DT = &AM.getResult<DominatorTreeAnalysis>(F);
   LoopInfo *LI = &AM.getResult<LoopAnalysis>(F);
   ScalarEvolution *SE = &AM.getResult<ScalarEvolutionAnalysis>(F);
   AssumptionCache *AC = &AM.getResult<AssumptionAnalysis>(F);
   OptimizationRemarkEmitter *ORE =
       &AM.getResult<OptimizationRemarkEmitterAnalysis>(F);
   const TargetTransformInfo *TTI = &AM.getResult<TargetIRAnalysis>(F);
 
-  LoopDataPrefetch LDP(AC, LI, SE, TTI, ORE);
+  LoopDataPrefetch LDP(AC, DT, LI, SE, TTI, ORE);
   bool Changed = LDP.run();
 
   if (Changed) {
     PreservedAnalyses PA;
     PA.preserve<DominatorTreeAnalysis>();
     PA.preserve<LoopAnalysis>();
     return PA;
   }
 
   return PreservedAnalyses::all();
 }
 
 bool LoopDataPrefetchLegacyPass::runOnFunction(Function &F) {
   if (skipFunction(F))
     return false;
 
+  DominatorTree *DT = &getAnalysis<DominatorTreeWrapperPass>().getDomTree();
   LoopInfo *LI = &getAnalysis<LoopInfoWrapperPass>().getLoopInfo();
   ScalarEvolution *SE = &getAnalysis<ScalarEvolutionWrapperPass>().getSE();
   AssumptionCache *AC =
       &getAnalysis<AssumptionCacheTracker>().getAssumptionCache(F);
   OptimizationRemarkEmitter *ORE =
       &getAnalysis<OptimizationRemarkEmitterWrapperPass>().getORE();
   const TargetTransformInfo *TTI =
       &getAnalysis<TargetTransformInfoWrapperPass>().getTTI(F);
 
-  LoopDataPrefetch LDP(AC, LI, SE, TTI, ORE);
+  LoopDataPrefetch LDP(AC, DT, LI, SE, TTI, ORE);
   return LDP.run();
 }
 
 bool LoopDataPrefetch::run() {
   // If PrefetchDistance is not set, don't run the pass.  This gives an
   // opportunity for targets to run this pass for selected subtargets only
   // (whose TTI sets PrefetchDistance).
   if (getPrefetchDistance() == 0)
bool LoopDataPrefetch::runOnLoop(Loop *L) {
   unsigned LoopConstantTripCount = SE->getSmallConstantTripCount(L);
   if (LoopConstantTripCount && LoopConstantTripCount < ItersAhead)
     return MadeChange;
 
   LLVM_DEBUG(dbgs() << "Prefetching " << ItersAhead
                     << " iterations ahead (loop size: " << LoopSize << ") in "
                     << L->getHeader()->getParent()->getName() << ": " << *L);
 
-  SmallVector<std::pair<Instruction *, const SCEVAddRecExpr *>, 16> PrefLoads;
+  // Pairs of (Prefetch, SCEVAddRecExpr) for each emitted prefetch.
+  SmallVector<std::pair<Instruction *, const SCEVAddRecExpr *>, 16> Prefetches;
   for (const auto BB : L->blocks()) {
     for (auto &I : *BB) {
       Value *PtrValue;
       Instruction *MemI;
 
       if (LoadInst *LMemI = dyn_cast<LoadInst>(&I)) {
         MemI = LMemI;
         PtrValue = LMemI->getPointerOperand();
for (auto &I : *BB) {
       if (!LSCEVAddRec)
         continue;
 
       // Check if the stride of the accesses is large enough to warrant a
       // prefetch.
       if (!isStrideLargeEnough(LSCEVAddRec))
         continue;
 
-      // We don't want to double prefetch individual cache lines. If this load
-      // is known to be within one cache line of some other load that has
-      // already been prefetched, then don't prefetch this one as well.
+      // We don't want to double prefetch individual cache lines. If this
+      // access is known to be within one cache line of some other one that
+      // has already been prefetched, then don't prefetch this one as well.
       bool DupPref = false;
-      for (const auto &PrefLoad : PrefLoads) {
-        const SCEV *PtrDiff = SE->getMinusSCEV(LSCEVAddRec, PrefLoad.second);
+      for (const auto &Pref : Prefetches) {
+        const SCEV *PtrDiff = SE->getMinusSCEV(LSCEVAddRec, Pref.second);
         if (const SCEVConstant *ConstPtrDiff =
             dyn_cast<SCEVConstant>(PtrDiff)) {
           int64_t PD = std::abs(ConstPtrDiff->getValue()->getSExtValue());
           if (PD < (int64_t) TTI->getCacheLineSize()) {
+            /// Make sure that the prefetch Pref already emitted is also
+            /// executed each iteration MemI is reached, by moving the
+            /// prefetch up in the CFG if needed.
+            BasicBlock *PrefetchBB = Pref.first->getParent();
+            if (PrefetchBB != BB) {
+              BasicBlock *DomBB = DT->findNearestCommonDominator(PrefetchBB, BB);
+              if (DomBB != PrefetchBB)
+                Pref.first->moveBefore(DomBB->getTerminator());
+            }
             DupPref = true;
             break;
           }
         }
       }
       if (DupPref)
         continue;
 
       const SCEV *NextLSCEV = SE->getAddExpr(LSCEVAddRec, SE->getMulExpr(
         SE->getConstant(LSCEVAddRec->getType(), ItersAhead),
         LSCEVAddRec->getStepRecurrence(*SE)));
       if (!isSafeToExpand(NextLSCEV, *SE))
         continue;
 
-      PrefLoads.push_back(std::make_pair(MemI, LSCEVAddRec));
-
       Type *I8Ptr = Type::getInt8PtrTy(BB->getContext(), PtrAddrSpace);
       SCEVExpander SCEVE(*SE, I.getModule()->getDataLayout(), "prefaddr");
       Value *PrefPtrValue = SCEVE.expandCodeFor(NextLSCEV, I8Ptr, MemI);
 
       IRBuilder<> Builder(MemI);
       Module *M = BB->getParent()->getParent();
       Type *I32 = Type::getInt32Ty(BB->getContext());
       Function *PrefetchFunc = Intrinsic::getDeclaration(
           M, Intrinsic::prefetch, PrefPtrValue->getType());
+      Instruction *Prefetch =
         Builder.CreateCall(
             PrefetchFunc,
             {PrefPtrValue,
              ConstantInt::get(I32, MemI->mayReadFromMemory() ? 0 : 1),
              ConstantInt::get(I32, 3), ConstantInt::get(I32, 1)});
+      Prefetches.push_back(std::make_pair(Prefetch, LSCEVAddRec));
       ++NumPrefetches;
       LLVM_DEBUG(dbgs() << "  Access: " << *PtrValue << ", SCEV: " << *LSCEV
                         << "\n");
       ORE->emit([&]() {
         return OptimizationRemark(DEBUG_TYPE, "Prefetched", MemI)
                << "prefetched memory access";
       });
 
       MadeChange = true;
     }
   }
 
   return MadeChange;
 }

test/CodeGen/SystemZ/prefetch-03.ll

+; RUN: llc < %s -mtriple=s390x-linux-gnu -mcpu=z14 -prefetch-distance=50 \
+; RUN:   -loop-prefetch-writes -stop-after=loop-data-prefetch | FileCheck %s
+;
+; Check that prefetches are emitted in a position that is executed each
+; iteration for each targeted memory instruction. The two stores in %true and
+; %false are within one cache line in memory, so they should get a single
+; prefetch in %for.body.
+;
+; CHECK-LABEL: for.body
+; CHECK: call void @llvm.prefetch.p0i8(i8* {{.*}}, i32 0
+; CHECK: call void @llvm.prefetch.p0i8(i8* {{.*}}, i32 1
+; CHECK-LABEL: true
+; CHECK-LABEL: false
+; CHECK-LABEL: latch
+
+define void @fun(i32* nocapture %Src, i32* nocapture readonly %Dst) {
+entry:
+  br label %for.body
+
+for.body:
+  %indvars.iv = phi i64 [ 0, %entry ], [ %indvars.iv.next.9, %latch ]
+  %arrayidx = getelementptr inbounds i32, i32* %Dst, i64 %indvars.iv
+  %0 = load i32, i32* %arrayidx, align 4
+  %cmp = icmp sgt i32 %0, 0
+  br i1 %cmp, label %true, label %false
+
+true:  
+  %arrayidx2 = getelementptr inbounds i32, i32* %Src, i64 %indvars.iv
+  store i32 %0, i32* %arrayidx2, align 4
+  br label %latch
+
+false:
+  %a = add i64 %indvars.iv, 8
+  %arrayidx3 = getelementptr inbounds i32, i32* %Src, i64 %a
+  store i32 %0, i32* %arrayidx3, align 4
+  br label %latch
+
+latch:
+  %indvars.iv.next.9 = add nuw nsw i64 %indvars.iv, 1600
+  %cmp.9 = icmp ult i64 %indvars.iv.next.9, 11200
+  br i1 %cmp.9, label %for.body, label %for.cond.cleanup
+
+for.cond.cleanup:
+  ret void
+}
+