[SandboxVec][DAG] Implement DGNode::isMem()

[vporpo]

DGNode::isMem() returns true if the node is a memory dependency candidate.

[llvmbot]

@llvm/pr-subscribers-llvm-transforms

Author: vporpo (vporpo)

Changes

DGNode::isMem() returns true if the node is a memory dependency candidate.

---

Full diff: https://github.com/llvm/llvm-project/pull/109504.diff

2 Files Affected:

• (modified) llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/DependencyGraph.h (+11-1)
• (modified) llvm/unittests/Transforms/Vectorize/SandboxVectorizer/DependencyGraphTest.cpp (+49)

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..5db6dc037f0d98 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::StoreInst>(&*It++);
+  auto *Add = cast<sandboxir::StoreInst>(&*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) {

[tschuett]

typo: constraings

[vporpo]

good catch :)

[tschuett]

same here: constraings

[tschuett]

isa + cast anti-pattern -> dyn_cast