[llvm] [SandboxVectorizer] Define SeedBundle: a set of instructions to be vectorized (PR #110696)

[via llvm-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-llvm-transforms

Author: None (Sterling-Augustine)

<details>
<summary>Changes</summary>

Seed collection will assemble instructions to be vectorized into SeedBundles. This data structure is not intended to be used directly, but will be the basis for load bundles, store bundles, and so on.

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


5 Files Affected:

- (added) llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/SeedCollector.h (+136) 
- (modified) llvm/lib/Transforms/Vectorize/CMakeLists.txt (+1) 
- (added) llvm/lib/Transforms/Vectorize/SandboxVectorizer/SeedCollector.cpp (+63) 
- (modified) llvm/unittests/Transforms/Vectorize/SandboxVectorizer/CMakeLists.txt (+2-1) 
- (added) llvm/unittests/Transforms/Vectorize/SandboxVectorizer/SeedCollectorTest.cpp (+101) 


``````````diff
diff --git a/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/SeedCollector.h b/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/SeedCollector.h
new file mode 100644
index 00000000000000..38088f30c8af61
--- /dev/null
+++ b/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/SeedCollector.h
@@ -0,0 +1,136 @@
+//===- SeedCollector.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
+//
+//===----------------------------------------------------------------------===//
+// This file contains the mechanism for collecting the seed instructions that
+// are used as starting points for forming the vectorization graph.
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_TRANSFORMS_VECTORIZE_SANDBOXVECTORIZER_SEEDCOLLECTOR_H
+#define LLVM_TRANSFORMS_VECTORIZE_SANDBOXVECTORIZER_SEEDCOLLECTOR_H
+
+#include "llvm/ADT/BitVector.h"
+#include "llvm/ADT/iterator_range.h"
+#include "llvm/Analysis/ScalarEvolution.h"
+#include "llvm/SandboxIR/Instruction.h"
+#include "llvm/SandboxIR/Utils.h"
+#include "llvm/SandboxIR/Value.h"
+#include <iterator>
+#include <memory>
+
+namespace llvm::sandboxir {
+class Instruction;
+class StoreInst;
+class BasicBlock;
+
+/// An ordered set of Instructions that can be vectorized.
+class SeedBundle {
+public:
+  using SeedList = SmallVector<sandboxir::Instruction *>;
+  /// Initialize a bundle with \p I.
+  explicit SeedBundle(sandboxir::Instruction *I, const DataLayout &DL) {
+    insertAt(begin(), I, DL);
+  }
+  explicit SeedBundle(SeedList &&L, const DataLayout &DL)
+      : Seeds(std::move(L)) {
+    for (auto &S : Seeds) {
+      NumUnusedBits += sandboxir::Utils::getNumBits(S, DL);
+    }
+  }
+  /// No need to allow copies.
+  SeedBundle(const SeedBundle &) = delete;
+  SeedBundle &operator=(const SeedBundle &) = delete;
+  virtual ~SeedBundle() {}
+
+  using iterator = SeedList::iterator;
+  using const_iterator = SeedList::const_iterator;
+  iterator begin() { return Seeds.begin(); }
+  iterator end() { return Seeds.end(); }
+  const_iterator begin() const { return Seeds.begin(); }
+  const_iterator end() const { return Seeds.end(); }
+
+  sandboxir::Instruction *operator[](unsigned Idx) const { return Seeds[Idx]; }
+
+  /// Insert \p I into position \p P. Clients should choose Pos
+  /// by symbol, symbol-offset, and program order (which depends if scheduling
+  /// bottom-up or top-down).
+  void insertAt(iterator Pos, sandboxir::Instruction *I, const DataLayout &DL) {
+#ifdef EXPENSIVE_CHECKS
+    for (auto Itr : Seeds) {
+      assert(*Itr != I && "Attempt to insert an instruction twice.");
+    }
+#endif
+    Seeds.insert(Pos, I);
+    NumUnusedBits += sandboxir::Utils::getNumBits(I, DL);
+  }
+
+  unsigned getFirstUnusedElementIdx() const {
+    for (unsigned ElmIdx : seq<unsigned>(0, Seeds.size()))
+      if (!isUsed(ElmIdx))
+        return ElmIdx;
+    return Seeds.size();
+  }
+  /// Marks elements as 'used' so that we skip them in `getSlice()`.
+  void setUsed(unsigned ElementIdx, const DataLayout &DL, unsigned Sz = 1,
+               bool VerifyUnused = true) {
+    if (ElementIdx + Sz >= UsedLanes.size())
+      UsedLanes.resize(ElementIdx + Sz);
+    for (unsigned Idx : seq<unsigned>(ElementIdx, ElementIdx + Sz)) {
+      assert((!VerifyUnused || !UsedLanes.test(Idx)) &&
+             "Already marked as used!");
+      UsedLanes.set(Idx);
+      UsedLaneCount++;
+    }
+    NumUnusedBits -= sandboxir::Utils::getNumBits(Seeds[ElementIdx], DL);
+  }
+
+  void setUsed(sandboxir::Instruction *V, const DataLayout &DL) {
+    auto It = std::find(begin(), end(), V);
+    assert(It != end() && "V not in the bundle!");
+    auto Idx = It - begin();
+    setUsed(Idx, DL, 1, /*VerifyUnused=*/false);
+  }
+  bool isUsed(unsigned Element) const {
+    return Element >= UsedLanes.size() ? false : UsedLanes.test(Element);
+  }
+  bool allUsed() const { return UsedLaneCount == Seeds.size(); }
+  unsigned getNumUnusedBits() const { return NumUnusedBits; }
+
+  /// \Returns a slice of seed elements, starting at the element \p StartIdx,
+  /// with a total size <= \p MaxVecRegBits. If \p ForcePowOf2 is true, then the
+  /// returned slice should have a total number of bits that is a power of 2.
+  MutableArrayRef<SeedList> getSlice(unsigned StartIdx, unsigned MaxVecRegBits,
+                                     bool ForcePowOf2, const DataLayout &DL);
+
+protected:
+  SeedList Seeds;
+  /// The lanes that we have already vectorized.
+  BitVector UsedLanes;
+  /// Tracks used lanes for constant-time accessor.
+  unsigned UsedLaneCount = 0;
+  /// Tracks the remaining bits available to vectorize
+  unsigned NumUnusedBits = 0;
+
+public:
+#ifndef NDEBUG
+  void dump(raw_ostream &OS) const {
+    for (auto [ElmIdx, I] : enumerate(*this)) {
+      OS.indent(2) << ElmIdx << ". ";
+      if (isUsed(ElmIdx))
+        OS << "[USED]";
+      else
+        OS << *I;
+      OS << "\n";
+    }
+  }
+  LLVM_DUMP_METHOD void dump() const {
+    dump(dbgs());
+    dbgs() << "\n";
+  }
+#endif // NDEBUG
+};
+} // namespace llvm::sandboxir
+#endif // LLVM_TRANSFORMS_VECTORIZE_SANDBOXVECTORIZER_SEEDCOLLECTOR_H
diff --git a/llvm/lib/Transforms/Vectorize/CMakeLists.txt b/llvm/lib/Transforms/Vectorize/CMakeLists.txt
index eeff4a9f6a8bae..887c2089c5a520 100644
--- a/llvm/lib/Transforms/Vectorize/CMakeLists.txt
+++ b/llvm/lib/Transforms/Vectorize/CMakeLists.txt
@@ -6,6 +6,7 @@ add_llvm_component_library(LLVMVectorize
   SandboxVectorizer/DependencyGraph.cpp
   SandboxVectorizer/Passes/BottomUpVec.cpp
   SandboxVectorizer/SandboxVectorizer.cpp
+  SandboxVectorizer/SeedCollector.cpp
   SLPVectorizer.cpp
   Vectorize.cpp
   VectorCombine.cpp
diff --git a/llvm/lib/Transforms/Vectorize/SandboxVectorizer/SeedCollector.cpp b/llvm/lib/Transforms/Vectorize/SandboxVectorizer/SeedCollector.cpp
new file mode 100644
index 00000000000000..668fea6904f866
--- /dev/null
+++ b/llvm/lib/Transforms/Vectorize/SandboxVectorizer/SeedCollector.cpp
@@ -0,0 +1,63 @@
+//===- SeedCollection.cpp  -0000000----------------------------------------===//
+//
+// 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/SeedCollector.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/Analysis/LoopAccessAnalysis.h"
+#include "llvm/Analysis/ValueTracking.h"
+#include "llvm/IR/BasicBlock.h"
+#include "llvm/IR/Type.h"
+#include "llvm/SandboxIR/Instruction.h"
+#include "llvm/SandboxIR/Utils.h"
+#include "llvm/Support/Debug.h"
+#include <span>
+
+using namespace llvm;
+
+MutableArrayRef<sandboxir::SeedBundle::SeedList>
+sandboxir::SeedBundle::getSlice(unsigned StartIdx, unsigned MaxVecRegBits,
+                                bool ForcePowerOf2, const DataLayout &DL) {
+  // Use uint32_t for counts to make it clear we are also using the proper
+  // isPowerOf2_[32|64].
+
+  // Count both the bits and the elements of the slice we are about to build.
+  // The bits tell us whether this is a legal slice (that is <= MaxVecRegBits),
+  // and the num of elements help us do the actual slicing.
+  uint32_t BitsSum = 0;
+  // As we are collecting slice elements we may go over the limit, so we need to
+  // remember the last legal one. This is used for the creation of the slice.
+  uint32_t LastGoodBitsSum = 0;
+  uint32_t LastGoodNumSliceElements = 0;
+  // Skip any used elements (which have already been handled) and all below
+  // `StartIdx`.
+  assert(StartIdx >= getFirstUnusedElementIdx() &&
+         "Expected unused at StartIdx");
+  uint32_t FirstGoodElementIdx = StartIdx;
+  // Go through elements starting at FirstGoodElementIdx.
+  for (auto [ElementCnt, S] : enumerate(make_range(
+           std::next(Seeds.begin(), FirstGoodElementIdx), Seeds.end()))) {
+    // Stop if we found a used element.
+    if (isUsed(FirstGoodElementIdx + ElementCnt))
+      break;
+    BitsSum += sandboxir::Utils::getNumBits(S, DL);
+    // Stop if the bits sum is over the limit.
+    if (BitsSum > MaxVecRegBits)
+      break;
+    // If forcing a power-of-2 bit-size we check if this bit size is accepted.
+    if (ForcePowerOf2 && !isPowerOf2_32(BitsSum))
+      continue;
+    LastGoodBitsSum = BitsSum;
+    LastGoodNumSliceElements = ElementCnt + 1;
+  }
+  if (LastGoodNumSliceElements < 2)
+    return {};
+  if (LastGoodBitsSum == 0)
+    return {};
+  return MutableArrayRef<sandboxir::SeedBundle::SeedList>(
+      &Seeds + FirstGoodElementIdx, LastGoodNumSliceElements);
+}
diff --git a/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/CMakeLists.txt b/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/CMakeLists.txt
index 9f1a3409c0c394..dcd7232db5f60c 100644
--- a/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/CMakeLists.txt
+++ b/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/CMakeLists.txt
@@ -11,4 +11,5 @@ add_llvm_unittest(SandboxVectorizerTests
   DependencyGraphTest.cpp
   IntervalTest.cpp
   LegalityTest.cpp
-  )
+  SeedCollectorTest.cpp	
+)
diff --git a/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/SeedCollectorTest.cpp b/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/SeedCollectorTest.cpp
new file mode 100644
index 00000000000000..6756ce944f9a91
--- /dev/null
+++ b/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/SeedCollectorTest.cpp
@@ -0,0 +1,101 @@
+//===- SeedCollectorTest.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/SeedCollector.h"
+#include "llvm/AsmParser/Parser.h"
+#include "llvm/SandboxIR/Function.h"
+#include "llvm/SandboxIR/Instruction.h"
+#include "llvm/Support/SourceMgr.h"
+#include "gtest/gtest.h"
+
+using namespace llvm;
+
+struct SeedBundleTest : 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("LegalityTest", errs());
+  }
+};
+
+TEST_F(SeedBundleTest, SeedBundle) {
+  parseIR(C, R"IR(
+define void @foo(float %v0, float %v1) {
+bb:
+  %add0 = fadd float %v0, %v1
+  %add1 = fadd float %v0, %v1
+  %add2 = fadd float %v0, %v1
+  %add3 = fadd float %v0, %v1
+  %add4 = fadd float %v0, %v1
+  ret void
+}
+)IR");
+  Function &LLVMF = *M->getFunction("foo");
+  sandboxir::Context Ctx(C);
+  auto &F = *Ctx.createFunction(&LLVMF);
+  DataLayout DL(M->getDataLayout());
+  auto *BB = &*F.begin();
+  auto It = BB->begin();
+  auto *I0 = &*It++;
+  auto *I1 = &*It++;
+  // Assume test instructions are identical in the number of bits.
+  const unsigned kFloatBits = sandboxir::Utils::getNumBits(I0, DL);
+  // Constructor
+  sandboxir::SeedBundle SBO(I0, DL);
+  EXPECT_EQ(*SBO.begin(), I0);
+  // getNumUnusedBits after constructor
+  EXPECT_EQ(SBO.getNumUnusedBits(), kFloatBits);
+  // setUsed
+  SBO.setUsed(I0, DL);
+  // allUsed
+  EXPECT_TRUE(SBO.allUsed());
+  // isUsed
+  EXPECT_TRUE(SBO.isUsed(0));
+  // getNumUnusedBits after setUsed
+  EXPECT_EQ(SBO.getNumUnusedBits(), 0u);
+  // insertAt
+  SBO.insertAt(SBO.end(), I1, DL);
+  EXPECT_NE(*SBO.begin(), I1);
+  // getNumUnusedBits after insertAt
+  EXPECT_EQ(SBO.getNumUnusedBits(), kFloatBits);
+  // allUsed
+  EXPECT_FALSE(SBO.allUsed());
+  // getFirstUnusedElement
+  EXPECT_EQ(SBO.getFirstUnusedElementIdx(), 1u);
+
+  sandboxir::SeedBundle::SeedList Seeds;
+  It = BB->begin();
+  Seeds.push_back(&*It++);
+  Seeds.push_back(&*It++);
+  Seeds.push_back(&*It++);
+  Seeds.push_back(&*It++);
+  Seeds.push_back(&*It++);
+  // Constructor
+  sandboxir::SeedBundle SB1(std::move(Seeds), DL);
+  // getNumUnusedBits after constructor
+  EXPECT_EQ(SB1.getNumUnusedBits(), 5 * kFloatBits);
+  // setUsed with index
+  SB1.setUsed(1, DL);
+  // getFirstUnusedElementIdx
+  EXPECT_EQ(SB1.getFirstUnusedElementIdx(), 0u);
+  EXPECT_EQ(SB1.getNumUnusedBits(), 4 * kFloatBits);
+  SB1.setUsed(unsigned(0), DL);
+  // getFirstUnusedElementIdx not at end
+  EXPECT_EQ(SB1.getFirstUnusedElementIdx(), 2u);
+  // getSlice
+  auto Slice0 = SB1.getSlice(2, /* MaxVecRegBits */ kFloatBits * 2,
+                             /* ForcePowerOf2 */ true, DL);
+  EXPECT_EQ(Slice0.size(), 2u);
+  auto Slice1 = SB1.getSlice(2, /* MaxVecRegBits */ kFloatBits * 3,
+                             /* ForcePowerOf2 */ false, DL);
+  EXPECT_EQ(Slice1.size(), 3u);
+}

``````````

</details>


https://github.com/llvm/llvm-project/pull/110696