[SandboxVec][BottomUpVec] Disable crossing BBs (#124039)

Crossing BBs is not currently supported by the structures of the
vectorizer. This patch fixes instances where this was happening,
including:
- a walk of use-def operands that updates the UnscheduledSuccs counter,
- the dead instruction removal is now done per BB,
- the scheduler, which will reject bundles that cross BBs.

Author: vporpo

llvm/lib/Transforms/Vectorize/SandboxVectorizer/DependencyGraph.cpp

@@ -232,6 +232,9 @@ void DependencyGraph::setDefUseUnscheduledSuccs(
       auto *OpI = dyn_cast<Instruction>(Op);
       if (OpI == nullptr)
         continue;
+      // TODO: For now don't cross BBs.
+      if (OpI->getParent() != I.getParent())
+        continue;
       if (!NewInterval.contains(OpI))
         continue;
       auto *OpN = getNode(OpI);

llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/BottomUpVec.cpp

@@ -169,14 +169,19 @@ Value *BottomUpVec::createVectorInstr(ArrayRef<Value *> Bndl,
 }
 
 void BottomUpVec::tryEraseDeadInstrs() {
-  // Visiting the dead instructions bottom-to-top.
-  SmallVector<Instruction *> SortedDeadInstrCandidates(
-      DeadInstrCandidates.begin(), DeadInstrCandidates.end());
-  sort(SortedDeadInstrCandidates,
-       [](Instruction *I1, Instruction *I2) { return I1->comesBefore(I2); });
-  for (Instruction *I : reverse(SortedDeadInstrCandidates)) {
-    if (I->hasNUses(0))
-      I->eraseFromParent();
+  DenseMap<BasicBlock *, SmallVector<Instruction *>> SortedDeadInstrCandidates;
+  // The dead instrs could span BBs, so we need to collect and sort them per BB.
+  for (auto *DeadI : DeadInstrCandidates)
+    SortedDeadInstrCandidates[DeadI->getParent()].push_back(DeadI);
+  for (auto &Pair : SortedDeadInstrCandidates)
+    sort(Pair.second,
+         [](Instruction *I1, Instruction *I2) { return I1->comesBefore(I2); });
+  for (const auto &Pair : SortedDeadInstrCandidates) {
+    for (Instruction *I : reverse(Pair.second)) {
+      if (I->hasNUses(0))
+        // Erase the dead instructions bottom-to-top.
+        I->eraseFromParent();
+    }
   }
   DeadInstrCandidates.clear();
 }

llvm/lib/Transforms/Vectorize/SandboxVectorizer/Scheduler.cpp

@@ -206,6 +206,13 @@ bool Scheduler::trySchedule(ArrayRef<Instruction *> Instrs) {
     // We start scheduling at the bottom instr of Instrs.
     ScheduleTopItOpt = std::next(VecUtils::getLowest(Instrs)->getIterator());
 
+    // TODO: For now don't cross BBs.
+    if (!DAG.getInterval().empty()) {
+      auto *BB = DAG.getInterval().top()->getParent();
+      if (any_of(Instrs, [BB](auto *I) { return I->getParent() != BB; }))
+        return false;
+    }
+
     // Extend the DAG to include Instrs.
     Interval<Instruction> Extension = DAG.extend(Instrs);
     // Add nodes to ready list.

llvm/test/Transforms/SandboxVectorizer/cross_bbs.ll

@@ -0,0 +1,28 @@
+; 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 @cross_bbs(ptr %ptr) {
+; CHECK-LABEL: define void @cross_bbs(
+; CHECK-SAME: ptr [[PTR:%.*]]) {
+; CHECK-NEXT:    [[PTR0:%.*]] = getelementptr i8, ptr [[PTR]], i32 0
+; CHECK-NEXT:    [[PTR1:%.*]] = getelementptr i8, ptr [[PTR]], i32 1
+; CHECK-NEXT:    [[L0:%.*]] = load i8, ptr [[PTR0]], align 1
+; CHECK-NEXT:    [[L1:%.*]] = load i8, ptr [[PTR1]], align 1
+; CHECK-NEXT:    [[PACK:%.*]] = insertelement <2 x i8> poison, i8 [[L0]], i32 0
+; CHECK-NEXT:    [[PACK1:%.*]] = insertelement <2 x i8> [[PACK]], i8 [[L1]], i32 1
+; CHECK-NEXT:    br label %[[BB:.*]]
+; CHECK:       [[BB]]:
+; CHECK-NEXT:    store <2 x i8> [[PACK1]], ptr [[PTR0]], align 1
+; CHECK-NEXT:    ret void
+;
+  %ptr0 = getelementptr i8, ptr %ptr, i32 0
+  %ptr1 = getelementptr i8, ptr %ptr, i32 1
+  %l0 = load i8, ptr %ptr0
+  %l1 = load i8, ptr %ptr1
+  br label %bb
+
+bb:
+  store i8 %l0, ptr %ptr0
+  store i8 %l1, ptr %ptr1
+  ret void
+}