[profi][NFC] Get rid of afdo_detail::TypeMap

[Amir]

Summary

Parametrize SampleProfileInference and SampleProfileLoaderBaseImpl by function
type (Function/MachineFunction) instead of block type
(BasicBlock/MachineBasicBlock). Move out specializations to appropriate
locations.

This change makes it possible to use GraphTraits instead of a custom TypeMap and
make SampleProfileInference not dependent on LLVM types, paving the way for
generalizing SampleProfileInference interfaces to BOLT IR types
(BinaryFunction/BinaryBasicBlock) in stale profile matching (D144500).

[hoy]

llvm/include/llvm/Transforms/Utils/SampleProfileInference.h

119

nit: rename BT to FT, which stands for function type

llvm/include/llvm/Transforms/Utils/SampleProfileInference.h

 
 #ifndef LLVM_TRANSFORMS_UTILS_SAMPLEPROFILEINFERENCE_H
 #define LLVM_TRANSFORMS_UTILS_SAMPLEPROFILEINFERENCE_H
 
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/DepthFirstIterator.h"
 #include "llvm/ADT/SmallVector.h"
 
-#include "llvm/IR/BasicBlock.h"
-#include "llvm/IR/Instruction.h"
-#include "llvm/IR/Instructions.h"
-
 namespace llvm {
 
-class Function;
-class MachineBasicBlock;
-class MachineFunction;
-
-namespace afdo_detail {
-
-template <class BlockT> struct TypeMap {};
-template <> struct TypeMap<BasicBlock> {
-  using BasicBlockT = BasicBlock;
-  using FunctionT = Function;
-};
-template <> struct TypeMap<MachineBasicBlock> {
-  using BasicBlockT = MachineBasicBlock;
-  using FunctionT = MachineFunction;
-};
-
-} // end namespace afdo_detail
-
 struct FlowJump;
 
 /// A wrapper of a binary basic block.
 struct FlowBlock {
   uint64_t Index;
   uint64_t Weight{0};
   bool HasUnknownWeight{true};
   bool IsUnlikely{false};
struct ProfiParams {
   /// The cost of taking an unlikely block/jump.
   const int64_t CostUnlikely = ((int64_t)1) << 30;
 };
 
 void applyFlowInference(const ProfiParams &Params, FlowFunction &Func);
 void applyFlowInference(FlowFunction &Func);
 
 /// Sample profile inference pass.
-template <typename BT> class SampleProfileInference {
+template <typename FT> class SampleProfileInference {
 public:
-  using BasicBlockT = typename afdo_detail::TypeMap<BT>::BasicBlockT;
-  using FunctionT = typename afdo_detail::TypeMap<BT>::FunctionT;
+  using NodeRef = typename GraphTraits<FT *>::NodeRef;
+  using BasicBlockT = typename std::remove_pointer<NodeRef>::type;
+  using FunctionT = FT;
   using Edge = std::pair<const BasicBlockT *, const BasicBlockT *>;
   using BlockWeightMap = DenseMap<const BasicBlockT *, uint64_t>;
   using EdgeWeightMap = DenseMap<Edge, uint64_t>;
   using BlockEdgeMap =
       DenseMap<const BasicBlockT *, SmallVector<const BasicBlockT *, 8>>;
 
   SampleProfileInference(FunctionT &F, BlockEdgeMap &Successors,
                          BlockWeightMap &SampleBlockWeights)
void SampleProfileInference<BT>::initFunction(
   }
 }
 
 template <typename BT>
 inline void SampleProfileInference<BT>::findUnlikelyJumps(
     const std::vector<const BasicBlockT *> &BasicBlocks,
     BlockEdgeMap &Successors, FlowFunction &Func) {}
 
-template <>
-inline void SampleProfileInference<BasicBlock>::findUnlikelyJumps(
-    const std::vector<const BasicBlockT *> &BasicBlocks,
-    BlockEdgeMap &Successors, FlowFunction &Func) {
-  for (auto &Jump : Func.Jumps) {
-    const auto *BB = BasicBlocks[Jump.Source];
-    const auto *Succ = BasicBlocks[Jump.Target];
-    const Instruction *TI = BB->getTerminator();
-    // Check if a block ends with InvokeInst and mark non-taken branch unlikely.
-    // In that case block Succ should be a landing pad
-    if (Successors[BB].size() == 2 && Successors[BB].back() == Succ) {
-      if (isa<InvokeInst>(TI)) {
-        Jump.IsUnlikely = true;
-      }
-    }
-    const Instruction *SuccTI = Succ->getTerminator();
-    // Check if the target block contains UnreachableInst and mark it unlikely
-    if (SuccTI->getNumSuccessors() == 0) {
-      if (isa<UnreachableInst>(SuccTI)) {
-        Jump.IsUnlikely = true;
-      }
-    }
-  }
-}
-
 template <typename BT>
 inline bool SampleProfileInference<BT>::isExit(const BasicBlockT *BB) {
   return BB->succ_empty();
 }
 
-template <>
-inline bool SampleProfileInference<BasicBlock>::isExit(const BasicBlock *BB) {
-  return succ_empty(BB);
-}
-
 } // end namespace llvm
 #endif // LLVM_TRANSFORMS_UTILS_SAMPLEPROFILEINFERENCE_H

llvm/include/llvm/Transforms/Utils/SampleProfileLoaderBaseImpl.h

bool profileIsValid(const Function &F, const FunctionSamples &Samples) const {
     return true;
   }
 };
 
 
 
 extern cl::opt<bool> SampleProfileUseProfi;
 
-template <typename BT> class SampleProfileLoaderBaseImpl {
+template <typename FT> class SampleProfileLoaderBaseImpl {
 public:
   SampleProfileLoaderBaseImpl(std::string Name, std::string RemapName,
                               IntrusiveRefCntPtr<vfs::FileSystem> FS)
       : Filename(Name), RemappingFilename(RemapName), FS(std::move(FS)) {}
   void dump() { Reader->dump(); }
 
+  using NodeRef = typename GraphTraits<FT *>::NodeRef;
+  using BT = typename std::remove_pointer<NodeRef>::type;
   using InstructionT = typename afdo_detail::IRTraits<BT>::InstructionT;
   using BasicBlockT = typename afdo_detail::IRTraits<BT>::BasicBlockT;
   using BlockFrequencyInfoT =
       typename afdo_detail::IRTraits<BT>::BlockFrequencyInfoT;
   using FunctionT = typename afdo_detail::IRTraits<BT>::FunctionT;
   using LoopT = typename afdo_detail::IRTraits<BT>::LoopT;
   using LoopInfoPtrT = typename afdo_detail::IRTraits<BT>::LoopInfoPtrT;
   using DominatorTreePtrT =
if (SampleProfileUseProfi) {
     // obviously wrong.
     Changed = true;
     while (Changed && I++ < SampleProfileMaxPropagateIterations) {
       Changed = propagateThroughEdges(F, true);
     }
   }
 }
 
-template <typename BT>
-void SampleProfileLoaderBaseImpl<BT>::applyProfi(
+template <typename FT>
+void SampleProfileLoaderBaseImpl<FT>::applyProfi(
     FunctionT &F, BlockEdgeMap &Successors, BlockWeightMap &SampleBlockWeights,
     BlockWeightMap &BlockWeights, EdgeWeightMap &EdgeWeights) {
-  auto Infer = SampleProfileInference<BT>(F, Successors, SampleBlockWeights);
+  auto Infer = SampleProfileInference<FT>(F, Successors, SampleBlockWeights);
   Infer.apply(BlockWeights, EdgeWeights);
 }
 
 /// Generate branch weight metadata for all branches in \p F.
 ///
 /// Branch weights are computed out of instruction samples using a
 /// propagation heuristic. Propagation proceeds in 3 phases:
 ///
unsigned SampleProfileLoaderBaseImpl<BT>::getFunctionLoc(FunctionT &F) {
   // about the missed opportunity.
   Func.getContext().diagnose(DiagnosticInfoSampleProfile(
       "No debug information found in function " + Func.getName() +
           ": Function profile not used",
       DS_Warning));
   return 0;
 }
 
-template <typename BT>
-void SampleProfileLoaderBaseImpl<BT>::computeDominanceAndLoopInfo(
-    FunctionT &F) {
-  DT.reset(new DominatorTree);
-  DT->recalculate(F);
-
-  PDT.reset(new PostDominatorTree(F));
-
-  LI.reset(new LoopInfo);
-  LI->analyze(*DT);
-}
-
 #undef DEBUG_TYPE
 
 } // namespace llvm
 #endif // LLVM_TRANSFORMS_UTILS_SAMPLEPROFILELOADERBASEIMPL_H

llvm/lib/CodeGen/MIRSampleProfile.cpp

template <> struct IRTraits<MachineBasicBlock> {
   }
   static SuccRangeT getSuccessors(MachineBasicBlock *BB) {
     return BB->successors();
   }
 };
 } // namespace afdo_detail
 
 class MIRProfileLoader final
