[llvm] [SandboxVec][DAG] Implement extend(ArrayRef) (PR #109493)

[via llvm-commits]

https://github.com/vporpo created https://github.com/llvm/llvm-project/pull/109493

This builds the DAG from an ArrayRef of Instructions.

>From f97959836a7be2d1c48a34e821c1e2506219bb4e Mon Sep 17 00:00:00 2001
From: Vasileios Porpodas <vporpodas at google.com>
Date: Tue, 17 Sep 2024 16:06:34 -0700
Subject: [PATCH] [SandboxVec][DAG] Implement extend(ArrayRef)

This builds the DAG from an ArrayRef of Instructions.
---
 .../SandboxVectorizer/DependencyGraph.h        |  7 ++++---
 .../SandboxVectorizer/DependencyGraph.cpp      | 18 ++++++++++++------
 .../SandboxVectorizer/DependencyGraphTest.cpp  |  5 ++++-
 3 files changed, 20 insertions(+), 10 deletions(-)

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