[SandboxVec][DAG] Implement UnscheduledSuccs (#112255)

This patch implements the UnscheduledSuccs counter in DGNode. It counts
the number of unscheduled successors and is used by the scheduler to
determine when a node is ready.

Author: vporpo

llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/DependencyGraph.h

@@ -96,16 +96,25 @@ class DGNode {
   // TODO: Use a PointerIntPair for SubclassID and I.
   /// For isa/dyn_cast etc.
   DGNodeID SubclassID;
+  /// The number of unscheduled successors.
+  unsigned UnscheduledSuccs = 0;
+  /// This is true if this node has been scheduled.
+  bool Scheduled = false;
 
   DGNode(Instruction *I, DGNodeID ID) : I(I), SubclassID(ID) {}
-  friend class MemDGNode; // For constructor.
+  friend class MemDGNode;       // For constructor.
+  friend class DependencyGraph; // For UnscheduledSuccs
 
 public:
   DGNode(Instruction *I) : I(I), SubclassID(DGNodeID::DGNode) {
     assert(!isMemDepNodeCandidate(I) && "Expected Non-Mem instruction, ");
   }
   DGNode(const DGNode &Other) = delete;
   virtual ~DGNode() = default;
+  /// \Returns the number of unscheduled successors.
+  unsigned getNumUnscheduledSuccs() const { return UnscheduledSuccs; }
+  /// \Returns true if this node has been scheduled.
+  bool scheduled() const { return Scheduled; }
   /// \Returns true if this is before \p Other in program order.
   bool comesBefore(const DGNode *Other) { return I->comesBefore(Other->I); }
   using iterator = PredIterator;
@@ -215,8 +224,16 @@ class MemDGNode final : public DGNode {
   MemDGNode *getPrevNode() const { return PrevMemN; }
   /// \Returns the next Mem DGNode in instruction order.
   MemDGNode *getNextNode() const { return NextMemN; }
-  /// Adds the mem dependency edge PredN->this.
-  void addMemPred(MemDGNode *PredN) { MemPreds.insert(PredN); }
+  /// Adds the mem dependency edge PredN->this. This also increments the
+  /// UnscheduledSuccs counter of the predecessor if this node has not been
+  /// scheduled.
+  void addMemPred(MemDGNode *PredN) {
+    auto Inserted = MemPreds.insert(PredN).second;
+    assert(Inserted && "PredN already exists!");
+    if (!Scheduled) {
+      ++PredN->UnscheduledSuccs;
+    }
+  }
   /// \Returns true if there is a memory dependency N->this.
   bool hasMemPred(DGNode *N) const {
     if (auto *MN = dyn_cast<MemDGNode>(N))
@@ -284,6 +301,10 @@ class DependencyGraph {
   /// \p DstN.
   void scanAndAddDeps(MemDGNode &DstN, const Interval<MemDGNode> &SrcScanRange);
 
+  /// Sets the UnscheduledSuccs of all DGNodes in \p NewInterval based on
+  /// def-use edges.
+  void setDefUseUnscheduledSuccs(const Interval<Instruction> &NewInterval);
+
   /// Create DAG nodes for instrs in \p NewInterval and update the MemNode
   /// chain.
   void createNewNodes(const Interval<Instruction> &NewInterval);

llvm/lib/Transforms/Vectorize/SandboxVectorizer/DependencyGraph.cpp

@@ -59,21 +59,17 @@ bool PredIterator::operator==(const PredIterator &Other) const {
 }
 
 #ifndef NDEBUG
-void DGNode::print(raw_ostream &OS, bool PrintDeps) const { I->dumpOS(OS); }
-void DGNode::dump() const {
-  print(dbgs());
-  dbgs() << "\n";
+void DGNode::print(raw_ostream &OS, bool PrintDeps) const {
+  OS << *I << " USuccs:" << UnscheduledSuccs << "\n";
 }
+void DGNode::dump() const { print(dbgs()); }
 void MemDGNode::print(raw_ostream &OS, bool PrintDeps) const {
-  I->dumpOS(OS);
+  DGNode::print(OS, false);
   if (PrintDeps) {
     // Print memory preds.
     static constexpr const unsigned Indent = 4;
-    for (auto *Pred : MemPreds) {
-      OS.indent(Indent) << "<-";
-      Pred->print(OS, false);
-      OS << "\n";
-    }
+    for (auto *Pred : MemPreds)
+      OS.indent(Indent) << "<-" << *Pred->getInstruction() << "\n";
   }
 }
 #endif // NDEBUG
@@ -215,6 +211,58 @@ void DependencyGraph::scanAndAddDeps(MemDGNode &DstN,
   }
 }
 
+void DependencyGraph::setDefUseUnscheduledSuccs(
+    const Interval<Instruction> &NewInterval) {
+  // +---+
+  // |   |  Def
+  // |   |   |
+  // |   |   v
+  // |   |  Use
+  // +---+
+  // Set the intra-interval counters in NewInterval.
+  for (Instruction &I : NewInterval) {
+    for (Value *Op : I.operands()) {
+      auto *OpI = dyn_cast<Instruction>(Op);
+      if (OpI == nullptr)
+        continue;
+      if (!NewInterval.contains(OpI))
+        continue;
+      auto *OpN = getNode(OpI);
+      if (OpN == nullptr)
+        continue;
+      ++OpN->UnscheduledSuccs;
+    }
+  }
+
+  // Now handle the cross-interval edges.
+  bool NewIsAbove = DAGInterval.empty() || NewInterval.comesBefore(DAGInterval);
+  const auto &TopInterval = NewIsAbove ? NewInterval : DAGInterval;
+  const auto &BotInterval = NewIsAbove ? DAGInterval : NewInterval;
+  // +---+
+  // |Top|
+  // |   |  Def
+  // +---+   |
+  // |   |   v
+  // |Bot|  Use
+  // |   |
+  // +---+
+  // Walk over all instructions in "BotInterval" and update the counter
+  // of operands that are in "TopInterval".
+  for (Instruction &BotI : BotInterval) {
+    for (Value *Op : BotI.operands()) {
+      auto *OpI = dyn_cast<Instruction>(Op);
+      if (OpI == nullptr)
+        continue;
+      if (!TopInterval.contains(OpI))
+        continue;
+      auto *OpN = getNode(OpI);
+      if (OpN == nullptr)
+        continue;
+      ++OpN->UnscheduledSuccs;
+    }
+  }
+}
+
 void DependencyGraph::createNewNodes(const Interval<Instruction> &NewInterval) {
   // Create Nodes only for the new sections of the DAG.
   DGNode *LastN = getOrCreateNode(NewInterval.top());
@@ -260,6 +308,8 @@ void DependencyGraph::createNewNodes(const Interval<Instruction> &NewInterval) {
     }
 #endif // NDEBUG
   }
+
+  setDefUseUnscheduledSuccs(NewInterval);
 }
 
 Interval<Instruction> DependencyGraph::extend(ArrayRef<Instruction *> Instrs) {

llvm/unittests/Transforms/Vectorize/SandboxVectorizer/DependencyGraphTest.cpp

@@ -249,6 +249,11 @@ define void @foo(ptr %ptr, i8 %v0, i8 %v1) {
   EXPECT_TRUE(N0->memPreds().empty());
   EXPECT_THAT(N1->memPreds(), testing::ElementsAre(N0));
   EXPECT_TRUE(N2->preds(DAG).empty());
+
+  // Check UnscheduledSuccs.
+  EXPECT_EQ(N0->getNumUnscheduledSuccs(), 1u); // N1
+  EXPECT_EQ(N1->getNumUnscheduledSuccs(), 0u);
+  EXPECT_EQ(N2->getNumUnscheduledSuccs(), 0u);
 }
 
 TEST_F(DependencyGraphTest, Preds) {
@@ -286,6 +291,14 @@ define i8 @foo(i8 %v0, i8 %v1) {
   EXPECT_THAT(StN->preds(DAG),
               testing::UnorderedElementsAre(CallN, CallN, AddN2));
   EXPECT_THAT(RetN->preds(DAG), testing::ElementsAre(AddN2));
+
+  // Check UnscheduledSuccs.
+  EXPECT_EQ(AddN0->getNumUnscheduledSuccs(), 1u); // AddN2
+  EXPECT_EQ(AddN1->getNumUnscheduledSuccs(), 2u); // AddN2, CallN
+  EXPECT_EQ(AddN2->getNumUnscheduledSuccs(), 2u); // StN, RetN
+  EXPECT_EQ(CallN->getNumUnscheduledSuccs(), 2u); // StN, StN
+  EXPECT_EQ(StN->getNumUnscheduledSuccs(), 0u);
+  EXPECT_EQ(RetN->getNumUnscheduledSuccs(), 0u);
 }
 
 TEST_F(DependencyGraphTest, MemDGNode_getPrevNode_getNextNode) {
@@ -711,6 +724,8 @@ define void @foo(ptr %ptr, i8 %v1, i8 %v2, i8 %v3, i8 %v4, i8 %v5) {
     EXPECT_EQ(DAG.getInterval().top(), S3);
     EXPECT_EQ(DAG.getInterval().bottom(), S3);
     [[maybe_unused]] auto *S3N = cast<sandboxir::MemDGNode>(DAG.getNode(S3));
+    // Check UnscheduledSuccs.
+    EXPECT_EQ(S3N->getNumUnscheduledSuccs(), 0u);
   }
   {
     // Scenario 2: Extend below
@@ -722,6 +737,10 @@ define void @foo(ptr %ptr, i8 %v1, i8 %v2, i8 %v3, i8 %v4, i8 %v5) {
     EXPECT_TRUE(S4N->hasMemPred(S3N));
     EXPECT_TRUE(S5N->hasMemPred(S4N));
     EXPECT_TRUE(S5N->hasMemPred(S3N));
+    // Check UnscheduledSuccs.
+    EXPECT_EQ(S3N->getNumUnscheduledSuccs(), 2u); // S4N, S5N
+    EXPECT_EQ(S4N->getNumUnscheduledSuccs(), 1u); // S5N
+    EXPECT_EQ(S5N->getNumUnscheduledSuccs(), 0u);
   }
   {
     // Scenario 3: Extend above
@@ -746,5 +765,12 @@ define void @foo(ptr %ptr, i8 %v1, i8 %v2, i8 %v3, i8 %v4, i8 %v5) {
     EXPECT_TRUE(S5N->hasMemPred(S3N));
     EXPECT_TRUE(S5N->hasMemPred(S2N));
     EXPECT_TRUE(S5N->hasMemPred(S1N));
+
+    // Check UnscheduledSuccs.
+    EXPECT_EQ(S1N->getNumUnscheduledSuccs(), 4u); // S2N, S3N, S4N, S5N
+    EXPECT_EQ(S2N->getNumUnscheduledSuccs(), 3u); // S3N, S4N, S5N
+    EXPECT_EQ(S3N->getNumUnscheduledSuccs(), 2u); // S4N, S5N
+    EXPECT_EQ(S4N->getNumUnscheduledSuccs(), 1u); // S5N
+    EXPECT_EQ(S5N->getNumUnscheduledSuccs(), 0u);
   }
 }