[SandboxVec][DAG] Implement extend(ArrayRef)

[vporpo]

This builds the DAG from an ArrayRef of Instructions.

[llvmbot]

@llvm/pr-subscribers-llvm-transforms

Author: vporpo (vporpo)

Changes

This builds the DAG from an ArrayRef of Instructions.

---

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

3 Files Affected:

• (modified) llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/DependencyGraph.h (+4-3)
• (modified) llvm/lib/Transforms/Vectorize/SandboxVectorizer/DependencyGraph.cpp (+12-6)
• (modified) llvm/unittests/Transforms/Vectorize/SandboxVectorizer/DependencyGraphTest.cpp (+4-1)

diff --git a/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/DependencyGraph.h b/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/DependencyGraph.h
index 8a2021a5e6ba60..3aa44781a2dd5c 100644
--- a/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/DependencyGraph.h
+++ b/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/DependencyGraph.h
@@ -25,6 +25,7 @@
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/iterator_range.h"
 #include "llvm/SandboxIR/SandboxIR.h"
+#include "llvm/Transforms/Vectorize/SandboxVectorizer/InstrInterval.h"
 
 namespace llvm::sandboxir {
 
@@ -72,9 +73,9 @@ class DependencyGraph {
       It->second = std::make_unique<DGNode>(I);
     return It->second.get();
   }
-  // TODO: extend() should work with intervals not the whole BB.
-  /// Build the dependency graph for \p BB.
-  void extend(BasicBlock *BB);
+  /// Build/extend the dependency graph such that it includes \p Instrs. Returns
+  /// the interval spanning \p Instrs.
+  InstrInterval extend(ArrayRef<Instruction *> Instrs);
 #ifndef NDEBUG
   void print(raw_ostream &OS) const;
   LLVM_DUMP_METHOD void dump() const;
diff --git a/llvm/lib/Transforms/Vectorize/SandboxVectorizer/DependencyGraph.cpp b/llvm/lib/Transforms/Vectorize/SandboxVectorizer/DependencyGraph.cpp
index 41e50953a4ec8a..fcac0c84f63a6c 100644
--- a/llvm/lib/Transforms/Vectorize/SandboxVectorizer/DependencyGraph.cpp
+++ b/llvm/lib/Transforms/Vectorize/SandboxVectorizer/DependencyGraph.cpp
@@ -7,6 +7,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "llvm/Transforms/Vectorize/SandboxVectorizer/DependencyGraph.h"
+#include "llvm/ADT/ArrayRef.h"
 
 using namespace llvm::sandboxir;
 
@@ -30,16 +31,21 @@ void DGNode::dump() const {
 }
 #endif // NDEBUG
 
-void DependencyGraph::extend(BasicBlock *BB) {
-  if (BB->empty())
-    return;
+InstrInterval DependencyGraph::extend(ArrayRef<Instruction *> Instrs) {
+  if (Instrs.empty())
+    return {};
   // TODO: For now create a chain of dependencies.
-  DGNode *LastN = getOrCreateNode(&*BB->begin());
-  for (auto &I : drop_begin(*BB)) {
-    auto *N = getOrCreateNode(&I);
+  InstrInterval Interval(Instrs);
+  auto *TopI = Interval.top();
+  auto *BotI = Interval.bottom();
+  DGNode *LastN = getOrCreateNode(BotI);
+  for (Instruction *I = BotI->getPrevNode(), *E = TopI->getPrevNode(); I != E;
+       I = I->getPrevNode()) {
+    auto *N = getOrCreateNode(I);
     N->addMemPred(LastN);
     LastN = N;
   }
+  return Interval;
 }
 
 #ifndef NDEBUG
diff --git a/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/DependencyGraphTest.cpp b/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/DependencyGraphTest.cpp
index dc85a22f7f4832..c62b2a4e508c06 100644
--- a/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/DependencyGraphTest.cpp
+++ b/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/DependencyGraphTest.cpp
@@ -44,7 +44,10 @@ define void @foo(ptr %ptr, i8 %v0, i8 %v1) {
   auto *S1 = cast<sandboxir::StoreInst>(&*It++);
   auto *Ret = cast<sandboxir::ReturnInst>(&*It++);
   sandboxir::DependencyGraph DAG;
-  DAG.extend(BB);
+  auto Span = DAG.extend({&*BB->begin(), BB->getTerminator()});
+  // Check extend().
+  EXPECT_EQ(Span.top(), &*BB->begin());
+  EXPECT_EQ(Span.bottom(), BB->getTerminator());
 
   sandboxir::DGNode *N0 = DAG.getNode(S0);
   sandboxir::DGNode *N1 = DAG.getNode(S1);