-    : public SampleProfileLoaderBaseImpl<MachineBasicBlock> {
+    : public SampleProfileLoaderBaseImpl<MachineFunction> {
 public:
   void setInitVals(MachineDominatorTree *MDT, MachinePostDominatorTree *MPDT,
                    MachineLoopInfo *MLI, MachineBlockFrequencyInfo *MBFI,
                    MachineOptimizationRemarkEmitter *MORE) {
     DT = MDT;
     PDT = MPDT;
     LI = MLI;
     BFI = MBFI;
if (FunctionSamples::ProfileIsProbeBased)
       return getProbeWeight(MI);
     if (ImprovedFSDiscriminator && MI.isMetaInstruction())
       return std::error_code();
     return getInstWeightImpl(MI);
   }
 };
 
 template <>
-void SampleProfileLoaderBaseImpl<
-    MachineBasicBlock>::computeDominanceAndLoopInfo(MachineFunction &F) {}
+void SampleProfileLoaderBaseImpl<MachineFunction>::computeDominanceAndLoopInfo(
+    MachineFunction &F) {}
 
 void MIRProfileLoader::setBranchProbs(MachineFunction &F) {
   LLVM_DEBUG(dbgs() << "\nPropagation complete. Setting branch probs\n");
   for (auto &BI : F) {
     MachineBasicBlock *BB = &BI;
     if (BB->succ_size() < 2)
       continue;
     const MachineBasicBlock *EC = EquivalenceClass[BB];

llvm/lib/Transforms/IPO/SampleProfile.cpp

void runStaleProfileMatching(
       LocToLocMap &IRToProfileLocationMap);
 };
 
 /// Sample profile pass.
 ///
 /// This pass reads profile data from the file specified by
 /// -sample-profile-file and annotates every affected function with the
 /// profile information found in that file.
-class SampleProfileLoader final
-    : public SampleProfileLoaderBaseImpl<BasicBlock> {
+class SampleProfileLoader final : public SampleProfileLoaderBaseImpl<Function> {
 public:
   SampleProfileLoader(
       StringRef Name, StringRef RemapName, ThinOrFullLTOPhase LTOPhase,
       IntrusiveRefCntPtr<vfs::FileSystem> FS,
       std::function<AssumptionCache &(Function &)> GetAssumptionCache,
       std::function<TargetTransformInfo &(Function &)> GetTargetTransformInfo,
       std::function<const TargetLibraryInfo &(Function &)> GetTLI)
       : SampleProfileLoaderBaseImpl(std::string(Name), std::string(RemapName),

 
 private:
   const char *getAnnotatedRemarkPassName() const {
     return AnnotatedPassName.c_str();
   }
 };
 } // end anonymous namespace
 
+namespace llvm {
+template <>
+inline bool SampleProfileInference<Function>::isExit(const BasicBlock *BB) {
+  return succ_empty(BB);
+}
+
+template <>
+inline void SampleProfileInference<Function>::findUnlikelyJumps(
+    const std::vector<const BasicBlockT *> &BasicBlocks,
+    BlockEdgeMap &Successors, FlowFunction &Func) {
+  for (auto &Jump : Func.Jumps) {
+    const auto *BB = BasicBlocks[Jump.Source];
+    const auto *Succ = BasicBlocks[Jump.Target];
+    const Instruction *TI = BB->getTerminator();
+    // Check if a block ends with InvokeInst and mark non-taken branch unlikely.
+    // In that case block Succ should be a landing pad
+    if (Successors[BB].size() == 2 && Successors[BB].back() == Succ) {
+      if (isa<InvokeInst>(TI)) {
+        Jump.IsUnlikely = true;
+      }
+    }
+    const Instruction *SuccTI = Succ->getTerminator();
+    // Check if the target block contains UnreachableInst and mark it unlikely
+    if (SuccTI->getNumSuccessors() == 0) {
+      if (isa<UnreachableInst>(SuccTI)) {
+        Jump.IsUnlikely = true;
+      }
+    }
+  }
+}
+
+template <>
+void SampleProfileLoaderBaseImpl<Function>::computeDominanceAndLoopInfo(
+    Function &F) {
+  DT.reset(new DominatorTree);
+  DT->recalculate(F);
+
+  PDT.reset(new PostDominatorTree(F));
+
+  LI.reset(new LoopInfo);
+  LI->analyze(*DT);
+}
+} // namespace llvm
+
 ErrorOr<uint64_t> SampleProfileLoader::getInstWeight(const Instruction &Inst) {
   if (FunctionSamples::ProfileIsProbeBased)
     return getProbeWeight(Inst);
 
   const DebugLoc &DLoc = Inst.getDebugLoc();
   if (!DLoc)
     return std::error_code();