[llvm] [SandboxVec][DAG] Add MemDGNode::MemSuccs (PR #127253)

[via llvm-commits]

https://github.com/vporpo updated https://github.com/llvm/llvm-project/pull/127253

>From 4b9fbf37630e551a462e9d97feae2a0d2e1e0c98 Mon Sep 17 00:00:00 2001
From: Vasileios Porpodas <vporpodas at google.com>
Date: Thu, 30 Jan 2025 14:25:56 -0800
Subject: [PATCH] [SandboxVec][DAG] Add MemDGNode::MemSuccs

This patch adds Memory successors to the memory nodes of the DAG.
This will help maintain the memory dependencies when nodes get removed.
---
 .../SandboxVectorizer/DependencyGraph.h       | 17 ++++++++
 .../SandboxVectorizer/DependencyGraphTest.cpp | 40 ++++++++++++++++++-
 2 files changed, 56 insertions(+), 1 deletion(-)

diff --git a/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/DependencyGraph.h b/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/DependencyGraph.h
index 6852d0b6714fb..b25f96571741e 100644
--- a/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/DependencyGraph.h
+++ b/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/DependencyGraph.h
@@ -117,6 +117,7 @@ class DGNode {
   virtual ~DGNode();
   /// \Returns the number of unscheduled successors.
   unsigned getNumUnscheduledSuccs() const { return UnscheduledSuccs; }
+  // TODO: Make this private?
   void decrUnscheduledSuccs() {
     assert(UnscheduledSuccs > 0 && "Counting error!");
     --UnscheduledSuccs;
@@ -214,6 +215,8 @@ class MemDGNode final : public DGNode {
   MemDGNode *NextMemN = nullptr;
   /// Memory predecessors.
   DenseSet<MemDGNode *> MemPreds;
+  /// Memory successors.
+  DenseSet<MemDGNode *> MemSuccs;
   friend class PredIterator; // For MemPreds.
   /// Creates both edges: this<->N.
   void setNextNode(MemDGNode *N) {
@@ -265,10 +268,20 @@ class MemDGNode final : public DGNode {
     [[maybe_unused]] auto Inserted = MemPreds.insert(PredN).second;
     assert(Inserted && "PredN already exists!");
     assert(PredN != this && "Trying to add a dependency to self!");
+    PredN->MemSuccs.insert(this);
     if (!Scheduled) {
       ++PredN->UnscheduledSuccs;
     }
   }
+  /// Removes the memory dependency PredN->this. This also updates the
+  /// UnscheduledSuccs counter of PredN if this node has not been scheduled.
+  void removeMemPred(MemDGNode *PredN) {
+    MemPreds.erase(PredN);
+    PredN->MemSuccs.erase(this);
+    if (!Scheduled) {
+      PredN->decrUnscheduledSuccs();
+    }
+  }
   /// \Returns true if there is a memory dependency N->this.
   bool hasMemPred(DGNode *N) const {
     if (auto *MN = dyn_cast<MemDGNode>(N))
@@ -279,6 +292,10 @@ class MemDGNode final : public DGNode {
   iterator_range<DenseSet<MemDGNode *>::const_iterator> memPreds() const {
     return make_range(MemPreds.begin(), MemPreds.end());
   }
+  /// \Returns all memory dependency successors.
+  iterator_range<DenseSet<MemDGNode *>::const_iterator> memSuccs() const {
+    return make_range(MemSuccs.begin(), MemSuccs.end());
+  }
 #ifndef NDEBUG
   virtual void print(raw_ostream &OS, bool PrintDeps = true) const override;
 #endif // NDEBUG
diff --git a/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/DependencyGraphTest.cpp b/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/DependencyGraphTest.cpp
index d81932dca4989..bb809bf33420e 100644
--- a/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/DependencyGraphTest.cpp
+++ b/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/DependencyGraphTest.cpp
@@ -250,6 +250,9 @@ define void @foo(ptr %ptr, i8 %v0, i8 %v1) {
   EXPECT_THAT(N1->memPreds(), testing::ElementsAre(N0));
   EXPECT_TRUE(N2->preds(DAG).empty());
 
+  // Check memSuccs().
+  EXPECT_THAT(N0->memSuccs(), testing::ElementsAre(N1));
+
   // Check UnscheduledSuccs.
   EXPECT_EQ(N0->getNumUnscheduledSuccs(), 1u); // N1
   EXPECT_EQ(N1->getNumUnscheduledSuccs(), 0u);
@@ -268,6 +271,41 @@ define void @foo(ptr %ptr, i8 %v0, i8 %v1) {
   EXPECT_TRUE(N0->scheduled());
 }
 
+TEST_F(DependencyGraphTest, AddRemoveMemPred) {
+  parseIR(C, R"IR(
+define void @foo(ptr %ptr, i8 %v0, i8 %v1) {
+  store i8 %v0, ptr %ptr
+  store i8 %v1, ptr %ptr
+  ret void
+}
+)IR");
+  llvm::Function *LLVMF = &*M->getFunction("foo");
+  sandboxir::Context Ctx(C);
+  auto *F = Ctx.createFunction(LLVMF);
+  auto *BB = &*F->begin();
+  auto It = BB->begin();
+  auto *S0 = cast<sandboxir::StoreInst>(&*It++);
+  auto *S1 = cast<sandboxir::StoreInst>(&*It++);
+
+  sandboxir::DependencyGraph DAG(getAA(*LLVMF), Ctx);
+  DAG.extend({&*BB->begin(), BB->getTerminator()});
+  auto *N0 = cast<sandboxir::MemDGNode>(DAG.getNode(S0));
+  auto *N1 = cast<sandboxir::MemDGNode>(DAG.getNode(S1));
+
+  // Check removeMemPred().
+  EXPECT_FALSE(N0->memSuccs().empty());
+  EXPECT_EQ(N0->getNumUnscheduledSuccs(), 1u);
+  N1->removeMemPred(N0);
+  EXPECT_TRUE(N1->memPreds().empty());
+  EXPECT_EQ(N0->getNumUnscheduledSuccs(), 0u);
+
+  // Check addMemPred().
+  N1->addMemPred(N0);
+  EXPECT_THAT(N1->memPreds(), testing::UnorderedElementsAre(N0));
+  EXPECT_THAT(N0->memSuccs(), testing::UnorderedElementsAre(N1));
+  EXPECT_THAT(N0->getNumUnscheduledSuccs(), 1u);
+}
+
 TEST_F(DependencyGraphTest, Preds) {
   parseIR(C, R"IR(
 declare ptr @bar(i8)
@@ -533,7 +571,7 @@ define void @foo(ptr noalias %ptr0, ptr noalias %ptr1) {
   EXPECT_TRUE(RetN->preds(DAG).empty());
 }
 
-TEST_F(DependencyGraphTest, VolatileSotres) {
+TEST_F(DependencyGraphTest, VolatileStores) {
   parseIR(C, R"IR(
 define void @foo(ptr noalias %ptr0, ptr noalias %ptr1, i8 %v) {
   store volatile i8 %v, ptr %ptr0