[SandboxVec][Interval] Implement getUnionInterval() and getSingleDiff()

[vporpo]

No description provided.

[llvmbot]

@llvm/pr-subscribers-llvm-transforms

Author: vporpo (vporpo)

Changes

---

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

2 Files Affected:

• (modified) llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Interval.h (+20)
• (modified) llvm/unittests/Transforms/Vectorize/SandboxVectorizer/IntervalTest.cpp (+61)

diff --git a/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Interval.h b/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Interval.h
index 8f25ad109f6a61..81f2ac9b3d0eb2 100644
--- a/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Interval.h
+++ b/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Interval.h
@@ -176,6 +176,26 @@ template <typename T> class Interval {
       Result.emplace_back(Intersection.To->getNextNode(), To);
     return Result;
   }
+  /// A difference that asserts that the result is a single interval.
+  Interval getSingleDiff(const Interval &Other) {
+    auto Diff = *this - Other;
+    assert(Diff.size() == 1 && "Expected a single interval!");
+    return Diff[0];
+  }
+  /// \Returns a single interval that spans across both this and \p Other.
+  // For example:
+  // |---|        this
+  //        |---| Other
+  // |----------| this->getUnionInterval(Other)
+  Interval getUnionInterval(const Interval &Other) {
+    if (empty())
+      return Other;
+    if (Other.empty())
+      return *this;
+    auto *NewFrom = From->comesBefore(Other.From) ? From : Other.From;
+    auto *NewTo = To->comesBefore(Other.To) ? Other.To : To;
+    return {NewFrom, NewTo};
+  }
 };
 
 } // namespace llvm::sandboxir
diff --git a/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/IntervalTest.cpp b/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/IntervalTest.cpp
index a697ce7727a9b0..7eaed91689fe8f 100644
--- a/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/IntervalTest.cpp
+++ b/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/IntervalTest.cpp
@@ -162,6 +162,9 @@ define void @foo(i8 %v0) {
     EXPECT_EQ(Diffs.size(), 1u);
     const sandboxir::Interval<sandboxir::Instruction> &Diff = Diffs[0];
     EXPECT_THAT(getPtrVec(Diff), testing::ElementsAre(I0, I1, I2, Ret));
+
+    // Check getSingleDiff().
+    EXPECT_EQ(I0Ret.getSingleDiff(Empty), Diff);
   }
   {
     // Check [] - [I0,Ret]
@@ -171,6 +174,9 @@ define void @foo(i8 %v0) {
     EXPECT_EQ(Diffs.size(), 1u);
     const sandboxir::Interval<sandboxir::Instruction> &Diff = Diffs[0];
     EXPECT_TRUE(Diff.empty());
+
+    // Check getSingleDiff().
+    EXPECT_EQ(Empty.getSingleDiff(I0Ret), Diff);
   }
   {
     // Check [I0,Ret] - [I0].
@@ -180,6 +186,9 @@ define void @foo(i8 %v0) {
     EXPECT_EQ(Diffs.size(), 1u);
     const sandboxir::Interval<sandboxir::Instruction> &Diff = Diffs[0];
     EXPECT_THAT(getPtrVec(Diff), testing::ElementsAre(I1, I2, Ret));
+
+    // Check getSingleDiff().
+    EXPECT_EQ(I0Ret.getSingleDiff(I0I0), Diff);
   }
   {
     // Check [I0,Ret] - [I1].
@@ -191,6 +200,9 @@ define void @foo(i8 %v0) {
     EXPECT_THAT(getPtrVec(Diff0), testing::ElementsAre(I0));
     const sandboxir::Interval<sandboxir::Instruction> &Diff1 = Diffs[1];
     EXPECT_THAT(getPtrVec(Diff1), testing::ElementsAre(I2, Ret));
+
+    // Check getSingleDiff().
+    EXPECT_DEATH(I0Ret.getSingleDiff(I1I1), ".*single.*");
   }
 }
 
@@ -249,3 +261,52 @@ define void @foo(i8 %v0) {
     EXPECT_THAT(getPtrVec(Intersection), testing::ElementsAre(I1));
   }
 }
+
+TEST_F(IntervalTest, UnionInterval) {
+  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++;
+  [[maybe_unused]] auto *I2 = &*It++;
+  auto *Ret = &*It++;
+
+  {
+    // Check [I0] unionInterval [I2].
+    sandboxir::Interval<sandboxir::Instruction> I0I0(I0, I0);
+    sandboxir::Interval<sandboxir::Instruction> I2I2(I2, I2);
+    auto SingleUnion = I0I0.getUnionInterval(I2I2);
+    EXPECT_THAT(getPtrVec(SingleUnion), testing::ElementsAre(I0, I1, I2));
+  }
+  {
+    // Check [I0] unionInterval Empty.
+    sandboxir::Interval<sandboxir::Instruction> I0I0(I0, I0);
+    sandboxir::Interval<sandboxir::Instruction> Empty;
+    auto SingleUnion = I0I0.getUnionInterval(Empty);
+    EXPECT_THAT(getPtrVec(SingleUnion), testing::ElementsAre(I0));
+  }
+  {
+    // Check [I0,I1] unionInterval [I1].
+    sandboxir::Interval<sandboxir::Instruction> I0I1(I0, I1);
+    sandboxir::Interval<sandboxir::Instruction> I1I1(I1, I1);
+    auto SingleUnion = I0I1.getUnionInterval(I1I1);
+    EXPECT_THAT(getPtrVec(SingleUnion), testing::ElementsAre(I0, I1));
+  }
+  {
+    // Check [I2,Ret] unionInterval [I0].
+    sandboxir::Interval<sandboxir::Instruction> I2Ret(I2, Ret);
+    sandboxir::Interval<sandboxir::Instruction> I0I0(I0, I0);
+    auto SingleUnion = I2Ret.getUnionInterval(I0I0);
+    EXPECT_THAT(getPtrVec(SingleUnion), testing::ElementsAre(I0, I1, I2, Ret));
+  }
+}

[Sterling-Augustine]

one small nit

[Sterling-Augustine]

\Returns a

[aeubanks]

this only asserts in asserts-enabled builds, perhaps drop the "that asserts"?

[vporpo]

I will change it to:
/// \Returns the interval difference 'this - Other'. This will crash in the Debug build if the result is not a single interval.
Wdyt ?

[alinas]

I'm not clear if there should be another union method API returning two intervals when there is no overlap (minimal instruction coverage vs maximal).
This looks good for this purpose.

[vporpo]

I'm not clear if there should be another union method API returning two intervals when there is no overlap (minimal instruction coverage vs maximal).

We can add one once it's needed.