From 36be3f2989f1b53863caabcf064bf206b87eccd0 Mon Sep 17 00:00:00 2001
From: Vasileios Porpodas <vporpodas@google.com>
Date: Wed, 13 Nov 2024 12:10:00 -0800
Subject: [PATCH] [SandboxVec][Interval] Implement Interval::notifyMoveInstr()

This patch implements the notifier for Instruction intervals.
It updates the interval's top/bottom.
---
 .../Vectorize/SandboxVectorizer/Interval.h    |  24 ++++
 .../SandboxVectorizer/IntervalTest.cpp        | 109 ++++++++++++++++++
 2 files changed, 133 insertions(+)

diff --git a/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Interval.h b/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Interval.h
index e2d0b82489ddc..922dd2c3a1f89 100644
--- a/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Interval.h
+++ b/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Interval.h
@@ -21,8 +21,10 @@
 #define LLVM_TRANSFORMS_VECTORIZE_SANDBOXVECTORIZER_INSTRINTERVAL_H
 
 #include "llvm/ADT/ArrayRef.h"
+#include "llvm/SandboxIR/Instruction.h"
 #include "llvm/Support/raw_ostream.h"
 #include <iterator>
+#include <type_traits>
 
 namespace llvm::sandboxir {
 
@@ -207,6 +209,28 @@ template <typename T> class Interval {
     return {NewTop, NewBottom};
   }
 
+  /// Update the interval when \p I is about to be moved before \p Before.
+  // SFINAE disables this for non-Instructions.
+  template <typename HelperT = T>
+  std::enable_if_t<std::is_same<HelperT, Instruction>::value, void>
+  notifyMoveInstr(HelperT *I, decltype(I->getIterator()) BeforeIt) {
+    assert(contains(I) && "Expect `I` in interval!");
+    assert(I->getIterator() != BeforeIt && "Can't move `I` before itself!");
+
+    // Nothing to do if the instruction won't move.
+    if (std::next(I->getIterator()) == BeforeIt)
+      return;
+
+    T *NewTop = Top->getIterator() == BeforeIt ? I
+                : I == Top                     ? Top->getNextNode()
+                                               : Top;
+    T *NewBottom = std::next(Bottom->getIterator()) == BeforeIt ? I
+                   : I == Bottom ? Bottom->getPrevNode()
+                                 : Bottom;
+    Top = NewTop;
+    Bottom = NewBottom;
+  }
+
 #ifndef NDEBUG
   void print(raw_ostream &OS) const {
     auto *Top = top();
diff --git a/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/IntervalTest.cpp b/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/IntervalTest.cpp
index b04e4fc7cffca..32521ed79a314 100644
--- a/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/IntervalTest.cpp
+++ b/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/IntervalTest.cpp
@@ -331,3 +331,112 @@ define void @foo(i8 %v0) {
     EXPECT_THAT(getPtrVec(SingleUnion), testing::ElementsAre(I0, I1, I2, Ret));
   }
 }
+
+TEST_F(IntervalTest, NotifyMoveInstr) {
+  parseIR(C, R"IR(
+define void @foo(i8 %v0) {
+  %I0 = add i8 %v0, %v0
+  %I1 = add i8 %v0, %v0
+  %I2 = add i8 %v0, %v0
+  ret void
+}
+)IR");
+  Function &LLVMF = *M->getFunction("foo");
+  sandboxir::Context Ctx(C);
+  auto &F = *Ctx.createFunction(&LLVMF);
+  auto *BB = &*F.begin();
+  auto It = BB->begin();
+  auto *I0 = &*It++;
+  auto *I1 = &*It++;
+  auto *I2 = &*It++;
+  auto *Ret = &*It++;
+  {
+    // Assert that we don't try to move external instr to the interval.
+    sandboxir::Interval<sandboxir::Instruction> I2Ret(I2, Ret);
+#ifndef NDEBUG
+    EXPECT_DEATH(I2Ret.notifyMoveInstr(I0, Ret->getIterator()), ".*interval.*");
+#endif // NDEBUG
+  }
+  {
+    // Assert that we don't move before self.
+    sandboxir::Interval<sandboxir::Instruction> I2Ret(I2, Ret);
+#ifndef NDEBUG
+    EXPECT_DEATH(I2Ret.notifyMoveInstr(Ret, Ret->getIterator()), ".*self.*");
+#endif // NDEBUG
+  }
+  {
+    // Single-element interval.
+    sandboxir::Interval<sandboxir::Instruction> I2I2(I2, I2);
+    I2I2.notifyMoveInstr(I2, Ret->getIterator());
+    EXPECT_EQ(I2I2.top(), I2);
+    EXPECT_EQ(I2I2.bottom(), I2);
+  }
+  {
+    // Two-element interval swap.
+    sandboxir::Interval<sandboxir::Instruction> I1I2(I1, I2);
+    I1I2.notifyMoveInstr(I2, I1->getIterator());
+    I2->moveBefore(I1);
+    EXPECT_EQ(I1I2.top(), I2);
+    EXPECT_EQ(I1I2.bottom(), I1);
+
+    I2->moveAfter(I1);
+  }
+  {
+    // Move to same position.
+    sandboxir::Interval<sandboxir::Instruction> I0Ret(I0, Ret);
+    I0Ret.notifyMoveInstr(I0, I1->getIterator());
+    I0->moveBefore(I1);
+    EXPECT_EQ(I0Ret.top(), I0);
+    EXPECT_EQ(I0Ret.bottom(), Ret);
+  }
+  {
+    // Move internal to internal.
+    sandboxir::Interval<sandboxir::Instruction> I0Ret(I0, Ret);
+    I0Ret.notifyMoveInstr(I2, I1->getIterator());
+    I2->moveBefore(I1);
+    EXPECT_EQ(I0Ret.top(), I0);
+    EXPECT_EQ(I0Ret.bottom(), Ret);
+
+    I2->moveAfter(I1);
+  }
+  {
+    // Move internal before top.
+    sandboxir::Interval<sandboxir::Instruction> I0Ret(I0, Ret);
+    I0Ret.notifyMoveInstr(I2, I0->getIterator());
+    I2->moveBefore(I0);
+    EXPECT_EQ(I0Ret.top(), I2);
+    EXPECT_EQ(I0Ret.bottom(), Ret);
+
+    I2->moveAfter(I1);
+  }
+  {
+    // Move internal to bottom.
+    sandboxir::Interval<sandboxir::Instruction> I0Ret(I0, Ret);
+    I0Ret.notifyMoveInstr(I2, BB->end());
+    I2->moveAfter(Ret);
+    EXPECT_EQ(I0Ret.top(), I0);
+    EXPECT_EQ(I0Ret.bottom(), I2);
+
+    I2->moveAfter(I1);
+  }
+  {
+    // Move bottom before internal.
+    sandboxir::Interval<sandboxir::Instruction> I0Ret(I0, Ret);
+    I0Ret.notifyMoveInstr(Ret, I2->getIterator());
+    Ret->moveBefore(I2);
+    EXPECT_EQ(I0Ret.top(), I0);
+    EXPECT_EQ(I0Ret.bottom(), I2);
+
+    Ret->moveAfter(I2);
+  }
+  {
+    // Move bottom before top.
+    sandboxir::Interval<sandboxir::Instruction> I0Ret(I0, Ret);
+    I0Ret.notifyMoveInstr(Ret, I0->getIterator());
+    Ret->moveBefore(I0);
+    EXPECT_EQ(I0Ret.top(), Ret);
+    EXPECT_EQ(I0Ret.bottom(), I2);
+
+    Ret->moveAfter(I2);
+  }
+}