[SandboxVec][DAG] Implement extend(ArrayRef)

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;

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(TopI);
+  for (Instruction *I = TopI->getNextNode(), *E = BotI->getNextNode(); I != E;
+       I = I->getNextNode()) {
+    auto *N = getOrCreateNode(I);
     N->addMemPred(LastN);
     LastN = N;
   }
+  return Interval;
 }
 
 #ifndef NDEBUG

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);