[TTI][Sink] Add target support for sinking more reads

[critson]

Summary

Allow target to indicate some reads may be sunk along critical
edges. This is useful for applications which have invariant
data loads or can establish apriori absence of memory writes.

[foad]

llvm/lib/Transforms/Scalar/Sink.cpp

84

Why not check MD_invariant_load directly here, rather than forcing every target to check it in their own canAlwaysSinkRead hook?

[arsenm]

llvm/lib/Transforms/Scalar/Sink.cpp

84

I agree, I don't see why this needs to be a target hook. We already have AA checks for invariant memory

[critson]

Logically I'd like some kind of control over the sinking process, but I guess that can be handled by removing invariant metadata as appropriate.

-> D135133

llvm/include/llvm/Analysis/TargetTransformInfo.h

 
   /// \returns How the target needs this vector-predicated operation to be
   /// transformed.
   VPLegalization getVPLegalizationStrategy(const VPIntrinsic &PI) const;
   /// @}
 
   /// @}
 
+  /// \returns True if the target has determined this read can be sunk
+  /// across critical edges.
+  bool canAlwaysSinkRead(const Instruction *Inst,
+                         const BasicBlock *Dest) const;
+
 private:
   /// The abstract base class used to type erase specific TTI
   /// implementations.
   class Concept;
 
   /// The template model for the base class which wraps a concrete
   /// implementation in a type erased interface.
   template <typename T> class Model;
public:
   virtual unsigned getGISelRematGlobalCost() const = 0;
   virtual unsigned getMinTripCountTailFoldingThreshold() const = 0;
   virtual bool enableScalableVectorization() const = 0;
   virtual bool supportsScalableVectors() const = 0;
   virtual bool hasActiveVectorLength(unsigned Opcode, Type *DataType,
                                      Align Alignment) const = 0;
   virtual VPLegalization
   getVPLegalizationStrategy(const VPIntrinsic &PI) const = 0;
+
+  virtual bool canAlwaysSinkRead(const Instruction *Inst,
+                                 const BasicBlock *Dest) const = 0;
 };
 
 template <typename T>
 class TargetTransformInfo::Model final : public TargetTransformInfo::Concept {
   T Impl;
 
 public:
   Model(T Impl) : Impl(std::move(Impl)) {}
bool hasActiveVectorLength(unsigned Opcode, Type *DataType,
                              Align Alignment) const override {
     return Impl.hasActiveVectorLength(Opcode, DataType, Alignment);
   }
 
   VPLegalization
   getVPLegalizationStrategy(const VPIntrinsic &PI) const override {
     return Impl.getVPLegalizationStrategy(PI);
   }
+
+  bool canAlwaysSinkRead(const Instruction *Inst,
+                         const BasicBlock *Dest) const override {
+    return Impl.canAlwaysSinkRead(Inst, Dest);
+  }
 };
 
 template <typename T>
 TargetTransformInfo::TargetTransformInfo(T Impl)
     : TTIImpl(new Model<T>(Impl)) {}
 
 /// Analysis pass providing the \c TargetTransformInfo.
 ///

llvm/include/llvm/Analysis/TargetTransformInfoImpl.h

public:
 
   TargetTransformInfo::VPLegalization
   getVPLegalizationStrategy(const VPIntrinsic &PI) const {
     return TargetTransformInfo::VPLegalization(
         /* EVLParamStrategy */ TargetTransformInfo::VPLegalization::Discard,
         /* OperatorStrategy */ TargetTransformInfo::VPLegalization::Convert);
   }
 
