[llvm] [SandboxVec][DAG] Implement DGNode::isMem() (PR #109504)

[via llvm-commits]

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

>From a753ab273d114f90f4c01356643ab93292a1a7ee Mon Sep 17 00:00:00 2001
From: Vasileios Porpodas <vporpodas at google.com>
Date: Wed, 18 Sep 2024 13:21:53 -0700
Subject: [PATCH] [SandboxVec][DAG] Implement DGNode::isMem()

DGNode::isMem() returns true if the node is a memory dependency candidate.
---
 .../SandboxVectorizer/DependencyGraph.h       | 12 ++++-
 .../SandboxVectorizer/DependencyGraphTest.cpp | 49 +++++++++++++++++++
 2 files changed, 60 insertions(+), 1 deletion(-)

diff --git a/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/DependencyGraph.h b/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/DependencyGraph.h
index 3aa44781a2dd5c..75b2073d0557c5 100644
--- a/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/DependencyGraph.h
+++ b/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/DependencyGraph.h
@@ -35,9 +35,16 @@ class DGNode {
   Instruction *I;
   /// Memory predecessors.
   DenseSet<DGNode *> MemPreds;
+  /// This is true if this may read/write memory, or if it has some ordering
+  /// constraings, like with stacksave/stackrestore and alloca/inalloca.
+  bool IsMem;
 
 public:
-  DGNode(Instruction *I) : I(I) {}
+  DGNode(Instruction *I) : I(I) {
+    IsMem = I->isMemDepCandidate() ||
+            (isa<AllocaInst>(I) && cast<AllocaInst>(I)->isUsedWithInAlloca()) ||
+            I->isStackSaveOrRestoreIntrinsic();
+  }
   Instruction *getInstruction() const { return I; }
   void addMemPred(DGNode *PredN) { MemPreds.insert(PredN); }
   /// \Returns all memory dependency predecessors.
@@ -46,6 +53,9 @@ class DGNode {
   }
   /// \Returns true if there is a memory dependency N->this.
   bool hasMemPred(DGNode *N) const { return MemPreds.count(N); }
+  /// \Returns true if this may read/write memory, or if it has some ordering
+  /// constraings, like with stacksave/stackrestore and alloca/inalloca.
+  bool isMem() const { return IsMem; }
 #ifndef NDEBUG
   void print(raw_ostream &OS, bool PrintDeps = true) const;
   friend raw_ostream &operator<<(DGNode &N, raw_ostream &OS) {
diff --git a/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/DependencyGraphTest.cpp b/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/DependencyGraphTest.cpp
index c62b2a4e508c06..f6bfd097f20a4e 100644
--- a/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/DependencyGraphTest.cpp
+++ b/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/DependencyGraphTest.cpp
@@ -27,6 +27,55 @@ struct DependencyGraphTest : public testing::Test {
   }
 };
 
+TEST_F(DependencyGraphTest, DGNode_IsMem) {
+  parseIR(C, R"IR(
+declare void @llvm.sideeffect()
+declare void @llvm.pseudoprobe(i64, i64, i32, i64)
+declare void @llvm.fake.use(...)
+declare void @bar()
+define void @foo(i8 %v1, ptr %ptr) {
+  store i8 %v1, ptr %ptr
+  %ld0 = load i8, ptr %ptr
+  %add = add i8 %v1, %v1
+  %stacksave = call ptr @llvm.stacksave()
+  call void @llvm.stackrestore(ptr %stacksave)
+  call void @llvm.sideeffect()
+  call void @llvm.pseudoprobe(i64 42, i64 1, i32 0, i64 -1)
+  call void @llvm.fake.use(ptr %ptr)
+  call void @bar()
+  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 *Store = cast<sandboxir::StoreInst>(&*It++);
+  auto *Load = cast<sandboxir::LoadInst>(&*It++);
+  auto *Add = cast<sandboxir::BinaryOperator>(&*It++);
+  auto *StackSave = cast<sandboxir::CallInst>(&*It++);
+  auto *StackRestore = cast<sandboxir::CallInst>(&*It++);
+  auto *SideEffect = cast<sandboxir::CallInst>(&*It++);
+  auto *PseudoProbe = cast<sandboxir::CallInst>(&*It++);
+  auto *FakeUse = cast<sandboxir::CallInst>(&*It++);
+  auto *Call = cast<sandboxir::CallInst>(&*It++);
+  auto *Ret = cast<sandboxir::ReturnInst>(&*It++);
+
+  sandboxir::DependencyGraph DAG;
+  DAG.extend({&*BB->begin(), BB->getTerminator()});
+  EXPECT_TRUE(DAG.getNode(Store)->isMem());
+  EXPECT_TRUE(DAG.getNode(Load)->isMem());
+  EXPECT_FALSE(DAG.getNode(Add)->isMem());
+  EXPECT_TRUE(DAG.getNode(StackSave)->isMem());
+  EXPECT_TRUE(DAG.getNode(StackRestore)->isMem());
+  EXPECT_FALSE(DAG.getNode(SideEffect)->isMem());
+  EXPECT_FALSE(DAG.getNode(PseudoProbe)->isMem());
+  EXPECT_TRUE(DAG.getNode(FakeUse)->isMem());
+  EXPECT_TRUE(DAG.getNode(Call)->isMem());
+  EXPECT_FALSE(DAG.getNode(Ret)->isMem());
+}
+
 TEST_F(DependencyGraphTest, Basic) {
   parseIR(C, R"IR(
 define void @foo(ptr %ptr, i8 %v0, i8 %v1) {