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: lib/CodeGen/MachineBasicBlock.cpp =================================================================== --- lib/CodeGen/MachineBasicBlock.cpp +++ lib/CodeGen/MachineBasicBlock.cpp @@ -519,6 +519,20 @@ Succ->addPredecessor(this); } +void MachineBasicBlock::addSuccessor(MachineBasicBlock *Succ, + BranchProbability Prob) { + // If we see non-zero value for the first time it means we actually use + // probability list, so we fill all Probs with 0's. + if (!Prob.isZero() && Probs.empty()) + Probs.resize(Successors.size()); + + if (!Prob.isZero() || !Probs.empty()) + 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 +544,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 +564,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 +612,12 @@ *getWeightIterator(NewI) += *OldWI; Weights.erase(OldWI); } + // Update its probability instead of adding a duplicate edge. + if (!Probs.empty()) { + probability_iterator OldWI = getProbabilityIterator(OldI); + *getProbabilityIterator(NewI) += *OldWI; + Probs.erase(OldWI); + } Successors.erase(OldI); } @@ -1116,6 +1150,15 @@ return *getWeightIterator(Succ); } +/// Return probability of the edge from this block to MBB. +BranchProbability +MachineBasicBlock::getSuccProbability(const_succ_iterator Succ) const { + if (Probs.empty()) + return BranchProbability(1, succ_size()); + + return *getProbabilityIterator(Succ); +} + /// Set successor weight of a given iterator. void MachineBasicBlock::setSuccWeight(succ_iterator I, uint32_t Weight) { if (Weights.empty()) @@ -1123,6 +1166,14 @@ *getWeightIterator(I) = Weight; } +/// Set successor probability of a given iterator. +void MachineBasicBlock::setSuccProbability(succ_iterator I, + BranchProbability Prob) { + 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 +1192,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". ///