[SandboxVec][BottomUpVec] Fix codegen when packing constants.

[vporpo]

Before this patch packing a bundle of constants would crash because getInsertPointAfterInstrs() expected instructions. This patch fixes this.

[llvmbot]

@llvm/pr-subscribers-vectorizers

@llvm/pr-subscribers-llvm-transforms

Author: vporpo (vporpo)

Changes

Before this patch packing a bundle of constants would crash because getInsertPointAfterInstrs() expected instructions. This patch fixes this.

---

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

3 Files Affected:

• (modified) llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Passes/BottomUpVec.h (+8-4)
• (modified) llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/BottomUpVec.cpp (+28-16)
• (added) llvm/test/Transforms/SandboxVectorizer/pack.ll (+16)

diff --git a/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Passes/BottomUpVec.h b/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Passes/BottomUpVec.h
index b463b8acf4c86e..147a86de4e34ec 100644
--- a/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Passes/BottomUpVec.h
+++ b/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Passes/BottomUpVec.h
@@ -38,16 +38,20 @@ class BottomUpVec final : public FunctionPass {
   /// collected during vectorization.
   void tryEraseDeadInstrs();
   /// Creates a shuffle instruction that shuffles \p VecOp according to \p Mask.
-  Value *createShuffle(Value *VecOp, const ShuffleMask &Mask);
-  /// Packs all elements of \p ToPack into a vector and returns that vector.
-  Value *createPack(ArrayRef<Value *> ToPack);
+  /// \p UserBB is the block of the user bundle.
+  Value *createShuffle(Value *VecOp, const ShuffleMask &Mask,
+                       BasicBlock *UserBB);
+  /// Packs all elements of \p ToPack into a vector and returns that vector. \p
+  /// UserBB is the block of the user bundle.
+  Value *createPack(ArrayRef<Value *> ToPack, BasicBlock *UserBB);
   /// After we create vectors for groups of instructions, the original
   /// instructions are potentially dead and may need to be removed. This
   /// function helps collect these instructions (along with the pointer operands
   /// for loads/stores) so that they can be cleaned up later.
   void collectPotentiallyDeadInstrs(ArrayRef<Value *> Bndl);
   /// Recursively try to vectorize \p Bndl and its operands.
-  Value *vectorizeRec(ArrayRef<Value *> Bndl, unsigned Depth);
+  Value *vectorizeRec(ArrayRef<Value *> Bndl, ArrayRef<Value *> UserBndl,
+                      unsigned Depth);
   /// Entry point for vectorization starting from \p Seeds.
   bool tryVectorize(ArrayRef<Value *> Seeds);
 
diff --git a/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/BottomUpVec.cpp b/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/BottomUpVec.cpp
index 8432b4c6c469ae..18c3b375c92a23 100644
--- a/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/BottomUpVec.cpp
+++ b/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/BottomUpVec.cpp
@@ -43,10 +43,15 @@ static SmallVector<Value *, 4> getOperand(ArrayRef<Value *> Bndl,
   return Operands;
 }
 
-static BasicBlock::iterator
-getInsertPointAfterInstrs(ArrayRef<Value *> Instrs) {
-  auto *BotI = VecUtils::getLowest(Instrs);
-  // If Bndl contains Arguments or Constants, use the beginning of the BB.
+/// \Returns the BB iterator after the lowest instruction in \p Vals, or the top
+/// of BB if no instruction found in \p Vals.
+static BasicBlock::iterator getInsertPointAfterInstrs(ArrayRef<Value *> Vals,
+                                                      BasicBlock *BB) {
+  auto *BotI = VecUtils::getLowest(Vals);
+  if (BotI == nullptr)
+    // We are using BB->begin() as the fallback insert point if `ToPack` did
+    // not contain instructions.
+    return BB->begin();
   return std::next(BotI->getIterator());
 }
 
@@ -61,7 +66,8 @@ Value *BottomUpVec::createVectorInstr(ArrayRef<Value *> Bndl,
     Type *ScalarTy = VecUtils::getElementType(Utils::getExpectedType(Bndl[0]));
     auto *VecTy = VecUtils::getWideType(ScalarTy, VecUtils::getNumLanes(Bndl));
 
-    BasicBlock::iterator WhereIt = getInsertPointAfterInstrs(Bndl);
+    BasicBlock::iterator WhereIt = getInsertPointAfterInstrs(
+        Bndl, cast<Instruction>(Bndl[0])->getParent());
 
     auto Opcode = cast<Instruction>(Bndl[0])->getOpcode();
     switch (Opcode) {
@@ -175,14 +181,15 @@ void BottomUpVec::tryEraseDeadInstrs() {
   DeadInstrCandidates.clear();
 }
 
-Value *BottomUpVec::createShuffle(Value *VecOp, const ShuffleMask &Mask) {
-  BasicBlock::iterator WhereIt = getInsertPointAfterInstrs({VecOp});
+Value *BottomUpVec::createShuffle(Value *VecOp, const ShuffleMask &Mask,
+                                  BasicBlock *UserBB) {
+  BasicBlock::iterator WhereIt = getInsertPointAfterInstrs({VecOp}, UserBB);
   return ShuffleVectorInst::create(VecOp, VecOp, Mask, WhereIt,
                                    VecOp->getContext(), "VShuf");
 }
 
-Value *BottomUpVec::createPack(ArrayRef<Value *> ToPack) {
-  BasicBlock::iterator WhereIt = getInsertPointAfterInstrs(ToPack);
+Value *BottomUpVec::createPack(ArrayRef<Value *> ToPack, BasicBlock *UserBB) {
+  BasicBlock::iterator WhereIt = getInsertPointAfterInstrs(ToPack, UserBB);
 
   Type *ScalarTy = VecUtils::getCommonScalarType(ToPack);
   unsigned Lanes = VecUtils::getNumLanes(ToPack);
@@ -258,8 +265,12 @@ void BottomUpVec::collectPotentiallyDeadInstrs(ArrayRef<Value *> Bndl) {
   }
 }
 
-Value *BottomUpVec::vectorizeRec(ArrayRef<Value *> Bndl, unsigned Depth) {
+Value *BottomUpVec::vectorizeRec(ArrayRef<Value *> Bndl,
+                                 ArrayRef<Value *> UserBndl, unsigned Depth) {
   Value *NewVec = nullptr;
+  auto *UserBB = !UserBndl.empty()
+                     ? cast<Instruction>(UserBndl.front())->getParent()
+                     : cast<Instruction>(Bndl[0])->getParent();
   const auto &LegalityRes = Legality->canVectorize(Bndl);
   switch (LegalityRes.getSubclassID()) {
   case LegalityResultID::Widen: {
@@ -272,7 +283,7 @@ Value *BottomUpVec::vectorizeRec(ArrayRef<Value *> Bndl, unsigned Depth) {
       break;
     case Instruction::Opcode::Store: {
       // Don't recurse towards the pointer operand.
-      auto *VecOp = vectorizeRec(getOperand(Bndl, 0), Depth + 1);
+      auto *VecOp = vectorizeRec(getOperand(Bndl, 0), Bndl, Depth + 1);
       VecOperands.push_back(VecOp);
       VecOperands.push_back(cast<StoreInst>(I)->getPointerOperand());
       break;
@@ -280,7 +291,7 @@ Value *BottomUpVec::vectorizeRec(ArrayRef<Value *> Bndl, unsigned Depth) {
     default:
       // Visit all operands.
       for (auto OpIdx : seq<unsigned>(I->getNumOperands())) {
-        auto *VecOp = vectorizeRec(getOperand(Bndl, OpIdx), Depth + 1);
+        auto *VecOp = vectorizeRec(getOperand(Bndl, OpIdx), Bndl, Depth + 1);
         VecOperands.push_back(VecOp);
       }
       break;
@@ -301,7 +312,7 @@ Value *BottomUpVec::vectorizeRec(ArrayRef<Value *> Bndl, unsigned Depth) {
     auto *VecOp = cast<DiamondReuseWithShuffle>(LegalityRes).getVector();
     const ShuffleMask &Mask =
         cast<DiamondReuseWithShuffle>(LegalityRes).getMask();
-    NewVec = createShuffle(VecOp, Mask);
+    NewVec = createShuffle(VecOp, Mask, UserBB);
     break;
   }
   case LegalityResultID::DiamondReuseMultiInput: {
@@ -315,7 +326,8 @@ Value *BottomUpVec::vectorizeRec(ArrayRef<Value *> Bndl, unsigned Depth) {
       if (auto *I = dyn_cast<Instruction>(ElmDescr.getValue()))
         DescrInstrs.push_back(I);
     }
-    auto WhereIt = getInsertPointAfterInstrs(DescrInstrs);
+    BasicBlock::iterator WhereIt =
+        getInsertPointAfterInstrs(DescrInstrs, UserBB);
 
     Value *LastV = PoisonValue::get(ResTy);
     for (auto [Lane, ElmDescr] : enumerate(Descr.getDescrs())) {
@@ -342,7 +354,7 @@ Value *BottomUpVec::vectorizeRec(ArrayRef<Value *> Bndl, unsigned Depth) {
     // If we can't vectorize the seeds then just return.
     if (Depth == 0)
       return nullptr;
-    NewVec = createPack(Bndl);
+    NewVec = createPack(Bndl, UserBB);
     break;
   }
   }
@@ -352,7 +364,7 @@ Value *BottomUpVec::vectorizeRec(ArrayRef<Value *> Bndl, unsigned Depth) {
 bool BottomUpVec::tryVectorize(ArrayRef<Value *> Bndl) {
   DeadInstrCandidates.clear();
   Legality->clear();
-  vectorizeRec(Bndl, /*Depth=*/0);
+  vectorizeRec(Bndl, {}, /*Depth=*/0);
   tryEraseDeadInstrs();
   return Change;
 }
diff --git a/llvm/test/Transforms/SandboxVectorizer/pack.ll b/llvm/test/Transforms/SandboxVectorizer/pack.ll
new file mode 100644
index 00000000000000..6607b31c021941
--- /dev/null
+++ b/llvm/test/Transforms/SandboxVectorizer/pack.ll
@@ -0,0 +1,16 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
+; RUN: opt -passes=sandbox-vectorizer -sbvec-vec-reg-bits=1024 -sbvec-allow-non-pow2 -sbvec-passes="bottom-up-vec<>" %s -S | FileCheck %s
+
+define void @pack_constants(ptr %ptr) {
+; CHECK-LABEL: define void @pack_constants(
+; CHECK-SAME: ptr [[PTR:%.*]]) {
+; CHECK-NEXT:    [[PTR0:%.*]] = getelementptr i8, ptr [[PTR]], i32 0
+; CHECK-NEXT:    store <2 x i8> <i8 0, i8 1>, ptr [[PTR0]], align 1
+; CHECK-NEXT:    ret void
+;
+  %ptr0 = getelementptr i8, ptr %ptr, i32 0
+  %ptr1 = getelementptr i8, ptr %ptr, i32 1
+  store i8 0, ptr %ptr0
+  store i8 1, ptr %ptr1
+  ret void
+}