From 31b85c6ead3e0e2979272046548f3ae8a51db289 Mon Sep 17 00:00:00 2001
From: vporpo <vporpodas@google.com>
Date: Fri, 11 Oct 2024 11:51:38 -0700
Subject: [PATCH] [SandboxVec][Interval] Implement Interval::comesBefore()
 (#112026)

This patch implements `Interval::comesBefore(const Interval &Other)`
which returns true if this interval is strictly before Other in program
order. The function asserts that the intervals are disjoint.
---
 .../Vectorize/SandboxVectorizer/Interval.h    |  6 ++++++
 .../SandboxVectorizer/DependencyGraph.cpp     |  6 +-----
 .../SandboxVectorizer/IntervalTest.cpp        | 19 +++++++++++++++++++
 3 files changed, 26 insertions(+), 5 deletions(-)

diff --git a/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Interval.h b/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Interval.h
index e0c581f1d50b40..0777d9b043eb62 100644
--- a/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Interval.h
+++ b/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Interval.h
@@ -127,6 +127,12 @@ template <typename T> class Interval {
   }
   /// Inequality.
   bool operator!=(const Interval &Other) const { return !(*this == Other); }
+  /// \Returns true if this interval comes before \p Other in program order.
+  /// This expects disjoint intervals.
+  bool comesBefore(const Interval &Other) const {
+    assert(disjoint(Other) && "Expect disjoint intervals!");
+    return bottom()->comesBefore(Other.top());
+  }
   /// \Returns true if this and \p Other have nothing in common.
   bool disjoint(const Interval &Other) const {
     if (Other.empty())
diff --git a/llvm/lib/Transforms/Vectorize/SandboxVectorizer/DependencyGraph.cpp b/llvm/lib/Transforms/Vectorize/SandboxVectorizer/DependencyGraph.cpp
index db58069de47051..a8e362571837cc 100644
--- a/llvm/lib/Transforms/Vectorize/SandboxVectorizer/DependencyGraph.cpp
+++ b/llvm/lib/Transforms/Vectorize/SandboxVectorizer/DependencyGraph.cpp
@@ -231,11 +231,7 @@ void DependencyGraph::createNewNodes(const Interval<Instruction> &NewInterval) {
   }
   // Link new MemDGNode chain with the old one, if any.
   if (!DAGInterval.empty()) {
-    // TODO: Implement Interval::comesBefore() to replace this check.
-    bool NewIsAbove = NewInterval.bottom()->comesBefore(DAGInterval.top());
-    assert(
-        (NewIsAbove || DAGInterval.bottom()->comesBefore(NewInterval.top())) &&
-        "Expected NewInterval below DAGInterval.");
+    bool NewIsAbove = NewInterval.comesBefore(DAGInterval);
     const auto &TopInterval = NewIsAbove ? NewInterval : DAGInterval;
     const auto &BotInterval = NewIsAbove ? DAGInterval : NewInterval;
     MemDGNode *LinkTopN =
diff --git a/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/IntervalTest.cpp b/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/IntervalTest.cpp
index ee461e48f0dc03..b04e4fc7cffcae 100644
--- a/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/IntervalTest.cpp
+++ b/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/IntervalTest.cpp
@@ -123,6 +123,25 @@ define void @foo(i8 %v0) {
     EXPECT_FALSE(Intvl1.disjoint(Intvl3));
     EXPECT_TRUE(Intvl1.disjoint(Empty));
   }
+  {
+    // Check comesBefore().
+    sandboxir::Interval<sandboxir::Instruction> Intvl1(I0, I0);
+    sandboxir::Interval<sandboxir::Instruction> Intvl2(I2, I2);
+    EXPECT_TRUE(Intvl1.comesBefore(Intvl2));
+    EXPECT_FALSE(Intvl2.comesBefore(Intvl1));
+
+    sandboxir::Interval<sandboxir::Instruction> Intvl12(I1, I2);
+    EXPECT_TRUE(Intvl1.comesBefore(Intvl12));
+    EXPECT_FALSE(Intvl12.comesBefore(Intvl1));
+    {
+#ifndef NDEBUG
+      // Check comesBefore() with non-disjoint intervals.
+      sandboxir::Interval<sandboxir::Instruction> Intvl1(I0, I2);
+      sandboxir::Interval<sandboxir::Instruction> Intvl2(I2, I2);
+      EXPECT_DEATH(Intvl1.comesBefore(Intvl2), ".*disjoint.*");
+#endif // NDEBUG
+    }
+  }
 }
 
 // Helper function for returning a vector of instruction pointers from a range