1 files changed, 58 insertions, 14 deletions

diff --git a/include/llvm/CodeGen/ScheduleDFS.h b/include/llvm/CodeGen/ScheduleDFS.h
index 1259c78b95..73ce99f471 100644
--- a/include/llvm/CodeGen/ScheduleDFS.h
+++ b/include/llvm/CodeGen/ScheduleDFS.h
@@ -27,6 +27,9 @@ class SUnit;
 
 /// \brief Represent the ILP of the subDAG rooted at a DAG node.
 ///
+/// ILPValues summarize the DAG subtree rooted at each node. ILPValues are
+/// valid for all nodes regardless of their subtree membership.
+///
 /// When computed using bottom-up DFS, this metric assumes that the DAG is a
 /// forest of trees with roots at the bottom of the schedule branching upward.
 struct ILPValue {
@@ -62,19 +65,30 @@ struct ILPValue {
 };
 
 /// \brief Compute the values of each DAG node for various metrics during DFS.
-///
-/// ILPValues summarize the DAG subtree rooted at each node up to
-/// SubtreeLimit. ILPValues are also valid for interior nodes of a subtree, not
-/// just the root.
 class SchedDFSResult {
   friend class SchedDFSImpl;
 
+  static const unsigned InvalidSubtreeID = ~0u;
+
   /// \brief Per-SUnit data computed during DFS for various metrics.
+  ///
+  /// A node's SubtreeID is set to itself when it is visited to indicate that it
+  /// is the root of a subtree. Later it is set to its parent to indicate an
+  /// interior node. Finally, it is set to a representative subtree ID during
+  /// finalization.
   struct NodeData {
     unsigned InstrCount;
     unsigned SubtreeID;
 
-    NodeData(): InstrCount(0), SubtreeID(0) {}
+    NodeData(): InstrCount(0), SubtreeID(InvalidSubtreeID) {}
+  };
+
+  /// \brief Per-Subtree data computed during DFS.
+  struct TreeData {
+    unsigned ParentTreeID;
+    unsigned SubInstrCount;
+
+    TreeData(): ParentTreeID(InvalidSubtreeID), SubInstrCount(0) {}
   };
 
   /// \brief Record a connection between subtrees and the connection level.
@@ -88,7 +102,10 @@ class SchedDFSResult {
   bool IsBottomUp;
   unsigned SubtreeLimit;
   /// DFS results for each SUnit in this DAG.
-  std::vector<NodeData> DFSData;
+  std::vector<NodeData> DFSNodeData;
+
+  // Store per-tree data indexed on tree ID,
+  SmallVector<TreeData, 16> DFSTreeData;
 
   // For each subtree discovered during DFS, record its connections to other
   // subtrees.
@@ -102,41 +119,68 @@ public:
   SchedDFSResult(bool IsBU, unsigned lim)
     : IsBottomUp(IsBU), SubtreeLimit(lim) {}
 
+  /// \brief Get the node cutoff before subtrees are considered significant.
+  unsigned getSubtreeLimit() const { return SubtreeLimit; }
+
+  /// \brief Return true if this DFSResult is uninitialized.
+  ///
+  /// resize() initializes DFSResult, while compute() populates it.
+  bool empty() const { return DFSNodeData.empty(); }
+
   /// \brief Clear the results.
   void clear() {
-    DFSData.clear();
+    DFSNodeData.clear();
+    DFSTreeData.clear();
     SubtreeConnections.clear();
     SubtreeConnectLevels.clear();
   }
 
   /// \brief Initialize the result data with the size of the DAG.
   void resize(unsigned NumSUnits) {
-    DFSData.resize(NumSUnits);
+    DFSNodeData.resize(NumSUnits);
   }
 
   /// \brief Compute various metrics for the DAG with given roots.
-  void compute(ArrayRef<SUnit *> Roots);
+  void compute(ArrayRef<SUnit> SUnits);
+
+  /// \brief Get the number of instructions in the given subtree and its
+  /// children.
+  unsigned getNumInstrs(const SUnit *SU) const {
+    return DFSNodeData[SU->NodeNum].InstrCount;
+  }
+
+  /// \brief Get the number of instructions in the given subtree not including
+  /// children.
+  unsigned getNumSubInstrs(unsigned SubtreeID) const {
+    return DFSTreeData[SubtreeID].SubInstrCount;
+  }
 
   /// \brief Get the ILP value for a DAG node.
   ///
   /// A leaf node has an ILP of 1/1.
-  ILPValue getILP(const SUnit *SU) {
-    return ILPValue(DFSData[SU->NodeNum].InstrCount, 1 + SU->getDepth());
+  ILPValue getILP(const SUnit *SU) const {
+    return ILPValue(DFSNodeData[SU->NodeNum].InstrCount, 1 + SU->getDepth());
   }
 
   /// \brief The number of subtrees detected in this DAG.
   unsigned getNumSubtrees() const { return SubtreeConnectLevels.size(); }
 
   /// \brief Get the ID of the subtree the given DAG node belongs to.
-  unsigned getSubtreeID(const SUnit *SU) {
-    return DFSData[SU->NodeNum].SubtreeID;
+  ///
+  /// For convenience, if DFSResults have not been computed yet, give everything
+  /// tree ID 0.
+  unsigned getSubtreeID(const SUnit *SU) const {
+    if (empty())
+      return 0;
+    assert(SU->NodeNum < DFSNodeData.size() &&  "New Node");
+    return DFSNodeData[SU->NodeNum].SubtreeID;
   }
 
   /// \brief Get the connection level of a subtree.
   ///
   /// For bottom-up trees, the connection level is the latency depth (in cycles)
   /// of the deepest connection to another subtree.
-  unsigned getSubtreeLevel(unsigned SubtreeID) {
+  unsigned getSubtreeLevel(unsigned SubtreeID) const {
     return SubtreeConnectLevels[SubtreeID];
   }