+  bool canAlwaysSinkRead(const Instruction *Inst,
+                         const BasicBlock *Dest) const {
+    return false;
+  }
+
 protected:
   // Obtain the minimum required size to hold the value (without the sign)
   // In case of a vector it returns the min required size for one element.
   unsigned minRequiredElementSize(const Value *Val, bool &isSigned) const {
     if (isa<ConstantDataVector>(Val) || isa<ConstantVector>(Val)) {
       const auto *VectorValue = cast<Constant>(Val);
 
       // In case of a vector need to pick the max between the min

llvm/lib/Analysis/TargetTransformInfo.cpp

bool TargetTransformInfo::enableScalableVectorization() const {
   return TTIImpl->enableScalableVectorization();
 }
 
 bool TargetTransformInfo::hasActiveVectorLength(unsigned Opcode, Type *DataType,
                                                 Align Alignment) const {
   return TTIImpl->hasActiveVectorLength(Opcode, DataType, Alignment);
 }
 
+bool TargetTransformInfo::canAlwaysSinkRead(const Instruction *Inst,
+                                            const BasicBlock *Dest) const {
+  return TTIImpl->canAlwaysSinkRead(Inst, Dest);
+}
+
 TargetTransformInfo::Concept::~Concept() = default;
 
 TargetIRAnalysis::TargetIRAnalysis() : TTICallback(&getDefaultTTI) {}
 
 TargetIRAnalysis::TargetIRAnalysis(
     std::function<Result(const Function &)> TTICallback)
     : TTICallback(std::move(TTICallback)) {}
 

llvm/lib/Transforms/Scalar/Sink.cpp

 // they aren't executed on paths where their results aren't needed.
 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/Transforms/Scalar/Sink.h"
 #include "llvm/ADT/Statistic.h"
 #include "llvm/Analysis/AliasAnalysis.h"
 #include "llvm/Analysis/LoopInfo.h"
+#include "llvm/Analysis/TargetTransformInfo.h"
 #include "llvm/IR/Dominators.h"
 #include "llvm/InitializePasses.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_ostream.h"
 #include "llvm/Transforms/Scalar.h"
 using namespace llvm;
 
 #define DEBUG_TYPE "sink"
static bool isSafeToMove(Instruction *Inst, AliasAnalysis &AA,
   }
 
   return true;
 }
 
 /// IsAcceptableTarget - Return true if it is possible to sink the instruction
 /// in the specified basic block.
 static bool IsAcceptableTarget(Instruction *Inst, BasicBlock *SuccToSinkTo,
-                               DominatorTree &DT, LoopInfo &LI) {
+                               DominatorTree &DT, LoopInfo &LI,
+                               TargetTransformInfo &TTI) {
   assert(Inst && "Instruction to be sunk is null");
   assert(SuccToSinkTo && "Candidate sink target is null");
 
   // It's never legal to sink an instruction into a block which terminates in an
   // EH-pad.
   if (SuccToSinkTo->getTerminator()->isExceptionalTerminator())
     return false;
 
   // If the block has multiple predecessors, this would introduce computation
   // on different code paths.  We could split the critical edge, but for now we
   // just punt.
   // FIXME: Split critical edges if not backedges.
   if (SuccToSinkTo->getUniquePredecessor() != Inst->getParent()) {
     // We cannot sink a load across a critical edge - there may be stores in
     // other code paths.
-    if (Inst->mayReadFromMemory())
+    if (Inst->mayReadFromMemory() && !TTI.canAlwaysSinkRead(Inst, SuccToSinkTo))
       return false;
 
     // We don't want to sink across a critical edge if we don't dominate the
     // successor. We could be introducing calculations to new code paths.
     if (!DT.dominates(Inst->getParent(), SuccToSinkTo))
       return false;
 
     // Don't sink instructions into a loop.
     Loop *succ = LI.getLoopFor(SuccToSinkTo);
     Loop *cur = LI.getLoopFor(Inst->getParent());
     if (succ != nullptr && succ != cur)
       return false;
   }
 
   return true;
 }
 
 /// SinkInstruction - Determine whether it is safe to sink the specified machine
 /// instruction out of its current block into a successor.
 static bool SinkInstruction(Instruction *Inst,
                             SmallPtrSetImpl<Instruction *> &Stores,
-                            DominatorTree &DT, LoopInfo &LI, AAResults &AA) {
+                            DominatorTree &DT, LoopInfo &LI, AAResults &AA,
+                            TargetTransformInfo &TTI) {
 
   // Don't sink static alloca instructions.  CodeGen assumes allocas outside the
   // entry block are dynamically sized stack objects.
   if (AllocaInst *AI = dyn_cast<AllocaInst>(Inst))
     if (AI->isStaticAlloca())
       return false;
 
   // Check if it's safe to move the instruction.
for (Use &U : Inst->uses()) {
     if (!DT.dominates(BB, SuccToSinkTo))
       return false;
   }
 
   if (SuccToSinkTo) {
     // The nearest common dominator may be in a parent loop of BB, which may not
     // be beneficial. Find an ancestor.
     while (SuccToSinkTo != BB &&
-           !IsAcceptableTarget(Inst, SuccToSinkTo, DT, LI))
+           !IsAcceptableTarget(Inst, SuccToSinkTo, DT, LI, TTI))
       SuccToSinkTo = DT.getNode(SuccToSinkTo)->getIDom()->getBlock();
     if (SuccToSinkTo == BB)
       SuccToSinkTo = nullptr;
   }
 
   // If we couldn't find a block to sink to, ignore this instruction.
   if (!SuccToSinkTo)
     return false;
 
   LLVM_DEBUG(dbgs() << "Sink" << *Inst << " (";
              Inst->getParent()->printAsOperand(dbgs(), false); dbgs() << " -> ";
              SuccToSinkTo->printAsOperand(dbgs(), false); dbgs() << ")\n");
 
   // Move the instruction.
   Inst->moveBefore(&*SuccToSinkTo->getFirstInsertionPt());
   return true;
 }
 
 static bool ProcessBlock(BasicBlock &BB, DominatorTree &DT, LoopInfo &LI,
-                         AAResults &AA) {
+                         AAResults &AA, TargetTransformInfo &TTI) {
   // Can't sink anything out of a block that has less than two successors.
   if (BB.getTerminator()->getNumSuccessors() <= 1) return false;
 
   // Don't bother sinking code out of unreachable blocks. In addition to being
   // unprofitable, it can also lead to infinite looping, because in an
   // unreachable loop there may be nowhere to stop.
   if (!DT.isReachableFromEntry(&BB)) return false;
 
do {
     // sinking.
     ProcessedBegin = I == BB.begin();
     if (!ProcessedBegin)
       --I;
 
     if (Inst->isDebugOrPseudoInst())
       continue;
 
-    if (SinkInstruction(Inst, Stores, DT, LI, AA)) {
+    if (SinkInstruction(Inst, Stores, DT, LI, AA, TTI)) {
       ++NumSunk;
       MadeChange = true;
     }
 
     // If we just processed the first instruction in the block, we're done.
   } while (!ProcessedBegin);
 
   return MadeChange;
 }
 
 static bool iterativelySinkInstructions(Function &F, DominatorTree &DT,
-                                        LoopInfo &LI, AAResults &AA) {
+                                        LoopInfo &LI, AAResults &AA,
+                                        TargetTransformInfo &TTI) {
   bool MadeChange, EverMadeChange = false;
 
   do {
     MadeChange = false;
     LLVM_DEBUG(dbgs() << "Sinking iteration " << NumSinkIter << "\n");
     // Process all basic blocks.
     for (BasicBlock &I : F)
-      MadeChange |= ProcessBlock(I, DT, LI, AA);
+      MadeChange |= ProcessBlock(I, DT, LI, AA, TTI);
     EverMadeChange |= MadeChange;
     NumSinkIter++;
   } while (MadeChange);
 
   return EverMadeChange;
 }
 
 PreservedAnalyses SinkingPass::run(Function &F, FunctionAnalysisManager &AM) {
   auto &DT = AM.getResult<DominatorTreeAnalysis>(F);
   auto &LI = AM.getResult<LoopAnalysis>(F);
   auto &AA = AM.getResult<AAManager>(F);
+  auto &TTI = AM.getResult<TargetIRAnalysis>(F);
 
-  if (!iterativelySinkInstructions(F, DT, LI, AA))
+  if (!iterativelySinkInstructions(F, DT, LI, AA, TTI))
     return PreservedAnalyses::all();
 
   PreservedAnalyses PA;
   PA.preserveSet<CFGAnalyses>();
   return PA;
 }
 
 namespace {
   class SinkingLegacyPass : public FunctionPass {
   public:
     static char ID; // Pass identification
     SinkingLegacyPass() : FunctionPass(ID) {
       initializeSinkingLegacyPassPass(*PassRegistry::getPassRegistry());
     }
 
     bool runOnFunction(Function &F) override {
       auto &DT = getAnalysis<DominatorTreeWrapperPass>().getDomTree();
       auto &LI = getAnalysis<LoopInfoWrapperPass>().getLoopInfo();
       auto &AA = getAnalysis<AAResultsWrapperPass>().getAAResults();
+      auto &TTI = getAnalysis<TargetTransformInfoWrapperPass>().getTTI(F);
 
-      return iterativelySinkInstructions(F, DT, LI, AA);
+      return iterativelySinkInstructions(F, DT, LI, AA, TTI);
     }
 
     void getAnalysisUsage(AnalysisUsage &AU) const override {
       AU.setPreservesCFG();
       FunctionPass::getAnalysisUsage(AU);
       AU.addRequired<AAResultsWrapperPass>();
       AU.addRequired<DominatorTreeWrapperPass>();
       AU.addRequired<LoopInfoWrapperPass>();
+      AU.addRequired<TargetTransformInfoWrapperPass>();
       AU.addPreserved<DominatorTreeWrapperPass>();
       AU.addPreserved<LoopInfoWrapperPass>();
     }
   };
 } // end anonymous namespace
 
 char SinkingLegacyPass::ID = 0;
 INITIALIZE_PASS_BEGIN(SinkingLegacyPass, "sink", "Code sinking", false, false)
+INITIALIZE_PASS_DEPENDENCY(TargetTransformInfoWrapperPass)
 INITIALIZE_PASS_DEPENDENCY(LoopInfoWrapperPass)
 INITIALIZE_PASS_DEPENDENCY(DominatorTreeWrapperPass)
 INITIALIZE_PASS_DEPENDENCY(AAResultsWrapperPass)
 INITIALIZE_PASS_END(SinkingLegacyPass, "sink", "Code sinking", false, false)
 
 FunctionPass *llvm::createSinkingPass() { return new SinkingLegacyPass(); }