[llvm] [SandboxVec] Simple Instruction Interval class (PR #108882)

[via llvm-commits]

https://github.com/vporpo updated https://github.com/llvm/llvm-project/pull/108882

>From cdfbd69009f79b6773f1ec31eba4ddd3f9da9943 Mon Sep 17 00:00:00 2001
From: Vasileios Porpodas <vporpodas at google.com>
Date: Mon, 16 Sep 2024 11:51:00 -0700
Subject: [PATCH] [SandboxVec] Simple Instruction Interval class

An InstrInterval is a range of instructions in a block.
The class will eventually have an API for set operations, like union,
intersection etc.
---
 .../SandboxVectorizer/InstrInterval.h         | 110 ++++++++++++++++++
 .../SandboxVectorizer/CMakeLists.txt          |   1 +
 .../SandboxVectorizer/InstrIntervalTest.cpp   |  68 +++++++++++
 3 files changed, 179 insertions(+)
 create mode 100644 llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/InstrInterval.h
 create mode 100644 llvm/unittests/Transforms/Vectorize/SandboxVectorizer/InstrIntervalTest.cpp

diff --git a/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/InstrInterval.h b/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/InstrInterval.h
new file mode 100644
index 00000000000000..ec1d126a6d3e3b
--- /dev/null
+++ b/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/InstrInterval.h
@@ -0,0 +1,110 @@
+//===- InstrInterval.h ------------------------------------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// The InstrInterval class is an interval of instructions in a block.
+// It provides an API for some basic operations on the interval, including some
+// simple set operations, like union, interseciton and others.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_TRANSFORMS_VECTORIZE_SANDBOXVECTORIZER_INSTRINTERVAL_H
+#define LLVM_TRANSFORMS_VECTORIZE_SANDBOXVECTORIZER_INSTRINTERVAL_H
+
+#include "llvm/SandboxIR/SandboxIR.h"
+#include <iterator>
+
+namespace llvm::sandboxir {
+
+/// A simple iterator for iterating the interval.
+template <typename DerefType, typename InstrIntervalType>
+class InstrIntervalIterator {
+  sandboxir::Instruction *I;
+  InstrIntervalType &R;
+
+public:
+  using difference_type = std::ptrdiff_t;
+  using value_type = sandboxir::Instruction;
+  using pointer = value_type *;
+  using reference = sandboxir::Instruction &;
+  using iterator_category = std::bidirectional_iterator_tag;
+
+  InstrIntervalIterator(sandboxir::Instruction *I, InstrIntervalType &R)
+      : I(I), R(R) {}
+  bool operator==(const InstrIntervalIterator &Other) const {
+    assert(&R == &Other.R && "Iterators belong to different regions!");
+    return Other.I == I;
+  }
+  bool operator!=(const InstrIntervalIterator &Other) const {
+    return !(*this == Other);
+  }
+  InstrIntervalIterator &operator++() {
+    assert(I != nullptr && "already at end()!");
+    I = I->getNextNode();
+    return *this;
+  }
+  InstrIntervalIterator operator++(int) {
+    auto ItCopy = *this;
+    ++*this;
+    return ItCopy;
+  }
+  InstrIntervalIterator &operator--() {
+    // `I` is nullptr for end() when ToI is the BB terminator.
+    I = I != nullptr ? I->getPrevNode() : R.ToI;
+    return *this;
+  }
+  InstrIntervalIterator operator--(int) {
+    auto ItCopy = *this;
+    --*this;
+    return ItCopy;
+  }
+  template <typename T =
+                std::enable_if<std::is_same<DerefType, Instruction *&>::value>>
+  sandboxir::Instruction &operator*() {
+    return *I;
+  }
+  DerefType operator*() const { return *I; }
+};
+
+class InstrInterval {
+  Instruction *FromI;
+  Instruction *ToI;
+
+public:
+  InstrInterval() : FromI(nullptr), ToI(nullptr) {}
+  InstrInterval(Instruction *FromI, Instruction *ToI) : FromI(FromI), ToI(ToI) {
+    assert(FromI->comesBefore(ToI) && "FromI should come before TopI!");
+  }
+  bool empty() const {
+    assert(((FromI == nullptr && ToI == nullptr) ||
+            (FromI != nullptr && ToI != nullptr)) &&
+           "Either none or both should be null");
+    return FromI == nullptr;
+  }
+  bool contains(Instruction *I) const {
+    if (empty())
+      return false;
+    return (FromI == I || FromI->comesBefore(I)) &&
+           (I == ToI || I->comesBefore(ToI));
+  }
+
+  using iterator =
+      InstrIntervalIterator<sandboxir::Instruction &, InstrInterval>;
+  using const_iterator = InstrIntervalIterator<const sandboxir::Instruction &,
+                                               const InstrInterval>;
+  iterator begin() { return iterator(FromI, *this); }
+  iterator end() {
+    return iterator(ToI != nullptr ? ToI->getNextNode() : nullptr, *this);
+  }
+  const_iterator begin() const { return const_iterator(FromI, *this); }
+  const_iterator end() const {
+    return const_iterator(ToI != nullptr ? ToI->getNextNode() : nullptr, *this);
+  }
+};
+} // namespace llvm::sandboxir
+
+#endif // LLVM_TRANSFORMS_VECTORIZE_SANDBOXVECTORIZER_INSTRINTERVAL_H
diff --git a/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/CMakeLists.txt b/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/CMakeLists.txt
index 2c7bf7d7e87541..b0ef71ba2114a1 100644
--- a/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/CMakeLists.txt
+++ b/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/CMakeLists.txt
@@ -9,5 +9,6 @@ set(LLVM_LINK_COMPONENTS
 
 add_llvm_unittest(SandboxVectorizerTests
   DependencyGraphTest.cpp
+  InstrIntervalTest.cpp
   LegalityTest.cpp
   )
diff --git a/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/InstrIntervalTest.cpp b/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/InstrIntervalTest.cpp
new file mode 100644
index 00000000000000..b8b8e4e982be56
--- /dev/null
+++ b/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/InstrIntervalTest.cpp
@@ -0,0 +1,68 @@
+//===- InstrIntervalTest.cpp ----------------------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/Transforms/Vectorize/SandboxVectorizer/InstrInterval.h"
+#include "llvm/AsmParser/Parser.h"
+#include "llvm/SandboxIR/SandboxIR.h"
+#include "llvm/Support/SourceMgr.h"
+#include "gmock/gmock-matchers.h"
+#include "gtest/gtest.h"
+
+using namespace llvm;
+
+struct InstrIntervalTest : public testing::Test {
+  LLVMContext C;
+  std::unique_ptr<Module> M;
+
+  void parseIR(LLVMContext &C, const char *IR) {
+    SMDiagnostic Err;
+    M = parseAssemblyString(IR, Err, C);
+    if (!M)
+      Err.print("InstrIntervalTest", errs());
+  }
+};
+
+TEST_F(InstrIntervalTest, Basic) {
+  parseIR(C, R"IR(
+define void @foo(i8 %v0) {
+  %add0 = add i8 %v0, %v0
+  %add1 = add i8 %v0, %v0
+  %add2 = 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++;
+
+  sandboxir::InstrInterval Interval(I0, Ret);
+  // Check empty().
+  EXPECT_FALSE(Interval.empty());
+  sandboxir::InstrInterval Empty;
+  EXPECT_TRUE(Empty.empty());
+  sandboxir::InstrInterval One(I0, I0);
+  EXPECT_FALSE(One.empty());
+  // Check contains().
+  for (auto &I : *BB) {
+    EXPECT_TRUE(Interval.contains(&I));
+    EXPECT_FALSE(Empty.contains(&I));
+  }
+  EXPECT_FALSE(One.contains(I1));
+  EXPECT_FALSE(One.contains(I2));
+  EXPECT_FALSE(One.contains(Ret));
+  // Check iterator.
+  auto BBIt = BB->begin();
+  for (auto &I : Interval)
+    EXPECT_EQ(&I, &*BBIt++);
+}