Index: include/llvm/CodeGen/MachineBasicBlock.h =================================================================== --- include/llvm/CodeGen/MachineBasicBlock.h +++ include/llvm/CodeGen/MachineBasicBlock.h @@ -16,6 +16,7 @@ #include "llvm/ADT/GraphTraits.h" #include "llvm/CodeGen/MachineInstr.h" +#include "llvm/Support/BranchProbability.h" #include "llvm/MC/MCRegisterInfo.h" #include "llvm/Support/DataTypes.h" #include @@ -96,6 +97,14 @@ typedef std::vector::iterator weight_iterator; typedef std::vector::const_iterator const_weight_iterator; + /// Keep track of the probabilities to the successors. This vector has the + /// same order as Successors, or it is empty if we don't use it (disable + /// optimization). + std::vector Probs; + typedef std::vector::iterator probability_iterator; + typedef std::vector::const_iterator + const_probability_iterator; + /// Keep track of the physical registers that are livein of the basicblock. typedef std::vector LiveInVector; LiveInVector LiveIns; @@ -425,9 +434,25 @@ /// Note that duplicate Machine CFG edges are not allowed. void addSuccessor(MachineBasicBlock *Succ, uint32_t Weight = 0); + /// Add Succ as a successor of this MachineBasicBlock. The Predecessors list + /// of Succ is automatically updated. PROB parameter is stored in + /// Probabilities list. + /// + /// Note that duplicate Machine CFG edges are not allowed. + void addSuccessor(MachineBasicBlock *Succ, BranchProbability Prob); + /// Set successor weight of a given iterator. void setSuccWeight(succ_iterator I, uint32_t Weight); + /// Set successor probability of a given iterator. + void setSuccProbability(succ_iterator I, BranchProbability Prob); + + /// Normalize probabilities of all successors so that the sum of them becomes + /// one. + void normalizeSuccProbs() { + BranchProbability::normalizeProbabilities(Probs); + } + /// Remove successor from the successors list of this MachineBasicBlock. The /// Predecessors list of Succ is automatically updated. void removeSuccessor(MachineBasicBlock *Succ); @@ -452,6 +477,9 @@ /// Return true if any of the successors have weights attached to them. bool hasSuccessorWeights() const { return !Weights.empty(); } + /// Return true if any of the successors have probabilities attached to them. + bool hasSuccessorProbabilities() const { return !Probs.empty(); } + /// Return true if the specified MBB is a predecessor of this block. bool isPredecessor(const MachineBasicBlock *MBB) const; @@ -709,6 +737,11 @@ weight_iterator getWeightIterator(succ_iterator I); const_weight_iterator getWeightIterator(const_succ_iterator I) const; + /// Return probability iterator corresponding to the I successor iterator. + probability_iterator getProbabilityIterator(succ_iterator I); + const_probability_iterator + getProbabilityIterator(const_succ_iterator I) const; + friend class MachineBranchProbabilityInfo; friend class MIPrinter; @@ -717,6 +750,10 @@ /// MachineBranchProbabilityInfo class. uint32_t getSuccWeight(const_succ_iterator Succ) const; + /// Return probability of the edge from this block to MBB. This method should + /// NOT be called directly, but by using getEdgeProbability method from + /// MachineBranchProbabilityInfo class. + BranchProbability getSuccProbability(const_succ_iterator Succ) const; // Methods used to maintain doubly linked list of blocks... friend struct ilist_traits; Index: include/llvm/Support/BranchProbability.h =================================================================== --- include/llvm/Support/BranchProbability.h +++ include/llvm/Support/BranchProbability.h @@ -31,6 +31,7 @@ // Denominator, which is a constant value. static const uint32_t D = 1u << 31; + static const uint32_t UnknownN = UINT32_MAX; // Construct a BranchProbability with only numerator assuming the denominator // is 1<<31. For internal use only. @@ -41,9 +42,11 @@ BranchProbability(uint32_t Numerator, uint32_t Denominator); bool isZero() const { return N == 0; } + bool isUnknown() const { return N == UnknownN; } static BranchProbability getZero() { return BranchProbability(0); } static BranchProbability getOne() { return BranchProbability(D); } + static BranchProbability getUnknown() { return BranchProbability(UnknownN); } // Create a BranchProbability object with the given numerator and 1<<31 // as denominator. static BranchProbability getRaw(uint32_t N) { return BranchProbability(N); } @@ -125,6 +128,10 @@ return Prob.print(OS); } +inline BranchProbability operator/(BranchProbability LHS, uint32_t RHS) { + return BranchProbability::getRaw(LHS.getNumerator() / RHS); +} + template void BranchProbability::normalizeProbabilities(ProbabilityList &Probs) { uint64_t Sum = 0; Index: lib/CodeGen/MachineBasicBlock.cpp =================================================================== --- lib/CodeGen/MachineBasicBlock.cpp +++ lib/CodeGen/MachineBasicBlock.cpp @@ -519,6 +519,23 @@ Succ->addPredecessor(this); } +void MachineBasicBlock::addSuccessor(MachineBasicBlock *Succ, + BranchProbability Prob) { + // If we see known probability for the first time it means we actually use + // probability list, so we fill all Probs with unknown probabilities. + if (!Prob.isUnknown() && Probs.empty()) { + Probs.assign(Successors.size(), BranchProbability::getUnknown()); + } + + if (!Probs.empty()) { + assert(!Prob.isUnknown()); + Probs.push_back(Prob); + } + + Successors.push_back(Succ); + Succ->addPredecessor(this); +} + void MachineBasicBlock::removeSuccessor(MachineBasicBlock *Succ) { Succ->removePredecessor(this); succ_iterator I = std::find(Successors.begin(), Successors.end(), Succ); @@ -530,6 +547,13 @@ Weights.erase(WI); } + // If probability list is empty it means we don't use it (disabled + // optimization). + if (!Probs.empty()) { + probability_iterator WI = getProbabilityIterator(I); + Probs.erase(WI); + } + Successors.erase(I); } @@ -543,6 +567,13 @@ Weights.erase(WI); } + // If probability list is empty it means we don't use it (disabled + // optimization). + if (!Probs.empty()) { + probability_iterator WI = getProbabilityIterator(I); + Probs.erase(WI); + } + (*I)->removePredecessor(this); return Successors.erase(I); } @@ -584,6 +615,12 @@ *getWeightIterator(NewI) += *OldWI; Weights.erase(OldWI); } + // Update its probability instead of adding a duplicate edge. + if (!Probs.empty()) { + probability_iterator OldPI = getProbabilityIterator(OldI); + *getProbabilityIterator(NewI) += *OldPI; + Probs.erase(OldPI); + } Successors.erase(OldI); } @@ -1116,6 +1153,31 @@ return *getWeightIterator(Succ); } +/// Return probability of the edge from this block to MBB. If probability list +/// is empty, return a default probability which is 1/N, where N is the number +/// of successors. If the probability of the given successor is unknown, then +/// sum up all known probabilities and return the complement of the sum divided +/// by the number of unknown probabilities. +BranchProbability +MachineBasicBlock::getSuccProbability(const_succ_iterator Succ) const { + if (Probs.empty()) + return BranchProbability(1, succ_size()); + + auto Prob = *getProbabilityIterator(Succ); + if (Prob.isUnknown()) { + uint32_t UnknownProbNum = 0; + BranchProbability KnownProb; + for (auto P : Probs) { + if (P.isUnknown()) + UnknownProbNum++; + else + KnownProb += P; + } + return KnownProb.getCompl() / UnknownProbNum; + } + return Prob; +} + /// Set successor weight of a given iterator. void MachineBasicBlock::setSuccWeight(succ_iterator I, uint32_t Weight) { if (Weights.empty()) @@ -1123,6 +1185,15 @@ *getWeightIterator(I) = Weight; } +/// Set successor probability of a given iterator. +void MachineBasicBlock::setSuccProbability(succ_iterator I, + BranchProbability Prob) { + assert(!Prob.isUnknown()); + if (Probs.empty()) + return; + *getProbabilityIterator(I) = Prob; +} + /// Return wight iterator corresonding to the I successor iterator. MachineBasicBlock::weight_iterator MachineBasicBlock:: getWeightIterator(MachineBasicBlock::succ_iterator I) { @@ -1141,6 +1212,25 @@ return Weights.begin() + index; } +/// Return probability iterator corresonding to the I successor iterator. +MachineBasicBlock::probability_iterator +MachineBasicBlock::getProbabilityIterator(MachineBasicBlock::succ_iterator I) { + assert(Probs.size() == Successors.size() && "Async probability list!"); + const size_t index = std::distance(Successors.begin(), I); + assert(index < Probs.size() && "Not a current successor!"); + return Probs.begin() + index; +} + +/// Return probability iterator corresonding to the I successor iterator +MachineBasicBlock::const_probability_iterator +MachineBasicBlock::getProbabilityIterator( + MachineBasicBlock::const_succ_iterator I) const { + assert(Probs.size() == Successors.size() && "Async probability list!"); + const size_t index = std::distance(Successors.begin(), I); + assert(index < Probs.size() && "Not a current successor!"); + return Probs.begin() + index; +} + /// Return whether (physical) register "Reg" has been ined and not ed /// as of just before "MI". ///